[None][perf] DSA: skip the indexer schedule rebuild for index-sharing draft steps and the duplicated kv-lens hook

tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py

@@ -972,9 +972,9 @@ def prepare_compressed_kv_metadata(
                 gen_output_offsets,
             )
 
-    def on_update_kv_lens(self):
+    def on_update_kv_lens(self, skip_indexer_schedule: bool = False):
         """Recompute kv-lens-dependent DeepSeek-V4 metadata on device."""
-        super().on_update_kv_lens()
+        super().on_update_kv_lens(skip_indexer_schedule=skip_indexer_schedule)
 
         batch_size = self.num_seqs
         num_tokens = self.num_tokens

tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py

@@ -1427,6 +1427,8 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata):
         """
         num_contexts = metadata.num_contexts
         num_generations = metadata.num_generations
+        # A full rebuild clears the stale marker.
+        metadata.indexer_schedule_stale = False
         gen_seq_lens = metadata.get_indexer_kv_lens(
             metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations]
         )
@@ -2012,6 +2014,15 @@ def sparse_attn_indexer(
             # A subclass filled the generation rows.
             pass
         elif has_decode and not metadata.skip_indexer_for_gen_reqs:
+            # SimpleNamespace doubles in test_indexer_gvr_prior reach this path without the field.
+            if getattr(metadata, "indexer_schedule_stale", False):
+                raise RuntimeError(
+                    "DSA indexer: the MQA-logits schedule is stale (its rebuild was skipped) "
+                    "but this pass computes TopK "
+                    f"(in_mtp_draft_loop={metadata.in_mtp_draft_loop}, "
+                    f"indexer_skip_topk={metadata.indexer_skip_topk}, "
+                    f"mtp_index_share={self.mtp_index_share})"
+                )
             # Get decode lengths per request (from seq_lens) for validation
             gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations]
             if not metadata.is_ragged_verify:

tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py

@@ -126,6 +126,10 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata):
     indexer_skip_topk: bool = False
     in_mtp_draft_loop: bool = False
     mtp_num_accepted: Optional[torch.Tensor] = None
+    # Set by on_update_kv_lens(skip_indexer_schedule=True): the DeepGEMM MQA-logits schedule, the
+    # 2D/expanded indexer KV lengths and the GVR row order were left untouched. Legal only if no
+    # MQA-logits + TopK pass runs before the next full rebuild; the indexer asserts via this marker.
+    indexer_schedule_stale: bool = False
     # Whether skip the indexer for context requests
     skip_indexer_for_ctx_reqs: bool = False
     # Whether skip the indexer for generation requests
@@ -638,7 +642,7 @@ def warmup_selfsampling_topk(
                 "will JIT-compile lazily on first touch instead."
             )
 
-    def on_update_kv_lens(self) -> None:
+    def on_update_kv_lens(self, skip_indexer_schedule: bool = False) -> None:
         # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas.
         # Especially for the changes in DecoderInputPreparer.preprocess_inputs().
         #
@@ -647,6 +651,10 @@ def on_update_kv_lens(self) -> None:
         # (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.
+        #
+        # skip_indexer_schedule leaves the indexer's MQA-logits schedule as is; see
+        # indexer_schedule_stale for what that covers and when it is legal.
+        self.indexer_schedule_stale = skip_indexer_schedule
 
         super().on_update_kv_lens()
 
@@ -740,59 +748,8 @@ def on_update_kv_lens(self) -> None:
                     dtype=torch.int64,
                     out=self.gen_cached_token_indptr[1 : self.num_generations + 1],
                 )
-            gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs]
-            gen_indexer_kv_lens = self.get_indexer_kv_lens(gen_kv_lens)
-            self.gen_indexer_kv_lens_cuda_runtime = gen_indexer_kv_lens
-            scheduler_context_lens = self.update_indexer_kv_lens_2d(gen_indexer_kv_lens)
-            scheduler_metadata_buffer = get_paged_mqa_logits_metadata(
-                gen_indexer_kv_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms
-            )
-            self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True)
-            if self.max_draft_tokens > 0 and not self.use_expanded_buffers_for_mtp:
-                scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata(
-                    scheduler_context_lens.contiguous(),
-                    _DG_SCHEDULE_BLOCK_KV,
-                    self.num_sms,
-                )
-                self.scheduler_metadata_buffer_full_next_n.copy_(
-                    scheduler_metadata_buffer_full_next_n, non_blocking=True
-                )
-            if self.use_expanded_buffers_for_mtp:
-                kv_lens_expanded, num_tokens = self.expand_per_gen_token(gen_indexer_kv_lens)
-                self.kv_lens_expanded_cuda[:num_tokens].copy_(kv_lens_expanded)
-                scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata(
-                    self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1),
-                    _DG_SCHEDULE_BLOCK_KV,
-                    self.num_sms,
-                )
-                self.scheduler_metadata_buffer_expanded.copy_(
-                    scheduler_metadata_buffer_expanded, non_blocking=True
-                )
-            if self.expand_for_dsl and self.dsl_kernel_atoms > 1:
-                num_rows = self.num_generations * self.dsl_kernel_atoms
-                atom_lens = gen_indexer_kv_lens.unsqueeze(1) - self.dsl_atom_ctx_offsets
-                self.kv_lens_expanded_cuda[:num_rows].copy_(atom_lens.reshape(-1))
-                scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata(
-                    self.kv_lens_expanded_cuda[:num_rows].view(-1, 1),
-                    _DG_SCHEDULE_BLOCK_KV,
-                    self.num_sms,
-                )
-                self.scheduler_metadata_buffer_expanded.copy_(
-                    scheduler_metadata_buffer_expanded, non_blocking=True
-                )
-            if self.expand_for_dsl and self.dsl_expand_factor > 1:
-                expand_factor = self.dsl_expand_factor
-                num_tokens = self.num_generations * expand_factor
-                gen_kv_lens_expanded = gen_indexer_kv_lens.repeat_interleave(expand_factor)
-                self.kv_lens_expanded_cuda[:num_tokens].copy_(gen_kv_lens_expanded)
-                scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata(
-                    self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1),
-                    _DG_SCHEDULE_BLOCK_KV,
-                    self.num_sms,
-                )
-                self.scheduler_metadata_buffer_expanded.copy_(
-                    scheduler_metadata_buffer_expanded, non_blocking=True
-                )
+        if self.num_generations > 0 and not skip_indexer_schedule:
+            self._rebuild_indexer_schedule()
 
         if fused_eligible:
             if not getattr(self, "_fused_dsa_meta_armed", False):
@@ -804,9 +761,67 @@ def on_update_kv_lens(self) -> None:
                 self._fused_dsa_meta_armed = True
             self._run_fused_dsa_decode_metadata()
 
-        self._compute_kv_lens_row_reorder()
+        if not skip_indexer_schedule:
+            self._compute_kv_lens_row_reorder()
         self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True)
 
+    def _rebuild_indexer_schedule(self) -> None:
+        """Rebuild what the indexer's MQA-logits + TopK pass reads from the gen KV lengths:
+        the DeepGEMM schedule(s), the 2D/expanded indexer KV lengths and their runtime copy."""
+        gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs]
+        gen_indexer_kv_lens = self.get_indexer_kv_lens(gen_kv_lens)
+        self.gen_indexer_kv_lens_cuda_runtime = gen_indexer_kv_lens
+        scheduler_context_lens = self.update_indexer_kv_lens_2d(gen_indexer_kv_lens)
+        scheduler_metadata_buffer = get_paged_mqa_logits_metadata(
+            gen_indexer_kv_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms
+        )
+        self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True)
+        if self.max_draft_tokens > 0 and not self.use_expanded_buffers_for_mtp:
+            scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata(
+                scheduler_context_lens.contiguous(),
+                _DG_SCHEDULE_BLOCK_KV,
+                self.num_sms,
+            )
+            self.scheduler_metadata_buffer_full_next_n.copy_(
+                scheduler_metadata_buffer_full_next_n, non_blocking=True
+            )
+        if self.use_expanded_buffers_for_mtp:
+            kv_lens_expanded, num_tokens = self.expand_per_gen_token(gen_indexer_kv_lens)
+            self.kv_lens_expanded_cuda[:num_tokens].copy_(kv_lens_expanded)
+            scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata(
+                self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1),
+                _DG_SCHEDULE_BLOCK_KV,
+                self.num_sms,
+            )
+            self.scheduler_metadata_buffer_expanded.copy_(
+                scheduler_metadata_buffer_expanded, non_blocking=True
+            )
+        if self.expand_for_dsl and self.dsl_kernel_atoms > 1:
+            num_rows = self.num_generations * self.dsl_kernel_atoms
+            atom_lens = gen_indexer_kv_lens.unsqueeze(1) - self.dsl_atom_ctx_offsets
+            self.kv_lens_expanded_cuda[:num_rows].copy_(atom_lens.reshape(-1))
+            scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata(
+                self.kv_lens_expanded_cuda[:num_rows].view(-1, 1),
+                _DG_SCHEDULE_BLOCK_KV,
+                self.num_sms,
+            )
+            self.scheduler_metadata_buffer_expanded.copy_(
+                scheduler_metadata_buffer_expanded, non_blocking=True
+            )
+        if self.expand_for_dsl and self.dsl_expand_factor > 1:
+            expand_factor = self.dsl_expand_factor
+            num_tokens = self.num_generations * expand_factor
+            gen_kv_lens_expanded = gen_indexer_kv_lens.repeat_interleave(expand_factor)
+            self.kv_lens_expanded_cuda[:num_tokens].copy_(gen_kv_lens_expanded)
+            scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata(
+                self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1),
+                _DG_SCHEDULE_BLOCK_KV,
+                self.num_sms,
+            )
+            self.scheduler_metadata_buffer_expanded.copy_(
+                scheduler_metadata_buffer_expanded, non_blocking=True
+            )
+
     def _run_fused_dsa_decode_metadata(self):
         """Fill req_idx_per_token + slot mappings + gen indptrs via one Triton
         launch (see kernels.py), replacing the eager chain for the
