[TRTLLM-11556][feat] Expand dynamic speculation to all spec decode algorithms

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1096,15 +1096,16 @@ def update_spec_dec_param(
                 # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length.
                 # So we create cache for position offsets and packed mask for each draft length to avoid reallocation.
                 assert max_draft_len == max_total_draft_tokens, "max_draft_len should be equal to max_total_draft_tokens for linear tree"
-                runtime_draft_len = (spec_metadata.runtime_draft_len
-                                     if spec_metadata is not None else
-                                     max_draft_len)
+                # For algos other than PARD, this equals runtime_draft_len (K); for PARD it's 2K-1.
+                runtime_draft_token_buffer_width = (
+                    spec_metadata.runtime_tokens_per_gen_step -
+                    1 if spec_metadata is not None else max_draft_len)
                 self.generate_spec_decoding_generation_length(
-                    runtime_draft_len=runtime_draft_len)
+                    runtime_draft_len=runtime_draft_token_buffer_width)
                 self.spec_decoding_position_offsets = generate_spec_decoding_position_offsets(
-                    self.max_num_requests, runtime_draft_len)
+                    self.max_num_requests, runtime_draft_token_buffer_width)
                 self.spec_decoding_packed_mask = generate_spec_decoding_packed_mask(
-                    self.max_num_requests, runtime_draft_len)
+                    self.max_num_requests, runtime_draft_token_buffer_width)
 
             self.update_position_offsets_for_cpp(cpp_query_len)
 

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -128,8 +128,10 @@ def __init__(self, config: CUDAGraphRunnerConfig):
 
     def _create_shared_static_tensors(self):
         """Allocates static tensors sized for the largest possible batch."""
-        max_draft_len = self.config.original_max_total_draft_tokens if self.config.spec_config is not None else 0
-        token_per_request = max_draft_len + 1
+        runtime_draft_token_buffer_width = (
+            self.config.original_max_total_draft_tokens
+            if self.config.spec_config is not None else 0)
+        token_per_request = runtime_draft_token_buffer_width + 1
         max_total_tokens = (self.max_supported_batch_size *
                             self.max_beam_width * token_per_request)
         max_total_tokens = min(max_total_tokens, self.config.max_num_tokens)
