diff --git a/prime_kernels/mxfp8_moe/grouped_gemm.py b/prime_kernels/mxfp8_moe/grouped_gemm.py index 338e820..709b5f6 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) @@ -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" 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