diff --git a/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp b/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp
index 861557e3f1e8..fea1e9ea6475 100644
--- a/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp
+++ b/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)
diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py
index b2294eb77c98..a03e84014109 100644
--- a/tensorrt_llm/_torch/compilation/utils.py
+++ b/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",
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index e8d66a568480..5234d30a1275 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/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(
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index dce32f886909..b198ec4d8be6 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/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.
diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py
index 7e1682366863..4e3ed16346f2 100644
--- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py
+++ b/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
diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py
index 823ecabfc4c6..345bea92147e 100644
--- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py
+++ b/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
diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
index 4c5c2e466fd8..9973d056432a 100644
--- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index 86b960c29d93..2ab9513e3ede 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/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)):
diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py
index 08eafe62b978..89dd6dab882e 100644
--- a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py
+++ b/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
     )
diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py
index 4e40bf419666..e91d7202e663 100644
--- a/tests/unittest/_torch/modules/test_mla_registry.py
+++ b/tests/unittest/_torch/modules/test_mla_registry.py
@@ -40,12 +40,12 @@ def update_quant_config(self, _quant_config: object) -> None:
         pass
 
 
-def _make_mla(config: ModelConfig) -> MLA:
+def _make_mla(config: ModelConfig, cls: type[MLA] = MLA) -> MLA:
     position_embedding = PositionalEmbeddingParams(
         type=PositionEmbeddingType.rope_gpt_neox,
         rope=RopeParams(dim=2, max_positions=8),
     )
-    return MLA(
+    return cls(
         hidden_size=8,
         num_attention_heads=2,
         num_key_value_heads=1,
@@ -65,6 +65,110 @@ def _make_mla(config: ModelConfig) -> MLA:
     )
 
 
+class _OutputGateStub:
+    """Minimal ``MLA`` stand-in that records what the output-gate hook sees."""
+
+    forward = MLA.forward
+
+    def __init__(self, register_to_config: bool) -> None:
+        self.mapping = SimpleNamespace(has_cp_helix=lambda: False, enable_attention_dp=False)
+        self.layer_idx = 0
+        self.layer_idx_str = "0"
+        self.register_to_config = register_to_config
+        self.attn_output = torch.zeros(2, 4)
+        self.gate_calls: list[tuple[torch.Tensor, torch.Tensor]] = []
+        self.projected: list[torch.Tensor] = []
+
+    def _create_outputs(
+        self, hidden_states: torch.Tensor, attn_metadata: object
+    ) -> list[torch.Tensor]:
+        return [self.attn_output]
+
+    def forward_impl(
+        self,
+        position_ids: object,
+        hidden_states: torch.Tensor,
+        attn_metadata: object,
+        attn_output: list[torch.Tensor],
+        latent_cache_gen: object = None,
+    ) -> None:
+        attn_output[0].fill_(2.0)
+
+    def _forward_custom_op(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: object,
+        attn_output: list[torch.Tensor],
+        latent_cache_gen: object,
+    ) -> None:
+        attn_output[0].fill_(2.0)
+
+    def _apply_output_gate(
+        self, hidden_states: torch.Tensor, attn_output: torch.Tensor
+    ) -> torch.Tensor:
+        self.gate_calls.append((hidden_states, attn_output.clone()))
+        return attn_output * 5.0
+
+    def _project_output(
+        self,
+        attn_output: list[torch.Tensor],
+        position_ids: object,
+        attn_metadata: object,
+        all_reduce_params: object,
+    ) -> torch.Tensor:
+        self.projected.append(attn_output[0])
+        return attn_output[0]
+
+
+def test_base_output_gate_is_identity() -> None:
+    attn_output = torch.randn(2, 4)
+
+    assert MLA._apply_output_gate(None, torch.randn(2, 4), attn_output) is attn_output
+
+
+@pytest.mark.parametrize("register_to_config", [False, True])
+def test_output_gate_runs_between_attention_and_output_projection(
+    register_to_config: bool,
+) -> None:
+    mla_layer = _OutputGateStub(register_to_config)
+    hidden_states = torch.randn(2, 4)
+    attn_metadata = SimpleNamespace(num_contexts=0, num_tokens=2)
+
+    with patch.object(torch.ops.trtllm, "create_mla_outputs", return_value=mla_layer.attn_output):
+        output = mla_layer.forward(None, hidden_states, attn_metadata)
+
+    # The hook receives the module input and the completed attention output.
+    assert len(mla_layer.gate_calls) == 1
+    gate_hidden_states, gate_attn_output = mla_layer.gate_calls[0]
+    assert gate_hidden_states is hidden_states
+    torch.testing.assert_close(gate_attn_output, torch.full((2, 4), 2.0))
+    # Its result, not the raw attention output, is what gets projected.
+    assert mla_layer.projected == [output]
+    torch.testing.assert_close(output, torch.full((2, 4), 10.0))
+
+
+def test_output_gate_override_rejected_with_sparse_hooks() -> None:
+    class _GatedMLA(MLA):
+        def _apply_output_gate(
+            self, hidden_states: torch.Tensor, attn_output: torch.Tensor
+        ) -> torch.Tensor:
+            return attn_output * 2.0
+
+    config = ModelConfig(skip_create_weights_in_init=True)
+    with (
+        patch(
+            "tensorrt_llm._torch.modules.mla.create_attention",
+            side_effect=lambda *args, **kwargs: _FakeAttention(),
+        ),
+        patch(
+            "tensorrt_llm._torch.modules.mla.get_sparse_mla_hooks",
+            return_value=Mock(),
+        ),
+        pytest.raises(NotImplementedError, match="_apply_output_gate"),
+    ):
+        _make_mla(config, cls=_GatedMLA)
+
+
 def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None:
     target_config = ModelConfig(skip_create_weights_in_init=True)
     draft_config = ModelConfig(skip_create_weights_in_init=True)
diff --git a/tests/unittest/_torch/thop/parallel/test_kda_decode.py b/tests/unittest/_torch/thop/parallel/test_kda_decode.py
index 75e05e49ef32..0a90c3a1332f 100644
--- a/tests/unittest/_torch/thop/parallel/test_kda_decode.py
+++ b/tests/unittest/_torch/thop/parallel/test_kda_decode.py
@@ -420,7 +420,14 @@ def test_kda_decode_matches_fla(
         gate_lower_bound=gate_lower_bound,
     )
 
-    actual_output = torch.ops.trtllm.kda_decode(
+    # kda_decode is inplace-only: the caller supplies the output buffer and the
+    # kernel writes into it (the op returns ``()``).
+    actual_output = torch.empty(
+        (batch_size, 1, num_heads, head_dim),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    torch.ops.trtllm.kda_decode(
         inputs.x_q,
         inputs.x_k,
         inputs.x_v,
@@ -449,6 +456,7 @@ def test_kda_decode_matches_fla(
         0.0 if gate_lower_bound is None else gate_lower_bound,
         head_dim**-0.5,
         OUTPUT_NORM_EPS,
+        actual_output,
     )
 
     _assert_parity("output", actual_output, expected_output)