From 40ec6227b2cc0b633dfaf72e267b4b01fb575357 Mon Sep 17 00:00:00 2001 From: Garrett Goon Date: Fri, 16 Oct 2026 21:19:31 +0000 Subject: [PATCH 1/4] fix(mxfp8_moe): zero the MXFP8 weight gradient of experts with no tokens `torch._scaled_grouped_mm` leaves an empty group's output unwritten, so with the default MXFP8 weight-gradient recipe an expert that received no tokens got uninitialized memory (NaN or ~1e38) as its gradient. The high-precision path through `torch._grouped_mm` already zero-fills. prime-rl's dispatchers pad empty experts to one alignment block, so it never hit this, but the public `grouped_gemm` did. Co-Authored-By: Claude Opus 5.5 (1M context) --- prime_kernels/mxfp8_moe/grouped_gemm.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/prime_kernels/mxfp8_moe/grouped_gemm.py b/prime_kernels/mxfp8_moe/grouped_gemm.py index 338e820..3778d52 100644 --- a/prime_kernels/mxfp8_moe/grouped_gemm.py +++ b/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) From dce217bd0e2d8ccc38753460c3a5f6e06b0ade6a Mon Sep 17 00:00:00 2001 From: Garrett Goon Date: Fri, 16 Oct 2026 21:19:31 +0000 Subject: [PATCH 2/4] fix(mxfp8_moe): accept non-contiguous output gradients in backward torchao's MXFP8 casts require contiguous input, so a reduction's expanded, stride-0 gradient (for example from `out.sum().backward()`) failed with "`x` must be contiguous". Co-Authored-By: Claude Opus 5.5 (1M context) --- prime_kernels/mxfp8_moe/grouped_gemm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/prime_kernels/mxfp8_moe/grouped_gemm.py b/prime_kernels/mxfp8_moe/grouped_gemm.py index 3778d52..ed62ea4 100644 --- a/prime_kernels/mxfp8_moe/grouped_gemm.py +++ b/prime_kernels/mxfp8_moe/grouped_gemm.py @@ -188,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( From 46ff49a6e28657947d0d31b5aaa72ad6c2bfa483 Mon Sep 17 00:00:00 2001 From: Garrett Goon Date: Fri, 16 Oct 2026 21:19:31 +0000 Subject: [PATCH 3/4] test(smoke): run mxfp8_moe forward and backward on the built wheel A release wheel without a usable mxfp8_moe now fails CI instead of shipping; prime-rl's MXFP8 expert compute loads it. The check covers a non-power-of-two group count and an empty expert. Co-Authored-By: Claude Opus 5.5 (1M context) --- scripts/smoke_test.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py index 52d06fc..69a3e84 100755 --- a/scripts/smoke_test.py +++ b/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 From 2018c76d6ab07faa1b0e3c0c20eb6fc81fe13473 Mon Sep 17 00:00:00 2001 From: Garrett Goon Date: Sat, 17 Oct 2026 02:51:25 +0000 Subject: [PATCH 4/4] docs(mxfp8_moe): state that group sizes must be multiples of the alignment The weight gradient's 32-row scale blocks assume aligned group boundaries; an unaligned group would straddle two groups' blocks and silently get wrong scales. Not checked at runtime, since a check would add a host sync to every call. Co-Authored-By: Claude Opus 5.5 (1M context) --- prime_kernels/mxfp8_moe/grouped_gemm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prime_kernels/mxfp8_moe/grouped_gemm.py b/prime_kernels/mxfp8_moe/grouped_gemm.py index ed62ea4..709b5f6 100644 --- a/prime_kernels/mxfp8_moe/grouped_gemm.py +++ b/prime_kernels/mxfp8_moe/grouped_gemm.py @@ -212,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"