diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 83fbab83feb8..f64febf02849 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1,3 +1,4 @@ +import inspect from dataclasses import replace from typing import Dict, Generic, List, Optional, Tuple @@ -24,6 +25,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, get_model_architecture, register_auto_model) @@ -984,6 +986,8 @@ def get_draft_model(model_config, draft_config, lm_head, model): return MTPDraftModelForCausalLM(model_config) elif spec_dec_mode.is_pard(): return PARDForCausalLM(draft_config) + elif spec_dec_mode.is_draft_target_one_model(): + return AutoModelForCausalLM.from_config(draft_config) else: raise NotImplementedError( f"get_draft_model does not support speculative decoding mode {spec_dec_mode}." @@ -1003,6 +1007,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): self.spec_worker = 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(): # Only create draft_model for modes MTP, Eagle3 (not SA) if not spec_config.spec_dec_mode.is_sa(): @@ -1037,7 +1042,7 @@ 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(): + elif spec_config.spec_dec_mode.is_external_drafter(): self.draft_config = ModelConfig.from_pretrained( model_config.spec_config.speculative_model, trust_remote_code=True, @@ -1160,10 +1165,15 @@ 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) - # PARD has independent weights; other methods share with target model - if not self.model_config.spec_config.spec_dec_mode.is_pard(): + 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_external_drafter( + ): self.draft_model.load_weights_from_target_model(self) def set_guided_decoder(self, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 691184c1fd02..60f2f8a60253 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -606,9 +606,9 @@ def _create_one_model_draft_kv_cache_manager( target_pretrained_config = self._model_engine.model.model_config.pretrained_config target_num_layers = target_pretrained_config.num_hidden_layers - # PARD: draft is a separate model, layers start from 0. + # PARD, External Drafter: 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(): + if self._speculative_config.spec_dec_mode.is_external_drafter(): num_draft_layers = self._draft_config.pretrained_config.num_hidden_layers spec_dec_layer_mask = [True] * num_draft_layers else: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 9306523d6b57..f4a2b9e0a2a9 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -602,6 +602,11 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): == self.mapping.cp_size - 1 else 0), req_beam_width, req) else: + # Chunked prefill may schedule the same request across multiple + # context chunks. Sequence allocation must happen only once. + if not req.is_first_context_chunk: + continue + if self.impl.add_sequence(req.py_request_id, req.prompt_len, req_beam_width, req): for _ in range(self.num_extra_kv_tokens): diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 3e938628a058..4771380ea3ba 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) @@ -18,6 +20,8 @@ get_spec_worker, update_spec_config_from_model_config) __all__ = [ + "DraftTargetOneModelSpecMetadata", + "DraftTargetOneModelWorker", "Eagle3SpecMetadata", "MTPEagleWorker", "MTPSampler", diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py new file mode 100644 index 000000000000..c026b6d5b290 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +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._utils import prefer_pinned +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=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.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 _prepare_attn_metadata_for_draft_target( + self, + attn_metadata: AttentionMetadata, + spec_metadata: DraftTargetOneModelSpecMetadata, + ): + """ + Save the attention metadata fields modified by DraftTarget. + + During CUDA-graph warmup, kv_lens_cuda is also saved/restored to avoid + cross-warmup accumulation. During capture and normal inference we keep + kv_lens_cuda live so the updates persist. + """ + is_capturing = torch.cuda.is_current_stream_capturing() + + if ( + spec_metadata.is_cuda_graph + and not is_capturing + and hasattr(attn_metadata, "kv_lens_cuda") + and isinstance(attn_metadata.kv_lens_cuda, torch.Tensor) + ): + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda", "kv_lens_cuda") + else: + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + + def _update_kv_after_first_draft_step( + self, + attn_metadata: AttentionMetadata, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + batch_size: 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] + ) + attn_metadata.kv_lens_cuda[:num_contexts] += 1 + + # Some attention backends keep extra indexing state derived from + # seq_lens / kv_lens that must be refreshed for chained drafting. + attn_metadata.update_for_spec_dec() + + def _update_kv_for_chained_draft_step( + self, + attn_metadata: AttentionMetadata, + batch_size: int, + ): + if hasattr(attn_metadata, "kv_lens_cuda"): + attn_metadata.kv_lens_cuda[:batch_size] += 1 + + attn_metadata.update_for_spec_dec() + + 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_draft_target(attn_metadata, spec_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 + self._update_kv_after_first_draft_step( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + else: + self._update_kv_for_chained_draft_step(attn_metadata, batch_size) + + 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 + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + + if num_contexts > 0: + input_ids_ctx = self._prepare_context_input_ids( + input_ids, + attn_metadata.num_ctx_tokens, + spec_metadata.gather_ids, + accepted_tokens, + num_contexts, + ).to(torch.int32) + else: + 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) + else: + input_ids_gen = torch.empty(0, dtype=torch.int32, device="cuda") + + input_ids = torch.cat([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 06b41c8b6f2a..a887cb22a595 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -65,6 +65,7 @@ class SpeculativeDecodingMode(IntEnum): NGRAM = auto() SA = auto() DRAFT_TARGET = auto() + DRAFT_TARGET_ONE_MODEL = auto() USER_PROVIDED = auto() SAVE_HIDDEN_STATES = auto() PARD = auto() @@ -88,7 +89,7 @@ def is_eagle3(self): def use_one_engine(self): return self.is_eagle3_one_model() or self.is_mtp_one_model( - ) or self.is_pard() or self.is_sa() + ) or self.is_external_drafter() or self.is_sa() def is_eagle3_one_model(self): return self == SpeculativeDecodingMode.EAGLE3_ONE_MODEL @@ -111,27 +112,34 @@ 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 is_external_drafter(self): + return self.is_pard() or self.is_draft_target_one_model() + def without_logits(self): return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.is_pard() or self.is_sa() + ) or self.is_external_drafter() or self.is_sa() 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_sa() or self.is_pard() + ) or self.is_ngram() or self.is_sa() or self.is_external_drafter() def support_overlap_scheduler(self): return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.is_sa() or self.has_draft_model() or self.is_pard() + ) or self.is_sa() or self.has_draft_model() or self.is_external_drafter( + ) 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( - ) or self.is_pard() or self.is_sa() + ) or self.is_external_drafter() or self.is_sa() def has_draft_model(self): return self.is_eagle3() or self.is_draft_target() or self.is_mtp_eagle() @@ -149,11 +157,12 @@ 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() or self.is_pard() + return self.is_eagle3_one_model() or self.is_external_drafter() 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_pard() or self.is_sa() + ) or self.is_eagle3_one_model() or self.is_external_drafter( + ) or self.is_sa() 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 da767444dd6a..17892bee8c3a 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -10,6 +10,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) @@ -101,6 +104,15 @@ def get_spec_metadata(spec_config, max_num_requests=max_num_requests, 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, @@ -217,6 +229,8 @@ def get_spec_decoder( nextn=spec_config.tokens_per_gen_step - 1) if spec_config.spec_dec_mode.is_sa(): return SASampler(sampler_args, max_draft_len=spec_config.max_draft_len) + 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}") @@ -278,6 +292,9 @@ def get_spec_worker(spec_config, return PARDWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_sa(): return SAWorker(spec_config, model_config) + 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 72e478d0c22b..c6b7f02b96ae 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 @@ -777,6 +777,9 @@ def tokens_per_gen_step(self) -> int: """Total tokens per gen request in one spec dec iteration (including golden token).""" return 1 + self.max_total_draft_tokens + def num_capture_layers(self) -> int: + return 0 + class KvCacheConnectorConfig(StrictBaseModel): """ @@ -1181,6 +1184,7 @@ def supports_backend(self, backend: str) -> bool: class DraftTargetDecodingConfig(DecodingBaseConfig): decoding_type: Literal["Draft_Target"] = "Draft_Target" + _draft_target_one_model: bool = PrivateAttr(True) @model_validator(mode="after") def validate_draft_target_config(self): @@ -1195,6 +1199,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" @@ -3272,6 +3284,11 @@ 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 else: self.decoding_config = None diff --git a/tests/unittest/_torch/speculative/test_draft_len_schedule.py b/tests/unittest/_torch/speculative/test_draft_len_schedule.py index 32c491460f34..67e5514cb9c2 100644 --- a/tests/unittest/_torch/speculative/test_draft_len_schedule.py +++ b/tests/unittest/_torch/speculative/test_draft_len_schedule.py @@ -29,7 +29,6 @@ def enforce_single_worker(monkeypatch): "drafter_type,schedule", [ ("ngram", {1: 3, 4: 2, 8: 1}), - ("model_drafter", {1: 3, 4: 2, 8: 1}), ], ) @pytest.mark.high_cuda_memory @@ -116,6 +115,7 @@ def test_correctness_across_batch_sizes(drafter_type: str, schedule: dict): is_public_pool=False, ) else: + # skipped for move to 1 model spec_config_fixed = DraftTargetDecodingConfig( max_draft_len=max_draft_len, speculative_model=str(draft_model), @@ -142,7 +142,6 @@ def test_correctness_across_batch_sizes(drafter_type: str, schedule: dict): "drafter_type,draft_schedule", [ ("ngram", {1: 5, 4: 4, 5: 3, 6: 2, 7: 1}), - ("model_drafter", {1: 5, 4: 4, 5: 3, 6: 2, 7: 1}), ], ) @pytest.mark.high_cuda_memory @@ -180,6 +179,7 @@ def test_draft_len_schedule_functionality( draft_len_schedule=draft_schedule, ) else: + # skipped for move to 1 model spec_config = DraftTargetDecodingConfig( max_draft_len=5, speculative_model=str(llm_models_root() / "llama-3.2-models" / "Llama-3.2-3B-Instruct"), diff --git a/tests/unittest/_torch/speculative/test_draft_target.py b/tests/unittest/_torch/speculative/test_draft_target.py index 6ba477051fd3..9f2b7d407d82 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, @@ -52,7 +52,7 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str): "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)