[TRTLLM-15160][feat] Add post-attention o_proj gate hook to shared MLA base and make trtllm::kda_decode inplace-only

cpp/tensorrt_llm/thop/kdaDecodeOp.cpp

@@ -213,30 +213,30 @@ void launch_selected_kernel(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::
     C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
 
-at::Tensor kda_decode_fusion_forward(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::Tensor w_q_t, at::Tensor w_k_t,
+// Inplace-only: the kernel writes the decode result into the caller-supplied
+// ``output`` and never allocates. Allocation lives in the Python wrapper, which
+// lets hot decode paths reuse one persistent, CUDA-graph-safe buffer.
+void kda_decode_fusion_forward(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::Tensor w_q_t, at::Tensor w_k_t,
     at::Tensor w_v_t, at::Tensor bias_q, at::Tensor bias_k, at::Tensor bias_v, at::Tensor cs_q, at::Tensor cs_k,
     at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, at::Tensor onorm_g,
     at::Tensor onorm_weight, std::optional<at::Tensor> ssm_state_indices, at::Tensor cu_seqlens, at::Tensor state,
     bool apply_onorm, bool update_conv_cache, bool use_lower_bound, bool apply_beta_sigmoid, double lower_bound,
-    double scale, double onorm_eps, std::optional<at::Tensor> output)
+    double scale, double onorm_eps, at::Tensor output)
 {
     validate_kda_decode_fusion_inputs(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v,
         a_log, g, dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, apply_onorm,
         update_conv_cache);
     int const B = static_cast<int>(x_q.size(1));
     int const HV = static_cast<int>(x_v.size(2));
-    auto out = output.has_value() ? *output : at::empty({B, 1, HV, kDimV}, x_q.options());
-    if (output.has_value())
-    {
-        TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::kBFloat16, "out must be a CUDA bfloat16 tensor");
-        TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
-        TORCH_CHECK(out.dim() == 4 && out.size(0) == B && out.size(1) == 1 && out.size(2) == HV && out.size(3) == kDimV,
-            "out must have shape [B, 1, HV, 128]");
-    }
+    TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16, "output must be a CUDA bfloat16 tensor");
+    TORCH_CHECK(output.device() == x_q.device(), "output must be on the same device as x_q");
+    TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
+    TORCH_CHECK(output.dim() == 4 && output.size(0) == B && output.size(1) == 1 && output.size(2) == HV
+            && output.size(3) == kDimV,
+        "output must have shape [B, 1, HV, 128]");
     launch_selected_kernel(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v, a_log, g,
-        dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, out, apply_onorm, update_conv_cache,
-        use_lower_bound, apply_beta_sigmoid, lower_bound, scale, onorm_eps);
-    return out;
+        dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, output, apply_onorm,
+        update_conv_cache, use_lower_bound, apply_beta_sigmoid, lower_bound, scale, onorm_eps);
 }
 
 } // namespace
@@ -256,7 +256,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "Tensor? ssm_state_indices, Tensor cu_seqlens, Tensor(d!) state, "
         "bool apply_onorm, bool update_conv_cache, bool use_lower_bound, "
         "bool apply_beta_sigmoid, float lower_bound, float scale, "
-        "float onorm_eps, Tensor(e!)? output=None) -> Tensor(e!)");
+        "float onorm_eps, Tensor(e!) output) -> ()");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)

tensorrt_llm/_torch/compilation/utils.py

@@ -171,6 +171,16 @@ def inplace_info():
         torch.ops.trtllm.inplace_slice_copy.default: {
             1: "dest"
         },
