[https://nvbugs/6535874][fix] Size DFlash/DSpark worker slot buffers from num_seq_slots

tensorrt_llm/_torch/speculative/dflash.py

@@ -80,6 +80,10 @@ class DFlashSpecMetadata(SpecMetadata):
     captured_hidden_states: Optional[torch.Tensor] = None
 
     def __post_init__(self):
+        # Preserve the initial slot capacity across CUDA graph copies, whose
+        # max_num_requests is narrowed to the captured graph bucket.
+        self.num_seq_slots = self.num_seq_slots or self.max_num_requests
+
         self.batch_indices_cuda = torch.empty(
             [self.max_num_requests],
             dtype=torch.int,
@@ -253,7 +257,14 @@ 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 a stable upper bound. _free_slots assigns
+        # rows by request ID and only needs the pre-graph max_num_requests;
+        # num_seq_slots is the surviving proxy after max_num_requests is
+        # narrowed to a captured graph bucket. Under disagg-ADP it can be
+        # 2 * max_num_requests, so this deliberately overallocates the large
+        # context K/V buffers to keep one safe capacity across graph buckets.
+        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

tensorrt_llm/_torch/speculative/dspark.py

@@ -62,6 +62,10 @@ class DSparkSpecMetadata(SpecMetadata):
     captured_hidden_states: Optional[torch.Tensor] = None
 
     def __post_init__(self):
+        # Preserve the initial slot capacity across CUDA graph copies, whose
+        # max_num_requests is narrowed to the captured graph bucket.
+        self.num_seq_slots = self.num_seq_slots or self.max_num_requests
+
         self.batch_indices_cuda = torch.empty(
             [self.max_num_requests],
             dtype=torch.int,
@@ -266,7 +270,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
         num_stages = draft_model.num_stages
         self._win = int(draft_model._attn_params["window_size"])
         head_dim = int(draft_model._attn_params["head_dim"])

tensorrt_llm/_torch/speculative/utils.py

@@ -210,6 +210,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():
@@ -225,6 +226,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():

tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py

@@ -16,13 +16,18 @@
 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 DFlashSpecMetadata, DFlashWorker
+from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode
+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__), ".."))
 
@@ -32,6 +37,93 @@
     "The future of AI is",
 ]
 
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+
+
+def test_dflash_metadata_preserves_default_seq_slot_pool_in_graph_copy():
+    metadata = DFlashSpecMetadata(
+        max_draft_len=4,
+        max_total_draft_tokens=4,
+        spec_dec_mode=SpeculativeDecodingMode.DFLASH,
+        max_num_requests=5,
+    )
+
+    graph_metadata = metadata.create_cuda_graph_metadata(max_batch_size=2)
+
+    assert metadata.num_seq_slots == 5
+    assert graph_metadata.max_num_requests == 2
+    assert graph_metadata.num_seq_slots == 5
+
+
+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
+    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[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))
+
+    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,

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"
@@ -93,6 +94,16 @@ def test_metadata_prepare_batch_indices():
     assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2]
 
 
+def test_metadata_preserves_default_seq_slot_pool_in_graph_copy():
+    metadata = _make_metadata(max_num_requests=5)
+
+    graph_metadata = metadata.create_cuda_graph_metadata(max_batch_size=2)
+
+    assert metadata.num_seq_slots == 5
+    assert graph_metadata.max_num_requests == 2
+    assert graph_metadata.num_seq_slots == 5
+
+
 def _make_worker():
     cfg = types.SimpleNamespace(
         max_draft_len=5,
@@ -134,6 +145,42 @@ 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
+    num_slots = num_seq_slots + 1
+    assert worker._kv_windows.shape[0] == num_slots
+    assert worker._ctx_len.shape[0] == num_slots
+    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 +239,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 +254,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