precommit

Author: Amir-19

flashinfer/cute_dsl/blockscaled_gemm.py

@@ -26,7 +26,8 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Optional, Tuple, Type, Union
+import functools
+from typing import Callable, Optional, Tuple, Type, Union, List
 
 import cuda.bindings.driver as cuda
 import cutlass
@@ -37,26 +38,21 @@
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.utils.blockscaled_layout as blockscaled_utils
 import cutlass.utils.distributed_helpers as distributed_helpers
-import torch
-import functools
 from cutlass._mlir import ir
 from cutlass.cute.nvgpu import cpasync, tcgen05
 from cutlass.cute.runtime import from_dlpack
-
 from cutlass.cutlass_dsl import (
-    Int32,
-    Int64,
-    Uint8,
-    Uint64,
     T,
     Integer,
     dsl_user_op,
     extract_mlir_values,
     new_from_mlir_values,
 )
-
 from cutlass.cute.typing import (
     Int32,
+    Int64,
+    Uint8,
+    Uint64,
     Float16,
     BFloat16,
     Float32,
@@ -65,10 +61,11 @@
     Tensor,
 )
 from cutlass._mlir.dialects import llvm
-from flashinfer.utils import get_compute_capability
 from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo
+import torch
+
+from flashinfer.utils import get_compute_capability
 from .utils import get_cutlass_dtype, cutlass_to_torch_dtype, get_num_sm, make_ptr
-from typing import Callable, List
 
 
 sizeof_i32 = 4
@@ -1865,7 +1862,6 @@ def kernel(
                 # Allreduce
                 #
                 if cutlass.const_expr(self.all_reduce == "two_shot"):
-                    
                     tile_id = Int32(
                         tile_sched._current_work_linear_idx
                         * cute.size(self.cluster_shape_mn)
@@ -2950,13 +2946,15 @@ def __call__(
         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,
+            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
@@ -2982,21 +2980,33 @@ def __call__(
                 order=(0, 1, 2) if self._c_major == "m" else (1, 0, 2),
             ),
         )
-        c_mc_tensor = cute.make_tensor(
-            c_mc_ptr,
-            layout=cute.make_ordered_layout(
-                (self._m, self._n, self._l),
-                order=(0, 1, 2) if self._c_major == "m" else (1, 0, 2),
-            ),
-        ) if c_mc_ptr is not None else None
-        barrier_flag_tensor = cute.make_tensor(
-            barrier_flag_ptr,
-            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((barrier_flag_size,), order=(0,)),
-        ) if barrier_flag_mc_ptr is not None else None
+        c_mc_tensor = (
+            cute.make_tensor(
+                c_mc_ptr,
+                layout=cute.make_ordered_layout(
+                    (self._m, self._n, self._l),
+                    order=(0, 1, 2) if self._c_major == "m" else (1, 0, 2),
+                ),
+            )
+            if c_mc_ptr is not None
+            else None
+        )
+        barrier_flag_tensor = (
+            cute.make_tensor(
+                barrier_flag_ptr,
+                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((barrier_flag_size,), order=(0,)),
+            )
+            if barrier_flag_mc_ptr is not None
+            else None
+        )
 
         # calculate sf_tensor shape and order
         def ceil_div(a, b):
@@ -3154,7 +3164,6 @@ def get_cute_pointers(
                 c_mc_data_ptr,
                 barrier_flag_data_ptr,
                 barrier_flag_mc_data_ptr,
-
             ) = (
                 a_tensor_gpu.data_ptr(),
                 b_tensor_gpu.data_ptr(),
@@ -3168,7 +3177,9 @@ def get_cute_pointers(
                 alpha_tensor_gpu.data_ptr() if alpha_tensor_gpu is not None else None,
                 c_mc_gpu.data_ptr() if c_mc_gpu is not None else None,
                 barrier_flag_gpu.data_ptr() if barrier_flag_gpu is not None else None,
-                barrier_flag_mc_gpu.data_ptr() if barrier_flag_mc_gpu is not None else None,
+                barrier_flag_mc_gpu.data_ptr()
+                if barrier_flag_mc_gpu is not None
+                else None,
             )
 
         a_ptr = make_ptr(

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -1,16 +1,12 @@
-import logging
 import multiprocessing as mp
 import pytest
 import socket
-from typing import Any, Tuple, Type
+from typing import Any, Tuple
 
-from cuda import cuda
 import cutlass
 import cutlass.cute as cute
-import cutlass.cute.testing as testing
 from cutlass.cute.runtime import from_dlpack
 import cutlass.torch as cutlass_torch
-import cutlass.utils as utils
 
 import torch
 import torch.distributed as dist
@@ -130,6 +126,7 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     enable_dst_signals: int,
     all_reduce: str,
     rank: int,
+    world_size: int,
 ):
     torch.manual_seed(42)
     device = torch.device("cuda", rank)
@@ -395,6 +392,7 @@ def _run_correctness_worker(
             enable_dst_signals=enable_dst_signals,
             all_reduce=all_reduce,
             rank=rank,
+            world_size=world_size,
         )
     except Exception as e:
         print(f"Rank {rank_id}: Exception during test: {e}")
@@ -430,9 +428,9 @@ def multi_process_parallel(
 
     for i in range(world_size):
         procs[i].join()
-        assert (
-            procs[i].exitcode == 0
-        ), f"Process {i} failed with exit code {procs[i].exitcode}"
+        assert procs[i].exitcode == 0, (
+            f"Process {i} failed with exit code {procs[i].exitcode}"
+        )
 
 
 @pytest.mark.skipif(