wip

Author: Amir-19

flashinfer/cute_dsl/blockscaled_gemm.py

@@ -2734,7 +2734,7 @@ def can_implement(
         can_implement = True
         # Skip unsupported types
         if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size(
-            ab_dtype, sf_dtype, sf_vec_size, c_dtype,all_reduce
+            ab_dtype, sf_dtype, sf_vec_size, c_dtype, all_reduce
         ):
             can_implement = False
         # Skip unsupported layouts

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -158,7 +158,7 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     l, m = lm
     k, n = kn
 
-    print(f"device: {device}")
+    #print(f"device: {device}")
 
     if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement(
         get_cutlass_dtype(ab_dtype),
@@ -201,7 +201,7 @@ def run_blockscaled_gemm_all_reduce_python_interface(
         init_type=cutlass_torch.TensorInitType.SCALAR,
         init_config=cutlass_torch.ScalarInitConfig(value=0.0),
     )
-    print(f"Rank {rank}: c_ref INITIAL shape={c_ref.shape}, stride={c_ref.stride()}")
+    #print(f"Rank {rank}: c_ref INITIAL shape={c_ref.shape}, stride={c_ref.stride()}")
     a_tensor, a_torch = cutlass_torch.cute_tensor_like(
         a_ref,
         get_cutlass_dtype(ab_dtype),
@@ -226,9 +226,9 @@ def run_blockscaled_gemm_all_reduce_python_interface(
         # (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
     )
@@ -279,7 +279,7 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     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}")
+    # print(f"Rank {rank}: masked_m = {masked_m_tensor}")
     for _ in range(iterations):
         dst_signals = (
             torch.zeros((l,), dtype=torch.uint32, device="cuda")
@@ -328,9 +328,9 @@ def run_blockscaled_gemm_all_reduce_python_interface(
     )
     # 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()}")
-    print(f"Rank {rank}: ref shape={ref.shape}, stride={ref.stride()}")
-    print(f"Rank {rank}: c_torch shape={c_torch.shape}, stride={c_torch.stride()}")
+    # print(f"Rank {rank}: c_ref shape={c_ref.shape}, stride={c_ref.stride()}")
+    # print(f"Rank {rank}: ref shape={ref.shape}, stride={ref.stride()}")
+    # print(f"Rank {rank}: c_torch shape={c_torch.shape}, stride={c_torch.stride()}")
     cute.testing.convert(
         c_tensor,
         from_dlpack(c_ref, assumed_align=16).mark_layout_dynamic(
@@ -472,70 +472,32 @@ def multi_process_parallel(
         ), 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`"
 )
 @pytest.mark.parametrize("world_size", [8])
 @pytest.mark.parametrize("lm", [(1, 1024), (2, 512), (4, 256)])
-@pytest.mark.parametrize("kn", [(7168, 4096)])
+@pytest.mark.parametrize("kn", [(7168, 4096), (2048, 7168)])
 @pytest.mark.parametrize(
     "ab_dtype,sf_dtype,c_dtype,sf_vec_size",
     [
         ("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),
+        ("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"])
@@ -576,6 +538,39 @@ def test_cute_dsl_blockscaled_gemm_allreduce_two_shot(
         pytest.skip(
             f"world_size {world_size} is greater than available_gpus {available_gpus}"
         )
+    #device = torch.device("cuda", rank)
+    major, minor = torch.cuda.get_device_capability(torch.device("cuda:0"))
+    if not (major == 10 and minor == 0):
+        pytest.skip("Cute-dsl backend is only supported on SM100.")
+    if enable_dst_signals and (sm_count is None):
+        pytest.skip("dst_signals require sm_count")
+
+    l, m = lm
+    k, n = kn
+    if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement(
+        get_cutlass_dtype(ab_dtype),
+        get_cutlass_dtype(sf_dtype),
+        sf_vec_size,
+        get_cutlass_dtype(c_dtype),
+        mma_tiler_mn,
+        cluster_shape_mn,
+        m,
+        n,
+        k,
+        l,
+        a_major,
+        b_major,
+        c_major,
+    ):
+        pytest.skip(
+            f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype},  {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}"
+        )
+
+    if not (a_major == "k" and b_major == "k" and c_major == "n"):
+        # not supported since we try to align deepgemm for now
+        pytest.skip(
+            f"Skip non deepgemm-like cases {a_major}, {b_major}, {c_major}. Might be added later"
+        )
     print(f"Running test for world_size={world_size}")
     multi_process_parallel(
         world_size,