wip

Author: Amir-19

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -354,7 +354,29 @@ def run_blockscaled_gemm_all_reduce_python_interface(
                 rtol=1e-02,
             )
 
-def _run_correctness_worker(world_size, rank, distributed_init_port):
+def _run_correctness_worker(
+    world_size, 
+    rank, 
+    distributed_init_port,
+    lm,
+    kn,
+    ab_dtype,
+    sf_dtype,
+    sf_vec_size,
+    c_dtype,
+    a_major,
+    b_major,
+    c_major,
+    fuse_alpha,
+    alpha_dtype,
+    mma_tiler_mn,
+    cluster_shape_mn,
+    sm_count,
+    tolerance,
+    iterations,
+    enable_dst_signals,
+    all_reduce,
+):
     assert rank >= 0
     torch.cuda.set_device(rank)
     device = torch.device("cuda", rank)
@@ -371,24 +393,24 @@ def _run_correctness_worker(world_size, rank, distributed_init_port):
 
     try:
         run_blockscaled_gemm_all_reduce_python_interface(
-            lm=(2, 512), # (1, 1024), (2, 512), (4, 256)
-            kn=(7168, 4096),
-            ab_dtype="float8_e5m2",
-            sf_dtype="float8_e8m0fnu",
-            sf_vec_size=32,
-            c_dtype="bfloat16",
-            a_major="k",
-            b_major="k",
-            c_major="n",
-            fuse_alpha=False,
-            alpha_dtype="float32",
-            mma_tiler_mn=(128, 128),
-            cluster_shape_mn=(1, 1),
-            tolerance=1e-01,
-            iterations=1,
-            sm_count=148,
-            enable_dst_signals=True,
-            all_reduce="two_shot",
+            lm=lm,
+            kn=kn,
+            ab_dtype=ab_dtype,
+            sf_dtype=sf_dtype,
+            sf_vec_size=sf_vec_size,
+            c_dtype=c_dtype,
+            a_major=a_major,
+            b_major=b_major,
+            c_major=c_major,
+            fuse_alpha=fuse_alpha,
+            alpha_dtype=alpha_dtype,
+            mma_tiler_mn=mma_tiler_mn,
+            cluster_shape_mn=cluster_shape_mn,
+            tolerance=tolerance,
+            iterations=iterations,
+            sm_count=sm_count,
+            enable_dst_signals=enable_dst_signals,
+            all_reduce=all_reduce,
             rank=rank,
         )
     except Exception as e:
@@ -433,7 +455,48 @@ def multi_process_parallel(
     not is_cute_dsl_available(), reason="Please `pip install nvidia-cutlass-dsl`"
 )
 @pytest.mark.parametrize("world_size", [8])
-def test_cute_dsl_blockscaled_gemm_allreduce_two_shot(world_size):
+@pytest.mark.parametrize("lm", [(1, 1024), (2, 512), (4, 256)])
+@pytest.mark.parametrize("kn", [(7168, 4096)])
+@pytest.mark.parametrize(
+    "ab_dtype,sf_dtype,c_dtype,sf_vec_size",
+    [
+        ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32),
+        # Add more combinations as needed
+    ],
+)
+@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("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("all_reduce", ["two_shot"])
+def test_cute_dsl_blockscaled_gemm_allreduce_two_shot(
+    world_size,
+    lm,
+    kn,
+    ab_dtype,
+    sf_dtype,
+    sf_vec_size,
+    c_dtype,
+    a_major,
+    b_major,
+    c_major,
+    fuse_alpha,
+    alpha_dtype,
+    mma_tiler_mn,
+    cluster_shape_mn,
+    sm_count,
+    tolerance,
+    iterations,
+    enable_dst_signals,
+    all_reduce,
+):
     available_gpus = torch.cuda.device_count()
     if world_size > available_gpus:
         pytest.skip(
@@ -443,6 +506,25 @@ def test_cute_dsl_blockscaled_gemm_allreduce_two_shot(world_size):
     multi_process_parallel(
         world_size,
         _run_correctness_worker,
-        target_args=(),
+        target_args=(
+            lm,
+            kn,
+            ab_dtype,
+            sf_dtype,
+            sf_vec_size,
+            c_dtype,
+            a_major,
+            b_major,
+            c_major,
+            fuse_alpha,
+            alpha_dtype,
+            mma_tiler_mn,
+            cluster_shape_mn,
+            sm_count,
+            tolerance,
+            iterations,
+            enable_dst_signals,
+            all_reduce,
+        ),
     )
     print(f"cute_dsl_blockscaled_gemm_allreduce_two_shot on {world_size} GPUs: OK")