@@ -486,6 +488,11 @@ def _get_padded_batch(self, batch: ScheduledRequests,
         if padding_size + batch.batch_size > self.config.batch_size:
             return 0
 
+        runtime_tokens_per_gen_step = (
+            self.spec_config.get_runtime_tokens_per_gen_step(runtime_draft_len)
+            if self.spec_config is not None else 1 + runtime_draft_len)
+        runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1
+
         # No padding if it would create too many concurrent requests.
         # This is not strictly required, but we should probably
         # respect the requirement just in case that changes in the future.
@@ -503,7 +510,7 @@ def _get_padded_batch(self, batch: ScheduledRequests,
             dummy_request = kv_cache_manager.add_dummy_requests(
                 [dummy_request_id],
                 is_gen=True,
-                max_num_draft_tokens=runtime_draft_len,
+                max_num_draft_tokens=runtime_draft_token_buffer_width,
                 use_mrope=self.config.use_mrope,
                 max_beam_width=self.config.max_beam_width,
                 draft_kv_cache_manager=draft_kv_cache_manager)

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -2521,26 +2521,53 @@ def _handle_dynamic_draft_len(self,
             from tensorrt_llm._torch.speculative.utils import \
                 get_draft_len_for_batch_size
 
+            spec_dec_mode = self.model_engine.spec_config.spec_dec_mode
+
             # 1. Resolve runtime draft length from schedule
             runtime_draft_len = get_draft_len_for_batch_size(
                 self.model_engine.spec_config.draft_len_schedule,
                 scheduled_batch.batch_size, self.model_engine.max_draft_len)
-
             # 2. Pad or truncate draft tokens to the resolved length
-            PADDING_TOKEN = 0
+            DRAFT_BUFFER_PAD = 0  # Buffer sentinel, not PARD mask_token_id.
             for request in scheduled_batch.generation_requests:
-                current_draft_len = len(request.py_draft_tokens)
-                if current_draft_len < runtime_draft_len:
-                    padding_needed = runtime_draft_len - current_draft_len
-                    request.py_draft_tokens.extend([PADDING_TOKEN] *
-                                                   padding_needed)
-                elif current_draft_len > runtime_draft_len:
-                    request.py_draft_tokens = request.py_draft_tokens[:
-                                                                      runtime_draft_len]
+                current_num_draft_tokens = len(request.py_draft_tokens)
+                if spec_dec_mode.is_pard():
+                    # special case: PARD carries 2K-1 draft tokens per request
+                    runtime_draft_token_buffer_width = (
+                        self.model_engine.spec_config.
+                        get_runtime_tokens_per_gen_step(runtime_draft_len) - 1)
+                    current_runtime_draft_len = (
+                        current_num_draft_tokens +
+                        1) // 2 if current_num_draft_tokens > 0 else 0
+                    real_draft_tokens = request.py_draft_tokens[:min(
+                        current_runtime_draft_len, runtime_draft_len)]
+                    real_draft_tokens.extend(
+                        [DRAFT_BUFFER_PAD] *
+                        (runtime_draft_len - len(real_draft_tokens)))
+                    request.py_draft_tokens = real_draft_tokens + [
+                        DRAFT_BUFFER_PAD
+                    ] * (runtime_draft_token_buffer_width -
+                         len(real_draft_tokens))
+                else:
+                    if current_num_draft_tokens < runtime_draft_len:
+                        padding_needed = (runtime_draft_len -
+                                          current_num_draft_tokens)
+                        request.py_draft_tokens.extend([DRAFT_BUFFER_PAD] *
+                                                       padding_needed)
+                    elif current_num_draft_tokens > runtime_draft_len:
+                        request.py_draft_tokens = request.py_draft_tokens[:
+                                                                          runtime_draft_len]
 
             self.model_engine.runtime_draft_len = runtime_draft_len
         else:
-            self.model_engine.runtime_draft_len = self.model_engine.max_total_draft_tokens
+            # Linear-tree modes (incl. PARD) use logical K; tree decoding
+            # (e.g. EAGLE3 dynamic tree) uses total tree tokens. Same
+            # selection as _prepare_tp_inputs and _get_graphs_to_capture.
+            spec_config = self.model_engine.spec_config
+            self.model_engine.runtime_draft_len = (
+                self.model_engine.max_draft_len
+                if spec_config is not None and spec_config.is_linear_tree else
+                self.model_engine.max_total_draft_tokens)
 
     def _can_queue(self, scheduled_batch):
 

tensorrt_llm/_torch/speculative/dflash.py

@@ -383,7 +383,18 @@ def forward(
         num_gens = batch_size - num_contexts
 
         raw_logits = logits
-        K = self.max_draft_len
+        K = spec_metadata.runtime_draft_len
+
+        if K == 0:
+            return self.skip_drafting(
+                input_ids,
+                position_ids,
+                hidden_states,
+                logits,
+                attn_metadata,
+                spec_metadata,
+                draft_model,
+            )
 
         # Lazy init buffers and attach worker reference for prepare()
         self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata)
@@ -485,7 +496,7 @@ def forward(
                 )
 
                 vocab_size = gen_logits.shape[-1]
-                gen_logits = gen_logits.reshape(num_gens, self.max_draft_len, vocab_size)
+                gen_logits = gen_logits.reshape(num_gens, K, vocab_size)
 
                 d2t = getattr(draft_model.model, "d2t", None)
                 gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long()
@@ -583,7 +594,7 @@ def prepare_1st_drafter_inputs(
             gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :]
 
             total_tokens_per_req = self._draft_tokens_per_req  # K+1
-            K = self.max_draft_len
+            K = spec_metadata.runtime_draft_len
 
             # Get captured multi-layer hidden states from spec_metadata
             captured_hs = spec_metadata.get_hidden_states(total_target_tokens)

tensorrt_llm/_torch/speculative/draft_target.py

@@ -70,7 +70,7 @@ def prepare(self):
             num_seqs, dtype=torch.int, device="cpu", pin_memory=prefer_pinned()
         )
         self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True)
-        self.num_tokens -= self.num_generations * self.max_draft_len
+        self.num_tokens -= self.num_generations * self.runtime_draft_len
         self.is_spec_dec_tree = False
         self.is_spec_dec_dynamic_tree = False
 
@@ -131,10 +131,11 @@ def _update_kv_after_first_draft_step(
         num_accepted_tokens: torch.Tensor,
         num_contexts: int,
         batch_size: int,
+        runtime_draft_len: int,
     ):
         if hasattr(attn_metadata, "kv_lens_cuda"):
             attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (
-                self.max_draft_len - num_accepted_tokens[num_contexts:batch_size]
+                runtime_draft_len - num_accepted_tokens[num_contexts:batch_size]
             )
             attn_metadata.kv_lens_cuda[:num_contexts] += 1
 
@@ -175,6 +176,18 @@ def forward(
         batch_size = attn_metadata.num_seqs
         num_contexts = attn_metadata.num_contexts
         num_gens = batch_size - num_contexts
+        runtime_draft_len = spec_metadata.runtime_draft_len
+
+        if runtime_draft_len == 0:
+            return self.skip_drafting(
+                input_ids,
+                position_ids,
+                hidden_states,
+                logits,
+                attn_metadata,
+                spec_metadata,
+                draft_model,
+            )
 
         raw_logits = logits
 
@@ -204,10 +217,10 @@ def forward(
         draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager)
 
         with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager):
-            for i in range(self.max_draft_len):
+            for i in range(runtime_draft_len):
                 if i == 0:
                     start_ids_gen = (
-                        spec_metadata.batch_indices_cuda[:num_gens] * (self.max_draft_len + 1)
+                        spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1)
                     ).long()
                     gather_ids_gen = (
                         start_ids_gen
@@ -260,7 +273,11 @@ def forward(
                         attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1)
                         attn_metadata.num_contexts = 0
                     self._update_kv_after_first_draft_step(
-                        attn_metadata, num_accepted_tokens, num_contexts, batch_size
+                        attn_metadata,
+                        num_accepted_tokens,
+                        num_contexts,
+                        batch_size,
+                        runtime_draft_len,
                     )
                 else:
                     self._update_kv_for_chained_draft_step(attn_metadata, batch_size)
@@ -306,13 +323,14 @@ def sample_and_accept_draft_tokens(
         batch_size = attn_metadata.num_seqs
         num_contexts = attn_metadata.num_contexts
         num_gens = batch_size - num_contexts
+        runtime_draft_len = spec_metadata.runtime_draft_len
 
         if spec_metadata.draft_tokens is None:
             draft_tokens = torch.zeros(
-                (num_gens, self.max_draft_len), dtype=torch.int, device=logits.device
+                (num_gens, runtime_draft_len), dtype=torch.int, device=logits.device
             )
         else:
-            draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, self.max_draft_len)
+            draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len)
 
         return self._sample_and_accept_draft_tokens_base(
             logits, draft_tokens, num_contexts, batch_size, spec_metadata
@@ -337,6 +355,7 @@ def prepare_1st_drafter_inputs(
         num_contexts = attn_metadata.num_contexts
         batch_size = attn_metadata.num_seqs
         num_gens = batch_size - num_contexts
+        runtime_draft_len = spec_metadata.runtime_draft_len
 
         if num_contexts > 0:
             input_ids_ctx = self._prepare_context_input_ids(
@@ -350,7 +369,9 @@ def prepare_1st_drafter_inputs(
             input_ids_ctx = torch.empty(0, dtype=torch.int32, device="cuda")
 
         if num_gens > 0:
-            input_ids_gen = accepted_tokens[num_contexts:, :].flatten().to(torch.int32)
+            input_ids_gen = (
+                accepted_tokens[num_contexts:, : runtime_draft_len + 1].flatten().to(torch.int32)
+            )
         else:
             input_ids_gen = torch.empty(0, dtype=torch.int32, device="cuda")
 

tensorrt_llm/_torch/speculative/eagle3.py

@@ -519,7 +519,7 @@ def prepare(self):
         if sa_manager is not None:
             gen_request_ids = self.request_ids[num_seqs - self.num_generations:]
             if gen_request_ids:
-                sa_manager.prepare(gen_request_ids, self.max_draft_len)
+                sa_manager.prepare(gen_request_ids, self.runtime_draft_len)
 
     def maybe_capture_hidden_states(
             self,
@@ -700,7 +700,7 @@ def forward(self,
                 num_accepted_tokens=num_accepted_tokens,
                 num_gens=num_gens,
                 num_contexts=num_contexts,
-                max_draft_len=self.max_draft_len,
+                max_draft_len=runtime_draft_len,
             )
 
         # Save the old attn_metadata and spec_metadata

tensorrt_llm/_torch/speculative/interface.py

@@ -315,8 +315,9 @@ def support_capturable_guided_decoder(self):
         ) or self.is_external_drafter() or self.is_sa()
 
     def support_dynamic_draft_len(self):
-        # TODO: expand to all one-model algorithms
-        return self.is_eagle3_one_model() or self.is_mtp_eagle_one_model()
+        return self.is_mtp_one_model() or self.is_eagle3_one_model(
+        ) or self.is_mtp_eagle_one_model() or self.is_pard() or self.is_dflash(
+        ) or self.is_draft_target_one_model() or self.is_sa()
 
     def has_draft_model(self):
         return self.is_eagle3() or self.is_draft_target() or self.is_mtp_eagle()
@@ -455,6 +456,9 @@ class SpecMetadata:
     # draft_len_schedule.  Otherwise it equals max_draft_len (the static max).
     # Always set by model_engine.forward() before any downstream code reads it.
     runtime_draft_len: int = 0
+    # Total runtime tokens per generation request for the current iteration,
+    # Normally, it equals 1 + runtime_draft_len. But for PARD, it equals 2 * runtime_draft_len.
+    runtime_tokens_per_gen_step: int = 1
 
     # Auto-detected per step from populated sampling params:
     # True if every request is greedy (no temp/top_k/top_p) and we can take
@@ -1113,9 +1117,8 @@ def _sample_and_accept_draft_tokens_base(
             num_accepted_tokens: [batch_size] - Number of accepted tokens per request
         """
         # Derive draft length from the actual draft_tokens shape rather than
-        # spec_metadata.runtime_draft_len, because they can differ: PARD sets
-        # runtime_draft_len = 2K-1 for input sizing but only passes K draft
-        # tokens for acceptance;
+        # spec_metadata.runtime_draft_len, because callers may slice a wider
+        # runtime token layout down to the K draft tokens used for acceptance.
         runtime_draft_len = draft_tokens.shape[-1]
         num_gens = batch_size - num_contexts