[TRTLLM-12373][feat] RMSNorm nvfp4 quant fusion for DS V3.2 / Kimi-K2.5 by JunyiXu-nv · Pull Request #14848 · NVIDIA/TensorRT-LLM

JunyiXu-nv · 2026-06-02T03:21:04Z

Summary by CodeRabbit

Release Notes

New Features
- Added FP4 quantization support for RMSNorm operations with optional residual and bias handling.
- Introduced fused RMSNorm + FP4 quantization kernels for improved performance.
- Enhanced attention module to support FP4-quantized activations with proper scaling factors.
Tests
- Added comprehensive test coverage for fused RMSNorm FP4 quantization operations and slicing functionality.

Description

On NVFP4 DeepSeek-V3.2 / Kimi-K2.5, every MLA RMSNorm that feeds a static-NVFP4
Linear is immediately followed by a standalone NVFP4 input-quantize. This PR
folds that RMSNorm -> fp4_quantize pair into a single kernel that computes the
RMSNorm in registers and emits the NVFP4 (packed E2M1 + per-block E4M3 scale)
representation in the same pass, eliminating the separate quantize kernel's full
read/write + launch.

Folded edges (attention-DP path):

layer-boundary residual_add -> RMSNorm -> next layer's kv_a_proj
layer-0 prologue (residual-less) input_layernorm -> kv_a_proj
dense post_attention_layernorm -> gate_up_proj
residual-less q_a_layernorm -> q_b_proj (reads the q column-slice of
kv_a_proj_with_mqa's output in place, no .contiguous() copy)

Kernel-time impact

Measured on GB200, on-GPU durations from nsys single-layer slices
(DSv3.2 DEP8, real MLA shapes), per fused edge:

Edge	unfused: RMSNorm + standalone quant	fused	saved
Dense layer	3.33 + 3.83 = 7.16 us	3.64 us	-3.52 us (~2x)
MoE layer	3.48 + 3.32 = 6.80 us	3.42 us	-3.38 us (~2x)

Per decoder layer that is ~7.3 us saved (dense, 3 edges) / ~3.5 us (MoE, 2 edges)
of GPU kernel time, plus the launch-count reduction (-4 kernels/dense-layer,
-3/MoE-layer).

Per-decoder-layer kernel composition (nsys single-layer slice)

Unfused = RMSNormKernel (flashinfer) + quantize_with_block_size (standalone
NVFP4); Fused = rms_norm_fp4_quant_kernel.

Layer	unfused norm+quant total	fused total	saved/layer
Dense	4 RMSNorm (13.31) + 5 quant (19.14) = 32.45 us	3 fused (10.91) + residual (14.24) = 25.15 us	~7.3 us
MoE	4 RMSNorm (13.92) + 6 quant (19.94) = 33.86 us	2 fused (6.85) + residual (23.56) = 30.40 us	~3.5 us

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

JunyiXu-nv · 2026-06-02T07:27:07Z

/bot run

tensorrt-cicd · 2026-06-02T07:33:40Z

PR_Github #51547 [ run ] triggered by Bot. Commit: 87e9fee Link to invocation

tensorrt-cicd · 2026-06-02T07:58:51Z

PR_Github #51547 [ run ] completed with state FAILURE. Commit: 87e9fee
/LLM/main/L0_MergeRequest_PR pipeline #40942 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

… MLA Fold the standalone NVFP4 input-quantize (FP16/BF16 -> NVFP4) into its producing RMSNorm at the DSv3.2/Kimi-K2.5 MLA sites, on the attention-DP path: - layer-boundary residual-add + RMSNorm -> next layer's kv_a_proj - layer-0 prologue (residual-less) input_layernorm -> kv_a_proj - dense post_attention_layernorm -> gate_up_proj - residual-less q_a_layernorm -> q_b_proj (reads a row-strided column slice of kv_a_proj_with_mqa's output without a contiguous copy) The fused kernel and its two thop ops (fused_add_rmsnorm_fp4_quantize / fused_rmsnorm_fp4_quantize) live in self-contained new files (kernels/rmsNormFp4QuantKernels.{h,cu}, thop/rmsNormFp4Quant.cpp) with their own RmsNormFp4QuantParams struct and a private copy of the small reduce_fusion device helpers, so the AllReduce files are untouched. RMSNorm.forward dispatches the fold via a single is_nvfp4/nvfp4_scale branch with a problem-size gate (_WS_M_THRESHOLD=4096, calibrated on GB200 at N=7168): the warp-specialized fused_add_rms_norm_quant serves the large contiguous residual edge, the one-CTA-per-row reduce_fusion ops serve small-M, residual-less, and row-strided edges. The fold is self-gating -- it fires only when the consuming Linear is static-NVFP4 (a scale is attached) -- so no env var or extra flag. Producing norms attach .nvfp4_scale at post_load_weights, matching the llama/nemotron convention. fused_add_rmsnorm_fp4_quantize returns the residual sum as a fresh tensor instead of mutating hidden_states in place: the op took hidden_states as Tensor(a!) and returned that same tensor, which torch.compile's functionalization pass rejects (an output may not alias an input). Under torch_compile=True this produced a graph-break that failed to pickle over the executor IPC, hanging the proxy. The thop now seeds a fresh empty_like buffer, so the schema drops the (a!) alias and the op is functionalizable; the kernel is unchanged. register_fake impls for both ops let torch.compile fake-trace them. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

