wip

Amir-19 · Amir-19 · commit 55ce566ebb0e · 2025-09-29T19:36:27.000-07:00
diff --git a/tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py b/tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py
@@ -171,7 +171,7 @@ def test_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()}")
     a_tensor, a_torch = cutlass_torch.cute_tensor_like(
         a_ref,
         get_cutlass_dtype(ab_dtype),
@@ -196,6 +196,7 @@ def test_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()}")
     alpha_tensor = (
         torch.randn(l, dtype=torch.float32, device=device) if fuse_alpha else None
     )

Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ def test_blockscaled_gemm_all_reduce_python_interface(`
`171`	`171`	`init_type=cutlass_torch.TensorInitType.SCALAR,`
`172`	`172`	`init_config=cutlass_torch.ScalarInitConfig(value=0.0),`
`173`	`173`	`)`
`174`		`-`
	`174`	`+ print(f"Rank {rank}: c_ref INITIAL shape={c_ref.shape}, stride={c_ref.stride()}")`
`175`	`175`	`a_tensor, a_torch = cutlass_torch.cute_tensor_like(`
`176`	`176`	`a_ref,`
`177`	`177`	`get_cutlass_dtype(ab_dtype),`
`@@ -196,6 +196,7 @@ def test_blockscaled_gemm_all_reduce_python_interface(`
`196`	`196`	`#(1 if c_major == "n" else 0),`
`197`	`197`	`is_dynamic_layout=True,`
`198`	`198`	`)`
	`199`	`+ print(f"Rank {rank}: c_torch INITIAL shape={c_torch.shape}, stride={c_torch.stride()}")`
`199`	`200`	`alpha_tensor = (`
`200`	`201`	`torch.randn(l, dtype=torch.float32, device=device) if fuse_alpha else None`
`201`	`202`	`)`