[None][feat] Draft Target One Model

tensorrt_llm/_torch/model_config.py

@@ -131,6 +131,13 @@ class ModelConfig(Generic[TConfig]):
     # If true, ONLY the vision encoder part of the full model is loaded/executed.
     mm_encoder_only: bool = False
 
+    # Offset applied to every attention layer_idx at construction time.
+    # In one-model speculative decoding the draft model's layers must use
+    # KV-cache indices that don't collide with the target model's layers.
+    # Setting this to num_target_layers when building the draft model avoids
+    # the need for a post-hoc index fixup.
+    layer_idx_offset: int = 0
+
     def __setattr__(self, key, value):
         """
         Prevent modification of frozen instance attributes.

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -1,3 +1,4 @@
+import inspect
 from typing import Dict, Generic, List, Optional, Tuple
 
 import torch
@@ -23,6 +24,7 @@
                            should_use_separate_draft_kv_cache)
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
+from .modeling_auto import AutoModelForCausalLM
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, TModel,
                              register_auto_model)
 
@@ -916,6 +918,16 @@ 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_draft_target_one_model():
+        # Set the layer index offset on the draft config so that attention
+        # layers are created with KV-cache indices that don't collide with the
+        # target model's layers (target uses [0, N), draft uses [N, N+M)).
+        num_target_layers = model_config.pretrained_config.num_hidden_layers
+        draft_config._frozen = False
+        draft_config.layer_idx_offset = num_target_layers
+        draft_config._frozen = True
+        draft_model = AutoModelForCausalLM.from_config(draft_config)
+        return draft_model
     else:
         raise NotImplementedError(
             f"get_draft_model does not support speculative decoding mode {spec_dec_mode}."
@@ -934,6 +946,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
         self.draft_config = None
         self.use_separate_draft_kv_cache = False
         spec_config = getattr(model_config, 'spec_config', None)
+        self.spec_config = spec_config
         if spec_config and spec_config.spec_dec_mode.use_one_engine():
             if spec_config.spec_dec_mode.is_eagle3_one_model():
                 if spec_config.eagle3_model_arch == "mistral_large3":
@@ -965,6 +978,18 @@ 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_draft_target_one_model():
+                self.draft_config = ModelConfig.from_pretrained(
+                    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,  # Draft model doesn't need spec_config
+                    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)
@@ -978,12 +1003,13 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                 use_separate_draft_kv_cache=self.use_separate_draft_kv_cache)
             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":
+            if self.draft_config is not None and (
+                    spec_config.spec_dec_mode.is_draft_target_one_model()
+                    or 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
-                    model_config.extra_attrs[key].update(value)
+                    if key in ('attn_layers', 'mla_layers'):
+                        assert key in model_config.extra_attrs
+                        model_config.extra_attrs[key].update(value)
         self.layer_idx = -1
 
     def forward(
@@ -1064,9 +1090,16 @@ def load_weights(self,
     def load_draft_weights(self,
                            weights: Dict,
                            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)
+        args = inspect.getfullargspec(self.draft_model.load_weights).args
+        if "weight_mapper" in args:
+            self.draft_model.load_weights(weights=weights,
+                                          weight_mapper=weight_mapper)
+        else:
+            self.draft_model.load_weights(weights=weights)
+
+        if self.spec_config and not self.spec_config.spec_dec_mode.is_draft_target_one_model(
+        ):
+            self.draft_model.load_weights_from_target_model(self)
 
     def set_guided_decoder(self,
                            guided_decoder: CapturableGuidedDecoder) -> bool:

tensorrt_llm/_torch/modules/attention.py

@@ -176,8 +176,14 @@ def __init__(
             attn_output_gate (Optional[bool]): Determines whether to use an output gate in the attention Op. If False, the decision is automatically handled by the attention backend based on its capabilities.
         """
         super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_idx_str = str(layer_idx)
+        # Apply the per-config layer index offset so that, e.g., draft-model
+        # attention layers are created with KV-cache indices that don't collide
+        # with the target model's layers.
+        offset = getattr(config, 'layer_idx_offset',
+                         0) if config is not None else 0
+        effective_layer_idx = layer_idx + offset if layer_idx is not None else layer_idx
+        self.layer_idx = effective_layer_idx
+        self.layer_idx_str = str(effective_layer_idx)
 
         self.register_to_config = False
         # We only register TRTLLM attention layers to config.
@@ -187,7 +193,7 @@ def __init__(
             suffix = 0
             # Makes sure there is no duplicate attention layer identifier.
             while self.layer_idx_str in config.extra_attrs["attn_layers"]:
-                self.layer_idx_str = str(layer_idx) + f"_{suffix}"
+                self.layer_idx_str = str(effective_layer_idx) + f"_{suffix}"
                 suffix += 1
             config.extra_attrs["attn_layers"][self.layer_idx_str] = weakref.ref(
                 self)
@@ -778,8 +784,12 @@ def __init__(
             enable_helix_test (bool): Whether to enable helix unit test.
         """
         super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_idx_str = str(layer_idx)
+        # Apply the per-config layer index offset (see Attention.__init__).
+        offset = getattr(config, 'layer_idx_offset',
+                         0) if config is not None else 0
+        effective_layer_idx = layer_idx + offset if layer_idx is not None else layer_idx
+        self.layer_idx = effective_layer_idx
+        self.layer_idx_str = str(effective_layer_idx)
         self.dtype = dtype
 
         self.hidden_size = hidden_size

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -379,6 +379,16 @@ def allocation_scope(current_stage: ExecutorMemoryType,
 
     validate_feature_combination(llm_args, model_engine, llm_args.sampler_type)
 
+    # Auto-populate num_draft_layers from the draft model's config for
+    # DraftTarget one-model mode, so get_num_spec_layers returns the correct
+    # value during KV cache creation.
+    if (spec_config is not None
+            and spec_config.spec_dec_mode.is_draft_target_one_model()):
+        draft_config = getattr(model_engine.model, 'draft_config', None)
+        if draft_config is not None and spec_config.num_draft_layers is None:
+            spec_config.num_draft_layers = \
+                draft_config.pretrained_config.num_hidden_layers
+
     calibrator = get_calibrator()
     layer_wise_benchmarks_config = llm_args.layer_wise_benchmarks_config
     calibrator.init(layer_wise_benchmarks_config.calibration_mode,

tensorrt_llm/_torch/speculative/__init__.py

@@ -1,4 +1,6 @@
 from .auto_heuristic import suggest_spec_config
+from .draft_target import (DraftTargetOneModelSpecMetadata,
+                           DraftTargetOneModelWorker)
 from .eagle3 import Eagle3SpecMetadata
 from .interface import (SpecMetadata, SpecWorkerBase,
                         should_use_separate_draft_kv_cache)
@@ -13,6 +15,8 @@
                     get_spec_worker, update_spec_config_from_model_config)
 
 __all__ = [
+    "DraftTargetOneModelSpecMetadata",
+    "DraftTargetOneModelWorker",
     "Eagle3SpecMetadata",
     "MTPEagleWorker",
     "MTPSpecMetadata",