wip

Author: Amir-19

flashinfer/cute_dsl/blockscaled_gemm.py

@@ -2625,6 +2625,60 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
 
         return is_valid
 
+    @staticmethod
+    def compute_barrier_flag_size(
+        m: int,
+        n: int,
+        l: int,
+        mma_tiler_mn: Tuple[int, int],
+        cluster_shape_mn: Tuple[int, int],
+        sm_count: int,
+    ) -> int:
+        """
+        Compute the required size for barrier flag tensors used in all-reduce synchronization.
+
+        The barrier flags are used for:
+        1. Per-tile synchronization during the all-reduce phase
+        2. Final inter-GPU synchronization barrier
+
+        :param m: Number of rows in the output matrix
+        :type m: int
+        :param n: Number of columns in the output matrix
+        :type n: int
+        :param l: Batch size
+        :type l: int
+        :param mma_tiler_mn: Shape of the MMA tiler (M, N)
+        :type mma_tiler_mn: Tuple[int, int]
+        :param cluster_shape_mn: Cluster dimensions (M, N)
+        :type cluster_shape_mn: Tuple[int, int]
+        :param sm_count: Number of SMs available
+        :type sm_count: int
+
+        :return: Total number of barrier flags needed
+        :rtype: int
+        """
+        # Calculate CTA tile shape accounting for 2-CTA instructions
+        use_2cta_instrs = mma_tiler_mn[0] == 256
+        cta_tile_shape_m = mma_tiler_mn[0] // (2 if use_2cta_instrs else 1)
+        cta_tile_shape_n = mma_tiler_mn[1]
+
+        # Calculate number of tiles per batch
+        num_tiles_m = (m + cta_tile_shape_m - 1) // cta_tile_shape_m
+        num_tiles_n = (n + cta_tile_shape_n - 1) // cta_tile_shape_n
+        num_tiles_per_batch = num_tiles_m * num_tiles_n
+
+        # Calculate number of clusters per batch
+        cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
+        num_ctas_per_tile = cluster_size
+
+        # Total tiles across all batches and clusters
+        num_tiles = num_tiles_per_batch * l * num_ctas_per_tile
+
+        # Add extra space for final barrier (one per SM)
+        total_barrier_size = num_tiles + sm_count
+
+        return total_barrier_size
+
     @staticmethod
     def can_implement(
         ab_dtype: Type[cutlass.Numeric],
@@ -2898,6 +2952,18 @@ def __call__(
         barrier_flag_mc_ptr: Optional[cute.Pointer],
         current_stream: cuda.CUstream,
     ):
+        if cutlass.const_expr(self._all_reduce != "none"):
+            barrier_flag_size = Sm100BlockScaledPersistentDenseGemmKernel.compute_barrier_flag_size(
+                self._m,
+                self._n,
+                self._l,
+                self._mma_tiler_mn,
+                self._cluster_shape_mn,
+                self._max_active_clusters,
+            )
+        else:
+            barrier_flag_size = 1  # Dummy size when not used
+
         a_tensor = cute.make_tensor(
             a_ptr,
             layout=cute.make_ordered_layout(
@@ -2931,11 +2997,11 @@ def __call__(
         # on
         barrier_flag_tensor = cute.make_tensor(
             barrier_flag_ptr,
-            layout=cute.make_ordered_layout((404,), order=(0,)),
+            layout=cute.make_ordered_layout((barrier_flag_size,), order=(0,)),
         ) if barrier_flag_ptr is not None else None
         barrier_flag_mc_tensor = cute.make_tensor(
             barrier_flag_mc_ptr,
-            layout=cute.make_ordered_layout((404,), order=(0,)),
+            layout=cute.make_ordered_layout((barrier_flag_size,), order=(0,)),
         ) if barrier_flag_mc_ptr is not None else None
 
         # calculate sf_tensor shape and order

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -81,23 +81,29 @@ def create_mc_tensor(torch_tensor_cpu, dtype, is_dynamic_layout=True):
     )
     return cute_tensor, cute_tensor_mc, torch_tensor_gpu, torch_tensor_mc
 
-def create_barrier_flags(m, n, l, mma_tiler_mn):
-        # NOTE: use_2cta_instrs from blockedscaled_gemm logic
-        use_2cta_instrs = mma_tiler_mn[0] == 256
-        cta_tile_shape_mn = (
-            mma_tiler_mn[0] // (2 if use_2cta_instrs else 1),
-            mma_tiler_mn[1],
+def create_barrier_flags(m, n, l, mma_tiler_mn, cluster_shape_mn, sm_count):
+        barrier_size = Sm100BlockScaledPersistentDenseGemmKernel.compute_barrier_flag_size(
+            m, n, l, mma_tiler_mn, cluster_shape_mn, sm_count
         )
-        problem_shape_ntile_mn = (m // cta_tile_shape_mn[0], n // cta_tile_shape_mn[1])
-        num_tiles_per_batch = problem_shape_ntile_mn[0] * problem_shape_ntile_mn[1]
-        num_tiles = num_tiles_per_batch * l
-        num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count
-
+        print("LOOK HERE",(barrier_size,))
+        # NOTE: use_2cta_instrs from blockedscaled_gemm logic
+        
+        # use_2cta_instrs = mma_tiler_mn[0] == 256
+        # cta_tile_shape_mn = (
+        #     mma_tiler_mn[0] // (2 if use_2cta_instrs else 1),
+        #     mma_tiler_mn[1],
+        # )
+        # problem_shape_ntile_mn = (m // cta_tile_shape_mn[0], n // cta_tile_shape_mn[1])
+        # num_tiles_per_batch = problem_shape_ntile_mn[0] * problem_shape_ntile_mn[1]
+        # num_tiles = num_tiles_per_batch * l
+        # num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count
         # +num_sms for final barrier
+        #num_tiles + num_sms
+
         barrier_flag = symm_mem.empty(
-            (num_tiles + num_sms,), device="cuda", dtype=torch.int32
+            (barrier_size,), device="cuda", dtype=torch.int32
         )
-        print("LOOK HERE",(num_tiles + num_sms,))
+        
         barrier_flag.fill_(0)
         symm = symm_mem.rendezvous(barrier_flag, group=dist.group.WORLD.group_name)
         barrier_flag_mc_ptr = symm.multicast_ptr
@@ -219,6 +225,8 @@ def run_blockscaled_gemm_all_reduce_python_interface(
         n,
         l,
         mma_tiler_mn,
+        cluster_shape_mn,
+        sm_count,
     )
     # for deepgemm-like python interface
     if ab_dtype == "float4_e2m1fn":