wip

Author: Amir-19

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -114,7 +114,7 @@ def create_barrier_flags(m, n, l, mma_tiler_mn):
         barrier_flag_torch = barrier_flag
         return barrier_flag_memref, barrier_flag_mc_memref, barrier_flag_torch, barrier_flag_mc_torch
 
-def test_blockscaled_gemm_all_reduce_python_interface(
+def run_blockscaled_gemm_all_reduce_python_interface(
     lm: Tuple[int, int],
     kn: Tuple[int, int],
     ab_dtype: cutlass.dtype,
@@ -362,7 +362,7 @@ def _run_correctness_worker(world_size, rank, distributed_init_port):
     rank_id = torch.distributed.get_rank()
 
     try:
-        test_blockscaled_gemm_all_reduce_python_interface(
+        run_blockscaled_gemm_all_reduce_python_interface(
             lm=(2, 512), # (1, 1024), (2, 512), (4, 256)
             kn=(7168, 4096),
             ab_dtype="float8_e5m2",