wip

Author: Amir-19

flashinfer/cute_dsl/blockscaled_gemm.py

@@ -2755,8 +2755,9 @@ def can_implement(
 
         # check for all reduce constraints
         # TODO(asamani): expand the logic for mnnvl support
-        if torch.distributed.get_world_size() not in [2, 4, 8] and all_reduce != "none":
-            can_implement = False
+        if all_reduce != "none":
+            if torch.distributed.get_world_size() not in [2, 4, 8]:
+                can_implement = False
         return can_implement
 
 

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -29,50 +29,50 @@
 
 def create_mc_tensor(torch_tensor_cpu, dtype, is_dynamic_layout=True):
     m, n, l = torch_tensor_cpu.shape
-    
+
     # Create flat symm_mem buffer
     total_elements = m * n * l
     torch_symm_flat = symm_mem.empty(
         (total_elements,), device="cuda", dtype=torch_tensor_cpu.dtype
     )
-    
+
     # Reshape to match input's stride pattern using as_strided
     torch_symm_tensor = torch_symm_flat.as_strided(
-        size=torch_tensor_cpu.shape,
-        stride=torch_tensor_cpu.stride()
+        size=torch_tensor_cpu.shape, stride=torch_tensor_cpu.stride()
     )
     torch_symm_tensor.copy_(torch_tensor_cpu)
-    
+
     symm = symm_mem.rendezvous(torch_symm_flat, group=dist.group.WORLD.group_name)
     mc_ptr = symm.multicast_ptr
-    
+
     # Create MC tensor with same stride
-    torch_tensor_mc_flat = cutlass_torch.as_tensor(mc_ptr, (total_elements,), torch_tensor_cpu.dtype)
+    torch_tensor_mc_flat = cutlass_torch.as_tensor(
+        mc_ptr, (total_elements,), torch_tensor_cpu.dtype
+    )
     torch_tensor_mc = torch_tensor_mc_flat.as_strided(
-        size=torch_tensor_cpu.shape,
-        stride=torch_tensor_cpu.stride()
+        size=torch_tensor_cpu.shape, stride=torch_tensor_cpu.stride()
     )
-    
+
     cute_tensor_mc = from_dlpack(torch_tensor_mc, assumed_align=16)
-    
+
     if is_dynamic_layout:
         for i, stride in enumerate(torch_tensor_mc.stride()):
             if stride == 1:
                 leading_dim = i
                 break
         cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim)
-    
+
     torch_tensor_gpu = torch_symm_tensor
     cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16)
     cute_tensor.element_type = dtype
-    
+
     if is_dynamic_layout:
         for i, stride in enumerate(torch_tensor_gpu.stride()):
             if stride == 1:
                 leading_dim = i
                 break
         cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim)
