From 69f73616a41771d38054757a6344d9afd7e8bda7 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:33:25 -0700 Subject: [PATCH 1/3] [nvbugs/6535874][fix] Size DFlash/DSpark ctx pools from the full seq-slot range The DFlash and DSpark workers allocate their per-request context pools once, lazily, on the first drafting forward and then reuse them for every later batch shape. Both sized the pool from spec_metadata.max_num_requests, which create_cuda_graph_metadata overwrites with the captured graph bucket size. Whichever bucket drafts first therefore pinned the pool to that bucket, leaving the scratch slot (_dummy_slot / _scratch_slot) inside the range real requests draw from, so CUDA-graph padding writes could overwrite a live request's context. Read num_seq_slots instead: it spans the full seq-slot pool and is left untouched by the graph copy, which is the same reason prepare_rejection_sampling_buffers derives its slot capacity from it. The DFlash and DSpark metadata constructions did not forward num_seq_slots, so the field defaulted to 0 and any bound derived from it silently collapsed back to max_num_requests; plumb it through as the MTP-Eagle path already does. Both halves are load-bearing. This does not change the pool size on either path today: num_seq_slots is non-None only under should_enable_dsv4_overlap_headroom, which requires is_mtp_eagle_one_model(), so it resolves to exactly max_num_requests here. test_seed_context_windows_preserves_state_across_prefill_chunks built its metadata as a bare SimpleNamespace carrying only three attributes, so it could not survive _lazy_init reading any further sizing field. Switch it to the real DSparkSpecMetadata via the file's own _make_metadata helper, as the other ten _lazy_init call sites in that file already do. The two test_dflash_qwen3_5_4b waivers are stale: the SSM-pool ValueError they were filed against was fixed by 7f7dccf991, which landed after the waivers, so they were never re-evaluated. Remove them. Leaves batch_indices_cuda (metadata-owned, reallocated by the graph copy's __post_init__) and the eagle3 padding computation (must track the current batch shape) unchanged. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 7 ++++++- tensorrt_llm/_torch/speculative/dspark.py | 5 ++++- tensorrt_llm/_torch/speculative/utils.py | 2 ++ tests/integration/test_lists/waives.txt | 2 -- .../hw_agnostic/test_dspark_worker.py | 16 ++++++---------- 5 files changed, 18 insertions(+), 14 deletions(-) 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_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 4625d1b0115a..81b1adb1f3b9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -192,13 +192,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 +207,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 From 258566c9dfaab3cae6b1f4c77cdc02ddf2356554 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:14:41 +0800 Subject: [PATCH 2/3] [nvbugs/6535874][test] Cover enlarged speculative slot pools (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../speculative/hw_agnostic/test_dflash.py | 73 +++++++++++++++++++ .../hw_agnostic/test_dspark_worker.py | 36 +++++++++ 2 files changed, 109 insertions(+) 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 81b1adb1f3b9..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() From ec208612bd3a7923daa855c075666642824d8f00 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:52:49 +0800 Subject: [PATCH 3/3] [Agent fix] Address speculative slot review feedback (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 2 +- tests/integration/test_lists/waives.txt | 2 ++ .../unittest/_torch/speculative/hw_agnostic/test_dflash.py | 7 +++++-- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 38e8b0b39515..32a866ef5e11 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -220,7 +220,7 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): # 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 + max_batch = spec_metadata.num_seq_slots # Prefer runtime max_seq_len over max_position_embeddings: YaRN # models advertise 100k+ positions, which would OOM the ctx buffer diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 0c88b7a01f2c..b2971766eacd 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -421,6 +421,8 @@ 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 716bc341a745..d2e4c808cf14 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py @@ -36,6 +36,8 @@ "The future of AI is", ] +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + def test_dflash_graph_bucket_uses_full_seq_slot_pool(): """A small graph bucket must not shrink the persistent context pool.""" @@ -85,9 +87,10 @@ def precompute_context_kv(self, hidden_states, position_ids): 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) + num_slots = num_seq_slots + 1 + assert worker._ctx_k_buf.shape[0] == num_slots assert worker._ctx_v_buf.shape == worker._ctx_k_buf.shape - assert worker._ctx_len.shape == (6,) + assert worker._ctx_len.shape[0] == num_slots 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))