From be3c298c64de3b874b9dbba3cd2cfdab9e29da8d Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Mon, 8 Jun 2026 22:42:22 -0700
Subject: [PATCH 1/9] [None][feat] NVFP4 RMSNorm+quant fusion for DeepSeek-V3.2
 / Kimi-K2.5 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>
---
 .../kernels/rmsNormFp4QuantKernels.cu         | 415 ++++++++++++++++++
 .../kernels/rmsNormFp4QuantKernels.h          |  88 ++++
 cpp/tensorrt_llm/thop/CMakeLists.txt          |   1 +
 cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp     | 251 +++++++++++
 .../_torch/attention_backend/sparse/dsa.py    |  17 +-
 .../_torch/custom_ops/cpp_custom_ops.py       |  59 +++
 .../_torch/models/modeling_deepseekv3.py      | 111 ++++-
 tensorrt_llm/_torch/modules/attention.py      |  48 +-
 tensorrt_llm/_torch/modules/linear.py         |  19 +
 tensorrt_llm/_torch/modules/mla.py            | 207 ++++++++-
 tensorrt_llm/_torch/modules/rms_norm.py       | 241 +++++++---
 tensorrt_llm/_torch/utils.py                  |   9 +-
 .../test_lists/test-db/l0_b200.yml            |   2 +
 .../modules/test_fp4_num_tokens_slice.py      | 105 +++++
 .../test_fused_rmsnorm_fp4_quantize.py        | 407 +++++++++++++++++
 15 files changed, 1879 insertions(+), 101 deletions(-)
 create mode 100644 cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
 create mode 100644 cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h
 create mode 100644 cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
 create mode 100644 tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py
 create mode 100644 tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py

diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
new file mode 100644
index 000000000000..dc287b0752c8
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
@@ -0,0 +1,415 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "rmsNormFp4QuantKernels.h"
+#include "tensorrt_llm/common/cudaBf16Fallbacks.cuh"
+#include "tensorrt_llm/common/cudaTypeUtils.cuh"
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/kernels/quantization.cuh"
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+// Self-contained device helpers for this kernel. These mirror the small
+// reduce_fusion utilities in customAllReduceKernels.cu but are copied here (in a
+// private namespace) so this translation unit does not depend on the AllReduce
+// files at all. Kept byte-identical to the originals.
+namespace rms_norm_fp4_quant
+{
+
+static constexpr int kBytesPerAccess = 16;
+static constexpr int kWarpSize = 32;
+static constexpr int kMaxCtaSize = 1024;
+
+// Type converter that packs data format to 128-bit data type.
+using PackedFloat = union
+{
+    int4 packed;
+    float unpacked[4];
+};
+
+using PackedHalf = union
+{
+    int4 packed;
+    half2 unpacked[4];
+};
+
+template <typename T>
+struct PackedOn16Bytes
+{
+};
+
+template <>
+struct PackedOn16Bytes<float>
+{
+    using Type = PackedFloat;
+};
+
+template <>
+struct PackedOn16Bytes<half>
+{
+    using Type = PackedHalf;
+};
+
+#ifdef ENABLE_BF16
+using PackedBFloat16 = union
+{
+    int4 packed;
+    __nv_bfloat162 unpacked[4];
+};
+
+template <>
+struct PackedOn16Bytes<__nv_bfloat16>
+{
+    using Type = PackedBFloat16;
+};
+#endif
+
+// add two 128b data
+template <typename T>
+inline __device__ int4 add128b(T& a, T& b)
+{
+    T c;
+    c.unpacked[0] = a.unpacked[0] + b.unpacked[0];
+    c.unpacked[1] = a.unpacked[1] + b.unpacked[1];
+    c.unpacked[2] = a.unpacked[2] + b.unpacked[2];
+    c.unpacked[3] = a.unpacked[3] + b.unpacked[3];
+    return c.packed;
+}
+
+inline __device__ float warp_reduce_sum(float val)
+{
+    val += __shfl_xor_sync(~0, val, 16);
+    val += __shfl_xor_sync(~0, val, 8);
+    val += __shfl_xor_sync(~0, val, 4);
+    val += __shfl_xor_sync(~0, val, 2);
+    val += __shfl_xor_sync(~0, val, 1);
+    return val;
+}
+
+inline __device__ float block_reduce_sum(float val)
+{
+    __shared__ float smem[kWarpSize];
+    int lane_id = threadIdx.x % kWarpSize, warp_id = threadIdx.x / kWarpSize, warp_num = blockDim.x / kWarpSize;
+    val = warp_reduce_sum(val);
+    if (lane_id == 0)
+    {
+        smem[warp_id] = val;
+    }
+    __syncthreads();
+    val = lane_id < warp_num ? smem[lane_id] : 0.f;
+    val = warp_reduce_sum(val);
+    return val;
+}
+
+template <typename T, typename PackedStruct>
+inline __device__ float accumulate(float acc, PackedStruct& vec)
+{
+    static constexpr int kLoopNum = sizeof(PackedStruct) / sizeof(T);
+#pragma unroll
+    for (int i = 0; i < kLoopNum; ++i)
+    {
+        float v = static_cast<float>(reinterpret_cast<T*>(vec.unpacked)[i]);
+        acc += v * v;
+    }
+    return acc;
+}
+
+template <typename T, bool Affine, typename PackedStruct>
+inline __device__ int4 rms_norm(float denom, PackedStruct& vec, PackedStruct& weight)
+{
+    static constexpr int kLoopNum = sizeof(PackedStruct) / sizeof(T);
+    PackedStruct ret;
+#pragma unroll
+    for (int i = 0; i < kLoopNum; ++i)
+    {
+        float v1 = static_cast<float>(reinterpret_cast<T*>(vec.unpacked)[i]);
+        if constexpr (Affine)
+        {
+            float v2 = static_cast<float>(reinterpret_cast<T*>(weight.unpacked)[i]);
+            reinterpret_cast<T*>(ret.unpacked)[i] = static_cast<T>(v1 * denom * v2);
+        }
+        else
+        {
+            reinterpret_cast<T*>(ret.unpacked)[i] = static_cast<T>(v1 * denom);
+        }
+    }
+    return ret.packed;
+}
+
+// Fused (optional residual-add +) RMSNorm + NVFP4 input-quantize. Invoked by
+// the standalone thop ops fused_add_rmsnorm_fp4_quantize (Residual=true) and
+// fused_rmsnorm_fp4_quantize (Residual=false) on the attention-DP path.
+// Performs the residual-add and reduction, then emits a per-block
+// (SF_VEC_SIZE=16) NVFP4 representation to (quant_out, scale_out). When
+// OutNorm=true, BF16 norm_out is also written (so a downstream consumer can read
+// the un-quantized value).
+//
+// Layout assumptions (match cvt_warp_fp16_to_fp4):
+//   - Each thread accesses kPackedSize=8 BF16/half elements (one int4 = 16 B).
+//   - Two adjacent threads cover one SF_VEC_SIZE=16 block; SF is computed
+//     warp-cooperatively via __shfl_xor_sync inside cvt_warp_fp16_to_fp4.
+//   - Caller guarantees hidden_size % SF_VEC_SIZE == 0.
+template <typename T, bool Bias = false, bool Residual = false, bool Affine = false, bool UseSmem = false,
+    bool OutNorm = false>
+__global__ void rms_norm_fp4_quant_kernel(RmsNormFp4QuantParams params, void* quant_out, void* scale_out,
+    void* norm_out_ptr, float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout,
+    int input_row_stride)
+{
+    static constexpr int kPackedSize = kBytesPerAccess / sizeof(T);
+    static constexpr int kSfVecSize = 16;
+    using PackedStruct = typename PackedOn16Bytes<T>::Type;
+
+    extern __shared__ uint8_t smem_ptr[];
+    T* smem = reinterpret_cast<T*>(smem_ptr);
+
+    int const bid = blockIdx.x;
+    int const tid = threadIdx.x;
+
+    T const* bias_buffer = reinterpret_cast<T const*>(params.bias_buffer);
+    T const* residual_buffer = reinterpret_cast<T const*>(params.residual_buffer);
+    T const* weight_buffer = reinterpret_cast<T const*>(params.weight_buffer);
+    T const* intermediate_buffer = reinterpret_cast<T const*>(params.intermediate_buffer);
+    T* residual_out_buffer = reinterpret_cast<T*>(params.residual_out_buffer);
+    T* norm_out = reinterpret_cast<T*>(norm_out_ptr);
+
+    int const block_offset = bid * params.hidden_size;
+    // Input rows may be strided (e.g. a column slice of a wider projection,
+    // such as the leading q_lora_rank columns of kv_a_proj_with_mqa). When
+    // input_row_stride <= 0 it defaults to hidden_size (packed rows), making
+    // this byte-identical to all existing callers. Only the INPUT read offset
+    // uses the stride; every output (residual/norm/quant/scale) stays packed.
+    int const input_block_offset = bid * (input_row_stride > 0 ? input_row_stride : params.hidden_size);
+    int const thread_offset = tid * kPackedSize;
+
+    if constexpr (Residual)
+    {
+        residual_buffer += block_offset;
+        // residual_out is packed [m, hidden_size], so it uses the dense block
+        // offset (not the possibly-strided input offset).
+        residual_out_buffer += block_offset;
+    }
+    intermediate_buffer += input_block_offset;
+    if constexpr (OutNorm)
+    {
+        norm_out += block_offset;
+    }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200))
+    cudaGridDependencySynchronize();
+#endif
+
+    PackedStruct inter_vec, weight_vec;
+    float acc = 0.f;
+    for (int offset = thread_offset; offset < params.hidden_size; offset += blockDim.x * kPackedSize)
+    {
+        inter_vec.packed = *reinterpret_cast<int4 const*>(intermediate_buffer + offset);
+        if constexpr (Bias)
+        {
+            PackedStruct bias_vec;
+            bias_vec.packed = *reinterpret_cast<int4 const*>(bias_buffer + offset);
+            inter_vec.packed = add128b(inter_vec, bias_vec);
+        }
+        if constexpr (Residual)
+        {
+            PackedStruct residual_vec;
+            residual_vec.packed = *reinterpret_cast<int4 const*>(residual_buffer + offset);
+            inter_vec.packed = add128b(inter_vec, residual_vec);
+            // Write the residual sum to a distinct output buffer (packed offset),
+            // leaving the input intermediate_buffer untouched.
+            *reinterpret_cast<int4*>(residual_out_buffer + offset) = inter_vec.packed;
+        }
+        acc = accumulate<T>(acc, inter_vec);
+        if constexpr (UseSmem)
+        {
+            *reinterpret_cast<int4*>(&smem[offset]) = inter_vec.packed;
+        }
+    }
+    acc = block_reduce_sum(acc);
+    float const denom = rsqrtf(acc / params.hidden_size + params.eps);
+
+    float const sf_scale = scale_factor_ptr ? *scale_factor_ptr : 1.f;
+    int const hidden_dim_packed = params.hidden_size / kSfVecSize;
+
+    for (int offset = thread_offset; offset < params.hidden_size; offset += blockDim.x * kPackedSize)
+    {
+        if constexpr (UseSmem)
+        {
+            inter_vec.packed = *reinterpret_cast<int4 const*>(&smem[offset]);
+        }
+        if constexpr (Affine)
+        {
+            weight_vec.packed = *reinterpret_cast<int4 const*>(weight_buffer + offset);
+        }
+        inter_vec.packed = rms_norm<T, Affine>(denom, inter_vec, weight_vec);
+        if constexpr (OutNorm)
+        {
+            *reinterpret_cast<int4*>(norm_out + offset) = inter_vec.packed;
+        }
+
+        // FP4 quantize this 8-element packed vec; warp-cooperate with the
+        // neighbour thread (offset ^ kPackedSize) to compute a single SF for
+        // their joint 16-element block.
+        ::tensorrt_llm::kernels::PackedVec<T> pv
+            = *reinterpret_cast<::tensorrt_llm::kernels::PackedVec<T>*>(&inter_vec);
+        int const access_id = bid * (params.hidden_size / kPackedSize) + (offset / kPackedSize);
+        int const access_id_in_token = offset / kPackedSize;
+        uint8_t* sf_out_ptr = ::tensorrt_llm::kernels::cvt_quant_get_sf_out_offset<uint32_t, 2>(std::nullopt, bid,
+            access_id_in_token, std::nullopt, hidden_dim_packed, reinterpret_cast<uint32_t*>(scale_out), sf_layout);
+        reinterpret_cast<uint32_t*>(quant_out)[access_id]
+            = ::tensorrt_llm::kernels::cvt_warp_fp16_to_fp4<T, kSfVecSize, /*UE8M0_SF=*/false>(
+                pv, sf_scale, sf_out_ptr);
+    }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200))
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+template <typename T, bool OutNorm = false>
+void launch_rms_norm_fp4_quant_kernel(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out,
+    void* norm_out_ptr, float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout,
+    cudaStream_t stream, int input_row_stride = 0)
+{
+    static constexpr int kPackedSize = kBytesPerAccess / sizeof(T);
+    TLLM_CHECK(params.hidden_size % kPackedSize == 0);
+    TLLM_CHECK(params.hidden_size % 16 == 0); // SF_VEC_SIZE
+    int need_threads = params.hidden_size / kPackedSize;
+    int cta_size = need_threads <= kMaxCtaSize ? (need_threads + kWarpSize - 1) / kWarpSize * kWarpSize : kMaxCtaSize;
+    int cta_num = params.elts_total / params.hidden_size;
+    bool const need_smem = (cta_size * kBytesPerAccess / sizeof(T) < params.hidden_size);
+    int smem_size = need_smem ? params.hidden_size * sizeof(T) : 0;
+    bool const use_smem = need_smem;
+
+    bool const has_bias = params.bias_buffer != nullptr;
+    bool const has_residual = params.residual_buffer != nullptr;
+    bool const has_weight = params.weight_buffer != nullptr;
+
+    // Macro-dispatch over Bias/Residual/Affine/UseSmem and the OutNorm template arg.
+#define DISPATCH_FP4_QUANT(BIAS, RESIDUAL, AFFINE, SMEM)                                                               \
+    if (use_smem == SMEM)                                                                                              \
+    {                                                                                                                  \
+        rms_norm_fp4_quant_kernel<T, BIAS, RESIDUAL, AFFINE, SMEM, OutNorm><<<cta_num, cta_size, smem_size, stream>>>( \
+            params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, input_row_stride);                \
+    }
+
+    auto launch = [&]()
+    {
+        if (has_bias && has_residual && has_weight)
+        {
+            DISPATCH_FP4_QUANT(true, true, true, true)
+            else DISPATCH_FP4_QUANT(true, true, true, false)
+        }
+        else if (!has_bias && has_residual && has_weight)
+        {
+            DISPATCH_FP4_QUANT(false, true, true, true)
+            else DISPATCH_FP4_QUANT(false, true, true, false)
+        }
+        else if (has_bias && !has_residual && has_weight)
+        {
+            DISPATCH_FP4_QUANT(true, false, true, true)
+            else DISPATCH_FP4_QUANT(true, false, true, false)
+        }
+        else if (!has_bias && !has_residual && has_weight)
+        {
+            DISPATCH_FP4_QUANT(false, false, true, true)
+            else DISPATCH_FP4_QUANT(false, false, true, false)
+        }
+        else if (has_bias && has_residual && !has_weight)
+        {
+            DISPATCH_FP4_QUANT(true, true, false, true)
+            else DISPATCH_FP4_QUANT(true, true, false, false)
+        }
+        else if (!has_bias && has_residual && !has_weight)
+        {
+            DISPATCH_FP4_QUANT(false, true, false, true)
+            else DISPATCH_FP4_QUANT(false, true, false, false)
+        }
+        else if (has_bias && !has_residual && !has_weight)
+        {
+            DISPATCH_FP4_QUANT(true, false, false, true)
+            else DISPATCH_FP4_QUANT(true, false, false, false)
+        }
+        else
+        {
+            DISPATCH_FP4_QUANT(false, false, false, true)
+            else DISPATCH_FP4_QUANT(false, false, false, false)
+        }
+    };
+    launch();
+#undef DISPATCH_FP4_QUANT
+}
+
+} // namespace rms_norm_fp4_quant
+
+void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out, void* norm_out_ptr,
+    float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout, nvinfer1::DataType dataType,
+    cudaStream_t stream, int input_row_stride)
+{
+    // The NVFP4 epilogue (cvt_warp_fp16_to_fp4) is compiled only for
+    // __CUDA_ARCH__ >= 1000 and emits zeros otherwise, so this kernel is correct
+    // only on SM 10.x (Blackwell). Fail fast on unsupported archs rather than
+    // silently producing wrong FP4 (the Python dispatch in rms_norm.py already
+    // routes those to the unfused path).
+    int const sm = tensorrt_llm::common::getSMVersion();
+    TLLM_CHECK_WITH_INFO(sm >= 100 && sm < 120,
+        "residualRmsNormFp4Quant requires SM 10.x (Blackwell); got SM %d. The fused NVFP4 epilogue is unsupported on "
+        "this arch.",
+        sm);
+    sync_check_cuda_error(stream);
+    bool const out_norm = (norm_out_ptr != nullptr);
+    if (out_norm)
+    {
+        switch (dataType)
+        {
+#ifdef ENABLE_BF16
+        case nvinfer1::DataType::kBF16:
+            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<__nv_bfloat16, /*OutNorm=*/true>(
+                params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            break;
+#endif
+        case nvinfer1::DataType::kHALF:
+            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<half, /*OutNorm=*/true>(
+                params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            break;
+        default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant");
+        }
+    }
+    else
+    {
+        switch (dataType)
+        {
+#ifdef ENABLE_BF16
+        case nvinfer1::DataType::kBF16:
+            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<__nv_bfloat16, /*OutNorm=*/false>(
+                params, quant_out, scale_out, nullptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            break;
+#endif
+        case nvinfer1::DataType::kHALF:
+            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<half, /*OutNorm=*/false>(
+                params, quant_out, scale_out, nullptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            break;
+        default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant");
+        }
+    }
+    sync_check_cuda_error(stream);
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h
new file mode 100644
index 000000000000..a41f54b83145
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h
@@ -0,0 +1,88 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/quantization.h"
+#include <NvInferRuntime.h>
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+// Parameters for the fused (optional residual-add +) RMSNorm + NVFP4
+// input-quantize kernel. Self-contained (does not borrow AllReduceParams) since
+// this kernel performs no allreduce — it backs the standalone thop ops
+// fused_add_rmsnorm_fp4_quantize / fused_rmsnorm_fp4_quantize on the
+// attention-DP path.
+struct RmsNormFp4QuantParams
+{
+    // Input values [m, hidden_size] (read with input_row_stride). Read-only: the
+    // residual sum is written to residual_out_buffer, never back into this.
+    void const* intermediate_buffer{nullptr};
+    // residual_in [m, hidden_size]; nullptr disables the residual add.
+    void const* residual_buffer{nullptr};
+    // residual_out [m, hidden_size] (packed): receives input + residual when
+    // residual_buffer != nullptr. A distinct buffer from intermediate_buffer so
+    // the input is never mutated (keeps the thop op functionalizable under
+    // torch.compile) and no pre-kernel copy is needed. nullptr when no residual.
+    void* residual_out_buffer{nullptr};
+    // optional bias add [hidden_size]; nullptr disables it.
+    void const* bias_buffer{nullptr};
+    // RMSNorm gamma [hidden_size]; nullptr selects the non-affine path.
+    void const* weight_buffer{nullptr};
+    int hidden_size{0};
+    float eps{0.f};
+    // Total element count (= m * hidden_size); used to derive the row count.
+    int64_t elts_total{0};
+};
+
+// Fused (optional residual-add +) RMSNorm + NVFP4 input-quantize. Folds RMSNorm
+// and the next op's NVFP4 input-quant so the (flashinfer RMSNorm + standalone
+// fp4_quantize) pair becomes one launch on the attention-DP path.
+//
+// Inputs (via params):
+//   - intermediate_buffer: the input values (read-only)
+//   - residual_buffer:     residual_in (set nullptr to disable)
+//   - residual_out_buffer: receives input + residual (packed); required when
+//                          residual_buffer != nullptr, else unused
+//   - bias_buffer:         optional bias add (set nullptr to disable)
+//   - weight_buffer:       RMSNorm gamma (set nullptr for no affine)
+//   - hidden_size, eps:    standard
+//
+// Outputs:
+//   - quant_out:        packed FP4 (E2M1) values, kSfVecSize=16 per byte
+//   - scale_out:        E4M3 scaling factors (one per 16-elem block)
+//   - norm_out_ptr:     optional BF16 normed value (matches the
+//                       OUT_QUANT_NVFP4 fusion shape); pass nullptr to skip
+//   - scale_factor_ptr: device pointer to the global per-tensor scale
+//                       (= 448*6 / amax for static-quant Linear)
+//   - sf_layout:        scaling-factor layout (typically Swizzled)
+//   - input_row_stride: element stride between input rows in intermediate_buffer.
+//                       Defaults to 0 meaning "== hidden_size" (packed rows), so
+//                       every existing caller is byte-identical. Set >0 to read a
+//                       strided slice (e.g. a column-slice of a wider projection)
+//                       without a preceding contiguous copy. Outputs stay packed.
+void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out, void* norm_out_ptr,
+    float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout, nvinfer1::DataType dataType,
+    cudaStream_t stream, int input_row_stride = 0);
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt
index 369bf721b099..c628a9a152c0 100644
--- a/cpp/tensorrt_llm/thop/CMakeLists.txt
+++ b/cpp/tensorrt_llm/thop/CMakeLists.txt
@@ -82,6 +82,7 @@ add_library(
   fusedAddRMSNormQuant.cpp
   fusedActivationQuant.cpp
   fusedGatedRMSNormQuant.cpp
+  rmsNormFp4Quant.cpp
   fusedTopkSoftmax.cpp
   gatherTreeOp.cpp
   groupRmsNormOp.cpp
diff --git a/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp b/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
new file mode 100644
index 000000000000..350b8b9600ff
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
@@ -0,0 +1,251 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/kernels/quantization.h"
+#include "tensorrt_llm/kernels/rmsNormFp4QuantKernels.h"
+#include "tensorrt_llm/runtime/torchUtils.h"
+#include "tensorrt_llm/thop/thUtils.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/EmptyTensor.h>
+
+#include <vector>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+// Fused residual-add + RMSNorm + NVFP4 input-quantize in one kernel.
+// Replaces the (flashinfer fused_add_rmsnorm + standalone fp4_quantize) pair on
+// the no-allreduce / attention-DP path. Each rank operates on its local tokens.
+//
+// Inputs:
+//   hidden_states : [..., hidden_size] BF16/FP16 — read-only. The residual sum
+//                   (hidden_states + residual) is returned as a fresh
+//                   residual_out tensor; hidden_states itself is not mutated, so
+//                   the op is functionalizable under torch.compile.
+//   residual      : [..., hidden_size] same dtype, read-only.
+//   norm_weight   : [hidden_size] same dtype, RMSNorm gamma.
+//   scale_factor  : [] float32, = (448 * 6) / amax for static-quant Linear.
+//   eps           : RMSNorm epsilon.
+//   return_norm_out : when true, also return the BF16 normed value (needed by
+//                     DSA indexer's pre_indexer_proj).
+//
+// Returns: [quant_out, scale_out, residual_out] or
+//          [norm_out, quant_out, scale_out, residual_out] when return_norm_out.
+std::vector<at::Tensor> fused_add_rmsnorm_fp4_quantize(at::Tensor const& hidden_states, at::Tensor residual,
+    at::Tensor const& norm_weight, at::Tensor const& scale_factor, double eps, bool return_norm_out)
+{
+    CHECK_TH_CUDA(hidden_states);
+    CHECK_CONTIGUOUS(hidden_states);
+    CHECK_TH_CUDA(residual);
+    CHECK_CONTIGUOUS(residual);
+    CHECK_TH_CUDA(norm_weight);
+    CHECK_CONTIGUOUS(norm_weight);
+    CHECK_INPUT(scale_factor, torch::kFloat32);
+    TORCH_CHECK(
+        hidden_states.scalar_type() == residual.scalar_type(), "hidden_states and residual must have matching dtype");
+    TORCH_CHECK(hidden_states.scalar_type() == norm_weight.scalar_type(),
+        "hidden_states and norm_weight must have matching dtype");
+
+    auto const& input_shape = hidden_states.sizes();
+    auto const rank = input_shape.size();
+    TORCH_CHECK(rank >= 2, "hidden_states should be >=2D");
+    int64_t m = 1;
+    for (size_t i = 0; i < rank - 1; i++)
+    {
+        m *= input_shape[i];
+    }
+    auto const k = input_shape[rank - 1];
+    int64_t const sf_vec_size = 16;
+    TORCH_CHECK(k % sf_vec_size == 0, "hidden_size must be divisible by 16");
+
+    std::vector<int64_t> quant_shape(input_shape.begin(), input_shape.end());
+    quant_shape[rank - 1] = k / 2;
+    at::Tensor quant_out = at::detail::empty_cuda(quant_shape, FLOAT4_E2M1X2, hidden_states.device(), std::nullopt);
+    at::Tensor scale_out = at::detail::empty_cuda({tensorrt_llm::computeSwizzledLayoutSFSize(m, k / sf_vec_size)},
+        SF_DTYPE, hidden_states.device(), std::nullopt);
+
+    at::Tensor norm_out;
+    void* norm_out_ptr = nullptr;
+    if (return_norm_out)
+    {
+        norm_out = torch::empty_like(hidden_states);
+        norm_out_ptr = norm_out.mutable_data_ptr();
+    }
+
+    // The kernel reads hidden_states (intermediate_buffer) and residual, adds
+    // them, and writes the sum to a distinct residual_out_buffer — hidden_states
+    // is never mutated, so the op is functionalizable under torch.compile (no
+    // output aliases an input) with no extra pre-kernel copy.
+    at::Tensor residual_out = torch::empty_like(hidden_states);
+
+    tensorrt_llm::kernels::RmsNormFp4QuantParams params{};
+    params.bias_buffer = nullptr;
+    params.residual_buffer = residual.data_ptr();
+    params.weight_buffer = norm_weight.data_ptr();
+    params.intermediate_buffer = hidden_states.data_ptr();
+    params.residual_out_buffer = residual_out.mutable_data_ptr();
+    params.hidden_size = static_cast<int>(k);
+    params.eps = static_cast<float>(eps);
+    params.elts_total = hidden_states.numel();
+
+    auto const stream = at::cuda::getCurrentCUDAStream(hidden_states.get_device());
+    auto const dtype = tensorrt_llm::runtime::TorchUtils::dataType(hidden_states.scalar_type());
+
+    tensorrt_llm::kernels::residualRmsNormFp4Quant(params, quant_out.mutable_data_ptr(), scale_out.mutable_data_ptr(),
+        norm_out_ptr, static_cast<float const*>(scale_factor.data_ptr()), tensorrt_llm::QuantizationSFLayout::SWIZZLED,
+        dtype, stream);
+
+    // residual_out holds the residual sum (= original hidden + original residual).
+    if (return_norm_out)
+    {
+        return {norm_out, quant_out, scale_out, residual_out};
+    }
+    return {quant_out, scale_out, residual_out};
+}
+
+// Residual-less variant of fused_add_rmsnorm_fp4_quantize. Replaces the
+// (flashinfer rmsnorm + standalone fp4_quantize) pair on intra-layer paths that
+// have NO residual add — e.g. DSv3.2/Kimi-K2.5 MLA's q_a_layernorm feeding the
+// static-NVFP4 q_b_proj. The kernel reads intermediate_buffer (=hidden_states),
+// skips the residual add (residual_buffer == nullptr selects Residual=false in
+// the launcher), RMSNorms it, and FP4-quantizes the result. hidden_states is
+// NOT modified (no residual write-back happens when Residual=false).
+//
+// Inputs:
+//   hidden_states : [..., hidden_size] BF16/FP16 — read-only.
+//   norm_weight   : [hidden_size] same dtype, RMSNorm gamma.
+//   scale_factor  : [] float32, = (448 * 6) / amax for static-quant Linear.
+//   eps           : RMSNorm epsilon.
+//   return_norm_out : when true, also return the BF16 normed value.
+//
+// Returns: [quant_out, scale_out] or [norm_out, quant_out, scale_out] when
+//          return_norm_out.
+std::vector<at::Tensor> fused_rmsnorm_fp4_quantize(at::Tensor const& hidden_states, at::Tensor const& norm_weight,
+    at::Tensor const& scale_factor, double eps, bool return_norm_out)
+{
+    CHECK_TH_CUDA(hidden_states);
+    CHECK_TH_CUDA(norm_weight);
+    CHECK_CONTIGUOUS(norm_weight);
+    CHECK_INPUT(scale_factor, torch::kFloat32);
+    TORCH_CHECK(hidden_states.scalar_type() == norm_weight.scalar_type(),
+        "hidden_states and norm_weight must have matching dtype");
+
+    auto const& input_shape = hidden_states.sizes();
+    auto const rank = input_shape.size();
+    TORCH_CHECK(rank >= 2, "hidden_states should be >=2D");
+    // hidden_states may be a column-slice of a wider projection (e.g. the leading
+    // q_lora_rank columns of kv_a_proj_with_mqa): its last dim is unit-stride but
+    // its row stride may exceed hidden_size. We read it in place via an input row
+    // stride and skip the otherwise-required contiguous copy. The kernel only
+    // reads with this stride; all outputs are written packed.
+    TORCH_CHECK(hidden_states.stride(rank - 1) == 1, "hidden_states last dim must be unit-stride");
+    // All leading dims must be densely packed on top of the row pitch so that a
+    // single per-row element stride describes the flattened [m, k] layout. The
+    // only permitted non-packing is a row pitch larger than k (a column slice).
+    for (size_t i = 0; i + 2 < rank; i++)
+    {
+        TORCH_CHECK(hidden_states.stride(i) == hidden_states.stride(i + 1) * input_shape[i + 1],
+            "hidden_states leading dims must be densely packed");
+    }
+    int64_t m = 1;
+    for (size_t i = 0; i < rank - 1; i++)
+    {
+        m *= input_shape[i];
+    }
+    auto const k = input_shape[rank - 1];
+    int64_t const sf_vec_size = 16;
+    TORCH_CHECK(k % sf_vec_size == 0, "hidden_size must be divisible by 16");
+    // Element stride between consecutive logical rows. For a contiguous tensor
+    // this equals k, so input_row_stride==0 (packed) and behavior is identical.
+    int64_t const row_stride = hidden_states.stride(rank - 2);
+    int const input_row_stride = (row_stride == k) ? 0 : static_cast<int>(row_stride);
+
+    std::vector<int64_t> quant_shape(input_shape.begin(), input_shape.end());
+    quant_shape[rank - 1] = k / 2;
+    at::Tensor quant_out = at::detail::empty_cuda(quant_shape, FLOAT4_E2M1X2, hidden_states.device(), std::nullopt);
+    at::Tensor scale_out = at::detail::empty_cuda({tensorrt_llm::computeSwizzledLayoutSFSize(m, k / sf_vec_size)},
+        SF_DTYPE, hidden_states.device(), std::nullopt);
+
+    at::Tensor norm_out;
+    void* norm_out_ptr = nullptr;
+    if (return_norm_out)
+    {
+        // The kernel writes norm_out packed (stride hidden_size), so allocate a
+        // packed (contiguous) tensor rather than mirroring a possibly-strided
+        // input layout.
+        norm_out = at::detail::empty_cuda(
+            input_shape.vec(), hidden_states.scalar_type(), hidden_states.device(), std::nullopt);
+        norm_out_ptr = norm_out.mutable_data_ptr();
+    }
+
+    // residual_buffer == nullptr selects the Residual=false kernel path: the
+    // kernel RMSNorms intermediate_buffer (=hidden_states) directly without any
+    // add or write-back, so hidden_states is left unmodified.
+    tensorrt_llm::kernels::RmsNormFp4QuantParams params{};
+    params.bias_buffer = nullptr;
+    params.residual_buffer = nullptr;
+    params.weight_buffer = norm_weight.data_ptr();
+    params.intermediate_buffer = hidden_states.data_ptr();
+    params.hidden_size = static_cast<int>(k);
+    params.eps = static_cast<float>(eps);
+    params.elts_total = hidden_states.numel();
+
+    auto const stream = at::cuda::getCurrentCUDAStream(hidden_states.get_device());
+    auto const dtype = tensorrt_llm::runtime::TorchUtils::dataType(hidden_states.scalar_type());
+
+    tensorrt_llm::kernels::residualRmsNormFp4Quant(params, quant_out.mutable_data_ptr(), scale_out.mutable_data_ptr(),
+        norm_out_ptr, static_cast<float const*>(scale_factor.data_ptr()), tensorrt_llm::QuantizationSFLayout::SWIZZLED,
+        dtype, stream, input_row_stride);
+
+    if (return_norm_out)
+    {
+        return {norm_out, quant_out, scale_out};
+    }
+    return {quant_out, scale_out};
+}
+
+} // namespace torch_ext
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "fused_add_rmsnorm_fp4_quantize("
+        "Tensor hidden_states,"
+        "Tensor residual,"
+        "Tensor norm_weight,"
+        "Tensor scale_factor,"
+        "float eps,"
+        "bool return_norm_out) -> Tensor[]");
+    m.def(
+        "fused_rmsnorm_fp4_quantize("
+        "Tensor hidden_states,"
+        "Tensor norm_weight,"
+        "Tensor scale_factor,"
+        "float eps,"
+        "bool return_norm_out) -> Tensor[]");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("fused_add_rmsnorm_fp4_quantize", &tensorrt_llm::torch_ext::fused_add_rmsnorm_fp4_quantize);
+    m.impl("fused_rmsnorm_fp4_quantize", &tensorrt_llm::torch_ext::fused_rmsnorm_fp4_quantize);
+}
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 0a31adaecc08..c82a19c31543 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -27,7 +27,7 @@
     maybe_execute_in_parallel
 from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
-from tensorrt_llm._torch.utils import maybe_compile
+from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile
 from tensorrt_llm._utils import (get_size_in_bytes, get_sm_version,
                                  maybe_pin_memory, prefer_pinned)
 from tensorrt_llm.bindings import DataType
@@ -2880,7 +2880,18 @@ def pre_indexer_proj(
         """
         assert self._fused_wk_wp_weight is not None, \
             "cache_derived_state() must be called before forward()"
-        hidden_float = _to_float(hidden_states)
+        # When the boundary fusion pre-quantized the next layer's kv_a_proj
+        # NVFP4 input, hidden_states is an Fp4QuantizedTensor that also carries
+        # the BF16 post-RMSNorm value. Use that BF16 view here -- the indexer
+        # weight is BF16 and the matmul needs a float input.
+        if isinstance(hidden_states, Fp4QuantizedTensor):
+            assert hidden_states.bf16_hidden_states is not None, (
+                "pre_indexer_proj received Fp4QuantizedTensor without bf16 view; "
+                "the producer fusion must request return_norm_out=True")
+            hidden_states_bf = hidden_states.bf16_hidden_states
+        else:
+            hidden_states_bf = hidden_states
+        hidden_float = _to_float(hidden_states_bf)
         with _tf32_matmul_enabled():
             # F.linear computes input @ weight.T internally; no explicit .t() needed.
             # _fused_wk_wp_weight is [head_dim + n_heads, hidden_size] (nn.Linear convention).
@@ -2891,7 +2902,7 @@ def pre_indexer_proj(
         indexer_k, weights = fused_out.split([self.head_dim, self.n_heads],
                                              dim=-1)
         # Cast indexer_k back to model dtype for downstream ops (k_norm, RoPE, FP8 quantize)
-        indexer_k = indexer_k.to(hidden_states.dtype)
+        indexer_k = indexer_k.to(hidden_states_bf.dtype)
 
         q_pe, q_nope, k_pe, k_nope = self._qk_projection_and_rope(
             qr, indexer_k, position_ids)
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index e82d81e219b2..cc868b577cb3 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1281,6 +1281,65 @@ def _(
             (m, n), dtype=input.dtype) if output_hp_norm else None
         return normed_output_fp4, output, sf_out, hp_output
 
+    def _swizzled_sf_size(m: int, n: int, sf_vec_size: int = 16) -> int:
+        # Mirrors tensorrt_llm::computeSwizzledLayoutSFSize: padUp(rows,128) *
+        # padUp(cols,4), with cols = n / sf_vec_size (SFMatrix columns).
+        cols = n // sf_vec_size
+        padded_row = ((m + 127) // 128) * 128
+        padded_col = ((cols + 3) // 4) * 4
+        return padded_row * padded_col
+
+    @torch.library.register_fake("trtllm::fused_add_rmsnorm_fp4_quantize")
+    def _(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        scale_factor: torch.Tensor,
+        eps: float,
+        return_norm_out: bool,
+    ) -> List[torch.Tensor]:
+        # quant_out: packed NVFP4 (E2M1x2) as uint8, last dim halved, leading
+        # dims preserved. scale_out: swizzled E4M3 scale bytes. residual_out:
+        # hidden+residual, a fresh tensor (the op does not mutate hidden_states).
+        # When return_norm_out, the BF16 post-RMSNorm value leads the tuple.
+        m = 1
+        for d in hidden_states.shape[:-1]:
+            m *= d
+        k = hidden_states.shape[-1]
+        quant_shape = (*hidden_states.shape[:-1], k // 2)
+        quant_out = hidden_states.new_empty(quant_shape, dtype=torch.uint8)
+        sf_out = hidden_states.new_empty((_swizzled_sf_size(m, k), ),
+                                         dtype=torch.uint8)
+        residual_out = torch.empty_like(hidden_states)
+        if return_norm_out:
+            norm_out = torch.empty_like(hidden_states)
+            return [norm_out, quant_out, sf_out, residual_out]
+        return [quant_out, sf_out, residual_out]
+
+    @torch.library.register_fake("trtllm::fused_rmsnorm_fp4_quantize")
+    def _(
+        hidden_states: torch.Tensor,
+        norm_weight: torch.Tensor,
+        scale_factor: torch.Tensor,
+        eps: float,
+        return_norm_out: bool,
+    ) -> List[torch.Tensor]:
+        # Residual-less variant. quant_out / sf_out as above; norm_out (packed,
+        # contiguous) leads the tuple when return_norm_out.
+        m = 1
+        for d in hidden_states.shape[:-1]:
+            m *= d
+        k = hidden_states.shape[-1]
+        quant_shape = (*hidden_states.shape[:-1], k // 2)
+        quant_out = hidden_states.new_empty(quant_shape, dtype=torch.uint8)
+        sf_out = hidden_states.new_empty((_swizzled_sf_size(m, k), ),
+                                         dtype=torch.uint8)
+        if return_norm_out:
+            norm_out = hidden_states.new_empty(tuple(hidden_states.shape),
+                                               dtype=hidden_states.dtype)
+            return [norm_out, quant_out, sf_out]
+        return [quant_out, sf_out]
+
     @torch.library.register_fake("trtllm::fused_gated_rmsnorm_quant")
     def _(
         x: torch.Tensor,
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index 57627b4650e7..c10fc2c81f9b 100755
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -66,7 +66,8 @@
 from ..modules.fused_moe.routing import Deepseekv3RoutingImpl
 # isort: on
 from ..modules.gated_mlp import GatedMLP
-from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig
+from ..modules.linear import (Linear, TensorParallelMode, WeightsLoadingConfig,
+                              is_static_nvfp4_input_eligible)
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
 from ..peft.lora.layer import LoraLayer
@@ -145,6 +146,18 @@ def moe_reduce_add_shared_output(routed_output, shared_output, out=None):
     return shared_output.add_(routed_reduced)
 
 
+def _static_nvfp4_input_scale(linear):
+    """Return ``linear``'s calibrated NVFP4 input_scale if it is eligible to be
+    folded into a producing RMSNorm's fused NVFP4 quantize, else None.
+
+    Eligibility is the shared ``is_static_nvfp4_input_eligible`` predicate, also
+    used by MLA._resolve_qa_fused_scale, so every fold site agrees on what
+    "static-NVFP4 eligible" means."""
+    if is_static_nvfp4_input_eligible(linear):
+        return linear.input_scale
+    return None
+
+
 class DeepseekV3WeightLoader:
 
     def __init__(self, model, is_draft_model: bool = False):
@@ -1336,9 +1349,18 @@ def __init__(self,
                 use_cute_dsl_blockscaling_mm,
             )
 
+        # On NVFP4 models, construct the boundary/prologue norms as NVFP4 norms.
+        # post_load_weights attaches each norm's .nvfp4_scale (the consuming
+        # Linear's static input_scale); the fused (add+)RMSNorm+NVFP4-quant then
+        # fires automatically from that attribute (RMSNorm.forward) only when a
+        # scale was attached, kernel chosen by the size gate. input_layernorm
+        # serves both layer 0's residual-less prologue and, as the previous
+        # layer's next_layer_layernorm, the layer-boundary add+RMSNorm edge.
+        nvfp4_norm = "nvfp4" if self.is_nvfp4 else None
         self.input_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                        eps=config.rms_norm_eps,
-                                       dtype=config.torch_dtype)
+                                       dtype=config.torch_dtype,
+                                       quantize_type=nvfp4_norm)
 
         # When enable_attention_dp is True, we normally skip attention all-reduce since each
         # DP rank works on different batch elements. However, with CP > 1, attention is split
@@ -1350,9 +1372,14 @@ def __init__(self,
                                        or self.mapping.tp_size == 1
                                        or can_skip_for_attention_dp)
 
+        # Dense (GatedMLP) layers fold post_attention_layernorm -> NVFP4
+        # gate_up_proj; MoE layers route expert-input quant differently and
+        # never get an .nvfp4_scale attached, so the fused path stays inert for
+        # them even when constructed as an NVFP4 norm.
         self.post_attention_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                                 eps=config.rms_norm_eps,
-                                                dtype=config.torch_dtype)
+                                                dtype=config.torch_dtype,
+                                                quantize_type=nvfp4_norm)
         self.next_layer_layernorm: RMSNorm = None
 
     def _get_decoder_layer_quant_config(
@@ -1419,7 +1446,16 @@ def forward(
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+            if getattr(self.input_layernorm, "nvfp4_scale", None) is not None:
+                # Layer-0 prologue: the residual-less input_layernorm folds the
+                # NVFP4 kv_a_proj input-quant via its attached .nvfp4_scale
+                # (RMSNorm.forward), returning an Fp4QuantizedTensor.
+                # return_norm_out stashes the BF16 view DSA's pre_indexer_proj
+                # needs on that tensor's bf16_hidden_states.
+                hidden_states = self.input_layernorm(hidden_states,
+                                                     return_norm_out=True)[0]
+            else:
+                hidden_states = self.input_layernorm(hidden_states)
         # Self Attention
         hidden_states = self.self_attn(
             position_ids=position_ids,
@@ -1536,12 +1572,28 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize):
                     self.layer_idx):
                 spec_metadata.maybe_capture_hidden_states(
                     self.layer_idx, hidden_states, residual)
-            if self.next_layer_layernorm is not None:
-                hidden_states, residual = self.next_layer_layernorm(
-                    hidden_states, residual)
+            hidden_states, residual = self._apply_next_layer_layernorm(
+                hidden_states, residual)
 
         return hidden_states, residual
 
+    def _apply_next_layer_layernorm(self, hidden_states, residual):
+        """Apply the next layer's input_layernorm at a layer boundary on the
+        attention-DP / no-allreduce path. When that norm has an .nvfp4_scale
+        attached, RMSNorm.forward folds the next layer's NVFP4 kv_a_proj
+        input-quant into the same kernel and returns (Fp4QuantizedTensor,
+        residual_out) with the BF16 view stashed for DSA's pre_indexer_proj
+        (return_norm_out); otherwise it applies the plain add+RMSNorm. Shared by
+        forward_MoE and forward_mlp so the boundary fold stays identical for MoE
+        and dense layers."""
+        if self.next_layer_layernorm is None:
+            return hidden_states, residual
+        if getattr(self.next_layer_layernorm, "nvfp4_scale", None) is not None:
+            return self.next_layer_layernorm(hidden_states,
+                                             residual,
+                                             return_norm_out=True)
+        return self.next_layer_layernorm(hidden_states, residual)
+
     def forward_mlp(
         self,
         hidden_states: torch.Tensor,
@@ -1573,6 +1625,14 @@ def forward_mlp(
                         eps=self.post_attention_layernorm.variance_epsilon,
                     ),
                 )
+        elif getattr(self.post_attention_layernorm, "nvfp4_scale",
+                     None) is not None:
+            # Attention-DP path (no allreduce): post_attention_layernorm folds
+            # the dense MLP's NVFP4 gate_up_proj input-quant into one kernel via
+            # its attached .nvfp4_scale (RMSNorm.forward), mirroring the
+            # PRE_MLP_FUSION quant but without the allreduce.
+            hidden_states, residual = self.post_attention_layernorm(
+                hidden_states, residual)
         else:
             # No fusion
             # We need to add twoshot allreduce here to avoid modifying MLA logic
@@ -1600,9 +1660,8 @@ def forward_mlp(
                     self.layer_idx):
                 spec_metadata.maybe_capture_hidden_states(
                     self.layer_idx, hidden_states, residual)
-            if self.next_layer_layernorm is not None:
-                hidden_states, residual = self.next_layer_layernorm(
-                    hidden_states, residual)
+            hidden_states, residual = self._apply_next_layer_layernorm(
+                hidden_states, residual)
 
         return hidden_states, residual
 
@@ -1947,5 +2006,33 @@ def setup_aliases(self) -> None:
             if idx == self.config.num_hidden_layers - 1:
                 layer.next_layer_layernorm = self.model.norm
             else:
-                layer.next_layer_layernorm = self.model.layers[
-                    idx + 1].input_layernorm
+                next_layer = self.model.layers[idx + 1]
+                layer.next_layer_layernorm = next_layer.input_layernorm
+
+            # Attach each norm's .nvfp4_scale (the consuming Linear's static
+            # input_scale) so RMSNorm.forward folds (add+)RMSNorm + NVFP4
+            # input-quant into one kernel on the attention-DP path. The fold is
+            # self-gating: _static_nvfp4_input_scale returns None (no fold)
+            # unless that Linear is NVFP4 with a static (calibrated) input_scale,
+            # so this is safe to run unconditionally on every layer.
+            #
+            # input_layernorm -> this layer's kv_a_proj covers BOTH the layer-0
+            # residual-less prologue AND the layer-boundary add+RMSNorm edge (the
+            # previous layer's next_layer_layernorm IS this input_layernorm).
+            # Both DSA (DSv3.2) and non-DSA (Kimi-K2.5) MLA are supported: the
+            # non-DSA forward_impl slices the Fp4QuantizedTensor by num_tokens
+            # via _slice_hidden_states_to_num_tokens, and kv_a_proj_with_mqa
+            # consumes the FP4 form directly.
+            self_attn = getattr(layer, "self_attn", None)
+            layer.input_layernorm.nvfp4_scale = _static_nvfp4_input_scale(
+                getattr(self_attn, "kv_a_proj_with_mqa", None))
+
+            # Intra-layer dense-MLP fold: post_attention_layernorm -> this
+            # layer's NVFP4 gate_up_proj. Only dense (GatedMLP) layers — MoE
+            # layers route their expert input quant differently and get no
+            # scale, so their fused path stays inert.
+            mlp = getattr(layer, "mlp", None)
+            if isinstance(mlp, GatedMLP):
+                layer.post_attention_layernorm.nvfp4_scale = (
+                    _static_nvfp4_input_scale(getattr(mlp, "gate_up_proj",
+                                                      None)))
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 80e9c67d74d3..57abfbd8291a 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -19,13 +19,57 @@
                            cp_allgather, reducescatter)
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
-from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
-                     is_nvfp4_marlin_enabled, is_torch_compiling)
+from ..utils import (Fp4QuantizedTensor, compute_swizzled_sf_shape,
+                     get_model_extra_attrs, is_nvfp4_marlin_enabled,
+                     is_torch_compiling)
 from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from .multi_stream_utils import maybe_execute_in_parallel
 from .rotary_embedding import MRotaryEmbedding, RotaryEmbedding
 
 
+def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int):
+    """Drop CUDA-graph padding by slicing hidden_states to num_tokens rows.
+
+    For a plain tensor this is the usual row slice. For an Fp4QuantizedTensor
+    (produced when the previous layer's boundary fusion pre-quantized this
+    layer's input) the slice must act on the packed FP4 form:
+      - fp4_tensor: row-major [m, k//2], so [:num_tokens] is a plain row slice;
+      - scaling_factor: 1D swizzled buffer of padUp(m,128)*padUp(cols,4). NVFP4
+        quant is per-row independent and the swizzle is laid out in independent
+        128-row tiles, so the leading padUp(num_tokens,128)*padUp(cols,4) bytes
+        are exactly the scale factors for the first num_tokens rows (verified by
+        tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py).
+      - bf16_hidden_states (when present): plain row slice.
+    """
+    if not isinstance(hidden_states, Fp4QuantizedTensor):
+        return hidden_states[:num_tokens, ...]
+
+    fp4 = hidden_states.fp4_tensor
+    if fp4.shape[0] == num_tokens:
+        return hidden_states  # no padding to strip
+    # The row-slice math below relies on the swizzled 128-row-tile layout; a
+    # non-swizzled (linear) scale buffer would be sliced incorrectly and then
+    # silently relabeled as swizzled. Reject it rather than corrupt the SF.
+    if not hidden_states.is_sf_swizzled:
+        raise ValueError(
+            "_slice_hidden_states_to_num_tokens only supports swizzled FP4 "
+            "scaling factors")
+    sf = hidden_states.scaling_factor
+    # fp4 packs 2 elements/byte, so the logical hidden dim is fp4_cols*2 and the
+    # scale-factor column count is that / SF_VEC_SIZE(16). The swizzled SF is
+    # laid out in independent 128-row tiles, so the leading num_tokens rows'
+    # scale factors occupy exactly the first padded_rows*padded_cols bytes.
+    sf_cols = fp4.shape[-1] * 2 // 16
+    padded_rows, padded_cols = compute_swizzled_sf_shape(num_tokens, sf_cols)
+    sf_len = padded_rows * padded_cols
+    sliced_bf16 = (hidden_states.bf16_hidden_states[:num_tokens]
+                   if hidden_states.bf16_hidden_states is not None else None)
+    return Fp4QuantizedTensor(fp4_tensor=fp4[:num_tokens].contiguous(),
+                              scaling_factor=sf.view(-1)[:sf_len].contiguous(),
+                              is_sf_swizzled=True,
+                              bf16_hidden_states=sliced_bf16)
+
+
 def extract_extra_attrs(layer_idx: str, attn_type: str):
     assert attn_type in ["mla", "attn"], "Invalid attention type"
     extra_attrs = get_model_extra_attrs()
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index e8da57a36c0e..ec1ea0057758 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3497,6 +3497,25 @@ def pre_reload_weights(self):
         self.quant_method.pre_reload_weights(self)
 
 
+def is_static_nvfp4_input_eligible(linear) -> bool:
+    """Whether ``linear`` consumes a static (calibrated) NVFP4 input, making it
+    eligible to have its input-quantize folded into a producing RMSNorm.
+
+    Eligible iff the Linear has NVFP4 weights, a calibrated (static)
+    ``input_scale``, no AWQ ``pre_quant_scale``, and is not forced to dynamic
+    quantization. This is the single canonical definition shared by every
+    NVFP4-fold site (the layer-boundary / dense folds in modeling_deepseekv3.py
+    and the q_a_layernorm -> q_b_proj fold in attention.py's MLA) so the gate
+    cannot drift between them.
+    """
+    if linear is None:
+        return False
+    return (getattr(linear, "has_nvfp4", False)
+            and not getattr(linear, "force_dynamic_quantization", False)
+            and getattr(linear, "input_scale", None) is not None
+            and getattr(linear, "pre_quant_scale", None) is None)
+
+
 class NVFP4ARCLinearMethod(NVFP4LinearMethod):
 
     supports_nccl_symmetric_memory_window_output: ClassVar[bool] = True
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index 4bee770842f7..6d19b3b106d4 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -47,15 +47,21 @@
 from ..attention_backend.utils import create_attention
 from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
-from ..utils import is_torch_compiling, maybe_compiled_cat, maybe_compiled_copy_
+from ..utils import (
+    Fp4QuantizedTensor,
+    is_torch_compiling,
+    maybe_compiled_cat,
+    maybe_compiled_copy_,
+)
 from .attention import (
     _helix_cp_allgather_input,
     _helix_cp_output_projection,
     _helix_post_process,
     _helix_zero_kv_mask,
+    _slice_hidden_states_to_num_tokens,
     extract_extra_attrs,
 )
-from .linear import Linear, TensorParallelMode
+from .linear import Linear, TensorParallelMode, is_static_nvfp4_input_eligible
 from .multi_stream_utils import do_multi_stream, maybe_execute_in_parallel
 from .rms_norm import RMSNorm
 from .rotary_embedding import RotaryEmbedding
@@ -113,9 +119,28 @@ def mla_custom_op_inplace(
     dsv4_output: Optional[torch.Tensor],
     dsv4_output_sf: Optional[torch.Tensor],
     enable_dsv4_epilogue_fusion: bool,
+    hidden_states_fp4: Optional[torch.Tensor] = None,
+    hidden_states_sf: Optional[torch.Tensor] = None,
 ) -> None:
+    # When hidden_states_fp4/_sf are provided, the previous layer's boundary
+    # fusion pre-quantized this layer's kv_a_proj NVFP4 input. Custom ops can't
+    # take dataclasses, so the Fp4QuantizedTensor is passed as explicit tensors
+    # and reconstructed here (mirrors mla_dsa_proj). bf16_hidden_states carries
+    # the un-quantized view for the o_proj output sizing in create_output.
     metadata, mla_layer = _extract_mla_extra_attrs(layer_idx)
+    if hidden_states_fp4 is not None or hidden_states_sf is not None:
+        assert hidden_states_fp4 is not None and hidden_states_sf is not None, (
+            "hidden_states_fp4 and hidden_states_sf must be passed together"
+        )
+        hidden_states = Fp4QuantizedTensor(
+            fp4_tensor=hidden_states_fp4,
+            scaling_factor=hidden_states_sf,
+            bf16_hidden_states=hidden_states,
+        )
     if mla_layer.is_deepseek_v4:
+        # DeepSeek-V4 uses MQA mode and has no residual-less RMSNorm+quant
+        # fusion entry point, so it cannot be reached with a pre-quantized
+        # Fp4QuantizedTensor input; the call site passes plain hidden_states.
         if enable_dsv4_epilogue_fusion:
             if dsv4_output is None or dsv4_output_sf is None:
                 raise RuntimeError(
@@ -148,6 +173,8 @@ def mla_dsa_proj(
     hidden_states: torch.Tensor,
     position_ids: Optional[torch.Tensor],
     layer_idx: str,
+    hidden_states_fp4: Optional[torch.Tensor] = None,
+    hidden_states_sf: Optional[torch.Tensor] = None,
 ) -> List[torch.Tensor]:
     """Token-wise projections for DSA MLA (CUDA-graph-capturable).
 
@@ -156,6 +183,14 @@ def mla_dsa_proj(
     update the indexer k cache — that happens in Op 2 (mla_dsa_attn_inplace)
     because the scatter kernel accesses batch-specific metadata.
 
+    When ``hidden_states_fp4`` / ``hidden_states_sf`` are provided, the
+    previous layer's boundary fusion already produced a fused NVFP4 view of
+    the post-RMSNorm hidden states (via fused_add_rmsnorm_fp4_quantize). The
+    op rebuilds an ``Fp4QuantizedTensor`` carrying the BF16 view in
+    ``hidden_states`` (for DSA's ``pre_indexer_proj``) plus the
+    pre-quantized FP4 form (consumed by ``kv_a_proj_with_mqa``).
+    Both must be passed together; passing either alone is rejected.
+
     Returns [q, compressed_kv, k_pe, latent_cache] when the short-MHA path
     handles all tokens, or [q, compressed_kv, k_pe, latent_cache, q_fp8,
     k_fp8, k_scale, weights, q_scale] when the indexer runs.  Under torch
@@ -165,7 +200,18 @@ def mla_dsa_proj(
     path ignores it in forward_dsa_attn.
     """
     metadata, mla_layer = _extract_mla_extra_attrs(layer_idx)
-    return mla_layer.forward_dsa_proj(position_ids, hidden_states, metadata)
+    if hidden_states_fp4 is not None or hidden_states_sf is not None:
+        assert hidden_states_fp4 is not None and hidden_states_sf is not None, (
+            "hidden_states_fp4 and hidden_states_sf must be passed together"
+        )
+        hs = Fp4QuantizedTensor(
+            fp4_tensor=hidden_states_fp4,
+            scaling_factor=hidden_states_sf,
+            bf16_hidden_states=hidden_states,
+        )
+    else:
+        hs = hidden_states
+    return mla_layer.forward_dsa_proj(position_ids, hs, metadata)
 
 
 @mla_dsa_proj.register_fake
@@ -173,6 +219,8 @@ def _mla_dsa_proj_fake(
     hidden_states: torch.Tensor,
     position_ids: Optional[torch.Tensor],
     layer_idx: str,
+    hidden_states_fp4: Optional[torch.Tensor] = None,
+    hidden_states_sf: Optional[torch.Tensor] = None,
 ) -> List[torch.Tensor]:
     # Under torch compile _should_use_short_mha is False, so the result is
     # always 9 tensors (4 attention inputs + 5 indexer intermediates, with
@@ -418,6 +466,14 @@ def __init__(
             "TRTLLM_DSV4_DISABLE_FMHA_EPILOGUE_FUSION"
         )
 
+        # Fold of the residual-less q_a_layernorm -> q_b_proj NVFP4 input
+        # quantize into a single fused (rmsnorm + fp4_quantize) kernel. Self-
+        # gating on q_b_proj being static-NVFP4 (see _resolve_qa_fused_scale);
+        # resolved lazily on first forward because q_b_proj.input_scale is only
+        # finalized after weight loading.
+        # None = not yet resolved; False = disabled; tensor = q_b_proj input_scale.
+        self._qa_fused_scale = None
+
         # tensor parallel
         if mapping_with_cp is not None:
             logger.warning_once(
@@ -1024,6 +1080,17 @@ def _attn_forward_gen(
             return attn_output
 
     def create_output(self, hidden_states: torch.Tensor, num_contexts: int):
+        # Upstream POST_MoE/MLP fusion (or attention-DP no-fusion fold) may pass
+        # an Fp4QuantizedTensor here; unpack to the BF16 view for sizing. The
+        # producing fold must have requested return_norm_out so the BF16 view
+        # is present (folds feeding an MLA always do).
+        if isinstance(hidden_states, Fp4QuantizedTensor):
+            assert hidden_states.bf16_hidden_states is not None, (
+                "MLA.create_output received an Fp4QuantizedTensor without a "
+                "bf16_hidden_states view; the producing fusion must use "
+                "return_norm_out=True"
+            )
+            hidden_states = hidden_states.bf16_hidden_states
         num_tokens = hidden_states.shape[0]
         if self.is_deepseek_v4:
             hidden_size = self.num_heads_tp_cp * self.v_head_dim
@@ -1210,6 +1277,69 @@ def _deepseek_v4_o_proj(
         output = self.o_b_proj(o_lora)
         return output
 
+    def _resolve_qa_fused_scale(self):
+        """Lazily decide whether the residual-less q_a_layernorm -> q_b_proj
+        NVFP4 fusion can run, caching q_b_proj's static input_scale.
+
+        Returns the input_scale tensor when eligible, else None. Eligible iff:
+          (1) model is non-lite (has a distinct q_a_layernorm / q_b_proj),
+          (2) q_b_proj is static-NVFP4 (calibrated input_scale, no
+              pre_quant_scale / AWQ, not forced-dynamic),
+          (3) q_a_layernorm is a plain (non-gemma, non-quantizing) RMSNorm.
+        """
+        if self._qa_fused_scale is not None:
+            # Already resolved: tensor (enabled) or False (disabled).
+            return self._qa_fused_scale if self._qa_fused_scale is not False else None
+        eligible = False
+        if (
+            not self.is_lite
+            and getattr(self, "q_a_layernorm", None) is not None
+            and getattr(self, "q_b_proj", None) is not None
+        ):
+            qb = self.q_b_proj
+            qa = self.q_a_layernorm
+            # q_b_proj must be static-NVFP4 (shared predicate) and q_a_layernorm
+            # a plain (non-gemma, non-quantizing) RMSNorm.
+            if (
+                is_static_nvfp4_input_eligible(qb)
+                and not getattr(qa, "use_gemma", False)
+                and not getattr(qa, "is_nvfp4", False)
+            ):
+                eligible = True
+        if eligible:
+            # Mark q_a_layernorm as an NVFP4 norm and attach q_b_proj's static
+            # input_scale, so RMSNorm.forward folds the residual-less
+            # q_a_layernorm + q_b_proj NVFP4 input-quant into one kernel (the
+            # size gate routes this N=q_lora_rank<2048 edge to reduce_fusion).
+            self.q_a_layernorm.is_nvfp4 = True
+            self.q_a_layernorm.nvfp4_scale = qb.input_scale
+        self._qa_fused_scale = qb.input_scale if eligible else False
+        return self._qa_fused_scale if eligible else None
+
+    def _q_a_layernorm_maybe_fused(self, q: torch.Tensor, return_norm_out: bool = False):
+        """Apply q_a_layernorm, optionally folding the q_b_proj NVFP4 input
+        quantize into the same kernel.
+
+        When fusion is disabled, returns the BF16/FP16 normed tensor (and, if
+        return_norm_out, the same tensor again so callers have a stable arity).
+
+        When fusion is enabled, returns an Fp4QuantizedTensor (consumed by
+        q_b_proj without re-quantizing). If return_norm_out is set, also returns
+        the BF16 post-RMSNorm value, needed by the DSA indexer's
+        pre_indexer_proj which requires a float q."""
+        scale = self._resolve_qa_fused_scale()
+        if scale is None:
+            normed = self.q_a_layernorm(q)
+            return (normed, normed) if return_norm_out else normed
+        # q is typically the leading q_lora_rank columns of the
+        # kv_a_proj_with_mqa output split, i.e. a column slice whose last dim is
+        # unit-stride but whose row pitch is the full projection width. The
+        # fused RMSNorm+NVFP4-quant kernel reads that strided layout directly,
+        # so no .contiguous() copy is needed. _resolve_qa_fused_scale has
+        # attached .nvfp4_scale to q_a_layernorm, so RMSNorm.forward folds the
+        # residual-less RMSNorm + NVFP4 quant automatically.
+        return self.q_a_layernorm(q, return_norm_out=return_norm_out)
+
     def forward_impl(
         self,
         position_ids: Optional[torch.Tensor],
@@ -1234,7 +1364,7 @@ def forward_impl(
         num_ctx_tokens = attn_metadata.num_ctx_tokens
         num_tokens = attn_metadata.num_tokens
 
-        hidden_states = hidden_states[:num_tokens, ...]
+        hidden_states = _slice_hidden_states_to_num_tokens(hidden_states, num_tokens)
         if position_ids is not None:
             position_ids = position_ids[..., :num_tokens]
 
@@ -1250,7 +1380,7 @@ def forward_impl(
             )
 
             q, compressed_kv = maybe_execute_in_parallel(
-                lambda: self.q_a_layernorm(q),
+                lambda: self._q_a_layernorm_maybe_fused(q),
                 lambda: self.kv_a_layernorm(compressed_kv),
                 self.ln_events[0],
                 self.ln_events[1],
@@ -1382,14 +1512,17 @@ def forward_dsa_proj(
             [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1
         )
 
-        q, compressed_kv = maybe_execute_in_parallel(
-            lambda: self.q_a_layernorm(q),
+        # When the q_a_layernorm -> q_b_proj NVFP4 fusion is active, q becomes an
+        # Fp4QuantizedTensor consumed by q_b_proj; qr (fed to the indexer) needs
+        # the BF16 post-RMSNorm value, so request return_norm_out here.
+        q_pair, compressed_kv = maybe_execute_in_parallel(
+            lambda: self._q_a_layernorm_maybe_fused(q, return_norm_out=True),
             lambda: self.kv_a_layernorm(compressed_kv),
             self.ln_events[0],
             self.ln_events[1],
             self.aux_stream,
         )
-        qr = q
+        q, qr = q_pair
         latent_cache = torch.concat([compressed_kv, k_pe], dim=-1)
 
         q = self.q_b_proj(q)
@@ -2906,9 +3039,22 @@ def forward(
             else:
                 attn_output = self.create_output(hidden_states, attn_metadata.num_contexts)
                 if self.is_dsa:
-                    proj_outputs = torch.ops.trtllm.mla_dsa_proj(
-                        hidden_states, position_ids, self.layer_idx_str
-                    )
+                    # When the previous layer's boundary fusion pre-quantized
+                    # this layer's kv_a_proj NVFP4 input, hidden_states is an
+                    # Fp4QuantizedTensor with both BF16 + FP4 views. Custom ops
+                    # can't take dataclasses, so pass tensors explicitly.
+                    if isinstance(hidden_states, Fp4QuantizedTensor):
+                        proj_outputs = torch.ops.trtllm.mla_dsa_proj(
+                            hidden_states.bf16_hidden_states,
+                            position_ids,
+                            self.layer_idx_str,
+                            hidden_states.fp4_tensor,
+                            hidden_states.scaling_factor,
+                        )
+                    else:
+                        proj_outputs = torch.ops.trtllm.mla_dsa_proj(
+                            hidden_states, position_ids, self.layer_idx_str
+                        )
                     q, compressed_kv, k_pe, latent_cache = proj_outputs[:4]
                     indexer_intermediates = proj_outputs[4:]
                     torch.ops.trtllm.mla_dsa_attn_inplace(
@@ -2922,16 +3068,35 @@ def forward(
                         attn_output,
                     )
                 else:
-                    torch.ops.trtllm.mla_custom_op_inplace(
-                        hidden_states,
-                        position_ids,
-                        self.layer_idx_str,
-                        attn_output,
-                        latent_cache_gen,
-                        None,
-                        None,
-                        False,
-                    )
+                    # Vanilla MLA (e.g. Kimi-K2.5): when the previous layer's
+                    # boundary fusion pre-quantized this layer's input,
+                    # hidden_states is an Fp4QuantizedTensor. Custom ops can't
+                    # take dataclasses, so pass the BF16 + FP4 + SF views as
+                    # explicit tensors.
+                    if isinstance(hidden_states, Fp4QuantizedTensor):
+                        torch.ops.trtllm.mla_custom_op_inplace(
+                            hidden_states.bf16_hidden_states,
+                            position_ids,
+                            self.layer_idx_str,
+                            attn_output,
+                            latent_cache_gen,
+                            None,
+                            None,
+                            False,
+                            hidden_states.fp4_tensor,
+                            hidden_states.scaling_factor,
+                        )
+                    else:
+                        torch.ops.trtllm.mla_custom_op_inplace(
+                            hidden_states,
+                            position_ids,
+                            self.layer_idx_str,
+                            attn_output,
+                            latent_cache_gen,
+                            None,
+                            None,
+                            False,
+                        )
         else:
             enable_dsv4_epilogue_fusion = (
                 self.is_deepseek_v4
diff --git a/tensorrt_llm/_torch/modules/rms_norm.py b/tensorrt_llm/_torch/modules/rms_norm.py
index 8575f006be7b..fb58478e3f6e 100644
--- a/tensorrt_llm/_torch/modules/rms_norm.py
+++ b/tensorrt_llm/_torch/modules/rms_norm.py
@@ -25,6 +25,17 @@
 from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
 from ..utils import Fp4QuantizedTensor, is_nvfp4_marlin_enabled
 
+# Hidden-dim bounds of the warp-specialized fused_add_rms_norm_quant kernel
+# (cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp).
+_WS_MIN_N = 2048
+_WS_MAX_N = 16384
+# Row-count crossover for the layer-boundary add+RMSNorm+quant edge: below this
+# the one-CTA-per-row reduce_fusion kernel (fused_add_rmsnorm_fp4_quantize) is
+# faster; at/above it the warp-specialized kernel's DMA-ahead pipeline wins.
+# Measured at N=7168 on GB200 (ws first overtakes at M=4096, ~6% faster; M<=3072
+# favors reduce_fusion). See analysis/bench_rmsnorm_threshold_n7168.py.
+_WS_M_THRESHOLD = 4096
+
 
 class RMSNorm(nn.Module):
 
@@ -75,14 +86,19 @@ def __init__(
         self.use_gemma = use_gemma
         self.use_cuda_tile = use_cuda_tile
 
-        # fused_add_rms_norm_quant only supports SM 9.x / 10.x because:
-        #  - Device code is guarded by is_major_v<9> || is_major_v<10>
-        #    (ws_layernorm.cuh:828, low_latency_layernorm.cuh:157).
-        # On unsupported SMs, fall back to flashinfer/generic RMSNorm and let
+        # The fused NVFP4 epilogue (cvt_warp_fp16_to_fp4 in quantization.cuh) is
+        # guarded by __CUDA_ARCH__ >= 1000, so it only produces correct FP4 on
+        # SM 10.x (Blackwell); SM 12.x lacks the warp-cooperative SF path. On
+        # unsupported SMs (incl. SM 9.x / Hopper, where the epilogue would
+        # silently emit zeros) fall back to flashinfer/generic RMSNorm and let
         # the downstream linear layer handle FP4 quantization.
         if self.is_nvfp4:
             sm_version = get_sm_version()
-            if not (90 <= sm_version < 120) or is_nvfp4_marlin_enabled():
+            # SM range narrowed to 10.x by this PR: the reduce_fusion kernels
+            # use cvt_warp_fp16_to_fp4 (guarded by __CUDA_ARCH__ >= 1000),
+            # which silently emits zeros on SM 9.x. Marlin is also
+            # incompatible with the fused NVFP4 path.
+            if not (100 <= sm_version < 120) or is_nvfp4_marlin_enabled():
                 self.is_nvfp4 = False
                 return_hp_output = False
         self.return_hp_output = return_hp_output
@@ -93,6 +109,8 @@ def forward(
         residual: Union[
             Optional[torch.Tensor],
             _ArgumentNotSpecifiedSentinelType] = _ARGUMENT_NOT_SPECIFIED_SENTINEL,
+        *,
+        return_norm_out: bool = False,
     ) -> Union[torch.Tensor, Fp4QuantizedTensor, Tuple[Union[
             torch.Tensor, Fp4QuantizedTensor], Optional[torch.Tensor]], Tuple[
                 Fp4QuantizedTensor, torch.Tensor, torch.Tensor]]:
@@ -100,63 +118,23 @@ def forward(
         if not has_residual:
             residual = None
 
-        if self.is_nvfp4 and has_residual and not self.use_gemma:
-            nvfp4_scale = getattr(self, "nvfp4_scale", None)
-            if nvfp4_scale is None:
-                raise ValueError(
-                    f"layeridx={getattr(self, 'layer_idx', None)} RMSNorm NVFP4 output requested "
-                    "but no `nvfp4_scale` is attached; ")
-
-            orig_shape = tuple(hidden_states.shape)
-            n = int(orig_shape[-1])
-            hs_2d = hidden_states.reshape(-1, n).contiguous()
-            res_2d = residual.reshape(-1, n)
-            gamma = self.weight
-
-            def _ensure_contiguous_with_dtype(t: torch.Tensor, key: str):
-                if t.dtype != hs_2d.dtype:
-                    raise ValueError(
-                        f"RMSNorm NVFP4 fused path: casting {key} from {t.dtype} to {hs_2d.dtype}."
-                    )
-                return t.contiguous()
-
-            res_2d = _ensure_contiguous_with_dtype(res_2d, "residual")
-            gamma = _ensure_contiguous_with_dtype(gamma, "gamma")
-
-            if hs_2d.device != res_2d.device or hs_2d.device != gamma.device:
-                raise RuntimeError(
-                    "RMSNorm NVFP4 fused path requires all tensors on the same device. "
-                    f"Got input={hs_2d.device}, residual={res_2d.device}, gamma={gamma.device}."
-                )
-
-            sf_scale = nvfp4_scale.contiguous()
-
-            results = torch.ops.trtllm.fused_add_rms_norm_quant(
-                hs_2d,
-                res_2d,
-                gamma,
-                sf_scale,
-                True,
-                eps=self.variance_epsilon,
-                output_hp_norm=self.return_hp_output,
-            )
-            normed_fp4_i32, residual_out_2d, sf_fused = results[:3]
-            normed_fp4_u8 = normed_fp4_i32.view(torch.uint8)
-            if len(orig_shape) != 2:
-                normed_fp4_u8 = normed_fp4_u8.reshape(*orig_shape[:-1], n // 2)
-                residual_out = residual_out_2d.reshape(orig_shape)
-            else:
-                residual_out = residual_out_2d
-
-            hidden_states_fused = Fp4QuantizedTensor(normed_fp4_u8, sf_fused)
-
-            outputs = [hidden_states_fused]
-            if has_residual:
-                outputs.append(residual_out)
-            if self.return_hp_output:
-                high_precision_normed_output = results[3].reshape(orig_shape)
-                outputs.append(high_precision_normed_output)
-            return outputs[0] if len(outputs) == 1 else tuple(outputs)
+        # Fused (add +) RMSNorm + NVFP4 input-quantize: when this norm feeds an
+        # NVFP4 Linear whose static input_scale is attached as self.nvfp4_scale,
+        # fold this norm and that Linear's input-quant into one kernel. With a
+        # residual this is the layer-boundary add+RMSNorm+quant; without, the
+        # residual-less variant (e.g. q_a_layernorm -> q_b_proj). The actual
+        # kernel is chosen by a size gate inside _fused_nvfp4_quant. The BF16
+        # post-RMSNorm value is also produced when return_norm_out (stashed on
+        # the Fp4QuantizedTensor for DSA consumers / returned for residual-less
+        # callers) or when self.return_hp_output (appended as an extra output
+        # for MoE-gate consumers). When is_nvfp4 but no scale is attached yet
+        # (e.g. a layer's first input_layernorm whose consumer has no static
+        # scale), fall through to the plain norm below.
+        nvfp4_scale = getattr(self, "nvfp4_scale",
+                              None) if self.is_nvfp4 else None
+        if nvfp4_scale is not None and not self.use_gemma:
+            return self._fused_nvfp4_quant(hidden_states, residual, nvfp4_scale,
+                                           return_norm_out)
 
         if self.return_hp_output:
             raise ValueError(
@@ -230,6 +208,145 @@ def _ensure_contiguous_with_dtype(t: torch.Tensor, key: str):
         else:
             return hidden_states
 
+    def _ws_kernel_eligible(self, hidden_states: torch.Tensor,
+                            residual: Optional[torch.Tensor]) -> bool:
+        """Whether the warp-specialized fused_add_rms_norm_quant kernel should
+        serve this call instead of the one-CTA-per-row reduce_fusion kernel.
+
+        ws only wins, and is only legal, for the contiguous rank-2 residual-add
+        edge with a large enough row count and an in-range hidden dim. The
+        residual-less and row-strided (column-slice) edges have no ws equivalent
+        and always use reduce_fusion. See _WS_M_THRESHOLD / _WS_MIN_N."""
+        if residual is None:
+            return False
+        n = hidden_states.shape[-1]
+        if not (_WS_MIN_N <= n <= _WS_MAX_N and n % 16 == 0):
+            return False
+        m = 1
+        for d in hidden_states.shape[:-1]:
+            m *= d
+        if m < _WS_M_THRESHOLD:
+            return False
+        # ws requires contiguous rank-2 input + residual (it cannot read a
+        # row-strided column slice and issues padded vectorized stores).
+        return hidden_states.is_contiguous() and residual.is_contiguous()
+
+    def _fused_nvfp4_quant(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        sf_scale: torch.Tensor,
+        return_norm_out: bool,
+    ):
+        """Fold this (add +) RMSNorm + the consuming NVFP4 Linear's input-quant
+        into one kernel. ``sf_scale`` is that Linear's static input_scale.
+
+        Picks the kernel by problem size (see _ws_kernel_eligible): the
+        warp-specialized ``fused_add_rms_norm_quant`` for the large contiguous
+        residual edge, else the one-CTA-per-row ``fused_add_rmsnorm_fp4_quantize``
+        / ``fused_rmsnorm_fp4_quantize`` (which also serve the residual-less and
+        row-strided column-slice edges ws cannot).
+
+        With a residual: returns ``(Fp4QuantizedTensor, residual_out)``. Without:
+        returns an ``Fp4QuantizedTensor`` (the residual-less q_a edge). When
+        ``return_norm_out`` the BF16 post-RMSNorm view is stashed on the
+        Fp4QuantizedTensor's ``bf16_hidden_states`` (and, residual-less, returned
+        alongside). When ``self.return_hp_output`` the BF16 normed value is
+        appended as a trailing output for MoE-gate consumers."""
+        # Either flag means "also produce the BF16 normed value".
+        want_norm = return_norm_out or self.return_hp_output
+
+        if self._ws_kernel_eligible(hidden_states, residual):
+            return self._fused_nvfp4_quant_ws(hidden_states, residual, sf_scale,
+                                              return_norm_out)
+
+        if residual is not None:
+            results = torch.ops.trtllm.fused_add_rmsnorm_fp4_quantize(
+                hidden_states,
+                residual,
+                self.weight,
+                sf_scale,
+                float(self.variance_epsilon),
+                want_norm,
+            )
+            if want_norm:
+                bf16_hs, act_fp4, act_sf, residual_out = results
+            else:
+                act_fp4, act_sf, residual_out = results
+                bf16_hs = None
+            fp4 = Fp4QuantizedTensor(act_fp4,
+                                     act_sf,
+                                     bf16_hidden_states=bf16_hs)
+            outputs = [fp4, residual_out]
+            if self.return_hp_output:
+                outputs.append(bf16_hs)
+            return tuple(outputs)
+
+        orig_shape = tuple(hidden_states.shape)
+        n = orig_shape[-1]
+        hs_2d = hidden_states.reshape(-1, n)
+        results = torch.ops.trtllm.fused_rmsnorm_fp4_quantize(
+            hs_2d,
+            self.weight,
+            sf_scale,
+            float(self.variance_epsilon),
+            want_norm,
+        )
+        if want_norm:
+            norm_out, act_fp4, act_sf = results
+            norm_out = norm_out.reshape(orig_shape)
+        else:
+            act_fp4, act_sf = results
+            norm_out = None
+        if len(orig_shape) != 2:
+            act_fp4 = act_fp4.reshape(*orig_shape[:-1], n // 2)
+        fp4 = Fp4QuantizedTensor(act_fp4, act_sf, bf16_hidden_states=norm_out)
+        return (fp4, norm_out) if return_norm_out else fp4
+
+    def _fused_nvfp4_quant_ws(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        sf_scale: torch.Tensor,
+        return_norm_out: bool,
+    ):
+        """Warp-specialized fused add+RMSNorm+NVFP4-quant
+        (torch.ops.trtllm.fused_add_rms_norm_quant) for the large contiguous
+        residual edge. Returns the same shape as the reduce_fusion residual
+        branch of _fused_nvfp4_quant."""
+        want_norm = return_norm_out or self.return_hp_output
+        orig_shape = tuple(hidden_states.shape)
+        n = int(orig_shape[-1])
+        hs_2d = hidden_states.reshape(-1, n).contiguous()
+        res_2d = residual.reshape(-1, n).contiguous()
+        gamma = self.weight.contiguous()
+
+        results = torch.ops.trtllm.fused_add_rms_norm_quant(
+            hs_2d,
+            res_2d,
+            gamma,
+            sf_scale.contiguous(),
+            True,
+            eps=self.variance_epsilon,
+            output_hp_norm=want_norm,
+        )
+        normed_fp4_i32, residual_out_2d, sf_fused = results[:3]
+        normed_fp4_u8 = normed_fp4_i32.view(torch.uint8)
+        if len(orig_shape) != 2:
+            normed_fp4_u8 = normed_fp4_u8.reshape(*orig_shape[:-1], n // 2)
+            residual_out = residual_out_2d.reshape(orig_shape)
+        else:
+            residual_out = residual_out_2d
+
+        bf16_hs = results[3].reshape(orig_shape) if want_norm else None
+        fp4 = Fp4QuantizedTensor(normed_fp4_u8,
+                                 sf_fused,
+                                 bf16_hidden_states=bf16_hs)
+        outputs = [fp4, residual_out]
+        if self.return_hp_output:
+            outputs.append(bf16_hs)
+        return tuple(outputs)
+
     def skip_forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py
index beb6c1c9f075..a1369917ad28 100644
--- a/tensorrt_llm/_torch/utils.py
+++ b/tensorrt_llm/_torch/utils.py
@@ -4,7 +4,7 @@
 import threading
 from dataclasses import dataclass
 from enum import Enum, IntEnum
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch.nn import functional as F
@@ -160,6 +160,13 @@ class Fp4QuantizedTensor:
     fp4_tensor: torch.Tensor
     scaling_factor: torch.Tensor
     is_sf_swizzled: bool = True
+    # Optional BF16/FP16 hidden-state view of the same logical activation.
+    # When the FP4 tensor is produced by a fused (add+)RMSNorm+NVFP4-quant that
+    # also returns the un-quantized (post-RMSNorm) value, this carries that
+    # BF16 tensor so downstream consumers needing the un-quantized form (e.g.
+    # DSv3.2's DSA indexer at sparse/dsa.py:pre_indexer_proj) can use it
+    # without dequantizing FP4.
+    bf16_hidden_states: Optional[torch.Tensor] = None
 
     @property
     def shape(self):
diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index fb7136397055..6a7c825d27e1 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -128,6 +128,8 @@ l0_b200:
   - unittest/_torch/modules/fused_ops/test_gelu_tanh_mul_fp4_quant.py
   - unittest/_torch/modules/fused_ops/test_rmsnorm_residual_add.py
   - unittest/_torch/modules/test_gemma4_fused_qkv_prep.py
+  - unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
+  - unittest/_torch/modules/test_fp4_num_tokens_slice.py
   - unittest/_torch/modules/test_awq_quantization.py
   - unittest/_torch/modules/test_triton_linear.py
   - unittest/_torch/modules/test_group_rmn_norm.py
diff --git a/tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py b/tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py
new file mode 100644
index 000000000000..433d7c684854
--- /dev/null
+++ b/tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py
@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Validates how to slice a padded NVFP4 activation by num_tokens so it can
+feed the NVFP4 GEMM -- the enabling step for folding the layer-boundary
+add+RMSNorm+NVFP4-quant on the **non-DSA** MLA path (e.g. Kimi-K2.5).
+
+On that path, MLA.forward_impl slices the next layer's input by num_tokens
+(``hidden_states = hidden_states[:num_tokens]``) to drop CUDA-graph padding,
+BEFORE feeding kv_a_proj_with_mqa. To let the boundary fusion pre-quantize that
+input (producing an Fp4QuantizedTensor sized for the padded batch), the slice
+must happen on the FP4 form.
+
+Key fact this test pins down: NVFP4 quantization is per-row independent, and
+the swizzled scale-factor (computeSwizzledLayoutSFSize = padUp(rows,128) x
+padUp(cols,4)) is laid out tile-by-tile in 128-row tiles. Therefore the slice
+is simply:
+
+  * fp4_tensor[:num_tokens]                                 (row-major, trivial)
+  * scaling_factor.view(-1)[:padUp(num_tokens,128)*padUp(cols,4)]
+        -- the LEADING swizzled bytes, which are byte-identical to freshly
+           swizzling only num_tokens rows because each 128-row tile's layout
+           is independent of the total row count.
+
+This test feeds that sliced FP4 input through the NVFP4 GEMM and asserts it
+matches the unpadded reference (quantize the real tokens, then GEMM). If it
+holds, forward_impl can slice the Fp4QuantizedTensor directly and the is_dsa
+gate on the boundary fold can be relaxed to cover non-DSA models.
+"""
+
+import pytest
+import torch
+
+import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils
+from tests.unittest.utils.util import skip_pre_blackwell
+
+SF_VEC = 16
+
+
+def _gemm(x_fp4, w_fp4, x_sf, w_sf, alpha, dtype):
+    return torch.ops.trtllm.nvfp4_gemm_cutlass(x_fp4, w_fp4, x_sf, w_sf, alpha, dtype)
+
+
+def _leading_sf_len(num_tokens, k):
+    return fp4_utils.pad_up(num_tokens, 128) * fp4_utils.pad_up(k // SF_VEC, 4)
+
+
+@skip_pre_blackwell
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+# (padded_m, num_tokens, n, k): padded_m = CUDA-graph batch, num_tokens = real
+# count. Includes num_tokens that straddle a 128-row swizzle-tile boundary
+# (130 with padded_m 256) so we exercise multi-tile slicing.
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (128, 40, 512, 1536),
+        (128, 100, 2048, 7168),
+        (256, 130, 512, 1536),
+        (32, 7, 512, 1536),
+    ],
+)
+def test_fp4_num_tokens_slice_feeds_gemm(dtype, shape):
+    padded_m, num_tokens, n, k = shape
+    assert num_tokens <= padded_m and k % SF_VEC == 0
+    torch.manual_seed(0)
+    dev = torch.device("cuda")
+
+    # Weight (static NVFP4), shared across paths.
+    w = torch.randn((n, k), dtype=dtype, device=dev)
+    w_sf_global = (448 * 6) / w.abs().max().float()
+    w_fp4, w_sf = torch.ops.trtllm.fp4_quantize(w, w_sf_global, SF_VEC, False)
+
+    # Padded activation: rows [num_tokens:padded_m] are graph padding (garbage).
+    # The per-tensor input scale is calibrated on the real tokens only, matching
+    # the static-NVFP4 Linear (input_scale comes from calibration, not runtime).
+    x_pad = torch.randn((padded_m, k), dtype=dtype, device=dev)
+    x_sf_global = (448 * 6) / x_pad[:num_tokens].abs().max().float()
+    x_pad_fp4, x_pad_sf = torch.ops.trtllm.fp4_quantize(x_pad, x_sf_global, SF_VEC, False)
+    alpha = (1.0 / (w_sf_global * x_sf_global)).to(torch.float32).view(1)
+
+    # ---- Reference: quantize ONLY the real tokens, then GEMM (what the unfused
+    # path produces today: slice BF16 first, then quantize). ----
+    x_real = x_pad[:num_tokens].contiguous()
+    xr_fp4, xr_sf = torch.ops.trtllm.fp4_quantize(x_real, x_sf_global, SF_VEC, False)
+    out_ref = _gemm(xr_fp4, w_fp4, xr_sf, w_sf, alpha, dtype)
+    assert out_ref.shape[0] == num_tokens
+
+    # ---- Sliced: fp4_tensor[:num_tokens] + leading swizzled SF bytes. This is
+    # the slice forward_impl would apply to a pre-quantized Fp4QuantizedTensor. ----
+    sliced_fp4 = x_pad_fp4[:num_tokens].contiguous()
+    sliced_sf = x_pad_sf.view(-1)[: _leading_sf_len(num_tokens, k)].contiguous()
+    out_sliced = _gemm(sliced_fp4, w_fp4, sliced_sf, w_sf, alpha, dtype)
+
+    torch.testing.assert_close(out_sliced.float(), out_ref.float(), rtol=2e-2, atol=0.2)
diff --git a/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py b/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
new file mode 100644
index 000000000000..35f421f69fc9
--- /dev/null
+++ b/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
@@ -0,0 +1,407 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for the fused (add +) RMSNorm + NVFP4-quantize kernels.
+
+Both ops fold the standalone NVFP4 input-quantize (FP16/BF16 -> NVFP4) into its
+producing RMSNorm at DeepSeek-V3.2 / Kimi-K2.5 MLA sites:
+
+  * ``fused_add_rmsnorm_fp4_quantize`` - layer-boundary
+    ``residual_add -> rms_norm -> nvfp4_quantize`` feeding the next layer's
+    ``kv_a_proj``.
+  * ``fused_rmsnorm_fp4_quantize`` - the residual-less intra-MLA
+    ``q_a_layernorm -> q_b_proj`` input quant, which additionally supports a
+    row-strided (column-slice) input read so the ``q`` slice from
+    ``kv_a_proj_with_mqa(...).split(...)`` needs no ``.contiguous()`` copy.
+
+Accuracy is checked against the unfused reference: a PyTorch RMSNorm followed by
+``torch.ops.trtllm.fp4_quantize`` (the exact kernel the fusion replaces). The
+fused kernel computes the RMSNorm in registers while already reading the tensor
+to quantize it, so a handful of values right at quantization boundaries can fall
+on the other side of a step due to FMA-vs-separate rounding; we require a >= 99%
+match on the packed FP4 bytes, mirroring ``test_fused_activation_quant.py``.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.utils import ceil_div, pad_up, unswizzle_sf
+from tests.unittest.utils.util import getSMVersion
+
+
+def fused_add_rmsnorm_fp4_quantize_available():
+    return hasattr(torch.ops, "trtllm") and hasattr(
+        torch.ops.trtllm, "fused_add_rmsnorm_fp4_quantize"
+    )
+
+
+def fused_rmsnorm_fp4_quantize_available():
+    return hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "fused_rmsnorm_fp4_quantize")
+
+
+def fp4_quantize_available():
+    return hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "fp4_quantize")
+
+
+skip_unless_add_rmsnorm = pytest.mark.skipif(
+    getSMVersion() < 100
+    or not fused_add_rmsnorm_fp4_quantize_available()
+    or not fp4_quantize_available(),
+    reason="Requires SM100+ (Blackwell) and trtllm "
+    "fused_add_rmsnorm_fp4_quantize + fp4_quantize ops",
+)
+
+skip_unless_rmsnorm = pytest.mark.skipif(
+    getSMVersion() < 100
+    or not fused_rmsnorm_fp4_quantize_available()
+    or not fp4_quantize_available(),
+    reason="Requires SM100+ (Blackwell) and trtllm fused_rmsnorm_fp4_quantize + fp4_quantize ops",
+)
+
+
+def rms_norm_ref(hidden_states: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """Reference RMSNorm matching tensorrt_llm/_torch/modules/rms_norm.py:
+    fp32 accumulate, rsqrt(var + eps), then weight * normed cast back to the
+    input dtype."""
+    input_dtype = hidden_states.dtype
+    hidden_states = hidden_states.to(torch.float32)
+    variance = hidden_states.pow(2).mean(-1, keepdim=True)
+    hidden_states = hidden_states * torch.rsqrt(variance + eps)
+    return weight * hidden_states.to(input_dtype)
+
+
+# NVFP4 packs SF_VEC_SIZE elements per E4M3 scale-factor block.
+SF_VEC = 16
+# NVFP4 numeric maxima: FP8 (E4M3) and E2M1.
+FP8_MAX, E2M1_MAX = 448.0, 6.0
+
+
+def fp4_quantize_ref(normed: torch.Tensor, sf_scale: torch.Tensor):
+    """Unfused baseline: the standalone NVFP4 quantize the fusion replaces."""
+    return torch.ops.trtllm.fp4_quantize(
+        normed.contiguous(),
+        sf_scale,
+        SF_VEC,
+        False,  # use_ue8m0
+        True,  # is_sf_swizzled_layout
+    )
+
+
+def make_sf_scale(normed: torch.Tensor) -> torch.Tensor:
+    """Per-tensor global SF scale fed to both the fused op and fp4_quantize
+    (both forward it to cvt_warp_fp16_to_fp4 as SFScaleVal). Matches the runtime
+    static NVFP4 input_scale convention in linear.py: FP8_MAX * E2M1_MAX / amax,
+    so the per-block E4M3 scale lands in range (SFValue = SFScaleVal*vecMax/6
+    <= 448)."""
+    amax = normed.abs().amax().float()
+    return (FP8_MAX * E2M1_MAX / amax).to(normed.device).view(1)
+
+
+# Packed-FP4 mantissa match threshold. The fused kernel computes RMSNorm in
+# registers (FMA) while reading the tensor to quantize it, vs the reference's
+# separate RMSNorm -> fp4_quantize; a few values right at an E2M1 step boundary
+# round the other way. With the correct (non-inverted) input scale this exercises
+# the full FP4 range, so bf16's 8-bit mantissa on small shapes can dip to ~0.984
+# (a real kernel error collapses the rate to < 0.5). fp16 stays >= 0.99.
+_FP4_MATCH_THRESHOLD = 0.97
+
+
+def assert_fp4_match(fp4_fused: torch.Tensor, fp4_ref: torch.Tensor, ctx: str):
+    assert fp4_fused.shape == fp4_ref.shape, (
+        f"shape mismatch {tuple(fp4_fused.shape)} vs {tuple(fp4_ref.shape)} ({ctx})"
+    )
+    match_rate = (fp4_fused == fp4_ref).float().mean().item()
+    assert match_rate >= _FP4_MATCH_THRESHOLD, (
+        f"FP4 packed match rate {match_rate:.4f} < {_FP4_MATCH_THRESHOLD} ({ctx})"
+    )
+
+
+def assert_sf_match(sf_fused: torch.Tensor, sf_ref: torch.Tensor, m: int, n: int, ctx: str):
+    """The per-block E4M3 scale factors must also match the unfused baseline,
+    not just the packed FP4 mantissas. The raw buffers are swizzled and padded
+    (rows -> 128, sf_cols -> 4) with uninitialized padding, so unswizzle both to
+    the logical [m, n/SF_VEC] grid and compare only the real region.
+
+    The SF is a per-16-elem-block amax. The fused kernel's FMA RMSNorm vs the
+    reference's separate RMSNorm shifts a block's amax by at most one E4M3
+    quantization step, so we require every block to match within +/-1 step
+    (exact value, not just a match rate -- on tiny shapes a 1-2 block diff out of
+    e.g. 32 blocks would fail a rate threshold despite being correct)."""
+    assert sf_fused.shape == sf_ref.shape, (
+        f"SF shape mismatch {tuple(sf_fused.shape)} vs {tuple(sf_ref.shape)} ({ctx})"
+    )
+    # unswizzle expects the padded dims; then slice the valid [m, num_sf_cols].
+    padded_rows = pad_up(m, 128)
+    num_sf_cols = ceil_div(n, SF_VEC)
+    padded_cols = pad_up(num_sf_cols, 4) * SF_VEC
+    u_fused = unswizzle_sf(sf_fused, padded_rows, padded_cols, SF_VEC)[:m, :num_sf_cols]
+    u_ref = unswizzle_sf(sf_ref, padded_rows, padded_cols, SF_VEC)[:m, :num_sf_cols]
+    # E4M3 is monotonic in its uint8 bit pattern over the positive range the SF
+    # uses, so adjacent representable scales differ by 1 in uint8 -- a +/-1 step
+    # tolerance is exactly |int(fused) - int(ref)| <= 1.
+    diff = (u_fused.to(torch.int16) - u_ref.to(torch.int16)).abs()
+    max_diff = int(diff.max().item()) if diff.numel() else 0
+    assert max_diff <= 1, (
+        f"SF differs by {max_diff} E4M3 steps (> 1) in {int((diff > 1).sum())} block(s) ({ctx})"
+    )
+
+
+# --------------------------------------------------------------------------- #
+# fused_add_rmsnorm_fp4_quantize: residual_add -> rms_norm -> nvfp4_quantize    #
+# --------------------------------------------------------------------------- #
+@skip_unless_add_rmsnorm
+@pytest.mark.parametrize("m", [1, 16, 64, 128])
+@pytest.mark.parametrize("n", [32, 128, 512, 7168])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_add_rmsnorm_fp4_quantize_vs_separate(m, n, dtype):
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+    eps = 1e-6
+
+    hidden = torch.randn(m, n, dtype=dtype, device=device)
+    residual = torch.randn(m, n, dtype=dtype, device=device)
+    weight = torch.randn(n, dtype=dtype, device=device)
+
+    # Reference: add in fp32 (matches kernel), then rmsnorm, then quantize.
+    added = hidden.float() + residual.float()
+    normed_ref = rms_norm_ref(added.to(dtype), weight, eps)
+    sf_scale = make_sf_scale(normed_ref)
+    fp4_ref, sf_ref = fp4_quantize_ref(normed_ref, sf_scale)
+
+    # The op returns residual_out as a fresh tensor and must NOT mutate
+    # hidden_states (required for torch.compile functionalization).
+    hidden_fused = hidden.clone()
+    hidden_before = hidden_fused.clone()
+    quant_out, scale_out, residual_out = torch.ops.trtllm.fused_add_rmsnorm_fp4_quantize(
+        hidden_fused, residual, weight, sf_scale, eps, False
+    )
+
+    assert_fp4_match(quant_out, fp4_ref, f"add-rmsnorm m={m} n={n} {dtype}")
+    # The per-block scale factors must match the unfused baseline too.
+    assert_sf_match(scale_out, sf_ref, m, n, f"add-rmsnorm m={m} n={n} {dtype}")
+
+    # hidden_states must be left untouched (no in-place alias).
+    torch.testing.assert_close(hidden_fused, hidden_before, rtol=0, atol=0)
+    # residual_out must equal hidden + residual (the next layer's residual).
+    torch.testing.assert_close(residual_out.float(), added, rtol=2e-2, atol=2e-2)
+
+
+@skip_unless_add_rmsnorm
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_add_rmsnorm_fp4_quantize_return_norm_out(dtype):
+    """return_norm_out=True must also yield the BF16/FP16 post-RMSNorm value
+    (consumed by the DSA indexer's pre_indexer_proj)."""
+    torch.manual_seed(7)
+    device = torch.device("cuda")
+    m, n, eps = 64, 7168, 1e-6
+
+    hidden = torch.randn(m, n, dtype=dtype, device=device)
+    residual = torch.randn(m, n, dtype=dtype, device=device)
+    weight = torch.randn(n, dtype=dtype, device=device)
+
+    added = hidden.float() + residual.float()
+    normed_ref = rms_norm_ref(added.to(dtype), weight, eps)
+    sf_scale = make_sf_scale(normed_ref)
+    fp4_ref, sf_ref = fp4_quantize_ref(normed_ref, sf_scale)
+
+    hidden_fused = hidden.clone()
+    norm_out, quant_out, scale_out, residual_out = torch.ops.trtllm.fused_add_rmsnorm_fp4_quantize(
+        hidden_fused, residual, weight, sf_scale, eps, True
+    )
+
+    assert_fp4_match(quant_out, fp4_ref, f"add-rmsnorm return_norm_out {dtype}")
+    assert_sf_match(scale_out, sf_ref, m, n, f"add-rmsnorm return_norm_out {dtype}")
+    torch.testing.assert_close(norm_out, normed_ref, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(residual_out.float(), added, rtol=2e-2, atol=2e-2)
+
+
+# --------------------------------------------------------------------------- #
+# fused_rmsnorm_fp4_quantize: residual-less rms_norm -> nvfp4_quantize          #
+# --------------------------------------------------------------------------- #
+@skip_unless_rmsnorm
+@pytest.mark.parametrize("m", [1, 16, 64, 128])
+@pytest.mark.parametrize("n", [32, 128, 512, 1536])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_rmsnorm_fp4_quantize_vs_separate(m, n, dtype):
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+    eps = 1e-6
+
+    hidden = torch.randn(m, n, dtype=dtype, device=device)
+    weight = torch.randn(n, dtype=dtype, device=device)
+
+    normed_ref = rms_norm_ref(hidden, weight, eps)
+    sf_scale = make_sf_scale(normed_ref)
+    fp4_ref, sf_ref = fp4_quantize_ref(normed_ref, sf_scale)
+
+    quant_out, scale_out = torch.ops.trtllm.fused_rmsnorm_fp4_quantize(
+        hidden, weight, sf_scale, eps, False
+    )
+
+    assert_fp4_match(quant_out, fp4_ref, f"rmsnorm m={m} n={n} {dtype}")
+    assert_sf_match(scale_out, sf_ref, m, n, f"rmsnorm m={m} n={n} {dtype}")
+
+
+@skip_unless_rmsnorm
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_rmsnorm_fp4_quantize_return_norm_out(dtype):
+    torch.manual_seed(11)
+    device = torch.device("cuda")
+    m, n, eps = 64, 1536, 1e-6
+
+    hidden = torch.randn(m, n, dtype=dtype, device=device)
+    weight = torch.randn(n, dtype=dtype, device=device)
+
+    normed_ref = rms_norm_ref(hidden, weight, eps)
+    sf_scale = make_sf_scale(normed_ref)
+    fp4_ref, sf_ref = fp4_quantize_ref(normed_ref, sf_scale)
+
+    norm_out, quant_out, scale_out = torch.ops.trtllm.fused_rmsnorm_fp4_quantize(
+        hidden, weight, sf_scale, eps, True
+    )
+
+    assert_fp4_match(quant_out, fp4_ref, f"rmsnorm return_norm_out {dtype}")
+    assert_sf_match(scale_out, sf_ref, m, n, f"rmsnorm return_norm_out {dtype}")
+    torch.testing.assert_close(norm_out, normed_ref, rtol=2e-2, atol=2e-2)
+
+
+@skip_unless_rmsnorm
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_rmsnorm_fp4_quantize_does_not_mutate_input(dtype):
+    """Residual=false path must not write back into hidden_states."""
+    torch.manual_seed(3)
+    device = torch.device("cuda")
+    m, n, eps = 32, 512, 1e-6
+
+    hidden = torch.randn(m, n, dtype=dtype, device=device)
+    weight = torch.randn(n, dtype=dtype, device=device)
+    hidden_before = hidden.clone()
+    sf_scale = make_sf_scale(rms_norm_ref(hidden, weight, eps))
+
+    torch.ops.trtllm.fused_rmsnorm_fp4_quantize(hidden, weight, sf_scale, eps, False)
+    torch.testing.assert_close(hidden, hidden_before, rtol=0, atol=0)
+
+
+# m starts at 2: a single-row (m=1) column slice is still contiguous in PyTorch
+# (the row-pitch stride is irrelevant when there is only one row), so it would
+# not exercise the input_row_stride read path at all.
+@skip_unless_rmsnorm
+@pytest.mark.parametrize("m", [2, 16, 64])
+@pytest.mark.parametrize("n", [512, 1536])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_rmsnorm_fp4_quantize_strided_input(m, n, dtype):
+    """The key F1 feature: read a row-strided column slice in place (no copy).
+
+    Mimics ``q = kv_a_proj_with_mqa(x).split(...)[0]``: a leading column slice
+    of a wider projection whose last dim is unit-stride but whose row pitch is
+    the full projection width. The fused op must read this strided layout
+    directly and match the contiguous reference exactly (the read offset is the
+    only thing the stride changes; all outputs are written packed)."""
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+    eps = 1e-6
+    extra = 2 * 16  # trailing columns (e.g. kv_lora + rope), keep 16-aligned
+
+    wide = torch.randn(m, n + extra, dtype=dtype, device=device)
+    q = wide[:, :n]  # column slice: unit-stride last dim, row pitch n+extra
+    assert not q.is_contiguous()
+    assert q.stride(0) == n + extra and q.stride(1) == 1
+
+    weight = torch.randn(n, dtype=dtype, device=device)
+    normed_ref = rms_norm_ref(q.contiguous(), weight, eps)
+    sf_scale = make_sf_scale(normed_ref)
+    fp4_ref, _ = fp4_quantize_ref(normed_ref, sf_scale)
+
+    quant_out, scale_out = torch.ops.trtllm.fused_rmsnorm_fp4_quantize(
+        q, weight, sf_scale, eps, False
+    )
+
+    assert quant_out.is_contiguous(), "outputs must be packed"
+    assert_fp4_match(quant_out, fp4_ref, f"strided m={m} n={n} {dtype}")
+
+
+@skip_unless_rmsnorm
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_rmsnorm_fp4_quantize_contiguous_matches_strided(dtype):
+    """A contiguous tensor (input_row_stride==0 path) and the same data viewed
+    as a column slice (input_row_stride>0 path) must produce identical FP4."""
+    torch.manual_seed(99)
+    device = torch.device("cuda")
+    m, n, eps = 48, 1536, 1e-6
+
+    base = torch.randn(m, n, dtype=dtype, device=device)
+    weight = torch.randn(n, dtype=dtype, device=device)
+    sf_scale = make_sf_scale(rms_norm_ref(base, weight, eps))
+
+    # Packed path.
+    fp4_packed, _ = torch.ops.trtllm.fused_rmsnorm_fp4_quantize(
+        base.contiguous(), weight, sf_scale, eps, False
+    )
+
+    # Strided path: embed the same rows into a wider buffer as a column slice.
+    wide = torch.randn(m, n + 32, dtype=dtype, device=device)
+    wide[:, :n] = base
+    strided = wide[:, :n]
+    fp4_strided, _ = torch.ops.trtllm.fused_rmsnorm_fp4_quantize(
+        strided, weight, sf_scale, eps, False
+    )
+
+    torch.testing.assert_close(fp4_packed, fp4_strided, rtol=0, atol=0)
+
+
+# ---------------------------------------------------------------------------
+# RMSNorm (M, N) kernel-selection gate (CPU-only logic; no kernel launch).
+# Verifies the routing in RMSNorm._ws_kernel_eligible: the warp-specialized
+# fused_add_rms_norm_quant ("ws") serves only the large, contiguous, rank-2
+# residual edge with an in-range hidden dim; everything else (small M,
+# residual-less, row-strided, out-of-range N) uses the reduce_fusion kernel.
+# ---------------------------------------------------------------------------
+
+
+def _make_rmsnorm(n: int):
+    from tensorrt_llm._torch.modules.rms_norm import RMSNorm
+
+    # device="meta" so no allocation / GPU is needed; the gate only reads
+    # tensor shapes and contiguity, never data.
+    return RMSNorm(hidden_size=n, eps=1e-6, dtype=torch.bfloat16, device="meta")
+
+
+@pytest.mark.parametrize(
+    "m,n,residual,contiguous,expect_ws",
+    [
+        # Large, contiguous, in-range N, M >= threshold (4096) -> ws.
+        (4096, 7168, True, True, True),
+        (8192, 8192, True, True, True),
+        # M below the threshold -> reduce_fusion.
+        (2048, 7168, True, True, False),
+        (1, 7168, True, True, False),
+        # Residual-less edge (q_a) -> reduce_fusion (no ws equivalent).
+        (8192, 7168, False, True, False),
+        # N below ws minimum (2048) -> reduce_fusion (q_a's 1536 lands here).
+        (8192, 1536, True, True, False),
+        # N above ws maximum -> reduce_fusion.
+        (8192, 20480, True, True, False),
+        # Row-strided (non-contiguous) input -> reduce_fusion.
+        (8192, 7168, True, False, False),
+    ],
+)
+def test_rmsnorm_ws_kernel_gate(m, n, residual, contiguous, expect_ws):
+    norm = _make_rmsnorm(n)
+    hs = torch.empty((m, n), dtype=torch.bfloat16, device="meta")
+    res = torch.empty((m, n), dtype=torch.bfloat16, device="meta") if residual else None
+    if not contiguous:
+        # A column slice of a wider buffer: unit-stride last dim, larger pitch.
+        hs = torch.empty((m, n + 32), dtype=torch.bfloat16, device="meta")[:, :n]
+        assert not hs.is_contiguous()
+    assert norm._ws_kernel_eligible(hs, res) is expect_ws

From 22634da945241d121f247feb3f6f7506459dae5d Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Tue, 16 Jun 2026 02:48:18 -0700
Subject: [PATCH 2/9] [None][fix] dsv32: allocate fused RMSNorm residual_out
 via empty_cuda

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>
---
 cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp      | 18 ++++++++++++------
 .../_torch/custom_ops/cpp_custom_ops.py        |  8 ++++++--
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp b/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
index 350b8b9600ff..ad61f2de2fef 100644
--- a/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
+++ b/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
@@ -85,15 +85,21 @@ std::vector<at::Tensor> fused_add_rmsnorm_fp4_quantize(at::Tensor const& hidden_
     void* norm_out_ptr = nullptr;
     if (return_norm_out)
     {
-        norm_out = torch::empty_like(hidden_states);
+        norm_out = at::detail::empty_cuda(
+            input_shape.vec(), hidden_states.scalar_type(), hidden_states.device(), std::nullopt);
         norm_out_ptr = norm_out.mutable_data_ptr();
     }
 
-    // The kernel reads hidden_states (intermediate_buffer) and residual, adds
-    // them, and writes the sum to a distinct residual_out_buffer — hidden_states
-    // is never mutated, so the op is functionalizable under torch.compile (no
-    // output aliases an input) with no extra pre-kernel copy.
-    at::Tensor residual_out = torch::empty_like(hidden_states);
+    // Freshly allocate residual_out (input + residual is written here by the
+    // kernel) the same way as quant_out/scale_out and the ws op
+    // (fusedAddRMSNormQuant.cpp): empty_cuda, not empty_like. empty_cuda is a
+    // plain allocation, whereas empty_like carries the input's layout and can be
+    // lowered to a separate node in the torch.compile trace. The kernel reads
+    // hidden_states (intermediate_buffer) read-only and writes the sum into this
+    // distinct buffer, so hidden_states is never mutated and no output aliases an
+    // input (the op stays functionalizable) -- with no pre-kernel copy.
+    at::Tensor residual_out
+        = at::detail::empty_cuda(input_shape.vec(), hidden_states.scalar_type(), hidden_states.device(), std::nullopt);
 
     tensorrt_llm::kernels::RmsNormFp4QuantParams params{};
     params.bias_buffer = nullptr;
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index cc868b577cb3..e2ff0338fc69 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1310,9 +1310,13 @@ def _(
         quant_out = hidden_states.new_empty(quant_shape, dtype=torch.uint8)
         sf_out = hidden_states.new_empty((_swizzled_sf_size(m, k), ),
                                          dtype=torch.uint8)
-        residual_out = torch.empty_like(hidden_states)
+        # Fresh allocations (new_empty), matching the real op's empty_cuda; not
+        # empty_like, which would carry the input's layout in the trace.
+        residual_out = hidden_states.new_empty(tuple(hidden_states.shape),
+                                               dtype=hidden_states.dtype)
         if return_norm_out:
-            norm_out = torch.empty_like(hidden_states)
+            norm_out = hidden_states.new_empty(tuple(hidden_states.shape),
+                                               dtype=hidden_states.dtype)
             return [norm_out, quant_out, sf_out, residual_out]
         return [quant_out, sf_out, residual_out]
 

From a0353efae0afc089d31ad87146d0283b1149981e Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Tue, 23 Jun 2026 22:26:59 -0700
Subject: [PATCH 3/9] [None][fix] dsv32: enable PDL and lowerCamelCase the
 fused RMSNorm-quant 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>
---
 .../kernels/rmsNormFp4QuantKernels.cu         | 26 ++++++++++++-------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
index dc287b0752c8..bbe804f920a9 100644
--- a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
+++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
@@ -18,6 +18,7 @@
 #include "tensorrt_llm/common/cudaBf16Fallbacks.cuh"
 #include "tensorrt_llm/common/cudaTypeUtils.cuh"
 #include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/kernels/quantization.cuh"
 #include <cstdint>
 
@@ -168,7 +169,7 @@ inline __device__ int4 rms_norm(float denom, PackedStruct& vec, PackedStruct& we
 //   - Caller guarantees hidden_size % SF_VEC_SIZE == 0.
 template <typename T, bool Bias = false, bool Residual = false, bool Affine = false, bool UseSmem = false,
     bool OutNorm = false>
-__global__ void rms_norm_fp4_quant_kernel(RmsNormFp4QuantParams params, void* quant_out, void* scale_out,
+__global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params, void* quant_out, void* scale_out,
     void* norm_out_ptr, float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout,
     int input_row_stride)
 {
@@ -282,9 +283,9 @@ __global__ void rms_norm_fp4_quant_kernel(RmsNormFp4QuantParams params, void* qu
 }
 
 template <typename T, bool OutNorm = false>
-void launch_rms_norm_fp4_quant_kernel(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out,
-    void* norm_out_ptr, float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout,
-    cudaStream_t stream, int input_row_stride = 0)
+void launchRmsNormFp4QuantKernel(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out, void* norm_out_ptr,
+    float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout, cudaStream_t stream,
+    int input_row_stride = 0)
 {
     static constexpr int kPackedSize = kBytesPerAccess / sizeof(T);
     TLLM_CHECK(params.hidden_size % kPackedSize == 0);
@@ -301,11 +302,16 @@ void launch_rms_norm_fp4_quant_kernel(RmsNormFp4QuantParams& params, void* quant
     bool const has_weight = params.weight_buffer != nullptr;
 
     // Macro-dispatch over Bias/Residual/Affine/UseSmem and the OutNorm template arg.
+    // Launch through launchWithPdlWhenEnabled so the kernel's PDL primitives
+    // (cudaGridDependencySynchronize / cudaTriggerProgrammaticLaunchCompletion)
+    // are actually enabled when TLLM_ENABLE_PDL is set.
 #define DISPATCH_FP4_QUANT(BIAS, RESIDUAL, AFFINE, SMEM)                                                               \
     if (use_smem == SMEM)                                                                                              \
     {                                                                                                                  \
-        rms_norm_fp4_quant_kernel<T, BIAS, RESIDUAL, AFFINE, SMEM, OutNorm><<<cta_num, cta_size, smem_size, stream>>>( \
-            params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, input_row_stride);                \
+        tensorrt_llm::common::launchWithPdlWhenEnabled("rmsNormFp4Quant",                                              \
+            rmsNormFp4QuantKernel<T, BIAS, RESIDUAL, AFFINE, SMEM, OutNorm>, dim3(cta_num), dim3(cta_size),            \
+            static_cast<size_t>(smem_size), stream, params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr,      \
+            sf_layout, input_row_stride);                                                                              \
     }
 
     auto launch = [&]()
@@ -379,12 +385,12 @@ void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, voi
         {
 #ifdef ENABLE_BF16
         case nvinfer1::DataType::kBF16:
-            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<__nv_bfloat16, /*OutNorm=*/true>(
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/true>(
                 params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
             break;
 #endif
         case nvinfer1::DataType::kHALF:
-            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<half, /*OutNorm=*/true>(
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<half, /*OutNorm=*/true>(
                 params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
             break;
         default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant");
@@ -396,12 +402,12 @@ void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, voi
         {
 #ifdef ENABLE_BF16
         case nvinfer1::DataType::kBF16:
-            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<__nv_bfloat16, /*OutNorm=*/false>(
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/false>(
                 params, quant_out, scale_out, nullptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
             break;
 #endif
         case nvinfer1::DataType::kHALF:
-            rms_norm_fp4_quant::launch_rms_norm_fp4_quant_kernel<half, /*OutNorm=*/false>(
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<half, /*OutNorm=*/false>(
                 params, quant_out, scale_out, nullptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
             break;
         default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant");

From 8bd890e50d380ee2d4a200265038116690141806 Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Tue, 23 Jun 2026 22:27:20 -0700
Subject: [PATCH 4/9] [None][test] dsv32: production-kernel reference +
 bit-exact fused-epilogue 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>
---
 .../test_fused_rmsnorm_fp4_quantize.py        | 62 ++++++++++++++++++-
 1 file changed, 59 insertions(+), 3 deletions(-)

diff --git a/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py b/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
index 35f421f69fc9..7102035fd3dd 100644
--- a/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
+++ b/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
@@ -36,6 +36,7 @@
 import pytest
 import torch
 
+from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE
 from tensorrt_llm._torch.utils import ceil_div, pad_up, unswizzle_sf
 from tests.unittest.utils.util import getSMVersion
 
@@ -71,9 +72,22 @@ def fp4_quantize_available():
 
 
 def rms_norm_ref(hidden_states: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
-    """Reference RMSNorm matching tensorrt_llm/_torch/modules/rms_norm.py:
-    fp32 accumulate, rsqrt(var + eps), then weight * normed cast back to the
-    input dtype."""
+    """Reference RMSNorm. Uses the *production* flashinfer RMSNorm kernel (the
+    exact kernel RMSNorm.forward runs on the unfused path the fusion replaces)
+    when available, so the comparison is apples-to-apples; otherwise falls back
+    to the fp32 PyTorch formula matching tensorrt_llm/_torch/modules/rms_norm.py.
+
+    Note the fused kernel still differs from this reference by ULPs: it does the
+    fp32 reduction + RMSNorm in registers in a single pass, whereas the unfused
+    path materializes the normed bf16 in a separate kernel before fp4_quantize.
+    Those ULP differences can flip a value across an E2M1/E4M3 step, so the
+    cross-path comparison is a high match rate (see _FP4_MATCH_THRESHOLD), while
+    the exact fused-epilogue equivalence is checked separately in
+    assert_fp4_bitexact_from_norm."""
+    if IS_FLASHINFER_AVAILABLE:
+        from tensorrt_llm._torch.custom_ops import flashinfer_rmsnorm
+
+        return flashinfer_rmsnorm(hidden_states.contiguous(), weight, eps)
     input_dtype = hidden_states.dtype
     hidden_states = hidden_states.to(torch.float32)
     variance = hidden_states.pow(2).mean(-1, keepdim=True)
@@ -157,6 +171,38 @@ def assert_sf_match(sf_fused: torch.Tensor, sf_ref: torch.Tensor, m: int, n: int
     )
 
 
+def assert_fp4_bitexact_from_norm(
+    fp4_fused: torch.Tensor,
+    sf_fused: torch.Tensor,
+    norm_out: torch.Tensor,
+    sf_scale: torch.Tensor,
+    m: int,
+    n: int,
+    ctx: str,
+):
+    """Bit-exact check of the fused NVFP4 epilogue. The fused kernel returns the
+    same post-RMSNorm BF16 value (norm_out) that it quantized, so quantizing that
+    exact tensor with the standalone fp4_quantize must reproduce the fused
+    quant_out/scale_out bit-for-bit -- there is no rounding-order freedom left
+    once the normed input is fixed. (This isolates the epilogue from the RMSNorm
+    ULP differences that the cross-path match-rate check tolerates.)
+
+    quant_out is dense row-major [m, n/2] so it's compared directly. scale_out is
+    swizzled+padded (rows -> 128, sf_cols -> 4) with uninitialized padding, so the
+    two separate allocations differ in the pad region only; compare the
+    unswizzled valid [m, n/SF_VEC] grid where bit-exactness must hold."""
+    fp4_req, sf_req = fp4_quantize_ref(norm_out, sf_scale)
+    assert torch.equal(fp4_fused, fp4_req), (
+        f"fused FP4 not bit-exact vs fp4_quantize(norm_out) ({ctx})"
+    )
+    padded_rows = pad_up(m, 128)
+    num_sf_cols = ceil_div(n, SF_VEC)
+    padded_cols = pad_up(num_sf_cols, 4) * SF_VEC
+    u_fused = unswizzle_sf(sf_fused, padded_rows, padded_cols, SF_VEC)[:m, :num_sf_cols]
+    u_req = unswizzle_sf(sf_req, padded_rows, padded_cols, SF_VEC)[:m, :num_sf_cols]
+    assert torch.equal(u_fused, u_req), f"fused SF not bit-exact vs fp4_quantize(norm_out) ({ctx})"
+
+
 # --------------------------------------------------------------------------- #
 # fused_add_rmsnorm_fp4_quantize: residual_add -> rms_norm -> nvfp4_quantize    #
 # --------------------------------------------------------------------------- #
@@ -222,6 +268,11 @@ def test_fused_add_rmsnorm_fp4_quantize_return_norm_out(dtype):
 
     assert_fp4_match(quant_out, fp4_ref, f"add-rmsnorm return_norm_out {dtype}")
     assert_sf_match(scale_out, sf_ref, m, n, f"add-rmsnorm return_norm_out {dtype}")
+    # Bit-exact: re-quantizing the fused op's own norm_out must reproduce its
+    # quant_out/scale_out exactly (isolates the epilogue from RMSNorm ULP noise).
+    assert_fp4_bitexact_from_norm(
+        quant_out, scale_out, norm_out, sf_scale, m, n, f"add-rmsnorm return_norm_out {dtype}"
+    )
     torch.testing.assert_close(norm_out, normed_ref, rtol=2e-2, atol=2e-2)
     torch.testing.assert_close(residual_out.float(), added, rtol=2e-2, atol=2e-2)
 
@@ -273,6 +324,11 @@ def test_fused_rmsnorm_fp4_quantize_return_norm_out(dtype):
 
     assert_fp4_match(quant_out, fp4_ref, f"rmsnorm return_norm_out {dtype}")
     assert_sf_match(scale_out, sf_ref, m, n, f"rmsnorm return_norm_out {dtype}")
+    # Bit-exact: re-quantizing the fused op's own norm_out must reproduce its
+    # quant_out/scale_out exactly (isolates the epilogue from RMSNorm ULP noise).
+    assert_fp4_bitexact_from_norm(
+        quant_out, scale_out, norm_out, sf_scale, m, n, f"rmsnorm return_norm_out {dtype}"
+    )
     torch.testing.assert_close(norm_out, normed_ref, rtol=2e-2, atol=2e-2)
 
 

From b63b06b4c19f8b6f45749b3431ab9c4b4c43c470 Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Wed, 24 Jun 2026 23:51:11 -0700
Subject: [PATCH 5/9] [None][test] dsv32: match production fp4_quantize call in
 bit-exact reference

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>
---
 .../_torch/modules/test_fused_rmsnorm_fp4_quantize.py    | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py b/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
index 7102035fd3dd..90585efd5c0d 100644
--- a/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
+++ b/tests/unittest/_torch/modules/test_fused_rmsnorm_fp4_quantize.py
@@ -102,13 +102,16 @@ def rms_norm_ref(hidden_states: torch.Tensor, weight: torch.Tensor, eps: float)
 
 
 def fp4_quantize_ref(normed: torch.Tensor, sf_scale: torch.Tensor):
-    """Unfused baseline: the standalone NVFP4 quantize the fusion replaces."""
+    """Unfused baseline: the exact standalone NVFP4 input-quantize the fusion
+    replaces. Mirrors the production static-NVFP4 call in
+    NVFP4LinearMethod._input_prepare (linear.py):
+        torch.ops.trtllm.fp4_quantize(input, input_scale, scaling_vector_size, False)
+    i.e. sfVecSize=16, sfUseUE8M0=False, isSfSwizzledLayout defaulting to True."""
     return torch.ops.trtllm.fp4_quantize(
         normed.contiguous(),
         sf_scale,
         SF_VEC,
-        False,  # use_ue8m0
-        True,  # is_sf_swizzled_layout
+        False,  # sfUseUE8M0
     )
 
 

From 4b59b8f80fd9a578f96572a82a77fc5d3e845dd8 Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Tue, 7 Jul 2026 00:43:23 -0700
Subject: [PATCH 6/9] [None][refactor] dsv32: make RMSNorm.nvfp4_scale an
 explicit attribute
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 tensorrt_llm/_torch/models/modeling_deepseekv3.py | 11 +++++------
 tensorrt_llm/_torch/models/modeling_nemotron_h.py |  4 ++--
 tensorrt_llm/_torch/modules/linear.py             |  4 ++--
 tensorrt_llm/_torch/modules/mla.py                | 13 ++-----------
 tensorrt_llm/_torch/modules/rms_norm.py           | 11 +++++++++--
 5 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index c10fc2c81f9b..f2344e6c3244 100755
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -147,10 +147,10 @@ def moe_reduce_add_shared_output(routed_output, shared_output, out=None):
 
 
 def _static_nvfp4_input_scale(linear):
-    """Return ``linear``'s calibrated NVFP4 input_scale if it is eligible to be
+    """Return `linear`'s calibrated NVFP4 input_scale if it is eligible to be
     folded into a producing RMSNorm's fused NVFP4 quantize, else None.
 
-    Eligibility is the shared ``is_static_nvfp4_input_eligible`` predicate, also
+    Eligibility is the shared `is_static_nvfp4_input_eligible` predicate, also
     used by MLA._resolve_qa_fused_scale, so every fold site agrees on what
     "static-NVFP4 eligible" means."""
     if is_static_nvfp4_input_eligible(linear):
@@ -1446,7 +1446,7 @@ def forward(
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if residual is None:
             residual = hidden_states
-            if getattr(self.input_layernorm, "nvfp4_scale", None) is not None:
+            if self.input_layernorm.nvfp4_scale is not None:
                 # Layer-0 prologue: the residual-less input_layernorm folds the
                 # NVFP4 kv_a_proj input-quant via its attached .nvfp4_scale
                 # (RMSNorm.forward), returning an Fp4QuantizedTensor.
@@ -1588,7 +1588,7 @@ def _apply_next_layer_layernorm(self, hidden_states, residual):
         and dense layers."""
         if self.next_layer_layernorm is None:
             return hidden_states, residual
-        if getattr(self.next_layer_layernorm, "nvfp4_scale", None) is not None:
+        if self.next_layer_layernorm.nvfp4_scale is not None:
             return self.next_layer_layernorm(hidden_states,
                                              residual,
                                              return_norm_out=True)
@@ -1625,8 +1625,7 @@ def forward_mlp(
                         eps=self.post_attention_layernorm.variance_epsilon,
                     ),
                 )
-        elif getattr(self.post_attention_layernorm, "nvfp4_scale",
-                     None) is not None:
+        elif self.post_attention_layernorm.nvfp4_scale is not None:
             # Attention-DP path (no allreduce): post_attention_layernorm folds
             # the dense MLP's NVFP4 gate_up_proj input-quant into one kernel via
             # its attached .nvfp4_scale (RMSNorm.forward), mirroring the
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index d931c3889d91..5c2e60c4706a 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -553,7 +553,7 @@ def __init__(
 
     def post_load_weights(self):
         """Post-process after loading weights."""
-        if self.norm.is_nvfp4 and not hasattr(self.norm, "nvfp4_scale"):
+        if self.norm.is_nvfp4 and self.norm.nvfp4_scale is None:
             self._try_attach_nvfp4_scale()
 
     def _try_attach_nvfp4_scale(self):
@@ -602,7 +602,7 @@ def forward(
 
         if hasattr(self, 'pre_allreduce'):
             norm = self.norm
-            has_nvfp4_scale = hasattr(norm, 'nvfp4_scale')
+            has_nvfp4_scale = norm.nvfp4_scale is not None
             if norm.is_nvfp4 and has_nvfp4_scale and norm.return_hp_output:
                 fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4
             elif norm.is_nvfp4 and has_nvfp4_scale:
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index ec1ea0057758..bb61df0b81b7 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3498,11 +3498,11 @@ def pre_reload_weights(self):
 
 
 def is_static_nvfp4_input_eligible(linear) -> bool:
-    """Whether ``linear`` consumes a static (calibrated) NVFP4 input, making it
+    """Whether `linear` consumes a static (calibrated) NVFP4 input, making it
     eligible to have its input-quantize folded into a producing RMSNorm.
 
     Eligible iff the Linear has NVFP4 weights, a calibrated (static)
-    ``input_scale``, no AWQ ``pre_quant_scale``, and is not forced to dynamic
+    `input_scale`, no AWQ `pre_quant_scale`, and is not forced to dynamic
     quantization. This is the single canonical definition shared by every
     NVFP4-fold site (the layer-boundary / dense folds in modeling_deepseekv3.py
     and the q_a_layernorm -> q_b_proj fold in attention.py's MLA) so the gate
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index 6d19b3b106d4..85d3f88071c2 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -47,12 +47,7 @@
 from ..attention_backend.utils import create_attention
 from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
-from ..utils import (
-    Fp4QuantizedTensor,
-    is_torch_compiling,
-    maybe_compiled_cat,
-    maybe_compiled_copy_,
-)
+from ..utils import Fp4QuantizedTensor, is_torch_compiling, maybe_compiled_cat, maybe_compiled_copy_
 from .attention import (
     _helix_cp_allgather_input,
     _helix_cp_output_projection,
@@ -1300,11 +1295,7 @@ def _resolve_qa_fused_scale(self):
             qa = self.q_a_layernorm
             # q_b_proj must be static-NVFP4 (shared predicate) and q_a_layernorm
             # a plain (non-gemma, non-quantizing) RMSNorm.
-            if (
-                is_static_nvfp4_input_eligible(qb)
-                and not getattr(qa, "use_gemma", False)
-                and not getattr(qa, "is_nvfp4", False)
-            ):
+            if is_static_nvfp4_input_eligible(qb) and not qa.use_gemma and not qa.is_nvfp4:
                 eligible = True
         if eligible:
             # Mark q_a_layernorm as an NVFP4 norm and attach q_b_proj's static
diff --git a/tensorrt_llm/_torch/modules/rms_norm.py b/tensorrt_llm/_torch/modules/rms_norm.py
index fb58478e3f6e..e1bb57874d31 100644
--- a/tensorrt_llm/_torch/modules/rms_norm.py
+++ b/tensorrt_llm/_torch/modules/rms_norm.py
@@ -103,6 +103,14 @@ def __init__(
                 return_hp_output = False
         self.return_hp_output = return_hp_output
 
+        # Static NVFP4 input scale of the Linear this norm feeds, enabling the
+        # fused (add +) RMSNorm + NVFP4-quantize path in forward. It cannot be
+        # set at construction time because the consuming Linear's calibrated
+        # input_scale only exists after weight loading; the owning model
+        # attaches it afterwards (DeepseekV3 post_load_weights /
+        # MLA._resolve_qa_fused_scale). None keeps the fusion disabled.
+        self.nvfp4_scale: Optional[torch.Tensor] = None
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -130,8 +138,7 @@ def forward(
         # for MoE-gate consumers). When is_nvfp4 but no scale is attached yet
         # (e.g. a layer's first input_layernorm whose consumer has no static
         # scale), fall through to the plain norm below.
-        nvfp4_scale = getattr(self, "nvfp4_scale",
-                              None) if self.is_nvfp4 else None
+        nvfp4_scale = self.nvfp4_scale if self.is_nvfp4 else None
         if nvfp4_scale is not None and not self.use_gemma:
             return self._fused_nvfp4_quant(hidden_states, residual, nvfp4_scale,
                                            return_norm_out)

From b87c4d2972bf3f5c6e393d234b4daebc448e28a7 Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Tue, 7 Jul 2026 00:47:11 -0700
Subject: [PATCH 7/9] [None][fix] dsv32: warp-uniform FP4 quantize loop; unify
 kernel args 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>
---
 .../kernels/rmsNormFp4QuantKernels.cu         | 104 +++++++++++-------
 .../kernels/rmsNormFp4QuantKernels.h          |  71 ++++++------
 cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp     |  19 +++-
 3 files changed, 111 insertions(+), 83 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
index bbe804f920a9..b70a9549df86 100644
--- a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
+++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu
@@ -169,9 +169,7 @@ inline __device__ int4 rms_norm(float denom, PackedStruct& vec, PackedStruct& we
 //   - Caller guarantees hidden_size % SF_VEC_SIZE == 0.
 template <typename T, bool Bias = false, bool Residual = false, bool Affine = false, bool UseSmem = false,
     bool OutNorm = false>
-__global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params, void* quant_out, void* scale_out,
-    void* norm_out_ptr, float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout,
-    int input_row_stride)
+__global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params)
 {
     static constexpr int kPackedSize = kBytesPerAccess / sizeof(T);
     static constexpr int kSfVecSize = 16;
@@ -188,7 +186,7 @@ __global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params, void* quant_
     T const* weight_buffer = reinterpret_cast<T const*>(params.weight_buffer);
     T const* intermediate_buffer = reinterpret_cast<T const*>(params.intermediate_buffer);
     T* residual_out_buffer = reinterpret_cast<T*>(params.residual_out_buffer);
-    T* norm_out = reinterpret_cast<T*>(norm_out_ptr);
+    T* norm_out = reinterpret_cast<T*>(params.norm_out);
 
     int const block_offset = bid * params.hidden_size;
     // Input rows may be strided (e.g. a column slice of a wider projection,
@@ -196,7 +194,7 @@ __global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params, void* quant_
     // input_row_stride <= 0 it defaults to hidden_size (packed rows), making
     // this byte-identical to all existing callers. Only the INPUT read offset
     // uses the stride; every output (residual/norm/quant/scale) stays packed.
-    int const input_block_offset = bid * (input_row_stride > 0 ? input_row_stride : params.hidden_size);
+    int const input_block_offset = bid * (params.input_row_stride > 0 ? params.input_row_stride : params.hidden_size);
     int const thread_offset = tid * kPackedSize;
 
     if constexpr (Residual)
@@ -245,37 +243,68 @@ __global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params, void* quant_
     acc = block_reduce_sum(acc);
     float const denom = rsqrtf(acc / params.hidden_size + params.eps);
 
-    float const sf_scale = scale_factor_ptr ? *scale_factor_ptr : 1.f;
+    float const sf_scale = params.scale_factor_ptr ? *params.scale_factor_ptr : 1.f;
     int const hidden_dim_packed = params.hidden_size / kSfVecSize;
 
-    for (int offset = thread_offset; offset < params.hidden_size; offset += blockDim.x * kPackedSize)
+    // cvt_warp_fp16_to_fp4 performs full-mask __shfl_xor_sync exchanges, which
+    // require every lane of the warp to execute the call. When hidden_size /
+    // kPackedSize is not a multiple of kWarpSize (e.g. hidden_size = 32, 128,
+    // 8208), the tail warp would otherwise be only partially active in this
+    // loop -- undefined behavior. Iterate with a warp-uniform bound (the warp's
+    // lane-0 offset) so all 32 lanes stay converged through the shuffle, and
+    // mask the per-lane loads/stores instead. Out-of-range lanes feed zeros to
+    // the shuffle; this is safe because hidden_size % kSfVecSize == 0 makes the
+    // active region end on an SF-pair boundary, so a padding lane's xor-1
+    // partner is always another padding lane.
+    int const lane_id = tid % kWarpSize;
+    for (int offset = thread_offset; offset - lane_id * kPackedSize < params.hidden_size;
+         offset += blockDim.x * kPackedSize)
     {
-        if constexpr (UseSmem)
-        {
-            inter_vec.packed = *reinterpret_cast<int4 const*>(&smem[offset]);
-        }
-        if constexpr (Affine)
+        bool const valid = offset < params.hidden_size;
+        if (valid)
         {
-            weight_vec.packed = *reinterpret_cast<int4 const*>(weight_buffer + offset);
+            if constexpr (UseSmem)
+            {
+                inter_vec.packed = *reinterpret_cast<int4 const*>(&smem[offset]);
+            }
+            if constexpr (Affine)
+            {
+                weight_vec.packed = *reinterpret_cast<int4 const*>(weight_buffer + offset);
+            }
+            inter_vec.packed = rms_norm<T, Affine>(denom, inter_vec, weight_vec);
+            if constexpr (OutNorm)
+            {
+                *reinterpret_cast<int4*>(norm_out + offset) = inter_vec.packed;
+            }
         }
-        inter_vec.packed = rms_norm<T, Affine>(denom, inter_vec, weight_vec);
-        if constexpr (OutNorm)
+        else
         {
-            *reinterpret_cast<int4*>(norm_out + offset) = inter_vec.packed;
+            // Benign values for the warp-cooperative SF exchange below (a
+            // padding lane's first-loop inter_vec may be uninitialized).
+            inter_vec.packed = make_int4(0, 0, 0, 0);
         }
 
         // FP4 quantize this 8-element packed vec; warp-cooperate with the
         // neighbour thread (offset ^ kPackedSize) to compute a single SF for
-        // their joint 16-element block.
+        // their joint 16-element block. All lanes (incl. padding) must execute
+        // this call; padding lanes pass a null SF pointer and drop the result.
         ::tensorrt_llm::kernels::PackedVec<T> pv
             = *reinterpret_cast<::tensorrt_llm::kernels::PackedVec<T>*>(&inter_vec);
-        int const access_id = bid * (params.hidden_size / kPackedSize) + (offset / kPackedSize);
-        int const access_id_in_token = offset / kPackedSize;
-        uint8_t* sf_out_ptr = ::tensorrt_llm::kernels::cvt_quant_get_sf_out_offset<uint32_t, 2>(std::nullopt, bid,
-            access_id_in_token, std::nullopt, hidden_dim_packed, reinterpret_cast<uint32_t*>(scale_out), sf_layout);
-        reinterpret_cast<uint32_t*>(quant_out)[access_id]
-            = ::tensorrt_llm::kernels::cvt_warp_fp16_to_fp4<T, kSfVecSize, /*UE8M0_SF=*/false>(
-                pv, sf_scale, sf_out_ptr);
+        uint8_t* sf_out_ptr = nullptr;
+        if (valid)
+        {
+            int const access_id_in_token = offset / kPackedSize;
+            sf_out_ptr = ::tensorrt_llm::kernels::cvt_quant_get_sf_out_offset<uint32_t, 2>(std::nullopt, bid,
+                access_id_in_token, std::nullopt, hidden_dim_packed, reinterpret_cast<uint32_t*>(params.scale_out),
+                params.sf_layout);
+        }
+        uint32_t const quant_val = ::tensorrt_llm::kernels::cvt_warp_fp16_to_fp4<T, kSfVecSize, /*UE8M0_SF=*/false>(
+            pv, sf_scale, sf_out_ptr);
+        if (valid)
+        {
+            int const access_id = bid * (params.hidden_size / kPackedSize) + (offset / kPackedSize);
+            reinterpret_cast<uint32_t*>(params.quant_out)[access_id] = quant_val;
+        }
     }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200))
     cudaTriggerProgrammaticLaunchCompletion();
@@ -283,9 +312,7 @@ __global__ void rmsNormFp4QuantKernel(RmsNormFp4QuantParams params, void* quant_
 }
 
 template <typename T, bool OutNorm = false>
-void launchRmsNormFp4QuantKernel(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out, void* norm_out_ptr,
-    float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout, cudaStream_t stream,
-    int input_row_stride = 0)
+void launchRmsNormFp4QuantKernel(RmsNormFp4QuantParams const& params, cudaStream_t stream)
 {
     static constexpr int kPackedSize = kBytesPerAccess / sizeof(T);
     TLLM_CHECK(params.hidden_size % kPackedSize == 0);
@@ -310,8 +337,7 @@ void launchRmsNormFp4QuantKernel(RmsNormFp4QuantParams& params, void* quant_out,
     {                                                                                                                  \
         tensorrt_llm::common::launchWithPdlWhenEnabled("rmsNormFp4Quant",                                              \
             rmsNormFp4QuantKernel<T, BIAS, RESIDUAL, AFFINE, SMEM, OutNorm>, dim3(cta_num), dim3(cta_size),            \
-            static_cast<size_t>(smem_size), stream, params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr,      \
-            sf_layout, input_row_stride);                                                                              \
+            static_cast<size_t>(smem_size), stream, params);                                                           \
     }
 
     auto launch = [&]()
@@ -363,9 +389,7 @@ void launchRmsNormFp4QuantKernel(RmsNormFp4QuantParams& params, void* quant_out,
 
 } // namespace rms_norm_fp4_quant
 
-void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out, void* norm_out_ptr,
-    float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout, nvinfer1::DataType dataType,
-    cudaStream_t stream, int input_row_stride)
+void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, nvinfer1::DataType dataType, cudaStream_t stream)
 {
     // The NVFP4 epilogue (cvt_warp_fp16_to_fp4) is compiled only for
     // __CUDA_ARCH__ >= 1000 and emits zeros otherwise, so this kernel is correct
@@ -377,21 +401,21 @@ void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, voi
         "residualRmsNormFp4Quant requires SM 10.x (Blackwell); got SM %d. The fused NVFP4 epilogue is unsupported on "
         "this arch.",
         sm);
+    TLLM_CHECK_WITH_INFO(params.quant_out != nullptr && params.scale_out != nullptr,
+        "residualRmsNormFp4Quant requires quant_out and scale_out output buffers.");
     sync_check_cuda_error(stream);
-    bool const out_norm = (norm_out_ptr != nullptr);
+    bool const out_norm = (params.norm_out != nullptr);
     if (out_norm)
     {
         switch (dataType)
         {
 #ifdef ENABLE_BF16
         case nvinfer1::DataType::kBF16:
-            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/true>(
-                params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/true>(params, stream);
             break;
 #endif
         case nvinfer1::DataType::kHALF:
-            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<half, /*OutNorm=*/true>(
-                params, quant_out, scale_out, norm_out_ptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<half, /*OutNorm=*/true>(params, stream);
             break;
         default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant");
         }
@@ -402,13 +426,11 @@ void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, voi
         {
 #ifdef ENABLE_BF16
         case nvinfer1::DataType::kBF16:
-            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/false>(
-                params, quant_out, scale_out, nullptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/false>(params, stream);
             break;
 #endif
         case nvinfer1::DataType::kHALF:
-            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<half, /*OutNorm=*/false>(
-                params, quant_out, scale_out, nullptr, scale_factor_ptr, sf_layout, stream, input_row_stride);
+            rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<half, /*OutNorm=*/false>(params, stream);
             break;
         default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant");
         }
diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h
index a41f54b83145..86e998fea209 100644
--- a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h
+++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h
@@ -26,62 +26,61 @@ TRTLLM_NAMESPACE_BEGIN
 namespace kernels
 {
 
-// Parameters for the fused (optional residual-add +) RMSNorm + NVFP4
-// input-quantize kernel. Self-contained (does not borrow AllReduceParams) since
-// this kernel performs no allreduce — it backs the standalone thop ops
-// fused_add_rmsnorm_fp4_quantize / fused_rmsnorm_fp4_quantize on the
-// attention-DP path.
+// All parameters for the fused (optional residual-add +) RMSNorm + NVFP4
+// input-quantize kernel: inputs, outputs, and layout config alike, so the
+// launcher takes a single struct. Self-contained (does not borrow
+// AllReduceParams) since this kernel performs no allreduce — it backs the
+// standalone thop ops fused_add_rmsnorm_fp4_quantize /
+// fused_rmsnorm_fp4_quantize on the attention-DP path.
 struct RmsNormFp4QuantParams
 {
+    // --- inputs ---
     // Input values [m, hidden_size] (read with input_row_stride). Read-only: the
     // residual sum is written to residual_out_buffer, never back into this.
     void const* intermediate_buffer{nullptr};
     // residual_in [m, hidden_size]; nullptr disables the residual add.
     void const* residual_buffer{nullptr};
+    // optional bias add [hidden_size]; nullptr disables it.
+    void const* bias_buffer{nullptr};
+    // RMSNorm gamma [hidden_size]; nullptr selects the non-affine path.
+    void const* weight_buffer{nullptr};
+    // Device pointer to the global per-tensor scale (= 448*6 / amax for
+    // static-quant Linear); nullptr means 1.0.
+    float const* scale_factor_ptr{nullptr};
+
+    // --- outputs ---
+    // Packed FP4 (E2M1) values, 2 per byte.
+    void* quant_out{nullptr};
+    // E4M3 scaling factors (one per SF_VEC_SIZE=16 block), laid out per sf_layout.
+    void* scale_out{nullptr};
+    // Optional BF16/FP16 post-RMSNorm value (packed rows); nullptr to skip.
+    void* norm_out{nullptr};
     // residual_out [m, hidden_size] (packed): receives input + residual when
     // residual_buffer != nullptr. A distinct buffer from intermediate_buffer so
     // the input is never mutated (keeps the thop op functionalizable under
     // torch.compile) and no pre-kernel copy is needed. nullptr when no residual.
     void* residual_out_buffer{nullptr};
-    // optional bias add [hidden_size]; nullptr disables it.
-    void const* bias_buffer{nullptr};
-    // RMSNorm gamma [hidden_size]; nullptr selects the non-affine path.
-    void const* weight_buffer{nullptr};
+
+    // --- config ---
     int hidden_size{0};
     float eps{0.f};
     // Total element count (= m * hidden_size); used to derive the row count.
     int64_t elts_total{0};
+    // Scaling-factor layout (typically SWIZZLED).
+    ::tensorrt_llm::QuantizationSFLayout sf_layout{::tensorrt_llm::QuantizationSFLayout::SWIZZLED};
+    // Element stride between input rows in intermediate_buffer. 0 means
+    // "== hidden_size" (packed rows). Set >0 to read a strided slice (e.g. a
+    // column-slice of a wider projection) without a preceding contiguous copy.
+    // Outputs are always written packed.
+    int input_row_stride{0};
 };
 
 // Fused (optional residual-add +) RMSNorm + NVFP4 input-quantize. Folds RMSNorm
 // and the next op's NVFP4 input-quant so the (flashinfer RMSNorm + standalone
-// fp4_quantize) pair becomes one launch on the attention-DP path.
-//
-// Inputs (via params):
-//   - intermediate_buffer: the input values (read-only)
-//   - residual_buffer:     residual_in (set nullptr to disable)
-//   - residual_out_buffer: receives input + residual (packed); required when
-//                          residual_buffer != nullptr, else unused
-//   - bias_buffer:         optional bias add (set nullptr to disable)
-//   - weight_buffer:       RMSNorm gamma (set nullptr for no affine)
-//   - hidden_size, eps:    standard
-//
-// Outputs:
-//   - quant_out:        packed FP4 (E2M1) values, kSfVecSize=16 per byte
-//   - scale_out:        E4M3 scaling factors (one per 16-elem block)
-//   - norm_out_ptr:     optional BF16 normed value (matches the
-//                       OUT_QUANT_NVFP4 fusion shape); pass nullptr to skip
-//   - scale_factor_ptr: device pointer to the global per-tensor scale
-//                       (= 448*6 / amax for static-quant Linear)
-//   - sf_layout:        scaling-factor layout (typically Swizzled)
-//   - input_row_stride: element stride between input rows in intermediate_buffer.
-//                       Defaults to 0 meaning "== hidden_size" (packed rows), so
-//                       every existing caller is byte-identical. Set >0 to read a
-//                       strided slice (e.g. a column-slice of a wider projection)
-//                       without a preceding contiguous copy. Outputs stay packed.
-void residualRmsNormFp4Quant(RmsNormFp4QuantParams& params, void* quant_out, void* scale_out, void* norm_out_ptr,
-    float const* scale_factor_ptr, ::tensorrt_llm::QuantizationSFLayout sf_layout, nvinfer1::DataType dataType,
-    cudaStream_t stream, int input_row_stride = 0);
+// fp4_quantize) pair becomes one launch on the attention-DP path. All inputs,
+// outputs, and layout configuration are carried in params (see the struct
+// field docs above); dataType selects the fp16/bf16 instantiation.
+void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, nvinfer1::DataType dataType, cudaStream_t stream);
 
 } // namespace kernels
 
diff --git a/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp b/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
index ad61f2de2fef..420cb80df461 100644
--- a/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
+++ b/cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp
@@ -106,17 +106,20 @@ std::vector<at::Tensor> fused_add_rmsnorm_fp4_quantize(at::Tensor const& hidden_
     params.residual_buffer = residual.data_ptr();
     params.weight_buffer = norm_weight.data_ptr();
     params.intermediate_buffer = hidden_states.data_ptr();
+    params.scale_factor_ptr = static_cast<float const*>(scale_factor.data_ptr());
+    params.quant_out = quant_out.mutable_data_ptr();
+    params.scale_out = scale_out.mutable_data_ptr();
+    params.norm_out = norm_out_ptr;
     params.residual_out_buffer = residual_out.mutable_data_ptr();
     params.hidden_size = static_cast<int>(k);
     params.eps = static_cast<float>(eps);
     params.elts_total = hidden_states.numel();
+    params.sf_layout = tensorrt_llm::QuantizationSFLayout::SWIZZLED;
 
     auto const stream = at::cuda::getCurrentCUDAStream(hidden_states.get_device());
     auto const dtype = tensorrt_llm::runtime::TorchUtils::dataType(hidden_states.scalar_type());
 
-    tensorrt_llm::kernels::residualRmsNormFp4Quant(params, quant_out.mutable_data_ptr(), scale_out.mutable_data_ptr(),
-        norm_out_ptr, static_cast<float const*>(scale_factor.data_ptr()), tensorrt_llm::QuantizationSFLayout::SWIZZLED,
-        dtype, stream);
+    tensorrt_llm::kernels::residualRmsNormFp4Quant(params, dtype, stream);
 
     // residual_out holds the residual sum (= original hidden + original residual).
     if (return_norm_out)
@@ -209,16 +212,20 @@ std::vector<at::Tensor> fused_rmsnorm_fp4_quantize(at::Tensor const& hidden_stat
     params.residual_buffer = nullptr;
     params.weight_buffer = norm_weight.data_ptr();
     params.intermediate_buffer = hidden_states.data_ptr();
+    params.scale_factor_ptr = static_cast<float const*>(scale_factor.data_ptr());
+    params.quant_out = quant_out.mutable_data_ptr();
+    params.scale_out = scale_out.mutable_data_ptr();
+    params.norm_out = norm_out_ptr;
     params.hidden_size = static_cast<int>(k);
     params.eps = static_cast<float>(eps);
     params.elts_total = hidden_states.numel();
+    params.sf_layout = tensorrt_llm::QuantizationSFLayout::SWIZZLED;
+    params.input_row_stride = input_row_stride;
 
     auto const stream = at::cuda::getCurrentCUDAStream(hidden_states.get_device());
     auto const dtype = tensorrt_llm::runtime::TorchUtils::dataType(hidden_states.scalar_type());
 
-    tensorrt_llm::kernels::residualRmsNormFp4Quant(params, quant_out.mutable_data_ptr(), scale_out.mutable_data_ptr(),
-        norm_out_ptr, static_cast<float const*>(scale_factor.data_ptr()), tensorrt_llm::QuantizationSFLayout::SWIZZLED,
-        dtype, stream, input_row_stride);
+    tensorrt_llm::kernels::residualRmsNormFp4Quant(params, dtype, stream);
 
     if (return_norm_out)
     {

From 8fd50ecb768debc5282a2d3daa38acfa1ee47d08 Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Fri, 10 Jul 2026 00:54:01 -0700
Subject: [PATCH 8/9] [None][chore] dsv32: name the NVFP4 SF vec size constant;
 clean up a 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>
---
 tensorrt_llm/_torch/modules/attention.py     | 12 +++++++-----
 tensorrt_llm/_torch/modules/rms_norm.py      |  4 ++--
 tensorrt_llm/quantization/utils/fp4_utils.py |  4 ++++
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 57abfbd8291a..98f287b9df01 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -7,6 +7,7 @@
 
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.quantization.utils.fp4_utils import NVFP4_SF_VEC_SIZE
 
 from ..attention_backend import (AttentionForwardArgs, AttentionMetadata,
                                  FlashInferAttentionMetadata,
@@ -55,11 +56,12 @@ def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int):
             "_slice_hidden_states_to_num_tokens only supports swizzled FP4 "
             "scaling factors")
     sf = hidden_states.scaling_factor
-    # fp4 packs 2 elements/byte, so the logical hidden dim is fp4_cols*2 and the
-    # scale-factor column count is that / SF_VEC_SIZE(16). The swizzled SF is
-    # laid out in independent 128-row tiles, so the leading num_tokens rows'
-    # scale factors occupy exactly the first padded_rows*padded_cols bytes.
-    sf_cols = fp4.shape[-1] * 2 // 16
+    # fp4 packs 2 elements/byte, so the logical hidden dim is fp4_cols*2 and
+    # the scale-factor column count is that / NVFP4_SF_VEC_SIZE. The swizzled
+    # SF is laid out in independent 128-row tiles, so the leading num_tokens
+    # rows' scale factors occupy exactly the first padded_rows*padded_cols
+    # bytes.
+    sf_cols = fp4.shape[-1] * 2 // NVFP4_SF_VEC_SIZE
     padded_rows, padded_cols = compute_swizzled_sf_shape(num_tokens, sf_cols)
     sf_len = padded_rows * padded_cols
     sliced_bf16 = (hidden_states.bf16_hidden_states[:num_tokens]
diff --git a/tensorrt_llm/_torch/modules/rms_norm.py b/tensorrt_llm/_torch/modules/rms_norm.py
index e1bb57874d31..5218c87a8425 100644
--- a/tensorrt_llm/_torch/modules/rms_norm.py
+++ b/tensorrt_llm/_torch/modules/rms_norm.py
@@ -32,8 +32,8 @@
 # Row-count crossover for the layer-boundary add+RMSNorm+quant edge: below this
 # the one-CTA-per-row reduce_fusion kernel (fused_add_rmsnorm_fp4_quantize) is
 # faster; at/above it the warp-specialized kernel's DMA-ahead pipeline wins.
-# Measured at N=7168 on GB200 (ws first overtakes at M=4096, ~6% faster; M<=3072
-# favors reduce_fusion). See analysis/bench_rmsnorm_threshold_n7168.py.
+# Measured at N=7168 on GB200 (ws first overtakes at M=4096, ~6% faster;
+# M<=3072 favors reduce_fusion).
 _WS_M_THRESHOLD = 4096
 
 
diff --git a/tensorrt_llm/quantization/utils/fp4_utils.py b/tensorrt_llm/quantization/utils/fp4_utils.py
index c62a4e2527b9..6142a4e333bc 100644
--- a/tensorrt_llm/quantization/utils/fp4_utils.py
+++ b/tensorrt_llm/quantization/utils/fp4_utils.py
@@ -6,6 +6,10 @@
 SF_DTYPE = torch.uint8
 FLOAT4_E2M1X2 = torch.uint8
 
+# Number of elements sharing one NVFP4 E4M3 scale factor
+# (SF_VEC_SIZE in cpp/tensorrt_llm/kernels/quantization.h).
+NVFP4_SF_VEC_SIZE = 16
+
 # For GEMM autotuning.
 # Taken from https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime//modelConfig.h#L38
 # TODO: move to model config, tune for blackwell hardware

From a101813eac313e264f7d4858a96879ec5d46421f Mon Sep 17 00:00:00 2001
From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Date: Mon, 13 Jul 2026 02:04:18 -0700
Subject: [PATCH 9/9] [None][refactor] dsv32: rename
 Fp4QuantizedTensor.bf16_hidden_states; 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>
---
 .../_torch/attention_backend/sparse/dsa.py    |  4 +-
 .../_torch/models/modeling_deepseekv3.py      |  2 +-
 tensorrt_llm/_torch/modules/attention.py      | 50 +------------
 tensorrt_llm/_torch/modules/mla.py            | 75 ++++++++++++++++---
 tensorrt_llm/_torch/modules/rms_norm.py       | 15 ++--
 tensorrt_llm/_torch/utils.py                  | 14 ++--
 6 files changed, 84 insertions(+), 76 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index c82a19c31543..af85390cf72e 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -2885,10 +2885,10 @@ def pre_indexer_proj(
         # the BF16 post-RMSNorm value. Use that BF16 view here -- the indexer
         # weight is BF16 and the matmul needs a float input.
         if isinstance(hidden_states, Fp4QuantizedTensor):
-            assert hidden_states.bf16_hidden_states is not None, (
+            assert hidden_states.unquantized_hidden_states is not None, (
                 "pre_indexer_proj received Fp4QuantizedTensor without bf16 view; "
                 "the producer fusion must request return_norm_out=True")
-            hidden_states_bf = hidden_states.bf16_hidden_states
+            hidden_states_bf = hidden_states.unquantized_hidden_states
         else:
             hidden_states_bf = hidden_states
         hidden_float = _to_float(hidden_states_bf)
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index f2344e6c3244..2ccc34a3cda4 100755
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -1451,7 +1451,7 @@ def forward(
                 # NVFP4 kv_a_proj input-quant via its attached .nvfp4_scale
                 # (RMSNorm.forward), returning an Fp4QuantizedTensor.
                 # return_norm_out stashes the BF16 view DSA's pre_indexer_proj
-                # needs on that tensor's bf16_hidden_states.
+                # needs on that tensor's unquantized_hidden_states.
                 hidden_states = self.input_layernorm(hidden_states,
                                                      return_norm_out=True)[0]
             else:
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 98f287b9df01..80e9c67d74d3 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -7,7 +7,6 @@
 
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.quantization.utils.fp4_utils import NVFP4_SF_VEC_SIZE
 
 from ..attention_backend import (AttentionForwardArgs, AttentionMetadata,
                                  FlashInferAttentionMetadata,
@@ -20,58 +19,13 @@
                            cp_allgather, reducescatter)
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
-from ..utils import (Fp4QuantizedTensor, compute_swizzled_sf_shape,
-                     get_model_extra_attrs, is_nvfp4_marlin_enabled,
-                     is_torch_compiling)
+from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
+                     is_nvfp4_marlin_enabled, is_torch_compiling)
 from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from .multi_stream_utils import maybe_execute_in_parallel
 from .rotary_embedding import MRotaryEmbedding, RotaryEmbedding
 
 
-def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int):
-    """Drop CUDA-graph padding by slicing hidden_states to num_tokens rows.
-
-    For a plain tensor this is the usual row slice. For an Fp4QuantizedTensor
-    (produced when the previous layer's boundary fusion pre-quantized this
-    layer's input) the slice must act on the packed FP4 form:
-      - fp4_tensor: row-major [m, k//2], so [:num_tokens] is a plain row slice;
-      - scaling_factor: 1D swizzled buffer of padUp(m,128)*padUp(cols,4). NVFP4
-        quant is per-row independent and the swizzle is laid out in independent
-        128-row tiles, so the leading padUp(num_tokens,128)*padUp(cols,4) bytes
-        are exactly the scale factors for the first num_tokens rows (verified by
-        tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py).
-      - bf16_hidden_states (when present): plain row slice.
-    """
-    if not isinstance(hidden_states, Fp4QuantizedTensor):
-        return hidden_states[:num_tokens, ...]
-
-    fp4 = hidden_states.fp4_tensor
-    if fp4.shape[0] == num_tokens:
-        return hidden_states  # no padding to strip
-    # The row-slice math below relies on the swizzled 128-row-tile layout; a
-    # non-swizzled (linear) scale buffer would be sliced incorrectly and then
-    # silently relabeled as swizzled. Reject it rather than corrupt the SF.
-    if not hidden_states.is_sf_swizzled:
-        raise ValueError(
-            "_slice_hidden_states_to_num_tokens only supports swizzled FP4 "
-            "scaling factors")
-    sf = hidden_states.scaling_factor
-    # fp4 packs 2 elements/byte, so the logical hidden dim is fp4_cols*2 and
-    # the scale-factor column count is that / NVFP4_SF_VEC_SIZE. The swizzled
-    # SF is laid out in independent 128-row tiles, so the leading num_tokens
-    # rows' scale factors occupy exactly the first padded_rows*padded_cols
-    # bytes.
-    sf_cols = fp4.shape[-1] * 2 // NVFP4_SF_VEC_SIZE
-    padded_rows, padded_cols = compute_swizzled_sf_shape(num_tokens, sf_cols)
-    sf_len = padded_rows * padded_cols
-    sliced_bf16 = (hidden_states.bf16_hidden_states[:num_tokens]
-                   if hidden_states.bf16_hidden_states is not None else None)
-    return Fp4QuantizedTensor(fp4_tensor=fp4[:num_tokens].contiguous(),
-                              scaling_factor=sf.view(-1)[:sf_len].contiguous(),
-                              is_sf_swizzled=True,
-                              bf16_hidden_states=sliced_bf16)
-
-
 def extract_extra_attrs(layer_idx: str, attn_type: str):
     assert attn_type in ["mla", "attn"], "Invalid attention type"
     extra_attrs = get_model_extra_attrs()
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index 85d3f88071c2..e2f708525af2 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -26,6 +26,7 @@
 from tensorrt_llm._utils import get_sm_version, is_sm_100f, nvtx_range, nvtx_range_debug
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.quantization.utils.fp4_utils import NVFP4_SF_VEC_SIZE
 
 from ..attention_backend import (
     AttentionForwardArgs,
@@ -47,13 +48,18 @@
 from ..attention_backend.utils import create_attention
 from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
-from ..utils import Fp4QuantizedTensor, is_torch_compiling, maybe_compiled_cat, maybe_compiled_copy_
+from ..utils import (
+    Fp4QuantizedTensor,
+    compute_swizzled_sf_shape,
+    is_torch_compiling,
+    maybe_compiled_cat,
+    maybe_compiled_copy_,
+)
 from .attention import (
     _helix_cp_allgather_input,
     _helix_cp_output_projection,
     _helix_post_process,
     _helix_zero_kv_mask,
-    _slice_hidden_states_to_num_tokens,
     extract_extra_attrs,
 )
 from .linear import Linear, TensorParallelMode, is_static_nvfp4_input_eligible
@@ -72,6 +78,55 @@ def _is_env_truthy(name: str) -> bool:
     return os.environ.get(name, "").strip().lower() in ("1", "true", "on")
 
 
+def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int):
+    """Drop CUDA-graph padding by slicing hidden_states to num_tokens rows.
+
+    For a plain tensor this is the usual row slice. For an Fp4QuantizedTensor
+    (produced when the previous layer's boundary fusion pre-quantized this
+    layer's input) the slice must act on the packed FP4 form:
+      - fp4_tensor: row-major [m, k//2], so [:num_tokens] is a plain row slice;
+      - scaling_factor: 1D swizzled buffer of padUp(m,128)*padUp(cols,4). NVFP4
+        quant is per-row independent and the swizzle is laid out in independent
+        128-row tiles, so the leading padUp(num_tokens,128)*padUp(cols,4) bytes
+        are exactly the scale factors for the first num_tokens rows (verified by
+        tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py).
+      - unquantized_hidden_states (when present): plain row slice.
+    """
+    if not isinstance(hidden_states, Fp4QuantizedTensor):
+        return hidden_states[:num_tokens, ...]
+
+    fp4 = hidden_states.fp4_tensor
+    if fp4.shape[0] == num_tokens:
+        return hidden_states  # no padding to strip
+    # The row-slice math below relies on the swizzled 128-row-tile layout; a
+    # non-swizzled (linear) scale buffer would be sliced incorrectly and then
+    # silently relabeled as swizzled. Reject it rather than corrupt the SF.
+    if not hidden_states.is_sf_swizzled:
+        raise ValueError(
+            "_slice_hidden_states_to_num_tokens only supports swizzled FP4 scaling factors"
+        )
+    sf = hidden_states.scaling_factor
+    # fp4 packs 2 elements/byte, so the logical hidden dim is fp4_cols*2 and
+    # the scale-factor column count is that / NVFP4_SF_VEC_SIZE. The swizzled
+    # SF is laid out in independent 128-row tiles, so the leading num_tokens
+    # rows' scale factors occupy exactly the first padded_rows*padded_cols
+    # bytes.
+    sf_cols = fp4.shape[-1] * 2 // NVFP4_SF_VEC_SIZE
+    padded_rows, padded_cols = compute_swizzled_sf_shape(num_tokens, sf_cols)
+    sf_len = padded_rows * padded_cols
+    sliced_unquant = (
+        hidden_states.unquantized_hidden_states[:num_tokens]
+        if hidden_states.unquantized_hidden_states is not None
+        else None
+    )
+    return Fp4QuantizedTensor(
+        fp4_tensor=fp4[:num_tokens].contiguous(),
+        scaling_factor=sf.view(-1)[:sf_len].contiguous(),
+        is_sf_swizzled=True,
+        unquantized_hidden_states=sliced_unquant,
+    )
+
+
 def _extract_mla_extra_attrs(layer_idx: str):
     metadata, mla_layer = extract_extra_attrs(layer_idx, "mla")
     assert isinstance(mla_layer, MLA), "MLA layer must be a subclass of MLA or an instance of MLA"
@@ -120,7 +175,7 @@ def mla_custom_op_inplace(
     # When hidden_states_fp4/_sf are provided, the previous layer's boundary
     # fusion pre-quantized this layer's kv_a_proj NVFP4 input. Custom ops can't
     # take dataclasses, so the Fp4QuantizedTensor is passed as explicit tensors
-    # and reconstructed here (mirrors mla_dsa_proj). bf16_hidden_states carries
+    # and reconstructed here (mirrors mla_dsa_proj). unquantized_hidden_states carries
     # the un-quantized view for the o_proj output sizing in create_output.
     metadata, mla_layer = _extract_mla_extra_attrs(layer_idx)
     if hidden_states_fp4 is not None or hidden_states_sf is not None:
@@ -130,7 +185,7 @@ def mla_custom_op_inplace(
         hidden_states = Fp4QuantizedTensor(
             fp4_tensor=hidden_states_fp4,
             scaling_factor=hidden_states_sf,
-            bf16_hidden_states=hidden_states,
+            unquantized_hidden_states=hidden_states,
         )
     if mla_layer.is_deepseek_v4:
         # DeepSeek-V4 uses MQA mode and has no residual-less RMSNorm+quant
@@ -202,7 +257,7 @@ def mla_dsa_proj(
         hs = Fp4QuantizedTensor(
             fp4_tensor=hidden_states_fp4,
             scaling_factor=hidden_states_sf,
-            bf16_hidden_states=hidden_states,
+            unquantized_hidden_states=hidden_states,
         )
     else:
         hs = hidden_states
@@ -1080,12 +1135,12 @@ def create_output(self, hidden_states: torch.Tensor, num_contexts: int):
         # producing fold must have requested return_norm_out so the BF16 view
         # is present (folds feeding an MLA always do).
         if isinstance(hidden_states, Fp4QuantizedTensor):
-            assert hidden_states.bf16_hidden_states is not None, (
+            assert hidden_states.unquantized_hidden_states is not None, (
                 "MLA.create_output received an Fp4QuantizedTensor without a "
-                "bf16_hidden_states view; the producing fusion must use "
+                "unquantized_hidden_states view; the producing fusion must use "
                 "return_norm_out=True"
             )
-            hidden_states = hidden_states.bf16_hidden_states
+            hidden_states = hidden_states.unquantized_hidden_states
         num_tokens = hidden_states.shape[0]
         if self.is_deepseek_v4:
             hidden_size = self.num_heads_tp_cp * self.v_head_dim
@@ -3036,7 +3091,7 @@ def forward(
                     # can't take dataclasses, so pass tensors explicitly.
                     if isinstance(hidden_states, Fp4QuantizedTensor):
                         proj_outputs = torch.ops.trtllm.mla_dsa_proj(
-                            hidden_states.bf16_hidden_states,
+                            hidden_states.unquantized_hidden_states,
                             position_ids,
                             self.layer_idx_str,
                             hidden_states.fp4_tensor,
@@ -3066,7 +3121,7 @@ def forward(
                     # explicit tensors.
                     if isinstance(hidden_states, Fp4QuantizedTensor):
                         torch.ops.trtllm.mla_custom_op_inplace(
-                            hidden_states.bf16_hidden_states,
+                            hidden_states.unquantized_hidden_states,
                             position_ids,
                             self.layer_idx_str,
                             attn_output,
diff --git a/tensorrt_llm/_torch/modules/rms_norm.py b/tensorrt_llm/_torch/modules/rms_norm.py
index 5218c87a8425..0e65940959a6 100644
--- a/tensorrt_llm/_torch/modules/rms_norm.py
+++ b/tensorrt_llm/_torch/modules/rms_norm.py
@@ -94,10 +94,7 @@ def __init__(
         # the downstream linear layer handle FP4 quantization.
         if self.is_nvfp4:
             sm_version = get_sm_version()
-            # SM range narrowed to 10.x by this PR: the reduce_fusion kernels
-            # use cvt_warp_fp16_to_fp4 (guarded by __CUDA_ARCH__ >= 1000),
-            # which silently emits zeros on SM 9.x. Marlin is also
-            # incompatible with the fused NVFP4 path.
+            # Marlin is also incompatible with the fused NVFP4 path.
             if not (100 <= sm_version < 120) or is_nvfp4_marlin_enabled():
                 self.is_nvfp4 = False
                 return_hp_output = False
@@ -257,7 +254,7 @@ def _fused_nvfp4_quant(
         With a residual: returns ``(Fp4QuantizedTensor, residual_out)``. Without:
         returns an ``Fp4QuantizedTensor`` (the residual-less q_a edge). When
         ``return_norm_out`` the BF16 post-RMSNorm view is stashed on the
-        Fp4QuantizedTensor's ``bf16_hidden_states`` (and, residual-less, returned
+        Fp4QuantizedTensor's ``unquantized_hidden_states`` (and, residual-less, returned
         alongside). When ``self.return_hp_output`` the BF16 normed value is
         appended as a trailing output for MoE-gate consumers."""
         # Either flag means "also produce the BF16 normed value".
@@ -283,7 +280,7 @@ def _fused_nvfp4_quant(
                 bf16_hs = None
             fp4 = Fp4QuantizedTensor(act_fp4,
                                      act_sf,
-                                     bf16_hidden_states=bf16_hs)
+                                     unquantized_hidden_states=bf16_hs)
             outputs = [fp4, residual_out]
             if self.return_hp_output:
                 outputs.append(bf16_hs)
@@ -307,7 +304,9 @@ def _fused_nvfp4_quant(
             norm_out = None
         if len(orig_shape) != 2:
             act_fp4 = act_fp4.reshape(*orig_shape[:-1], n // 2)
-        fp4 = Fp4QuantizedTensor(act_fp4, act_sf, bf16_hidden_states=norm_out)
+        fp4 = Fp4QuantizedTensor(act_fp4,
+                                 act_sf,
+                                 unquantized_hidden_states=norm_out)
         return (fp4, norm_out) if return_norm_out else fp4
 
     def _fused_nvfp4_quant_ws(
@@ -348,7 +347,7 @@ def _fused_nvfp4_quant_ws(
         bf16_hs = results[3].reshape(orig_shape) if want_norm else None
         fp4 = Fp4QuantizedTensor(normed_fp4_u8,
                                  sf_fused,
-                                 bf16_hidden_states=bf16_hs)
+                                 unquantized_hidden_states=bf16_hs)
         outputs = [fp4, residual_out]
         if self.return_hp_output:
             outputs.append(bf16_hs)
diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py
index a1369917ad28..4f40e12c1ded 100644
--- a/tensorrt_llm/_torch/utils.py
+++ b/tensorrt_llm/_torch/utils.py
@@ -160,13 +160,13 @@ class Fp4QuantizedTensor:
     fp4_tensor: torch.Tensor
     scaling_factor: torch.Tensor
     is_sf_swizzled: bool = True
-    # Optional BF16/FP16 hidden-state view of the same logical activation.
-    # When the FP4 tensor is produced by a fused (add+)RMSNorm+NVFP4-quant that
-    # also returns the un-quantized (post-RMSNorm) value, this carries that
-    # BF16 tensor so downstream consumers needing the un-quantized form (e.g.
-    # DSv3.2's DSA indexer at sparse/dsa.py:pre_indexer_proj) can use it
-    # without dequantizing FP4.
-    bf16_hidden_states: Optional[torch.Tensor] = None
+    # Optional un-quantized (BF16/FP16) hidden-state view of the same logical
+    # activation. When the FP4 tensor is produced by a fused
+    # (add+)RMSNorm+NVFP4-quant that also returns the un-quantized
+    # (post-RMSNorm) value, this carries that tensor so downstream consumers
+    # needing the un-quantized form (e.g. DSv3.2's DSA indexer at
+    # sparse/dsa.py:pre_indexer_proj) can use it without dequantizing FP4.
+    unquantized_hidden_states: Optional[torch.Tensor] = None
 
     @property
     def shape(self):