-    
+
     cute_tensor = cutlass_torch.convert_cute_tensor(
         torch_tensor_gpu,
         cute_tensor,
@@ -81,44 +81,49 @@ 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, cluster_shape_mn, sm_count):
-        barrier_size = Sm100BlockScaledPersistentDenseGemmKernel.compute_barrier_flag_size(
-            m, n, l, mma_tiler_mn, cluster_shape_mn, sm_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(
-            (barrier_size,), device="cuda", dtype=torch.int32
-        )
-        
-        barrier_flag.fill_(0)
-        symm = symm_mem.rendezvous(barrier_flag, group=dist.group.WORLD.group_name)
-        barrier_flag_mc_ptr = symm.multicast_ptr
-
-        barrier_flag_memref = from_dlpack(barrier_flag)
-        barrier_flag_memref = barrier_flag_memref.mark_layout_dynamic()
-        barrier_flag_mc_torch = cutlass_torch.as_tensor(
-                barrier_flag_mc_ptr, barrier_flag.shape, barrier_flag.dtype
-        )
-        barrier_flag_mc_memref = from_dlpack(
-            barrier_flag_mc_torch,
-        )
-        barrier_flag_mc_memref = barrier_flag_mc_memref.mark_layout_dynamic()
-        barrier_flag_torch = barrier_flag
-        return barrier_flag_memref, barrier_flag_mc_memref, barrier_flag_torch, barrier_flag_mc_torch
+    barrier_size = Sm100BlockScaledPersistentDenseGemmKernel.compute_barrier_flag_size(
+        m, n, l, mma_tiler_mn, cluster_shape_mn, sm_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((barrier_size,), device="cuda", dtype=torch.int32)
+
+    barrier_flag.fill_(0)
+    symm = symm_mem.rendezvous(barrier_flag, group=dist.group.WORLD.group_name)
+    barrier_flag_mc_ptr = symm.multicast_ptr
+
+    barrier_flag_memref = from_dlpack(barrier_flag)
+    barrier_flag_memref = barrier_flag_memref.mark_layout_dynamic()
+    barrier_flag_mc_torch = cutlass_torch.as_tensor(
+        barrier_flag_mc_ptr, barrier_flag.shape, barrier_flag.dtype
+    )
+    barrier_flag_mc_memref = from_dlpack(
+        barrier_flag_mc_torch,
+    )
+    barrier_flag_mc_memref = barrier_flag_mc_memref.mark_layout_dynamic()
+    barrier_flag_torch = barrier_flag
+    return (
+        barrier_flag_memref,
+        barrier_flag_mc_memref,
+        barrier_flag_torch,
+        barrier_flag_mc_torch,
+    )
+
 
 def run_blockscaled_gemm_all_reduce_python_interface(
     lm: Tuple[int, int],
@@ -139,7 +144,7 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     iterations: int,
     enable_dst_signals: int,
     all_reduce: str,
-    rank:int,
+    rank: int,
 ):
     torch.manual_seed(42)
     device = torch.device("cuda", rank)
@@ -187,7 +192,12 @@ def run_blockscaled_gemm_all_reduce_python_interface(
         l, n, k, b_major == "n", cutlass.Float32, device=device
     )
     c_ref = cutlass_torch.matrix(
-        l, m, n, c_major == "m", cutlass.Float32, device=device,
+        l,
+        m,
+        n,
+        c_major == "m",
+        cutlass.Float32,
+        device=device,
         init_type=cutlass_torch.TensorInitType.SCALAR,
         init_config=cutlass_torch.ScalarInitConfig(value=0.0),
     )
@@ -213,14 +223,21 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     c_tensor, c_tensor_mc, c_torch, c_torch_mc = create_mc_tensor(
         c_ref,
         get_cutlass_dtype(c_dtype),
-        #(1 if c_major == "n" else 0),
+        # (1 if c_major == "n" else 0),
         is_dynamic_layout=True,
     )
-    print(f"Rank {rank}: c_torch INITIAL shape={c_torch.shape}, stride={c_torch.stride()}")
+    print(
+        f"Rank {rank}: c_torch INITIAL shape={c_torch.shape}, stride={c_torch.stride()}"
+    )
     alpha_tensor = (
         torch.randn(l, dtype=torch.float32, device=device) if fuse_alpha else None
     )
-    barrier_flag_memref, barrier_flag_mc_memref, barrier_flag_torch, barrier_flag_mc_torch = create_barrier_flags(
+    (
+        barrier_flag_memref,
+        barrier_flag_mc_memref,
+        barrier_flag_torch,
+        barrier_flag_mc_torch,
+    ) = create_barrier_flags(
         m,
         n,
         l,
@@ -254,15 +271,15 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor(
         l, n, k, sf_vec_size, get_cutlass_dtype(sf_dtype), device
     )
-    #masked_m_tensor = torch.randint(0, m, (l,), dtype=torch.int32, device=device)
-    # if rank == 0:
-    #     masked_m_tensor = torch.randint(0, m, (l,), dtype=torch.int32, device=device)
-    # else:
-    #     masked_m_tensor = torch.empty((l,), dtype=torch.int32, device=device)
-    # torch.distributed.broadcast(masked_m_tensor, src=0)
+    # masked_m_tensor = torch.randint(0, m, (l,), dtype=torch.int32, device=device)
+    if rank == 0:
+        masked_m_tensor = torch.randint(0, m, (l,), dtype=torch.int32, device=device)
+    else:
+        masked_m_tensor = torch.empty((l,), dtype=torch.int32, device=device)
+    torch.distributed.broadcast(masked_m_tensor, src=0)
     # to hack and test:
-    masked_m_tensor = torch.full((l,), m, dtype=torch.int32, device=device)    
-    print(f"Rank {rank}: masked_m = {masked_m_tensor}") 
+    # masked_m_tensor = torch.full((l,), m, dtype=torch.int32, device=device)
+    print(f"Rank {rank}: masked_m = {masked_m_tensor}")
     for _ in range(iterations):
         dst_signals = (
             torch.zeros((l,), dtype=torch.uint32, device="cuda")
@@ -306,7 +323,9 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     ref = torch.einsum("mkl,nkl->mnl", res_a, res_b)
     ref = torch.einsum("mnl,l->mnl", ref, alpha_tensor)
     ref = ref.contiguous()
-    torch.distributed.all_reduce(ref, op=torch.distributed.ReduceOp.SUM, group=dist.group.WORLD)
+    torch.distributed.all_reduce(
+        ref, op=torch.distributed.ReduceOp.SUM, group=dist.group.WORLD
+    )
     # Convert c back to f32 for comparison.
     ref = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0)
     print(f"Rank {rank}: c_ref shape={c_ref.shape}, stride={c_ref.stride()}")
@@ -354,9 +373,10 @@ def run_blockscaled_gemm_all_reduce_python_interface(
                 rtol=1e-02,
             )
 
+
 def _run_correctness_worker(
-    world_size, 
-    rank, 
+    world_size,
+    rank,
     distributed_init_port,
     lm,
     kn,
@@ -447,9 +467,48 @@ 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.parametrize("lm", [(1, 1024), (2, 512), (4, 256)])
+# @pytest.mark.parametrize("kn", [(7168, 4096), (2048, 7168)])
+# @pytest.mark.parametrize(
+#     "ab_dtype,sf_dtype,c_dtype,sf_vec_size",
+#     [
+#         ("float4_e2m1fn", "float8_e8m0fnu", "float16", 16),
+#         ("float4_e2m1fn", "float8_e8m0fnu", "bfloat16", 16),
+#         ("float4_e2m1fn", "float8_e8m0fnu", "float32", 16),
+#         ("float4_e2m1fn", "float8_e4m3fn", "float16", 16),
+#         ("float4_e2m1fn", "float8_e4m3fn", "bfloat16", 16),
+#         ("float4_e2m1fn", "float8_e4m3fn", "float32", 16),
+#         ("float8_e4m3fn", "float8_e8m0fnu", "bfloat16", 32),
+#         ("float8_e4m3fn", "float8_e8m0fnu", "float16", 32),
+#         ("float8_e4m3fn", "float8_e8m0fnu", "float32", 32),
+#         ("float8_e4m3fn", "float8_e8m0fnu", "float8_e4m3fn", 32),
+#         ("float8_e4m3fn", "float8_e8m0fnu", "float8_e5m2", 32),
+#         ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32),
+#         ("float8_e5m2", "float8_e8m0fnu", "float16", 32),
+#         ("float8_e5m2", "float8_e8m0fnu", "float32", 32),
+#         ("float8_e5m2", "float8_e8m0fnu", "float8_e4m3fn", 32),
+#         ("float8_e5m2", "float8_e8m0fnu", "float8_e5m2", 32),
+#     ],
+# )
+# @pytest.mark.parametrize("a_major", ["k"])
+# @pytest.mark.parametrize("b_major", ["k"])
+# @pytest.mark.parametrize("c_major", ["n"])
+# @pytest.mark.parametrize("fuse_alpha", [False, True])
+# @pytest.mark.parametrize("alpha_dtype", ["float32"])
+# @pytest.mark.parametrize("mma_tiler_mn", [(128, 128)])
+# @pytest.mark.parametrize("cluster_shape_mn", [(1, 1)])
+# @pytest.mark.parametrize("sm_count", [132, None])
+# @pytest.mark.parametrize("tolerance", [1e-01])
+# @pytest.mark.parametrize("iterations", [3])
+# @pytest.mark.parametrize("enable_dst_signals", [False, True])
+
+# ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32)
+
 
 @pytest.mark.skipif(
     not is_cute_dsl_available(), reason="Please `pip install nvidia-cutlass-dsl`"
@@ -460,21 +519,36 @@ def multi_process_parallel(
 @pytest.mark.parametrize(
     "ab_dtype,sf_dtype,c_dtype,sf_vec_size",
     [
-        ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32),
-        # Add more combinations as needed
+        ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32)
+        # ("float4_e2m1fn", "float8_e8m0fnu", "float16", 16),
+        # ("float4_e2m1fn", "float8_e8m0fnu", "bfloat16", 16),
+        # ("float4_e2m1fn", "float8_e8m0fnu", "float32", 16),
+        # ("float4_e2m1fn", "float8_e4m3fn", "float16", 16),
+        # ("float4_e2m1fn", "float8_e4m3fn", "bfloat16", 16),
+        # ("float4_e2m1fn", "float8_e4m3fn", "float32", 16),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "bfloat16", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float16", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float32", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float8_e4m3fn", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float8_e5m2", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float16", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float32", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float8_e4m3fn", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float8_e5m2", 32),
     ],
 )
 @pytest.mark.parametrize("a_major", ["k"])
 @pytest.mark.parametrize("b_major", ["k"])
 @pytest.mark.parametrize("c_major", ["n"])
-@pytest.mark.parametrize("fuse_alpha", [False])
+@pytest.mark.parametrize("fuse_alpha", [False, True])
 @pytest.mark.parametrize("alpha_dtype", ["float32"])
 @pytest.mark.parametrize("mma_tiler_mn", [(128, 128)])
 @pytest.mark.parametrize("cluster_shape_mn", [(1, 1)])
 @pytest.mark.parametrize("sm_count", [148])
 @pytest.mark.parametrize("tolerance", [1e-01])
 @pytest.mark.parametrize("iterations", [1])
-@pytest.mark.parametrize("enable_dst_signals", [True])
+@pytest.mark.parametrize("enable_dst_signals", [False, True])
 @pytest.mark.parametrize("all_reduce", ["two_shot"])
 def test_cute_dsl_blockscaled_gemm_allreduce_two_shot(
     world_size,
@@ -527,4 +601,4 @@ def test_cute_dsl_blockscaled_gemm_allreduce_two_shot(
             all_reduce,
         ),
     )
-    print(f"cute_dsl_blockscaled_gemm_allreduce_two_shot on {world_size} GPUs: OK")
+    print(f"cute_dsl_blockscaled_gemm_allreduce_two_shot on {world_size} GPUs: OK")