[KERNELS] fixes a couple of issues in matmul_ogs.py (triton-lang#7177)

Author: ptillet

python/triton_kernels/triton_kernels/matmul_ogs.py

@@ -584,8 +584,7 @@ def apply_allocation(allocation: MatmulAllocation, output):
 
 def _create_tma_descriptors(
     x: torch.Tensor,
-    x_tensor: torch.Tensor,
-    w_tensor: torch.Tensor,
+    w: torch.Tensor,
     mx_tensor: Optional[torch.Tensor],
     routing_data: RoutingData,
     mx_ctx: MicroscalingCtx,
@@ -611,21 +610,21 @@ def _create_tma_descriptors(
     if (use_host_tma_descriptors):
         if USE_GATHER_TMA or X_USE_LOAD_TMA:
             x_desc = TensorDescriptorBuilder.create_input_descriptor(
-                    x_tensor, K, x.stride(1), x.stride(2),
+                    x, K, x.stride(1), x.stride(2),
                     opt_flags.block_k, opt_flags.block_m,
                     USE_GATHER_TMA, X_USE_LOAD_TMA
                 )
         descriptors.append(x_desc)
         if (expt_data is not None and len(expt_data.block_pid_map) > 0):
             w_desc = TensorDescriptorBuilder.create_weight_descriptor(
-                    w_tensor, opt_flags.block_k, opt_flags.block_n, w_transpose
+                    w, opt_flags.block_k, opt_flags.block_n, w_transpose
                 )
-            is_microscaled_format = (mx_ctx.weight_scale is not None) and (w_tensor.dtype == torch.uint8)
+            is_microscaled_format = (mx_ctx.weight_scale is not None) and (w.dtype == torch.uint8)
             if is_microscaled_format:
                 # Pad the inner shape to 128 for mxfp4 weights
                 # for mixed precision fp8 x mxfp4 compute
                 pad = 128
-                dim_to_pad = -1 if w_transpose else -2
+                dim_to_pad = -1
                 old_size = w_desc.shape[dim_to_pad]
                 padded_size = math.ceil(old_size / pad) * pad
                 if padded_size != old_size:
@@ -645,7 +644,7 @@ def _create_tma_descriptors(
     # TODO: Currently all or none, instead should support a mixture
     # of host and device descriptors
     if None in descriptors or len(descriptors) == 0:
-        descriptors = [x_tensor, w_tensor, mx_tensor]
+        descriptors = [x, w, mx_tensor]
         use_host_tma_descriptors = False
     if opt_flags.is_persistent:
         opt_flags.target_kernel_kwargs["USE_HOST_TMA_DESCRIPTORS"] = use_host_tma_descriptors
@@ -759,9 +758,7 @@ def matmul_ogs(x, w, bias,
     USE_GATHER_TMA = HAS_TMA_GS and gather_indx is not None
     X_USE_LOAD_TMA = gather_indx is None and not USE_GATHER_TMA
     _, x_tensor, w_tensor, mx_tensor = _create_tma_descriptors(
-        x=x,
-        x_tensor=flex.lhs_data.reinterpret(x),
-        w_tensor=flex.rhs_data.reinterpret(w),
+        x=x, w=w,
         mx_tensor=mx_ctx.weight_scale,
         routing_data=routing_data,
         mx_ctx=mx_ctx,
@@ -777,7 +774,10 @@ def matmul_ogs(x, w, bias,
         w_transpose=w.stride(2) != 1,
         mx_transpose=mx_scale_stride_n != 1,
     )
-
+    if isinstance(x_tensor, torch.Tensor):
+        x_tensor = flex.lhs_data.reinterpret(x)
+    if isinstance(w_tensor, torch.Tensor):
+        w_tensor = flex.rhs_data.reinterpret(w)
     (kernels._p_matmul_ogs if opt_flags.is_persistent else kernels._matmul_ogs)[(n_cta,)](
                    flex.out_data.reinterpret(memory["output"]),
                    flex.out_data.reinterpret(out0), *out0.stride(),

python/triton_kernels/triton_kernels/matmul_ogs_details/_common.py

@@ -52,11 +52,14 @@ def matmul_repr(specialization):
 
         def convert_dtype(dtype):
             if "tensordesc" in dtype:
-                return dtype.split("<")[1].split("[")[0]
+                ret = convert_dtype(dtype.split("<")[1].split("[")[0])
+                return ret
             elif "u8" in dtype:
                 return "mxfp4"
-            else:
+            elif dtype[0] == "*":
                 return dtype[1:]
+            else:
+                return dtype
 
         dtypes = "x".join([convert_dtype(f"{signature[i]}") for i in reorder(["Y", "X", "W"])])
         layouts = "".join([f"{layout(i)}" for i in reorder(["stride_y_n", "stride_x_k", "stride_w_n"])])

python/triton_kernels/triton_kernels/matmul_ogs_details/_finalize_matmul.py

@@ -291,7 +291,7 @@ def _finalize_matmul(
                         if src_idx != -1:
                             As = A + src_idx.to(tl.int64) * stride_a_m + offs_n
                             for ki in tl.static_range(K):
-                                acc += tl.load(As, mask=(src_idxs != -1)[:, None] & n_mask[None, :], other=0.0)
+                                acc += tl.load(As, mask=n_mask, other=0.0)
                                 As += stride_a_k
                 else:
                     As = A + src_idxs.to(tl.int64)[:, None] * stride_a_m + offs_n[None, :]