diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 39a7289fee60..7a2b2bf0c6fa 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/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. diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 1d062e8413a8..9d9fa98b3f32 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/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: diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 59f7f1493c32..1c04ecf3a1a6 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f79395c822ee..8d807e3e6e35 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index c305a058d1b6..2dfdd9df9d93 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/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", diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py new file mode 100644 index 000000000000..5b471048d2c1 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -0,0 +1,289 @@ +""" +DraftTarget One-Model Speculative Decoding Implementation. + +This module implements a one-model approach for DraftTarget speculative decoding, +where the draft and target models share the same model engine. The draft model +layers are integrated into the target model's KV cache and run in a single forward pass. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch +from torch import nn + +from tensorrt_llm.mapping import Mapping + +from ..attention_backend import AttentionMetadata +from ..pyexecutor.sampler import TorchSampler +from .interface import SpecMetadata, SpecWorkerBase +from .mtp import MTPSampler + +if TYPE_CHECKING: + from ...llmapi.llm_args import DraftTargetDecodingConfig + + +@dataclass +class DraftTargetOneModelSpecMetadata(SpecMetadata): + """ + Metadata for DraftTarget one-model speculative decoding. + + This class manages the batch information needed for the one-model DraftTarget + approach where draft and target models share the same model engine. + Unlike Eagle3/MTP, DraftTarget does not require capturing hidden states + from the target model to pass to the draft model. + """ + + # The max number of tokens + max_num_tokens: int = 0 + # The index of the batch inputs + batch_indices_cuda: Optional[torch.Tensor] = None + + def __post_init__(self): + self.batch_indices_cuda = torch.empty( + [self.max_num_requests], + dtype=torch.int, + device="cuda", + ) + + def prepare(self): + """Prepare the metadata before model forward.""" + assert self.request_ids is not None + # Update batch indices + num_seqs = len(self.request_ids) + batch_indices = torch.arange(num_seqs, dtype=torch.int, device="cpu", pin_memory=True) + self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) + self.num_tokens -= self.num_generations * self.max_draft_len + self.is_spec_dec_tree = False + self.is_spec_dec_dynamic_tree = False + + +class DraftTargetOneModelSampler(MTPSampler): + """ + Sampler for DraftTarget one-model speculative decoding. + + Inherits from MTPSampler to reuse the speculative decoding sampling logic. + """ + + def __init__(self, args: TorchSampler.Args): + super().__init__(args, nextn=args.max_draft_len) + + +class DraftTargetOneModelWorker(SpecWorkerBase): + def __init__( + self, + spec_config: "DraftTargetDecodingConfig", + mapping: Mapping, + use_separate_draft_kv_cache: bool = False, + ): + super().__init__(use_separate_draft_kv_cache) + self.spec_config = spec_config + self.mapping = mapping + + @property + def max_draft_len(self) -> int: + return self.spec_config.max_draft_len + + def forward( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata: AttentionMetadata, + spec_metadata: DraftTargetOneModelSpecMetadata, + draft_model: nn.Module, + resource_manager=None, + ): + """ + Technically incorrect at the moment. + Leverages Eagle3/MTP setup that does this for the context + input_ids_ctx[:-1].copy_(input_prompt_ids[1:]) + In DraftTarget, we do not want to shift, which necessitates increasing the final chunk of each request by 1 + for the final accepted token. This creates a big headache since then the kv lens, seq_lens, token counts all + have to be updated and then reverted when heading back to the target. TODO: non trivially fix this issue. + """ + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._execute_guided_decoder_if_present(logits) + + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + # Prepare attention metadata for speculative decoding and save state for restore + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the first draft forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + ) + + next_draft_tokens = [] + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + 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): + if i == 0: + start_ids_gen = ( + spec_metadata.batch_indices_cuda[:num_gens] * (self.max_draft_len + 1) + ).long() + gather_ids_gen = ( + start_ids_gen + + num_accepted_tokens[num_contexts:] + - 1 + + attn_metadata.num_ctx_tokens + ) + gather_ids = torch.concat( + [spec_metadata.gather_ids[:num_contexts], gather_ids_gen], dim=0 + ) + else: + gather_ids = spec_metadata.batch_indices_cuda[:batch_size] + + if self.guided_decoder is not None: + new_tokens = inputs["input_ids"][gather_ids] + self.guided_decoder.add_draft_batch( + new_tokens, num_accepted_tokens, draft_step=i + ) + + if original_all_rank_num_tokens is not None: + if i == 0: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + elif spec_metadata.all_rank_num_seqs is not None: + attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs + + hidden_states = draft_model.model(**inputs) + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + # Disable spec-dec mode for chained draft steps + attn_metadata.use_spec_decoding = False + + logits = draft_model.logits_processor( + hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True + ) + if self.guided_decoder is not None: + d2t = getattr(draft_model.model, "d2t", None) + self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i) + + new_draft_token = self.draft_decoder(logits, draft_model) + next_draft_tokens.append(new_draft_token) + + # Update inputs and metadata for next draft step + position_ids = inputs["position_ids"][gather_ids] + 1 + if i == 0: + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata._seq_lens_cuda[:batch_size].fill_(1) + attn_metadata.on_update() + if inputs["attn_metadata"].kv_cache_manager is not None: + attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) + attn_metadata.num_contexts = 0 + 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:] + ) + attn_metadata.kv_lens_cuda[:num_contexts] += 1 + elif hasattr(attn_metadata, "kv_lens_cuda"): + attn_metadata.kv_lens_cuda[:batch_size] += 1 + + inputs = { + "input_ids": new_draft_token, + "position_ids": position_ids, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } + + next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + + # Restore attention metadata to original state + self._restore_attn_metadata_from_spec_dec(attn_metadata) + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def sample_and_accept_draft_tokens( + self, + logits: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DraftTargetOneModelSpecMetadata, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + if spec_metadata.draft_tokens is None: + draft_tokens = torch.zeros( + (num_gens, self.max_draft_len), dtype=torch.int, device=logits.device + ) + else: + draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, self.max_draft_len) + + return self._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + + def draft_decoder( + self, + logits: torch.Tensor, + draft_model: nn.Module, + ): + d2t = getattr(draft_model.model, "d2t", None) + return self._draft_sampler_greedy(logits, d2t) + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DraftTargetOneModelSpecMetadata, + ): + num_contexts = attn_metadata.num_contexts + + input_ids_ctx = self._prepare_context_input_ids( + input_ids, + attn_metadata.num_ctx_tokens, + spec_metadata.gather_ids, + accepted_tokens, + num_contexts, + ) + + input_ids_gen = accepted_tokens[num_contexts:, :].flatten() + input_ids = torch.concat([input_ids_ctx, input_ids_gen], dim=0) + + return { + "input_ids": input_ids, + "position_ids": position_ids, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 8ab48b55d1fc..d9a3da444c69 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -64,6 +64,7 @@ class SpeculativeDecodingMode(IntEnum): EAGLE3_ONE_MODEL = auto() NGRAM = auto() DRAFT_TARGET = auto() + DRAFT_TARGET_ONE_MODEL = auto() USER_PROVIDED = auto() SAVE_HIDDEN_STATES = auto() NONE = auto() @@ -85,7 +86,8 @@ def is_eagle3(self): return self == SpeculativeDecodingMode.EAGLE3 def use_one_engine(self): - return self.is_eagle3_one_model() or self.is_mtp_one_model() + return self.is_eagle3_one_model() or self.is_mtp_one_model( + ) or self.is_draft_target_one_model() def is_eagle3_one_model(self): return self == SpeculativeDecodingMode.EAGLE3_ONE_MODEL @@ -102,25 +104,30 @@ def is_none(self): def is_draft_target(self): return self == SpeculativeDecodingMode.DRAFT_TARGET + def is_draft_target_one_model(self): + return self == SpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL + def is_save_hidden_states(self): return self == SpeculativeDecodingMode.SAVE_HIDDEN_STATES def without_logits(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model() + return self.is_mtp_one_model() or self.is_eagle3_one_model( + ) or self.is_draft_target_one_model() def needs_kv_cache_rewind(self): return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.is_ngram() + ) or self.is_ngram() or self.is_draft_target_one_model() def support_overlap_scheduler(self): return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.has_draft_model() + ) or self.is_draft_target_one_model() or self.has_draft_model() def support_guided_decoder(self): return self.is_none() or self.has_spec_drafter() def support_capturable_guided_decoder(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model() + return self.is_mtp_one_model() or self.is_eagle3_one_model( + ) or self.is_draft_target_one_model() def has_draft_model(self): return self.is_eagle3() or self.is_draft_target() or self.is_mtp_eagle() @@ -138,11 +145,11 @@ def need_load_draft_weights(self): Whether the draft model and target model are in the same model engine, and the draft model needs to load weights from the separate checkpoint. """ - return self.is_eagle3_one_model() + return self.is_eagle3_one_model() or self.is_draft_target_one_model() def has_spec_decoder(self): return self.is_mtp_one_model() or self.is_mtp_eagle() or self.is_eagle3( - ) or self.is_eagle3_one_model() + ) or self.is_eagle3_one_model() or self.is_draft_target_one_model() def has_spec_drafter(self): return self.is_eagle3() or self.is_draft_target() or self.is_ngram( diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 8fad9beb2610..44bb8e1f904d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -7,6 +7,9 @@ from ..pyexecutor.sampler import TorchSampler from ..pyexecutor.seq_slot_manager import SeqSlotManager from ..speculative.interface import SpecMetadata +from .draft_target import (DraftTargetOneModelSampler, + DraftTargetOneModelSpecMetadata, + DraftTargetOneModelWorker) from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata) @@ -80,6 +83,15 @@ def get_spec_metadata(spec_config, layers_to_capture=spec_config.eagle3_layers_to_capture, allow_advanced_sampling=spec_config.allow_advanced_sampling, ) + if spec_config.spec_dec_mode.is_draft_target_one_model(): + return DraftTargetOneModelSpecMetadata( + max_draft_len=spec_config.max_draft_len, + max_total_draft_tokens=spec_config.max_total_draft_tokens, + spec_dec_mode=spec_config.spec_dec_mode, + max_num_requests=max_num_requests, + max_num_tokens=max_num_tokens, + allow_advanced_sampling=spec_config.allow_advanced_sampling, + ) if spec_config.spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesSpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -167,6 +179,8 @@ def get_spec_decoder(sampler_args: TorchSampler.Args, return TorchSampler(sampler_args) if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelSampler(sampler_args) + if spec_config.spec_dec_mode.is_draft_target_one_model(): + return DraftTargetOneModelSampler(sampler_args) raise ValueError( f"Unsupported speculative decoding mode: {spec_config.spec_dec_mode}") @@ -208,6 +222,11 @@ def get_num_spec_layers(spec_config): if spec_config.spec_dec_mode.is_eagle3_one_model(): num_eagle_layers = spec_config.num_eagle_layers return num_eagle_layers if num_eagle_layers is not None else 1 + if spec_config.spec_dec_mode.is_draft_target_one_model(): + # For DraftTarget one-model, the number of spec layers equals the + # number of draft model layers that need separate KV cache entries. + num_draft_layers = spec_config.num_draft_layers + return num_draft_layers if num_draft_layers is not None else 1 return 0 @@ -224,6 +243,9 @@ def get_spec_worker(spec_config, if spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelWorker(spec_config, mapping, use_separate_draft_kv_cache) + if spec_dec_mode.is_draft_target_one_model(): + return DraftTargetOneModelWorker(spec_config, mapping, + use_separate_draft_kv_cache) return None diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index eb9cfe8b68e4..b8e58edc9e8a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -758,7 +758,7 @@ def supports_backend(self, backend: str) -> bool: """ return True - @functools.cached_property + @property def spec_dec_mode(self): # spec_dec_mode has more functionality than the raw decoding_mode string. # Use an alias for the import here to avoid name collisions with the one for the @@ -772,6 +772,10 @@ def spec_dec_mode(self): def is_linear_tree(self) -> bool: return self.max_draft_len == self.max_total_draft_tokens + @functools.cached_property + def num_capture_layers(self) -> int: + return 0 + class KvCacheConnectorConfig(StrictBaseModel): """ @@ -1148,6 +1152,8 @@ def supports_backend(self, backend: str) -> bool: class DraftTargetDecodingConfig(DecodingBaseConfig): decoding_type: Literal["Draft_Target"] = "Draft_Target" + num_draft_layers: Optional[int] = None + _draft_target_one_model: bool = True @model_validator(mode="after") def validate_draft_target_config(self): @@ -1162,6 +1168,14 @@ def validate_draft_target_config(self): def supports_backend(self, backend: str) -> bool: return backend == "pytorch" or backend == "_autodeploy" + @functools.cached_property + def spec_dec_mode(self): + from tensorrt_llm._torch.speculative.interface import \ + SpeculativeDecodingMode as TorchSpeculativeDecodingMode + if self._draft_target_one_model: + return TorchSpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL + return TorchSpeculativeDecodingMode.DRAFT_TARGET + class MTPDecodingConfig(DecodingBaseConfig): decoding_type: Literal["MTP"] = "MTP" @@ -3102,6 +3116,13 @@ def validate_speculative_config(self): self.disable_overlap_scheduler = True self.cuda_graph_config = None self.speculative_config.max_draft_len = 1 + elif isinstance(self.speculative_config, DraftTargetDecodingConfig): + assert self.speculative_config.max_draft_len > 0 + assert self.speculative_config.speculative_model is not None, "Draft model must be specified." + if self.backend == "_autodeploy": + self.speculative_config._draft_target_one_model = False + # Invalidate cached spec_dec_mode in case it was already accessed + self.speculative_config.__dict__.pop('spec_dec_mode', None) else: self.decoding_config = None diff --git a/tests/unittest/_torch/speculative/test_draft_target.py b/tests/unittest/_torch/speculative/test_draft_target.py index 6ba477051fd3..0f963c4fcf0c 100644 --- a/tests/unittest/_torch/speculative/test_draft_target.py +++ b/tests/unittest/_torch/speculative/test_draft_target.py @@ -30,7 +30,7 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str): max_draft_len = 4 kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) cuda_graph_config = CudaGraphConfig( - batch_sizes=[1]) if use_cuda_graph else None + batch_sizes=[1, max_batch_size]) if use_cuda_graph else None llm_common_config = dict( model=target_model_dir, @@ -46,13 +46,14 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str): spec_config = DraftTargetDecodingConfig( max_draft_len=max_draft_len, speculative_model=draft_model_dir, + num_draft_layers=32, ) prompts = [ "The capital of France is", "The president of the United States is", ] - sampling_params = SamplingParams(max_tokens=32) + sampling_params = SamplingParams(max_tokens=32, temperature=0.0) llm_spec = LLM(**llm_common_config, speculative_config=spec_config) results_spec = llm_spec.generate(prompts, sampling_params)