fix(mxfp8_moe): zero empty-expert weight gradients, accept non-contiguous gradients, smoke-test the kernel

prime_kernels/mxfp8_moe/grouped_gemm.py

@@ -165,6 +165,10 @@ def _wgrad(
         offs=offsets,
         out_dtype=torch.bfloat16,
     )
+    # `_scaled_grouped_mm` leaves an empty group's output unwritten, unlike `_grouped_mm`, so an
+    # expert that received no tokens would get uninitialized memory instead of a zero gradient.
+    empty_groups = torch.diff(offsets, prepend=offsets.new_zeros(1)) == 0
+    grad_weight = torch.where(empty_groups[:, None, None], 0.0, grad_weight)
     return grad_weight.transpose(-2, -1)
 
 
@@ -184,6 +188,9 @@ def forward(
     @staticmethod
     def backward(ctx, grad_output: torch.Tensor):
         input_act, weight_t, offsets = ctx.saved_tensors
+        # torchao's MXFP8 casts require contiguous input, but autograd may pass a non-contiguous gradient
+        # (e.g. the stride-0 expansion from `out.sum().backward()`).
+        grad_output = grad_output.contiguous()
         grad_input = _dgrad(grad_output, weight_t, offsets) if ctx.needs_input_grad[0] else None
         grad_weight = (
             _wgrad(
@@ -205,7 +212,7 @@ def grouped_gemm(
     *,
     high_precision_wgrad: bool = False,
 ) -> torch.Tensor:
-    """Run differentiable MXFP8 grouped GEMM on already aligned token groups."""
+    """Run differentiable MXFP8 grouped GEMM on token groups; every group size must be a multiple of `TOKEN_GROUP_ALIGNMENT`, or the weight gradient is silently wrong."""
     if input_act.ndim != 2 or weight_t.ndim != 3:
         raise ValueError(
             f"MXFP8 grouped GEMM expects 2D activations and 3D weights, got {input_act.ndim}D and {weight_t.ndim}D"

scripts/smoke_test.py

@@ -34,6 +34,24 @@ def smoke_prime_kernels():
     flash_moe = prime_kernels.load("flash_moe")
     print("flash_moe loaded:", flash_moe)
 
+    if "torchao" in status.get("mxfp8_moe", ""):
+        print("mxfp8_moe needs torchao, which this environment lacks; skipping its check")
+        return
+    assert status.get("mxfp8_moe") == "available", status
+    mxfp8_moe = prime_kernels.load("mxfp8_moe")
+    # 48 groups take the Triton scale rearrange with power-of-two padding; empty groups at the start,
+    # middle and end must get zero weight gradients, so unwritten memory is made NaN first.
+    counts = torch.full((48,), 32, device="cuda")
+    counts[[0, 23, 47]] = 0
+    offsets = torch.cumsum(counts, 0, dtype=torch.int32)
+    x = torch.randn(int(counts.sum()), 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    weight = torch.randn(48, 128, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    torch.full((64 << 20,), float("nan"), device="cuda")
+    mxfp8_moe.grouped_gemm(x, weight.transpose(-2, -1), offsets).sum().backward()
+    assert torch.isfinite(x.grad).all() and torch.isfinite(weight.grad[counts > 0]).all()
+    assert torch.all(weight.grad[counts == 0] == 0), "experts with no tokens must get a zero gradient"
+    print("mxfp8_moe grouped_gemm forward + backward ran")
+
 
 def smoke_deep_gemm():
     import deep_gemm