@@ -861,8 +876,13 @@ def update_for_spec_dec(self):
         self.num_ctx_cached_tokens = 0
         self.max_gen_seq_len = 1
 
-        # device
-        self.on_update_kv_lens()
+        # device. Inside an index-sharing draft loop every call here prepares either a
+        # TopK-reuse draft step or the post-loop target forward, which rebuilds in
+        # _preprocess_inputs: the indexer schedule is never read before its next full rebuild.
+        skip_indexer_schedule = (
+            self.in_mtp_draft_loop and self.sparse_metadata_params.mtp_index_share
+        )
+        self.on_update_kv_lens(skip_indexer_schedule=skip_indexer_schedule)
 
     # Create buffers for mla_rope_append_paged_kv_assign_q
     def create_buffers_for_mla_rope_append(self, capture_graph=False):
@@ -1003,6 +1023,7 @@ def create_buffers_for_indexer(self, capture_graph=False):
         self.indexer_skip_topk = False
         self.in_mtp_draft_loop = False
         self.mtp_num_accepted = None
+        self.indexer_schedule_stale = False
 
         # Indexer metadata
         # Separate slot mappings for non-interleaved layout (flat byte indices)

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py

@@ -2916,8 +2916,19 @@ def preprocess_inputs(
             if self.guided_decoder is not None:
                 self.guided_decoder.token_event.record()
             return inputs
+        # Overlap scheduling corrects kv_lens_cuda below from the runtime
+        # accepted-token counts and calls on_update_kv_lens() again after it.
+        has_kv_lens_correction = (
+            enable_spec_decode
+            and not self._config.disable_overlap_scheduler
+            and attn_meta is not None
+            and attn_meta.kv_cache_manager is not None
+            and hasattr(attn_meta, "kv_lens_cuda")
+        )
         # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step.
-        if attn_meta is not None:
+        # With a correction, the call after it does this: every on_update_kv_lens rebuilds
+        # from the current lengths and nothing reads its outputs in between.
+        if attn_meta is not None and not has_kv_lens_correction:
             attn_meta.on_update_kv_lens()
 
         if enable_spec_decode and not self._config.disable_overlap_scheduler:
@@ -2943,7 +2954,7 @@ def preprocess_inputs(
                         self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens]
                     )
 
