fix(norm): convert float2 to e4m3 directly in packed cast

[elwhyjay]

📌 Description

layernorm_quant can miss the FP8 e4m3 correctness bound when the quantization scale is small. Reported for A10/SM86 in the nightly pipeline.

The packed float2 to __nv_fp8x2_e4m3 cast in cudaTypeUtils.cuh rounded through bf16 before quantizing:

return __nv_fp8x2_e4m3(bf1622float2(float22bf162(val)));

generalLayerNorm applies the per-tensor scale before this cast, so the mantissa bits dropped by the bf16 step are scaled up by 1/scale. At scale=0.01 that is a 100x amplification, enough to push values into a neighbouring e4m3 bucket. The test reference models the bf16 rounding that the kernel does before scaling, but not this second rounding after it, so the two diverge.

This PR converts directly with __nv_cvt_float2_to_fp8x2, matching the style already used in vec_dtypes.cuh. The e5m2 packed cast already converts without a bf16 round trip, so it is unchanged, which matches the failures being e4m3 only. The cast is only reachable through generalLayerNorm, whose sole caller is LayerNormQuant, so layernorm_quant is the only affected API.

This PR also seeds test_layernorm_quant. The bf16 round trip is architecture independent, so the error was always present on the packed path; the check bounds the fraction of drifting elements, and unseeded inputs left that fraction varying per run, which is why it surfaced as an intermittent SM86 failure.

🔍 Related Issues

#4160

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Verified on an RTX A6000 (SM86, same architecture as the A10 in the nightly pipeline).

Before and after, over 30 seeds per shape at scale=0.01, e4m3, batch 1:

hidden	seeds over 1% bound (before)	after	worst mismatch (before)	after
500	4/30	0/30	1.20%	0.00%
1024	3/30	0/30	1.27%	0.00%

The two cases named in the report, hidden=500 and hidden=1024 at scale=0.01, pass after the fix.

• Full tests/utils/test_norm.py: 1543 passed, 1344 skipped. Skips are PDL cases that need Hopper or newer.
• pytest tests/trace/ -k "norm": 123 passed, 21 skipped.
• pre-commit run on the changed files: all hooks pass.

Unrelated pre-existing failures: pytest tests/trace/ -k "rope" reports 24 failures on this machine both with and without this change, so they are not caused by it.

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

• Bug Fixes

  • Improved FP8 conversion accuracy and handling by converting values directly with saturation.

• Tests

  • Stabilized quantized layer normalization test results by making input generation deterministic.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 36745bb8-63ee-490a-b90e-81b12aaa3ef1

📥 Commits

Reviewing files that changed from the base of the PR and between 4b969c9 and b8406cf.

📒 Files selected for processing (2)

• include/flashinfer/trtllm/common/cudaTypeUtils.cuh
• tests/utils/test_norm.py

---

📝 Walkthrough

Walkthrough

The FP8x2 E4M3 conversion now uses direct finite-saturated conversion instead of a BF16 intermediate. The layernorm quantization test initializes a fixed Torch RNG seed.

Changes

FP8 conversion stabilization

Layer / File(s)	Summary
Direct FP8 conversion and deterministic validation include/flashinfer/trtllm/common/cudaTypeUtils.cuh, tests/utils/test_norm.py	float2 values convert directly to __nv_fp8x2_e4m3 with finite saturation and E4M3 formatting; test_layernorm_quant uses a fixed Torch RNG seed.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related issues

• [Bug]layernorm_quant FP8 e4m3 correctness failure on A10/SM86 with small scale #4160 — Covers the BF16-intermediate conversion change and deterministic test_layernorm_quant setup.

Possibly related PRs

• flashinfer-ai/flashinfer#3962 — Extends the same FP8 conversion and layernorm_quant test path.

Suggested labels: op: misc

Suggested reviewers: saltyminty, yzh119, bkryu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly states the main change: direct float2-to-e4m3 conversion in the packed cast.
Description check	✅ Passed	The description includes the required sections, related issue, checklist, and test results, with only optional reviewer notes left blank.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 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.}

[kahyunnam]

/bot run tests/utils

[flashinfer-bot]

GitLab MR !1051 has been created, and the CI pipeline #59823527 is currently running. I'll report back once the pipeline job completes.

[kahyunnam]

Thanks for the follow up fix @elwhyjay ! Approving, I can help merge after internal CI run clears.

[bkryu]

Thank you @elwhyjay, the internal CI looks good to me