From 3d04d5cdf288418ea3dc971b190979db5c7e2b4b Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:25:52 -0700 Subject: [PATCH 1/5] [None][perf] DSA: skip the indexer schedule rebuild for index-sharing draft steps and the duplicated kv-lens hook Two dead-work elisions on the DSA decode path with bit-identical device results. 1. Indexer schedule rebuilds inside an index-sharing draft loop. update_for_spec_dec() -> on_update_kv_lens() rebuilt the DeepGEMM MQA-logits schedule (the schedule launches, the 2D/expanded indexer KV-length copies and the GVR row order) before every draft step and once more after the last one. With mtp_index_share every update_for_spec_dec() made inside the draft loop prepares either a TopK-reuse draft step, which never reads the schedule, or the post-loop target forward, which rebuilds it in _preprocess_inputs. - Invariant: those buffers are read only by the indexer's MQA-logits + TopK pass (the decode branch of sparse_attn_indexer). Slot mappings and the gen indptrs are still refreshed on every call. - Decision: DSAtrtllmAttentionMetadata.update_for_spec_dec() decides the skip itself (in_mtp_draft_loop and sparse_metadata_params.mtp_index_share) and calls on_update_kv_lens(skip_indexer_schedule=True); the public parameterless call stays the full rebuild. The schedule block is extracted into _rebuild_indexer_schedule(). eagle3.py is untouched. - Stale assert: on_update_kv_lens() records the skip in indexer_schedule_stale until the next full rebuild (on_update_kv_lens() or Indexer.prepare_scheduler_metadata) and the TopK pass raises if it finds the marker set, so a misplaced skip fails loudly. - DeepseekV4TrtllmAttentionMetadata overrides the hook and forwards the argument, so its own rebuild still runs on every call. 2. Duplicated on_update_kv_lens() at the step boundary. Under the overlap scheduler with speculative decoding _preprocess_inputs called the hook before and again after the kv-length correction with nothing reading its results in between. Idempotence travels with the method: interface.py adds the idempotent_kv_lens_hook decorator (contract: a pure rebuild from the current lengths, calling it twice equals calling it once) and kv_lens_hook_is_idempotent(metadata), which reads the marker off the class's hook so any override without it is not idempotent. DSA decorates its hook; the engine drops the first call when a correction follows and the hook is idempotent, and the correction block uses the same has_kv_lens_correction predicate. Effect, structurally: 5-12 launches per skipped call; two skipped calls per decode step at MTP-1, six at MTP-5. Tests: - sparse/dsa/test_dsa_indexer.py::test_indexer_decode_mtp_topk_reuse: the reuse steps always run with the schedule marked stale and poisoned; reuse must copy the stash 1:1 and a TopK pass on it must raise. - sparse/dsa/test_mtp_schedule_skip.py (new, GPU): a real DSAtrtllmAttentionMetadata over a real DSACacheManager. update_for_spec_dec() inside an index-sharing draft loop leaves the poisoned schedule buffers untouched and sets the marker, on_update_kv_lens() rebuilds and clears it, and outside the draft loop or without index share update_for_spec_dec() rebuilds. - executor/test_preprocess_inputs_kv_lens_hook.py (new, CPU): DecoderRunner._preprocess_inputs on a bare runner calls a decorated hook once and an undecorated or subclass-overridden hook twice when a correction applies, and once when spec decode or the overlap scheduler is off. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../_torch/attention/backends/interface.py | 21 +++ .../backends/sparse/deepseek_v4/metadata.py | 4 +- .../attention/backends/sparse/dsa/indexer.py | 11 ++ .../attention/backends/sparse/dsa/metadata.py | 113 ++++++----- .../engine/runners/decoder/runner.py | 22 ++- .../attention/sparse/dsa/test_dsa_indexer.py | 11 ++ .../sparse/dsa/test_mtp_schedule_skip.py | 176 ++++++++++++++++++ .../test_preprocess_inputs_kv_lens_hook.py | 108 +++++++++++ 8 files changed, 416 insertions(+), 50 deletions(-) create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py create mode 100644 tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 50f08a55569e..2e136b9d8837 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -63,6 +63,27 @@ class AttentionInputType(IntEnum): generation_only = 2 +def idempotent_kv_lens_hook(fn): + """Mark an ``on_update_kv_lens`` override as a pure rebuild from the current lengths. + + Contract: everything the hook writes is derived from the metadata's current + seq/kv lengths and block offsets, so calling it twice equals calling it once. + A caller about to call it again with nothing reading its results in between + may therefore drop the first call (see ``kv_lens_hook_is_idempotent``). + """ + fn.__idempotent_kv_lens_hook__ = True + return fn + + +def kv_lens_hook_is_idempotent(metadata: "AttentionMetadata") -> bool: + """Whether the ``on_update_kv_lens`` that ``metadata`` dispatches to carries the marker. + + Read from the class: a subclass override without the marker is not idempotent. + """ + return getattr( + type(metadata).on_update_kv_lens, "__idempotent_kv_lens_hook__", False) + + @dataclass(kw_only=True) class AttentionMetadata: """ diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 316897a406a6..f5d9597fbebc 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -922,9 +922,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 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index f59fed9faf9f..2a1a3ffbb3c2 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -1179,6 +1179,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] ) @@ -1680,6 +1682,15 @@ def sparse_attn_indexer( metadata.shared_topk_indices[:num_generations, :] ) 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: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index f4707d530b14..edbfa591b907 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -13,6 +13,7 @@ import tensorrt_llm import tensorrt_llm.bindings +from tensorrt_llm._torch.attention.backends.interface import idempotent_kv_lens_hook from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.utils import maybe_compile @@ -120,6 +121,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 @@ -611,7 +616,8 @@ def warmup_selfsampling_topk( "will JIT-compile lazily on first touch instead." ) - def on_update_kv_lens(self) -> None: + @idempotent_kv_lens_hook + 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 the _preprocess_inputs() of model_engine.py. # @@ -620,6 +626,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() @@ -712,47 +722,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_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): @@ -764,9 +735,55 @@ 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_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 @@ -821,8 +838,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): @@ -963,6 +985,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) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 556501ebabb8..5b5ff532fc83 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -15,7 +15,10 @@ import torch -from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionMetadata, + kv_lens_hook_is_idempotent, +) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.compilation.backend import Backend @@ -2712,8 +2715,21 @@ def _preprocess_inputs( Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get("attn_metadata") + # 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: + # An idempotent hook that the correction repeats right after, with + # nothing reading its results in between, is called once. + if attn_meta is not None and not ( + has_kv_lens_correction and kv_lens_hook_is_idempotent(attn_meta) + ): attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: @@ -2739,7 +2755,7 @@ def _preprocess_inputs( :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 diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 9f1ba6b50c6d..24044015a8c3 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -3960,6 +3960,10 @@ def test_indexer_decode_mtp_topk_reuse(step0_mode, batch_size): 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 @@ -4100,12 +4104,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") diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py new file mode 100644 index 000000000000..53ad276b5d20 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""update_for_spec_dec() skips the indexer schedule rebuild only inside an index-sharing draft loop. + +Runs on a real DSAtrtllmAttentionMetadata over a real DSACacheManager: prepare() +fills the DeepGEMM MQA-logits schedule buffers and the 2D indexer KV lengths, the +tests poison them and check which kv-lens updates rebuild them. +""" + +import pytest +import torch +from utils.util import skip_pre_hopper + +from tensorrt_llm._torch.attention.backends.sparse.dsa import ( + DSACacheManager, + DSAtrtllmAttentionMetadata, +) +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig +from tensorrt_llm.mapping import Mapping + +HEAD_DIM = 128 +TOKENS_PER_BLOCK = 64 +MAX_SEQ_LEN = 4096 +KV_LENS = [1000, 1500] +# One draft token so the full-next_n schedule is part of the rebuild. +MAX_DRAFT_TOKENS = 1 + + +def _has_deep_gemm() -> bool: + try: + from tensorrt_llm import deep_gemm + + return deep_gemm is not None + except Exception: + return False + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"), + pytest.mark.skipif(not _has_deep_gemm(), reason="DeepGEMM not available"), + skip_pre_hopper, +] + + +def _make_decode_metadata(index_share: bool) -> DSAtrtllmAttentionMetadata: + """A prepared one-token-per-request generation batch, as draft steps >= 1 see it.""" + batch_size = len(KV_LENS) + sparse_config = DeepSeekSparseAttentionConfig( + index_n_heads=32, + index_head_dim=HEAD_DIM, + index_topk=2048, + indexer_k_dtype="fp8", + skip_indexer_for_short_seqs=False, + index_share_for_mtp_iteration=index_share, + ) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + kv_cache_manager = DSACacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, max_tokens=MAX_SEQ_LEN * batch_size + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=batch_size, + mapping=mapping, + dtype=DataType.HALF, + sparse_attention_config=sparse_config, + ) + request_ids = list(range(batch_size)) + kv_cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=KV_LENS, is_gen=False, prepare_resource=True + ) + cached = [kv_len - 1 for kv_len in KV_LENS] + metadata = DSAtrtllmAttentionMetadata( + seq_lens=torch.ones(batch_size, dtype=torch.int), + request_ids=request_ids, + max_num_requests=batch_size, + num_contexts=0, + prompt_lens=cached, + max_num_tokens=batch_size * (1 + MAX_DRAFT_TOKENS), + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached), + mapping=mapping, + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), + ) + # The draft depth is set the way the engine sets it; this also sizes the + # next_n-dependent buffers. Spec-dec attention stays off, as in draft steps >= 1. + metadata.update_spec_dec_param( + batch_size=batch_size, + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=MAX_DRAFT_TOKENS, + max_total_draft_tokens=MAX_DRAFT_TOKENS, + ) + metadata.prepare() + torch.cuda.synchronize() + assert metadata.sparse_metadata_params.mtp_index_share is index_share + assert (metadata.shared_topk_indices is not None) is index_share + return metadata + + +def _schedule_buffers(metadata): + """What on_update_kv_lens(skip_indexer_schedule=True) leaves untouched.""" + return { + "scheduler_metadata_buffer": metadata.scheduler_metadata_buffer, + "scheduler_metadata_buffer_full_next_n": metadata.scheduler_metadata_buffer_full_next_n, + "kv_lens_cuda_2d": metadata.kv_lens_cuda_2d[: metadata.num_generations], + } + + +def _snapshot(metadata): + torch.cuda.synchronize() + reference = {name: buf.clone() for name, buf in _schedule_buffers(metadata).items()} + # The poison value must be distinguishable from a real rebuild. + for name, buf in reference.items(): + assert not torch.equal(buf, torch.full_like(buf, -1)), name + return reference + + +def _poison(metadata): + for buf in _schedule_buffers(metadata).values(): + buf.fill_(-1) + + +def _assert_untouched(metadata): + torch.cuda.synchronize() + for name, buf in _schedule_buffers(metadata).items(): + assert torch.equal(buf, torch.full_like(buf, -1)), name + + +def _assert_rebuilt(metadata, reference): + torch.cuda.synchronize() + for name, buf in _schedule_buffers(metadata).items(): + assert torch.equal(buf, reference[name]), name + + +def test_index_sharing_draft_loop_skips_the_schedule_until_the_next_full_rebuild(): + metadata = _make_decode_metadata(index_share=True) + reference = _snapshot(metadata) + + metadata.set_in_mtp_draft_loop(True) + _poison(metadata) + metadata.update_for_spec_dec() + assert metadata.indexer_schedule_stale is True + _assert_untouched(metadata) + + # The public hook (the pre-draft refresh and the target forward) is the full rebuild. + metadata.on_update_kv_lens() + assert metadata.indexer_schedule_stale is False + _assert_rebuilt(metadata, reference) + + +@pytest.mark.parametrize( + "in_draft_loop,index_share", + [(False, True), (True, False)], + ids=["outside_the_draft_loop", "draft_loop_without_index_share"], +) +def test_update_for_spec_dec_rebuilds_outside_an_index_sharing_draft_loop( + in_draft_loop, index_share +): + metadata = _make_decode_metadata(index_share=index_share) + reference = _snapshot(metadata) + + metadata.set_in_mtp_draft_loop(in_draft_loop) + _poison(metadata) + metadata.update_for_spec_dec() + assert metadata.indexer_schedule_stale is False + _assert_rebuilt(metadata, reference) diff --git a/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py b/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py new file mode 100644 index 000000000000..88965a7b6e02 --- /dev/null +++ b/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""_preprocess_inputs calls an idempotent on_update_kv_lens once instead of twice. + +Under the overlap scheduler with speculative decoding the engine corrects +kv_lens_cuda and calls the hook again right after. A hook marked with +idempotent_kv_lens_hook is then called once, after the correction; an unmarked +hook, or a run without a correction, keeps the original calls. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.interface import ( + idempotent_kv_lens_hook, + kv_lens_hook_is_idempotent, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + +KV_LENS = torch.tensor([10, 20, 30], dtype=torch.int32) +KV_LENS_OFFSETS = torch.tensor([1, 2, 0], dtype=torch.int32) +TOKENS_PER_REQUEST = 3 +NUM_TOKENS = KV_LENS.numel() * TOKENS_PER_REQUEST + + +class _Metadata: + """The slice of TrtllmAttentionMetadata that _preprocess_inputs touches.""" + + def __init__(self): + self.kv_cache_manager = object() + self.num_seqs = KV_LENS.numel() + self.num_contexts = 0 + self.num_generations = KV_LENS.numel() + self.num_ctx_tokens = 0 + self.num_chunked_ctx_requests = 0 + self.kv_lens_cuda = KV_LENS.clone() + # kv_lens_cuda as seen by each on_update_kv_lens() call. + self.kv_lens_at_hook = [] + + def on_update_kv_lens(self): + self.kv_lens_at_hook.append(self.kv_lens_cuda.clone()) + + +class _IdempotentMetadata(_Metadata): + @idempotent_kv_lens_hook + def on_update_kv_lens(self): + super().on_update_kv_lens() + + +class _UnmarkedSubclassMetadata(_IdempotentMetadata): + """Overrides a marked hook without the marker (DeepSeek-V4 over DSA).""" + + def on_update_kv_lens(self): + super().on_update_kv_lens() + + +def _run_preprocess_inputs(metadata, *, spec_decode, overlap): + engine = object.__new__(DecoderRunner) + engine._config = SimpleNamespace(disable_overlap_scheduler=not overlap) + engine.previous_pos_id_offsets_cuda = torch.arange(NUM_TOKENS, dtype=torch.int32) + engine.previous_kv_lens_offsets_cuda = KV_LENS_OFFSETS.clone() + engine.mapping = SimpleNamespace(has_cp_helix=lambda: False) + engine.guided_decoder = None + inputs = { + "input_ids": torch.zeros(NUM_TOKENS, dtype=torch.int32), + "position_ids": torch.zeros(1, NUM_TOKENS, dtype=torch.int32), + "attn_metadata": metadata, + } + engine._preprocess_inputs(inputs, enable_spec_decode=spec_decode, runtime_draft_len=0) + return inputs + + +def test_kv_lens_hook_is_idempotent_reads_the_dispatched_hook(): + assert not kv_lens_hook_is_idempotent(_Metadata()) + assert kv_lens_hook_is_idempotent(_IdempotentMetadata()) + assert not kv_lens_hook_is_idempotent(_UnmarkedSubclassMetadata()) + + +@pytest.mark.parametrize( + "metadata_cls", [_Metadata, _IdempotentMetadata, _UnmarkedSubclassMetadata] +) +def test_correction_repeats_the_hook_only_when_it_is_not_idempotent(metadata_cls): + metadata = metadata_cls() + inputs = _run_preprocess_inputs(metadata, spec_decode=True, overlap=True) + corrected = KV_LENS + KV_LENS_OFFSETS + # The correction itself is untouched. + assert torch.equal(metadata.kv_lens_cuda, corrected) + assert torch.equal(inputs["position_ids"][0], torch.arange(NUM_TOKENS, dtype=torch.int32)) + if kv_lens_hook_is_idempotent(metadata): + # Only the call after the correction survives. + assert len(metadata.kv_lens_at_hook) == 1 + assert torch.equal(metadata.kv_lens_at_hook[0], corrected) + else: + # Once before and once after the correction. + assert len(metadata.kv_lens_at_hook) == 2 + assert torch.equal(metadata.kv_lens_at_hook[0], KV_LENS) + assert torch.equal(metadata.kv_lens_at_hook[1], corrected) + + +@pytest.mark.parametrize("spec_decode,overlap", [(False, True), (True, False), (False, False)]) +def test_single_hook_call_is_kept_when_no_correction_follows(spec_decode, overlap): + metadata = _IdempotentMetadata() + _run_preprocess_inputs(metadata, spec_decode=spec_decode, overlap=overlap) + assert len(metadata.kv_lens_at_hook) == 1 + assert torch.equal(metadata.kv_lens_at_hook[0], KV_LENS) + assert torch.equal(metadata.kv_lens_cuda, KV_LENS) From 76d46a8aaf58d5124a655ca8085c5bf84b165889 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:44:49 -0700 Subject: [PATCH 2/5] [None][perf] Drop the first on_update_kv_lens() call before the kv-lens correction unconditionally Replace the idempotent_kv_lens_hook marker and kv_lens_hook_is_idempotent() with a plain rule: when DecoderRunner._preprocess_inputs is about to correct kv_lens_cuda and call on_update_kv_lens() again, it skips the call before the correction for every backend. Every implementation on main overwrites its outputs from the current lengths and nothing reads them between the two calls: AttentionMetadata (no-op), TrtllmAttentionMetadata (flag invalidation + Fp4MlaState), DSA, DeepSeek-V4 (compressed-KV metadata and indices), QSA (token mapping), and both MiniMax-M3 metadata classes. The CPU test now checks one call after the correction, and the original single call when no correction applies. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../_torch/attention/backends/interface.py | 21 -------- .../attention/backends/sparse/dsa/metadata.py | 2 - .../engine/runners/decoder/runner.py | 13 ++--- .../test_preprocess_inputs_kv_lens_hook.py | 52 ++++--------------- 4 files changed, 13 insertions(+), 75 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 2e136b9d8837..50f08a55569e 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -63,27 +63,6 @@ class AttentionInputType(IntEnum): generation_only = 2 -def idempotent_kv_lens_hook(fn): - """Mark an ``on_update_kv_lens`` override as a pure rebuild from the current lengths. - - Contract: everything the hook writes is derived from the metadata's current - seq/kv lengths and block offsets, so calling it twice equals calling it once. - A caller about to call it again with nothing reading its results in between - may therefore drop the first call (see ``kv_lens_hook_is_idempotent``). - """ - fn.__idempotent_kv_lens_hook__ = True - return fn - - -def kv_lens_hook_is_idempotent(metadata: "AttentionMetadata") -> bool: - """Whether the ``on_update_kv_lens`` that ``metadata`` dispatches to carries the marker. - - Read from the class: a subclass override without the marker is not idempotent. - """ - return getattr( - type(metadata).on_update_kv_lens, "__idempotent_kv_lens_hook__", False) - - @dataclass(kw_only=True) class AttentionMetadata: """ diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index edbfa591b907..98025fab8a49 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -13,7 +13,6 @@ import tensorrt_llm import tensorrt_llm.bindings -from tensorrt_llm._torch.attention.backends.interface import idempotent_kv_lens_hook from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.utils import maybe_compile @@ -616,7 +615,6 @@ def warmup_selfsampling_topk( "will JIT-compile lazily on first touch instead." ) - @idempotent_kv_lens_hook 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 the _preprocess_inputs() of model_engine.py. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 5b5ff532fc83..deb47da1c324 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -15,10 +15,7 @@ import torch -from tensorrt_llm._torch.attention.backends.interface import ( - AttentionMetadata, - kv_lens_hook_is_idempotent, -) +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.compilation.backend import Backend @@ -2725,11 +2722,9 @@ def _preprocess_inputs( and hasattr(attn_meta, "kv_lens_cuda") ) # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - # An idempotent hook that the correction repeats right after, with - # nothing reading its results in between, is called once. - if attn_meta is not None and not ( - has_kv_lens_correction and kv_lens_hook_is_idempotent(attn_meta) - ): + # 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: diff --git a/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py b/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py index 88965a7b6e02..6c764ddba502 100644 --- a/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py +++ b/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py @@ -1,11 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""_preprocess_inputs calls an idempotent on_update_kv_lens once instead of twice. +"""_preprocess_inputs calls on_update_kv_lens once per step. -Under the overlap scheduler with speculative decoding the engine corrects -kv_lens_cuda and calls the hook again right after. A hook marked with -idempotent_kv_lens_hook is then called once, after the correction; an unmarked -hook, or a run without a correction, keeps the original calls. +Under the overlap scheduler with speculative decoding the runner corrects +kv_lens_cuda and calls the hook after the correction; the call before it is +dropped. Without a correction the single call before it is kept. """ from types import SimpleNamespace @@ -13,10 +12,6 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.interface import ( - idempotent_kv_lens_hook, - kv_lens_hook_is_idempotent, -) from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner KV_LENS = torch.tensor([10, 20, 30], dtype=torch.int32) @@ -43,19 +38,6 @@ def on_update_kv_lens(self): self.kv_lens_at_hook.append(self.kv_lens_cuda.clone()) -class _IdempotentMetadata(_Metadata): - @idempotent_kv_lens_hook - def on_update_kv_lens(self): - super().on_update_kv_lens() - - -class _UnmarkedSubclassMetadata(_IdempotentMetadata): - """Overrides a marked hook without the marker (DeepSeek-V4 over DSA).""" - - def on_update_kv_lens(self): - super().on_update_kv_lens() - - def _run_preprocess_inputs(metadata, *, spec_decode, overlap): engine = object.__new__(DecoderRunner) engine._config = SimpleNamespace(disable_overlap_scheduler=not overlap) @@ -72,36 +54,20 @@ def _run_preprocess_inputs(metadata, *, spec_decode, overlap): return inputs -def test_kv_lens_hook_is_idempotent_reads_the_dispatched_hook(): - assert not kv_lens_hook_is_idempotent(_Metadata()) - assert kv_lens_hook_is_idempotent(_IdempotentMetadata()) - assert not kv_lens_hook_is_idempotent(_UnmarkedSubclassMetadata()) - - -@pytest.mark.parametrize( - "metadata_cls", [_Metadata, _IdempotentMetadata, _UnmarkedSubclassMetadata] -) -def test_correction_repeats_the_hook_only_when_it_is_not_idempotent(metadata_cls): - metadata = metadata_cls() +def test_hook_runs_once_after_the_correction(): + metadata = _Metadata() inputs = _run_preprocess_inputs(metadata, spec_decode=True, overlap=True) corrected = KV_LENS + KV_LENS_OFFSETS # The correction itself is untouched. assert torch.equal(metadata.kv_lens_cuda, corrected) assert torch.equal(inputs["position_ids"][0], torch.arange(NUM_TOKENS, dtype=torch.int32)) - if kv_lens_hook_is_idempotent(metadata): - # Only the call after the correction survives. - assert len(metadata.kv_lens_at_hook) == 1 - assert torch.equal(metadata.kv_lens_at_hook[0], corrected) - else: - # Once before and once after the correction. - assert len(metadata.kv_lens_at_hook) == 2 - assert torch.equal(metadata.kv_lens_at_hook[0], KV_LENS) - assert torch.equal(metadata.kv_lens_at_hook[1], corrected) + assert len(metadata.kv_lens_at_hook) == 1 + assert torch.equal(metadata.kv_lens_at_hook[0], corrected) @pytest.mark.parametrize("spec_decode,overlap", [(False, True), (True, False), (False, False)]) def test_single_hook_call_is_kept_when_no_correction_follows(spec_decode, overlap): - metadata = _IdempotentMetadata() + metadata = _Metadata() _run_preprocess_inputs(metadata, spec_decode=spec_decode, overlap=overlap) assert len(metadata.kv_lens_at_hook) == 1 assert torch.equal(metadata.kv_lens_at_hook[0], KV_LENS) From 2008747de7390788f6c4cbbfc4b97933de383f37 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:49:21 -0700 Subject: [PATCH 3/5] [None][test] Trim the schedule-skip and kv-lens-hook tests to the essential checks test_mtp_schedule_skip.py: one test on one index-sharing metadata. Outside the draft loop update_for_spec_dec() rebuilds the schedule; inside it the schedule is left as is and marked stale; on_update_kv_lens() rebuilds it and clears the marker. test_preprocess_inputs_kv_lens_hook.py: one parametrized test. The hook runs exactly once per step, after the correction when one applies. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../sparse/dsa/test_mtp_schedule_skip.py | 84 ++++++------------- .../test_preprocess_inputs_kv_lens_hook.py | 23 ++--- 2 files changed, 32 insertions(+), 75 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py index 53ad276b5d20..4e0bf5e995eb 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py @@ -1,11 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""update_for_spec_dec() skips the indexer schedule rebuild only inside an index-sharing draft loop. - -Runs on a real DSAtrtllmAttentionMetadata over a real DSACacheManager: prepare() -fills the DeepGEMM MQA-logits schedule buffers and the 2D indexer KV lengths, the -tests poison them and check which kv-lens updates rebuild them. -""" +"""update_for_spec_dec() skips the indexer schedule rebuild only inside an index-sharing draft loop.""" import pytest import torch @@ -46,7 +41,7 @@ def _has_deep_gemm() -> bool: ] -def _make_decode_metadata(index_share: bool) -> DSAtrtllmAttentionMetadata: +def _make_decode_metadata() -> DSAtrtllmAttentionMetadata: """A prepared one-token-per-request generation batch, as draft steps >= 1 see it.""" batch_size = len(KV_LENS) sparse_config = DeepSeekSparseAttentionConfig( @@ -55,7 +50,7 @@ def _make_decode_metadata(index_share: bool) -> DSAtrtllmAttentionMetadata: index_topk=2048, indexer_k_dtype="fp8", skip_indexer_for_short_seqs=False, - index_share_for_mtp_iteration=index_share, + index_share_for_mtp_iteration=True, ) mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) kv_cache_manager = DSACacheManager( @@ -102,75 +97,46 @@ def _make_decode_metadata(index_share: bool) -> DSAtrtllmAttentionMetadata: ) metadata.prepare() torch.cuda.synchronize() - assert metadata.sparse_metadata_params.mtp_index_share is index_share - assert (metadata.shared_topk_indices is not None) is index_share + assert metadata.sparse_metadata_params.mtp_index_share return metadata def _schedule_buffers(metadata): """What on_update_kv_lens(skip_indexer_schedule=True) leaves untouched.""" - return { - "scheduler_metadata_buffer": metadata.scheduler_metadata_buffer, - "scheduler_metadata_buffer_full_next_n": metadata.scheduler_metadata_buffer_full_next_n, - "kv_lens_cuda_2d": metadata.kv_lens_cuda_2d[: metadata.num_generations], - } - - -def _snapshot(metadata): - torch.cuda.synchronize() - reference = {name: buf.clone() for name, buf in _schedule_buffers(metadata).items()} - # The poison value must be distinguishable from a real rebuild. - for name, buf in reference.items(): - assert not torch.equal(buf, torch.full_like(buf, -1)), name - return reference + return [ + metadata.scheduler_metadata_buffer, + metadata.scheduler_metadata_buffer_full_next_n, + metadata.kv_lens_cuda_2d[: metadata.num_generations], + ] def _poison(metadata): - for buf in _schedule_buffers(metadata).values(): + for buf in _schedule_buffers(metadata): buf.fill_(-1) -def _assert_untouched(metadata): +def _is_poisoned(metadata): torch.cuda.synchronize() - for name, buf in _schedule_buffers(metadata).items(): - assert torch.equal(buf, torch.full_like(buf, -1)), name + return [bool((buf == -1).all()) for buf in _schedule_buffers(metadata)] -def _assert_rebuilt(metadata, reference): - torch.cuda.synchronize() - for name, buf in _schedule_buffers(metadata).items(): - assert torch.equal(buf, reference[name]), name - +def test_schedule_is_skipped_only_inside_the_draft_loop(): + metadata = _make_decode_metadata() -def test_index_sharing_draft_loop_skips_the_schedule_until_the_next_full_rebuild(): - metadata = _make_decode_metadata(index_share=True) - reference = _snapshot(metadata) + # Outside the draft loop: full rebuild. + _poison(metadata) + metadata.update_for_spec_dec() + assert not metadata.indexer_schedule_stale + assert _is_poisoned(metadata) == [False, False, False] + # Inside the index-sharing draft loop: schedule left as is and marked stale. metadata.set_in_mtp_draft_loop(True) _poison(metadata) metadata.update_for_spec_dec() - assert metadata.indexer_schedule_stale is True - _assert_untouched(metadata) + assert metadata.indexer_schedule_stale + assert _is_poisoned(metadata) == [True, True, True] - # The public hook (the pre-draft refresh and the target forward) is the full rebuild. + # The next full rebuild (target forward) restores it and clears the marker. metadata.on_update_kv_lens() - assert metadata.indexer_schedule_stale is False - _assert_rebuilt(metadata, reference) - - -@pytest.mark.parametrize( - "in_draft_loop,index_share", - [(False, True), (True, False)], - ids=["outside_the_draft_loop", "draft_loop_without_index_share"], -) -def test_update_for_spec_dec_rebuilds_outside_an_index_sharing_draft_loop( - in_draft_loop, index_share -): - metadata = _make_decode_metadata(index_share=index_share) - reference = _snapshot(metadata) - - metadata.set_in_mtp_draft_loop(in_draft_loop) - _poison(metadata) - metadata.update_for_spec_dec() - assert metadata.indexer_schedule_stale is False - _assert_rebuilt(metadata, reference) + assert not metadata.indexer_schedule_stale + assert _is_poisoned(metadata) == [False, False, False] diff --git a/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py b/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py index 6c764ddba502..465c7caa64ec 100644 --- a/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py +++ b/tests/unittest/_torch/executor/test_preprocess_inputs_kv_lens_hook.py @@ -51,24 +51,15 @@ def _run_preprocess_inputs(metadata, *, spec_decode, overlap): "attn_metadata": metadata, } engine._preprocess_inputs(inputs, enable_spec_decode=spec_decode, runtime_draft_len=0) - return inputs -def test_hook_runs_once_after_the_correction(): - metadata = _Metadata() - inputs = _run_preprocess_inputs(metadata, spec_decode=True, overlap=True) - corrected = KV_LENS + KV_LENS_OFFSETS - # The correction itself is untouched. - assert torch.equal(metadata.kv_lens_cuda, corrected) - assert torch.equal(inputs["position_ids"][0], torch.arange(NUM_TOKENS, dtype=torch.int32)) - assert len(metadata.kv_lens_at_hook) == 1 - assert torch.equal(metadata.kv_lens_at_hook[0], corrected) - - -@pytest.mark.parametrize("spec_decode,overlap", [(False, True), (True, False), (False, False)]) -def test_single_hook_call_is_kept_when_no_correction_follows(spec_decode, overlap): +@pytest.mark.parametrize( + "spec_decode,overlap,corrected", + [(True, True, True), (True, False, False), (False, True, False)], +) +def test_hook_runs_once_per_step(spec_decode, overlap, corrected): metadata = _Metadata() _run_preprocess_inputs(metadata, spec_decode=spec_decode, overlap=overlap) + expected = KV_LENS + KV_LENS_OFFSETS if corrected else KV_LENS assert len(metadata.kv_lens_at_hook) == 1 - assert torch.equal(metadata.kv_lens_at_hook[0], KV_LENS) - assert torch.equal(metadata.kv_lens_cuda, KV_LENS) + assert torch.equal(metadata.kv_lens_at_hook[0], expected) From e3361c2e77d62fd5bde39b75d40b6ec743ca2cd1 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Sun, 1 Nov 2026 12:22:10 -0500 Subject: [PATCH 4/5] [None][test] Check refreshed KV-lens buffers on schedule skip and shut down the cache manager - test_mtp_schedule_skip: advance the KV lengths inside the draft loop and check that slot mappings and gen indptrs are still rewritten (equal to a full rebuild, different from the pre-advance values) while the schedule buffers stay untouched. Build the metadata in a fixture that shuts the DSACacheManager down on teardown. - test_preprocess_inputs_kv_lens_hook: mark cpu_only so the CPU stage runs it. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../sparse/dsa/test_mtp_schedule_skip.py | 110 ++++++++++++------ .../test_preprocess_inputs_kv_lens_hook.py | 2 + 2 files changed, 75 insertions(+), 37 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py index 4e0bf5e995eb..1b00ee952e3d 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py @@ -41,8 +41,13 @@ def _has_deep_gemm() -> bool: ] -def _make_decode_metadata() -> DSAtrtllmAttentionMetadata: - """A prepared one-token-per-request generation batch, as draft steps >= 1 see it.""" +@pytest.fixture +def decode_metadata(): + """A prepared one-token-per-request generation batch, as draft steps >= 1 see it. + + The cache manager is shut down on teardown: it owns CUDA reservations that a + garbage-collector free would release on a thread without a current CUDA context. + """ batch_size = len(KV_LENS) sparse_config = DeepSeekSparseAttentionConfig( index_n_heads=32, @@ -68,37 +73,40 @@ def _make_decode_metadata() -> DSAtrtllmAttentionMetadata: dtype=DataType.HALF, sparse_attention_config=sparse_config, ) - request_ids = list(range(batch_size)) - kv_cache_manager.add_dummy_requests( - request_ids=request_ids, token_nums=KV_LENS, is_gen=False, prepare_resource=True - ) - cached = [kv_len - 1 for kv_len in KV_LENS] - metadata = DSAtrtllmAttentionMetadata( - seq_lens=torch.ones(batch_size, dtype=torch.int), - request_ids=request_ids, - max_num_requests=batch_size, - num_contexts=0, - prompt_lens=cached, - max_num_tokens=batch_size * (1 + MAX_DRAFT_TOKENS), - kv_cache_manager=kv_cache_manager, - kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached), - mapping=mapping, - sparse_metadata_params=sparse_config.to_sparse_metadata_params(), - ) - # The draft depth is set the way the engine sets it; this also sizes the - # next_n-dependent buffers. Spec-dec attention stays off, as in draft steps >= 1. - metadata.update_spec_dec_param( - batch_size=batch_size, - is_spec_decoding_enabled=False, - is_spec_dec_tree=False, - is_spec_dec_dynamic_tree=False, - max_draft_len=MAX_DRAFT_TOKENS, - max_total_draft_tokens=MAX_DRAFT_TOKENS, - ) - metadata.prepare() - torch.cuda.synchronize() - assert metadata.sparse_metadata_params.mtp_index_share - return metadata + try: + request_ids = list(range(batch_size)) + kv_cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=KV_LENS, is_gen=False, prepare_resource=True + ) + cached = [kv_len - 1 for kv_len in KV_LENS] + metadata = DSAtrtllmAttentionMetadata( + seq_lens=torch.ones(batch_size, dtype=torch.int), + request_ids=request_ids, + max_num_requests=batch_size, + num_contexts=0, + prompt_lens=cached, + max_num_tokens=batch_size * (1 + MAX_DRAFT_TOKENS), + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=cached), + mapping=mapping, + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), + ) + # The draft depth is set the way the engine sets it; this also sizes the + # next_n-dependent buffers. Spec-dec attention stays off, as in draft steps >= 1. + metadata.update_spec_dec_param( + batch_size=batch_size, + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=MAX_DRAFT_TOKENS, + max_total_draft_tokens=MAX_DRAFT_TOKENS, + ) + metadata.prepare() + torch.cuda.synchronize() + assert metadata.sparse_metadata_params.mtp_index_share + yield metadata + finally: + kv_cache_manager.shutdown() def _schedule_buffers(metadata): @@ -120,23 +128,51 @@ def _is_poisoned(metadata): return [bool((buf == -1).all()) for buf in _schedule_buffers(metadata)] -def test_schedule_is_skipped_only_inside_the_draft_loop(): - metadata = _make_decode_metadata() +def _kv_lens_buffers(metadata): + """What on_update_kv_lens() refreshes from the KV lengths even when it skips the schedule.""" + num_tokens = metadata.num_tokens + num_generations = metadata.num_generations + return [ + metadata.slot_mapping_fp8[:num_tokens], + metadata.slot_mapping_scale[:num_tokens], + metadata.gen_kv_indptr[1 : num_generations + 1], + metadata.gen_cached_token_indptr[1 : num_generations + 1], + ] + + +def _snapshot(buffers): + torch.cuda.synchronize() + return [buf.clone() for buf in buffers] + + +def test_schedule_is_skipped_only_inside_the_draft_loop(decode_metadata): + metadata = decode_metadata # Outside the draft loop: full rebuild. _poison(metadata) metadata.update_for_spec_dec() assert not metadata.indexer_schedule_stale assert _is_poisoned(metadata) == [False, False, False] + before_advance = _snapshot(_kv_lens_buffers(metadata)) - # Inside the index-sharing draft loop: schedule left as is and marked stale. + # Inside the index-sharing draft loop, after a draft step appends a token: the schedule + # is left as is and marked stale, but the KV-lens-dependent buffers are still rewritten. metadata.set_in_mtp_draft_loop(True) + metadata.kv_lens_cuda[: metadata.num_generations] += 1 _poison(metadata) + for buf in _kv_lens_buffers(metadata): + buf.fill_(-1) metadata.update_for_spec_dec() assert metadata.indexer_schedule_stale assert _is_poisoned(metadata) == [True, True, True] + in_draft_loop = _snapshot(_kv_lens_buffers(metadata)) - # The next full rebuild (target forward) restores it and clears the marker. + # The next full rebuild (target forward) restores the schedule, clears the marker and + # writes the same slot mappings and gen indptrs, which follow the advanced lengths. metadata.on_update_kv_lens() assert not metadata.indexer_schedule_stale assert _is_poisoned(metadata) == [False, False, False] + full_rebuild = _snapshot(_kv_lens_buffers(metadata)) + for skipped, rebuilt, stale in zip(in_draft_loop, full_rebuild, before_advance): + assert torch.equal(skipped, rebuilt) + assert not torch.equal(skipped, stale) diff --git a/tests/unittest/_torch/pyexecutor/test_preprocess_inputs_kv_lens_hook.py b/tests/unittest/_torch/pyexecutor/test_preprocess_inputs_kv_lens_hook.py index ce91652fb53d..964725a6a423 100644 --- a/tests/unittest/_torch/pyexecutor/test_preprocess_inputs_kv_lens_hook.py +++ b/tests/unittest/_torch/pyexecutor/test_preprocess_inputs_kv_lens_hook.py @@ -14,6 +14,8 @@ from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer +pytestmark = pytest.mark.cpu_only + KV_LENS = torch.tensor([10, 20, 30], dtype=torch.int32) KV_LENS_OFFSETS = torch.tensor([1, 2, 0], dtype=torch.int32) TOKENS_PER_REQUEST = 3 From 3ac132e04ff683ea62e302a8063640656451bc6d Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Sun, 1 Nov 2026 23:06:40 -0500 Subject: [PATCH 5/5] [None][test] Drop the removed prepare_resource argument from test_mtp_schedule_skip add_dummy_requests() lost prepare_resource in #20253 and now always allocates the KV cache eagerly, so the fixture failed during setup with a TypeError on every GPU stage. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../_torch/attention/sparse/dsa/test_mtp_schedule_skip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py index 1b00ee952e3d..dbc3e055b9ef 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py @@ -76,7 +76,7 @@ def decode_metadata(): try: request_ids = list(range(batch_size)) kv_cache_manager.add_dummy_requests( - request_ids=request_ids, token_nums=KV_LENS, is_gen=False, prepare_resource=True + request_ids=request_ids, token_nums=KV_LENS, is_gen=False ) cached = [kv_len - 1 for kv_len in KV_LENS] metadata = DSAtrtllmAttentionMetadata(