-                if hasattr(inputs["attn_metadata"], "kv_lens_cuda"):
+                if has_kv_lens_correction:
                     if (
                         num_ctx_requests >= num_chunked_ctx_requests
                         and num_chunked_ctx_requests > 0

tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py

@@ -4570,6 +4570,10 @@ def test_indexer_decode_mtp_topk_reuse(step0_mode, batch_size, dsa_cache_manager
     step. Covers both step-0 shapes: "gen" (steady state, next_n > 1) and
     "context" (the first gen round after prefill, which runs the context path).
     Step-0 Top-K correctness is covered by test_indexer_decode_custom_vs_fallback.
+
+    The reuse steps run with the MQA-logits schedule marked stale and poisoned,
+    as the draft loop leaves it before an index-shared step; reuse must not read
+    it, and a step that computes Top-K on it must raise instead.
     """
     torch.manual_seed(7)
     heads, head_dim, block_size = 32, 128, 64
@@ -4709,12 +4713,19 @@ def make_inputs(n_tokens):
         meta.in_mtp_draft_loop = True
         meta.shared_topk_indices = stash
         meta.indexer_skip_topk = True
+        meta.indexer_schedule_stale = True
+        meta.scheduler_metadata_buffer.fill_(-1)
+        meta.scheduler_metadata_buffer_full_next_n.fill_(-1)
+        meta.kv_lens_cuda_2d.fill_(-1)
         hs, qs, ks_fp8, ks_scale, ws = make_inputs(batch_size)
         indexer._update_k_cache(ks_fp8, ks_scale, meta)
         topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws)
         assert torch.equal(topk, stash[:batch_size, :]), (
             f"{step0_mode} draft reuse step {step} should copy the stash 1:1 (next_n=1)"
         )
+        meta.indexer_skip_topk = False
+        with pytest.raises(RuntimeError, match="stale"):
+            indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws)
 
 
 @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")