[elwhyjay]

📌 Description

E5M2 output is rarely used in inference, but the current code fails silently instead of rejecting it. This PR is a small fix to close that gap.

rmsnorm_quant and fused_add_rmsnorm_quant quantize to the dtype of the out tensor and the CUDA dispatch accepts e5m2, but:

• CuTe-DSL backend (default): the norm dtype map has no float8_e5m2 entry, so an e5m2 out tensor raises KeyError; the store epilogue and the clamp were e4m3 only anyway.
• CUDA fallback (FLASHINFER_USE_CUDA_NORM=1): the kernels clamp to [-448, 448] for any output dtype, silently collapsing the e5m2 range (max 57344).

Changes:

• norm/utils.py: FLOAT8_E5M2_MAX, e5m2 conversion intrinsics (hw viacvt.rn.satfinite.e5m2x2.f32, sw via bit manipulation), thin fp8 wrappers that select the variant from a constexpr dtype, and a
  float8_e5m2 dtype map entry.
• norm/kernels/{rmsnorm,fused_add_rmsnorm}.py: quant epilogues dispatch on the output element type; the clamp bound follows the output dtype. Kernel signatures and compile cache keys are unchanged.
• include/flashinfer/norm.cuh: both quant kernels clamp with QuantTypeStaticVals<O>::MAX_VAL, as layernorm_quant already does. The trait moves above the kernels and drops its ENABLE_FP8 guard:
  the fp8 types are always visible (vec_dtypes.cuh includes cuda_fp8.h unconditionally), and these kernels are instantiated for fp8 through DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8 even in builds without ENABLE_FP8, so keeping the guard would break the no-fp8 build once they reference the trait (exercised by test_norm_compilation_without_fp8).
• Tests: test_norm_quant and test_fused_add_rmsnorm_quant parametrized over quant_dtype (e4m3fn, e5m2), same style as test_layernorm_quant.

Verified on RTX 5090: the full quant grids pass on both backends (3072 CuTe-DSL, 3085 CUDA cases). Both hw and sw conversion paths were checked bit-exact against torch fp8 casts over all finite fp16 values plus random and boundary floats, for both formats; since sm_120 always dispatches the hw path, the sw path was force-compiled with use_hw_fp8=False for this.

🔍 Related Issues

🚀 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

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

Reviewer Notes

Summary by CodeRabbit

• New Features

  • Added FP8 E5M2 support for RMSNorm quantization and fused add + RMSNorm quantization.
  • Quantized outputs now support both FP8 E4M3 and E5M2 formats.
  • Clamping automatically uses the selected output format’s valid numeric range.

• Documentation

  • Clarified supported FP8 output formats and quantization behavior.

• Tests

  • Expanded coverage across both FP8 formats, including stride and large-tensor cases.

[gemini-code-assist]

Caution

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

[coderabbitai]

📝 Walkthrough

Walkthrough

RMSNorm and fused add RMSNorm quantization now support FP8 E4M3 and E5M2 output dtypes, using dtype-specific limits and conversion helpers. Utilities, CUDA traits, DSL kernels, documentation, and tests were updated accordingly.

Changes

FP8 norm quantization

Layer / File(s)	Summary
E5M2 conversion support flashinfer/norm/utils.py	Adds E5M2 maximum-range metadata, scalar and vector conversion stores, dtype dispatch, dtype-string mapping, and exports.
Dtype-specific quantization bounds include/flashinfer/norm.cuh	Moves quantization traits earlier in the header and replaces fixed ±448.0f clipping with output-type-specific limits.
Generic FP8 kernel stores flashinfer/norm/kernels/rmsnorm.py, flashinfer/norm/kernels/fused_add_rmsnorm.py	Uses output-dtype maxima and generic hardware/software conversion helpers across vectorized and fallback quantization paths.
Public documentation and dtype coverage flashinfer/norm/__init__.py, tests/utils/test_norm.py	Documents both supported FP8 output types and parametrizes reference implementations and tests over E4M3 and E5M2.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related issues

• [Bug]: Legacy CUDA RMSNorm quantization clips E5M2 output to the E4M3 range #4053: Replaces hard-coded FP8 clipping with output-dtype-specific limits and adds E5M2 coverage.

Possibly related PRs

