[RFC] Testing Methodology for FP32 Convergence Tests on Fused SwiGLU/dSwiGLU Kernels

[Pearblossom-M]

Background

LigerMLP (#1357) co-optimizes the full forward and backward dataflow of the SwiGLU MLP, fusing the gate/up GEMM + SwiGLU into a single Triton kernel (rather than just the element-wise SiLU-and-gating stage), in order to reduce the memory I/O of intermediate tensors and speed up training.

Current validation status:

• Unit tests (test/transformers/test_mlp.py, judged by cosine similarity) all pass.
• bf16 convergence tests pass overall, except for mini_gpt_oss, mini_qwen3_moe, and mini_pixtral: the first two also fail on the unmodified branch and are unrelated to this change; mini_pixtral fails due to a rounding-path difference (~3.18e-2 relative error), not a formula bug.
• The one systematic blocker is the fp32 convergence tests: across the three files under test/convergence/fp32/ (test_mini_models.py, test_mini_models_with_logits.py, test_mini_models_multimodal.py), about 25 cases — including mini_llama3, mini_llava, mini_qwen2, mini_qwen2_vl, mini_qwen2_5_vl, mini_qwen3, mini_mistral, mini_ministral, mini_granite3, mini_exaone4, and mini_mllama — all fail on Loss / Top-k logprobs mismatches.

Root Cause

By default, tl.dot computes fp32 matmuls via the TF32 Tensor Core path, which is less precise than true IEEE FP32. The reference implementation (nn.Linear) uses full FP32. This creates a systematic numerical gap on fp32 inputs, which is the root cause of the widespread convergence test failures — not a bug in the kernel logic.

Approaches I tried but didn't work

Forcing input_precision="ieee" in tl.dot to get true FP32: triggers a shared memory overflow (reports needing 114712 bytes vs. a hardware limit of 101376 bytes). The reason is that TF32 compiles to Tensor Core MMA instructions with a compact data layout, while IEEE cannot use Tensor Cores and falls back to a software FMA path that needs more shared memory. Even the smallest tile configuration in the current autotune search space overflows, so there is no fallback option.

Candidate Directions (open for discussion)

1. Relax fp32 tolerances: simple to implement, but lacks a principled basis for "how much error is acceptable" — treats the symptom, not the cause.
2. Force IEEE precision + redesign a finer-grained autotune search space: fixes the precision issue but comes with a clear performance cost, and it's unverified whether it can fully avoid the shared memory limit.
3. Change the testing methodology itself

Question for Maintainers

For operators like this one — inherently precision-sensitive due to their reliance on Tensor Cores — does the community/maintainers already have a preferred testing methodology?

References

• Implementation PR: feat: add LigerMLP module #1357
• Full failure logs and attempted approaches: feat: add LigerMLP module #1357 (comment)
• Original RFC for this kernel: [RFC] Full-pipeline SwiGLU MLP with forward/backward co-optimization #1347