fix(moe): keep per-token FP8 quant scales finite for tiny activations

[yilin-void]

📌 Description

The CUTLASS fused-MoE per-token FP8 quantization paths currently compute the
E4M3 quantization scale as 448.0f / row_amax. For a positive activation amax
smaller than 448.0f / FLT_MAX (approximately 1.317e-36), that division
overflows to infinity. Multiplying the row by the infinite scale can then
produce non-finite FP8 values and poison the entire MoE output row.

Qwen3.5-397B model-level impact

We found this failure while evaluating Qwen3.5-397B-A17B with MXFP4 expert
weights and FP8 activations on 8 GPUs. All results below use the same 1,314
GSM8K examples, 5-shot prompts, raw completion API, and deterministic decoding:

Configuration	Correct	Accuracy
BF16	1264 / 1314	96.1948%
Official FP8 checkpoint	1247 / 1314	94.9011%
Calibrated WINT4 x FP8	1251 / 1314	95.2055%
MXFP4 x FP8, overflowing kernel	1155 / 1314	87.8995%
MXFP4 x FP8, kernel-safe overflow guard	1244 / 1314	94.6728%

Preventing the overflow recovered 89 answers, or 6.7732 percentage points.
The corrected MXFP4 result was within 3 answers of FP8 and 7 answers of the
calibrated WINT4 run. The remaining MXFP4-vs-WINT4 difference was not
significant in the paired comparison (exact McNemar p=0.44263).

The failure was traced to a routed, effectively dormant expert in layer 5. Its
gate/up and down-projection weights had RMS values around 1e-20, producing a
post-SwiGLU row around 1e-38. The subsequent 448 / row_amax overflow
generated 4,096 non-finite values, poisoning one complete hidden row. Replaying
that saved row with a finite-scale kernel reduced the non-finite count from
4,096 to zero; the local-MoE output had 1.455% relative RMSE and 0.999894 cosine
similarity against the safe reference.

The full-model 1244 / 1314 result above used an earlier kernel-safe guard at
the same two quantization sites that maps the overflowing tiny row to a finite
zero-quantization path. This PR uses a less destructive denominator clamp and
saturates the scale at FLT_MAX. The focused H200 tests below validate the
exact implementation in this PR; the full-model result is included as
model-level evidence for the root cause and the impact of removing the
non-finite path.

This PR:

• adds one device helper that clamps the denominator to
  448.0f / FLT_MAX, guaranteeing a finite quantization scale;
• uses the helper for both the expanded GEMM1 input and the post-activation
  GEMM2 input quantization paths; and
• adds Hopper regression coverage for tiny nonzero rows at both stages.

Zero rows retain the existing scale of 1.0f, and rows above the overflow
threshold are numerically unchanged. Tiny nonzero rows use the largest finite
FP32 scale, avoiding non-finite values in the downstream MoE computation.

🔍 Related Issues

N/A

🚀 Pull Request Checklist

✅ Pre-commit Checks

• Changed Python file passes ruff check and ruff format --check.
• Changed CUDA header passes the repository's clang-format 19.1.1 hook.
• git diff --check passes.

🧪 Tests

• Tests have been added or updated as needed.
• H200: test_moe_fp8_mxfp4_humming_tiny_amax_stays_finite — 2 passed
  (input and post_activation).
• H200: existing
  test_moe_fp8_mxfp4_humming_prescale_hopper_correctness[False-small] —
  passed.

The new post-activation case has a reference SwiGLU amax of approximately
3.806e-37. With the original division, all 512 output elements are NaN and
the regression test fails. With this patch, the output remains finite and the
test passes.

🔬 Experimental Track

Not applicable; this PR does not add or change an experimental API/backend.

Reviewer Notes

The failure was originally isolated from a dormant expert in a
Qwen3.5-397B-A17B W4A8 run. Its near-zero post-SwiGLU row exercised the same
overflow path reproduced by the focused test in this PR.

Summary by CodeRabbit

• Bug Fixes

  • Improved FP8 quantization scaling for extremely small or non-positive activation values.
  • Prevented scale underflow and overflow, preserving finite, nonzero, numerically useful outputs in FP8/MXFP4 mixture-of-experts computations.
  • Improved handling of residual values during input and post-activation processing.

• Tests

  • Expanded regression coverage with deterministic tiny-value scenarios and independent reference comparisons.
  • Verified finite, nonzero results and numerical agreement for affected computations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: fd69111a-a279-4cd3-96d5-dc0eba297da5

📥 Commits

Reviewing files that changed from the base of the PR and between a70bda3 and 0501a31.

📒 Files selected for processing (2)

• csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh
• tests/moe/test_trtllm_cutlass_fused_moe.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

The change replaces single-value FP8 scale computation with residual-aware quantization and dequantization scale pairs. Both fused MoE FP8 paths use the new helper. The Humming regression test verifies finite, nonzero, and numerically accurate output for tiny activations.

Changes

FP8 quantization safety

Layer / File(s)	Summary
Safe FP8 scale integration csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh	Adds FP8ScalePair and computeSafeFP8ScalePair. The helper handles non-positive amax and zero residual cases, clamps scale bounds, and computes the residual-aware dequantization scale. expandInputRowsKernel and doActivationKernel use the paired scales.
Humming tiny-amax signal regression tests/moe/test_trtllm_cutlass_fused_moe.py	Uses deterministic FP4 payloads and per-stage E8M0 scales. It builds residual expert scales, computes an independent PyTorch reference, and checks finite, nonzero, and numerically close output.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Merge Risk: ⚪ Minimal · up to 0501a

This change preserves recoverable tiny FP8 activation signals in fused-MoE quantization and adds regression coverage for both quantization stages. No concrete current-head merge risk remains.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the main change: preventing non-finite per-token FP8 quantization scales for tiny activations in the MoE path.
Description check	✅ Passed	The description is complete and follows the repository template. It explains the problem, impact, implementation, related issues, checks, tests, experimental status, and reviewer notes. The checklist …

Full details: Docstring Coverage

Explanation

Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}