[None][fix] Harden dynamic tree kernels and clean up one-model speculative decoding

Fix OOB read in CUDA tree-walk kernels by bounding the while loop and
guarding against missing parent tokens. Remove duplicate SpeculativeConfig
type alias that shadowed the discriminated union. Replace assert-False
error handling with proper RuntimeError, propagate runtime max_batch_size
from executor creator to one-model worker, and clean up dead code.

Signed-off-by: Qianyi Guan <qguan@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>

Author: sunnyqgg

cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu

@@ -112,7 +112,7 @@ __global__ void buildDynamicTreeKernel(int64_t const* parentList, int64_t const*
     {
         // Walk up to root, setting treeMask ancestor bits and counting depth
         int32_t curPosition = tid - 1;
-        while (true)
+        while (position < depth + 1)
         {
             position += 1;
             treeMask[tokenTreeIdx + curPosition] = 1;
@@ -131,6 +131,10 @@ __global__ void buildDynamicTreeKernel(int64_t const* parentList, int64_t const*
                     break;
                 }
             }
+            if (curPosition == draftTokenNum)
+            {
+                break;
+            }
         }
         positions[bid * draftTokenNum + tid] = position + seqLen;
     }
@@ -204,7 +208,7 @@ __global__ void buildDynamicTreeKernelPacked(int64_t const* parentList, int64_t
     else
     {
         int32_t curPosition = tid - 1;
-        while (true)
+        while (position < depth + 1)
         {
             position += 1;
 
@@ -231,6 +235,10 @@ __global__ void buildDynamicTreeKernelPacked(int64_t const* parentList, int64_t
                     break;
                 }
             }
+            if (curPosition == draftTokenNum)
+            {
+                break;
+            }
         }
         positions[bid * draftTokenNum + tid] = position + seqLen;
     }

examples/llm-api/quickstart_advanced.py

@@ -13,7 +13,6 @@
     "Hello, my name is",
     "The capital of France is",
     "The future of AI is",
-    # "A conversation between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: Hello! ASSISTANT:",
 ]
 
 

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -361,6 +361,12 @@ def create_py_executor(
         has_draft_model_engine = spec_config.spec_dec_mode.has_draft_model()
         has_spec_drafter = spec_config.spec_dec_mode.has_spec_drafter()
 
+        # Propagate runtime max_batch_size so one-model workers can pre-allocate
+        # buffers without hardcoding. Pydantic forbids extra fields, so use
+        # object.__setattr__ to bypass validation.
+        object.__setattr__(spec_config, '_runtime_max_batch_size',
+                           max_batch_size)
+
         # WAR for https://nvbugs/5807902
         # Disable separate draft KV cache in disaggregated mode
         # Enable separate pool for None DI + Non-KVBM and Aggregated + KVBM

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -3307,7 +3307,7 @@ def update_requests(
             state.sampler_event.synchronize()
 
         assert state.host is not None
-        new_tokens = state.host.new_tokens  # torch.Size([5, 2048, 1])
+        new_tokens = state.host.new_tokens
         finish_reasons = state.host.finish_reasons_list()
 
         new_tokens_list = new_tokens.tolist()

tensorrt_llm/_torch/speculative/drafting_loops.py

@@ -676,9 +676,7 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
 
         # Build dynamic tree structure using CUDA kernel (in-place, writes to pre-allocated buffers)
         tree_structure = None
-        if hasattr(
-                self,
-                'tree_ops_converter') and self.tree_ops_converter is not None:
+        if self.tree_ops_converter is not None:
             try:
                 # Write directly to spec_tree_manager buffers (target model reads from there).
                 # use_packed_mask=True: kernel outputs bit-packed mask directly,
@@ -701,7 +699,8 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
                 )
 
             except Exception as e:
-                assert False, f"Dynamic tree CUDA kernel failed: {e}"
+                raise RuntimeError(
+                    f"Dynamic tree CUDA kernel failed: {e}") from e
 
         # return_new_draft_tokens: [max_total_draft_tokens, batch_size]
         return_new_draft_tokens = torch.transpose(real_draft_tokens, 0, 1)
@@ -735,7 +734,8 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
             }
         }
 
