[TRTLLM-10886][feat] Support PARD(Parallel Draft Model) in one-model spec dec (NVIDIA#11438)

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Author: ziyixiong-nv

.pre-commit-config.yaml

@@ -1441,7 +1441,7 @@ repos:
         additional_dependencies:
         - tomli
         # add ignore words list
-        args: ["-L", "Mor,ans,thirdparty,subtiles", "--skip", "ATTRIBUTIONS-*.md,*.svg", "--skip", "security_scanning/*"]
+        args: ["-L", "Mor,ans,thirdparty,subtiles,PARD,pard", "--skip", "ATTRIBUTIONS-*.md,*.svg", "--skip", "security_scanning/*"]
         exclude: 'scripts/attribution/data/cas/.*$'
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.9.4

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -518,7 +518,7 @@ def __init__(
 
         # check for max total draft tokens
         if self.spec_config is not None:
-            self.max_total_draft_tokens = self.spec_config.max_total_draft_tokens
+            self.max_total_draft_tokens = self.spec_config.tokens_per_gen_step - 1
         else:
             self.max_total_draft_tokens = 0
 
@@ -1063,7 +1063,7 @@ def create_autodeploy_executor(ad_config: LlmArgs, tokenizer: Optional[Tokenizer
     max_total_draft_tokens = (
         0
         if ad_config.speculative_config is None
-        else ad_config.speculative_config.max_total_draft_tokens
+        else ad_config.speculative_config.tokens_per_gen_step - 1
     )
 
     # initialize model engine

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -718,7 +718,7 @@ def __init__(
         reduce_output: bool = True,
     ):
         config = model_config.pretrained_config
-        predicted_tokens_per_seq = model_config.spec_config.max_total_draft_tokens + 1 if model_config.spec_config is not None else 1
+        predicted_tokens_per_seq = model_config.spec_config.tokens_per_gen_step if model_config.spec_config is not None else 1
         super().__init__(hidden_size=config.hidden_size,
                          num_attention_heads=config.num_attention_heads,
                          num_key_value_heads=config.num_key_value_heads,
@@ -766,7 +766,7 @@ def __init__(
         reduce_output: bool = True,
     ):
         config = model_config.pretrained_config
-        predicted_tokens_per_seq = model_config.spec_config.max_total_draft_tokens + 1 if model_config.spec_config is not None else 1
+        predicted_tokens_per_seq = model_config.spec_config.tokens_per_gen_step if model_config.spec_config is not None else 1
 
         super().__init__(hidden_size=config.hidden_size,
                          num_attention_heads=config.num_attention_heads,

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 from typing import Dict, Generic, List, Optional, Tuple
 
 import torch
@@ -24,7 +25,7 @@
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, TModel,
-                             register_auto_model)
+                             get_model_architecture, register_auto_model)
 
 
 def _ensure_draft_vocab_size(config: PretrainedConfig) -> None:
@@ -108,9 +109,9 @@ def __init__(
         config = model_config.pretrained_config
         self._next_layer_regular = next_layer_regular
 
-        predicted_tokens_per_seq = (
-            model_config.spec_config.max_total_draft_tokens +
-            1 if model_config.spec_config is not None else 1)
+        predicted_tokens_per_seq = (model_config.spec_config.tokens_per_gen_step
+                                    if model_config.spec_config is not None else
+                                    1)
 
         super().__init__(
             hidden_size=config.hidden_size,
@@ -702,6 +703,70 @@ def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return hidden_states
 
 
+class PARDForCausalLM(nn.Module):
+    """Draft model wrapper for PARD (Parallel Draft) speculative decoding.
+
+    See PARDWorker for the full algorithm description.
+    """
+
+    def __init__(self, draft_config):
+        super().__init__()
+        DraftModelClass, _ = get_model_architecture(
+            draft_config.pretrained_config)
+
+        # Remove spec_config to prevent recursive spec-dec initialization
+        draft_config_no_spec = replace(draft_config, spec_config=None)
+
+        # Weights will be loaded later by ModelLoader.load_draft_weights()
+        self.draft_model_full = DraftModelClass(draft_config_no_spec)
+        self.model = self.draft_model_full.model
+        self.lm_head = self.draft_model_full.lm_head
+
+        # Required by weight mappers
+        self.model_config = draft_config_no_spec
+        self.config = draft_config_no_spec.pretrained_config
+
+        # Fall back: pard_token -> mask_token_id -> vocab_size
+        pretrained_config = draft_config.pretrained_config
+        self.mask_token_id = getattr(
+            pretrained_config, 'pard_token',
+            getattr(pretrained_config, 'mask_token_id',
+                    pretrained_config.vocab_size))
+        logger.info(
+            f"PARD draft model initialized with mask_token_id: {self.mask_token_id}"
+        )
+
+        self.logits_processor = None  # Set by caller after construction
+
+    def load_weights(self, weights: Dict, weight_mapper=None, **kwargs):
+        """Load weights into the PARD draft model."""
+        self.draft_model_full.load_weights(weights=weights,
+                                           weight_mapper=weight_mapper,
+                                           **kwargs)
+
+    def forward(
+        self,
+        attn_metadata,
+        input_ids: torch.LongTensor = None,
+        position_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        return_context_logits: bool = False,
+        spec_metadata=None,
+        hidden_states: torch.Tensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden_states_out = self.model(
+            input_ids=input_ids,
+            attn_metadata=attn_metadata,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            spec_metadata=spec_metadata,
+            **kwargs,
+        )
+
+        return hidden_states_out, hidden_states_out
+
+
 class MTPForCausalLM(nn.Module):
 
     def __init__(
@@ -917,6 +982,8 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                               lm_head, model)
     elif spec_dec_mode.is_mtp_eagle():
         return MTPDraftModelForCausalLM(model_config)
+    elif spec_dec_mode.is_pard():
+        return PARDForCausalLM(draft_config)
     else:
         raise NotImplementedError(
             f"get_draft_model does not support speculative decoding mode {spec_dec_mode}."
@@ -967,11 +1034,27 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                 self.draft_config.quant_config.kv_cache_quant_algo = \
                 model_config.quant_config.kv_cache_quant_algo
 
+            elif spec_config.spec_dec_mode.is_pard():
+                self.draft_config = ModelConfig.from_pretrained(
+                    model_config.spec_config.speculative_model,
+                    trust_remote_code=True,
+                    attn_backend=model_config.attn_backend,
+                    moe_backend=model_config.moe_backend,
+                    mapping=model_config.mapping,
+                    spec_config=None,  # Avoid recursive spec-dec
+                    max_num_tokens=model_config.max_num_tokens,
+                    moe_max_num_tokens=model_config.moe_max_num_tokens)
+                self.draft_config.quant_config.kv_cache_quant_algo = \
+                    model_config.quant_config.kv_cache_quant_algo
+
             self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(
                 spec_config)
 
             self.draft_model = get_draft_model(model_config, self.draft_config,
                                                self.lm_head, self.model)
+            if spec_config.spec_dec_mode.is_pard(
+            ) and self.draft_model is not None:
+                self.draft_model.logits_processor = self.logits_processor
             self.spec_worker = get_spec_worker(
                 model_config.spec_config,
                 model_config,
@@ -980,7 +1063,10 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
             self.epilogue.append(self.draft_model)
             self.epilogue.append(self.spec_worker)
 
-            if self.draft_config is not None and model_config.spec_config.eagle3_model_arch == "llama3":
+            # EAGLE3-specific logic: merge extra_attrs from draft model for Llama3
+            if (self.draft_config is not None and model_config.spec_config.
+                    spec_dec_mode.is_eagle3_one_model()
+                    and model_config.spec_config.eagle3_model_arch == "llama3"):
                 for key, value in self.draft_config.extra_attrs.items():
                     assert key in ('attn_layers', 'mla_layers')
                     assert key in model_config.extra_attrs
@@ -1067,7 +1153,9 @@ def load_draft_weights(self,
                            weight_mapper: Optional[BaseWeightMapper] = None):
         self.draft_model.load_weights(weights=weights,
                                       weight_mapper=weight_mapper)
-        self.draft_model.load_weights_from_target_model(self)
+        # PARD has independent weights; other methods share with target model
+        if not self.model_config.spec_config.spec_dec_mode.is_pard():
+            self.draft_model.load_weights_from_target_model(self)
 
     def set_guided_decoder(self,
                            guided_decoder: CapturableGuidedDecoder) -> bool:

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -277,10 +277,10 @@ def _get_token_num_for_estimation(self) -> int:
         num_extra_tokens_per_seq = 1  # account for generated tokens
         spec_cfg = self._speculative_config
         if not self._llm_args.disable_overlap_scheduler and spec_cfg is not None:
-            num_extra_tokens_per_seq += spec_cfg.max_total_draft_tokens
+            num_extra_tokens_per_seq += spec_cfg.tokens_per_gen_step - 1
 
         if spec_cfg is not None:
-            num_extra_tokens_per_seq += spec_cfg.max_total_draft_tokens
+            num_extra_tokens_per_seq += spec_cfg.tokens_per_gen_step - 1
             num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg)
 
         if self._dummy_reqs is None:
@@ -570,15 +570,17 @@ def _create_one_model_draft_kv_cache_manager(
         # Draft model layers in one-model mode start at target_num_layers.
         target_pretrained_config = self._model_engine.model.model_config.pretrained_config
         target_num_layers = target_pretrained_config.num_hidden_layers
-        # Use get_num_spec_layers to get the correct number of draft layers
-        # for the speculative decoding mode (e.g., num_eagle_layers for Eagle3)
-        num_draft_layers = get_num_spec_layers(self._speculative_config)
 
-        # Create layer_mask: False for target layers, True for draft layers.
-        # This ensures the draft KV cache manager uses the correct layer indices
-        # (e.g., layers 32, 33, ... instead of 0, 1, ...).
-        spec_dec_layer_mask = [False
-                               ] * target_num_layers + [True] * num_draft_layers
+        # PARD: draft is a separate model, layers start from 0.
+        # Other methods (EAGLE3, MTP): draft layers are appended after target layers.
+        if self._speculative_config.spec_dec_mode.is_pard():
+            num_draft_layers = self._draft_config.pretrained_config.num_hidden_layers
+            spec_dec_layer_mask = [True] * num_draft_layers
+        else:
+            num_draft_layers = get_num_spec_layers(self._speculative_config)
+            spec_dec_layer_mask = [False] * target_num_layers + [
+                True
+            ] * num_draft_layers
 
         # Get the effective draft config (explicit draft_config if available,
         # otherwise fall back to target model config for MTP).
@@ -1091,8 +1093,8 @@ def create_py_executor_instance(
         max_beam_width=max_beam_width,
         max_draft_len=spec_config.max_draft_len
         if spec_config is not None else 0,
-        max_total_draft_tokens=spec_config.max_total_draft_tokens
-        if spec_config is not None else 0,
+        max_total_draft_tokens=(spec_config.tokens_per_gen_step -
+                                1) if spec_config is not None else 0,
         kv_cache_transceiver=kv_cache_transceiver,
         guided_decoder=guided_decoder,
         start_worker=start_worker,
@@ -1120,7 +1122,7 @@ def create_torch_sampler_args(
     max_draft_len = (0 if speculative_config is None else
                      speculative_config.max_draft_len)
     max_total_draft_tokens = (0 if speculative_config is None else
-                              speculative_config.max_total_draft_tokens)
+                              speculative_config.tokens_per_gen_step - 1)
 
     return TorchSampler.Args(
         max_seq_len=max_seq_len,

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -173,7 +173,11 @@ def __init__(
             ExpertStatistic.create(self.dist.rank)
         self.llm_args = llm_args
         self.original_max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
-        self.original_max_total_draft_tokens = spec_config.max_total_draft_tokens if spec_config is not None else 0
+        self.original_max_total_draft_tokens = (
+            spec_config.tokens_per_gen_step -
+            1) if spec_config is not None else 0
+        # Saved before zeroing for draft models; used by update_spec_dec_param.
+        self._spec_dec_max_total_draft_tokens = spec_config.max_total_draft_tokens if spec_config is not None else 0
 
         # The draft model won't have any draft tokens attached to
         # generation requests when we invoke it autoregressively
@@ -342,7 +346,7 @@ def __init__(
             self.without_logits = self.spec_config.spec_dec_mode.without_logits(
             ) or self.model_is_wrapped
             self.max_draft_len = spec_config.max_draft_len
-            self.max_total_draft_tokens = spec_config.max_total_draft_tokens
+            self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1
         else:
             self.without_logits = False
             self.max_draft_len = 0
@@ -389,8 +393,9 @@ def __init__(
         # Pre-allocated buffers for draft model to avoid implicit synchronization
         # These are used to build index tensors without creating tensors from Python lists
         max_first_draft_tokens = self.batch_size * (
-            self.original_max_draft_len + 1) if spec_config else self.batch_size
-        tokens_per_draft = self.original_max_draft_len + 1
+            self.original_max_total_draft_tokens +
+            1) if spec_config else self.batch_size
+        tokens_per_draft = self.original_max_total_draft_tokens + 1
         self.idx_accepted_tokens_cache = None
         self.draft_token_positions_cache = None
         if spec_config:
@@ -1892,7 +1897,7 @@ def _apply_incremental_update_target(
         # Pre-compute constants
         extend_requests = scheduled_requests.generation_requests
         num_extend_requests = len(extend_requests)
-        num_tokens_per_extend_request = self.original_max_draft_len + 1
+        num_tokens_per_extend_request = self.original_max_total_draft_tokens + 1
         spec_config = self.spec_config
 
         prompt_lengths = torch.empty(num_extend_requests,
@@ -3480,7 +3485,7 @@ def forward(self,
                 is_spec_dec_tree=spec_metadata.is_spec_dec_tree,
                 is_spec_dec_dynamic_tree=spec_metadata.is_spec_dec_dynamic_tree,
                 max_draft_len=self.original_max_draft_len,
-                max_total_draft_tokens=self.original_max_total_draft_tokens,
+                max_total_draft_tokens=self._spec_dec_max_total_draft_tokens,
                 model_is_wrapped=self.model_is_wrapped,
                 spec_metadata=spec_metadata,
                 spec_tree_manager=spec_tree_manager,

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -415,12 +415,12 @@ def drafting_loop_wrapper(model):
                     if use_tree_drafter:
                         return TreeDraftingLoopWrapper(
                             spec_config.max_draft_len,
-                            spec_config.max_total_draft_tokens, max_batch_size,
+                            spec_config.tokens_per_gen_step - 1, max_batch_size,
                             model)
                     else:
                         return LinearDraftingLoopWrapper(
                             spec_config.max_draft_len,
-                            spec_config.max_total_draft_tokens, model)
+                            spec_config.tokens_per_gen_step - 1, model)
             else:
                 drafting_loop_wrapper = None
 
@@ -460,11 +460,11 @@ def drafting_loop_wrapper(model):
     model_engine_max_seq_len = model_engine.max_seq_len
     net_max_seq_len = model_engine_max_seq_len
     if not llm_args.disable_overlap_scheduler and spec_config is not None:
-        model_engine_max_seq_len += spec_config.max_total_draft_tokens
+        model_engine_max_seq_len += spec_config.tokens_per_gen_step - 1
 
     if spec_config is not None:
         model_engine_max_seq_len += get_num_extra_kv_tokens(spec_config)
-        model_engine_max_seq_len += spec_config.max_total_draft_tokens
+        model_engine_max_seq_len += spec_config.tokens_per_gen_step - 1
 
     if has_draft_model_engine and not llm_args.disable_overlap_scheduler:
         logger.warning(
@@ -546,7 +546,7 @@ def drafting_loop_wrapper(model):
                 }
                 if spec_config is not None:
                     kwargs[
-                        "max_num_draft_tokens"] = spec_config.max_total_draft_tokens
+                        "max_num_draft_tokens"] = spec_config.tokens_per_gen_step - 1
 
                 if spec_config is None or spec_config.spec_dec_mode.support_guided_decoder(
                 ):

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -351,7 +351,8 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         self.attention_dp_events_gather_period_ms = kv_cache_config.attention_dp_events_gather_period_ms
         self.max_num_tokens = max_num_tokens
         self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
-        self.max_total_draft_tokens = spec_config.max_total_draft_tokens if spec_config is not None else 0
+        self.max_total_draft_tokens = (spec_config.tokens_per_gen_step -
+                                       1) if spec_config is not None else 0
 
         # Determine max_attention_window_vec
         if kv_cache_config.max_attention_window is None:

tensorrt_llm/_torch/speculative/__init__.py

@@ -4,6 +4,7 @@
                         should_use_separate_draft_kv_cache)
 from .mtp import MTPEagleWorker, MTPSpecMetadata, MTPWorker
 from .ngram import NGramDrafter, NGramPoolManager
+from .pard import PARDSpecMetadata, PARDWorker
 from .save_hidden_state import (SaveHiddenStatesResourceManager,
                                 SaveHiddenStatesSpecMetadata)
 from .spec_tree_manager import SpecTreeManager
@@ -19,6 +20,8 @@
     "MTPWorker",
     "NGramDrafter",
     "NGramPoolManager",
+    "PARDSpecMetadata",
+    "PARDWorker",
     "SaveHiddenStatesResourceManager",
     "SaveHiddenStatesSpecMetadata",
     "SpecMetadata",

tensorrt_llm/_torch/speculative/eagle3.py

@@ -43,7 +43,7 @@ def __init__(self, config: "EagleDecodingConfig", dtype: torch.dtype,
         from ...llmapi.llm_args import EagleDecodingConfig
 
         if isinstance(config, EagleDecodingConfig):
-            self.max_total_draft_tokens = config.max_total_draft_tokens
+            self.max_total_draft_tokens = config.tokens_per_gen_step - 1
         else:
             self.max_total_draft_tokens = self.max_draft_len