+        # kda_decode mutates three conv caches, the recurrent state and its
+        # output. Keys follow the ``Tensor(a!)..Tensor(e!)`` declaration order,
+        # not the positional argument index.
+        torch.ops.trtllm.kda_decode.default: {
+            1: "conv_state_q",
+            2: "conv_state_k",
+            3: "conv_state_v",
+            4: "state",
+            5: "output"
+        },
         torch.ops.trtllm.verify_dynamic_tree_rejection_out_op.default: {
             5: "acceptIndex",
             6: "acceptTokenNum",

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -361,41 +361,20 @@ def _(
                                 dtype=torch.int32)
 
     @torch.library.register_fake("trtllm::kda_decode")
-    def _(x_q: torch.Tensor,
-          x_k: torch.Tensor,
-          x_v: torch.Tensor,
-          w_q_t: torch.Tensor,
-          w_k_t: torch.Tensor,
-          w_v_t: torch.Tensor,
-          bias_q: torch.Tensor,
-          bias_k: torch.Tensor,
-          bias_v: torch.Tensor,
-          conv_state_q: torch.Tensor,
-          conv_state_k: torch.Tensor,
-          conv_state_v: torch.Tensor,
-          a_log: torch.Tensor,
-          g: torch.Tensor,
-          dt_bias: torch.Tensor,
-          beta: torch.Tensor,
-          onorm_g: torch.Tensor,
-          onorm_weight: torch.Tensor,
-          ssm_state_indices: Optional[torch.Tensor],
-          cu_seqlens: torch.Tensor,
-          state: torch.Tensor,
-          apply_onorm: bool,
-          update_conv_cache: bool,
-          use_lower_bound: bool,
-          apply_beta_sigmoid: bool,
-          lower_bound: float,
-          scale: float,
-          onorm_eps: float,
-          output: Optional[torch.Tensor] = None) -> torch.Tensor:
-        # Mirror the CUDA impl: write into the caller-provided output when
-        # given (schema returns Tensor(e!)), else allocate.
-        if output is not None:
-            return output
-        # x_q is [1, tokens, H, 128]; the kernel emits one row per token.
-        return x_q.new_empty((x_q.size(1), 1, x_v.size(2), x_v.size(3)))
+    def _(x_q: torch.Tensor, x_k: torch.Tensor, x_v: torch.Tensor,
+          w_q_t: torch.Tensor, w_k_t: torch.Tensor, w_v_t: torch.Tensor,
+          bias_q: torch.Tensor, bias_k: torch.Tensor, bias_v: torch.Tensor,
+          conv_state_q: torch.Tensor, conv_state_k: torch.Tensor,
+          conv_state_v: torch.Tensor, a_log: torch.Tensor, g: torch.Tensor,
+          dt_bias: torch.Tensor, beta: torch.Tensor, onorm_g: torch.Tensor,
+          onorm_weight: torch.Tensor, ssm_state_indices: Optional[torch.Tensor],
+          cu_seqlens: torch.Tensor, state: torch.Tensor, apply_onorm: bool,
+          update_conv_cache: bool, use_lower_bound: bool,
+          apply_beta_sigmoid: bool, lower_bound: float, scale: float,
+          onorm_eps: float, output: torch.Tensor) -> None:
+        # Inplace-only: the kernel writes into ``output``, so there is nothing
+        # to allocate and nothing to return.
+        return None
 
     @torch.library.register_fake("trtllm::minimax_m3_fp8_indexer_qk_norm_rope")
     def minimax_m3_fp8_indexer_qk_norm_rope_fake(

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -1548,7 +1548,8 @@ def __init__(
     def forward(
         self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata
     ) -> torch.Tensor:
-        out = self.mixer(hidden_states, attn_metadata)
+        # MLA.forward takes position_ids first; K3 is NoPE, so pass None.
+        out = self.mixer(None, hidden_states, attn_metadata)
         if self._o_allreduce is not None:
             # Head-sharded TP: sum the row-sharded o_proj partials across
             # the head-shard group.

tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py

@@ -20,7 +20,7 @@
 from ....logger import logger
 from ....mapping import Mapping
 from ....models.modeling_utils import QuantConfig
-from ...attention_backend import AttentionMetadata, TrtllmAttention, TrtllmAttentionMetadata
+from ...attention_backend import TrtllmAttention, TrtllmAttentionMetadata
 from ...attention_backend.interface import PositionalEmbeddingParams, RopeParams
 from ...model_config import ModelConfig
 from ..linear import Linear, TensorParallelMode
@@ -291,9 +291,9 @@ def __init__(
             rms_norm_eps=rms_norm_eps,
             flashinfer_mla_backend=_select_mla_generation_backend(model_config.get_quant_config()),
         )
-        # K3 calls forward_impl() directly to insert its output gate before
-        # the base row-parallel o_proj. The original executor metadata remains
-        # intact, so MLA performs its native mixed context/generation split.
+        # Run MLA eagerly instead of through the registered custom op: K3's
+        # accuracy is validated against the eager path. The output gate is a
+        # base hook (_apply_output_gate) and works on either branch.
         self.register_to_config = False
 
         self.use_output_gate = use_output_gate
@@ -341,29 +341,14 @@ def __init__(
         if dtype is not None:
             _meta_safe_cast_dtype(self, dtype)
 
-    def _apply_output_gate_and_o_proj(
+    def _apply_output_gate(
         self,
         hidden_states: torch.Tensor,
-        attn_out: torch.Tensor,
+        attn_output: torch.Tensor,
     ) -> torch.Tensor:
+        # Sigmoid gate on o_proj's input. g_proj matches o_proj's input
+        # sharding, so the multiply composes with the helix-CP output
+        # projection.
         if self.use_output_gate:
-            attn_out = attn_out * self.g_proj(hidden_states).sigmoid()
-        return self.o_proj(attn_out)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        # _create_outputs() rather than create_output(): the base implementation
-        # takes a list so a sparse-attention backend can append its own buffers,
-        # and it routes through the sparse hooks when they are installed. The
-        # dense path this module uses is element 0.
-        attn_outputs = self._create_outputs(hidden_states, attn_metadata)
-        super().forward_impl(
-            None,
-            hidden_states,
-            attn_metadata,
-            attn_output=attn_outputs,
-        )
-        return self._apply_output_gate_and_o_proj(hidden_states, attn_outputs[0])
+            return attn_output * self.g_proj(hidden_states).sigmoid()
+        return attn_output

tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py

@@ -133,7 +133,11 @@ def run_kda_decode_fusion_cuda(
         )
     if ssm_state_indices is None and not state.is_contiguous():
         raise ValueError("state must be contiguous because it is updated in place")
-    if out is not None:
+    if out is None:
+        # The op is inplace-only and never allocates, so supply a buffer here.
+        # Hot decode paths pass a persistent one instead.
+        out = x_q.new_empty((B, 1, HV, 128))
+    else:
         _require_cuda_bf16("out", out)
         if not out.is_contiguous():
             raise ValueError("out must be contiguous")
@@ -226,4 +230,5 @@ def run_kda_decode_fusion_cuda(
         float(scale),
         float(onorm_eps),
     )
-    return torch.ops.trtllm.kda_decode(*args, *launch_args, output=out)
+    torch.ops.trtllm.kda_decode(*args, *launch_args, output=out)
+    return out

tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py

@@ -194,6 +194,8 @@ def __init__(
         self._projection_aux_stream = aux_stream
         self._projection_fork_event = torch.cuda.Event()
         self._projection_join_event = torch.cuda.Event()
+        # Output buffer for the inplace-only ``trtllm::kda_decode`` op.
+        self._o_dense: Optional[torch.Tensor] = None
         self._packed_conv_weight: Optional[torch.Tensor] = None
         self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None
 
@@ -561,6 +563,24 @@ def forward_decode(
         H = self.num_heads
         B = x2d.shape[0]
 
+        # kda_decode writes its [B, 1, H, hd] result into this buffer. It is
+        # sized to the pool slot count on the first decode and never
+        # reallocated, because captured graphs bind this pointer.
+        if self._o_dense is None:
+            if torch.cuda.is_current_stream_capturing():
+                return self.forward_decode_fallback(
+                    x2d, conv_pool, ssm_pool, slot_indices, layer_cache, ssm_state_indices
+                )
+            self._o_dense = torch.empty(
+                max(conv_pool.shape[0], B), 1, H, hd, dtype=torch.bfloat16, device=x2d.device
+            )
+        else:
+            assert self._o_dense.shape[0] >= B, (
+                f"KDA decode output buffer holds {self._o_dense.shape[0]} rows "
+                f"but the decode batch is {B}; reallocating would corrupt "
+                f"previously captured CUDA graphs"
+            )
+
         def _project_qkvg() -> torch.Tensor:
             if self._qkvg_proj_weight is not None:
                 return torch.nn.functional.linear(x2d, self._qkvg_proj_weight)
@@ -611,7 +631,7 @@ def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]:
             state=ssm_pool,
             onorm_g=x_qkvg[:, 3 * d :].unflatten(-1, (H, hd)).unsqueeze(0),
             onorm_weight=self._onorm_w_f32,
-            out=None,
+            out=self._o_dense[:B],
             ssm_state_indices=ssm_state_indices,
             cu_seqlens=mamba_metadata._arange_buffer[: B + 1],
             scale=hd**-0.5,

tensorrt_llm/_torch/modules/mla.py

@@ -337,6 +337,18 @@ def __init__(
                 "fuse_qkv_a_proj=True; the separate q_a_proj layout is not "
                 "supported with sparse MLA."
             )
+        if self.sparse_attn_hooks is not None and type(self)._apply_output_gate is not (
+            MLA._apply_output_gate
+        ):
+            # _apply_output_gate is defined against the dense output path, where
+            # attn_output[0] is o_proj's input. Sparse hooks own _create_outputs
+            # and _project_output, so the tensor's rank, shape and projection
+            # differ and the gate would not compose.
+            raise NotImplementedError(
+                f"{type(self).__name__} overrides _apply_output_gate, which is "
+                "only defined for the dense MLA output path; it cannot be "
+                "combined with sparse MLA hooks."
+            )
 
         # Fold the residual-less q_a_layernorm -> q_b_proj NVFP4 input
         # quantization into one fused RMSNorm + FP4-quantize kernel. Resolve
@@ -1753,6 +1765,27 @@ def _forward_custom_op(
                 latent_cache_gen,
             )
 
+    def _apply_output_gate(
+        self,
+        hidden_states: torch.Tensor,
+        attn_output: torch.Tensor,
+    ) -> torch.Tensor:
+        """Transform the attention output before the output projection.
+
+        ``attn_output`` is the row-parallel ``o_proj``'s input and
+        ``hidden_states`` is the unquantized module input, so an override that
+        derives a gate from ``hidden_states`` must match ``o_proj``'s input
+        sharding. Runs on both the ``register_to_config`` and eager branches,
+        and ahead of the helix-CP reduce-scatter.
+
+        Dense output path only: ``__init__`` rejects an override combined with
+        sparse MLA hooks, which replace ``o_proj`` and may hand
+        ``attn_output[0]`` on in a different rank and layout.
+
+        The base is the identity.
+        """
+        return attn_output
+
     def _project_output(
         self,
         attn_output: list[torch.Tensor],
@@ -1806,14 +1839,17 @@ def forward(
             hidden_states, attn_metadata, self.mapping, self.layer_idx
         )
 
+        # Unquantized view of the module input, used by create_mla_outputs and
+        # by _apply_output_gate.
+        output_hidden_states = hidden_states
+        if isinstance(hidden_states, Fp4QuantizedTensor):
+            assert hidden_states.unquantized_hidden_states is not None, (
+                "MLA.forward received an Fp4QuantizedTensor without a "
+                "unquantized_hidden_states view"
+            )
+            output_hidden_states = hidden_states.unquantized_hidden_states
+
         if self.register_to_config:
-            output_hidden_states = hidden_states
-            if isinstance(hidden_states, Fp4QuantizedTensor):
-                assert hidden_states.unquantized_hidden_states is not None, (
-                    "MLA.forward received an Fp4QuantizedTensor without a "
-                    "unquantized_hidden_states view"
-                )
-                output_hidden_states = hidden_states.unquantized_hidden_states
             attn_output = [
                 torch.ops.trtllm.create_mla_outputs(output_hidden_states, self.layer_idx_str)
             ]
@@ -1833,6 +1869,7 @@ def forward(
                 latent_cache_gen=latent_cache_gen,
             )
 
+        attn_output[0] = self._apply_output_gate(output_hidden_states, attn_output[0])
         return self._project_output(attn_output, position_ids, attn_metadata, all_reduce_params)
 
     def resmooth_parameters(self, module_weight, module_weight_scale, recipe=(1, 128, 128)):

tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py

@@ -36,9 +36,10 @@ def mark_ranges():
         Qwen3NextSparseMoeBlock.forward
     )
     # Kimi K3. KDA runs directly through `KimiKDALinearAttention`.
-    # `KimiK3MLAAttention` overrides `MLA.forward`, so its range is on the
-    # `KimiMLARuntime` wrapper. The gate is entered through `compute_logits`,
-    # not `forward`. Its MLPs are the shared `GatedMLP`.
+    # `KimiK3MLAAttention` runs the shared `MLA.forward`, annotated below; the
+    # `KimiMLARuntime` range bounds the attention-plus-reduction around it. The
+    # gate is entered through `compute_logits`, not `forward`. Its MLPs are the
+    # shared `GatedMLP`.
     KimiKDALinearAttention.forward = nvtx.annotate("KimiKDALinearAttention")(
         KimiKDALinearAttention.forward
     )