From acefe293aa02904a0301b2c640e677cf8530c03c Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Tue, 24 Mar 2026 00:32:01 -0700
Subject: [PATCH 01/14] [None][perf] Split MLA DSA custom op for piecewise CUDA
 graph capture
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Split the monolithic mla_custom_op_inplace into two ops for DSA models:
- mla_dsa_proj (Op 1): Token-wise projections (cublas_mm, rope, FP8
  quantize, weight scaling). CUDA-graph-capturable — no batch metadata
  access, no tensor slicing by num_tokens.
- mla_dsa_attn_inplace (Op 2): Batch-dependent k cache update,
  sparse_attn_indexer, and context/generation attention dispatch.
  Excluded from CUDA graph capture.

This enables the piecewise CUDA graph optimizer to capture the
compute-heavy projection portion of DSA MLA while keeping the
batch-structure-dependent attention dispatch outside the graph.

Key design decisions:
- Indexer split into pre_indexer_proj (graph-safe) and _update_k_cache
  (moved to Op 2) to avoid capturing metadata-dependent scatter ops.
- All num_tokens slicing deferred to Op 2 so graph capture sees
  fixed-shape tensors.
- Indexer intermediates (q_fp8, k_fp8, k_scale, weights) returned from
  Op 1 as List[Tensor] and passed explicitly to Op 2 — no stashing on
  self to avoid CUDA graph address aliasing.
- _should_use_short_mha disabled under torch compile for straight-line
  control flow in Op 1.
- Non-DSA MLA unchanged (still uses mla_custom_op_inplace).

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  61 +++--
 .../_torch/compilation/piecewise_optimizer.py |   7 +-
 tensorrt_llm/_torch/compilation/utils.py      |   3 +
 tensorrt_llm/_torch/modules/attention.py      | 230 +++++++++++++++---
 4 files changed, 257 insertions(+), 44 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 5fdbae704172..2a9573bf29dc 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1670,14 +1670,43 @@ def _prep_q_or_k(self, qk_pe: torch.Tensor, qk_nope: torch.Tensor):
         return fp8_out, scale
 
     @torch.inference_mode()
