wip

Author: Amir-19

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -28,45 +28,51 @@
 
 
 def create_mc_tensor(torch_tensor_cpu, dtype, is_dynamic_layout=True):
-    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
+    m, n, l = torch_tensor_cpu.shape
+    
+    # Create flat symm_mem buffer
+    total_elements = m * n * l
+    torch_symm_flat = symm_mem.empty(
+        (total_elements,), device="cuda", dtype=torch_tensor_cpu.dtype
     )
-
-    torch_symm_tensor_lmn.copy_(torch_tensor_cpu_lmn)
-    torch_symm_tensor = torch_symm_tensor_lmn.permute(1, 2, 0)
-    symm = symm_mem.rendezvous(torch_symm_tensor, group=dist.group.WORLD.group_name)
+    
+    # Reshape to match input's stride pattern using as_strided
+    torch_symm_tensor = torch_symm_flat.as_strided(
+        size=torch_tensor_cpu.shape,
+        stride=torch_tensor_cpu.stride()
+    )
+    torch_symm_tensor.copy_(torch_tensor_cpu)
+    
+    symm = symm_mem.rendezvous(torch_symm_flat, group=dist.group.WORLD.group_name)
     mc_ptr = symm.multicast_ptr
-    # create MC tensor memref
-    torch_tensor_mc = cutlass_torch.as_tensor(mc_ptr, torch_tensor_cpu.shape, torch_tensor_cpu.dtype)
-    cute_tensor_mc = from_dlpack(
-        torch_tensor_mc,
-        assumed_align=16,
+    
+    # Create MC tensor with same stride
+    torch_tensor_mc_flat = cutlass_torch.as_tensor(mc_ptr, (total_elements,), torch_tensor_cpu.dtype)
+    torch_tensor_mc = torch_tensor_mc_flat.as_strided(
+        size=torch_tensor_cpu.shape,
+        stride=torch_tensor_cpu.stride()
     )
-    # if is_dynamic_layout:
-    #     cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim)
+    
+    cute_tensor_mc = from_dlpack(torch_tensor_mc, assumed_align=16)
+    
     if is_dynamic_layout:
         for i, stride in enumerate(torch_tensor_mc.stride()):
             if stride == 1:
                 leading_dim = i
                 break
         cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim)
+    
     torch_tensor_gpu = torch_symm_tensor
     cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16)
     cute_tensor.element_type = dtype
-    # if is_dynamic_layout:
-    #     cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim)
+    
     if is_dynamic_layout:
-        for i, stride in enumerate(torch_tensor_mc.stride()):
+        for i, stride in enumerate(torch_tensor_gpu.stride()):
             if stride == 1:
                 leading_dim = i
                 break
         cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim)
+    
     cute_tensor = cutlass_torch.convert_cute_tensor(
         torch_tensor_gpu,
         cute_tensor,