Address review feedback: allocate residual_out (and the add-op's norm_out) with at::detail::empty_cuda instead of torch::empty_like, matching quant_out/scale_out and the ws op (fusedAddRMSNormQuant.cpp). empty_like carries the input's layout and can lower to a separate node in the torch.compile trace; empty_cuda is a plain fresh allocation. The register_fake impl uses new_empty for the same reason. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

…ant kernel Address review feedback: - Launch rmsNormFp4QuantKernel through tensorrt_llm::common::launchWithPdlWhenEnabled (envUtils.h) so the kernel's existing PDL primitives (cudaGridDependencySynchronize / cudaTriggerProgrammaticLaunchCompletion) are actually enabled under TLLM_ENABLE_PDL. - Rename rms_norm_fp4_quant_kernel -> rmsNormFp4QuantKernel and launch_rms_norm_fp4_quant_kernel -> launchRmsNormFp4QuantKernel to match the codebase lowerCamelCase kernel-naming convention. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

…logue check Address review feedback on the fused RMSNorm-quant tests: - rms_norm_ref now calls the production flashinfer RMSNorm kernel (the exact unfused-path kernel the fusion replaces) when available, instead of a PyTorch formula, so the cross-path comparison is apples-to-apples. - Add assert_fp4_bitexact_from_norm: re-quantizing the fused op's own returned norm_out with fp4_quantize must reproduce its quant_out/scale_out bit-for-bit (FP4 compared directly; SF compared on the unswizzled valid [m, n/16] region to skip the uninitialized swizzle padding). This proves the epilogue is exact once the normed input is fixed; the remaining RMSNorm ULP differences across the two paths stay covered by the match-rate check. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

…eference Call torch.ops.trtllm.fp4_quantize with the exact production static-NVFP4 signature from NVFP4LinearMethod._input_prepare (linear.py): fp4_quantize(input, input_scale, scaling_vector_size, False), relying on isSfSwizzledLayout=True default. The bit-exact fused-epilogue check (assert_fp4_bitexact_from_norm) thus compares against the same quant op production runs. Behavior unchanged (the dropped explicit True equals the schema default). Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

Address review feedback on readability of the fused-quant scale attachment: - nvfp4_scale is now declared as an Optional attribute initialized to None in RMSNorm.__init__, with a docstring explaining why it cannot be set at construction time (the consuming Linear's calibrated input_scale only exists after weight loading; the owning model attaches it in post_load_weights / MLA._resolve_qa_fused_scale). - Replace all getattr(..., "nvfp4_scale")/hasattr probing with direct attribute access (rms_norm.py, modeling_deepseekv3.py, mla.py) — including the two pre-existing hasattr checks in modeling_nemotron_h.py, which the explicit attribute would otherwise turn into always-True. - Docstring nit: double -> single backticks in _static_nvfp4_input_scale / is_static_nvfp4_input_eligible. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

…into params struct Address review feedback on rmsNormFp4QuantKernels: 1. Fix partial-warp __shfl_xor_sync UB in the quantize epilogue: for hidden_size values where the last iteration leaves a warp partially active (e.g. 32, 128, 8208), lanes past the bound exited the loop while their neighbors entered cvt_warp_fp16_to_fp4's full-mask shuffles. The loop bound is now warp-uniform (offset - lane_id * kPackedSize), so all 32 lanes stay converged; inactive lanes are masked on loads, stores, and the SF output pointer and feed zeros to the shuffle. Safe because hidden_size % 16 == 0 pair-aligns the active region (an inactive lane's xor-1 partner is always inactive too). 2. Move the remaining direct kernel arguments (quant_out, scale_out, norm_out, scale_factor_ptr, sf_layout, input_row_stride) into RmsNormFp4QuantParams so the kernel, launcher, and thop callers all use a single params struct; launcher signature is now residualRmsNormFp4Quant(params, dtype, stream). Also add an explicit null-check on the mandatory quant/scale output pointers. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

…stale comment link Address review nits: - Replace the magic number 16 in the FP4 padding-slice helper with a named NVFP4_SF_VEC_SIZE constant in quantization/utils/fp4_utils.py (aligned with SF_VEC_SIZE in cpp/tensorrt_llm/kernels/quantization.h). - Remove a stale reference to a local benchmark script from the _WS_M_THRESHOLD comment in rms_norm.py; the measurement summary stays. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

… move slice helper into mla.py Address review nits: - Rename Fp4QuantizedTensor.bf16_hidden_states to unquantized_hidden_states since the field can carry either a BF16 or FP16 tensor. - Move _slice_hidden_states_to_num_tokens from attention.py to mla.py (its only user) and drop the imports attention.py no longer needs. - Drop a comment in RMSNorm.__init__ that referred to "this PR"; the block comment above it already documents the SM-support behavior. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

JunyiXu-nv · 2026-07-14T03:41:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-14T03:48:12Z

PR_Github #59092 [ run ] triggered by Bot. Commit: a101813 Link to invocation

tensorrt-cicd · 2026-07-14T07:13:13Z

PR_Github #59092 [ run ] completed with state SUCCESS. Commit: a101813
/LLM/main/L0_MergeRequest_PR pipeline #47607 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

JunyiXu-nv · 2026-07-14T07:42:52Z

/bot run

juney-nvidia

Approved since this PR has already received the approvals before and gets invalidated due to the latest changes of Code owner rule.

tensorrt-cicd · 2026-07-14T07:54:43Z

PR_Github #59147 [ run ] triggered by Bot. Commit: a101813 Link to invocation

tensorrt-cicd · 2026-07-14T10:40:43Z

PR_Github #59147 [ run ] completed with state SUCCESS. Commit: a101813
/LLM/main/L0_MergeRequest_PR pipeline #47655 completed with status: 'SUCCESS'

CI Report

Link to invocation

github-actions Bot assigned JunyiXu-nv Jun 2, 2026

JunyiXu-nv force-pushed the junyix/f1-residualless-rmsnorm-quant-fusion branch from 1914278 to 87e9fee Compare June 2, 2026 05:46

JunyiXu-nv changed the title ~~[None][feat] NVFP4 input-quant fusions for DeepSeek-V3.2 / Kimi-K2.5 MLA~~ [TRTLLM-12373][feat] NVFP4 input-quant fusions for DeepSeek-V3.2 / Kimi-K2.5 MLA Jun 2, 2026

JunyiXu-nv force-pushed the junyix/f1-residualless-rmsnorm-quant-fusion branch 11 times, most recently from 7b79444 to 10513bb Compare June 9, 2026 05:43

JunyiXu-nv changed the title ~~[TRTLLM-12373][feat] NVFP4 input-quant fusions for DeepSeek-V3.2 / Kimi-K2.5 MLA~~ [TRTLLM-12373][feat] RMSNorm nvfp4 quant fusion for DeepSeek-V3.2 / Kimi-K2.5 MLA Jun 9, 2026

JunyiXu-nv changed the title ~~[TRTLLM-12373][feat] RMSNorm nvfp4 quant fusion for DeepSeek-V3.2 / Kimi-K2.5 MLA~~ [TRTLLM-12373][feat] RMSNorm nvfp4 quant fusion for DS V3.2 / Kimi-K2.5 Jun 9, 2026

JunyiXu-nv force-pushed the junyix/f1-residualless-rmsnorm-quant-fusion branch from 10513bb to 26b70b6 Compare June 9, 2026 07:54

JunyiXu-nv marked this pull request as ready for review June 9, 2026 08:16

JunyiXu-nv requested review from a team as code owners June 9, 2026 08:16

JunyiXu-nv requested review from 2ez4bz, hlu1, pengbowang-nv and xxi-nv June 9, 2026 08:16

JunyiXu-nv added 9 commits July 13, 2026 20:05

yuxianq approved these changes Jul 14, 2026

View reviewed changes

JunyiXu-nv force-pushed the junyix/f1-residualless-rmsnorm-quant-fusion branch from cbe6091 to a101813 Compare July 14, 2026 03:19

JunyiXu-nv requested review from a team as code owners July 14, 2026 03:19

JunyiXu-nv requested review from chienchunhung and nv-xtf July 14, 2026 03:19

juney-nvidia approved these changes Jul 14, 2026

View reviewed changes

JunyiXu-nv removed request for chienchunhung and nv-xtf July 15, 2026 02:51

JunyiXu-nv merged commit 501777a into NVIDIA:main Jul 15, 2026
7 checks passed

Wanli-Jiang mentioned this pull request Jul 15, 2026

[TRTLLM-14026][feat] BREAKING: Remove C++ modules for legacy TRT backend #16369

Merged

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[TRTLLM-12373][feat] RMSNorm nvfp4 quant fusion for DS V3.2 / Kimi-K2.5#14848

[TRTLLM-12373][feat] RMSNorm nvfp4 quant fusion for DS V3.2 / Kimi-K2.5#14848
JunyiXu-nv merged 9 commits into
NVIDIA:mainfrom
JunyiXu-nv:junyix/f1-residualless-rmsnorm-quant-fusion

JunyiXu-nv commented Jun 2, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

JunyiXu-nv commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

juney-nvidia left a comment

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants

Uh oh!

Conversation

JunyiXu-nv commented Jun 2, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Release Notes

Description

Kernel-time impact

Per-decoder-layer kernel composition (nsys single-layer slice)

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

JunyiXu-nv commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

juney-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants

JunyiXu-nv commented Jun 2, 2026 •

edited by coderabbitai Bot

Loading