feat: fused FP8 quantization output for LayerNorm - #3962
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds ChangesLayerNorm FP8 quantization
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant PythonAPI as flashinfer.norm.layernorm_quant
participant FFIBinding as norm.cu FFI binding
participant CUDAKernel as norm::LayerNormQuant
Caller->>PythonAPI: input, gamma, beta, scale, eps
PythonAPI->>FFIBinding: dispatch layernorm_quant
FFIBinding->>CUDAKernel: normalize and quantize to FP8
CUDAKernel-->>Caller: write quantized output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the layernorm_quant API, which performs Layer Normalization followed by FP8 quantization, including its CUDA kernels, Python bindings, benchmark routines, trace templates, and tests. The review feedback highlights two important issues: a potential signed/unsigned comparison warning in the C++ code (csrc/norm.cu) when comparing input.stride(0) with hidden_size, and a runtime AttributeError in the benchmark script (benchmarks/routines/norm.py) caused by referencing torch.float32.itemsize instead of element_size().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| TVM_FFI_ICHECK_EQ(input.stride(0), hidden_size) << "input must be contiguous"; | ||
| TVM_FFI_ICHECK_EQ(output.stride(0), hidden_size) << "output must be contiguous"; |
There was a problem hiding this comment.
Comparing input.stride(0) (which is int64_t) with hidden_size (which is unsigned int) can trigger signed/unsigned comparison warnings (e.g., -Wsign-compare). If -Werror is enabled, this will cause compilation to fail. It is safer to cast hidden_size to int64_t to avoid any compiler warnings.
TVM_FFI_ICHECK_EQ(input.stride(0), static_cast<int64_t>(hidden_size)) << "input must be contiguous";
TVM_FFI_ICHECK_EQ(output.stride(0), static_cast<int64_t>(hidden_size)) << "output must be contiguous";
| num_elements = np.prod(input_shape) | ||
| problem_bytes = ( | ||
| num_elements * input_dtype.itemsize # input read | ||
| + 2 * hidden_size * torch.float32.itemsize # gamma and beta read |
There was a problem hiding this comment.
In standard PyTorch, torch.dtype objects (like torch.float32) do not have an itemsize attribute. Using torch.float32.itemsize will raise an AttributeError at runtime. It is safer and more standard to use torch.float32.element_size() or simply the constant 4 (since float32 is always 4 bytes).
| + 2 * hidden_size * torch.float32.itemsize # gamma and beta read | |
| + 2 * hidden_size * 4 # gamma and beta read (float32 is 4 bytes) |
There was a problem hiding this comment.
torch.dtype does not seem to have element_size(). I also ran this routine end to end with --refcheck on torch 2.9 and it completed without errors.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
flashinfer/norm/__init__.py (1)
534-536: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExport
layernorm_quantin__all__.The new public API
layernorm_quantis missing from the module's__all__list (located around line 1783 in the file). Please export it sofrom flashinfer.norm import *behaves correctly and documentation engines capture it organically.🤖 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/__init__.py` around lines 534 - 536, Update the module’s __all__ list to include the public layernorm_quant function defined by the layernorm_quant declaration, preserving the existing export style and ordering so wildcard imports and documentation discovery expose it.
🤖 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.
Inline comments:
In `@csrc/norm.cu`:
- Around line 301-304: Add an explicit float32 dtype assertion for the scale
tensor alongside the existing checks in the layernorm binding, before its data
pointer is cast to float*. Use the scale symbol and match the validation style
and error behavior of the input, gamma, and beta checks.
In `@flashinfer/trace/templates/norm.py`:
- Around line 661-664: The quantized output tensor in the norm trace template is
missing its dtype source. In flashinfer/trace/templates/norm.py lines 661-664,
update the “out” Tensor constructor to use dtype_from="out"; then regenerate
tests/trace/fi_trace_out/layernorm_quant_h768.json at line 53 so its dtype is
float8_e4m3fn.
---
Nitpick comments:
In `@flashinfer/norm/__init__.py`:
- Around line 534-536: Update the module’s __all__ list to include the public
layernorm_quant function defined by the layernorm_quant declaration, preserving
the existing export style and ordering so wildcard imports and documentation
discovery expose it.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9b902070-3262-4b60-bf63-4680c9c2ac6b
📒 Files selected for processing (13)
benchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/norm.pycsrc/flashinfer_norm_binding.cucsrc/norm.cudocs/api/norm.rstflashinfer/__init__.pyflashinfer/norm/__init__.pyflashinfer/trace/templates/norm.pyinclude/flashinfer/norm.cuhinclude/flashinfer/trtllm/common/cudaTypeUtils.cuhtests/trace/example.pytests/trace/fi_trace_out/layernorm_quant_h768.jsontests/utils/test_norm.py
|
Gentle ping @kahyunnam @yzh119 @bkryu @yongwww .Would you be able to review this when you get a chance? Thank you. |
|
/bot run tests/utils |
|
Thank you @elwhyjay for the PR. Adding the motivation in the description helps us understand. @kahyunnam , the PR looks generally good to me, but can you take another pass? |
kahyunnam
added
the
op: misc
|
[FAILED] Pipeline #58198596: 13/20 passed |
kahyunnam
left a comment
kahyunnam
left a comment
There was a problem hiding this comment.
small nit but overall looks good! approving
| y = torch.empty_like(x, dtype=quant_dtype) | ||
| flashinfer.norm.layernorm_quant(y, x, gamma, beta, scale) | ||
|
|
||
| torch.testing.assert_close(y_ref.float(), y.float(), rtol=1, atol=1) |
There was a problem hiding this comment.
rtol/atol=1 seems a bit loose?
There was a problem hiding this comment.
Thank you for the review @kahyunnam . I agree with it. rtol/atol=1 only bounds a per-element FP8 step, so a systematic off-by-one would still pass. Since ref and kernel share the same quantize path, how about checking the fraction of mismatched elements instead?
tol = torch.finfo(quant_dtype).eps * y_ref.float().abs() # one FP8 step at |ref|
mismatched = ((y_ref.float() - y.float()).abs() > tol).float().mean().item()
assert mismatched < 0.01This keeps the FP8 rounding slack but catches a systematic error. Let me know what you think.
There was a problem hiding this comment.
finfo(fp8).eps can change a bit depending on e4m3 vs e5m2; let's reuse an existing helper: flashinfer.trace.default_check + default_tolerances (FP8 defaults: rtol=0.1, atol=0.1) supports max_mismatch_pct. Trace templates for FP8 norm use the same family of checks.
also, the ref and kernel slightly differ: right now the pytest ref is fp32-only, let's round normalized output to bf16 before / scale and FP8 cast (same as the benchmark refcheck)
|
Thank you for your efforts here @elwhyjay ! Setting this to automerge when pre-merge passes. |
4 participants
📌 Description
FlashInfer provides fused norm + FP8 quantization for the RMSNorm family (
rmsnorm_quant,fused_add_rmsnorm_quant), but not for LayerNorm. Models that use LayerNorm (OPT, Falcon, BERT family, ViT/DiT variants, Whisper/CLIP style encoders) currently have to run layernorm and quantization as two kernels, which costs about 7 bytes of traffic per element instead of 3.The kernel side is already there:
generalLayerNormininclude/flashinfer/norm.cuhimplements per-tensor and per-token quantization paths, but theLayerNormhost launcher passes nullptr for all quant arguments, with a note:This PR connects the existing per-tensor path rather than writing a new kernel, so the diff stays small.
flashinfer.layernorm_quant(out, input, gemma, beta, scale, eps). Output dtype (float8_e4m3fn or float8_e5m2) is taken from the preallocatedouttensor, same asrmsnorm_quant.QuantTypeStaticValsspecializations (constants match TRT-LLMquantTypeUtils.cuh) and the missing e5m2cuda_castspecializations. The FP8 dispatch macro instantiates both e4m3 and e5m2, so the e5m2 casts are needed for compilation.LayerNormQuanthost launcher, guarded byENABLE_FP8. The csrc launcher keeps a non-FP8 build working (test_norm_compilation_without_fp8passes).generalLayerNormto be applied as division (out = normed / scale). The kernel originally multiplied, but this path was unreachable from every existing call site (the pointer is always nullptr), so no current behavior changes. Keeping the multiply semantics would makelayernorm_quantandrmsnorm_quantdisagree on whatscalemeans, which seemed worse for fusion passes that target both. Happy to adjust if there is a concern I am missing here.layernorm_quantbenchmark routine.Performance
RTX 5090, bf16 input, e4m3 output, eps 1e-6. Fused kernel vs. unfused baseline (
flashinfer.layernormfollowed by eager torch div/clamp/cast), median over 100 iterations under CUDA graph:The eager baseline materializes intermediate tensors for div/clamp/cast, so large shapes gain more than the pure traffic ratio (7B vs 3B per element, about 2.3x). At batch 4096 / hidden 16384 the fused kernel reaches about 1.35 TB/s on this card.
Reproduce with:
🔍 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
tests/utils/test_norm.py::test_layernorm_quant: e4m3 and e5m2, batch 1 to 989, hidden 111 to 16384 (odd sizes cover the non-vectorized path), scale 0.01/1.0/10.0, tolerance rtol=1/atol=1 following the existing FP8 norm tests. 168 cases pass.test_layernorm_quant_invalid_inputs: rejects fp16 input, non-contiguous input, and non-scalar scale.test_norm_compilation_without_fp8passes unchanged (norm module still builds withoutENABLE_FP8).Full
tests/utils/test_norm.pypasses (2887 tests).pytest tests/trace/template consistency and end-to-end tests pass for the new trace.CUDA graph capture/replay smoke test passes; scale is read on device at kernel execution time, so updating the scale tensor between replays works.
Tests have been added or updated as needed.
All tests are passing (
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
flashinfer.layernorm_quant(LayerNorm + FP8 quantization) with CUDA support across supported compute capabilities.flashinfer.normand the top-levelflashinfernamespace.layernorm_quant.layernorm_quant.