-    def sample(self, logits: torch.Tensor, max_top_k: int) -> torch.Tensor:
+    def sample(self, logits: torch.Tensor,
+               max_top_k: int) -> tuple[torch.Tensor, torch.Tensor]:
         # TODO: inject the sampler here so we can support non-greedy
 
         # for draft_layer_idx == 0, logits is of shape [batch_size, vocab_size]

tensorrt_llm/_torch/speculative/dynamic_tree_ops.py

@@ -81,21 +81,6 @@ def __init__(
         self.max_batch_size = max_batch_size
         self.device = device
 
-        # Pre-allocate buffers for tree building
-        self._preallocate_buffers()
-
-    def _preallocate_buffers(self):
-        """Pre-allocate reusable buffers to minimize runtime allocation."""
-        # Preallocate parent_list buffer (max size)
-        # Size: [max_batch_size, K * (depth - 1) + 1]
-        # Note: Only first max_total_draft_tokens are used
-        self.parent_list_buffer = torch.full(
-            (self.max_batch_size, self.K * (self.depth - 1) + 1),
-            -1,
-            dtype=torch.int32,
-            device=self.device,
-        )
-
     def build_dynamic_tree(
         self,
         history_draft_tokens_parent_buffer: torch.Tensor,

tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py

@@ -105,7 +105,9 @@ def __init__(
         K = spec_config.dynamic_tree_max_topK
         max_draft_len = spec_config.max_draft_len
         max_total_draft_tokens = spec_config.tokens_per_gen_step - 1
-        max_batch_size = 256
+        # Read runtime max_batch_size set by py_executor_creator; fall back to 256
+        max_batch_size = getattr(spec_config, "_runtime_max_batch_size", 256)
+        self._max_batch_size = max_batch_size
 
         self.K = K
         self.max_total_draft_tokens = max_total_draft_tokens
@@ -207,6 +209,7 @@ def _ensure_spec_tree_manager(self, resource_manager):
 
     # ---- Overridden dispatch methods ----
 
+    @override
     def _forward_draft_loop(
         self,
         inputs,
@@ -236,6 +239,7 @@ def _forward_draft_loop(
             resource_manager,
         )
 
+    @override
     def forward(
         self,
         input_ids,
@@ -267,6 +271,7 @@ def forward(
             ]
         return output
 
+    @override
     def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata):
         """Override to handle dynamic tree verification."""
         batch_size = attn_metadata.num_seqs
@@ -277,6 +282,7 @@ def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata):
             logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens
         )
 
+    @override
     def prepare_1st_drafter_inputs(
         self,
         input_ids,
@@ -434,6 +440,10 @@ def _forward_dynamic_tree_draft_loop(
         self._ensure_spec_tree_manager(resource_manager)
         spec_tree_manager = self.spec_tree_manager
 
+        assert batch_size <= self._max_batch_size, (
+            f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}"
+        )
+
         # === Step 0: Initial forward ===
         # Inputs already in uniform-padded layout from prepare_1st_drafter_inputs:
         # max_draft_len + 1 tokens per gen request (contiguous accepted + zero padding).
@@ -703,7 +713,9 @@ def _sample_and_accept_dynamic_tree(
 
     # ---- Dynamic tree helper methods (matching two-model naming) ----
 
-    def sample(self, logits: torch.Tensor, max_top_k: int, draft_model=None) -> torch.Tensor:
+    def sample(
+        self, logits: torch.Tensor, max_top_k: int, draft_model=None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """TopK sampling with log softmax for dynamic tree."""
         last_p = self.logsoftmax(logits)
         topk_values, topk_indices = torch.topk(last_p, k=max_top_k, dim=-1)

tensorrt_llm/_torch/speculative/spec_tree_manager.py

@@ -317,41 +317,10 @@ def compute_spec_dec_mask_matrix(self, tree_idx=0):
             indices = path[path > -1]
             self.spec_dec_mask_matrix[actual_idx][i, indices] = 1
 
-    # # Compute the packed mask according to the mask matrix
-    # def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask):
-    #     # mask_matrix: shape: [bs, num_tokens, num_tokens]
-    #     # packed_mask: output buffer, shape: [bs, max_total_draft_tokens + 1, num_blocks]
-    #     # This function pads last dim to padded_len and packs each row into 32-bit integers
-    #     # Only the first num_tokens rows of packed_mask are filled
-
-    #     num_tokens = mask_matrix.shape[1]
-    #     padded_len = self.max_total_draft_tokens + 1
-    #     num_blocks = math.ceil(padded_len / 32)
-
-    #     # Pad last dim: [bs, num_tokens, num_tokens] -> [bs, num_tokens, padded_len]
-
-    #     pad_cols = padded_len - num_tokens
-    #     padded_matrix = torch.nn.functional.pad(mask_matrix, (0, pad_cols), value=0)
-
-    #     # Flatten to [bs * num_tokens, padded_len] for packing
-    #     int_tensor = padded_matrix.reshape(-1, padded_len)
-    #     packed_mask = packed_mask[:,:num_tokens,:].reshape(-1, num_blocks)
-
-    #     # Pack each 32-bit block
-    #     for block_idx in range(num_blocks):
-    #         start_idx = block_idx * 32
-    #         end_idx = min(start_idx + 32, padded_len)
-    #         block_bits = int_tensor[:, start_idx:end_idx]
-    #         weight = torch.pow(
-    #             2,
-    #             torch.arange(end_idx - start_idx,
-    #                          dtype=torch.int32,
-    #                          device=int_tensor.device))
-    #         block_value = torch.sum(block_bits * weight, dim=-1)
-    #         packed_mask[:, block_idx] = block_value
-
     def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask):
         bs, num_tokens, num_tokens_attend = mask_matrix.shape
+        assert mask_matrix.ndim == 3, f"Expected 3D mask_matrix, got {mask_matrix.ndim}D"
+        assert packed_mask.ndim == 3, f"Expected 3D packed_mask, got {packed_mask.ndim}D"
         num_blocks = packed_mask.shape[-1]
 
         # 1. Prepare bit weights (1, 1, 32)

tensorrt_llm/llmapi/llm_args.py

@@ -1821,8 +1821,8 @@ def supports_backend(self, backend: str) -> bool:
 SpeculativeConfig: TypeAlias = Annotated[
     Union[
         DraftTargetDecodingConfig,
+        Eagle3DecodingConfig,  # Must be before EagleDecodingConfig since it's a subclass
         EagleDecodingConfig,
-        Eagle3DecodingConfig,
         LookaheadDecodingConfig,
         MedusaDecodingConfig,
         MTPDecodingConfig,
@@ -1836,19 +1836,6 @@ def supports_backend(self, backend: str) -> bool:
     Field(discriminator="decoding_type"),
 ]
 
-SpeculativeConfig: TypeAlias = Optional[Union[
-    DraftTargetDecodingConfig,
-    Eagle3DecodingConfig,  # Must be before EagleDecodingConfig since it's a subclass
-    EagleDecodingConfig,
-    LookaheadDecodingConfig,
-    MedusaDecodingConfig,
-    MTPDecodingConfig,
-    NGramDecodingConfig,
-    UserProvidedDecodingConfig,
-    SaveHiddenStatesDecodingConfig,
-    AutoDecodingConfig,
-]]
-
 SparseAttentionConfig: TypeAlias = Annotated[
     Union[
         RocketSparseAttentionConfig,

tests/unittest/_torch/speculative/test_draft_token_prepare_for_generation.py

@@ -825,7 +825,7 @@ def run_test(
         ref_hidden_states_read_indices,
     )
 
-    ##### CASE 2 dynamic tree, batch size = 1, cur_draft_idx = 1 #############
+    ##### CASE 3 dynamic tree, batch size = 1, cur_draft_idx = 2 #############
     max_batch_size = 1
     max_draft_len = 3
     max_total_draft_tokens = 15