• flashinfer-ai/flashinfer#2243: Earlier FP8 norm kernel work that this change generalizes beyond E4M3.
• flashinfer-ai/flashinfer#3758: Updates FP8 quantization reference math in related norm paths.
• flashinfer-ai/flashinfer#4056: Related dtype-specific FP8 clipping and test coverage changes.

Suggested labels: run-ci, 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 summarizes the main change: adding FP8 E5M2 output support to the two quantized norm kernels.
Description check	✅ Passed	The description follows the template well, includes a detailed change summary, checklist status, and updated tests.
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.}

[coderabbitai]

🧹 Nitpick comments (2)

flashinfer/norm/utils.py (1)

544-602: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Dispatch wrappers silently default unknown dtypes to E4M3.

cvt_and_store_f32_to_fp8_hw/_sw and the vectorized variants only special-case cutlass.Float8E5M2; anything else silently routes to the E4M3 path instead of raising, unlike get_fp8_max which raises ValueError for unsupported dtypes. This is currently safe because callers always call get_fp8_max(mY.element_type) first (which would raise), but the wrappers themselves don't enforce that invariant independently.

♻️ Optional: make the fallback explicit

 def cvt_and_store_f32_to_fp8_hw(val: Float32, addr: Int64, dtype: cutlass.Constexpr):
     if cutlass.const_expr(dtype is cutlass.Float8E5M2):
         cvt_and_store_f32_to_e5m2_hw(val, addr)
-    else:
+    elif cutlass.const_expr(dtype is cutlass.Float8E4M3FN):
         cvt_and_store_f32_to_e4m3_hw(val, addr)
+    else:
+        raise ValueError(f"Unsupported FP8 output dtype: {dtype}")

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/norm/utils.py` around lines 544 - 602, Update the FP8 dispatch
wrappers cvt_and_store_f32_to_fp8_hw, cvt_and_store_f32_to_fp8_sw,
cvt_and_store_8xf32_to_fp8_hw, cvt_and_store_4xf32_to_fp8_hw, and
cvt_and_store_2xf32_to_fp8_hw to explicitly accept only cutlass.Float8E5M2 or
cutlass.Float8E4M3; raise ValueError for any other dtype instead of silently
selecting the E4M3 path.

include/flashinfer/norm.cuh (1)

36-61: 🗄️ Data Integrity & Integration | 🔵 Trivial

CI coverage reminder for ENABLE_FP8-disabled builds.

The trait relocation is needed for the fp8 output instantiations, but this repo no longer has an ENABLE_FP8 build toggle and a related norm.cuh/reduceKernelUtils.cuh build issue still exists. Add coverage of non-ENABLE_FP8 builds with fp8 output types, or update the build matrix to ensure this path does not regress.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@include/flashinfer/norm.cuh` around lines 36 - 61, Update the build/test
matrix to cover configurations without ENABLE_FP8 while still instantiating fp8
output types, including the norm.cuh and reduceKernelUtils.cuh path. Ensure the
non-ENABLE_FP8 configuration builds successfully and remains protected against
regressions.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@flashinfer/norm/utils.py`:
- Around line 544-602: Update the FP8 dispatch wrappers
cvt_and_store_f32_to_fp8_hw, cvt_and_store_f32_to_fp8_sw,
cvt_and_store_8xf32_to_fp8_hw, cvt_and_store_4xf32_to_fp8_hw, and
cvt_and_store_2xf32_to_fp8_hw to explicitly accept only cutlass.Float8E5M2 or
cutlass.Float8E4M3; raise ValueError for any other dtype instead of silently
selecting the E4M3 path.

In `@include/flashinfer/norm.cuh`:
- Around line 36-61: Update the build/test matrix to cover configurations
without ENABLE_FP8 while still instantiating fp8 output types, including the
norm.cuh and reduceKernelUtils.cuh path. Ensure the non-ENABLE_FP8 configuration
builds successfully and remains protected against regressions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: b5439de7-ed72-474d-bbd4-34907bb17e2c

📥 Commits

Reviewing files that changed from the base of the PR and between f90e9c4 and bbaf5cd.

📒 Files selected for processing (6)

• flashinfer/norm/__init__.py
• flashinfer/norm/kernels/fused_add_rmsnorm.py
• flashinfer/norm/kernels/rmsnorm.py
• flashinfer/norm/utils.py
• include/flashinfer/norm.cuh
• tests/utils/test_norm.py

[bkryu]

/bot run tests/utils

[flashinfer-bot]

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