fix(tests): fused MoE GeGLU parity threshold is tighter than the gather_qmm reference's own jitter #964

Summary

fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape fails deterministically on main. The failure is not a kernel defect: the assertion threshold is tighter than the reference implementation's own numerical jitter, which makes it unsatisfiable regardless of how correct the fused kernel is.

Reproduction

On main at commit 14412c13, Apple M1 Ultra (macOS, Metal backend):

cargo test -p mlxcel-core --release --features metal,accelerate

Result: 1050 passed; 1 failed. The failing test is fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape.

This is not a flake. The test iterates for seed in [886u64, 887, 888] and calls random_seed(seed), so the inputs are fully deterministic. It fails on the first seed, 886.

Measured values at seed 886

These are printed by the test's own diagnostic line:

• fused vs dense f32 (bf16-rounded): nrms=0e0, nmax=0e0
• fused vs gather_qmm: nrms=1.1697728227157964e-2, nmax=4.522727947528232e-2
• gather_qmm vs dense baseline (nrms_ref_jitter): nrms=1.1574492404669066e-2

Failing assertion

src/lib/mlxcel-core/src/fused_moe_parity_tests.rs, the second assertion in the test (around line 369):

assert!(
    nrms_gather < 5e-3 && nmax_gather < 5e-2,
    "fused vs gather_qmm reference deviates (seed {seed}): nrms={nrms_gather:e} nmax={nmax_gather:e}"
);

The nrms_gather < 5e-3 bound is the one that fails: actual 1.1698e-2, roughly 2.3x over the limit.

Root cause

The threshold sits below the reference's own noise floor, so the assertion cannot be satisfied by any implementation.

gather_qmm rounds gate/up/activation through bf16 between GEMMs, which gives it an intrinsic ~1.157e-2 nrms deviation from the dense f32 reference. The test measures exactly this and stores it in nrms_ref_jitter. Meanwhile the fused kernel is bit-exact against the bf16-rounded dense reference (nrms=0e0, nmax=0e0). Because fused is numerically identical to dense, the fused-vs-gather deviation is necessarily approximately equal to the gather-vs-dense deviation, that is ~1.157e-2. A 5e-3 bound is therefore mathematically unreachable, and this failure says nothing about the fused kernel's correctness.

Note that the test already computes nrms_ref_jitter and prints it in the diagnostic output, but never uses it in an assertion.

Suggested fix direction

Final call is the implementer's, but the natural fix is to make the gather comparison relative to the measured reference jitter instead of an absolute constant. For example, assert that nrms_gather is within a small factor of nrms_ref_jitter, since that is the quantity which actually bounds legitimate disagreement between the two paths. Alternatively, raise the absolute bound above the documented fp16-jitter class and record the measured magnitude in the accompanying comment.

Either way, keep the dense-reference assertion (nrms_dense < 5e-4 && nmax_dense < 2e-2) unchanged. That is the meaningful correctness gate, and it currently passes with a very wide margin, which is the positive evidence that the kernel math is fine.

Impact

Every sub-issue of epic #909 (the serving-performance epic, #898 through #908) carries a mandatory regression gate worded as "full cargo test -p mlxcel-core passes". With this failure present on main, that gate cannot be satisfied cleanly, and each implementer has to independently rediscover that the failure is pre-existing and unrelated to their own change. Fixing it, or explicitly quarantining it with a recorded reason, removes a recurring source of misattributed regressions.

Related: #909