-    def forward(self, qr: torch.Tensor, hidden_states: torch.Tensor,
-                metadata: DSAtrtllmAttentionMetadata,
-                position_ids: torch.Tensor):
-        quant_block_size = metadata.kv_cache_manager.quant_block_size
-        assert quant_block_size == 128, "Only support quant_block_size = 128 for now"
+    def pre_indexer(
+        self, qr: torch.Tensor, hidden_states: torch.Tensor,
+        metadata: DSAtrtllmAttentionMetadata, position_ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Token-wise projections, FP8 quantize, weight scaling, and k cache update.
+
+        Runs the full indexer pre-computation including k cache update.
+        Used by the eager path (Indexer.forward) where everything runs
+        outside CUDA graph capture.
+
+        Returns (q_fp8, k_fp8, k_scale, weights).
+        """
+        q_fp8, k_fp8, k_scale, weights = self.pre_indexer_proj(
+            qr, hidden_states, position_ids)
+
+        weights, _ = maybe_execute_in_parallel(
+            lambda: weights,
+            lambda: self._update_k_cache(k_fp8, k_scale, metadata),
+            self.ln_events[0],
+            self.ln_events[1],
+            self.aux_stream,
+        )
+
+        return q_fp8, k_fp8, k_scale, weights
+
+    def pre_indexer_proj(
+        self, qr: torch.Tensor, hidden_states: torch.Tensor,
+        position_ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Pure token-wise projections (CUDA-graph-capturable).
 
-        # Single fused FP32 cuBLAS GEMM (TF32 on Ampere+): wk + weights_proj
-        # hidden_states (FP32) @ fused_weight^T -> [N, head_dim + n_heads]
+        Runs cublas_mm, qk_projection_and_rope, FP8 quantize, and weight
+        scaling.  Does NOT touch the k cache or any batch-specific metadata,
+        so this can safely run inside a captured CUDA graph partition.
+
+        Returns (q_fp8, k_fp8, k_scale, weights).
+        """
         assert self._fused_wk_wp_weight is not None, \
             "post_load_weights() must be called before forward()"
         hidden_float = _to_float(hidden_states)
@@ -1704,14 +1733,16 @@ def forward(self, qr: torch.Tensor, hidden_states: torch.Tensor,
         q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim)
         q_scale = q_scale.view(-1, self.n_heads, 1)
 
-        weights, _ = maybe_execute_in_parallel(
-            lambda: self._weight_scale(weights, q_scale),
-            lambda: self._update_k_cache(
-                k_fp8, k_scale, metadata),  # store k_fp8 and k_scale in k cache
-            self.ln_events[0],
-            self.ln_events[1],
-            self.aux_stream,
-        )
+        weights = self._weight_scale(weights, q_scale)
+
+        return q_fp8, k_fp8, k_scale, weights
+
+    @torch.inference_mode()
+    def forward(self, qr: torch.Tensor, hidden_states: torch.Tensor,
+                metadata: DSAtrtllmAttentionMetadata,
+                position_ids: torch.Tensor):
+        q_fp8, k_fp8, k_scale, weights = self.pre_indexer(
+            qr, hidden_states, metadata, position_ids)
 
         # Return topk indices buffer for sparse attention [num_tokens, index_topk]
         return self.sparse_attn_indexer(metadata, hidden_states, q_fp8, k_fp8,
diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py
index bca4a4715ea2..d4172201eb13 100644
--- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py
+++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py
@@ -262,13 +262,18 @@ def piecewise_optimizer(
         if (not stop_partition and is_call_function(node, [
                 torch.ops.trtllm.attn_custom_op_inplace.default,
                 torch.ops.trtllm.mla_custom_op_inplace.default,
+                torch.ops.trtllm.mla_dsa_attn_inplace.default,
                 torch.ops.aten.index.Tensor,
                 torch.ops.aten.cumsum.default,
         ])):
             idx += 1
             node_to_graph_id[node] = idx
             exclude_modules_id.append(idx)
-            if node.target != torch.ops.trtllm.attn_custom_op_inplace.default and node.target != torch.ops.trtllm.mla_custom_op_inplace.default:
+            if (node.target != torch.ops.trtllm.attn_custom_op_inplace.default
+                    and node.target
+                    != torch.ops.trtllm.mla_custom_op_inplace.default
+                    and node.target
+                    != torch.ops.trtllm.mla_dsa_attn_inplace.default):
                 # We only know it is safe to continue splitting after attention
                 stop_partition = True
             else:
diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py
index 2f8b7b19deb5..649b4ebef073 100644
--- a/tensorrt_llm/_torch/compilation/utils.py
+++ b/tensorrt_llm/_torch/compilation/utils.py
@@ -69,6 +69,9 @@ def inplace_info():
         torch.ops.trtllm.mla_custom_op_inplace.default: {
             1: "output"
         },
+        torch.ops.trtllm.mla_dsa_attn_inplace.default: {
+            1: "output"
+        },
         torch.ops.trtllm.fused_qk_norm_rope.default: {
             1: "qkv"
         },
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 22b455ab3756..cba89c58f1f9 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -950,6 +950,84 @@ def mla_custom_op_inplace(
                                latent_cache_gen=latent_cache_gen)
 
 
+@torch.library.custom_op("trtllm::mla_dsa_proj", mutates_args=())
+def mla_dsa_proj(
+    hidden_states: torch.Tensor,
+    position_ids: Optional[torch.Tensor],
+    layer_idx: str,
+) -> List[torch.Tensor]:
+    """Token-wise projections for DSA MLA (CUDA-graph-capturable).
+
+    Runs kv_a_proj, layernorms, q_b_proj, and conditionally
+    indexer.pre_indexer (which updates the indexer k cache).
+
+    Returns [q, compressed_kv, k_pe, latent_cache] when the short-MHA path
+    handles all tokens, or [q, compressed_kv, k_pe, latent_cache, q_fp8,
+    k_fp8, k_scale, weights] when the indexer runs.  Under torch compile,
+    _should_use_short_mha returns False so the result is always length 8,
+    keeping control flow straight-line for CUDA graph capture.
+    """
+    metadata, mla_layer = extract_extra_attrs(layer_idx, "mla")
+    return mla_layer.forward_dsa_proj(position_ids, hidden_states, metadata)
+
+
+@mla_dsa_proj.register_fake
+def _mla_dsa_proj_fake(
+    hidden_states: torch.Tensor,
+    position_ids: Optional[torch.Tensor],
+    layer_idx: str,
+) -> List[torch.Tensor]:
+    # Under torch compile _should_use_short_mha is False, so always 8 tensors.
+    metadata, mla_layer = extract_extra_attrs(layer_idx, "mla")
+    num_tokens = hidden_states.shape[0]
+    indexer = mla_layer.mqa.indexer
+    q = hidden_states.new_empty(
+        [num_tokens, mla_layer.num_heads_tp * mla_layer.qk_head_dim])
+    compressed_kv = hidden_states.new_empty(
+        [num_tokens, mla_layer.kv_lora_rank])
+    k_pe = hidden_states.new_empty([num_tokens, mla_layer.qk_rope_head_dim])
+    latent_cache = hidden_states.new_empty(
+        [num_tokens, mla_layer.kv_lora_rank + mla_layer.qk_rope_head_dim])
+    # Indexer intermediates: q_fp8, k_fp8, k_scale, weights
+    q_fp8 = hidden_states.new_empty(
+        [num_tokens, indexer.n_heads, indexer.head_dim],
+        dtype=torch.float8_e4m3fn)
+    k_fp8 = hidden_states.new_empty([num_tokens, indexer.head_dim],
+                                    dtype=torch.float8_e4m3fn)
+    k_scale = hidden_states.new_empty([num_tokens, indexer.head_dim // 128],
+                                      dtype=torch.float32)
+    weights = hidden_states.new_empty([num_tokens, indexer.n_heads],
+                                      dtype=torch.float32)
+    return [
+        q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, weights
+    ]
+
+
+@torch.library.custom_op("trtllm::mla_dsa_attn_inplace",
+                         mutates_args=("output", ))
+def mla_dsa_attn_inplace(
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    k_pe: torch.Tensor,
+    latent_cache: torch.Tensor,
+    indexer_intermediates: List[torch.Tensor],
+    position_ids: Optional[torch.Tensor],
+    layer_idx: str,
+    output: torch.Tensor,
+) -> None:
+    """Batch-structure-dependent attention dispatch for DSA MLA.
+
+    indexer_intermediates is [q_fp8, k_fp8, k_scale, weights] when the
+    indexer ran in Op 1, or [] when short-MHA handled all tokens.
+    Runs sparse_attn_indexer then dispatches context/generation attention.
+    This op is excluded from CUDA graph capture.
+    """
+    metadata, mla_layer = extract_extra_attrs(layer_idx, "mla")
+    mla_layer.forward_dsa_attn(q, compressed_kv, k_pe, latent_cache,
+                               indexer_intermediates, position_ids, metadata,
+                               output)
+
+
 def fp8_block_scaling_bmm_out(
     mat1: torch.Tensor,
     mat2_fp8: torch.Tensor,
@@ -1584,27 +1662,50 @@ def forward_impl_with_dsa(self, position_ids: Optional[torch.Tensor],
         Forward pass for the MLA module with DSA (always in MQA mode).
         Writes result into output tensor in-place.
 
+        Delegates to forward_dsa_proj (token-wise projections) followed by
+        forward_dsa_attn (batch-dependent attention dispatch).
+
         Args:
             position_ids (Optional[torch.IntTensor]): The position IDs.
             hidden_states (torch.Tensor): The hidden states.
             attn_metadata (AttentionMetadata): The attention metadata.
             output (torch.Tensor): The output tensor to write results into.
         """
+        proj_outputs = self.forward_dsa_proj(position_ids, hidden_states,
+                                             attn_metadata)
+        q, compressed_kv, k_pe, latent_cache = proj_outputs[:4]
+        indexer_intermediates = proj_outputs[4:]
+        self.forward_dsa_attn(q, compressed_kv, k_pe, latent_cache,
+                              indexer_intermediates, position_ids,
+                              attn_metadata, output)
+
+    def forward_dsa_proj(
+        self,
+        position_ids: Optional[torch.Tensor],
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> List[torch.Tensor]:
+        """Token-wise projections for DSA MLA (CUDA-graph-capturable Op 1).
+
+        Runs kv_a_proj, layernorms, q_b_proj, and conditionally
+        indexer.pre_indexer_proj().
+
+        IMPORTANT: This method must NOT slice tensors by num_tokens or
+        access batch-specific metadata, so that all operations are
+        unconditionally straight-line for CUDA graph capture.  Slicing
+        to num_tokens happens in forward_dsa_attn (Op 2, outside graph).
+
+        Returns [q, compressed_kv, k_pe, latent_cache] when short-MHA
+        handles all tokens (eager only), or
+        [q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale,
+        weights] when the indexer runs.  Under torch compile
+        _should_use_short_mha returns False so it is always length 8.
+        """
         assert self.mqa is not None, "DSA is only supported in MQA mode"
-        # split q, k, v into context and gen batches
-        num_contexts = attn_metadata.num_contexts
-        num_generations = attn_metadata.num_generations
-        num_ctx_tokens = attn_metadata.num_ctx_tokens
-        num_tokens = attn_metadata.num_tokens
-
-        hidden_states = hidden_states[:num_tokens, ...]
-        if position_ids is not None:
-            position_ids = position_ids[..., :num_tokens]
 
         q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split(
             [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1)
 
-        # TODO: possibly overlap/fuse q_a_rmsnorm + kv_a_rmsnorm + indexer.k_layernorm?
         q, compressed_kv = maybe_execute_in_parallel(
             lambda: self.q_a_layernorm(q),
             lambda: self.kv_a_layernorm(compressed_kv),
@@ -1615,33 +1716,86 @@ def forward_impl_with_dsa(self, position_ids: Optional[torch.Tensor],
         qr = q
         latent_cache = torch.concat([compressed_kv, k_pe], dim=-1)
 
-        # TODO: fuse wq_b + (indexer) wlq here
         q = self.q_b_proj(q)
 
-        # Check if the short-seq MHA path will handle context, in which case
-        # the indexer (topk_indices) is not needed for context tokens.
-        # The MHA path handles cached tokens via forward_context(), which
-        # dispatches to forward_context_with_cached_kv or
-        # forward_context_with_chunked_prefill as needed.
+        use_short_mha_for_ctx = self._should_use_short_mha(
+            attn_metadata, position_ids)
+
+        # Skip the indexer when the short MHA path handles all context
+        # tokens and there are no generation tokens.
+        if use_short_mha_for_ctx and attn_metadata.num_generations == 0:
+            return [q, compressed_kv, k_pe, latent_cache]
+
+        # pre_indexer_proj is the CUDA-graph-safe portion: pure token-wise
+        # compute (cublas_mm, rope, FP8 quantize, weight scaling) with no
+        # access to batch-specific metadata or the k cache.
+        q_fp8, k_fp8, k_scale, weights = self.mqa.indexer.pre_indexer_proj(
+            qr, hidden_states, position_ids)
+
+        return [
+            q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, weights
+        ]
+
+    def forward_dsa_attn(
+        self,
+        q: torch.Tensor,
+        compressed_kv: torch.Tensor,
+        k_pe: torch.Tensor,
+        latent_cache: torch.Tensor,
+        indexer_intermediates: List[torch.Tensor],
+        position_ids: Optional[torch.Tensor],
+        attn_metadata: AttentionMetadata,
+        output: torch.Tensor,
+    ) -> None:
+        """Batch-structure-dependent attention for DSA MLA (Op 2, not graph-captured).
+
+        indexer_intermediates is [q_fp8, k_fp8, k_scale, weights] when the
+        indexer ran in Op 1, or [] when short-MHA handled all tokens.
+
+        All num_tokens slicing happens here (not in Op 1) because
+        num_tokens comes from batch-specific metadata and must not be
+        baked into CUDA graph capture.
+        """
+        num_contexts = attn_metadata.num_contexts
+        num_generations = attn_metadata.num_generations
+        num_ctx_tokens = attn_metadata.num_ctx_tokens
+        num_tokens = attn_metadata.num_tokens
+
+        # Slice Op 1 outputs to actual num_tokens (Op 1 operates on the
+        # full padded tensor for CUDA graph compatibility).
+        q = q[:num_tokens, ...]
+        compressed_kv = compressed_kv[:num_tokens, ...]
+        k_pe = k_pe[:num_tokens, ...]
+        latent_cache = latent_cache[:num_tokens, ...]
+        if position_ids is not None:
+            position_ids = position_ids[..., :num_tokens]
+
         use_short_mha_for_ctx = (num_contexts > 0
                                  and self._should_use_short_mha(
                                      attn_metadata, position_ids))
 
-        # Skip the indexer entirely when the short MHA path handles all
-        # context tokens and there are no generation tokens.
         if use_short_mha_for_ctx and num_generations == 0:
             topk_indices = None
         else:
-            topk_indices = self.indexer(
-                qr,
-                hidden_states,
+            q_fp8, k_fp8, k_scale, weights = indexer_intermediates
+            # Slice indexer intermediates to actual num_tokens (they were
+            # computed on the full padded tensor in Op 1).
+            q_fp8 = q_fp8[:num_tokens, ...]
+            k_fp8 = k_fp8[:num_tokens, ...]
+            k_scale = k_scale[:num_tokens, ...]
+            weights = weights[:num_tokens, ...]
+            # Update the indexer k cache here (outside CUDA graph) because
+            # it accesses batch-specific metadata (slot_mapping_fp8/scale).
+            self.mqa.indexer._update_k_cache(k_fp8, k_scale, attn_metadata)
+            topk_indices = self.mqa.indexer.sparse_attn_indexer(
                 attn_metadata,
-                position_ids,
+                q,  # only used for shape/device in buffer allocation
+                q_fp8,
+                k_fp8,
+                k_scale,
+                weights,
             )
 
-        assert q.shape[
-            0] == num_tokens, f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}"
-
         assert output is not None, "output must be provided"
 
         if num_contexts > 0:
@@ -1741,7 +1895,13 @@ def _should_use_short_mha(self, attn_metadata: AttentionMetadata,
         chunked context where the full attention span exceeds the threshold
         even if the new token count is small.  Falls back to num_ctx_tokens
         (total new context tokens) when max_ctx_kv_len is not set.
+
+        Disabled under torch compile so that the split DSA custom ops
+        (mla_dsa_proj / mla_dsa_attn_inplace) have unconditionally
+        straight-line control flow for CUDA graph capture.
         """
+        if is_torch_compiling():
+            return False
         if not (self.short_seq_mha_threshold > 0 and not self.apply_rotary_emb
                 and self.mapping.cp_size == 1 and position_ids is not None):
             return False
@@ -2599,9 +2759,23 @@ def forward(
         attn_output = self.create_output(hidden_states,
                                          attn_metadata.num_contexts)
         if self.register_to_config:
-            torch.ops.trtllm.mla_custom_op_inplace(
-                hidden_states, position_ids, self.layer_idx_str, attn_output,
-                None if self.is_dsa else latent_cache_gen)
+            if self.is_dsa:
+                proj_outputs = torch.ops.trtllm.mla_dsa_proj(
+                    hidden_states, position_ids, self.layer_idx_str)
+                q, compressed_kv, k_pe, latent_cache = (proj_outputs[0],
+                                                        proj_outputs[1],
+                                                        proj_outputs[2],
+                                                        proj_outputs[3])
+                indexer_intermediates = proj_outputs[4:]
+                torch.ops.trtllm.mla_dsa_attn_inplace(
+                    q, compressed_kv, k_pe, latent_cache, indexer_intermediates,
+                    position_ids, self.layer_idx_str, attn_output)
+            else:
+                torch.ops.trtllm.mla_custom_op_inplace(hidden_states,
+                                                       position_ids,
+                                                       self.layer_idx_str,
+                                                       attn_output,
+                                                       latent_cache_gen)
         elif self.is_dsa:
             self.forward_impl_with_dsa(position_ids,
                                        hidden_states,

From c8466f084d1565c7f8c238d74482a798fb17e97f Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Tue, 24 Mar 2026 07:30:52 -0700
Subject: [PATCH 02/14] [None][fix] Fix pre_indexer bogus parallel pattern and
 mla_dsa_proj docstring

- Remove no-op `lambda: weights` in pre_indexer's maybe_execute_in_parallel;
  _weight_scale already ran in pre_indexer_proj, so just call _update_k_cache
  directly.
- Fix mla_dsa_proj docstring: k cache update happens in Op 2
  (mla_dsa_attn_inplace), not Op 1.

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 10 +---------
 tensorrt_llm/_torch/modules/attention.py            |  4 +++-
 2 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 2a9573bf29dc..4656bb496703 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1684,15 +1684,7 @@ def pre_indexer(
         """
         q_fp8, k_fp8, k_scale, weights = self.pre_indexer_proj(
             qr, hidden_states, position_ids)
-
-        weights, _ = maybe_execute_in_parallel(
-            lambda: weights,
-            lambda: self._update_k_cache(k_fp8, k_scale, metadata),
-            self.ln_events[0],
-            self.ln_events[1],
-            self.aux_stream,
-        )
-
+        self._update_k_cache(k_fp8, k_scale, metadata)
         return q_fp8, k_fp8, k_scale, weights
 
     def pre_indexer_proj(
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index cba89c58f1f9..22d01bc6108c 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -959,7 +959,9 @@ def mla_dsa_proj(
     """Token-wise projections for DSA MLA (CUDA-graph-capturable).
 
     Runs kv_a_proj, layernorms, q_b_proj, and conditionally
-    indexer.pre_indexer (which updates the indexer k cache).
+    indexer.pre_indexer_proj (FP8 quantize, weight scaling).  Does NOT
+    update the indexer k cache — that happens in Op 2 (mla_dsa_attn_inplace)
+    because the scatter kernel accesses batch-specific metadata.
 
     Returns [q, compressed_kv, k_pe, latent_cache] when the short-MHA path
     handles all tokens, or [q, compressed_kv, k_pe, latent_cache, q_fp8,

From f0dc05f20b76faeb9cb374b0a8bb0bee00b34a26 Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Tue, 24 Mar 2026 07:41:34 -0700
Subject: [PATCH 03/14] [None][refactor] Move _update_k_cache into
 sparse_attn_indexer

Move _update_k_cache call to the top of sparse_attn_indexer so
the k cache is populated right before prefill chunks gather from it.
Remove pre_indexer (now redundant); forward() and forward_dsa_proj
both call pre_indexer_proj directly.

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    | 25 ++++---------------
 tensorrt_llm/_torch/modules/attention.py      |  3 ---
 2 files changed, 5 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 4656bb496703..f85c06a51db3 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1395,6 +1395,9 @@ def sparse_attn_indexer(
         weights: torch.Tensor,
         use_custom_topk: bool = True,
     ) -> torch.Tensor:
+        # Update the indexer k cache before prefill chunks gather from it.
+        self._update_k_cache(k_fp8, k_scale, metadata)
+
         num_contexts = metadata.num_contexts
         num_generations = metadata.num_generations
         num_ctx_tokens = metadata.num_ctx_tokens
@@ -1669,24 +1672,6 @@ def _prep_q_or_k(self, qk_pe: torch.Tensor, qk_nope: torch.Tensor):
             qk_pe, qk_nope, self.scale_fmt == "ue8m0")
         return fp8_out, scale
 
-    @torch.inference_mode()
-    def pre_indexer(
-        self, qr: torch.Tensor, hidden_states: torch.Tensor,
-        metadata: DSAtrtllmAttentionMetadata, position_ids: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Token-wise projections, FP8 quantize, weight scaling, and k cache update.
-
-        Runs the full indexer pre-computation including k cache update.
-        Used by the eager path (Indexer.forward) where everything runs
-        outside CUDA graph capture.
-
-        Returns (q_fp8, k_fp8, k_scale, weights).
-        """
-        q_fp8, k_fp8, k_scale, weights = self.pre_indexer_proj(
-            qr, hidden_states, position_ids)
-        self._update_k_cache(k_fp8, k_scale, metadata)
-        return q_fp8, k_fp8, k_scale, weights
-
     def pre_indexer_proj(
         self, qr: torch.Tensor, hidden_states: torch.Tensor,
         position_ids: torch.Tensor
@@ -1733,8 +1718,8 @@ def pre_indexer_proj(
     def forward(self, qr: torch.Tensor, hidden_states: torch.Tensor,
                 metadata: DSAtrtllmAttentionMetadata,
                 position_ids: torch.Tensor):
-        q_fp8, k_fp8, k_scale, weights = self.pre_indexer(
-            qr, hidden_states, metadata, position_ids)
+        q_fp8, k_fp8, k_scale, weights = self.pre_indexer_proj(
+            qr, hidden_states, position_ids)
 
         # Return topk indices buffer for sparse attention [num_tokens, index_topk]
         return self.sparse_attn_indexer(metadata, hidden_states, q_fp8, k_fp8,
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 22d01bc6108c..ca4487e90363 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -1786,9 +1786,6 @@ def forward_dsa_attn(
             k_fp8 = k_fp8[:num_tokens, ...]
             k_scale = k_scale[:num_tokens, ...]
             weights = weights[:num_tokens, ...]
-            # Update the indexer k cache here (outside CUDA graph) because
-            # it accesses batch-specific metadata (slot_mapping_fp8/scale).
-            self.mqa.indexer._update_k_cache(k_fp8, k_scale, attn_metadata)
             topk_indices = self.mqa.indexer.sparse_attn_indexer(
                 attn_metadata,
                 q,  # only used for shape/device in buffer allocation

From 66627a1b85108c92bc948a44e2913d67049dcb0c Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Mon, 30 Mar 2026 22:36:58 -0700
Subject: [PATCH 04/14] [None][chore] Clean up MLA DSA custom op dispatch

- Remove dead is_dsa branch from mla_custom_op_inplace since DSA is now
  exclusively handled by the split mla_dsa_proj/mla_dsa_attn_inplace ops
- Use literal 1 for k_scale shape to match C++ fusedCatFp8 kernel output
- Simplify proj_outputs unpacking

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/attention.py | 24 +++++++-----------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index ca4487e90363..118d60bbe701 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -937,17 +937,11 @@ def mla_custom_op_inplace(
     latent_cache_gen: Optional[torch.Tensor],
 ) -> None:
     metadata, mla_layer = extract_extra_attrs(layer_idx, "mla")
-    if mla_layer.is_dsa:
-        mla_layer.forward_impl_with_dsa(position_ids,
-                                        hidden_states,
-                                        metadata,
-                                        output=output)
-    else:
-        mla_layer.forward_impl(position_ids,
-                               hidden_states,
-                               metadata,
-                               output=output,
-                               latent_cache_gen=latent_cache_gen)
+    mla_layer.forward_impl(position_ids,
+                           hidden_states,
+                           metadata,
+                           output=output,
+                           latent_cache_gen=latent_cache_gen)
 
 
 @torch.library.custom_op("trtllm::mla_dsa_proj", mutates_args=())
@@ -996,8 +990,7 @@ def _mla_dsa_proj_fake(
         dtype=torch.float8_e4m3fn)
     k_fp8 = hidden_states.new_empty([num_tokens, indexer.head_dim],
                                     dtype=torch.float8_e4m3fn)
-    k_scale = hidden_states.new_empty([num_tokens, indexer.head_dim // 128],
-                                      dtype=torch.float32)
+    k_scale = hidden_states.new_empty([num_tokens, 1], dtype=torch.float32)
     weights = hidden_states.new_empty([num_tokens, indexer.n_heads],
                                       dtype=torch.float32)
     return [
@@ -2761,10 +2754,7 @@ def forward(
             if self.is_dsa:
                 proj_outputs = torch.ops.trtllm.mla_dsa_proj(
                     hidden_states, position_ids, self.layer_idx_str)
-                q, compressed_kv, k_pe, latent_cache = (proj_outputs[0],
-                                                        proj_outputs[1],
-                                                        proj_outputs[2],
-                                                        proj_outputs[3])
+                q, compressed_kv, k_pe, latent_cache = proj_outputs[:4]
                 indexer_intermediates = proj_outputs[4:]
                 torch.ops.trtllm.mla_dsa_attn_inplace(
                     q, compressed_kv, k_pe, latent_cache, indexer_intermediates,

From c21c41cb6120740712af283b1ee554f172950cc4 Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Mon, 30 Mar 2026 22:54:33 -0700
Subject: [PATCH 05/14] [None][chore] Restore quant_block_size assertion in
 sparse_attn_indexer

The assertion was dropped when the old Indexer.forward was split into
pre_indexer_proj and sparse_attn_indexer. Restore it in
sparse_attn_indexer which has access to metadata.kv_cache_manager.

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index f85c06a51db3..7605210a02e6 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1395,6 +1395,9 @@ def sparse_attn_indexer(
         weights: torch.Tensor,
         use_custom_topk: bool = True,
     ) -> torch.Tensor:
+        assert metadata.kv_cache_manager is None or \
+            metadata.kv_cache_manager.quant_block_size == 128, \
+            "Only support quant_block_size = 128 for now"
         # Update the indexer k cache before prefill chunks gather from it.
         self._update_k_cache(k_fp8, k_scale, metadata)
 

From 7556e67ff799d355644fedc0b954df9e26371019 Mon Sep 17 00:00:00 2001
From: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Date: Thu, 2 Apr 2026 15:46:40 +0800
Subject: [PATCH 06/14] [None][fix] Fix DSACacheManager and RocketCacheManager
 KV cache estimation ignoring num_layers for draft models (#12683)

Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>
Co-authored-by: Lanyu Liao <lancelly@users.noreply.github.com>
---
 .../_torch/attention_backend/sparse/dsa.py    |  9 ++--
 .../_torch/attention_backend/sparse/rocket.py |  9 ++--
 .../_torch/pyexecutor/resource_manager.py     | 48 +++++++++----------
 3 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 7605210a02e6..9cc38da0f1ec 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1931,7 +1931,9 @@ def shutdown(self):
         super().shutdown()
 
     @staticmethod
-    def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
+    def get_cache_size_per_token(model_config: ModelConfig,
+                                 mapping: Mapping,
+                                 num_layers: Optional[int] = None,
                                  **kwargs):
         config = model_config.pretrained_config
         sparse_attn_config = model_config.sparse_attention_config
@@ -1948,9 +1950,8 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
         # get head dim
         head_dim = config.kv_lora_rank + config.qk_rope_head_dim
 
-        # provide at least 1 layer to prevent division by zero cache size
-        num_attention_layers = max(
-            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         mem_per_token *= num_attention_layers * head_dim
 
         # 1 for K, others for indexer K cache
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket.py
index a76ba4d2e569..fa7e8e4af659 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket.py
@@ -1037,7 +1037,9 @@ def free_resources(self, request):
         self.kt_cache_manager.free_resources(request)
 
     @staticmethod
-    def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
+    def get_cache_size_per_token(model_config: ModelConfig,
+                                 mapping: Mapping,
+                                 num_layers: Optional[int] = None,
                                  **kwargs):
         # get kv cache dtype bytes
         mem_per_token = 2
@@ -1061,9 +1063,8 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
             head_dim = config.hidden_size // config.num_attention_heads
         head_dim = head_dim * num_key_value_heads // tp_size
 
-        # provide at least 1 layer to prevent division by zero cache size
-        num_attention_layers = max(
-            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         mem_per_token *= num_attention_layers * head_dim
 
         # K and V
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 43a7503655b6..f2f74381660c 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -842,6 +842,26 @@ def calculate_scaling_factor_size_bytes(
         return get_size_in_bytes(cache_size // quant_vector_size,
                                  scaling_factor_dtype)
 
+    @staticmethod
+    def _resolve_num_attention_layers(
+        model_config: ModelConfigPython,
+        mapping: Mapping,
+        num_layers: Optional[int] = None,
+    ) -> int:
+        """Compute the effective number of attention layers for cache sizing.
+
+        When *num_layers* is explicitly provided (e.g. for draft models whose
+        HF config layer count differs from runtime), it is used directly
+        without PP distribution.  Otherwise the layer count is derived from
+        the model config and distributed evenly across PP ranks via
+        ``mapping.pp_layers``.
+        """
+        if num_layers is not None:
+            return max(num_layers, 1)
+        # provide at least 1 layer to prevent division by zero cache size
+        return max(
+            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
+
     # TODO: refactor get_cache_size_per_token and get_cache_bytes_per_token to use the same logic
     @staticmethod
     def get_cache_size_per_token(model_config: ModelConfigPython,
@@ -870,18 +890,8 @@ def get_cache_size_per_token(model_config: ModelConfigPython,
             head_dim = head_dim * num_key_value_heads // tp_size
             kv_factor = 2
 
-        # When num_layers is explicitly provided (e.g. for draft models
-        # where the HF config layer count differs from runtime), use it
-        # directly without PP distribution.  Draft layers have their own
-        # PP assignment logic (see get_pp_layers) that doesn't match the
-        # standard uniform split, so pp_layers() would give wrong results.
-        if num_layers is not None:
-            num_attention_layers = max(num_layers, 1)
-        else:
-            # provide at least 1 layer to prevent division by zero cache size
-            num_attention_layers = max(
-                len(mapping.pp_layers(model_config.get_num_attention_layers())),
-                1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         # K and V
         mem_per_token = kv_factor * num_attention_layers * head_dim
         # The data type bytes.
@@ -2521,18 +2531,8 @@ def get_cache_size_per_token(model_config: ModelConfigPython,
             head_dim = head_dim * num_key_value_heads // tp_size
             kv_factor = 2
 
-        # When num_layers is explicitly provided (e.g. for draft models
-        # where the HF config layer count differs from runtime), use it
-        # directly without PP distribution.  Draft layers have their own
-        # PP assignment logic (see get_pp_layers) that doesn't match the
-        # standard uniform split, so pp_layers() would give wrong results.
-        if num_layers is not None:
-            num_attention_layers = max(num_layers, 1)
-        else:
-            # provide at least 1 layer to prevent division by zero cache size
-            num_attention_layers = max(
-                len(mapping.pp_layers(model_config.get_num_attention_layers())),
-                1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         mem_per_token *= num_attention_layers * head_dim
 
         # K and V

From d9038d3b904dae99c6d1cbe73058fe9f99ab4616 Mon Sep 17 00:00:00 2001
From: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Date: Thu, 2 Apr 2026 17:38:36 +0800
Subject: [PATCH 07/14] [None][fix] Fix compute token accounting for KV cache
 reuse with context chunking (#12682)

Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>
Co-authored-by: Lanyu Liao <lancelly@users.noreply.github.com>
---
 .../batch_manager/microBatchScheduler.cpp     |  75 ++++++-----
 .../_torch/pyexecutor/model_engine.py         |  18 +++
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 125 +++++++++++++++++-
 .../_torch/pyexecutor/resource_manager.py     |  76 ++++++++++-
 4 files changed, 252 insertions(+), 42 deletions(-)

diff --git a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp
index 92d4589e0b66..4e1ee9d9c874 100644
--- a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp
+++ b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp
@@ -24,6 +24,25 @@ namespace tensorrt_llm::batch_manager
 
 using SizeType32 = MicroBatchScheduler::SizeType32;
 
+/// Return the forward-pass token cost for a context chunk with KV cache reuse.
+///
+/// setPrepopulatedPromptLen shifts the chunk window right by the prepopulated
+/// amount rather than shrinking it.  For non-last chunks the model still
+/// processes approximately @p chunkSize tokens; only for the last chunk is the
+/// cost @p contextRemaining - @p reusable.
+static SizeType32 reuse_adjusted_compute(SizeType32 chunkSize, SizeType32 reusable, SizeType32 contextRemaining)
+{
+    if (reusable <= 0)
+    {
+        return chunkSize;
+    }
+    if (reusable + chunkSize < contextRemaining)
+    {
+        return chunkSize;
+    }
+    return std::max<SizeType32>(0, contextRemaining - reusable);
+}
+
 MicroBatchScheduler::MicroBatchScheduler(std::optional<batch_scheduler::ContextChunkingConfig> ctxChunkConfig,
     std::optional<SizeType32> maxContextLength, LlmRequestState noScheduleUntilState,
     LlmRequestState noScheduleAfterState)
@@ -38,16 +57,15 @@ void MicroBatchScheduler::fitDraftTokens(RequestVector& contextsToBeChunked,
     std::optional<SizeType32> ctxTokensCapacity, SizeType32 const chunkUnitSize,
     std::optional<SizeType32> const& maxContextLength)
 {
-    // How many compute tokens (chunk - reusable) are in this batch already?
+    // How many compute tokens are in this batch already?
     SizeType32 numCtxTokens{0};
     for (auto const& llmReq : contextsToBeChunked)
     {
         SizeType32 const chunkSize = llmReq->getContextChunkSize();
-        // contextRemaining = P for first chunk; used to compute actual model token count.
         SizeType32 const contextRemaining = llmReq->getContextRemainingLength();
         SizeType32 const reusable
             = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), contextRemaining) : 0;
-        numCtxTokens += std::min(chunkSize, std::max<SizeType32>(0, contextRemaining - reusable));
+        numCtxTokens += reuse_adjusted_compute(chunkSize, reusable, contextRemaining);
     }
 
     // Discard draft tokens that won't fit into the existing chunk unit, max
@@ -125,14 +143,13 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
             SizeType32 actualChunkSize = llmReq->getContextChunkSize();
             SizeType32 actualIncrement = actualChunkSize - pastChunkSize;
 
-            // Compute-aware budget: reusable tokens are served from cache and do not
-            // consume forward-pass capacity. Only the tokens beyond the reusable prefix count.
-            SizeType32 const reusable = llmReq->isFirstContextChunk()
-                ? std::min(llmReq->getEstimatedReusableTokens(), llmReq->getContextRemainingLength())
-                : 0;
-            SizeType32 const pastCompute = std::max<SizeType32>(0, pastChunkSize - std::min(reusable, pastChunkSize));
-            SizeType32 const actualCompute
-                = std::max<SizeType32>(0, actualChunkSize - std::min(reusable, actualChunkSize));
+            // Compute-aware budget accounting for setPrepopulatedPromptLen's
+            // chunk-shift behaviour (non-last chunks keep their full size).
+            SizeType32 const contextRemaining = llmReq->getContextRemainingLength();
+            SizeType32 const reusable
+                = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), contextRemaining) : 0;
+            SizeType32 const pastCompute = reuse_adjusted_compute(pastChunkSize, reusable, contextRemaining);
+            SizeType32 const actualCompute = reuse_adjusted_compute(actualChunkSize, reusable, contextRemaining);
             SizeType32 const computeIncrement = actualCompute - pastCompute;
 
             if ((ctxTokensCapacity && numCtxTokens + computeIncrement > ctxTokensCapacity.value())
@@ -181,24 +198,21 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
     for (auto& llmReq : contextsToBeChunked)
     {
         SizeType32 const suggestedChunkSize = llmReq->getContextRemainingLength();
-        // Reusable tokens are "free" — they don't consume forward-pass compute budget.
         SizeType32 const reusable
             = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), suggestedChunkSize) : 0;
-        SizeType32 const computeCost = suggestedChunkSize - reusable;
+        SizeType32 const computeCost = reuse_adjusted_compute(suggestedChunkSize, reusable, suggestedChunkSize);
         SizeType32 actualChunkSize = suggestedChunkSize;
         if (ctxTokensCapacity && computeCost > ctxTokensCapacity.value())
         {
-            // Model processes min(chunk_size, P - reusable) tokens starting from position reusable.
-            // To keep model tokens within budget: chunk_size <= capacity (not reusable + capacity).
             actualChunkSize = ctxTokensCapacity.value();
         }
         if (maxContextLength)
         {
-            // maxContextLength limits compute tokens, not total tokens.
-            SizeType32 const actualCompute = std::max<SizeType32>(0, actualChunkSize - reusable);
+            SizeType32 const actualCompute = reuse_adjusted_compute(actualChunkSize, reusable, suggestedChunkSize);
             if (actualCompute > maxContextLength.value())
             {
-                actualChunkSize = std::min<SizeType32>(reusable + maxContextLength.value(), suggestedChunkSize);
+                actualChunkSize = maxContextLength.value();
+                actualChunkSize = std::min(actualChunkSize, suggestedChunkSize);
             }
         }
         if (actualChunkSize != suggestedChunkSize)
@@ -208,10 +222,7 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
         llmReq->setContextChunkSize(actualChunkSize);
         if (ctxTokensCapacity)
         {
-            // Decrement by actual model token count: min(chunk_size, P - reusable).
-            // This equals min(actualChunkSize, computeCost) since computeCost = suggestedChunkSize - reusable.
-            SizeType32 const modelCost
-                = std::min(actualChunkSize, std::max<SizeType32>(0, suggestedChunkSize - reusable));
+            SizeType32 const modelCost = reuse_adjusted_compute(actualChunkSize, reusable, suggestedChunkSize);
             ctxTokensCapacity = ctxTokensCapacity.value() - modelCost;
         }
     }
@@ -309,10 +320,12 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
             if (!mCtxChunkConfig) // skip chunking
             {
                 constexpr SizeType32 beam{0};
-                reqNumTokens
-                    = llmReq->getNumTokens(beam) + (llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0);
-                // Compute tokens = total - reusable (at least 1 to make progress)
-                SizeType32 const computeTokens = std::max(1, reqNumTokens - reusable);
+                SizeType32 const contextTokens = llmReq->getNumTokens(beam);
+                SizeType32 const draftTokens = llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0;
+                reqNumTokens = contextTokens + draftTokens;
+                SizeType32 const contextCompute
+                    = reuse_adjusted_compute(contextTokens, reusable, llmReq->getContextRemainingLength());
+                SizeType32 const computeTokens = std::max(1, contextCompute + draftTokens);
                 TLLM_CHECK_WITH_INFO(!mMaxContextLength || computeTokens <= mMaxContextLength.value(),
                     "Context compute tokens (%d) exceeds the limit value (%d)", computeTokens,
                     mMaxContextLength.value());
@@ -329,9 +342,8 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
                 llmReq->setContextChunkSize(llmReq->getContextRemainingLength());
                 auto const draftTokens
                     = (llmReq->isLastContextChunk() && llmReq->hasDraftTokens()) ? llmReq->getNumDraftTokens() : 0;
-                // Compute cost: context compute + draft tokens
-                // (reusable tokens only offset context tokens, not draft tokens)
-                SizeType32 const contextCompute = std::max(0, llmReq->getContextChunkSize() - reusable);
+                SizeType32 const contextCompute = reuse_adjusted_compute(
+                    llmReq->getContextChunkSize(), reusable, llmReq->getContextRemainingLength());
                 SizeType32 computeTokens = contextCompute + draftTokens;
 
                 if (mMaxContextLength)
@@ -398,10 +410,9 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
         if (llmReq->getContextChunkSize() > 0)
         {
             contextRequests.emplace_back(llmReq);
-            // Only count compute tokens (total - reusable).
-            // Reusable credit only applies to the first context chunk.
             SizeType32 const reusable = llmReq->isFirstContextChunk() ? llmReq->getEstimatedReusableTokens() : 0;
-            SizeType32 const computeTokens = std::max(0, llmReq->getContextChunkSize() - reusable);
+            SizeType32 const computeTokens
+                = reuse_adjusted_compute(llmReq->getContextChunkSize(), reusable, llmReq->getContextRemainingLength());
             batchNumTokens += computeTokens;
             TLLM_LOG_DEBUG("context request scheduled: ID %lu, chunk size %d%s", llmReq->mRequestId,
                 llmReq->getContextChunkSize(), reusable > 0 ? (", reusable " + std::to_string(reusable)).c_str() : "");
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 88ec88a5115e..e52a9eb40743 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -2636,6 +2636,24 @@ def previous_seq_slots_device():
         num_tokens = len(input_ids)
         num_draft_tokens = len(draft_tokens)
         total_num_tokens = len(position_ids)
+        if total_num_tokens > self.max_num_tokens:
+            ctx_details = []
+            for r in scheduled_requests.context_requests:
+                pos = r.context_current_position
+                csz = r.context_chunk_size
+                full = len(r.get_tokens(0))
+                tokens = min(csz, max(0, full - pos))
+                ctx_details.append(
+                    f"rid={r.py_request_id} pos={pos} chunk={csz} "
+                    f"full={full} tokens={tokens}")
+            gen_count = len(scheduled_requests.generation_requests)
+            from tensorrt_llm.logger import logger as _mnt_logger
+            _mnt_logger.error(
+                f"MNT overflow: total={total_num_tokens} "
+                f"max={self.max_num_tokens} "
+                f"ctx_reqs={len(scheduled_requests.context_requests)} "
+                f"gen_reqs={gen_count} "
+                f"ctx_breakdown=[{'; '.join(ctx_details)}]")
         assert total_num_tokens <= self.max_num_tokens, (
             f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})"
         )
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 130769730689..bc8872307d87 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2731,6 +2731,116 @@ def _balance_adp_requests(self, context_requests: list[LlmRequest],
                     balanced_context_requests = context_requests
         return balanced_context_requests
 
+    @staticmethod
+    def _compute_scheduled_tokens(context_requests, generation_requests):
+        """Compute the total number of scheduled tokens for batch waiting decisions.
+
+        For context requests, we estimate the actual compute tokens for this
+        iteration (excluding tokens served from KV cache).
+
+        For generation requests, each contributes 1 + num_draft_tokens.
+
+        Note on reusable token handling:
+        estimated_reusable_tokens is an absolute count from position 0.
+        Depending on the scheduler, context_current_position may or may not
+        have been advanced past the reusable prefix by the time this method
+        is called:
+        - V1 scheduler: prepare_context runs after scheduling, so
+          context_current_position is still 0.
+        - V2 scheduler: prepare_context runs during scheduling, so
+          context_current_position is already advanced to the reused offset.
+        To handle both correctly, the reusable credit applied to the current
+        chunk is max(0, reusable - context_current_position), i.e. only the
+        portion of the reusable range that falls within this chunk's span.
+        """
+        num_scheduled_ctx_tokens = 0
+        for ctx_req in context_requests:
+            reusable = (ctx_req.estimated_reusable_tokens
+                        if ctx_req.is_first_context_chunk else 0)
+            # Credit only the reusable tokens that overlap with the current
+            # chunk: if context_current_position has already been advanced past
+            # the reusable prefix (V2), the credit is 0; if not (V1), the full
+            # reusable count is subtracted.
+            reusable_in_chunk = max(0,
+                                    reusable - ctx_req.context_current_position)
+            remaining = ctx_req.context_remaining_length
+            if (reusable_in_chunk > 0 and
+                    reusable_in_chunk + ctx_req.context_chunk_size < remaining):
+                compute = ctx_req.context_chunk_size
+            else:
+                compute = max(1, ctx_req.context_chunk_size - reusable_in_chunk)
+            num_scheduled_ctx_tokens += compute
+        num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens
+                                       for gen_req in generation_requests)
+        return num_scheduled_ctx_tokens + num_scheduled_gen_tokens
+
+    def _maybe_log_batch_wait_decision(
+        self,
+        context_requests: list[LlmRequest],
+        generation_requests: list[LlmRequest],
+        num_scheduled_tokens: int,
+        wait_threshold: float,
+        should_waiting: bool,
+    ) -> None:
+        """Diagnostics for batch_wait: set TLLM_LOG_BATCH_WAIT=1 (rank 0 only)."""
+        if self.dist.rank != 0:
+            return
+
+        num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens
+                                       for gen_req in generation_requests)
+        num_scheduled_ctx_formula = num_scheduled_tokens - num_scheduled_gen_tokens
+
+        chunk_ctx_sum = 0
+        ctx_summaries: List[str] = []
+        max_detail = 4
+        for i, ctx_req in enumerate(context_requests):
+            full_len = len(ctx_req.get_tokens(0))
+            begin = ctx_req.context_current_position
+            chunk_sz = ctx_req.context_chunk_size
+            this_chunk = min(chunk_sz, max(0, full_len - begin))
+            chunk_ctx_sum += this_chunk
+            reusable = (ctx_req.estimated_reusable_tokens
+                        if ctx_req.is_first_context_chunk else 0)
+            reusable_in_chunk = max(0, reusable - begin)
+            remaining = ctx_req.context_remaining_length
+            if (reusable_in_chunk > 0
+                    and reusable_in_chunk + chunk_sz < remaining):
+                formula_contrib = chunk_sz
+            else:
+                formula_contrib = max(1, chunk_sz - reusable_in_chunk)
+            if i < max_detail:
+                ctx_summaries.append(
+                    f"rid={ctx_req.py_request_id} full={full_len} pos={begin} "
+                    f"chunk_sz={chunk_sz} this_chunk={this_chunk} "
+                    f"reusable={reusable} formula_contrib={formula_contrib}")
+        n_ctx = len(context_requests)
+        if n_ctx > max_detail:
+            ctx_summaries.append(f"... +{n_ctx - max_detail} more ctx req(s)")
+
+        logger.info(
+            "batch_wait: formula_total=",
+            num_scheduled_tokens,
+            " formula_ctx=",
+            num_scheduled_ctx_formula,
+            " formula_gen=",
+            num_scheduled_gen_tokens,
+            " chunk_ctx_sum=",
+            chunk_ctx_sum,
+            " threshold=",
+            wait_threshold,
+            " wait_iter=",
+            self.batch_wait_iters_count,
+            "/",
+            self.batch_wait_timeout_iters,
+            " should_defer_ctx=",
+            should_waiting,
+            " num_gen=",
+            len(generation_requests),
+            " ctx_detail=[",
+            "; ".join(ctx_summaries),
+            "]",
+        )
+
     def _waiting_requests(self, context_requests: list[LlmRequest],
                           generation_requests: list[LlmRequest]):
         """
@@ -2740,13 +2850,16 @@ def _waiting_requests(self, context_requests: list[LlmRequest],
         - The number of waiting iterations is smaller than `self.batch_wait_timeout_iters`.
         """
 
-        num_scheduled_ctx_tokens = sum(
-            len(ctx_req.get_tokens(0)) for ctx_req in context_requests)
-        num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens
-                                       for gen_req in generation_requests)
-        num_scheduled_tokens = num_scheduled_ctx_tokens + num_scheduled_gen_tokens
+        num_scheduled_tokens = self._compute_scheduled_tokens(
+            context_requests, generation_requests)
+        wait_threshold = (self.batch_wait_max_tokens_ratio *
+                          self.max_num_tokens)
 
-        should_waiting = self.batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < self.batch_wait_max_tokens_ratio * self.max_num_tokens
+        should_waiting = self.batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < wait_threshold
+        self._maybe_log_batch_wait_decision(context_requests,
+                                            generation_requests,
+                                            num_scheduled_tokens,
+                                            wait_threshold, should_waiting)
         if should_waiting:
             self.batch_wait_iters_count += 1
             return []
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index f2f74381660c..b1e7390306bb 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -616,6 +616,22 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
             # wait for all pending work to finish before launching offload/onboarding/partial copy
             self.impl.sync_transfer_manager_with_buffer_manager()
 
+            # Pre-addSequence budget re-validation.  The C++ scheduler
+            # should already account for the chunk-shift cost, but under
+            # heavy KV-cache eviction the actual reuse may be lower than
+            # estimated.  We re-probe the radix tree and estimate the
+            # true forward cost; if it exceeds the remaining budget the
+            # request is skipped (re-scheduled next iteration).
+            remaining_budget = None
+            if self.enable_block_reuse and not self.is_draft:
+                gen_tokens = sum(
+                    req.get_beam_width_by_iter(for_next_iteration=False) +
+                    get_draft_token_length(req)
+                    for req in scheduled_batch.generation_requests)
+                remaining_budget = self.max_num_tokens - gen_tokens
+
+            accepted_ctx_requests = []
+
             # allocate KV Cache
             for req in scheduled_batch.context_requests:
                 req_beam_width = req.sampling_config.beam_width
@@ -632,9 +648,37 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                 else:
                     if req.is_first_context_chunk and self._kv_connector_should_add_sequence(
                             req):
-                        self.impl.add_sequence(req.py_request_id,
-                                               req.prompt_len, req_beam_width,
-                                               req)
+                        if remaining_budget is not None:
+                            unique_tokens = req.get_unique_tokens(0)
+                            reusable_blocks = self.impl.count_reusable_blocks(
+                                unique_tokens, req, False)
+                            actual_reuse = (reusable_blocks *
+                                            self.tokens_per_block)
+                            req_compute = self._estimate_post_reuse_compute(
+                                actual_reuse, req.context_chunk_size,
+                                req.prompt_len)
+                            if req_compute > remaining_budget:
+                                logger.warning(
+                                    f"Reuse budget: skip req "
+                                    f"{req.py_request_id} "
+                                    f"(compute={req_compute}, "
+                                    f"chunk={req.context_chunk_size}, "
+                                    f"reuse={actual_reuse}, "
+                                    f"remaining={remaining_budget})")
+                                continue
+                            remaining_budget -= req_compute
+
+                        try:
+                            self.impl.add_sequence(req.py_request_id,
+                                                   req.prompt_len,
+                                                   req_beam_width, req)
+                        except RuntimeError:
+                            logger.warning(
+                                f"add_sequence: req "
+                                f"{req.py_request_id} already exists, "
+                                f"skipping")
+                            accepted_ctx_requests.append(req)
+                            continue
                         for _ in range(self.num_extra_kv_tokens):
                             self.impl.add_token(req.py_request_id)
                         for _ in range(get_draft_token_length(req)):
@@ -644,9 +688,16 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                             block_ids = self.get_cache_indices(req)
                             self.kv_connector_manager.update_state_after_alloc(
                                 req, block_ids)
+                    elif remaining_budget is not None:
+                        reusable = (req.estimated_reusable_tokens
+                                    if req.is_first_context_chunk else 0)
+                        remaining_budget -= self._estimate_post_reuse_compute(
+                            reusable, req.context_chunk_size, req.prompt_len)
+
+                accepted_ctx_requests.append(req)
 
             # A request may change from `context_requests_chunking` to `context_requests_last_chunk` in `add_sequence` due to KV cache reuse, so we rebuild the context request lists here.
-            scheduled_batch.reset_context_requests()
+            scheduled_batch.reset_context_requests(accepted_ctx_requests)
 
             for req in scheduled_batch.generation_requests:
                 if self.mapping.has_cp_helix():
@@ -671,6 +722,23 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
             self.kv_connector_manager.build_scheduler_output(
                 scheduled_batch, self)
 
+    def _estimate_post_reuse_compute(self, reuse_tokens: int, chunk_size: int,
+                                     prompt_len: int) -> int:
+        """Estimate forward compute tokens after setPrepopulatedPromptLen.
+
+        For non-last chunks the chunk window shifts right by the reused
+        amount and the forward cost is approximately chunk_size.  For
+        last chunks the cost is prompt_len - reuse (original formula).
+        """
+        P = reuse_tokens
+        if P > 0 and P < prompt_len:
+            if P + chunk_size < prompt_len:
+                aligned_end = ((P + chunk_size) // self.tokens_per_block *
+                               self.tokens_per_block)
+                return max(1, aligned_end - P)
+            return max(1, prompt_len - P)
+        return chunk_size
+
     def _kv_connector_should_add_sequence(self, request: LlmRequest) -> bool:
         return self.kv_connector_manager is None or self.kv_connector_manager.should_add_sequence(
             request)

From ae2139f7a45b609c47aab0754699ba67beeb8859 Mon Sep 17 00:00:00 2001
From: Yukun He <23156053+hyukn@users.noreply.github.com>
Date: Thu, 2 Apr 2026 18:04:35 +0800
Subject: [PATCH 08/14] [https://nvbugs/5983390][perf] Cherry-pick #12581:
 Multiple host perf optimizations for DSA part  (#12681)

Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
---
 .../kernels/IndexerKCacheGather.h             |  35 ++
 .../kernels/convertReqIndexToGlobal.cu        |  97 ++++
 .../kernels/convertReqIndexToGlobal.h         |  34 ++
 .../kernels/indexerKCacheGather.cu            | 163 +++++++
 cpp/tensorrt_llm/thop/CMakeLists.txt          |   2 +
 .../thop/IndexerKCacheGatherOp.cpp            | 121 +++++
 .../thop/convertReqIndexToGlobalOp.cpp        |  84 ++++
 .../_torch/attention_backend/sparse/dsa.py    | 192 ++++----
 .../_torch/attention_backend/sparse/kernel.py |  12 +-
 .../_torch/custom_ops/cpp_custom_ops.py       |  14 +
 tensorrt_llm/_torch/speculative/mtp.py        |  26 +-
 .../attention/sparse/test_cpp_custom_ops.py   | 421 ++++++++++++++++++
 .../sparse/test_triton_gather_k_cache.py      |  28 --
 13 files changed, 1098 insertions(+), 131 deletions(-)
 create mode 100644 cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
 create mode 100644 cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu
 create mode 100644 cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h
 create mode 100644 cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
 create mode 100644 cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
 create mode 100644 cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp
 create mode 100644 tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py

diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
new file mode 100644
index 000000000000..60422dea5910
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mapping_fp8,
+    int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
+    int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
+    int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
+    cudaStream_t stream = 0);
+
+}
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu
new file mode 100644
index 000000000000..aaf792e33c1d
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu
@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "convertReqIndexToGlobal.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+// Each thread handles one element at (token_id, col).
+// Grid: (num_tokens, ceil(numTopkTokens / blockDim.x))
+__global__ void convertReqIndexToGlobalKernel(int32_t const* __restrict__ reqId, int32_t const* __restrict__ blockTable,
+    int32_t const* __restrict__ tokenIndices, int32_t* __restrict__ output, int32_t numTopkTokens,
+    int32_t maxNumBlocksPerReq, int32_t blockSize, int32_t strideFactor, int32_t layerId, int64_t btStride0,
+    int64_t btStride1, int64_t tiStride0, int64_t tiStride1, int64_t outStride0, int64_t outStride1)
+{
+    int32_t const tokenId = blockIdx.x;
+    int32_t const col = blockIdx.y * blockDim.x + threadIdx.x;
+
+    if (col >= numTopkTokens)
+    {
+        return;
+    }
+
+    // Load request id for this token
+    int32_t const req = reqId[tokenId];
+
+    // Load token index
+    int32_t const tok = tokenIndices[tokenId * tiStride0 + col * tiStride1];
+
+    // Invalid token → output -1
+    if (tok < 0)
+    {
+        output[tokenId * outStride0 + col * outStride1] = -1;
+        return;
+    }
+
+    // Compute block id and in-block offset
+    int32_t const blockId = tok / blockSize;
+    int32_t const inblockOff = tok % blockSize + layerId * blockSize;
+
+    // Guard block_table access
+    if (blockId >= maxNumBlocksPerReq)
+    {
+        output[tokenId * outStride0 + col * outStride1] = -1;
+        return;
+    }
+
+    int32_t const base = blockTable[req * btStride0 + blockId * btStride1];
+
+    // Padding entry in block table
+    if (base < 0)
+    {
+        output[tokenId * outStride0 + col * outStride1] = -1;
+        return;
+    }
+
+    output[tokenId * outStride0 + col * outStride1] = base * strideFactor + inblockOff;
+}
+
+void invokeConvertReqIndexToGlobal(int32_t const* reqId, int32_t const* blockTable, int32_t const* tokenIndices,
+    int32_t* output, int32_t numTokens, int32_t numTopkTokens, int32_t maxNumBlocksPerReq, int32_t blockSize,
+    int32_t strideFactor, int32_t layerId, int64_t btStride0, int64_t btStride1, int64_t tiStride0, int64_t tiStride1,
+    int64_t outStride0, int64_t outStride1, cudaStream_t stream)
+{
+    if (numTokens == 0 || numTopkTokens == 0)
+    {
+        return;
+    }
+
+    constexpr int32_t kThreadsPerBlock = 256;
+    int32_t const tilesPerRow = (numTopkTokens + kThreadsPerBlock - 1) / kThreadsPerBlock;
+    dim3 const grid(numTokens, tilesPerRow);
+    dim3 const block(kThreadsPerBlock);
+
+    convertReqIndexToGlobalKernel<<<grid, block, 0, stream>>>(reqId, blockTable, tokenIndices, output, numTopkTokens,
+        maxNumBlocksPerReq, blockSize, strideFactor, layerId, btStride0, btStride1, tiStride0, tiStride1, outStride0,
+        outStride1);
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h
new file mode 100644
index 000000000000..74aa7783ee61
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h
@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+void invokeConvertReqIndexToGlobal(int32_t const* reqId, int32_t const* blockTable, int32_t const* tokenIndices,
+    int32_t* output, int32_t numTokens, int32_t numTopkTokens, int32_t maxNumBlocksPerReq, int32_t blockSize,
+    int32_t strideFactor, int32_t layerId, int64_t btStride0, int64_t btStride1, int64_t tiStride0, int64_t tiStride1,
+    int64_t outStride0, int64_t outStride1, cudaStream_t stream = 0);
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
new file mode 100644
index 000000000000..dcb5a212b319
--- /dev/null
+++ b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
@@ -0,0 +1,163 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "IndexerKCacheGather.h"
+#include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+namespace
+{
+/**
+ * Given a flat element index and tensor shape [d0, d1, d2, d3] with strides [s0, s1, s2, s3],
+ * find the actual memory offset within the given k cache pool using the strides.
+ */
+__device__ __forceinline__ int64_t flatIndexToMemoryOffset(
+    int64_t flat_idx, int32_t d0, int32_t d1, int32_t d2, int32_t d3, int64_t s0, int64_t s1, int64_t s2, int64_t s3)
+{
+    // Unravel from innermost to outermost dimension
+    int32_t i3 = flat_idx % d3;
+    flat_idx /= d3;
+
+    int32_t i2 = flat_idx % d2;
+    flat_idx /= d2;
+
+    int32_t i1 = flat_idx % d1;
+    flat_idx /= d1;
+
+    int32_t i0 = flat_idx;
+
+    // Compute memory offset using strides
+    return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3;
+}
+
+} // anonymous namespace
+
+/**
+ * CUDA kernel to gather both FP8 K values and scales from the indexer k cache pool.
+ * This is the inverse of indexerKCacheScatterUnifiedKernel.
+ *
+ * @param k_cache             Indexer k cache pool with shape [num_blocks, block_size, 1, per_token_size]
+ *                            (can be non-contiguous)
+ * @param slot_mapping_fp8    Flat element index for FP8 data start position [total_kv_len]
+ * @param slot_mapping_scale  Flat element index for scale data start position [total_kv_len]
+ * @param out_fp8             Output FP8 data [num_tokens, head_dim] contiguous
+ * @param out_scale           Output scale data [num_tokens, scale_size] contiguous
+ * @param k_token_start       Start offset into slot_mapping arrays
+ * @param num_tokens          Number of tokens to gather
+ * @param head_dim            Head dimension (must be 128)
+ * @param scale_size          Scale size in bytes (must be 4)
+ * @param cache_stride_0      Stride for k_cache dimension 0 (in bytes)
+ * @param cache_stride_1      Stride for k_cache dimension 1 (in bytes)
+ * @param cache_stride_2      Stride for k_cache dimension 2 (in bytes)
+ * @param cache_stride_3      Stride for k_cache dimension 3 (in bytes)
+ * @param cache_dim_0         Size of k_cache dimension 0
+ * @param cache_dim_1         Size of k_cache dimension 1
+ * @param cache_dim_2         Size of k_cache dimension 2
+ * @param cache_dim_3         Size of k_cache dimension 3
+ */
+__global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_cache,
+    int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale,
+    uint8_t* __restrict__ out_fp8, uint8_t* __restrict__ out_scale, int32_t k_token_start, int32_t num_tokens,
+    int32_t head_dim, int32_t scale_size, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2,
+    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3)
+{
+    // For head_dim=128, each thread handles 4 bytes/elements per read/write instruction
+    constexpr int VEC_SIZE = 4;
+
+    // Token index from block.x
+    int32_t token_idx = blockIdx.x;
+
+    if (token_idx >= num_tokens)
+    {
+        return;
+    }
+
+    // Index into slot_mapping with k_token_start offset
+    int32_t slot_idx = k_token_start + token_idx;
+
+    int64_t flat_idx_fp8_base = slot_mapping_fp8[slot_idx];
+    int64_t flat_idx_scale_base = slot_mapping_scale[slot_idx];
+
+    if (flat_idx_fp8_base < 0 || flat_idx_scale_base < 0)
+    {
+        return;
+    }
+
+    int32_t head_dim_idx = threadIdx.x * VEC_SIZE;
+    int64_t flat_idx = flat_idx_fp8_base + head_dim_idx;
+
+    // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
+    int64_t src_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
+        cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
+    int64_t dst_offset = token_idx * head_dim + head_dim_idx;
+
+    // 4 bytes read from non-contiguous cache, write to contiguous output
+    *reinterpret_cast<uint32_t*>(&out_fp8[dst_offset]) = *reinterpret_cast<uint32_t const*>(&k_cache[src_offset]);
+
+    // Only thread 0 reads the single 4 bytes scale value
+    if (threadIdx.x == 0)
+    {
+        int64_t src_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, cache_dim_0, cache_dim_1, cache_dim_2,
+            cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
+        int64_t dst_offset_scale = token_idx * scale_size; // scale_size = 4
+
+        // 4 bytes read for scale
+        *reinterpret_cast<uint32_t*>(&out_scale[dst_offset_scale])
+            = *reinterpret_cast<uint32_t const*>(&k_cache[src_offset_scale]);
+    }
+}
+
+void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mapping_fp8,
+    int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
+    int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
+    int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
+    cudaStream_t stream)
+{
+    if (num_tokens == 0)
+    {
+        return;
+    }
+
+    // Assertions for DeepSeek-V3.2 configuration
+    constexpr int32_t QUANT_BLOCK_SIZE = 128;
+    TLLM_CHECK_WITH_INFO(
+        head_dim == QUANT_BLOCK_SIZE, "head_dim must equal 128 for DeepSeek-V3 indexer cache (got %d)", head_dim);
+    TLLM_CHECK_WITH_INFO(
+        scale_size == 4, "scale_size must equal 4 bytes (1 float32 scale per token, got %d)", scale_size);
+
+    // For head_dim=128, we use 32 threads to handle 128 bytes per token and extra 4 bytes for scale
+    constexpr int32_t THREADS_PER_BLOCK = 32;
+
+    dim3 block(THREADS_PER_BLOCK);
+    dim3 grid(num_tokens);
+
+    indexerKCacheGatherUnifiedKernel<<<grid, block, 0, stream>>>(k_cache, slot_mapping_fp8, slot_mapping_scale, out_fp8,
+        out_scale, k_token_start, num_tokens, head_dim, scale_size, cache_stride_0, cache_stride_1, cache_stride_2,
+        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3);
+
+    // Check for kernel launch errors
+    TLLM_CUDA_CHECK(cudaGetLastError());
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt
index 10f2a0eee4e7..39b11feeed90 100644
--- a/cpp/tensorrt_llm/thop/CMakeLists.txt
+++ b/cpp/tensorrt_llm/thop/CMakeLists.txt
@@ -89,6 +89,7 @@ add_library(
   fp4BlockScaleMoe.cpp
   noAuxTcOp.cpp
   fusedCatFp8Op.cpp
+  IndexerKCacheGatherOp.cpp
   IndexerKCacheScatterOp.cpp
   IndexerTopKOp.cpp
   ncclCommunicatorOp.cpp
@@ -110,6 +111,7 @@ add_library(
   tinygemm2.cpp
   dsv3RopeOp.cpp
   fusedGemmAllreduceOp.cpp
+  convertReqIndexToGlobalOp.cpp
   trtllmGenQKVProcessOp.cpp)
 set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON)
 target_link_libraries(
diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
new file mode 100644
index 000000000000..27fbc33bd1be
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
@@ -0,0 +1,121 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tensorrt_llm/common/opUtils.h"
+#include "tensorrt_llm/runtime/torchUtils.h"
+
+#include "tensorrt_llm/kernels/IndexerKCacheGather.h"
+
+namespace th = torch;
+namespace tl = tensorrt_llm;
+namespace tk = tensorrt_llm::kernels;
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k_cache,
+    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t k_token_start, int64_t num_tokens)
+{
+    constexpr int32_t HEAD_DIM = 128;
+    constexpr int32_t SCALE_SIZE = 4;
+
+    auto device = k_cache.device();
+
+    // Early return for empty gather
+    if (num_tokens == 0)
+    {
+        auto k_fp8 = th::empty({0, HEAD_DIM}, th::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(device));
+        auto k_scale = th::empty({0, 1}, th::TensorOptions().dtype(torch::kFloat32).device(device));
+        return std::make_tuple(std::move(k_fp8), std::move(k_scale));
+    }
+
+    // Validate all tensors are CUDA tensors
+    TORCH_CHECK(k_cache.is_cuda() && slot_mapping_fp8.is_cuda() && slot_mapping_scale.is_cuda(),
+        "All tensors must be CUDA tensors");
+
+    // Validate tensor dimensions
+    TORCH_CHECK(slot_mapping_fp8.dim() == 1, "slot_mapping_fp8 must be a 1D Tensor");
+    TORCH_CHECK(slot_mapping_scale.dim() == 1, "slot_mapping_scale must be a 1D Tensor");
+
+    // Enforce k_cache is 4D tensor
+    TORCH_CHECK(k_cache.dim() == 4,
+        "k_cache must be a 4D Tensor [num_blocks, block_size, 1, per_token_size], got %d dimensions",
+        static_cast<int>(k_cache.dim()));
+
+    // Validate tensor dtypes
+    TORCH_CHECK(slot_mapping_fp8.scalar_type() == torch::kInt64, "slot_mapping_fp8 must be int64");
+    TORCH_CHECK(slot_mapping_scale.scalar_type() == torch::kInt64, "slot_mapping_scale must be int64");
+
+    // Validate tensors are contiguous (except k_cache which may be non-contiguous)
+    // k_cache can be non-contiguous - we handle this via strides
+    TORCH_CHECK(slot_mapping_fp8.is_contiguous(), "slot_mapping_fp8 must be contiguous");
+    TORCH_CHECK(slot_mapping_scale.is_contiguous(), "slot_mapping_scale must be contiguous");
+
+    // Validate slot_mapping has enough elements
+    TORCH_CHECK(slot_mapping_fp8.size(0) >= k_token_start + num_tokens,
+        "slot_mapping_fp8 too short for k_token_start + num_tokens");
+    TORCH_CHECK(slot_mapping_scale.size(0) >= k_token_start + num_tokens,
+        "slot_mapping_scale too short for k_token_start + num_tokens");
+
+    int32_t cache_dim_0 = static_cast<int32_t>(k_cache.size(0)); // num_blocks
+    int32_t cache_dim_1 = static_cast<int32_t>(k_cache.size(1)); // block_size
+    int32_t cache_dim_2 = static_cast<int32_t>(k_cache.size(2)); // num_kv_heads
+    int32_t cache_dim_3 = static_cast<int32_t>(k_cache.size(3)); // per_token_size
+
+    // Validation for indexer k cache pool for DeepSeek-V3.2 constraints
+    TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1 for DeepSeek-V3.2, got %d", cache_dim_2);
+
+    int64_t cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
+    int64_t cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
+    int64_t cache_stride_2 = static_cast<int64_t>(k_cache.stride(2));
+    int64_t cache_stride_3 = static_cast<int64_t>(k_cache.stride(3));
+
+    // Allocate output buffers
+    auto num_tokens_i32 = static_cast<int32_t>(num_tokens);
+    auto out_fp8 = th::empty({num_tokens, HEAD_DIM}, th::TensorOptions().dtype(torch::kUInt8).device(device));
+    auto out_scale = th::empty({num_tokens, SCALE_SIZE}, th::TensorOptions().dtype(torch::kUInt8).device(device));
+
+    auto stream = at::cuda::getCurrentCUDAStream(k_cache.get_device());
+
+    tk::invokeIndexerKCacheGather(k_cache.data_ptr<uint8_t>(), slot_mapping_fp8.data_ptr<int64_t>(),
+        slot_mapping_scale.data_ptr<int64_t>(), out_fp8.data_ptr<uint8_t>(), out_scale.data_ptr<uint8_t>(),
+        static_cast<int32_t>(k_token_start), num_tokens_i32, HEAD_DIM, SCALE_SIZE, cache_dim_0, cache_dim_1,
+        cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3, stream);
+
+    // View-cast to final dtypes (no copy)
+    auto k_fp8 = out_fp8.view(torch::kFloat8_e4m3fn);
+    auto k_scale = out_scale.view(torch::kFloat32).view({num_tokens, 1});
+
+    return std::make_tuple(std::move(k_fp8), std::move(k_scale));
+}
+
+} // namespace torch_ext
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "indexer_k_cache_gather_op(Tensor k_cache, Tensor slot_mapping_fp8, "
+        "Tensor slot_mapping_scale, int k_token_start, int num_tokens) -> (Tensor, Tensor)");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("indexer_k_cache_gather_op", &tensorrt_llm::torch_ext::indexer_k_cache_gather_op);
+}
diff --git a/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp b/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp
new file mode 100644
index 000000000000..45b862289de5
--- /dev/null
+++ b/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp
@@ -0,0 +1,84 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tensorrt_llm/common/opUtils.h"
+#include "tensorrt_llm/kernels/convertReqIndexToGlobal.h"
+#include "tensorrt_llm/runtime/torchUtils.h"
+
+namespace th = torch;
+namespace tk = tensorrt_llm::kernels;
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+th::Tensor convertReqIndexToGlobal(th::Tensor const& reqId, th::Tensor const& blockTable,
+    th::Tensor const& tokenIndices, int64_t blockSize, int64_t numTopkTokens, int64_t strideFactor, int64_t layerId)
+{
+    TORCH_CHECK(reqId.is_cuda() && blockTable.is_cuda() && tokenIndices.is_cuda(), "All tensors must be CUDA tensors");
+    TORCH_CHECK(reqId.scalar_type() == th::kInt32, "req_id must be int32");
+    TORCH_CHECK(blockTable.scalar_type() == th::kInt32, "block_table must be int32");
+    TORCH_CHECK(tokenIndices.scalar_type() == th::kInt32, "token_indices must be int32");
+
+    TORCH_CHECK(reqId.dim() == 1, "req_id must be 1D");
+    TORCH_CHECK(blockTable.dim() == 2, "block_table must be 2D");
+    TORCH_CHECK(tokenIndices.dim() == 2, "token_indices must be 2D");
+
+    // Ensure contiguous
+    auto reqIdC = reqId.contiguous();
+    auto blockTableC = blockTable.contiguous();
+    auto tokenIndicesC = tokenIndices.contiguous();
+
+    int32_t const numTokens = static_cast<int32_t>(reqIdC.size(0));
+    int32_t const maxNumBlocksPerReq = static_cast<int32_t>(blockTableC.size(1));
+
+    // Allocate output
+    auto out = th::empty_like(tokenIndicesC);
+
+    // Extract strides
+    int64_t const btStride0 = blockTableC.stride(0);
+    int64_t const btStride1 = blockTableC.stride(1);
+    int64_t const tiStride0 = tokenIndicesC.stride(0);
+    int64_t const tiStride1 = tokenIndicesC.stride(1);
+    int64_t const outStride0 = out.stride(0);
+    int64_t const outStride1 = out.stride(1);
+
+    auto stream = at::cuda::getCurrentCUDAStream(reqIdC.get_device()).stream();
+
+    tk::invokeConvertReqIndexToGlobal(reqIdC.data_ptr<int32_t>(), blockTableC.data_ptr<int32_t>(),
+        tokenIndicesC.data_ptr<int32_t>(), out.data_ptr<int32_t>(), numTokens, static_cast<int32_t>(numTopkTokens),
+        maxNumBlocksPerReq, static_cast<int32_t>(blockSize), static_cast<int32_t>(strideFactor),
+        static_cast<int32_t>(layerId), btStride0, btStride1, tiStride0, tiStride1, outStride0, outStride1, stream);
+
+    return out;
+}
+
+} // namespace torch_ext
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "convert_req_index_to_global(Tensor req_id, Tensor block_table, Tensor token_indices, int block_size, int "
+        "num_topk_tokens, int stride_factor, int layer_id) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("convert_req_index_to_global", &tensorrt_llm::torch_ext::convertReqIndexToGlobal);
+}
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index 9cc38da0f1ec..b3a5fbeb3136 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -32,9 +32,6 @@
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
-from .kernel import (triton_convert_req_index_to_global_index,
-                     triton_gather_k_cache)
-
 ModelConfig = tensorrt_llm.bindings.ModelConfig
 
 if TYPE_CHECKING:
@@ -119,59 +116,33 @@ def transform_local_topk_and_prepare_pool_view(
     layer_idx: int,
     is_generation: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Convert local topk indices to global pool indices and prepare KV pool.
-    Auto-detects stride and handles both contiguous/strided layouts.
+    """Convert local topk indices to global pool indices and prepare KV pool.
 
-    Args:
-        topk_indices: [num_tokens, NUM_TOPK]
-        attn_metadata: Metadata with block_table and request mappings
-        kv_cache_manager: KV cache manager
-        layer_idx: Layer index
-        is_generation: Generation vs context phase
-
-    Returns:
-        (global_indices, kv_pool):
-            - global_indices: [num_tokens, NUM_TOPK]
-            - kv_pool: [total_tokens, 1, head_dim]
+    Uses cached values from attn_metadata._ensure_pool_view_cached()
+    to avoid redundant Python/CUDA overhead across layers.
     """
     assert topk_indices.dtype == torch.int32
 
-    # Get all layer KV cache pool: [num_blocks, num_layers, kv_factor, blockSize]
-    kv_cache_manager = attn_metadata.kv_cache_manager
-    all_layer_kv_pool = kv_cache_manager.get_unique_primary_pool(
-    )  # [num_blocks, num_layers, kv_factor, blockSize]
-    num_blocks, num_layers, _, _ = all_layer_kv_pool.shape
-    tokens_per_block = kv_cache_manager.tokens_per_block
-    head_dim = kv_cache_manager.head_dim
-    assert all_layer_kv_pool.is_contiguous(
-    ), "all_layer_kv_pool should be contiguous"
-    all_layer_kv_pool = all_layer_kv_pool.squeeze(2).view(-1, 1, head_dim)
-    stride_factor = num_layers * tokens_per_block
-
-    # Get block_table and request indices for this phase
+    attn_metadata._ensure_pool_view_cached()
+
     if is_generation:
-        block_table = attn_metadata.block_table[
-            attn_metadata.num_contexts:attn_metadata.num_seqs]
-        req_idx = attn_metadata.req_idx_per_token[
-            attn_metadata.num_ctx_tokens:attn_metadata.num_tokens]
-        req_idx = req_idx - attn_metadata.num_contexts
+        block_table = attn_metadata._cached_block_table_gen
+        req_idx = attn_metadata._cached_req_idx_gen
     else:
-        block_table = attn_metadata.block_table[:attn_metadata.num_contexts]
-        req_idx = attn_metadata.req_idx_per_token[:attn_metadata.num_ctx_tokens]
+        block_table = attn_metadata._cached_block_table_ctx
+        req_idx = attn_metadata._cached_req_idx_ctx
 
-    # Convert to global indices
-    global_indices = triton_convert_req_index_to_global_index(
+    global_indices = torch.ops.trtllm.convert_req_index_to_global(
         req_idx,
         block_table,
         topk_indices,
-        BLOCK_SIZE=tokens_per_block,
-        NUM_TOPK_TOKENS=topk_indices.shape[1],
-        stride_factor=stride_factor,
-        layer_id=layer_idx,
+        attn_metadata._cached_tokens_per_block,
+        topk_indices.shape[1],
+        attn_metadata._cached_stride_factor,
+        layer_idx,
     )
 
-    return global_indices, all_layer_kv_pool
+    return global_indices, attn_metadata._cached_pool_view
 
 
 def split_prefill_chunks(
@@ -346,6 +317,20 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata):
 
     def __init__(self, *args, **kwargs):
         self.num_sms = tensorrt_llm.deep_gemm.get_num_sms()
+        # Cached step-invariant values for transform_local_topk_and_prepare_pool_view.
+        # These are recomputed once per step in _ensure_pool_view_cached() and
+        # reused across all layers to avoid redundant Python/CUDA overhead.
+        # Initialized here as plain instance attributes (not class-level
+        # annotations) to stay invisible to dataclass/torch.compile introspection.
+        self._pool_cache_valid = False
+        self._cached_kv_mgr_id = 0
+        self._cached_pool_view = None
+        self._cached_stride_factor = 0
+        self._cached_tokens_per_block = 0
+        self._cached_block_table_ctx = None
+        self._cached_block_table_gen = None
+        self._cached_req_idx_ctx = None
+        self._cached_req_idx_gen = None
         super().__init__(*args, **kwargs)
         if self.sparse_attention_config.indexer_max_chunk_size is not None:
             self.indexer_max_chunk_size = self.sparse_attention_config.indexer_max_chunk_size
@@ -582,40 +567,83 @@ def update_spec_dec_param(
             capture_graph = self.is_cuda_graph
             self.create_expanded_buffers(capture_graph=capture_graph)
 
+    def _invalidate_pool_view_cache(self):
+        """Invalidate the cached pool view and related step-invariant values.
+
+        Must be called at the start of each forward step (in prepare()) so that
+        _ensure_pool_view_cached() recomputes them for the new batch.
+        """
+        self._pool_cache_valid = False
+
+    def _ensure_pool_view_cached(self):
+        """Compute and cache values used by
+        transform_local_topk_and_prepare_pool_view().
+
+        These values (pool view, stride factor, block table slices, request
+        index slices) are constant across all layers sharing the same KV pool
+        and batch dimensions within a forward pass. Caching them avoids
+        redundant Python/CUDA overhead per layer.
+
+        Safety: _invalidate_pool_view_cache() is called unconditionally at the
+        start of every step (prepare() and on_update_kv_lens()), so the boolean
+        flag is always cleared before the first per-layer call within a step.
+        """
+        if self._pool_cache_valid and self._cached_kv_mgr_id == id(
+                self.kv_cache_manager):
+            return
+
+        pool = self.kv_cache_manager.get_unique_primary_pool()
+        kv_cache_manager = self.kv_cache_manager
+        num_blocks, num_layers, _, _ = pool.shape
+        self._cached_tokens_per_block = kv_cache_manager.tokens_per_block
+        head_dim = kv_cache_manager.head_dim
+        self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim)
+        self._cached_stride_factor = (num_layers *
+                                      self._cached_tokens_per_block)
+        self._cached_block_table_ctx = self.block_table[:self.num_contexts]
+        self._cached_block_table_gen = self.block_table[self.num_contexts:self.
+                                                        num_seqs]
+        self._cached_req_idx_ctx = self.req_idx_per_token[:self.num_ctx_tokens]
+        self._cached_req_idx_gen = (
+            self.req_idx_per_token[self.num_ctx_tokens:self.num_tokens] -
+            self.num_contexts)
+        self._cached_kv_mgr_id = id(kv_cache_manager)
+        self._pool_cache_valid = True
+
+    @maybe_compile(dynamic=True)
+    def _get_dense_topk_indices(self, seq_lens, kv_lens, num_tokens):
+        device = kv_lens.device
+        past_kv_lens = kv_lens - seq_lens
+        # get position ids
+        seq_ends = torch.cumsum(seq_lens, dim=0)
+        seq_starts = seq_ends - seq_lens
+        per_seq_offsets = past_kv_lens - seq_starts  # Shape: [batch_size]
+        global_indices = torch.arange(num_tokens, device=device)
+        batch_indices = torch.searchsorted(seq_ends,
+                                           global_indices,
+                                           side='right')
+        repeated_offsets = per_seq_offsets[batch_indices]
+        position_ids = global_indices + repeated_offsets
+        # get the dense topk indices with causal mask
+        range_row = torch.arange(self.sparse_mla_topk, device=device)
+        mask = range_row <= position_ids.unsqueeze(1)
+        return torch.where(mask, range_row, -1)
+
     def prepare_dense_topk_indices(self,
                                    kv_lens,
                                    device=False):  # device=False means use CPU
 
-        @maybe_compile(dynamic=True)
-        def _get_dense_topk_indices(seq_lens, kv_lens, num_tokens):
-            device = kv_lens.device
-            past_kv_lens = kv_lens - seq_lens
-            # get position ids
-            seq_ends = torch.cumsum(seq_lens, dim=0)
-            seq_starts = seq_ends - seq_lens
-            per_seq_offsets = past_kv_lens - seq_starts  # Shape: [batch_size]
-            global_indices = torch.arange(num_tokens, device=device)
-            batch_indices = torch.searchsorted(seq_ends,
-                                               global_indices,
-                                               side='right')
-            repeated_offsets = per_seq_offsets[batch_indices]
-            position_ids = global_indices + repeated_offsets
-            # get the dense topk indices with causal mask
-            range_row = torch.arange(self.sparse_mla_topk, device=device)
-            mask = range_row <= position_ids.unsqueeze(1)
-            return torch.where(mask, range_row, -1)
-
         if self.num_contexts > 0 and self.skip_indexer_for_ctx_reqs:
             ctx_range = slice(self.num_ctx_tokens)
             if device:
                 self.topk_indices_buffer[ctx_range, :].copy_(
-                    _get_dense_topk_indices(
+                    self._get_dense_topk_indices(
                         self.seq_lens_cuda[:self.num_contexts],
                         kv_lens[:self.num_contexts], self.num_ctx_tokens),
                     non_blocking=True)
             else:
                 self.host_topk_indices_buffer[
-                    ctx_range, :] = _get_dense_topk_indices(
+                    ctx_range, :] = self._get_dense_topk_indices(
                         self.seq_lens[:self.num_contexts],
                         kv_lens[:self.num_contexts], self.num_ctx_tokens)
                 self.topk_indices_buffer[ctx_range, :].copy_(
@@ -626,14 +654,14 @@ def _get_dense_topk_indices(seq_lens, kv_lens, num_tokens):
             gen_range = slice(self.num_ctx_tokens, self.num_tokens)
             if device:
                 self.topk_indices_buffer[gen_range, :].copy_(
-                    _get_dense_topk_indices(
+                    self._get_dense_topk_indices(
                         self.seq_lens_cuda[self.num_contexts:self.num_seqs],
                         kv_lens[self.num_contexts:self.num_seqs],
                         self.num_tokens - self.num_ctx_tokens),
                     non_blocking=True)
             else:
                 self.host_topk_indices_buffer[
-                    gen_range, :] = _get_dense_topk_indices(
+                    gen_range, :] = self._get_dense_topk_indices(
                         self.seq_lens[self.num_contexts:self.num_seqs],
                         kv_lens[self.num_contexts:self.num_seqs],
                         self.num_tokens - self.num_ctx_tokens)
@@ -671,6 +699,7 @@ def _get_pool_block_indices(self) -> torch.Tensor:
 
     def prepare(self):
         super().prepare()
+        self._invalidate_pool_view_cache()
 
         # Get kv lengths
         assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True"
@@ -856,6 +885,13 @@ def on_update_kv_lens(self):
         # (inside _preprocess_inputs) to account for variable accepted tokens. The indexer
         # slot_mapping_* buffers also depend on these effective cached lengths. If we do not
         # refresh slot mappings here, indexer K-cache updates can be written with stale offsets.
+
+        # _preprocess_inputs() also uses this as a general hook to "invalidate per-forward-pass
+        # caches so they are recomputed (and captured) on every _forward_step". Invalidate the
+        # pool_view cache here so it is recomputed on the next
+        # transform_local_topk_and_prepare_pool_view() call.
+        self._invalidate_pool_view_cache()
+
         if self.kv_cache_manager is not None and self.num_tokens > 0:
             seq_lens = self.seq_lens_cuda[:self.num_seqs]
             # Runtime cached lengths after overlap/spec-dec correction.
@@ -1429,23 +1465,15 @@ def sparse_attn_indexer(
                     tp_rank = metadata.mapping.tp_rank
                     tp_size = metadata.mapping.tp_size
 
-                # Use the 2D pool data directly (contiguous) instead of the
-                # 4D view, because the 4D view may have strides that
-                # prevent flattening via .view(-1).
-                layer_offset = metadata.kv_cache_manager.layer_offsets[
-                    self.layer_idx]
-                gather_k_cache_pool = metadata.kv_cache_manager.indexer_k_cache_pool_per_layer[
-                    layer_offset]
+                k_cache_4d = metadata.kv_cache_manager.get_indexer_k_cache_buffers(
+                    self.layer_idx)
 
                 for chunk in metadata.indexer_prefill_chunks:
-                    chunk_k_fp8, chunk_k_scale = triton_gather_k_cache(
-                        gather_k_cache_pool,
-                        metadata.slot_mapping_fp8_fullkv,
-                        metadata.slot_mapping_scale_fullkv,
-                        chunk.k_token_start,
-                        chunk.k_token_end,
-                        self.head_dim,
-                    )
+                    num_k_tokens = chunk.k_token_end - chunk.k_token_start
+                    chunk_k_fp8, chunk_k_scale = torch.ops.trtllm.indexer_k_cache_gather_op(
+                        k_cache_4d, metadata.slot_mapping_fp8_fullkv,
+                        metadata.slot_mapping_scale_fullkv, chunk.k_token_start,
+                        num_k_tokens)
 
                     chunk_num_token = chunk.token_end - chunk.token_start
                     apply_q_split = q_split_eligible and chunk_num_token >= q_split_threshold
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py
index 9daec9b0add0..6412610064fa 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py
@@ -1799,12 +1799,12 @@ def _convert_req_index_to_global_index_kernel_with_stride_factor(
     block_table_ptr,  # int32 [num_requests, max_num_blocks_per_req]
     token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
     out_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
-    # shapes (compile-time where possible)
-    max_num_blocks_per_req: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-    BLOCK_N: tl.constexpr,  # tile width along columns # strides (in elements)
-    stride_factor: tl.constexpr,  # for strided memory layout adjustment
-    layer_id: tl.constexpr,  # for layer interleaving layout
+    # shapes
+    max_num_blocks_per_req,
+    BLOCK_SIZE,
+    BLOCK_N: tl.constexpr,  # tile width along columns (used in tl.arange)
+    stride_factor,  # for strided memory layout adjustment
+    layer_id,  # for layer interleaving layout
     bt_stride0,
     bt_stride1,
     ti_stride0,
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index 21b7de23c373..ff2ca7b65017 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1105,3 +1105,17 @@ def _(seq_q_offsets, seq_kv_offsets, padding_offsets, tokens_info,
           rotary_embedding_scale, rotary_embedding_base, rotary_embedding_dim,
           rotary_scaling_type, rotary_embedding_max_positions):
         return True
+
+    @torch.library.register_fake("trtllm::convert_req_index_to_global")
+    def _(req_id: torch.Tensor, block_table: torch.Tensor,
+          token_indices: torch.Tensor, block_size: int, num_topk_tokens: int,
+          stride_factor: int, layer_id: int) -> torch.Tensor:
+        return torch.empty_like(token_indices)
+
+    @torch.library.register_fake("trtllm::indexer_k_cache_gather_op")
+    def _(k_cache: torch.Tensor, slot_mapping_fp8: torch.Tensor,
+          slot_mapping_scale: torch.Tensor, k_token_start: int,
+          num_tokens: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        k_fp8 = k_cache.new_empty([num_tokens, 128], dtype=torch.float8_e4m3fn)
+        k_scale = k_cache.new_empty([num_tokens, 1], dtype=torch.float32)
+        return k_fp8, k_scale
diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py
index 1ac648cadf52..cda368e972cc 100644
--- a/tensorrt_llm/_torch/speculative/mtp.py
+++ b/tensorrt_llm/_torch/speculative/mtp.py
@@ -1132,6 +1132,13 @@ def update_draft_tokens(self, next_draft_tokens, new_draft_token,
         position_ids = inputs["position_ids"][gather_ids] + 1
         return hidden_states, position_ids
 
+    @torch.compile(options={"max-autotune": True})
+    def prepare_position_ids_and_last_tokens(self, position_ids, seq_lens_cuda):
+        position_ids = position_ids.squeeze(0)
+        last_tokens_idx = torch.cumsum(seq_lens_cuda, dim=0,
+                                       dtype=torch.long) - 1
+        return position_ids, last_tokens_idx
+
     def forward(
         self,
         input_ids,
@@ -1167,16 +1174,8 @@ def forward(
         # Save the old attn_metadata and spec_metadata
         self._prepare_attn_metadata_for_spec_dec(attn_metadata)
 
-        # Prepare inputs for the 1st MTP layer
-        @torch.compile(options={"max-autotune": True})
-        def prepare_position_ids_and_last_tokens(position_ids, attn_metadata):
-            position_ids = position_ids.squeeze(0)
-            last_tokens_idx = torch.cumsum(
-                attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
-            return position_ids, last_tokens_idx
-
-        position_ids, last_tokens_idx = prepare_position_ids_and_last_tokens(
-            position_ids, attn_metadata)
+        position_ids, last_tokens_idx = self.prepare_position_ids_and_last_tokens(
+            position_ids, attn_metadata.seq_lens_cuda)
         inputs = self.prepare_drafter_inputs(input_ids=input_ids,
                                              position_ids=position_ids,
                                              last_tokens_idx=last_tokens_idx,
@@ -1301,12 +1300,9 @@ def prepare_position_ids_and_last_tokens(position_ids, attn_metadata):
                     # as draft model only infer 1 token for the subsequent inference.
                     attn_metadata.use_spec_decoding = False
                 elif hasattr(attn_metadata, 'kv_lens_cuda'):
+                    # update kv_lens_cuda
+                    attn_metadata.kv_lens_cuda[:batch_size] += 1
 
-                    @torch.compile(options={"max-autotune": True})
-                    def update_kv_lens(kv_lens_cuda, batch_size):
-                        kv_lens_cuda[:batch_size] += 1
-
-                    update_kv_lens(attn_metadata.kv_lens_cuda, batch_size)
                     # update metadata
                     # some attention metadata needs to be updated when changing kv_lens
                     attn_metadata.update_for_spec_dec()
diff --git a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
new file mode 100644
index 000000000000..fd1a00b34c62
--- /dev/null
+++ b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
@@ -0,0 +1,421 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for DSA C++ custom ops:
+
+- ``torch.ops.trtllm.indexer_k_cache_gather_op``
+- ``torch.ops.trtllm.convert_req_index_to_global``
+"""
+
+import pytest
+import torch
+
+# Import tensorrt_llm to load C++ custom operators
+import tensorrt_llm  # noqa: F401
+
+# ---------------------------------------------------------------------------
+# Constants matching the C++ kernel (DeepSeek-V3.2 indexer config)
+# ---------------------------------------------------------------------------
+HEAD_DIM = 128
+SCALE_BYTES = 4
+BYTES_PER_TOKEN = HEAD_DIM + SCALE_BYTES
+
+
+# ===================================================================
+# Test 1: indexer_k_cache_gather_op
+# ===================================================================
+
+
+def _reference_indexer_k_cache_gather(
+    k_cache: torch.Tensor,
+    slot_mapping_fp8: torch.Tensor,
+    slot_mapping_scale: torch.Tensor,
+    k_token_start: int,
+    num_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Python reference for indexer_k_cache_gather_op.
+
+    The C++ op treats the 4D k_cache as a flat byte buffer (via strides)
+    and gathers HEAD_DIM bytes for FP8 data and SCALE_BYTES bytes for scales
+    at positions given by the slot mappings.
+    """
+    device = k_cache.device
+
+    if num_tokens == 0:
+        return (
+            torch.empty(0, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device),
+            torch.empty(0, 1, dtype=torch.float32, device=device),
+        )
+
+    # Flatten the cache to a 1D byte view using a contiguous copy
+    k_cache_flat = k_cache.contiguous().reshape(-1)
+
+    fp8_bases = slot_mapping_fp8[k_token_start : k_token_start + num_tokens]
+    byte_offsets_fp8 = torch.arange(HEAD_DIM, device=device, dtype=torch.int64)
+    gather_fp8 = fp8_bases.unsqueeze(1) + byte_offsets_fp8.unsqueeze(0)
+    out_fp8 = k_cache_flat[gather_fp8]
+
+    scale_bases = slot_mapping_scale[k_token_start : k_token_start + num_tokens]
+    byte_offsets_scale = torch.arange(SCALE_BYTES, device=device, dtype=torch.int64)
+    gather_scale = scale_bases.unsqueeze(1) + byte_offsets_scale.unsqueeze(0)
+    out_scale = k_cache_flat[gather_scale]
+
+    k_fp8 = out_fp8.view(torch.float8_e4m3fn)
+    k_scale = out_scale.view(torch.float32).view(num_tokens, 1)
+    return k_fp8, k_scale
+
+
+def _create_4d_cache_and_mappings(
+    total_kv_len: int,
+    num_blocks: int,
+    block_size: int,
+    per_token_size: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Build a random 4D k_cache with valid, non-overlapping slot mappings.
+
+    k_cache shape: [num_blocks, block_size, 1, per_token_size]
+    Each token occupies ``per_token_size`` bytes at a unique (block, slot)
+    position. Slot mappings point into the *contiguous flat* byte layout.
+    """
+    k_cache = torch.randint(
+        0,
+        256,
+        (num_blocks, block_size, 1, per_token_size),
+        dtype=torch.uint8,
+        device=device,
+    )
+
+    total_slots = num_blocks * block_size
+    assert total_kv_len <= total_slots, (
+        f"total_kv_len ({total_kv_len}) exceeds total slots ({total_slots})"
+    )
+
+    # Pick random unique (block, slot) positions
+    perm = torch.randperm(total_slots, device=device)[:total_kv_len]
+    # Flat byte offset of each token in the contiguous view
+    # Each token starts at perm[i] * per_token_size (within the [total_slots, 1, per_token_size] view)
+    # But the cache is [num_blocks, block_size, 1, per_token_size], so flat offset is:
+    # perm[i] * (1 * per_token_size) since dim2=1
+    flat_starts = perm.to(torch.int64) * per_token_size
+
+    slot_mapping_fp8 = flat_starts
+    slot_mapping_scale = flat_starts + HEAD_DIM
+
+    return k_cache, slot_mapping_fp8, slot_mapping_scale
+
+
+@pytest.mark.parametrize(
+    "total_kv_len,num_blocks,block_size,k_token_start,num_tokens",
+    [
+        # Gather all tokens
+        (64, 8, 8, 0, 64),
+        # Sub-range from the middle
+        (128, 16, 8, 32, 64),
+        # Single token
+        (10, 2, 8, 3, 1),
+        # Larger test
+        (512, 64, 8, 100, 300),
+        # Non-power-of-2 token count
+        (100, 16, 8, 10, 37),
+        # Very small
+        (3, 1, 4, 0, 3),
+        # Larger block_size
+        (64, 4, 16, 0, 64),
+    ],
+)
+def test_indexer_k_cache_gather_contiguous(
+    total_kv_len,
+    num_blocks,
+    block_size,
+    k_token_start,
+    num_tokens,
+):
+    """Test C++ gather op on a contiguous 4D cache against Python reference."""
+    device = torch.device("cuda")
+    per_token_size = BYTES_PER_TOKEN  # 132
+
+    k_cache, slot_fp8, slot_scale = _create_4d_cache_and_mappings(
+        total_kv_len, num_blocks, block_size, per_token_size, device
+    )
+
+    # C++ op
+    cpp_fp8, cpp_scale = torch.ops.trtllm.indexer_k_cache_gather_op(
+        k_cache, slot_fp8, slot_scale, k_token_start, num_tokens
+    )
+
+    # Reference
+    ref_fp8, ref_scale = _reference_indexer_k_cache_gather(
+        k_cache, slot_fp8, slot_scale, k_token_start, num_tokens
+    )
+
+    assert cpp_fp8.shape == ref_fp8.shape
+    assert cpp_scale.shape == ref_scale.shape
+    assert torch.equal(cpp_fp8.view(torch.uint8), ref_fp8.view(torch.uint8)), "FP8 data mismatch"
+    assert torch.equal(cpp_scale.view(torch.uint8), ref_scale.view(torch.uint8)), (
+        "Scale data mismatch"
+    )
+
+
+@pytest.mark.parametrize(
+    "total_kv_len,num_blocks,block_size,k_token_start,num_tokens",
+    [
+        (64, 8, 8, 0, 64),
+        (128, 16, 8, 32, 64),
+        (512, 64, 8, 100, 300),
+    ],
+)
+def test_indexer_k_cache_gather_noncontiguous(
+    total_kv_len,
+    num_blocks,
+    block_size,
+    k_token_start,
+    num_tokens,
+):
+    """Test C++ gather op on a non-contiguous 4D cache.
+
+    The real KV cache pool is typically a large buffer viewed with
+    non-contiguous strides (layer interleaving). Simulate this by
+    allocating a larger buffer and taking a strided view.
+    """
+    device = torch.device("cuda")
+    per_token_size = BYTES_PER_TOKEN
+    num_layers = 3  # simulate multi-layer interleaving
+    target_layer = 1
+
+    # Allocate the full interleaved pool:
+    # [num_blocks, block_size, num_layers, per_token_size]
+    full_pool = torch.randint(
+        0,
+        256,
+        (num_blocks, block_size, num_layers, per_token_size),
+        dtype=torch.uint8,
+        device=device,
+    )
+
+    # Select a single layer → shape [num_blocks, block_size, 1, per_token_size]
+    # This slice is non-contiguous because stride(1) spans across all layers.
+    k_cache_nc = full_pool[:, :, target_layer : target_layer + 1, :]
+    assert not k_cache_nc.is_contiguous(), "Expected non-contiguous cache"
+
+    # Build slot mappings based on the *contiguous* copy (what the reference sees)
+    k_cache_contig = k_cache_nc.contiguous()
+    total_slots = num_blocks * block_size
+    perm = torch.randperm(total_slots, device=device)[:total_kv_len]
+    flat_starts = perm.to(torch.int64) * per_token_size
+    slot_fp8 = flat_starts
+    slot_scale = flat_starts + HEAD_DIM
+
+    # C++ op (handles non-contiguous strides internally)
+    cpp_fp8, cpp_scale = torch.ops.trtllm.indexer_k_cache_gather_op(
+        k_cache_nc, slot_fp8, slot_scale, k_token_start, num_tokens
+    )
+
+    # Reference uses contiguous copy
+    ref_fp8, ref_scale = _reference_indexer_k_cache_gather(
+        k_cache_contig, slot_fp8, slot_scale, k_token_start, num_tokens
+    )
+
+    assert cpp_fp8.shape == ref_fp8.shape
+    assert cpp_scale.shape == ref_scale.shape
+    assert torch.equal(cpp_fp8.view(torch.uint8), ref_fp8.view(torch.uint8)), (
+        "FP8 data mismatch (non-contiguous)"
+    )
+    assert torch.equal(cpp_scale.view(torch.uint8), ref_scale.view(torch.uint8)), (
+        "Scale data mismatch (non-contiguous)"
+    )
+
+
+def test_indexer_k_cache_gather_empty():
+    """Zero-length gather should return correctly shaped empty tensors."""
+    device = torch.device("cuda")
+    k_cache = torch.randint(0, 256, (4, 8, 1, BYTES_PER_TOKEN), dtype=torch.uint8, device=device)
+    slot_fp8 = torch.zeros(10, dtype=torch.int64, device=device)
+    slot_scale = torch.zeros(10, dtype=torch.int64, device=device)
+
+    k_fp8, k_scale = torch.ops.trtllm.indexer_k_cache_gather_op(
+        k_cache, slot_fp8, slot_scale, k_token_start=5, num_tokens=0
+    )
+
+    assert k_fp8.shape == (0, HEAD_DIM)
+    assert k_scale.shape == (0, 1)
+    assert k_fp8.dtype == torch.float8_e4m3fn
+    assert k_scale.dtype == torch.float32
+
+
+# ===================================================================
+# Test 2: convert_req_index_to_global
+# ===================================================================
+
+
+def _reference_convert_req_index_to_global(
+    req_id: torch.Tensor,
+    block_table: torch.Tensor,
+    token_indices: torch.Tensor,
+    block_size: int,
+    stride_factor: int,
+    layer_id: int,
+) -> torch.Tensor:
+    """Python reference for convert_req_index_to_global.
+
+    For each (token_id, indice_id):
+        tok = token_indices[token_id, indice_id]
+        if tok == -1:
+            out = -1
+        else:
+            block_id = tok // block_size
+            inblock_off = tok % block_size + layer_id * block_size
+            req = req_id[token_id]
+            base = block_table[req, block_id]
+            if block_id >= max_blocks_per_req or base < 0:
+                out = -1
+            else:
+                out = base * stride_factor + inblock_off
+    """
+    num_tokens, num_topk = token_indices.shape
+    max_blocks_per_req = block_table.shape[1]
+
+    out = torch.empty_like(token_indices)
+    # Move to CPU for the reference loop
+    req_id_cpu = req_id.cpu()
+    block_table_cpu = block_table.cpu()
+    token_indices_cpu = token_indices.cpu()
+    out_cpu = out.cpu()
+
+    for i in range(num_tokens):
+        req = req_id_cpu[i].item()
+        for j in range(num_topk):
+            tok = token_indices_cpu[i, j].item()
+            if tok < 0:
+                out_cpu[i, j] = -1
+                continue
+            block_id = tok // block_size
+            inblock_off = tok % block_size + layer_id * block_size
+            if block_id >= max_blocks_per_req:
+                out_cpu[i, j] = -1
+                continue
+            base = block_table_cpu[req, block_id].item()
+            if base < 0:
+                out_cpu[i, j] = -1
+                continue
+            out_cpu[i, j] = base * stride_factor + inblock_off
+
+    return out_cpu.to(token_indices.device)
+
+
+@pytest.mark.parametrize(
+    "num_tokens,num_requests,num_topk,block_size,stride_factor,layer_id",
+    [
+        # Basic case: contiguous pool (stride_factor == block_size), layer 0
+        (4, 2, 128, 64, 64, 0),
+        # Layer interleaving: stride_factor > block_size
+        (4, 2, 128, 64, 192, 1),
+        # Single token
+        (1, 1, 128, 64, 64, 0),
+        # Larger batch
+        (32, 8, 256, 64, 64, 0),
+        # Small block_size
+        (8, 4, 128, 16, 16, 0),
+        # layer_id > 0 with contiguous pool
+        (4, 2, 128, 64, 64, 2),
+        # Large topk
+        (4, 2, 2048, 64, 192, 1),
+    ],
+)
+def test_convert_req_index_to_global(
+    num_tokens,
+    num_requests,
+    num_topk,
+    block_size,
+    stride_factor,
+    layer_id,
+):
+    """Test C++ op against Python reference."""
+    device = torch.device("cuda")
+    torch.manual_seed(42)
+
+    max_kv_len = 4096
+    max_blocks_per_req = (max_kv_len + block_size - 1) // block_size
+
+    # req_id: which request each token belongs to
+    req_id = torch.randint(0, num_requests, (num_tokens,), dtype=torch.int32, device=device)
+
+    # block_table: maps (request, block_id) → physical block index
+    # Use positive values; some entries can be -1 (padding)
+    block_table = torch.randint(
+        0, 1000, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device
+    )
+    # Add some padding (-1) in later blocks
+    block_table[:, max_blocks_per_req // 2 :] = -1
+
+    # token_indices: request-local token positions, some -1 (invalid)
+    max_valid_token = (max_blocks_per_req // 2) * block_size - 1
+    token_indices = torch.randint(
+        0, max(max_valid_token, 1), (num_tokens, num_topk), dtype=torch.int32, device=device
+    )
+    # Sprinkle -1s for invalid tokens (~20%)
+    invalid_mask = torch.rand(num_tokens, num_topk, device=device) < 0.2
+    token_indices[invalid_mask] = -1
+
+    # C++ op
+    cpp_out = torch.ops.trtllm.convert_req_index_to_global(
+        req_id, block_table, token_indices, block_size, num_topk, stride_factor, layer_id
+    )
+
+    # Reference
+    ref_out = _reference_convert_req_index_to_global(
+        req_id, block_table, token_indices, block_size, stride_factor, layer_id
+    )
+
+    assert cpp_out.shape == ref_out.shape
+    assert torch.equal(cpp_out, ref_out), (
+        f"Mismatch found.\n"
+        f"  First diff at: {(cpp_out != ref_out).nonzero(as_tuple=False)[0].tolist()}\n"
+        f"  cpp: {cpp_out[(cpp_out != ref_out)][:5]}\n"
+        f"  ref: {ref_out[(cpp_out != ref_out)][:5]}"
+    )
+
+
+def test_convert_req_index_to_global_all_invalid():
+    """All token_indices are -1 → entire output should be -1."""
+    device = torch.device("cuda")
+    num_tokens, num_topk, block_size = 4, 128, 64
+
+    req_id = torch.zeros(num_tokens, dtype=torch.int32, device=device)
+    block_table = torch.ones(1, 16, dtype=torch.int32, device=device) * 100
+    token_indices = torch.full((num_tokens, num_topk), -1, dtype=torch.int32, device=device)
+
+    out = torch.ops.trtllm.convert_req_index_to_global(
+        req_id, block_table, token_indices, block_size, num_topk, block_size, 0
+    )
+
+    assert (out == -1).all(), "All outputs should be -1 for all-invalid input"
+
+
+def test_convert_req_index_to_global_block_table_padding():
+    """block_table entries of -1 should produce -1 in output."""
+    device = torch.device("cuda")
+    num_tokens, num_topk, block_size = 2, 64, 16
+
+    req_id = torch.zeros(num_tokens, dtype=torch.int32, device=device)
+    # All block_table entries are -1 (no valid blocks)
+    block_table = torch.full((1, 8), -1, dtype=torch.int32, device=device)
+    # Valid-looking token indices that would map to these blocks
+    token_indices = torch.randint(0, 128, (num_tokens, num_topk), dtype=torch.int32, device=device)
+
+    out = torch.ops.trtllm.convert_req_index_to_global(
+        req_id, block_table, token_indices, block_size, num_topk, block_size, 0
+    )
+
+    assert (out == -1).all(), "All outputs should be -1 when block_table is all padding"
diff --git a/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py b/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py
index b21b481d62a0..ed0ff4a10357 100644
--- a/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py
+++ b/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py
@@ -149,31 +149,3 @@ def test_triton_gather_k_cache_empty():
     assert k_scale.shape == (0, 1)
     assert k_fp8.dtype == torch.float8_e4m3fn
     assert k_scale.dtype == torch.float32
-
-
-if __name__ == "__main__":
-    print("\n" + "=" * 80)
-    print("Testing Triton Gather K Cache Kernel")
-    print("=" * 80)
-
-    print("\n--- Empty tokens ---")
-    test_triton_gather_k_cache_empty()
-
-    print("\n--- Gather all tokens ---")
-    test_triton_gather_k_cache(64, 0, 64, 128, 16)
-
-    print("\n--- Sub-range from middle ---")
-    test_triton_gather_k_cache(128, 32, 96, 128, 32)
-
-    print("\n--- Single token ---")
-    test_triton_gather_k_cache(10, 3, 4, 128, 4)
-
-    print("\n--- Larger test ---")
-    test_triton_gather_k_cache(512, 100, 400, 128, 64)
-
-    print("\n--- Non-power-of-2 tokens ---")
-    test_triton_gather_k_cache(100, 10, 47, 128, 32)
-
-    print("\n" + "=" * 80)
-    print("All tests passed!")
-    print("=" * 80)

From 9c7984e7ac7b862b6abe9bef8de25ed72da6101c Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Thu, 2 Apr 2026 20:37:54 +0800
Subject: [PATCH 09/14] =?UTF-8?q?[https://nvbugs/6029220][fix]=20Disable?=
 =?UTF-8?q?=20multi-stream=20in=20maybe=5Fexecute=5Fi=E2=80=A6=20(#12692)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_deepseekv3.py     |  8 ++++++--
 tensorrt_llm/_torch/models/modeling_exaone_moe.py     |  1 +
 tensorrt_llm/_torch/models/modeling_glm.py            |  2 ++
 tensorrt_llm/_torch/models/modeling_llama.py          |  7 ++++++-
 .../_torch/models/modeling_llama_min_latency.py       |  7 ++++++-
 tensorrt_llm/_torch/models/modeling_nemotron_h.py     |  1 +
 tensorrt_llm/_torch/models/modeling_qwen3_next.py     |  2 ++
 tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py     |  6 ++++--
 tensorrt_llm/_torch/modules/multi_stream_utils.py     | 11 +++++++++--
 9 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index 138c6b8af6d7..8b2da5404489 100755
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -1159,9 +1159,12 @@ def _compute_routed_output():
         # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames
 
         routed_output, shared_output = maybe_execute_in_parallel(
-            _compute_routed_output, _compute_shared_output,
+            _compute_routed_output,
+            _compute_shared_output,
             self.event_dict[EventType.Main],
-            self.event_dict[EventType.MoeShared], self.aux_stream)
+            self.event_dict[EventType.MoeShared],
+            self.aux_stream,
+            disable_on_compile=True)
 
         if not do_finalize:
             return [shared_output, *routed_output]
@@ -1644,6 +1647,7 @@ def norm_hidden():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
         # Split hidden_states columnwise based on TP
diff --git a/tensorrt_llm/_torch/models/modeling_exaone_moe.py b/tensorrt_llm/_torch/models/modeling_exaone_moe.py
index fe420178558e..5a796f926c4d 100644
--- a/tensorrt_llm/_torch/models/modeling_exaone_moe.py
+++ b/tensorrt_llm/_torch/models/modeling_exaone_moe.py
@@ -542,6 +542,7 @@ def norm_hidden():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
         # Split hidden_states columnwise based on TP
diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py
index 3698c273afd6..c6a635d66b3e 100644
--- a/tensorrt_llm/_torch/models/modeling_glm.py
+++ b/tensorrt_llm/_torch/models/modeling_glm.py
@@ -432,6 +432,7 @@ def _compute_routed_output():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
 
         if not do_finalize:
@@ -864,6 +865,7 @@ def norm_hidden():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
         # Split hidden_states columnwise based on TP
diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py
index 743e0b8ef502..aab9d3992c66 100644
--- a/tensorrt_llm/_torch/models/modeling_llama.py
+++ b/tensorrt_llm/_torch/models/modeling_llama.py
@@ -338,7 +338,12 @@ def forward(
         fn1 = lambda: self.compute_routed_output(
             hidden_states, all_rank_num_tokens, cutlass_min_latency_mode)
         shared_output, routed_output = maybe_execute_in_parallel(
-            fn0, fn1, self.moe_event[0], self.moe_event[1], self.aux_stream)
+            fn0,
+            fn1,
+            self.moe_event[0],
+            self.moe_event[1],
+            self.aux_stream,
+            disable_on_compile=True)
         if cutlass_min_latency_mode:
             return [shared_output, *routed_output]
 
diff --git a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py
index ae3d1601fbe5..a0c45ca11277 100644
--- a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py
+++ b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py
@@ -612,7 +612,12 @@ def forward(
         fn1 = lambda: self.compute_routed_output(
             hidden_states, all_rank_num_tokens, hidden_states_high)
         shared_output, routed_output = maybe_execute_in_parallel(
-            fn0, fn1, self.moe_event[0], self.moe_event[1], self.aux_stream)
+            fn0,
+            fn1,
+            self.moe_event[0],
+            self.moe_event[1],
+            self.aux_stream,
+            disable_on_compile=True)
 
         assert shared_output.size() == routed_output.size(
         ), f'unmatched tensor shape'
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index ca19d1da4fd3..e416edd450c9 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -322,6 +322,7 @@ def _compute_routed_output():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream_shared,
+            disable_on_compile=True,
         )
 
         final_hidden_states = shared_output + routed_output
diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py
index 68e9c5a70f5f..cf17233cf075 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py
@@ -254,6 +254,7 @@ def _compute_shared_output():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         if not do_finalize:
             return final_hidden_states
@@ -808,6 +809,7 @@ def _compute_projected_states_ba():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.Attention],
             self.aux_stream,
+            disable_on_compile=True,
         )
 
         # Use fused kernel when possible to avoid elementwise ops
diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py
index 42c1df6a5d4f..e0028af9b26f 100644
--- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py
+++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py
@@ -421,10 +421,12 @@ def convert_dt():
             # convert will go second and we lose PDL, but we're using cuda
             # graphs for low latency so that seems ok.
             # If any of the contiguous calls below actually fire, that also breaks PDL.
-            xbc_d, dt_d = maybe_execute_in_parallel(conv1d, convert_dt,
+            xbc_d, dt_d = maybe_execute_in_parallel(conv1d,
+                                                    convert_dt,
                                                     self.events[0],
                                                     self.events[1],
-                                                    self.aux_steram)
+                                                    self.aux_steram,
+                                                    disable_on_compile=True)
 
             x_d, B_d, C_d = torch.split(
                 xbc_d,
diff --git a/tensorrt_llm/_torch/modules/multi_stream_utils.py b/tensorrt_llm/_torch/modules/multi_stream_utils.py
index c7b58c0896bb..f3e60a208646 100644
--- a/tensorrt_llm/_torch/modules/multi_stream_utils.py
+++ b/tensorrt_llm/_torch/modules/multi_stream_utils.py
@@ -37,7 +37,8 @@ def maybe_execute_in_parallel(
         fn1: Callable,
         event0: torch.cuda.Event,
         event1: torch.cuda.Event,
-        aux_stream: Optional[torch.cuda.Stream] = None) -> tuple[Any, Any]:
+        aux_stream: Optional[torch.cuda.Stream] = None,
+        disable_on_compile: bool = False) -> tuple[Any, Any]:
     """Utility function to run two functions in two cuda streams in parallel. Multi-stream is
     only enabled when cuda graph is turned on because switch stream has extra host overhead.
 
@@ -52,12 +53,18 @@ def maybe_execute_in_parallel(
         event1 (torch.cuda.Event): cuda event for fn1
         aux_stream (Optional[torch.cuda.Stream]): the second cuda stream for fn1.
             Multi-stream is disabled when aux_stream is None.
+        disable_on_compile (bool): if True, disable multi-stream when
+            torch.compile is tracing. Callers that are not inside a custom op
+            should set this to True so that stream/event ops are not captured
+            by dynamo. Callers inside custom ops (e.g. attention, MoE) should
+            leave this as False since custom ops are opaque to the compiler.
 
     Returns:
         tuple[Any, Any]: the return values of fn0() and fn1()
     """
 
-    multi_stream = do_multi_stream() and aux_stream is not None
+    multi_stream = (do_multi_stream() and aux_stream is not None and
+                    not (disable_on_compile and torch.compiler.is_compiling()))
 
     if multi_stream:
         event0.record()

From ca4ace150355e999a419a43155d9abd60275ea9e Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Thu, 2 Apr 2026 22:30:53 -0700
Subject: [PATCH 10/14] [https://nvbugs/5983390][perf] Reduce host overhead in
 DSA MLA attention path (#12691)

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 cpp/tensorrt_llm/nanobind/thop/bindings.cpp   |   4 +-
 .../thop/IndexerKCacheScatterOp.cpp           | 101 +++++++++---------
 cpp/tensorrt_llm/thop/attentionOp.cpp         |  18 +---
 cpp/tensorrt_llm/thop/attentionOp.h           |   2 +-
 .../_torch/attention_backend/sparse/dsa.py    |  32 ++----
 .../_torch/attention_backend/trtllm.py        |  10 +-
 .../_torch/attention_backend/trtllm_gen.py    |  23 +---
 .../custom_ops/attention/trtllm_attention.py  |   9 ++
 .../attention/sparse/test_dsa_indexer.py      |   9 +-
 9 files changed, 90 insertions(+), 118 deletions(-)

diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp
index fc161ab4a6ca..b71c39d40874 100644
--- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp
+++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp
@@ -70,8 +70,8 @@ void initBindings(nb::module_& m)
         nb::arg("cu_kv_seqlens") = std::nullopt, nb::arg("fmha_scheduler_counter") = std::nullopt,
         nb::arg("mla_bmm1_scale") = std::nullopt, nb::arg("mla_bmm2_scale") = std::nullopt,
         nb::arg("quant_q_buffer") = std::nullopt, nb::arg("flash_mla_tile_scheduler_metadata") = std::nullopt,
-        nb::arg("flash_mla_num_splits") = std::nullopt, "Multi-head attention operation",
-        nb::call_guard<nb::gil_scoped_release>());
+        nb::arg("flash_mla_num_splits") = std::nullopt, nb::arg("num_contexts") = 0, nb::arg("num_ctx_tokens") = 0,
+        "Multi-head attention operation", nb::call_guard<nb::gil_scoped_release>());
 
     m.def(
         "get_helix_workspace_size_per_rank",
diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
index 940d59258ca8..f5a1336ea3e1 100644
--- a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
+++ b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
@@ -28,69 +28,66 @@ TRTLLM_NAMESPACE_BEGIN
 namespace torch_ext
 {
 
-void indexer_k_cache_scatter_op(th::Tensor const& k_fp8_bytes, th::Tensor const& k_scale_bytes, th::Tensor& k_cache,
-    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale)
+void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_scale, th::Tensor& k_cache,
+    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t num_tokens)
 {
-    // Validate all tensors are CUDA tensors
-    TORCH_CHECK(k_fp8_bytes.is_cuda() && k_scale_bytes.is_cuda() && k_cache.is_cuda() && slot_mapping_fp8.is_cuda()
+    // k_fp8: [>=num_tokens, head_dim] in FP8 (1 byte/element) — reinterpreted as uint8
+    // k_scale: [>=num_tokens, head_dim // quant_block_size] in float32 — reinterpreted as uint8 bytes
+    // slot_mapping_fp8, slot_mapping_scale: [>=num_tokens] int64 — only first num_tokens used
+    // k_cache: [num_blocks, block_size, 1, per_token_size] uint8
+
+    TORCH_CHECK(k_fp8.is_cuda() && k_scale.is_cuda() && k_cache.is_cuda() && slot_mapping_fp8.is_cuda()
             && slot_mapping_scale.is_cuda(),
         "All tensors must be CUDA tensors");
 
     // Validate tensor dimensions
-    TORCH_CHECK(k_fp8_bytes.dim() == 2, "k_fp8_bytes must be a 2D Tensor [num_tokens, head_dim]");
-    TORCH_CHECK(k_scale_bytes.dim() == 2, "k_scale_bytes must be a 2D Tensor [num_tokens, scale_size]");
-    TORCH_CHECK(slot_mapping_fp8.dim() == 1, "slot_mapping_fp8 must be a 1D Tensor [num_tokens]");
-    TORCH_CHECK(slot_mapping_scale.dim() == 1, "slot_mapping_scale must be a 1D Tensor [num_tokens]");
-
-    // Enforce k_cache is 4D tensor
-    TORCH_CHECK(k_cache.dim() == 4,
-        "k_cache must be a 4D Tensor [num_blocks, block_size, 1, per_token_size], got %d dimensions",
+    TORCH_CHECK(k_fp8.dim() == 2, "k_fp8 must be 2D [num_tokens, head_dim]");
+    TORCH_CHECK(k_scale.dim() == 2, "k_scale must be 2D [num_tokens, scale_elements]");
+    TORCH_CHECK(slot_mapping_fp8.dim() == 1, "slot_mapping_fp8 must be 1D [num_tokens]");
+    TORCH_CHECK(slot_mapping_scale.dim() == 1, "slot_mapping_scale must be 1D [num_tokens]");
+    TORCH_CHECK(k_cache.dim() == 4, "k_cache must be 4D [num_blocks, block_size, 1, per_token_size], got %d dims",
         static_cast<int>(k_cache.dim()));
 
-    // Validate tensor dtypes
-    TORCH_CHECK(k_fp8_bytes.scalar_type() == torch::kUInt8, "k_fp8_bytes must be uint8");
-    TORCH_CHECK(k_scale_bytes.scalar_type() == torch::kUInt8, "k_scale_bytes must be uint8");
+    // Validate tensor dtypes — reinterpret_cast below assumes specific element sizes
+    TORCH_CHECK(k_fp8.element_size() == 1, "k_fp8 must have 1-byte elements (e.g. FP8), got %d", k_fp8.element_size());
+    TORCH_CHECK(k_scale.element_size() == 4, "k_scale must have 4-byte elements (e.g. float32), got %d",
+        k_scale.element_size());
     TORCH_CHECK(slot_mapping_fp8.scalar_type() == torch::kInt64, "slot_mapping_fp8 must be int64");
     TORCH_CHECK(slot_mapping_scale.scalar_type() == torch::kInt64, "slot_mapping_scale must be int64");
 
-    // Validate tensor shapes are consistent
-    auto num_tokens = static_cast<int32_t>(k_fp8_bytes.size(0));
-    TORCH_CHECK(
-        k_scale_bytes.size(0) == num_tokens, "k_scale_bytes first dimension must equal k_fp8_bytes first dimension");
-    TORCH_CHECK(slot_mapping_fp8.size(0) == num_tokens, "slot_mapping_fp8 length must equal num_tokens");
-    TORCH_CHECK(slot_mapping_scale.size(0) == num_tokens, "slot_mapping_scale length must equal num_tokens");
-
-    // Validate tensors are contiguous (except k_cache which may be non-contiguous)
-    TORCH_CHECK(k_fp8_bytes.is_contiguous(), "k_fp8_bytes must be contiguous");
-    TORCH_CHECK(k_scale_bytes.is_contiguous(), "k_scale_bytes must be contiguous");
-    // k_cache can be non-contiguous - we handle this via strides
+    TORCH_CHECK(k_fp8.is_contiguous(), "k_fp8 must be contiguous");
+    TORCH_CHECK(k_scale.is_contiguous(), "k_scale must be contiguous");
     TORCH_CHECK(slot_mapping_fp8.is_contiguous(), "slot_mapping_fp8 must be contiguous");
     TORCH_CHECK(slot_mapping_scale.is_contiguous(), "slot_mapping_scale must be contiguous");
 
-    int32_t head_dim = static_cast<int32_t>(k_fp8_bytes.size(1));     // head_dim = quant_block_size = 128
-    int32_t scale_size = static_cast<int32_t>(k_scale_bytes.size(1)); // scale_size = 4 bytes
-
-    int32_t cache_dim_0 = static_cast<int32_t>(k_cache.size(0));      // num_blocks
-    int32_t cache_dim_1 = static_cast<int32_t>(k_cache.size(1));      // block_size
-    int32_t cache_dim_2 = static_cast<int32_t>(k_cache.size(2));      // num_kv_heads
-    int32_t cache_dim_3 = static_cast<int32_t>(k_cache.size(3));      // per_token_size
-
-    // Validation for indexer k cache pool for DeepSeek-V3.2 constraints
-    TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1 for DeepSeek-V3.2, got %d", cache_dim_2);
-    TORCH_CHECK(head_dim == 128, "k_fp8_bytes head_dim must be 128 for DeepSeek-V3.2, got %d", head_dim);
-    TORCH_CHECK(scale_size == 4, "k_scale_bytes scale_size must be 4 bytes for DeepSeek-V3.2, got %d", scale_size);
-
-    int64_t cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
-    int64_t cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
-    int64_t cache_stride_2 = static_cast<int64_t>(k_cache.stride(2));
-    int64_t cache_stride_3 = static_cast<int64_t>(k_cache.stride(3));
-
-    auto stream = at::cuda::getCurrentCUDAStream(k_fp8_bytes.get_device());
-
-    tk::invokeIndexerKCacheScatter(k_fp8_bytes.data_ptr<uint8_t>(), k_scale_bytes.data_ptr<uint8_t>(),
-        k_cache.data_ptr<uint8_t>(), slot_mapping_fp8.data_ptr<int64_t>(), slot_mapping_scale.data_ptr<int64_t>(),
-        num_tokens, head_dim, scale_size, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0,
-        cache_stride_1, cache_stride_2, cache_stride_3, stream);
+    // FP8 is 1 byte per element, so head_dim in elements == head_dim in bytes.
+    int32_t const head_dim = static_cast<int32_t>(k_fp8.size(1));
+    // Scale size in bytes: num_scale_elements * bytes_per_element.
+    int32_t const scale_size = static_cast<int32_t>(k_scale.size(1)) * static_cast<int32_t>(k_scale.element_size());
+
+    int32_t const cache_dim_0 = static_cast<int32_t>(k_cache.size(0));
+    int32_t const cache_dim_1 = static_cast<int32_t>(k_cache.size(1));
+    int32_t const cache_dim_2 = static_cast<int32_t>(k_cache.size(2));
+    int32_t const cache_dim_3 = static_cast<int32_t>(k_cache.size(3));
+
+    TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1, got %d", cache_dim_2);
+    TORCH_CHECK(head_dim == 128, "k_fp8 head_dim must be 128, got %d", head_dim);
+    TORCH_CHECK(scale_size == 4, "k_scale scale_size must be 4 bytes, got %d", scale_size);
+
+    int64_t const cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
+    int64_t const cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
+    int64_t const cache_stride_2 = static_cast<int64_t>(k_cache.stride(2));
+    int64_t const cache_stride_3 = static_cast<int64_t>(k_cache.stride(3));
+
+    auto stream = at::cuda::getCurrentCUDAStream(k_fp8.get_device());
+
+    // Reinterpret k_fp8 as uint8 bytes and k_scale as raw bytes via data_ptr.
+    // For slot mappings, use data_ptr directly — only the first num_tokens entries are read.
+    tk::invokeIndexerKCacheScatter(reinterpret_cast<uint8_t const*>(k_fp8.data_ptr()),
+        reinterpret_cast<uint8_t const*>(k_scale.data_ptr()), k_cache.data_ptr<uint8_t>(),
+        slot_mapping_fp8.data_ptr<int64_t>(), slot_mapping_scale.data_ptr<int64_t>(), static_cast<int32_t>(num_tokens),
+        head_dim, scale_size, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1,
+        cache_stride_2, cache_stride_3, stream);
 }
 
 } // namespace torch_ext
@@ -100,8 +97,8 @@ TRTLLM_NAMESPACE_END
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
-        "indexer_k_cache_scatter_op(Tensor k_fp8_bytes, Tensor k_scale_bytes, Tensor(a!) k_cache, "
-        "Tensor slot_mapping_fp8, Tensor slot_mapping_scale) -> ()");
+        "indexer_k_cache_scatter_op(Tensor k_fp8, Tensor k_scale, Tensor(a!) k_cache, "
+        "Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens) -> ()");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp
index 9a7af4da49f6..b526310564e9 100644
--- a/cpp/tensorrt_llm/thop/attentionOp.cpp
+++ b/cpp/tensorrt_llm/thop/attentionOp.cpp
@@ -630,7 +630,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> cu_q_seqlens, std::optional<torch::Tensor> cu_kv_seqlens,
     std::optional<torch::Tensor> fmha_scheduler_counter, std::optional<torch::Tensor> mla_bmm1_scale,
     std::optional<torch::Tensor> mla_bmm2_scale, std::optional<torch::Tensor> quant_q_buffer,
-    std::optional<torch::Tensor> flash_mla_tile_scheduler_metadata, std::optional<torch::Tensor> flash_mla_num_splits)
+    std::optional<torch::Tensor> flash_mla_tile_scheduler_metadata, std::optional<torch::Tensor> flash_mla_num_splits,
+    int64_t num_contexts, int64_t num_ctx_tokens)
 {
     TLLM_LOG_TRACE("Attention op starts at layer %d", layer_idx);
     // Use these tensors to infer if the attention is using KV cache
@@ -833,20 +834,9 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     }
     bool const is_gen_only = attn_input_type == AttentionInputType::GenerationOnly;
 
-    int32_t num_contexts = 0;
-    // count context requests
-    for (int32_t idx = 0; idx < num_seqs; idx++)
-    {
-        if (request_types[idx] != RequestType::kCONTEXT)
-        {
-            break;
-        }
-        ++num_contexts;
-    }
-    int32_t const num_generations = num_seqs - num_contexts;
+    int32_t const num_generations = num_seqs - static_cast<int32_t>(num_contexts);
     int32_t const num_tokens = qkv_or_q.size(0);
-    int32_t const num_ctx_tokens = host_context_lengths.slice(0, 0, num_contexts).sum().item<int32_t>();
-    int32_t const num_gen_tokens = is_gen_only ? num_tokens : num_tokens - num_ctx_tokens;
+    int32_t const num_gen_tokens = is_gen_only ? num_tokens : num_tokens - static_cast<int32_t>(num_ctx_tokens);
     auto const ctx_total_kv_len = host_total_kv_lens.index({0}).item<int32_t>();
     auto const gen_total_kv_len = host_total_kv_lens.index({1}).item<int32_t>();
 
diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h
index 0fc4788d6f0b..854de818292f 100644
--- a/cpp/tensorrt_llm/thop/attentionOp.h
+++ b/cpp/tensorrt_llm/thop/attentionOp.h
@@ -78,7 +78,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> fmha_scheduler_counter, std::optional<torch::Tensor> mla_bmm1_scale,
     std::optional<torch::Tensor> mla_bmm2_scale, std::optional<torch::Tensor> quant_q_buffer,
     std::optional<torch::Tensor> flash_mla_tile_scheduler_metadata = std::nullopt,
-    std::optional<torch::Tensor> flash_mla_num_splits = std::nullopt);
+    std::optional<torch::Tensor> flash_mla_num_splits = std::nullopt, int64_t num_contexts, int64_t num_ctx_tokens);
 
 struct KvCachePoolPointers
 {
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
index b3a5fbeb3136..566dbaf388f2 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py
@@ -1392,34 +1392,18 @@ def _update_k_cache(self, k_fp8: torch.Tensor, k_scale: torch.Tensor,
         if metadata.kv_cache_manager is None or metadata.slot_mapping_fp8 is None:
             return
 
-        # [num_blocks, block_size, 1, per_token_size ]
         k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(
             self.layer_idx)
 
         num_tokens = k_fp8.shape[0]
-        head_dim = k_fp8.shape[1]
-        scale_size = k_scale.shape[1] * 4  # Convert to bytes (float32 = 4 bytes)
-
-        # Convert to bytes: flatten first, then view as uint8, then reshape
-        k_fp8_bytes = k_fp8.view(-1).view(torch.uint8).view(
-            num_tokens, head_dim)
-
-        # k_scale: for single-element tensors, contiguous() may be no-op
-        # Fix stride(-1) for byte-level view
-        k_scale_flat = k_scale.view(-1)
-        if k_scale_flat.stride(-1) != 1:
-            k_scale_flat = torch.as_strided(k_scale_flat.contiguous(),
-                                            size=(k_scale_flat.numel(), ),
-                                            stride=(1, ))
-        k_scale_bytes = k_scale_flat.view(torch.uint8).view(
-            num_tokens, scale_size)
-
-        # Use CUDA kernel to scatter FP8 and scale bytes into cache
-        flat_indices_fp8 = metadata.slot_mapping_fp8[:num_tokens]
-        flat_indices_scale = metadata.slot_mapping_scale[:num_tokens]
-        torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8_bytes, k_scale_bytes,
-                                                    k_cache, flat_indices_fp8,
-                                                    flat_indices_scale)
+
+        # The C++ op reinterprets k_fp8 (FP8) and k_scale (float32) as raw
+        # bytes internally and only reads the first num_tokens entries from
+        # the slot mapping buffers, avoiding Python-side view/slice overhead.
+        torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache,
+                                                    metadata.slot_mapping_fp8,
+                                                    metadata.slot_mapping_scale,
+                                                    num_tokens)
 
     def sparse_attn_indexer(
         self,
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 271538190c8a..d0b710cae5e2 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -407,6 +407,8 @@ def run(
         mla_bmm1_scale: Optional[torch.Tensor] = None,
         mla_bmm2_scale: Optional[torch.Tensor] = None,
         quant_q_buffer: Optional[torch.Tensor] = None,
+        num_contexts: int = 0,
+        num_ctx_tokens: int = 0,
     ):
         """
         Run the attention operation.
@@ -638,6 +640,8 @@ def run(
                 quant_q_buffer,
                 self.quant_config,
                 self.kv_cache_manager,
+                num_contexts,
+                num_ctx_tokens,
                 global_layer_idx=self.global_layer_idx,
             )
         else:
@@ -722,6 +726,8 @@ def run(
                 quant_q_buffer,
                 self.flash_mla_tile_scheduler_metadata,
                 self.flash_mla_num_splits,
+                num_contexts,
+                num_ctx_tokens,
             )
 
         if self.print_skip_softmax_stat:
@@ -2049,7 +2055,9 @@ def forward(
                          fmha_scheduler_counter=fmha_scheduler_counter,
                          mla_bmm1_scale=mla_bmm1_scale,
                          mla_bmm2_scale=mla_bmm2_scale,
-                         quant_q_buffer=quant_q_buffer)
+                         quant_q_buffer=quant_q_buffer,
+                         num_contexts=metadata.num_contexts,
+                         num_ctx_tokens=metadata.num_ctx_tokens)
 
         if output_sf is None:
             return output
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py
index a11d8f51fe19..edb45df74059 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py
@@ -1362,23 +1362,6 @@ def run_generation(self, params: EnqueueGenerationParams):
             )
 
 
-def _parse_request_types(host_request_types: torch.Tensor) -> Tuple[int, int]:
-    """
-    Parse request types to count context and generation requests.
-
-    Args:
-        host_request_types: Request types tensor (0=context, 1=generation).
-        num_seqs: Total number of sequences.
-
-    Returns:
-        Tuple of (num_contexts, num_generations).
-    """
-
-    num_generations = host_request_types.sum().item()
-    num_contexts = host_request_types.size(0) - num_generations
-    return num_contexts, num_generations
-
-
 def is_supported(
     q: torch.Tensor,
     num_heads: int,
@@ -1561,6 +1544,8 @@ def trtllm_gen_attention(
     quant_q_buffer: Optional[torch.Tensor],
     quant_config: Optional[QuantConfig],
     kv_cache_manager: Optional[KVCacheManager],
+    num_contexts: int,
+    num_ctx_tokens: int,
     global_layer_idx: Optional[int] = None,
 ) -> None:
     """
@@ -1691,9 +1676,7 @@ def trtllm_gen_attention(
     if attention_input_type is not None:
         attn_input_type = AttentionInputType(attention_input_type)
 
-    num_contexts, num_generations = _parse_request_types(host_request_types)
-
-    num_ctx_tokens = int(host_context_lengths[:num_contexts].sum()) if num_contexts > 0 else 0
+    num_generations = host_request_types.size(0) - num_contexts
     num_gen_tokens = num_tokens - num_ctx_tokens
 
     # Prepare Workspace
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py
index d0c8a1dd8c0f..d6d4b75da2af 100644
--- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py
@@ -82,6 +82,9 @@ def __init__(self):
         # keeping a separate copy here since we sometimes have to overwrite the original values
         self.host_past_kv_lengths: Optional[torch.Tensor] = None  # [max_batch] int32 pinned
         self.host_context_lengths: Optional[torch.Tensor] = None  # [max_batch] int32 pinned
+        # Batch counts for thop.attention (updated every forward in plan_host)
+        self.num_contexts: int = 0
+        self.num_ctx_tokens: int = 0
         # Persistent block_offsets buffer for CUDA graph compatibility.
         # Pre-allocated to max size so the tensor address is stable across replays.
         self.block_offsets: Optional[torch.Tensor] = None
@@ -171,6 +174,10 @@ def plan_host(
         """
         num_seq = num_prefill + num_decode
 
+        # Batch counts for thop.attention
+        self.num_contexts = num_prefill
+        self.num_ctx_tokens = int(seq_len_host[:num_prefill].sum()) if num_prefill > 0 else 0
+
         # host_request_types: 0 = prefill (context), 1 = decode (generation)
         self.host_request_types[:num_prefill].fill_(0)
         self.host_request_types[num_prefill:num_seq].fill_(1)
@@ -500,6 +507,8 @@ def trtllm_mha_with_cache(
         None,  # mla_bmm1_scale
         None,  # mla_bmm2_scale
         None,  # quant_q_buffer
+        num_contexts=_GlobalTrtllmPlanner.num_contexts,
+        num_ctx_tokens=_GlobalTrtllmPlanner.num_ctx_tokens,
     )
 
     if out is not None:
diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py
index 46bb102fcf7a..244c4d5dbc2a 100644
--- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py
@@ -703,7 +703,7 @@ def test_indexer_k_cache_scatter_custom_op():
                              dtype=torch.bfloat16)
     k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_original)
 
-    # Prepare byte-level data
+    # Prepare byte-level data for the Python reference path
     scale_size = k_scale.shape[1] * 4
     k_fp8_bytes = k_fp8.view(-1).view(torch.uint8).view(num_tokens, head_dim)
     k_scale_flat = k_scale.view(-1)
@@ -742,9 +742,10 @@ def test_indexer_k_cache_scatter_custom_op():
 
     # ========== Path 1: CUDA Kernel ==========
     print(f"\n=== Path 1: CUDA Kernel ===")
-    torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8_bytes, k_scale_bytes,
-                                                k_cache_cuda, flat_indices_fp8,
-                                                flat_indices_scale)
+    torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache_cuda,
+                                                metadata.slot_mapping_fp8,
+                                                metadata.slot_mapping_scale,
+                                                num_tokens)
     torch.cuda.synchronize()
     print(f"✓ CUDA kernel completed")
 

From 4140cccd45cac2b1042b32fa5669ac0551cbcdf3 Mon Sep 17 00:00:00 2001
From: v-shobhit <161510941+v-shobhit@users.noreply.github.com>
Date: Sun, 5 Apr 2026 17:27:59 -0700
Subject: [PATCH 11/14] Add default values for num_contexts and num_ctx_tokens

Signed-off-by: v-shobhit <161510941+v-shobhit@users.noreply.github.com>
---
 cpp/tensorrt_llm/thop/attentionOp.h | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h
index 854de818292f..77f9a965e48e 100644
--- a/cpp/tensorrt_llm/thop/attentionOp.h
+++ b/cpp/tensorrt_llm/thop/attentionOp.h
@@ -78,7 +78,9 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> fmha_scheduler_counter, std::optional<torch::Tensor> mla_bmm1_scale,
     std::optional<torch::Tensor> mla_bmm2_scale, std::optional<torch::Tensor> quant_q_buffer,
     std::optional<torch::Tensor> flash_mla_tile_scheduler_metadata = std::nullopt,
-    std::optional<torch::Tensor> flash_mla_num_splits = std::nullopt, int64_t num_contexts, int64_t num_ctx_tokens);
+    std::optional<torch::Tensor> flash_mla_num_splits = std::nullopt,
+    int64_t num_contexts = 0,
+    int64_t num_ctx_tokens = 0);
 
 struct KvCachePoolPointers
 {

From 01c494799025e1713df74bffd55ec1dbc1de9408 Mon Sep 17 00:00:00 2001
From: Jin Li <59594262+liji-nv@users.noreply.github.com>
Date: Tue, 7 Apr 2026 06:12:42 -0700
Subject: [PATCH 12/14] =?UTF-8?q?[None][fix]=20Pre-subtract=20non-first-ch?=
 =?UTF-8?q?unk=20context=20costs=20in=20reuse=20budge=E2=80=A6=20(#12806)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
---
 .../_torch/pyexecutor/resource_manager.py      | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index b1e7390306bb..ad606440ce0b 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -630,6 +630,14 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                     for req in scheduled_batch.generation_requests)
                 remaining_budget = self.max_num_tokens - gen_tokens
 
+                # Pre-subtract the fixed cost of non-first-chunk context
+                # requests.  These have no reuse to re-validate and their
+                # compute cost is committed, so first-chunk budget checks
+                # must see the budget with these costs already removed.
+                for req in scheduled_batch.context_requests:
+                    if not req.is_first_context_chunk:
+                        remaining_budget -= req.context_chunk_size
+
             accepted_ctx_requests = []
 
             # allocate KV Cache
@@ -688,9 +696,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                             block_ids = self.get_cache_indices(req)
                             self.kv_connector_manager.update_state_after_alloc(
                                 req, block_ids)
-                    elif remaining_budget is not None:
-                        reusable = (req.estimated_reusable_tokens
-                                    if req.is_first_context_chunk else 0)
+                    elif remaining_budget is not None and req.is_first_context_chunk:
+                        # First-chunk request that skipped add_sequence
+                        # (e.g. kv_connector said not to).  Subtract its
+                        # estimated cost so later first-chunk checks see a
+                        # correct budget.  Non-first-chunk costs were
+                        # already pre-subtracted above.
+                        reusable = req.estimated_reusable_tokens
                         remaining_budget -= self._estimate_post_reuse_compute(
                             reusable, req.context_chunk_size, req.prompt_len)
 

From f24dcb33e1d001fdf93ac520ddbfb2178fc795de Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Fri, 10 Apr 2026 10:33:33 +0800
Subject: [PATCH 13/14] [None][fix] Fix disagg KV cache transfer crash on reuse
 budget skip (#12909)

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 25 ++++++++++++-------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index bc8872307d87..8f1a983db9d2 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3018,8 +3018,16 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
                         resource_mgr_type].prepare_resources(
                             disagg_gen_init_to_prepare)
 
-            # Trigger KV cache exchange for new disagg_gen_init_requests
-            self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests)
+            # Use the filtered batch after prepare_resources, not the original
+            # list.  prepare_resources may skip add_sequence for some requests
+            # (e.g. reuse budget exhaustion) and remove them from the batch via
+            # reset_context_requests.  Passing skipped requests to KV transfer
+            # causes unordered_map::at in getSequence because the sequence was
+            # never registered.  Skipped requests remain in active_requests
+            # with DISAGG_GENERATION_INIT state and will be rescheduled.
+            accepted_requests = list(
+                disagg_gen_init_to_prepare.context_requests)
+            self._recv_disagg_gen_cache(accepted_requests)
 
     @nvtx_range("_prepare_disagg_gen_transmission_complete")
     def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
@@ -3642,14 +3650,13 @@ def _handle_responses(self):
                                     f"Request {request.py_request_id} has no avg_decoded_tokens_per_iter"
                                 )
 
-                # If partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1,
-                # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse
-                # path is fixed.
-                if self.enable_partial_reuse_for_disagg and not self.kv_cache_manager.is_vswa and self.dist.pp_size == 1:
+                # Never terminate a request while KV cache transfer is
+                # still in progress — the GEN worker may still be receiving
+                # data from the pinned blocks.  The request will be
+                # terminated by _end_transfer_and_maybe_terminate once the
+                # transfer completes (or times out).
+                if not request.is_disagg_context_transmission_state:
                     requests_to_terminate.append(request)
-                else:
-                    if not request.is_disagg_context_transmission_state:
-                        requests_to_terminate.append(request)
             else:
                 new_active_requests.append(request)
 

From 4867bc922109d6487bceb6afdcf5c05e073e534a Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Sun, 5 Apr 2026 18:24:07 +0800
Subject: [PATCH 14/14] [TRTLLM-11597][fix] fix disagg kvcache router for chat
 API and add disagg benchmark for ai_perf (#12337)

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 .../slurm/benchmark/run_benchmark_aiperf.sh   | 102 ++++++++++++
 .../disaggregated/slurm/benchmark/submit.py   |  22 ++-
 tensorrt_llm/serve/router.py                  |  42 ++++-
 tests/unittest/disaggregated/test_router.py   | 156 ++++++++++++++++++
 4 files changed, 313 insertions(+), 9 deletions(-)
 create mode 100755 examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh

diff --git a/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh b/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh
new file mode 100755
index 000000000000..1f2d19b18417
--- /dev/null
+++ b/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh
@@ -0,0 +1,102 @@
+#!/bin/bash
+
+# aiperf-based benchmark script for disaggregated serving
+# Args: model_name dataset_file multi_round num_gen_servers concurrency_list streaming log_path hostname port ucx_warmup_requests
+
+set -e
+set -u
+trap 'echo "Error occurred at line $LINENO"; exit 1' ERR
+
+if [ "$#" -lt 10 ]; then
+    echo "Error: Missing required arguments, got $# arguments, args: $@"
+    echo "Usage: $0 model_name dataset_file multi_round num_gen_servers concurrency_list streaming log_path hostname port ucx_warmup_requests"
+    exit 1
+fi
+
+model_name=$1
+dataset_file=$2
+multi_round=$3
+num_gen_servers=$4
+concurrency_list=$5
+streaming=$6
+log_path=$7
+hostname=$8
+port=$9
+ucx_warmup_requests=${10}
+
+# check process id is not 0
+if [[ ${SLURM_PROCID} != "0" ]]; then
+    echo "Process id is ${SLURM_PROCID} for loadgen, exiting"
+    exit 0
+fi
+
+# Always install/upgrade aiperf to ensure we have the version with trust_remote_code fix
+# (container may have an older version with parallel_decode.py that lacks trust_remote_code)
+echo "Installing aiperf..."
+pip install --force-reinstall --no-deps 'aiperf @ git+https://github.com/ai-dynamo/aiperf.git@ac3d91652e5e024bfb4ac38d48603423aad666bc'
+
+# warmup requests for ucx connections
+if [ "${ucx_warmup_requests}" -gt 0 ]; then
+    echo "warming up ucx connections with small requests... ${ucx_warmup_requests}"
+    python -m tensorrt_llm.serve.scripts.benchmark_serving \
+        --model ${model_name} \
+        --dataset-name random \
+        --random-ids \
+        --random-input-len 100 \
+        --random-output-len 10 \
+        --num-prompts ${ucx_warmup_requests} \
+        --host ${hostname} \
+        --port ${port} \
+        --ignore-eos \
+        --trust-remote-code \
+        --non-streaming
+    echo "UCX warmup done"
+fi
+
+# Trust remote code globally for custom tokenizers in parallel workers
+export HF_HUB_TRUST_REMOTE_CODE=1
+
+echo "Hostname: ${hostname}, Port: ${port}"
+echo "Starting aiperf benchmark..."
+
+concurrency_list=$(echo "${concurrency_list}" | tr ',' ' ')
+for concurrency in ${concurrency_list}; do
+    concurrency=$((concurrency))
+    request_count=$((concurrency * multi_round))
+    # benchmark_duration: 20min per round
+    benchmark_duration=$((multi_round * 1200))
+    echo "Benchmarking with concurrency ${concurrency} ... ${request_count} requests, duration ${benchmark_duration}s"
+    mkdir -p ${log_path}/concurrency_${concurrency}
+
+    aiperf profile \
+        -m ${model_name} \
+        --tokenizer ${model_name} \
+        --tokenizer-trust-remote-code \
+        --url http://${hostname}:${port} \
+        --streaming \
+        --ui simple \
+        --input-file ${dataset_file} \
+        --artifact-dir ${log_path}/concurrency_${concurrency} \
+        --concurrency ${concurrency} \
+        --concurrency-ramp-duration 60 \
+        --custom-dataset-type mooncake_trace \
+        --benchmark-duration ${benchmark_duration} \
+        --benchmark-grace-period 60 \
+        --workers-max 200 \
+        --request-timeout-seconds 1200 \
+        --profile-export-level records \
+        --extra-inputs ignore_eos:true \
+        --request-count ${request_count} \
+        --record-processors 8
+
+    echo "Benchmark with concurrency ${concurrency} done"
+done
+
+# Fetch perf metrics from disagg server
+echo "Fetching perf metrics from http://${hostname}:${port}/perf_metrics ..."
+curl -s "http://${hostname}:${port}/perf_metrics" > ${log_path}/perf_metrics.json 2>&1 || true
+if [ -s "${log_path}/perf_metrics.json" ]; then
+    echo "Perf metrics saved to ${log_path}/perf_metrics.json"
+else
+    echo "Warning: perf_metrics response was empty or endpoint not available"
+fi
diff --git a/examples/disaggregated/slurm/benchmark/submit.py b/examples/disaggregated/slurm/benchmark/submit.py
index 716260fea333..d026d1ac362b 100644
--- a/examples/disaggregated/slurm/benchmark/submit.py
+++ b/examples/disaggregated/slurm/benchmark/submit.py
@@ -113,7 +113,9 @@ def assign_servers(
     return allocations
 
 
-def convert_allocations_to_server_config(allocations, server_port=8333):
+def convert_allocations_to_server_config(allocations,
+                                         server_port=8333,
+                                         router_config=None):
     generation_servers = {}
     context_servers = {}
     server_hostname = None
@@ -127,6 +129,8 @@ def convert_allocations_to_server_config(allocations, server_port=8333):
                 f"{list(instance['nodes'].keys())[0]}:{instance['port']}")
 
         server_config_entry = {'num_instances': num_servers, 'urls': urls}
+        if router_config:
+            server_config_entry['router'] = router_config.copy()
 
         if server_type == "GEN":
             generation_servers = server_config_entry
@@ -481,7 +485,12 @@ def submit_job(config, log_dir, dry_run):
         json.dump(allocations, f, indent=2)
 
     # Generate disagg server config
-    server_config = convert_allocations_to_server_config(allocations)
+    router_config = config.get('router_config', None)
+    server_config = convert_allocations_to_server_config(
+        allocations, router_config=router_config)
+    # Merge server_config_extra into disagg server config
+    if 'server_config_extra' in config:
+        server_config.update(config['server_config_extra'])
     with open(os.path.join(log_dir, "server_config_base.yaml"), "w") as f:
         yaml.dump(server_config, f)
     disagg_server_hostname = server_config['hostname']
@@ -608,7 +617,14 @@ def submit_job(config, log_dir, dry_run):
         benchmark_prefix = client_slurm_prefix + [
             f"--export \"{convert_envs_to_str(env_var)}\""
         ]
-        if benchmark_config['use_nv_sa_benchmark']:
+        if benchmark_config.get('use_aiperf', False):
+            benchmark_cmd = [
+                f"bash {os.path.join(script_dir, 'run_benchmark_aiperf.sh')}",
+                f"'{env_config['model_path']}' '{benchmark_config['dataset_file']}' {benchmark_config['multi_round']} {gen_num} '{benchmark_config['concurrency_list']}' {benchmark_config['streaming']} '{log_dir}' {disagg_server_hostname} {disagg_server_port} {ucx_warmup_requests}",
+                f"&> {log_dir}/6_bench.log"
+            ]
+            client_cmds.append(" ".join(benchmark_prefix + benchmark_cmd))
+        elif benchmark_config['use_nv_sa_benchmark']:
             if benchmark_config['mode'] == "gen_only":
                 print(
                     f"[ERROR] SA benchmark client script is not supported for gen_only mode"
diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py
index 04336403ea6d..e72c58b1abb1 100644
--- a/tensorrt_llm/serve/router.py
+++ b/tensorrt_llm/serve/router.py
@@ -1,5 +1,6 @@
 import asyncio
 import heapq
+import os
 from abc import ABC, abstractmethod
 from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Union
 
@@ -626,9 +627,38 @@ def __init__(self,
         self._tokenizers = {}
         # TODO: use max_num_tokens? per server?
         self._max_batch_size = max_batch_size
+        env_tokens_per_block = os.environ.get(
+            "TRTLLM_KVCACHE_AWARE_ROUTER_HASH_TOKENS_PER_BLOCK")
+        if env_tokens_per_block is not None:
+            tokens_per_block = int(env_tokens_per_block)
         self._tokens_per_block = tokens_per_block
+        logger.info(
+            f"KvCacheAwareRouter: tokens_per_block={self._tokens_per_block}")
+
+    def _get_tokenizer(self, model: str):
+        if model not in self._tokenizers:
+            self._tokenizers[model] = AutoTokenizer.from_pretrained(model)
+        return self._tokenizers[model]
 
     def _tokenize(self, request: OpenAIRequest) -> list[list[int]]:
+        # Handle ChatCompletionRequest (has messages, not prompt)
+        if isinstance(request, ChatCompletionRequest):
+            if request.prompt_token_ids is not None:
+                return [request.prompt_token_ids]
+            tokenizer = self._get_tokenizer(request.model)
+            token_ids = tokenizer.apply_chat_template(
+                [
+                    msg if isinstance(msg, dict) else dict(msg)
+                    for msg in request.messages
+                ],
+                add_generation_prompt=request.add_generation_prompt,
+                tokenize=True,
+            )
+            # Set prompt_token_ids so the worker server skips re-tokenization
+            request.prompt_token_ids = token_ids
+            return [token_ids]
+
+        # Handle CompletionRequest (has prompt)
         prompts = request.prompt
         if isinstance(prompts, list) and isinstance(prompts[0], list):
             return prompts
@@ -639,12 +669,12 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]:
         else:
             assert isinstance(prompts, list) and isinstance(prompts[0], str)
 
-        # TODO: send tokenize-only request instead of tokenizing locally
-        if request.model not in self._tokenizers:
-            self._tokenizers[request.model] = AutoTokenizer.from_pretrained(
-                request.model)
-        tokenizer = self._tokenizers[request.model]
-        return [tokenizer(prompt)["input_ids"] for prompt in prompts]
+        tokenizer = self._get_tokenizer(request.model)
+        token_lists = [tokenizer(prompt)["input_ids"] for prompt in prompts]
+        # Replace string prompts with token IDs so the worker server
+        # skips re-tokenization
+        request.prompt = token_lists if len(token_lists) > 1 else token_lists[0]
+        return token_lists
 
     async def get_next_server(
             self,
diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py
index 73310f5f03ab..c016be080e83 100644
--- a/tests/unittest/disaggregated/test_router.py
+++ b/tests/unittest/disaggregated/test_router.py
@@ -319,6 +319,162 @@ async def test_kv_cache_aware_router(servers):
     assert servers == ("server3", "server1", "server2")
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("api_type", ["completion", "chat"])
+async def test_kv_cache_aware_router_multi_turn_conversation(api_type):
+    """Test that consecutive turns of a multi-turn conversation route to the
+    same server due to KV cache prefix hits.
+
+    Simulates two concurrent sessions inspired by
+    agentic_data/dataset_sample2000.jsonl session sess-fca58a1f44cd:
+      Turn 0: 68 hash_ids (system prompt + first user input)
+      Turn 1:  9 hash_ids (second user input, accumulated with turn 0)
+      Turn 2:  6 hash_ids (third user input, accumulated with turn 1)
+
+    Scaled down to 10, 3, 2 blocks for test manageability.  Each hash_id
+    maps to a deterministic block of tokens (mirroring aiperf's
+    HashIdRandomGenerator).  The router should prefer the server that
+    already caches the conversation's prefix.
+    """
+    server_list = ["server1", "server2", "server3"]
+    tokens_per_block = 32
+
+    router = KvCacheAwareRouter(
+        server_role=None,
+        servers=server_list,
+        use_tokens=False,
+        max_batch_size=64,
+        tokens_per_block=tokens_per_block,
+    )
+
+    # -- helpers ----------------------------------------------------------
+    def hash_id_to_block(hash_id: int) -> list[int]:
+        """Deterministic token block per hash_id (mirrors aiperf corpus sampling)."""
+        return [(hash_id * 7 + i) % 50000 for i in range(tokens_per_block)]
+
+    def build_tokens(hash_ids: list[int]) -> list[int]:
+        tokens = []
+        for hid in hash_ids:
+            tokens.extend(hash_id_to_block(hid))
+        # Append one extra token so the last full block is included in hashing.
+        # (KvCacheManager excludes the very last token from block keys.)
+        tokens.append(0)
+        return tokens
+
+    def make_request(token_ids: list[int]):
+        """Create a CompletionRequest or ChatCompletionRequest with pre-tokenized IDs."""
+        if api_type == "completion":
+            return CompletionRequest(model="TinyLlama", prompt=[token_ids])
+        else:
+            # Use prompt_token_ids to skip tokenizer (no real model needed)
+            return ChatCompletionRequest(
+                model="TinyLlama",
+                messages=[{
+                    "role": "user",
+                    "content": "dummy"
+                }],
+                prompt_token_ids=token_ids,
+            )
+
+    # -- dataset-inspired hash_ids per turn (new blocks only) -------------
+    # Session A (the conversation under test)
+    sess_a_turn0_hids = list(range(10))  # 10 blocks
+    sess_a_turn1_hids = list(range(100, 103))  # 3 new blocks
+    sess_a_turn2_hids = list(range(200, 202))  # 2 new blocks
+
+    # Session B (competing traffic on a different server)
+    sess_b_turn0_hids = list(range(500, 510))  # 10 completely different blocks
+
+    # -- build accumulated token sequences --------------------------------
+    # Turn 0: just the first turn's tokens
+    sess_a_turn0_tokens = build_tokens(sess_a_turn0_hids)
+
+    # Turn 1 accumulated: turn 0 tokens + simulated assistant reply + new user tokens
+    sess_a_turn1_tokens = build_tokens(sess_a_turn0_hids + [9990, 9991] +
+                                       sess_a_turn1_hids)
+    # (hash_ids 9990/9991 stand in for the assistant-reply blocks)
+
+    # Turn 2 accumulated: extends turn 1 further
+    sess_a_turn2_tokens = build_tokens(sess_a_turn0_hids + [9990, 9991] +
+                                       sess_a_turn1_hids + [9992, 9993] +
+                                       sess_a_turn2_hids)
+
+    sess_b_tokens = build_tokens(sess_b_turn0_hids)
+
+    # -- Round 1: initial routing (empty caches) --------------------------
+    # Route both sessions concurrently so load-balancing spreads them to
+    # different servers (with equal KV cache misses, ties are broken by load).
+    req_a0 = make_request(sess_a_turn0_tokens)
+    server_a, info_a0 = await router.get_next_server(req_a0)
+    # Do NOT finish req_a0 yet — keep its load active so session B avoids server_a
+
+    req_b0 = make_request(sess_b_tokens)
+    server_b, info_b0 = await router.get_next_server(req_b0)
+
+    # Now finish both and populate caches
+    await router.finish_request(req_a0)
+    await router.finish_request(req_b0)
+    router._server_state[server_a].add_blocks(info_a0["block_hashes"][0])
+    router._server_state[server_b].add_blocks(info_b0["block_hashes"][0])
+
+    # Sanity: two sessions should land on different servers
+    assert server_a != server_b, "Disjoint sessions should land on different servers"
+
+    # Verify block hashes are disjoint between sessions
+    blocks_a = set(info_a0["block_hashes"][0])
+    blocks_b = set(info_b0["block_hashes"][0])
+    assert blocks_a.isdisjoint(
+        blocks_b), "Different sessions must not share block hashes"
+
+    # -- Round 2: turn 1 of session A (prefix extends turn 0) ------------
+    req_a1 = make_request(sess_a_turn1_tokens)
+    server_a1, info_a1 = await router.get_next_server(req_a1)
+    await router.finish_request(req_a1)
+
+    assert server_a1 == server_a, (
+        f"Turn 1 must route to the same server as turn 0 ({server_a}) "
+        f"due to KV cache prefix hit, but got {server_a1}. "
+        f"Matches: {info_a1['matches']}")
+
+    # The match count on server_a must equal the prefix overlap
+    server_a_idx = list(router._server_state.keys()).index(server_a)
+    expected_prefix_match = len(sess_a_turn0_hids) * tokens_per_block
+    assert info_a1["matches"][server_a_idx] == expected_prefix_match, (
+        f"Expected {expected_prefix_match} matched tokens on server_a, "
+        f"got {info_a1['matches'][server_a_idx]}")
+
+    # Update server_a cache with new blocks from turn 1
+    router._server_state[server_a].add_blocks(info_a1["block_hashes"][0])
+
+    # -- Round 3: turn 2 of session A (prefix extends turn 1) ------------
+    req_a2 = make_request(sess_a_turn2_tokens)
+    server_a2, info_a2 = await router.get_next_server(req_a2)
+    await router.finish_request(req_a2)
+
+    assert server_a2 == server_a, (
+        f"Turn 2 must route to the same server as turns 0-1 ({server_a}) "
+        f"due to KV cache prefix hit, but got {server_a2}. "
+        f"Matches: {info_a2['matches']}")
+
+    # Turn 2 should match all of turn 0 + turn 1 prefix blocks
+    expected_full_match = (
+        len(sess_a_turn0_hids) + 2 +
+        len(sess_a_turn1_hids)  # turn0 + reply + turn1
+    ) * tokens_per_block
+    assert info_a2["matches"][server_a_idx] == expected_full_match, (
+        f"Expected {expected_full_match} matched tokens on turn 2, "
+        f"got {info_a2['matches'][server_a_idx]}")
+
+    # -- Verify session B still routes to its own server ------------------
+    req_b1 = make_request(sess_b_tokens)
+    server_b1, info_b1 = await router.get_next_server(req_b1)
+    await router.finish_request(req_b1)
+
+    assert server_b1 == server_b, (
+        f"Session B should route to its original server ({server_b}), "
+        f"but got {server_b1}")
+
+
 def test_create_router(servers):
     default_router = create_router(None, servers)
     assert isinstance(default_router, RoundRobinRouter)