From 2082dd6056b0c06b0f6f6961a192c1303649a63b Mon Sep 17 00:00:00 2001
From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
Date: Thu, 22 Oct 2026 17:58:38 +0000
Subject: [PATCH 1/6] [TRTLLM-17121][feat] Support phase-dependent MLA heads
 for Helix prefill

Project H/(TP*CP) MLA query heads for context tokens while retaining
H/TP query heads and the existing Helix exchange for generation.
Keep context KV expansion, RoPE and attention buffers consistent.

Reuse loaded BF16/FP8 projection weights and the existing GEMM paths.
Preserve the 128-row scale grid and prepare movable, nonpersistent
slice-scale buffers with invalidation on weight reload or transform.

Cover projection slices, scale layouts, cache lifecycle and complete
K3 MLA outputs across context, cached/chunked context and mixed decode.

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention/mla.py          | 165 +++++++----
 .../_torch/models/modeling_kimi_linear.py     |  34 ++-
 tensorrt_llm/_torch/modules/linear.py         | 124 ++++++--
 tensorrt_llm/quantization/utils/fp8_utils.py  |  73 +++++
 .../integration/test_lists/test-db/l0_cpu.yml |   1 +
 .../test_lists/test-db/l0_gb300.yml           |   2 +
 .../attention/multi_gpu/test_mla_helix.py     |   8 +-
 .../multi_gpu/test_mla_helix_context.py       | 278 ++++++++++++++++++
 .../modules/test_linear_output_slice.py       | 167 +++++++++++
 9 files changed, 767 insertions(+), 85 deletions(-)
 create mode 100644 tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py
 create mode 100644 tests/unittest/_torch/modules/test_linear_output_slice.py

diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py
index 328bb37accbf..882fb780bb6c 100644
--- a/tensorrt_llm/_torch/attention/mla.py
+++ b/tensorrt_llm/_torch/attention/mla.py
@@ -418,6 +418,11 @@ def __init__(
         assert self.num_heads % (tp_size * cp_size) == 0
         self.num_heads_tp = self.num_heads // tp_size
         self.num_heads_tp_cp = self.num_heads_tp // cp_size
+        self.num_context_heads = (
+            self.num_heads_tp_cp
+            if self.mapping.has_cp_helix() and self.sparse_attn_hooks is None
+            else self.num_heads_tp
+        )
         self.num_key_value_heads_tp = (self.num_key_value_heads + tp_size - 1) // tp_size
 
         if rms_norm_eps is None:
@@ -633,9 +638,13 @@ def yarn_get_mscale(scale=1, mscale=1):
             self.mha = create_attention(
                 config.attn_backend,
                 self.layer_idx,
-                self.num_heads_tp,
+                self.num_context_heads,
                 head_dim=self.qk_head_dim,
-                num_kv_heads=self.num_key_value_heads_tp,
+                num_kv_heads=(
+                    self.num_context_heads
+                    if self.num_context_heads != self.num_heads_tp
+                    else self.num_key_value_heads_tp
+                ),
                 pos_embd_params=self.pos_embd_params,
                 quant_config=quant_config,
                 q_scaling=q_scaling,
@@ -807,13 +816,13 @@ def apply_rope(
         q: torch.Tensor,
         k_pe: torch.Tensor,
         position_ids: torch.Tensor,
+        num_heads: Optional[int] = None,
     ) -> torch.Tensor:
-        q = q.view(-1, self.num_heads_tp, self.qk_head_dim)
-        q_pe = q[..., self.qk_nope_head_dim :].reshape(
-            -1, self.num_heads_tp * self.qk_rope_head_dim
-        )
+        num_heads = self.num_heads_tp if num_heads is None else num_heads
+        q = q.view(-1, num_heads, self.qk_head_dim)
+        q_pe = q[..., self.qk_nope_head_dim :].reshape(-1, num_heads * self.qk_rope_head_dim)
         q_pe, k_pe = self.rotary_emb(position_ids, [q_pe, k_pe])
-        q[..., self.qk_nope_head_dim :] = q_pe.view(-1, self.num_heads_tp, self.qk_rope_head_dim)
+        q[..., self.qk_nope_head_dim :] = q_pe.view(-1, num_heads, self.qk_rope_head_dim)
         return k_pe
 
     def _attn_forward_gen(
@@ -981,6 +990,38 @@ def forward_impl(
                 latent_cache_gen=latent_cache_gen,
             )
 
+    def _project_q_by_phase(
+        self, q: torch.Tensor, num_ctx_tokens: int
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if self.num_context_heads == self.num_heads_tp or num_ctx_tokens == 0:
+            projected = self.q_b_proj(q)
+            return projected[:num_ctx_tokens], projected[num_ctx_tokens:]
+        first_head = self.mapping.cp_rank * self.num_context_heads
+        first = first_head * self.qk_head_dim
+        last = first + self.num_context_heads * self.qk_head_dim
+        context = self.q_b_proj.forward_slice(q[:num_ctx_tokens], first, last)
+        generation = self.q_b_proj(q[num_ctx_tokens:]) if q.shape[0] > num_ctx_tokens else None
+        return context, generation
+
+    def _project_context_kv(self, compressed_kv: torch.Tensor) -> torch.Tensor:
+        if self.num_context_heads == self.num_heads_tp:
+            return self.kv_b_proj(compressed_kv)
+        first_head = self.mapping.cp_rank * self.num_context_heads
+        last_head = first_head + self.num_context_heads
+        key = self.kv_b_proj.forward_slice(
+            compressed_kv,
+            first_head * self.qk_nope_head_dim,
+            last_head * self.qk_nope_head_dim,
+        )
+        value_start = self.num_heads_tp * self.qk_nope_head_dim
+        value = self.kv_b_proj.forward_slice(
+            compressed_kv,
+            value_start + first_head * self.v_head_dim,
+            value_start + last_head * self.v_head_dim,
+        )
+        # Preserve the [all K | all V] row stride consumed by context FP8 preprocessing.
+        return torch.cat((key, value), dim=-1)
+
     def _forward_impl(
         self,
         position_ids: Optional[torch.Tensor],
@@ -1034,22 +1075,19 @@ def _forward_impl(
                 self.aux_stream,
             )
 
-        q, latent_cache = maybe_execute_in_parallel(
-            lambda: self.q_b_proj(q),
+        projected_q, latent_cache = maybe_execute_in_parallel(
+            lambda: self._project_q_by_phase(q, num_ctx_tokens),
             lambda: torch.concat([compressed_kv, k_pe], dim=-1),
             self.ln_events[0],
             self.ln_events[1],
             self.aux_stream,
         )
 
-        assert q.shape[0] == num_tokens, (
-            f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}"
-        )
+        q_ctx, q_gen = projected_q
 
         assert output is not None, "output must be provided"
 
         if num_contexts > 0:
-            q_ctx = q[:num_ctx_tokens, ...]
             compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...]
             k_pe_ctx = k_pe[:num_ctx_tokens, ...]
             latent_cache_ctx = latent_cache[:num_ctx_tokens, ...]
@@ -1058,7 +1096,9 @@ def _forward_impl(
                 # position_ids spans [ctx..., gen...] in mixed batches; slice to
                 # match q_ctx/k_pe_ctx so external RoPE uses ctx positions.
                 ctx_position_ids = position_ids[..., :num_ctx_tokens]
-                k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids)
+                k_pe_ctx = self.apply_rope(
+                    q_ctx, k_pe_ctx, ctx_position_ids, self.num_context_heads
+                )
                 # External RoPE is only used by backends that do not handle
                 # fused RoPE internally, so keep latent_cache in sync.
                 latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], dim=-1)
@@ -1077,7 +1117,7 @@ def _forward_impl(
             )
 
         if num_generations > 0:
-            q_gen = q[num_ctx_tokens:, ...]
+            assert q_gen is not None
             compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...]
             k_pe_gen = k_pe[num_ctx_tokens:, ...]
             if latent_cache_gen is None:
@@ -1130,11 +1170,17 @@ def _get_kvbproj_head_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
         cached = getattr(self, "_kvbproj_head_weights_cache", None)
         w = self.kv_b_proj.weight
         if cached is not None and cached[0] is w:
-            return cached[1], cached[2]
-        nope_rows = self.num_heads_tp * self.qk_nope_head_dim
-        w_k = w[:nope_rows].view(self.num_heads_tp, self.qk_nope_head_dim, self.kv_lora_rank)
-        w_v = w[nope_rows:].view(self.num_heads_tp, self.v_head_dim, self.kv_lora_rank)
-        self._kvbproj_head_weights_cache = (w, w_k, w_v)
+            w_k, w_v = cached[1], cached[2]
+        else:
+            nope_rows = self.num_heads_tp * self.qk_nope_head_dim
+            w_k = w[:nope_rows].view(self.num_heads_tp, self.qk_nope_head_dim, self.kv_lora_rank)
+            w_v = w[nope_rows:].view(self.num_heads_tp, self.v_head_dim, self.kv_lora_rank)
+            self._kvbproj_head_weights_cache = (w, w_k, w_v)
+        if self.num_context_heads != self.num_heads_tp:
+            first = self.mapping.cp_rank * self.num_context_heads
+            return w_k[first : first + self.num_context_heads], w_v[
+                first : first + self.num_context_heads
+            ]
         return w_k, w_v
 
     def forward_context_default(
@@ -1148,7 +1194,7 @@ def forward_context_default(
         latent_cache: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Dense MHA context path: expand KV via kv_b_proj and run attention."""
-        k = torch.empty_like(q).view(-1, self.num_heads_tp, self.qk_head_dim)
+        k = torch.empty_like(q).view(-1, self.num_context_heads, self.qk_head_dim)
         if self._use_kvbproj_strided_out():
             # Expand KV via per-head batched GEMMs that write k_nope directly
             # into the strided K buffer, skipping the [t, H*256] kv
@@ -1156,7 +1202,7 @@ def forward_context_default(
             num_tokens = compressed_kv.shape[0]
             w_k, w_v = self._get_kvbproj_head_weights()
             ckv = compressed_kv.unsqueeze(0).expand(
-                self.num_heads_tp, num_tokens, self.kv_lora_rank
+                self.num_context_heads, num_tokens, self.kv_lora_rank
             )
             # [H, t, qk_nope] view with strides (qk_head_dim, H*qk_head_dim, 1)
             k_nope_out = k[..., : self.qk_nope_head_dim].transpose(0, 1)
@@ -1165,33 +1211,36 @@ def forward_context_default(
             # H*(qk_nope+v) -- the layout of the kv.split() view. Keep that
             # contract: allocate a kv-shaped buffer and only write its v
             # half; the k_nope half is never written or read.
-            nope_cols = self.num_heads_tp * self.qk_nope_head_dim
-            v_cols = self.num_heads_tp * self.v_head_dim
+            nope_cols = self.num_context_heads * self.qk_nope_head_dim
+            v_cols = self.num_context_heads * self.v_head_dim
             kv_buf = torch.empty(num_tokens, nope_cols + v_cols, dtype=q.dtype, device=q.device)
             v = kv_buf[:, nope_cols:]
             # [H, t, v_head] view with strides (v_head, H*(qk_nope+v), 1)
-            v_out = v.unflatten(-1, (self.num_heads_tp, self.v_head_dim)).transpose(0, 1)
+            v_out = v.unflatten(-1, (self.num_context_heads, self.v_head_dim)).transpose(0, 1)
             # cuBLAS consumes the transposed weights and writes strided
             # outputs directly, without materializing contiguous copies.
             torch.ops.trtllm.bmm_out(ckv, w_k.transpose(1, 2), k_nope_out)
             torch.ops.trtllm.bmm_out(ckv, w_v.transpose(1, 2), v_out)
         else:
-            kv = self.kv_b_proj(compressed_kv)
+            kv = self._project_context_kv(compressed_kv)
             k_nope, v = kv.split(
-                [self.num_heads_tp * self.qk_nope_head_dim, self.num_heads_tp * self.v_head_dim],
+                [
+                    self.num_context_heads * self.qk_nope_head_dim,
+                    self.num_context_heads * self.v_head_dim,
+                ],
                 -1,
             )
 
             maybe_compiled_copy_(
                 k[..., : self.qk_nope_head_dim],
-                k_nope.view(-1, self.num_heads_tp, self.qk_nope_head_dim),
+                k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim),
             )
         # When rope_fusion=True (apply_rotary_emb=False), the rope portion
         # of k is left uninitialized here; the fused attention kernel
         # handles k_pe RoPE via latent_cache instead.
         if self.apply_rotary_emb:
             k[..., self.qk_nope_head_dim :] = k_pe.view(-1, 1, self.qk_rope_head_dim)
-        k = k.view(-1, self.num_heads_tp * self.qk_head_dim)
+        k = k.view(-1, self.num_context_heads * self.qk_head_dim)
 
         attn_output = self.mha.forward(
             q,
@@ -1238,21 +1287,21 @@ def forward_context_with_cached_kv(
         assert full_k_pe.is_contiguous()
 
         # compute full_k_nope and full_v from full_compressed_kv
-        full_kv = self.kv_b_proj(full_compressed_kv)
+        full_kv = self._project_context_kv(full_compressed_kv)
         full_k_nope, full_v = full_kv.split(
             [
-                self.num_heads_tp * self.qk_nope_head_dim,
-                self.num_heads_tp * self.v_head_dim,
+                self.num_context_heads * self.qk_nope_head_dim,
+                self.num_context_heads * self.v_head_dim,
             ],
             -1,
         )
 
-        full_k_nope = full_k_nope.view(-1, self.num_heads_tp, self.qk_nope_head_dim)
+        full_k_nope = full_k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim)
         full_k_pe = full_k_pe.view(-1, 1, self.qk_rope_head_dim)
         full_k = maybe_compiled_cat(
-            (full_k_nope, full_k_pe.expand(-1, self.num_heads_tp, -1)), dim=-1
+            (full_k_nope, full_k_pe.expand(-1, self.num_context_heads, -1)), dim=-1
         )
-        full_k = full_k.view(-1, self.num_heads_tp * self.qk_head_dim)
+        full_k = full_k.view(-1, self.num_context_heads * self.qk_head_dim)
 
         # release pytorch activation memory
         full_compressed_kv = None
@@ -1295,19 +1344,19 @@ def forward_context_with_chunked_prefill(
 
         # [total_token_q, num_heads, 2] -> [total_token_q, num_heads] float2
         self.softmax_stats_tensor = torch.empty(
-            (attn_metadata.num_ctx_tokens, self.num_heads_tp, 2),
+            (attn_metadata.num_ctx_tokens, self.num_context_heads, 2),
             dtype=torch.float,
             device="cuda",
         )
         self.temp_softmax_stats_tensor = torch.empty(
-            (attn_metadata.num_ctx_tokens, self.num_heads_tp, 2),
+            (attn_metadata.num_ctx_tokens, self.num_context_heads, 2),
             dtype=torch.float,
             device="cuda",
         )
 
         attn_output = output
         temp_attn_output = q.new_empty(
-            (q.size(0), self.num_heads_tp * self.v_head_dim), dtype=q.dtype
+            (q.size(0), self.num_context_heads * self.v_head_dim), dtype=q.dtype
         )
 
         # use fake cached_cu_seq_len for chunked loop
@@ -1337,21 +1386,21 @@ def forward_context_with_chunked_prefill(
 
             # up proj to uncompressed kv
             # [tokens, 2, h, kv_dim], without rope_dim
-            chunked_kv = self.kv_b_proj(chunked_compressed_kv)
+            chunked_kv = self._project_context_kv(chunked_compressed_kv)
             chunked_k_nope, chunked_v = chunked_kv.split(
                 [
-                    self.num_heads_tp * self.qk_nope_head_dim,
-                    self.num_heads_tp * self.v_head_dim,
+                    self.num_context_heads * self.qk_nope_head_dim,
+                    self.num_context_heads * self.v_head_dim,
                 ],
                 -1,
             )
 
-            chunked_k_nope = chunked_k_nope.view(-1, self.num_heads_tp, self.qk_nope_head_dim)
+            chunked_k_nope = chunked_k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim)
             chunked_k_pe = chunked_k_pe.view(-1, 1, self.qk_rope_head_dim)
             chunked_k = maybe_compiled_cat(
-                (chunked_k_nope, chunked_k_pe.expand(-1, self.num_heads_tp, -1)), dim=-1
+                (chunked_k_nope, chunked_k_pe.expand(-1, self.num_context_heads, -1)), dim=-1
             )
-            chunked_k = chunked_k.view(-1, self.num_heads_tp * self.qk_head_dim)
+            chunked_k = chunked_k.view(-1, self.num_context_heads * self.qk_head_dim)
 
             # release pytorch activation memory
             chunked_compressed_kv = None
@@ -1394,7 +1443,7 @@ def forward_context_with_chunked_prefill(
             )
 
         # deal with the uncached kv
-        kv = self.kv_b_proj(compressed_kv)
+        kv = self._project_context_kv(compressed_kv)
         _, k_pe = latent_cache.view([-1, self.kv_lora_rank + self.qk_rope_head_dim]).split(
             [self.kv_lora_rank, self.qk_rope_head_dim], -1
         )
@@ -1402,16 +1451,16 @@ def forward_context_with_chunked_prefill(
 
         k_nope, v = kv.split(
             [
-                self.num_heads_tp * self.qk_nope_head_dim,
-                self.num_heads_tp * self.v_head_dim,
+                self.num_context_heads * self.qk_nope_head_dim,
+                self.num_context_heads * self.v_head_dim,
             ],
             -1,
         )
 
-        k_nope = k_nope.view(-1, self.num_heads_tp, self.qk_nope_head_dim)
+        k_nope = k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim)
         k_pe = k_pe.view(-1, 1, self.qk_rope_head_dim)
-        k = maybe_compiled_cat((k_nope, k_pe.expand(-1, self.num_heads_tp, -1)), dim=-1)
-        k = k.view(-1, self.num_heads_tp * self.qk_head_dim)
+        k = maybe_compiled_cat((k_nope, k_pe.expand(-1, self.num_context_heads, -1)), dim=-1)
+        k = k.view(-1, self.num_context_heads * self.qk_head_dim)
 
         # copy q_lens to replace kv_lens_runtime
         attn_metadata.kv_lens_runtime = attn_metadata.prompt_lens_cpu_runtime
@@ -1538,7 +1587,7 @@ def forward_context(
             # may stall inside inductor's compile worker.
             if trtllm_attention.has_cached_kv_for_mla_context_warmup(attn_metadata):
                 self.cached_warmup_forward_context_with_cached_kv(
-                    self.num_heads_tp,
+                    self.num_context_heads,
                     self.qk_nope_head_dim,
                     self.qk_rope_head_dim,
                     self.kv_lora_rank,
@@ -2213,6 +2262,20 @@ def transform_weights(self) -> None:
         self._weights_transformed = True
 
     def cache_derived_state(self) -> None:
+        if self.num_context_heads != self.num_heads_tp:
+            first_head = self.mapping.cp_rank * self.num_context_heads
+            last_head = first_head + self.num_context_heads
+            self.q_b_proj.prepare_output_slice(
+                first_head * self.qk_head_dim, last_head * self.qk_head_dim
+            )
+            self.kv_b_proj.prepare_output_slice(
+                first_head * self.qk_nope_head_dim, last_head * self.qk_nope_head_dim
+            )
+            value_start = self.num_heads_tp * self.qk_nope_head_dim
+            self.kv_b_proj.prepare_output_slice(
+                value_start + first_head * self.v_head_dim,
+                value_start + last_head * self.v_head_dim,
+            )
         self._weights_transformed = True
 
     def post_load_weights(self) -> None:
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index ec7aee2a0f05..539beed494f8 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -138,6 +138,7 @@
     from transformers import PretrainedConfig
 
     from ...llmapi.llm_args import DecodingBaseConfig
+    from ...quantization.utils.fp8_utils import Fp8BlockScaleWeightSliceCache
 
 # Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table
 # holds cos=1/sin=0), but the chunked-context path indexes the table by
@@ -835,6 +836,7 @@ def __init__(
         self.out_features = out_features
         self._weights_transformed = True
         self._fused_projection = None
+        self._output_slice_cache: Optional[Fp8BlockScaleWeightSliceCache] = None
         # Buffers (not parameters): these are the module's weights post-load;
         # there is nothing further to load into them and they must not be
         # touched by any later autocast/dtype move.
@@ -917,6 +919,7 @@ def from_linear(cls, linear: nn.Linear | TrtllmLinear) -> "_Fp8BlockScaleWeightR
 
     def transform_weights(self) -> None:
         if not self._weights_transformed:
+            self._output_slice_cache = None
             self.weight, self.weight_scale = self._prepare_weights(self.weight, self.weight_scale)
             if self._fused_projection is not None:
                 owner_ref, start = self._fused_projection
@@ -930,6 +933,22 @@ def transform_weights(self) -> None:
     def post_load_weights(self) -> None:
         self.transform_weights()
 
+    def forward_slice(self, x: torch.Tensor, start: int, end: int) -> torch.Tensor:
+        self.prepare_output_slice(start, end)
+        assert self._output_slice_cache is not None
+        weight, scale, crop = self._output_slice_cache.get(self.weight, start, end)
+        out = self._apply_with_weights(x, weight, scale)
+        return out[..., crop : crop + end - start].contiguous()
+
+    def prepare_output_slice(self, start: int, end: int) -> None:
+        """Prepare fixed output-slice scales after checkpoint transformation."""
+        from ...quantization.utils.fp8_utils import Fp8BlockScaleWeightSliceCache
+
+        self.transform_weights()
+        if self._output_slice_cache is None:
+            self._output_slice_cache = Fp8BlockScaleWeightSliceCache()
+        self._output_slice_cache.prepare(self.weight, self.weight_scale, start, end)
+
     def forward(
         self,
         x: torch.Tensor,
@@ -938,23 +957,28 @@ def forward(
         lora_params: Optional[dict] = None,
         layer_idx: Optional[int] = None,
     ) -> torch.Tensor:
-        out_shape = (*x.shape[:-1], self.out_features)
+        return self._apply_with_weights(x, self.weight, self.weight_scale)
+
+    def _apply_with_weights(
+        self, x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor
+    ) -> torch.Tensor:
+        out_shape = (*x.shape[:-1], weight.shape[0])
         x_flat = x.reshape(-1, x.shape[-1])
         if _rubin_mxfp8_is_available():
             act_fp8, act_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x_flat)
             out = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin(
                 act_fp8,
-                self.weight,
+                weight,
                 act_scale,
-                self.weight_scale,
+                weight_scale,
                 output_dtype=torch.bfloat16,
                 fine_grained_m=True,
             )
         else:
             out = torch.ops.trtllm.fp8_swap_ab_gemm(
                 x_flat,
-                self.weight,
-                self.weight_scale,
+                weight,
+                weight_scale,
                 output_dtype=x.dtype,
                 disable_ue8m0_cast=True,
             )
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 655ae90476cc..e4e713875c99 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -32,8 +32,8 @@
     preprocess_weights_for_mixed_gemm
 from tensorrt_llm.quantization.mode import QuantAlgo
 from tensorrt_llm.quantization.utils.fp8_utils import (
-    per_token_quant_and_transform, resmooth_to_fp8_e8m0,
-    transform_k128_scales_to_cutedsl_mxfp8_layout,
+    Fp8BlockScaleWeightSliceCache, per_token_quant_and_transform,
+    resmooth_to_fp8_e8m0, transform_k128_scales_to_cutedsl_mxfp8_layout,
     transform_sf_into_required_layout)
 
 from ..._utils import get_sm_version, is_sm_100f
@@ -438,6 +438,11 @@ def apply(self, module: Linear, input: torch.Tensor,
               bias: Optional[torch.Tensor], *args, **kwargs):
         raise NotImplementedError
 
+    def apply_slice(self, module: Linear, input: torch.Tensor, start: int,
+                    end: int) -> torch.Tensor:
+        raise NotImplementedError(
+            "This quantization method does not support output slicing")
+
     def apply_linear_allreduce(self, module: Linear, input: torch.Tensor,
                                bias: Optional[torch.Tensor], tp_rank: int,
                                tp_group: List[int], *args, **kwargs):
@@ -615,8 +620,12 @@ def create_weights(self, module: Linear, in_features: int,
         else:
             module.register_parameter("bias", None)
 
-    def apply(self, module: Linear, input: torch.Tensor,
-              bias: Optional[torch.Tensor]):
+    def apply(self,
+              module: Linear,
+              input: torch.Tensor,
+              bias: Optional[torch.Tensor],
+              weight: Optional[torch.Tensor] = None):
+        weight = module.weight if weight is None else weight
         # The opt-in low-M dispatcher routes to the built-in CuTe-DSL low-m GEMM
         # kernel and returns None to fall through to the normal GEMM path.
         # Skip when use_custom_cublas_mm is set: that path may allocate output
@@ -624,16 +633,16 @@ def apply(self, module: Linear, input: torch.Tensor,
         # the low-M dispatcher does not support.
         if _should_apply_low_m_gemm(input) and not getattr(
                 module, "use_custom_cublas_mm", False):
-            output = apply_low_m_gemm(module, input, module.weight, bias)
+            output = apply_low_m_gemm(module, input, weight, bias)
             if output is not None:
                 return output
         # CuTe DSL BF16 GEMM path for Blackwell / SM107
         if (module.use_cute_dsl_bf16_gemm and is_sm_100f()
-                and module.weight.dtype == torch.bfloat16):
+                and weight.dtype == torch.bfloat16):
             # input: [*, K], weight: [N, K], output: [*, N]
             input_2d = input.view(-1, input.shape[-1])  # [M, K]
             m, k = input_2d.shape
-            n = module.weight.shape[0]
+            n = weight.shape[0]
             output = torch.empty(m,
                                  n,
                                  dtype=torch.bfloat16,
@@ -643,7 +652,7 @@ def apply(self, module: Linear, input: torch.Tensor,
                             torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell)
             bf16_gemm_op(
                 input_2d.contiguous(),
-                module.weight,
+                weight,
                 output,
             )
             # Reshape output back to match input batch dims
@@ -662,15 +671,20 @@ def apply(self, module: Linear, input: torch.Tensor,
                      and module.mapping is not None else None)
             output = torch.ops.trtllm.cublas_mm(
                 input,
-                module.weight.t(),
+                weight.t(),
                 bias,
                 out_dtype=None,
                 output_buffer_kind=output_buffer_kind,
                 group=group)
         else:
-            output = F.linear(input, module.weight, bias)
+            output = F.linear(input, weight, bias)
         return output
 
+    def apply_slice(self, module: Linear, input: torch.Tensor, start: int,
+                    end: int) -> torch.Tensor:
+        bias = module.bias[start:end] if module.bias is not None else None
+        return self.apply(module, input, bias, weight=module.weight[start:end])
+
     def load_weights_vanilla(self,
                              module: Linear,
                              weights: List[Dict],
@@ -1304,6 +1318,12 @@ def create_weights(self, module: Linear, in_features: int,
 
     def apply(self, module: Linear, input: torch.Tensor,
               bias: Optional[torch.Tensor]):
+        return self._apply_with_weights(module, input, bias, module.weight,
+                                        module.weight_scale)
+
+    def _apply_with_weights(self, module: Linear, input: torch.Tensor,
+                            bias: Optional[torch.Tensor], weight: torch.Tensor,
+                            weight_scale: torch.Tensor):
         # fp8_block_scaling_gemm does not support writing into an NCCL window
         # buffer; supports_nccl_symmetric_memory_window_output is False so the window path is bypassed.
         if isinstance(input, tuple):
@@ -1319,15 +1339,14 @@ def apply(self, module: Linear, input: torch.Tensor,
                                         or module.disable_deep_gemm))
             if uses_cute_dsl_rubin:
                 output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin(
-                    activation, module.weight, activation_scale,
-                    module.weight_scale)
+                    activation, weight, activation_scale, weight_scale)
             elif (activation_scale.dtype == torch.int32 and is_sm_100f()
                   and not module.disable_deep_gemm):
                 output = torch.ops.trtllm.fp8_prequantized_swap_ab_gemm(
                     activation,
                     activation_scale,
-                    module.weight,
-                    module.weight_scale,
+                    weight,
+                    weight_scale,
                     disable_ue8m0_cast=True,
                 )
             else:
@@ -1357,15 +1376,13 @@ def apply(self, module: Linear, input: torch.Tensor,
                     act_input_fp8, act_input_sf = (
                         torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(input))
                     output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin(
-                        act_input_fp8, module.weight, act_input_sf,
-                        module.weight_scale)
+                        act_input_fp8, weight, act_input_sf, weight_scale)
                 else:
                     act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
                         input)
                     if sm_version in (100, 103):
                         output = torch.ops.trtllm.cute_dsl_fp8_gemm_blackwell(
-                            act_input_fp8, module.weight, act_input_sf,
-                            module.weight_scale)
+                            act_input_fp8, weight, act_input_sf, weight_scale)
                     else:
                         # cute_dsl_fp8_gemm_blackwell runs only on sm100/103. On
                         # other sm_100f GPUs (e.g. sm107 without the CuTe DSL
@@ -1379,24 +1396,23 @@ def apply(self, module: Linear, input: torch.Tensor,
                                 "trtllm-gen kernel instead.",
                                 key="cute_dsl_fp8_blockscale_unsupported_sm")
                         output = torch.ops.trtllm.fp8_block_scaling_gemm(
-                            act_input_fp8, module.weight, act_input_sf,
-                            module.weight_scale)
+                            act_input_fp8, weight, act_input_sf, weight_scale)
             else:
                 output = torch.ops.trtllm.fp8_swap_ab_gemm(
                     input,
-                    module.weight,
-                    module.weight_scale,
+                    weight,
+                    weight_scale,
                     disable_ue8m0_cast=True,
                 )
         elif sm_version == 120:
             act_input_fp8, act_input_sf = per_token_quant_and_transform(input)
             output = torch.ops.trtllm.fp8_block_scaling_gemm(
-                act_input_fp8, module.weight, act_input_sf, module.weight_scale)
+                act_input_fp8, weight, act_input_sf, weight_scale)
         else:
             act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
                 input)
             output = torch.ops.trtllm.fp8_block_scaling_gemm(
-                act_input_fp8, module.weight, act_input_sf, module.weight_scale)
+                act_input_fp8, weight, act_input_sf, weight_scale)
 
         # Reshape output back to original shape (with out_features as last dim)
         if len(original_shape) > 2:
@@ -1406,6 +1422,20 @@ def apply(self, module: Linear, input: torch.Tensor,
             output = output + bias
         return output
 
+    def apply_slice(self, module: Linear, input: torch.Tensor, start: int,
+                    end: int) -> torch.Tensor:
+        # Keep the original N-block quantization grid when a head boundary
+        # falls inside a 128-row scale block. The extra rows are cropped
+        # after GEMM; their cost is bounded by two scale blocks.
+        assert module._output_slice_cache is not None
+        weight, scale, crop = module._output_slice_cache.get(
+            module.weight, start, end)
+        output = self._apply_with_weights(module, input, None, weight, scale)
+        output = output[..., crop:crop + end - start].contiguous()
+        if module.bias is not None:
+            output = output + module.bias[start:end]
+        return output
+
     def _get_scale_name(self, weights: List[Dict]):
         # `weight_scale_inv` for DS recipe and `weight_scale` for ModelOpt recipe.
         # Actually they hold identical values of data_amax / 448.
@@ -4102,6 +4132,7 @@ def __init__(
         self._weights_created = False
         self._weights_transformed = False
         self.reduce_output = reduce_output
+        self._output_slice_cache: Optional[Fp8BlockScaleWeightSliceCache] = None
         self.use_custom_cublas_mm = use_custom_cublas_mm
         self.use_cute_dsl_bf16_gemm = use_cute_dsl_bf16_gemm
         self.lora = lora
@@ -4480,6 +4511,48 @@ def _maybe_fuse_bias_into_allreduce(
             assert all_reduce_params is None or all_reduce_params.enable_allreduce is False, "Cannot fuse norm/residual/bias ops into allreduce op since we do not call allreduce op when tp_size is 1."
             return False
 
+    def forward_slice(self, input: torch.Tensor, start: int,
+                      end: int) -> torch.Tensor:
+        """Project a contiguous range of this rank's output features.
+
+        The range is local to the existing COLUMN shard. This reuses its
+        loaded weights and quantization without materializing all outputs.
+
+        Args:
+            input: Activations with shape [..., in_features].
+            start: First output feature in the local weight shard.
+            end: Exclusive end of the local output feature range.
+
+        Returns:
+            A contiguous tensor with shape [..., end - start].
+        """
+        if self.tp_mode == TensorParallelMode.ROW or self.gather_output:
+            raise ValueError(
+                "Output slicing requires a local COLUMN projection")
+        if not 0 <= start < end <= self.weight.shape[0]:
+            raise ValueError(f"Invalid output feature range [{start}, {end})")
+        if self._locality_domain_weight_shards is not None or self.lora is not None:
+            raise NotImplementedError(
+                "Output slicing does not support locality-domain shards or LoRA"
+            )
+        if type(self.quant_method) not in (UnquantizedLinearMethod,
+                                           FP8BlockScalesLinearMethod):
+            raise NotImplementedError(
+                "Output slicing supports unquantized and FP8 block-scale projections"
+            )
+        self.prepare_output_slice(start, end)
+        return self.quant_method.apply_slice(self, input, start, end)
+
+    def prepare_output_slice(self, start: int, end: int) -> None:
+        """Prepare FP8 slice scales once, after weight transformation."""
+        if type(self.quant_method) is not FP8BlockScalesLinearMethod:
+            return
+        self.transform_weights()
+        if self._output_slice_cache is None:
+            self._output_slice_cache = Fp8BlockScaleWeightSliceCache()
+        self._output_slice_cache.prepare(self.weight, self.weight_scale, start,
+                                         end)
+
     def forward(
         self,
         input: Union[torch.Tensor, Fp4QuantizedTensor],
@@ -4554,6 +4627,7 @@ def load_weights(self,
                 f"{type(self.quant_method).__name__} does not support "
                 "allow_partial_loading")
         self._weights_transformed = False
+        self._output_slice_cache = None
         self.quant_method.load_weights(
             self,
             weights,
@@ -4566,6 +4640,7 @@ def process_weights_after_loading(self):
     def transform_weights(self) -> None:
         if self._weights_transformed:
             return
+        self._output_slice_cache = None
         self.quant_method.transform_weights(self)
         # Localized allocations and full-weight release belong to the mutating
         # stage so both staged ModelLoader walks and post_load_weights run them.
@@ -4593,6 +4668,7 @@ def post_load_weights(self) -> None:
         self.transform_weights()
 
     def pre_reload_weights(self):
+        self._output_slice_cache = None
         assert hasattr(
             self.quant_method, "pre_reload_weights"
         ), "pre_reload_weights is not supported for this quant method"
diff --git a/tensorrt_llm/quantization/utils/fp8_utils.py b/tensorrt_llm/quantization/utils/fp8_utils.py
index 26933a19b958..67d57cad7220 100644
--- a/tensorrt_llm/quantization/utils/fp8_utils.py
+++ b/tensorrt_llm/quantization/utils/fp8_utils.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from typing import Optional, Tuple
 
 import torch
@@ -199,6 +202,76 @@ def get_tma_aligned_size(x: int, element_size: int) -> int:
     return align(x, alignment)
 
 
+def slice_fp8_blockscale_weight(
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    start: int,
+    end: int,
+) -> Tuple[torch.Tensor, torch.Tensor, int]:
+    """Select output rows while retaining the original 128-row scale grid.
+
+    Returns the weight slice, its scales and the leading output rows to crop.
+    Supports raw FP32, DeepGEMM packed int32 and CuTe R128c4 uint8 scales.
+    """
+    if not 0 <= start < end <= weight.shape[0]:
+        raise ValueError(f"Invalid output feature range [{start}, {end})")
+    block_size = 128
+    first = start // block_size * block_size
+    last = min(ceil_div(end, block_size) * block_size, weight.shape[0])
+    if weight_scale.dtype == torch.float32:
+        scale = weight_scale[first // block_size:ceil_div(last, block_size)]
+    elif weight_scale.dtype == torch.int32:
+        # A row slice keeps the original leading dimension. Repack the small
+        # scale tensor to the column-major TMA layout for this GEMM's N.
+        rows = last - first
+        aligned_rows = get_tma_aligned_size(rows, weight_scale.element_size())
+        scale = weight_scale.new_empty(
+            (weight_scale.shape[1], aligned_rows)).t()[:rows]
+        scale.copy_(weight_scale[first:last])
+    elif weight_scale.dtype == torch.uint8:
+        # R128c4 stores each 128-row tile contiguously across padded K32
+        # columns. The rounded weight slice selects whole scale tiles.
+        if weight_scale.dim() != 1 or weight.shape[1] % 128:
+            raise ValueError(
+                "R128c4 scales require a flat tensor and K divisible by 128")
+        scale_cols = align(weight.shape[1] // 32, 4)
+        expected_size = align(weight.shape[0], 128) * scale_cols
+        if weight_scale.numel() != expected_size:
+            raise ValueError(f"Expected {expected_size} R128c4 scale bytes")
+        scale = weight_scale.narrow(0, first * scale_cols,
+                                    align(last - first, 128) * scale_cols)
+    else:
+        raise NotImplementedError(
+            "Sliced FP8 projections require FP32, packed int32 or R128c4 uint8 scales"
+        )
+    return weight[first:last], scale, start - first
+
+
+class Fp8BlockScaleWeightSliceCache(torch.nn.Module):
+    """Retain prepared output-slice scales as movable, nonpersistent buffers."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._ranges: dict[tuple[int, int], tuple[int, int, int]] = {}
+
+    def prepare(self, weight: torch.Tensor, weight_scale: torch.Tensor,
+                start: int, end: int) -> None:
+        """Prepare a fixed slice after the source weights are transformed."""
+        if (start, end) in self._ranges:
+            return
+        sliced_weight, scale, crop = slice_fp8_blockscale_weight(
+            weight, weight_scale, start, end)
+        first = start - crop
+        self.register_buffer(f"scale_{start}_{end}", scale, persistent=False)
+        self._ranges[start, end] = (first, first + sliced_weight.shape[0], crop)
+
+    def get(self, weight: torch.Tensor, start: int,
+            end: int) -> tuple[torch.Tensor, torch.Tensor, int]:
+        """Return a view of live weights and the prepared scale buffer."""
+        first, last, crop = self._ranges[start, end]
+        return weight[first:last], self.get_buffer(f"scale_{start}_{end}"), crop
+
+
 def get_col_major_tma_aligned_packed_tensor(x: torch.Tensor) -> torch.Tensor:
     # NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA
     assert x.dtype == torch.float and x.dim() in (2, 3)
diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml
index 3ed9f4085675..79360e6b1308 100644
--- a/tests/integration/test_lists/test-db/l0_cpu.yml
+++ b/tests/integration/test_lists/test-db/l0_cpu.yml
@@ -12,6 +12,7 @@ l0_cpu:
       backend: generic
       orchestrator: mpi
   tests:
+  - unittest/_torch/modules/test_linear_output_slice.py::test_fp8_weight_slice_cache
   - disaggregated/test_aiperf_gate.py::test_fires_on_error_storm
   - disaggregated/test_aiperf_gate.py::test_passes_healthy_run_with_cancellations
   - disaggregated/test_aiperf_gate.py::test_missing_export_raises
diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml
index 1908fd1db4ce..ce21a01c9ddc 100644
--- a/tests/integration/test_lists/test-db/l0_gb300.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300.yml
@@ -19,6 +19,8 @@ l0_gb300:
       backend: pytorch
   tests:
   # ------------- PyTorch tests ---------------
+  # MLA context head slices, including FP8 scale-block boundaries.
+  - unittest/_torch/modules/test_linear_output_slice.py
   # ------------- DeepSeek-V4.1 model, Engram and DSpark ---------------
   - unittest/_torch/modeling/test_modeling_deepseekv41.py
   - unittest/_torch/modules/test_engram_decode_hash.py
diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py
index 741b2511bbb2..f6842d9dc471 100644
--- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py
+++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py
@@ -377,13 +377,11 @@ def _run_mla_distributed(
         input_ctx, position_ids_ctx, scenario, rank, world_size
     )
 
-    # Context step — populate KV cache only; output is discarded.
-    # We call forward_impl directly instead of mla.forward() because
-    # o_proj is sized for helix-gen (num_heads_tp_cp * v_head_dim) and would
-    # crash on the full-width context attention output.
+    # This decode fixture populates each rank's cache from a different token
+    # chunk. Complete, replicated-input prefill outputs are checked separately.
     activate_all_ranks_for_context(attn_metadata, position_ids_ctx_rank)
     ctx_output = input_ctx_rank.new_empty(
-        [input_ctx_rank.shape[0], mla.num_heads_tp * mla.v_head_dim], dtype=input_ctx_rank.dtype
+        [input_ctx_rank.shape[0], mla.mha.num_heads * mla.v_head_dim], dtype=input_ctx_rank.dtype
     )
     mla.forward_impl(position_ids_ctx_rank, input_ctx_rank, attn_metadata, attn_output=[ctx_output])
 
diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py
new file mode 100644
index 000000000000..91de893ee466
--- /dev/null
+++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py
@@ -0,0 +1,278 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Check Helix's phase-dependent head split with complete Kimi K3 MLA outputs.
+
+These tests deliberately replicate the KV history to isolate projection and
+head-layout correctness. Owner-only persistence and distributed cached-context
+loading need separate tests; this fixture does not validate their handoff.
+"""
+
+import pickle
+import sys
+import weakref
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+import _torch.attention.multi_gpu.helix_test_utils as helix_utils
+import cloudpickle
+import pytest
+import torch
+from _torch.attention.multi_gpu.helix_test_utils import (
+    CACHE_TYPE_SELFKONLY,
+    run_helix_test,
+    setup_kv_and_metadata,
+)
+from _torch.modeling.kimi_k3_mla_reference import KimiMLAAttention
+from _torch.modeling.test_kimi_k3_checkpoint import load_mla_checkpoint
+from _torch.modeling.test_kimi_k3_mla import _assert_matches_reference, _checkpoint_weights
+from mpi4py import MPI
+from utils.util import skip_pre_blackwell
+
+from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures, KVCacheParams
+from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM
+from tensorrt_llm._torch.utils import model_extra_attrs
+from tensorrt_llm.functional import AllReduceStrategy
+from tensorrt_llm.mapping import CpType, Mapping
+from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
+
+cloudpickle.register_pickle_by_value(sys.modules[__name__])
+cloudpickle.register_pickle_by_value(helix_utils)
+MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL)
+
+
+@dataclass(frozen=True)
+class _Scenario:
+    tp_size: int = 1
+    cp_size: int = 2
+    num_heads: int = 32
+    fp8: bool = False
+    gate: bool = True
+
+
+@torch.no_grad()
+def _check_context_heads(rank, world_size, scenario, use_nccl_for_alltoall, fifo_version):
+    torch.manual_seed(42)
+    mapping = Mapping(
+        world_size=world_size,
+        rank=rank,
+        tp_size=scenario.tp_size,
+        cp_size=scenario.cp_size,
+        cp_config={
+            "cp_type": CpType.HELIX,
+            "use_nccl_for_alltoall": use_nccl_for_alltoall,
+            "fifo_version": fifo_version,
+        },
+    )
+    config = KimiLinearConfig(
+        vocab_size=128,
+        hidden_size=256,
+        intermediate_size=256,
+        num_hidden_layers=1,
+        attn_res_block_size=1,
+        torch_dtype=torch.bfloat16,
+        linear_attn_config={"kda_layers": [], "full_attn_layers": [1], "num_heads": 32},
+        num_attention_heads=scenario.num_heads,
+        num_key_value_heads=scenario.num_heads,
+        q_lora_rank=1536,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        rms_norm_eps=1e-6,
+        mla_use_nope=True,
+        mla_use_output_gate=scenario.gate,
+        max_position_embeddings=512,
+    )
+    config._attn_implementation = "eager"
+    reference = KimiMLAAttention(config, layer_idx=0).bfloat16().eval()
+    weights = _checkpoint_weights(reference, scenario.fp8)
+    model_config = ModelConfig(
+        pretrained_config=config,
+        mapping=mapping,
+        quant_config=QuantConfig(),
+        quant_config_dict={
+            f"model.layers.0.self_attn.{name}": QuantConfig(
+                quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128
+            )
+            for name in weights
+        }
+        if scenario.fp8
+        else None,
+        allreduce_strategy=AllReduceStrategy.NCCL,
+        use_cute_dsl_blockscaling_mm=True,
+        use_cute_dsl_blockscaling_bmm=True,
+    )
+    with torch.device("cuda"):
+        model = KimiLinearForCausalLM(model_config).eval()
+    projections = {
+        f"{name}.{suffix}": value
+        for name, pair in weights.items()
+        for suffix, value in pair.items()
+    }
+    projections.update(
+        {name: value for name, value in reference.named_parameters() if value.ndim == 1}
+    )
+    load_mla_checkpoint(model, projections, "")
+    runtime = model.model.layers[0].self_attn
+    for module in runtime.modules():
+        if callable(getattr(module, "post_load_weights", None)):
+            module.post_load_weights()
+    mixer = runtime.mixer
+    expected_heads = config.num_attention_heads // world_size
+    assert mixer.mha.num_heads == expected_heads, (
+        f"context backend has {mixer.mha.num_heads} heads, expected {expected_heads}"
+    )
+    assert mixer.mqa.num_heads == config.num_attention_heads // scenario.tp_size
+
+    reference.cuda()
+    batch, context_length = 2, 128
+    generator = torch.Generator(device="cuda").manual_seed(1234)
+    hidden = torch.randn(
+        batch,
+        context_length + 1,
+        config.hidden_size,
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    mask = torch.full((context_length + 1, context_length + 1), float("-inf"), device="cuda").triu(
+        1
+    )[None, None]
+    expected = reference(hidden, attention_mask=mask)
+    cache_scenario = SimpleNamespace(
+        ctx_len=context_length,
+        batch=batch,
+        kv_cache_tokens_per_block=64,
+        num_layers=1,
+        kv_cache_dtype=torch.bfloat16,
+    )
+    # A complete local cache on every rank intentionally isolates item 2 from
+    # owner-only writes (item 3) and remote cached-KV gathering (item 4).
+    manager, metadata = setup_kv_and_metadata(
+        cache_scenario,
+        Mapping(world_size=1, rank=0, tp_size=1),
+        cache_type=CACHE_TYPE_SELFKONLY,
+        num_kv_heads=1,
+        head_dim=config.kv_lora_rank + config.qk_rope_head_dim,
+    )
+    full_q_projection_rows = []
+    handle = mixer.q_b_proj.register_forward_hook(
+        lambda module, inputs, output: full_q_projection_rows.append(output.shape[0])
+    )
+    extra_attrs = {
+        "attention_metadata": weakref.ref(metadata),
+        "mla_layers": {mixer.layer_idx_str: weakref.ref(mixer)},
+    }
+    try:
+        context = hidden[:, :context_length].reshape(-1, config.hidden_size).contiguous()
+        with model_extra_attrs(extra_attrs):
+            actual = runtime(context, metadata).view(batch, context_length, -1)
+        _assert_matches_reference(actual, expected[:, :context_length], scenario.fp8)
+        if scenario.cp_size > 1:
+            assert full_q_projection_rows == [], (
+                "prefill must not run the full TP-local Q projection"
+            )
+
+        # Revisit the second half using a complete local prefix. This checks
+        # head layouts in both cached-context implementations without needing
+        # distributed KV gathering. Rewriting the same tokens is intentional.
+        half = context_length // 2
+        for chunked in (False, True):
+            print(f"Checking cached context: CP={scenario.cp_size}, chunked={chunked}", flush=True)
+            metadata = get_attention_backend("TRTLLM").Metadata(
+                seq_lens=torch.tensor([half] * batch, dtype=torch.int),
+                request_ids=[0, 1],
+                max_num_requests=batch,
+                num_contexts=batch,
+                # Context runtime lengths describe this query chunk. The
+                # cached lengths separately account for its 64-token prefix.
+                prompt_lens=[half] * batch,
+                max_num_tokens=half * batch,
+                kv_cache_manager=manager,
+                kv_cache_params=KVCacheParams(
+                    use_cache=True, num_cached_tokens_per_seq=[half] * batch
+                ),
+                runtime_features=AttentionRuntimeFeatures(
+                    chunked_prefill=chunked,
+                    chunk_size=half,
+                    chunked_prefill_buffer_batch_size=1,
+                ),
+                enable_context_mla_with_cached_kv=True,
+            )
+            metadata.prepare()
+            extra_attrs["attention_metadata"] = weakref.ref(metadata)
+            suffix = hidden[:, half:context_length].reshape(-1, config.hidden_size).contiguous()
+            with model_extra_attrs(extra_attrs):
+                actual = runtime(suffix, metadata).view(batch, half, -1)
+            _assert_matches_reference(actual, expected[:, half:context_length], scenario.fp8)
+        if scenario.cp_size > 1:
+            assert full_q_projection_rows == [], "cached prefill must retain the CP head split"
+
+        metadata = get_attention_backend("TRTLLM").Metadata(
+            seq_lens=torch.tensor([context_length, 1], dtype=torch.int),
+            request_ids=[0, 1],
+            max_num_requests=batch,
+            num_contexts=1,
+            prompt_lens=[context_length, context_length],
+            max_num_tokens=context_length + 1,
+            kv_cache_manager=manager,
+            kv_cache_params=KVCacheParams(
+                use_cache=True, num_cached_tokens_per_seq=[0, context_length]
+            ),
+            mapping=mapping,
+            enable_context_mla_with_cached_kv=True,
+        )
+        # The generation kernel emits CP softmax statistics only when its
+        # Helix fields are present. Every rank appends to its full replica in
+        # this item 2 fixture; ownership is tested separately in item 3.
+        metadata.update_helix_param([context_length], [False])
+        metadata.prepare()
+        extra_attrs["attention_metadata"] = weakref.ref(metadata)
+        mixed = torch.cat((hidden[0, :context_length], hidden[1, -1:]), dim=0)
+        with model_extra_attrs(extra_attrs):
+            actual = runtime(mixed, metadata)
+        _assert_matches_reference(
+            actual[:context_length], expected[0, :context_length], scenario.fp8
+        )
+        _assert_matches_reference(actual[-1], expected[1, -1], scenario.fp8)
+        if scenario.cp_size > 1:
+            assert full_q_projection_rows == [1], (
+                "only decode tokens use the full TP-local Q projection"
+            )
+        return 0.0
+    finally:
+        handle.remove()
+        manager.shutdown()
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs")
+@skip_pre_blackwell
+@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"])
+@pytest.mark.parametrize("gate", [False, True], ids=["no_gate", "gate"])
+@pytest.mark.parametrize(
+    "tp_size,cp_size,num_heads",
+    [
+        pytest.param(1, 1, 32, id="tp1-cp1"),
+        pytest.param(1, 2, 32, id="tp1-cp2"),
+        pytest.param(
+            2,
+            2,
+            32,
+            id="tp2-cp2",
+            marks=pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires four GPUs"),
+        ),
+        pytest.param(1, 2, 6, id="tp1-cp2-unaligned-q-scale-block"),
+    ],
+)
+def test_k3_mla_helix_context_heads(fp8, gate, tp_size, cp_size, num_heads):
+    scenario = _Scenario(tp_size=tp_size, cp_size=cp_size, num_heads=num_heads, fp8=fp8, gate=gate)
+    run_helix_test(
+        _check_context_heads,
+        scenario,
+        "nccl",
+        world_size=scenario.tp_size * scenario.cp_size,
+        max_mismatch_ratio=0.0,
+    )
diff --git a/tests/unittest/_torch/modules/test_linear_output_slice.py b/tests/unittest/_torch/modules/test_linear_output_slice.py
new file mode 100644
index 000000000000..9abb57243abe
--- /dev/null
+++ b/tests/unittest/_torch/modules/test_linear_output_slice.py
@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Output slices preserve loaded weights and FP8 scale-block boundaries."""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
+from tensorrt_llm.quantization.utils.fp8_utils import Fp8BlockScaleWeightSliceCache
+
+
+def _make_linear(backend, bias):
+    torch.manual_seed(42)
+    in_features, out_features = 512, 1152
+    quantized = backend != "bf16"
+    with torch.device("cuda"):
+        linear = Linear(
+            in_features,
+            out_features,
+            bias=bias,
+            dtype=torch.bfloat16,
+            tensor_parallel_mode=TensorParallelMode.COLUMN,
+            quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128)
+            if quantized
+            else QuantConfig(),
+            use_cute_dsl_blockscaling_mm=backend == "fp8_cute",
+        )
+    weight = torch.randn(out_features, in_features, device="cuda") * 0.02
+    checkpoint = {}
+    if quantized:
+        blocks = weight.view(out_features // 128, 128, in_features // 128, 128)
+        scale = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-8) / 448.0
+        expanded = scale.repeat_interleave(128, 0).repeat_interleave(128, 1)
+        checkpoint["weight"] = (weight / expanded).to(torch.float8_e4m3fn)
+        checkpoint["weight_scale"] = scale
+    else:
+        checkpoint["weight"] = weight.bfloat16()
+    if bias:
+        checkpoint["bias"] = torch.randn(out_features, dtype=torch.bfloat16, device="cuda")
+    linear.load_weights([checkpoint])
+    linear.post_load_weights()
+    return linear, checkpoint
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or get_sm_version() < 100, reason="requires Blackwell or later"
+)
+@pytest.mark.parametrize("backend", ["bf16", "fp8_cute", "fp8_deepgemm"])
+@pytest.mark.parametrize("bias", [False, True])
+@pytest.mark.parametrize("start,end", [(0, 576), (576, 1152), (513, 1107)])
+@torch.inference_mode()
+def test_linear_output_slice(backend, bias, start, end):
+    if get_sm_version() == 107 and backend == "fp8_deepgemm":
+        pytest.skip("SM107 native FP8 slices are covered by fp8_cute")
+    linear, _ = _make_linear(backend, bias)
+    quantized = backend != "bf16"
+    original_weight = linear.weight.clone()
+    original_scale = linear.weight_scale.clone() if quantized else None
+    inputs = torch.randn(31, linear.in_features, dtype=torch.bfloat16, device="cuda")
+    expected = linear(inputs)[:, start:end]
+    linear.prepare_output_slice(start, end)
+    cached_scale = (
+        linear._output_slice_cache.get(linear.weight, start, end)[1] if quantized else None
+    )
+    actual = linear.forward_slice(inputs, start, end)
+    assert actual.is_contiguous()
+    assert torch.isfinite(actual).all()
+    torch.testing.assert_close(actual, expected, atol=0.005, rtol=0.03)
+    torch.testing.assert_close(
+        linear.weight.view(torch.uint8), original_weight.view(torch.uint8), atol=0, rtol=0
+    )
+    if quantized:
+        torch.testing.assert_close(linear.weight_scale, original_scale, atol=0, rtol=0)
+        assert linear._output_slice_cache.get(linear.weight, start, end)[1] is cached_scale
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or get_sm_version() < 100, reason="requires Blackwell or later"
+)
+@pytest.mark.parametrize("start,end", [(0, 576), (576, 1152), (513, 1107)])
+@torch.inference_mode()
+def test_kimi_fp8_output_slice(start, end):
+    from tensorrt_llm._torch.models.modeling_kimi_linear import _Fp8BlockScaleWeightReadLinear
+
+    backend = "fp8_cute" if get_sm_version() == 107 else "fp8_deepgemm"
+    linear, _ = _make_linear(backend, False)
+    projection = _Fp8BlockScaleWeightReadLinear(
+        linear.weight.detach(), linear.weight_scale.detach(), linear.out_features
+    )
+    inputs = torch.randn(2, 17, linear.in_features, dtype=torch.bfloat16, device="cuda")
+    expected = projection(inputs)[..., start:end]
+    projection.prepare_output_slice(start, end)
+    cached_scale = projection._output_slice_cache.get(projection.weight, start, end)[1]
+    actual = projection.forward_slice(inputs, start, end)
+    torch.testing.assert_close(actual, expected, atol=0.005, rtol=0.03)
+    assert actual.is_contiguous()
+    assert projection._output_slice_cache.get(projection.weight, start, end)[1] is cached_scale
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or get_sm_version() < 100, reason="requires Blackwell or later"
+)
+@pytest.mark.parametrize("backend", ["fp8_cute", "fp8_deepgemm"])
+@torch.inference_mode()
+def test_linear_output_slice_reload(backend):
+    if get_sm_version() == 107 and backend == "fp8_deepgemm":
+        pytest.skip("SM107 native FP8 slices are covered by fp8_cute")
+    linear, checkpoint = _make_linear(backend, False)
+    inputs = torch.randn(7, linear.in_features, dtype=torch.bfloat16, device="cuda")
+    linear.prepare_output_slice(513, 1107)
+    old_cache = linear._output_slice_cache
+    linear.pre_reload_weights()
+    assert linear._output_slice_cache is None
+    reloaded = {name: value.clone() for name, value in checkpoint.items()}
+    reloaded["weight_scale"].mul_(2)
+    linear.load_weights([reloaded])
+    linear.post_load_weights()
+    actual = linear.forward_slice(inputs, 513, 1107)
+    torch.testing.assert_close(actual, linear(inputs)[:, 513:1107], atol=0.005, rtol=0.03)
+    assert linear._output_slice_cache is not old_cache
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("layout", ["fp32", "packed", "r128c4"])
+@pytest.mark.parametrize("start,end", [(0, 576), (576, 1152), (513, 1107)])
+@pytest.mark.parametrize("rows", [1152, 1107])
+def test_fp8_weight_slice_cache(layout, start, end, rows):
+    cols = 512
+    end = min(end, rows)
+    weight = torch.empty(rows, cols, dtype=torch.float8_e4m3fn)
+    if layout == "fp32":
+        scales = torch.arange((rows + 127) // 128 * (cols // 128), dtype=torch.float32).view(
+            -1, cols // 128
+        )
+    elif layout == "packed":
+        scales = torch.arange(rows * 2, dtype=torch.int32).view(2, rows).t()
+    else:
+        scales = torch.randint(
+            0,
+            256,
+            (rows, cols // 32),
+            dtype=torch.uint8,
+            generator=torch.Generator().manual_seed(42),
+        )
+        scales = torch.ops.trtllm.block_scale_interleave(scales)
+    cache = Fp8BlockScaleWeightSliceCache()
+    cache.prepare(weight, scales, start, end)
+    sliced, prepared_scale, crop = cache.get(weight, start, end)
+    first, last = start // 128 * 128, min((end + 127) // 128 * 128, rows)
+    assert sliced.data_ptr() == weight[first:last].data_ptr()
+    assert crop == start - first
+    if layout == "fp32":
+        expected = scales[first // 128 : (last + 127) // 128]
+    elif layout == "packed":
+        expected = scales[first:last]
+    else:
+        padded_rows = (rows + 127) // 128 * 128
+        unswizzled = torch.ops.trtllm.block_scale_interleave_reverse(scales.view(padded_rows, -1))
+        expected = torch.ops.trtllm.block_scale_interleave(unswizzled[first:last].contiguous())
+    torch.testing.assert_close(prepared_scale, expected, atol=0, rtol=0)
+    cache.prepare(weight, scales, start, end)
+    assert cache.get(weight, start, end)[1] is prepared_scale
+    assert not cache.state_dict()
+    cache.to("meta")
+    assert cache.get(weight, start, end)[1].is_meta

From 8972f3ddf9aafd32e3e7ec26d5d2bd9ea7f8ed2d Mon Sep 17 00:00:00 2001
From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
Date: Fri, 23 Oct 2026 09:05:28 -0700
Subject: [PATCH 2/6] [TRTLLM-17121][fix] Update partial MLA test fixtures

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
---
 tests/unittest/_torch/attention/test_attention_mla.py  | 3 +++
 tests/unittest/_torch/executor/test_model_loader_mx.py | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py
index 83133dc0c924..0934be31eaee 100644
--- a/tests/unittest/_torch/attention/test_attention_mla.py
+++ b/tests/unittest/_torch/attention/test_attention_mla.py
@@ -1466,6 +1466,7 @@ def _context_projection(device: str,
                         bias: bool = False) -> types.SimpleNamespace:
     model = types.SimpleNamespace(
         num_heads_tp=4,
+        num_context_heads=4,
         qk_nope_head_dim=16,
         qk_rope_head_dim=8,
         qk_head_dim=24,
@@ -1483,6 +1484,8 @@ def _context_projection(device: str,
         mla_module.MLA._use_kvbproj_strided_out, model)
     model._get_kvbproj_head_weights = types.MethodType(
         mla_module.MLA._get_kvbproj_head_weights, model)
+    model._project_context_kv = types.MethodType(
+        mla_module.MLA._project_context_kv, model)
     return model
 
 
diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py
index c00cf02efd62..2d276f47cf91 100644
--- a/tests/unittest/_torch/executor/test_model_loader_mx.py
+++ b/tests/unittest/_torch/executor/test_model_loader_mx.py
@@ -2562,6 +2562,8 @@ def test_mla_transform_weights_is_idempotent(monkeypatch):
     monkeypatch.setattr(mla_mod, "get_sm_version", lambda: 120)
     quant_mode = SimpleNamespace(has_fp8_block_scales=lambda: True)
     mla = MLA.__new__(MLA)
+    mla.num_heads_tp = 1
+    mla.num_context_heads = 1
     mla._weights_transformed = False
     mla.kv_b_proj = SimpleNamespace(quant_config=SimpleNamespace(quant_mode=quant_mode))
     mla.k_b_proj_trans = "k_weight"

From f7dbcdce2328a6c4b844cdee70894bea66adda6e Mon Sep 17 00:00:00 2001
From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
Date: Fri, 23 Oct 2026 10:33:10 -0700
Subject: [PATCH 3/6] [TRTLLM-17121][fix] Reject unsupported dense Helix
 attention backends

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention/mla.py          |  5 ++
 .../_torch/attention/test_attention_mla.py    | 55 +++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py
index 882fb780bb6c..195709be2327 100644
--- a/tensorrt_llm/_torch/attention/mla.py
+++ b/tensorrt_llm/_torch/attention/mla.py
@@ -423,6 +423,11 @@ def __init__(
             if self.mapping.has_cp_helix() and self.sparse_attn_hooks is None
             else self.num_heads_tp
         )
+        if self.num_context_heads != self.num_heads_tp and config.attn_backend.upper() != "TRTLLM":
+            raise NotImplementedError(
+                "Dense MLA with Helix context parallelism requires the TRTLLM attention backend; "
+                f"got {config.attn_backend}."
+            )
         self.num_key_value_heads_tp = (self.num_key_value_heads + tp_size - 1) // tp_size
 
         if rms_norm_eps is None:
diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py
index 0934be31eaee..e6057f3dff7f 100644
--- a/tests/unittest/_torch/attention/test_attention_mla.py
+++ b/tests/unittest/_torch/attention/test_attention_mla.py
@@ -1461,6 +1461,61 @@ def test_context_prefix_counts_tokens_not_rows():
     assert state.mla_prepare_ctx_cu_seqlens().tolist() == [0, 2, 5, 12]
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("mapping_override", [False, True])
+@pytest.mark.parametrize(
+    "backend,cp_size,sparse,rejected",
+    [
+        ("FLASHINFER", 2, False, True),
+        ("VANILLA", 2, False, True),
+        ("TRTLLM", 2, False, False),
+        ("trtllm", 2, False, False),
+        ("FLASHINFER", 1, False, False),
+        ("VANILLA", 1, False, False),
+        ("FLASHINFER", 2, True, False),
+    ],
+)
+def test_mla_helix_backend_validation(monkeypatch, mapping_override, backend,
+                                      cp_size, sparse, rejected):
+    """Reject unsupported dense Helix backends before allocating projections."""
+    mapping = Mapping(world_size=cp_size,
+                      cp_size=cp_size,
+                      cp_config={"cp_type": "HELIX"})
+    config = mla_module.ModelConfig(
+        mapping=Mapping() if mapping_override else mapping,
+        attn_backend=backend,
+    )
+    hooks = types.SimpleNamespace() if sparse else None
+    monkeypatch.setattr(mla_module, "get_sparse_mla_hooks", lambda _: hooks)
+    # Stop accepted configurations before projection/backend allocation.
+    quant_config = mock.Mock(side_effect=RuntimeError("quantization setup"))
+    monkeypatch.setattr(mla_module.ModelConfig, "get_quant_config",
+                        quant_config)
+    expected_error = NotImplementedError if rejected else RuntimeError
+    expected_message = "Dense MLA.*TRTLLM" if rejected else "quantization setup"
+    with pytest.raises(expected_error, match=expected_message):
+        mla_module.MLA(
+            hidden_size=128,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=8,
+            v_head_dim=16,
+            q_lora_rank=32,
+            kv_lora_rank=32,
+            predicted_tokens_per_seq=1,
+            max_position_embeddings=128,
+            bias=False,
+            pos_embd_params=PositionalEmbeddingParams(
+                type=PositionEmbeddingType.rope_gpt_neox,
+                rope=RopeParams(dim=8, max_positions=128),
+            ),
+            config=config,
+            mapping_with_cp=mapping if mapping_override else None,
+        )
+    assert quant_config.call_count == (0 if rejected else 1)
+
+
 def _context_projection(device: str,
                         dtype: torch.dtype,
                         bias: bool = False) -> types.SimpleNamespace:

From f7ef219f2096b44ce0829215a262cb664b3965f9 Mon Sep 17 00:00:00 2001
From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
Date: Fri, 23 Oct 2026 11:56:03 -0700
Subject: [PATCH 4/6] [TRTLLM-17121][fix] Validate Helix slices during load and
 preserve GMS transforms

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/linear.py         | 10 +++----
 .../_torch/pyexecutor/model_loader.py         |  7 +++--
 .../integration/test_lists/test-db/l0_cpu.yml |  1 +
 .../_torch/executor/test_model_loader_gms.py  | 29 +++++++++++++++++++
 .../modules/test_linear_output_slice.py       | 24 ++++++++++++++-
 5 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index e4e713875c99..04a3a28256e9 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -4526,6 +4526,11 @@ def forward_slice(self, input: torch.Tensor, start: int,
         Returns:
             A contiguous tensor with shape [..., end - start].
         """
+        self.prepare_output_slice(start, end)
+        return self.quant_method.apply_slice(self, input, start, end)
+
+    def prepare_output_slice(self, start: int, end: int) -> None:
+        """Validate output slicing and prepare FP8 scales after weight transformation."""
         if self.tp_mode == TensorParallelMode.ROW or self.gather_output:
             raise ValueError(
                 "Output slicing requires a local COLUMN projection")
@@ -4540,11 +4545,6 @@ def forward_slice(self, input: torch.Tensor, start: int,
             raise NotImplementedError(
                 "Output slicing supports unquantized and FP8 block-scale projections"
             )
-        self.prepare_output_slice(start, end)
-        return self.quant_method.apply_slice(self, input, start, end)
-
-    def prepare_output_slice(self, start: int, end: int) -> None:
-        """Prepare FP8 slice scales once, after weight transformation."""
         if type(self.quant_method) is not FP8BlockScalesLinearMethod:
             return
         self.transform_weights()
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 67a3fbba4a6f..6212c445c5bd 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -1476,9 +1476,9 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                         #      compatibility check (GMS has no disk fallback).
                         #   4. `materialize_module`: zero-copy bind GMS
                         #      pool storage onto the model parameters.
-                        #   5. Per-module `cache_derived_state`: recompute
-                        #      Python-side state from real, materialized
-                        #      tensors without re-running one-shot transforms.
+                        #   5. Mark materialized weights as transformed, then
+                        #      recompute Python-side derived state without
+                        #      re-running one-shot transforms.
                         #   6. `post_load_publish`: apply the format-specific
                         #      receiver publication policy.
                         with timing_metric(
@@ -1496,6 +1496,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                             self._check_gms_source_identity(gms_backend)
 
                             gms_backend.materialize_module(model)
+                            self._mark_weights_transformed(model)
                             self._walk_cache_state(model)
 
                             self._post_load_publish(
diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml
index 79360e6b1308..e1a43fe4d654 100644
--- a/tests/integration/test_lists/test-db/l0_cpu.yml
+++ b/tests/integration/test_lists/test-db/l0_cpu.yml
@@ -13,6 +13,7 @@ l0_cpu:
       orchestrator: mpi
   tests:
   - unittest/_torch/modules/test_linear_output_slice.py::test_fp8_weight_slice_cache
+  - unittest/_torch/modules/test_linear_output_slice.py::test_prepare_output_slice_rejects_unsupported_projection_before_forward
   - disaggregated/test_aiperf_gate.py::test_fires_on_error_storm
   - disaggregated/test_aiperf_gate.py::test_passes_healthy_run_with_cancellations
   - disaggregated/test_aiperf_gate.py::test_missing_export_raises
diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py
index 5e16c03a904c..c3c4e1926dd2 100644
--- a/tests/unittest/_torch/executor/test_model_loader_gms.py
+++ b/tests/unittest/_torch/executor/test_model_loader_gms.py
@@ -369,6 +369,35 @@ def _append(*_args, **_kwargs):
     backend.materialize_module.assert_called_once()
 
 
+def test_gms_ro_does_not_retransform_materialized_slice_weights(monkeypatch):
+    events = []
+    loader = _make_loader(monkeypatch, events=events)
+    backend = _build_gms_backend(is_rw=False, events=events)
+    _install_gms_backend(monkeypatch, backend)
+
+    class _AlreadyTransformedProjection(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self._weights_transformed = False
+            self.cache_rebuilt = False
+
+        def cache_derived_state(self):
+            if not self._weights_transformed:
+                raise AssertionError("GMS read-only weights must not be transformed again")
+            self.cache_rebuilt = True
+
+    model = model_loader_mod.AutoModelForCausalLM.from_config.return_value
+    model.projection = _AlreadyTransformedProjection()
+    checkpoint_loader = MagicMock(name="checkpoint_loader")
+    checkpoint_loader.checkpoint_format = "HF"
+
+    loaded_model, _ = loader.load("/ckpt", checkpoint_loader)
+
+    assert loaded_model is model
+    assert model.projection.cache_rebuilt
+    assert model.projection._weights_transformed
+
+
 def test_gms_rw_post_load_runs_inside_pool_before_finalize(monkeypatch):
     """Every step that may allocate or rebind tensors must run inside the GMS pool.
 
diff --git a/tests/unittest/_torch/modules/test_linear_output_slice.py b/tests/unittest/_torch/modules/test_linear_output_slice.py
index 9abb57243abe..adf6680511a9 100644
--- a/tests/unittest/_torch/modules/test_linear_output_slice.py
+++ b/tests/unittest/_torch/modules/test_linear_output_slice.py
@@ -5,7 +5,7 @@
 import pytest
 import torch
 
-from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
+from tensorrt_llm._torch.modules.linear import Linear, NVFP4LinearMethod, TensorParallelMode
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
 from tensorrt_llm.quantization.utils.fp8_utils import Fp8BlockScaleWeightSliceCache
@@ -122,6 +122,28 @@ def test_linear_output_slice_reload(backend):
     assert linear._output_slice_cache is not old_cache
 
 
+@pytest.mark.cpu_only
+def test_prepare_output_slice_rejects_unsupported_projection_before_forward():
+    linear = Linear(
+        8,
+        16,
+        bias=False,
+        dtype=torch.bfloat16,
+        tensor_parallel_mode=TensorParallelMode.COLUMN,
+        quant_config=QuantConfig(),
+        reduce_output=False,
+    )
+
+    linear.lora = object()
+    with pytest.raises(NotImplementedError, match="LoRA"):
+        linear.prepare_output_slice(0, 8)
+
+    linear.lora = None
+    linear.quant_method = NVFP4LinearMethod()
+    with pytest.raises(NotImplementedError, match="FP8 block-scale"):
+        linear.prepare_output_slice(0, 8)
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("layout", ["fp32", "packed", "r128c4"])
 @pytest.mark.parametrize("start,end", [(0, 576), (576, 1152), (513, 1107)])

From 45fe0b70588432ac5518b04578b70b070ab312af Mon Sep 17 00:00:00 2001
From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
Date: Fri, 23 Oct 2026 12:57:50 -0700
Subject: [PATCH 5/6] [TRTLLM-17121][test] Focus Helix prefill regression
 matrix

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
---
 .../multi_gpu/test_mla_helix_context.py         | 15 ++++++++-------
 .../_torch/attention/test_attention_mla.py      | 17 ++++++++---------
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py
index 91de893ee466..54ec875f4951 100644
--- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py
+++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py
@@ -250,21 +250,22 @@ def _check_context_heads(rank, world_size, scenario, use_nccl_for_alltoall, fifo
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs")
 @skip_pre_blackwell
-@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"])
-@pytest.mark.parametrize("gate", [False, True], ids=["no_gate", "gate"])
 @pytest.mark.parametrize(
-    "tp_size,cp_size,num_heads",
+    "fp8,gate,tp_size,cp_size,num_heads",
     [
-        pytest.param(1, 1, 32, id="tp1-cp1"),
-        pytest.param(1, 2, 32, id="tp1-cp2"),
+        pytest.param(False, True, 1, 1, 32, id="bf16-control-tp1-cp1"),
+        pytest.param(False, False, 1, 2, 32, id="bf16-no-gate-tp1-cp2"),
+        pytest.param(True, True, 1, 2, 32, id="fp8-gate-tp1-cp2"),
         pytest.param(
+            True,
+            True,
             2,
             2,
             32,
-            id="tp2-cp2",
+            id="fp8-gate-tp2-cp2",
             marks=pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires four GPUs"),
         ),
-        pytest.param(1, 2, 6, id="tp1-cp2-unaligned-q-scale-block"),
+        pytest.param(True, False, 1, 2, 6, id="fp8-no-gate-tp1-cp2-unaligned-q-scale-block"),
     ],
 )
 def test_k3_mla_helix_context_heads(fp8, gate, tp_size, cp_size, num_heads):
diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py
index e6057f3dff7f..9c3c83e284d7 100644
--- a/tests/unittest/_torch/attention/test_attention_mla.py
+++ b/tests/unittest/_torch/attention/test_attention_mla.py
@@ -1462,17 +1462,16 @@ def test_context_prefix_counts_tokens_not_rows():
 
 
 @pytest.mark.cpu_only
-@pytest.mark.parametrize("mapping_override", [False, True])
 @pytest.mark.parametrize(
-    "backend,cp_size,sparse,rejected",
+    "mapping_override,backend,cp_size,sparse,rejected",
     [
-        ("FLASHINFER", 2, False, True),
-        ("VANILLA", 2, False, True),
-        ("TRTLLM", 2, False, False),
-        ("trtllm", 2, False, False),
-        ("FLASHINFER", 1, False, False),
-        ("VANILLA", 1, False, False),
-        ("FLASHINFER", 2, True, False),
+        (False, "FLASHINFER", 2, False, True),
+        (False, "VANILLA", 2, False, True),
+        (False, "TRTLLM", 2, False, False),
+        (True, "FLASHINFER", 2, False, True),
+        (True, "TRTLLM", 2, False, False),
+        (False, "FLASHINFER", 1, False, False),
+        (False, "FLASHINFER", 2, True, False),
     ],
 )
 def test_mla_helix_backend_validation(monkeypatch, mapping_override, backend,

From 306f5f589d6a377d8cf51e8fbedafb9081e43f61 Mon Sep 17 00:00:00 2001
From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
Date: Sat, 24 Oct 2026 10:55:26 -0700
Subject: [PATCH 6/6] [TRTLLM-17121][test] Initialize context heads in MLA
 locality fixture

Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com>
---
 tests/unittest/_torch/attention/test_mla_locality_domain.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unittest/_torch/attention/test_mla_locality_domain.py b/tests/unittest/_torch/attention/test_mla_locality_domain.py
index 5fefc8ff40e4..6bceb277357e 100644
--- a/tests/unittest/_torch/attention/test_mla_locality_domain.py
+++ b/tests/unittest/_torch/attention/test_mla_locality_domain.py
@@ -81,6 +81,7 @@ def _make_bare_absorption_mla(
     )
     mla.num_heads_tp = 2
     mla.num_heads_tp_cp = 2
+    mla.num_context_heads = mla.num_heads_tp
     mla.qk_nope_head_dim = qk_nope_head_dim
     mla.kv_lora_rank = 512
     mla.v_head_dim = 128