wip

Author: Amir-19

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -28,10 +28,18 @@
 
 
 def create_mc_tensor(torch_tensor_cpu, dtype, is_dynamic_layout=True):
-    torch_symm_tensor = symm_mem.empty(
-        torch_tensor_cpu.shape, device="cuda", dtype=torch_tensor_cpu.dtype
+    torch_tensor_cpu_lmn = torch_tensor_cpu.permute(2, 0, 1).contiguous()
+    # torch_symm_tensor = symm_mem.empty(
+    #     torch_tensor_cpu_lmn, device="cuda", dtype=torch_tensor_cpu.dtype
+    # )
+    torch_symm_tensor_lmn = symm_mem.empty(
+        torch_tensor_cpu_lmn.shape,  # (l, m, n)
+        device="cuda", 
+        dtype=torch_tensor_cpu.dtype
     )
-    torch_symm_tensor.copy_(torch_tensor_cpu)
+
+    torch_symm_tensor_lmn.copy_(torch_tensor_cpu)
+    torch_symm_tensor = torch_symm_tensor_lmn.permute(1, 2, 0)
     symm = symm_mem.rendezvous(torch_symm_tensor, group=dist.group.WORLD.group_name)
     mc_ptr = symm.multicast_ptr
     # create MC tensor memref