[kernels] fix tma construction bug when M == 1 (#7148)

Author: ptillet

python/triton_kernels/triton_kernels/matmul_ogs.py

@@ -144,11 +144,17 @@ def create_block_scale_descriptor(mx_tensor: torch.Tensor, block_k: int, block_n
                                                                    block_shape=[1, MX_SCALE_BLOCK_K,
                                                                                 block_n], transpose=transpose)
 
+    @staticmethod
+    def squeeze_after_dim(x, dim=2):
+        shape = list(x.shape)
+        new_shape = [s for s in shape[:dim - 1] if s != 1] + shape[dim - 1:]
+        return x.view(*new_shape)
+
     @staticmethod
     def create_input_descriptor_gather(x_tensor: torch.Tensor, K: int, x_stride_1: int, x_stride_2: int,
                                        block_k: int) -> TensorDescriptor:
         """Create a tensor descriptor for input matrix X via TMA gather"""
-        x_desc = x_tensor.squeeze()
+        x_desc = TensorDescriptorBuilder.squeeze_after_dim(x_tensor)
         assert x_desc.ndim == 2, "TMA gather descriptor requires 2D input"
         INT_MAX = 2147483647
         return TensorDescriptor(base=x_desc, shape=[INT_MAX, K], strides=[x_stride_1, x_stride_2],
@@ -158,7 +164,7 @@ def create_input_descriptor_gather(x_tensor: torch.Tensor, K: int, x_stride_1: i
     def create_input_descriptor_load(x_tensor: torch.Tensor, K: int, x_stride_1: int, x_stride_2: int, block_m: int,
                                      block_k: int) -> TensorDescriptor:
         """Create a tensor descriptor for input matrix X via TMA"""
-        x_desc = x_tensor.squeeze()
+        x_desc = TensorDescriptorBuilder.squeeze_after_dim(x_tensor)
         assert x_desc.ndim in [2, 3], "LHS input TMA descriptor builder expects 2D or 3D input"
         return TensorDescriptor(base=x_desc, shape=[x_desc.shape[0], K], strides=[x_stride_1, x_stride_2],
                                 block_shape=[block_m, block_k])