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

cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h

@@ -0,0 +1,87 @@
+/*
+ * 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
+{
+
+// 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};
+
+    // --- 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. 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
+
+TRTLLM_NAMESPACE_END

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

cpp/tensorrt_llm/thop/rmsNormFp4Quant.cpp

@@ -0,0 +1,264 @@
+/*
+ * 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 = at::detail::empty_cuda(
+            input_shape.vec(), hidden_states.scalar_type(), hidden_states.device(), std::nullopt);
+        norm_out_ptr = norm_out.mutable_data_ptr();
+    }
+
+    // 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;
+    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, 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.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, dtype, stream);
+
+    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);
+}

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.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.unquantized_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)