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 4ed9a04e968b..e36e198abb6a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/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 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index 70d63a9f199f..9f1645b5c32a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/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: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index e0773630dcb3..2509ae65f9ac 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index a9fc0625740e..350a978af852 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/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 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 44b74c5fa735..21f09988aa6f 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/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") 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..dbc3e055b9ef --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_mtp_schedule_skip.py @@ -0,0 +1,178 @@ +# 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.""" + +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, +] + + +@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, + index_head_dim=HEAD_DIM, + index_topk=2048, + indexer_k_dtype="fp8", + skip_indexer_for_short_seqs=False, + index_share_for_mtp_iteration=True, + ) + 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, + ) + try: + request_ids = list(range(batch_size)) + kv_cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=KV_LENS, is_gen=False + ) + 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): + """What on_update_kv_lens(skip_indexer_schedule=True) leaves untouched.""" + 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): + buf.fill_(-1) + + +def _is_poisoned(metadata): + torch.cuda.synchronize() + return [bool((buf == -1).all()) for buf in _schedule_buffers(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, 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 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 new file mode 100644 index 000000000000..964725a6a423 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_preprocess_inputs_kv_lens_hook.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""preprocess_inputs calls on_update_kv_lens once per step. + +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 + +import pytest +import torch + +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 +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()) + + +def _run_preprocess_inputs(metadata, *, spec_decode, overlap): + engine = object.__new__(DecoderInputPreparer) + engine._config = SimpleNamespace(disable_overlap_scheduler=not overlap) + engine._buffers = SimpleNamespace( + previous_pos_id_offsets_cuda=torch.arange(NUM_TOKENS, dtype=torch.int32), + 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) + + +@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], expected)