diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 585b796a6dfb..38e8b0b39515 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -215,7 +215,12 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): if self._ctx_buf_inited: return - max_batch = spec_metadata.max_num_requests + # Worker-owned and allocated once, then reused for every later batch + # shape, so this must span the full seq-slot pool. max_num_requests is + # shrunk to the captured graph bucket by create_cuda_graph_metadata, + # which would pin the pool to whichever bucket drafts first and leave + # _dummy_slot aliasing a live request's row. + max_batch = spec_metadata.num_seq_slots or spec_metadata.max_num_requests # Prefer runtime max_seq_len over max_position_embeddings: YaRN # models advertise 100k+ positions, which would OOM the ctx buffer diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 7f43f103ec4e..9c8edfcb7a53 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -266,7 +266,10 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: if self._win_inited: return - max_batch = spec_metadata.max_num_requests + # Worker-owned and allocated once, so this must span the full seq-slot + # pool rather than max_num_requests, which create_cuda_graph_metadata + # shrinks to the captured graph bucket (see the DFlash counterpart). + max_batch = spec_metadata.num_seq_slots or spec_metadata.max_num_requests num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) head_dim = int(draft_model._attn_params["head_dim"]) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index bffa8833058c..a3b6e638686a 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -206,6 +206,7 @@ def get_spec_metadata(spec_config, dtype=model_config.torch_dtype, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, + num_seq_slots=num_seq_slots, draft_vocab_size=draft_vocab_size, ) if spec_config.spec_dec_mode.is_dspark(): @@ -221,6 +222,7 @@ def get_spec_metadata(spec_config, dtype=model_config.torch_dtype, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, + num_seq_slots=num_seq_slots, draft_vocab_size=draft_vocab_size, ) if spec_config.spec_dec_mode.is_draft_target_one_model(): diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index acc7b260a7cb..91bb998e9516 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -406,8 +406,6 @@ unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py: unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TRTLLMSampler-cuda_graph_and_overlap-None-1-1-True-True-False] SKIP (https://nvbugs/6463819) unittest/_torch/sampler/test_trtllm_sampler.py::test_trtllm_sampler_best_of_with_logprobs SKIP (https://nvbugs/6487837) unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py::test_no_topk_matches_full[0.9] SKIP (https://nvbugs/6550099) -unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[False] SKIP (https://nvbugs/6535767) -unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[True] SKIP (https://nvbugs/6535767) unittest/_torch/speculative/hw_agnostic/test_ngram.py::test_llama_ngram[True-True-TRTLLM] SKIP (https://nvbugs/6507102) unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py::test_fp8_rowwise_linear[dtype1] SKIP (https://nvbugs/6301807) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingDSv3-swiglu-1024-1024-1] SKIP (https://nvbugs/5908070) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py index 3258e9065b56..716bc341a745 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py @@ -16,13 +16,17 @@ import os import sys import unittest +from types import SimpleNamespace import pytest import torch from utils.llm_data import llm_models_root from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch.speculative.dflash import DFlashWorker +from tensorrt_llm._torch.speculative.utils import get_spec_metadata from tensorrt_llm.llmapi import CudaGraphConfig, DFlashDecodingConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping sys.path.append(os.path.join(os.path.dirname(__file__), "..")) @@ -33,6 +37,75 @@ ] +def test_dflash_graph_bucket_uses_full_seq_slot_pool(): + """A small graph bucket must not shrink the persistent context pool.""" + num_seq_slots = 5 + spec_config = DFlashDecodingConfig( + max_draft_len=4, + target_layer_ids=[0], + ) + metadata = get_spec_metadata( + spec_config, + SimpleNamespace(hidden_size=4, torch_dtype=torch.bfloat16, vocab_size=32), + max_num_requests=num_seq_slots, + max_num_tokens=8, + num_seq_slots=num_seq_slots, + ).create_cuda_graph_metadata(max_batch_size=2) + + class DraftModel: + block_size = 5 + config = SimpleNamespace(max_position_embeddings=8) + fc = SimpleNamespace(weight=torch.empty(0, dtype=torch.bfloat16)) + hidden_norm = object() + _num_attn_layers = 1 + _num_kv_heads = 2 + _head_dim = 4 + + def _build_fused_kv_buffers(self): + pass + + def project_target_hidden(self, hidden_states): + return hidden_states + + def precompute_context_kv(self, hidden_states, position_ids): + shape = (hidden_states.shape[0], 1, 2, 4) + return ( + torch.zeros(shape, dtype=torch.bfloat16, device="cuda"), + torch.zeros(shape, dtype=torch.bfloat16, device="cuda"), + ) + + worker = DFlashWorker(spec_config, Mapping()) + draft_model = DraftModel() + attn_metadata = SimpleNamespace( + max_seq_len=8, + num_ctx_tokens=1, + num_contexts=1, + _seq_lens=[1], + ) + worker._lazy_init_ctx_buffers(draft_model, metadata, attn_metadata) + + assert metadata.max_num_requests == 2 < metadata.num_seq_slots + assert worker._ctx_k_buf.shape == (6, 1, 13, 2, 4) + assert worker._ctx_v_buf.shape == worker._ctx_k_buf.shape + assert worker._ctx_len.shape == (6,) + assert worker._batch_to_slot.shape == (num_seq_slots,) + assert worker._dummy_slot == num_seq_slots + assert list(worker._free_slots) == list(range(num_seq_slots)) + + metadata.request_ids = [42] + worker._store_prefill_context( + draft_model, + metadata, + attn_metadata, + torch.tensor([0], device="cuda"), + total_target_tokens=1, + ) + live_slot = worker._req_to_slot[42] + assert live_slot != worker._dummy_slot + assert worker._dummy_slot not in worker._free_slots + assert list(worker._free_slots) == [1, 2, 3, 4] + + def _make_llm_config( target_model_dir: str, dflash_model_dir: str, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 4625d1b0115a..8af124afd4c5 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -28,6 +28,7 @@ from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSparkWorker from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.speculative.utils import get_spec_metadata pytestmark = pytest.mark.skipif( not torch.cuda.is_available(), reason="DSpark metadata/worker allocate CUDA buffers" @@ -134,6 +135,41 @@ def test_worker_lazy_init_window_buffers(): assert id(worker._kv_windows) == buf_id +def test_worker_graph_bucket_uses_full_seq_slot_pool(): + """A small graph bucket must not shrink the persistent rolling-window pool.""" + num_seq_slots = 5 + spec_config = types.SimpleNamespace( + max_draft_len=5, + tokens_per_gen_step=6, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + target_layer_ids=[], + ) + metadata = get_spec_metadata( + spec_config, + types.SimpleNamespace(hidden_size=HIDDEN, torch_dtype=torch.bfloat16, vocab_size=32), + max_num_requests=num_seq_slots, + max_num_tokens=8, + num_seq_slots=num_seq_slots, + ).create_cuda_graph_metadata(max_batch_size=2) + worker = _make_worker() + worker._lazy_init( + _fake_draft_model(num_stages=1, window_size=8, head_dim=4), + metadata, + ) + + assert metadata.max_num_requests == 2 < metadata.num_seq_slots + assert worker._kv_windows.shape == (6, 1, 8, 4) + assert worker._ctx_len.shape == (6,) + assert worker._batch_to_slot.shape == (num_seq_slots,) + assert worker._scratch_slot == num_seq_slots + assert list(worker._free_slots) == list(range(num_seq_slots)) + + live_slots = {worker._assign_slot(100 + i, reset=False) for i in range(num_seq_slots)} + assert live_slots == set(range(num_seq_slots)) + assert worker._scratch_slot not in live_slots + assert list(worker._free_slots) == [] + + def test_worker_rejects_mismatched_block_size(): worker = _make_worker() draft_model = _fake_draft_model() @@ -192,13 +228,12 @@ def write_context_windows(self, hidden, positions, windows): worker = _make_worker() draft_model = DraftModel() - metadata = types.SimpleNamespace( - max_num_requests=1, - request_ids=[100], - get_hidden_states=lambda _num_tokens: torch.zeros( - 3, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 - ), - ) + # Real metadata rather than a bare SimpleNamespace: _lazy_init reads + # slot-pool sizing fields off it, and a sparse stub silently omits any + # field added later. Its own get_hidden_states serves the per-chunk + # captures, sized by the total_target_tokens passed below. + metadata = _make_metadata(max_num_requests=1) + metadata.request_ids = [100] worker._lazy_init(draft_model, metadata) first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) @@ -208,9 +243,6 @@ def write_context_windows(self, hidden, positions, windows): slot = worker._req_to_slot[100] assert int(worker._ctx_len[slot]) == 3 - metadata.get_hidden_states = lambda _num_tokens: torch.zeros( - 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 - ) second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) worker._seed_context_windows( draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2