Mamba3 SISO kernel: 3-7.5% deviation from reference (heavy-tailed, masked by bf16 rtol=1e-1)

Mamba-3 SISO kernel deviates from reference — heavy-tailed numerical errors

Summary

mamba3_siso_combined (Triton) does not match mamba3_siso_fwd_ref (the reference implementation in tests/ops/triton/test_mamba3_siso.py) at any tested config. Relative error is 3–7.5% across all settings; the deviation is heavy-tailed (few outliers up to ~0.3 absolute) rather than a constant offset. Current CI tolerance (rtol=1e-1 at bf16) masks the discrepancy.

Reproduction

Script: reconcile_m3_kernel.py (RankAdaptation repo). Uses the repo's own create_mamba3_siso_inputs + mamba3_siso_fwd_ref, then compares the combined kernel output on identical inputs.

Config tested: batch=4, seqlen=256, nheads=1, nheads_qk=1, headdim_qk=64, headdim_v=64 (paper-format parity config); plus hd64/hd128 × fp32/bf16 sweeps.

Results (relative error):

Case	fp32, hd64
D+Z+states	3.4e-2
D+Z no-states	3.4e-2
D only	4.5e-2
Z only	5.2e-2
plain	7.4e-2

Earlier sweeps (hd128, fp32 and bf16): 3.4–4.2e-2 consistently.

Error structure

• mean(diff) ≈ 0 in ALL cases (−8.7e-6 … +5.8e-6) → NOT a constant/DC offset (no missing D-skip or single missing term).
• mean(|diff|) = 2–4e-3 but max|diff| = 0.07–0.30 → heavy-tailed per-element errors, concentrated on a small fraction of outputs.
• Relative errors on near-zero reference values reach 22–95% (small absolute diffs, large relative) — consistent with precision loss (tf32/accumulation order / low-rank approximations) inside the recurrence, not a missing term.

Why it matters

State-tracking tasks (parity, counting) require EXACT state dynamics. In our testing, Mamba-3 does not stably learn parity in the paper's own format at 1-layer/10K (best transient 62.5%, collapses) nor 2-layer/20K (52.1% best), even though the module trains fine on next-token copy (loss → 0). A systematic per-step numerical deviation inside the scan is a plausible contributor to this instability, and would not be visible to the loose bf16 CI tolerance.

Suggested investigation

1. Confirm whether tf32 is enabled on the kernel's tl.dot paths (fp32 inputs) and whether disabling tf32 / forcing bf16 accumulation changes the error.
2. Check per-chunk boundary state passing for precision loss compounding.
3. Consider tightening the upstream test tolerance for fp32 runs (e.g., rtol=1e-3) so regressions of this class are caught.

Environment

• triton 3.6.0, torch 2.11.0+cu128, RTX 4060 (sm_89)
• repo commit: state-spaces/mamba @ ed6ce09 (vendored mamba_src)