From e0ea6e78d2ddd29e0eafe71d8cf38268d8d529c7 Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Tue, 1 Jul 2025 17:13:12 +0800 Subject: [PATCH 1/3] v0.1: cherry-pick and follow Fanrong suggestions Signed-off-by: wili-65535 --- examples/llm-api/quickstart_advanced.py | 5 +- tensorrt_llm/_torch/model_config.py | 2 +- .../_torch/models/modeling_speculative.py | 2 +- tensorrt_llm/_torch/pyexecutor/_util.py | 12 +- tensorrt_llm/_torch/pyexecutor/config.py | 3 +- .../_torch/pyexecutor/model_engine.py | 26 ++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 +- .../_torch/pyexecutor/py_executor_creator.py | 13 +- .../_torch/pyexecutor/resource_manager.py | 15 +- tensorrt_llm/_torch/speculative/__init__.py | 24 ++- .../_torch/speculative/draft_target.py | 24 +-- tensorrt_llm/_torch/speculative/eagle3.py | 43 +----- tensorrt_llm/_torch/speculative/interface.py | 39 +---- tensorrt_llm/_torch/speculative/mtp.py | 52 +------ tensorrt_llm/_torch/speculative/ngram.py | 40 +---- tensorrt_llm/_torch/speculative/utils.py | 72 ++++++--- tensorrt_llm/llmapi/llm_args.py | 142 ++++++++++++------ .../defs/accuracy/accuracy_core.py | 3 - .../defs/accuracy/test_llm_api_pytorch.py | 3 +- .../_torch/speculative/test_draft_target.py | 4 +- .../_torch/speculative/test_eagle3.py | 2 +- tests/unittest/_torch/speculative/test_mtp.py | 11 +- .../unittest/_torch/speculative/test_ngram.py | 8 +- .../_torch/speculative/test_user_provided.py | 31 ++-- 24 files changed, 259 insertions(+), 325 deletions(-) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index cfecf9c4aff2..22a79a7eb6e6 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -169,14 +169,15 @@ def setup_llm(args): elif spec_decode_algo == "EAGLE3": spec_config = EagleDecodingConfig( max_draft_len=args.spec_decode_nextn, - pytorch_weights_path=args.draft_model_dir, + speculative_model=args.draft_model_dir, eagle3_one_model=args.use_one_model) elif spec_decode_algo == "DRAFT_TARGET": spec_config = DraftTargetDecodingConfig( max_draft_len=args.spec_decode_nextn, - pytorch_weights_path=args.draft_model_dir) + speculative_model=args.draft_model_dir) elif spec_decode_algo == "NGRAM": spec_config = NGramDecodingConfig( + max_draft_len=args.spec_decode_nextn, prompt_lookup_num_tokens=args.spec_decode_nextn, max_matching_ngram_size=args.max_matching_ngram_size, is_keep_all=True, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index dab7bc460143..830cd5bda6a7 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -70,7 +70,7 @@ class ModelConfig(Generic[TConfig]): # to support mixed quantization. skip_create_weights_in_init: bool = False - spec_config: Optional["SpecConfig"] = None + spec_config: Optional["DecodingBaseConfig"] = None lora_config: Optional["LoraConfig"] = None is_generation: bool = True diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 55591cc354f4..11f360965e26 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -340,7 +340,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): model_config, 'spec_config', None ) and model_config.spec_config.spec_dec_mode.use_one_engine(): draft_config = ModelConfig.from_pretrained( - model_config.spec_config.draft_model_path, + model_config.spec_config.speculative_model, trust_remote_code=True, attn_backend=model_config.attn_backend, moe_backend=model_config.moe_backend, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 91a7d3ec00e8..39d253a96d86 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -17,7 +17,7 @@ from tensorrt_llm.mapping import Mapping from ..model_config import ModelConfig -from ..speculative import get_spec_decoder +from ..speculative import get_num_extra_kv_tokens, get_spec_decoder from .config_utils import is_mla, is_nemotron_hybrid from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver from .llm_request import ExecutorResponse @@ -157,11 +157,11 @@ def _get_token_num_for_estimation(self) -> int: if not pytorch_backend_config.disable_overlap_scheduler: num_extra_tokens_per_seq = num_extra_tokens_per_seq + 1 if spec_cfg is not None: - num_extra_tokens_per_seq += spec_cfg.max_draft_tokens + num_extra_tokens_per_seq += spec_cfg.max_draft_len if spec_cfg is not None: - num_extra_tokens_per_seq += spec_cfg.max_draft_tokens - num_extra_tokens_per_seq += spec_cfg.num_extra_kv_tokens + num_extra_tokens_per_seq += spec_cfg.max_draft_len + num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg) for req in self._dummy_reqs: num_req_tokens = len(req.input_token_ids) + num_extra_tokens_per_seq # Requests cannot share KV cache blocks. Round up to nearest integer multiple of block size. @@ -538,7 +538,7 @@ def create_py_executor_instance( disable_overlap_scheduler, max_batch_size=executor_config.max_batch_size, max_beam_width=executor_config.max_beam_width, - max_draft_tokens=spec_config.max_draft_tokens + max_draft_tokens=spec_config.max_draft_len if spec_config is not None else 0, kv_cache_transceiver=kv_cache_transceiver, draft_model_engine=draft_model_engine, @@ -550,7 +550,7 @@ def create_torch_sampler_args(executor_config: ExecutorConfig, mapping: Mapping, *, max_seq_len: int, enable_mixed_sampler: bool): max_num_sequences = executor_config.max_batch_size * mapping.pp_size max_draft_tokens = (0 if executor_config.speculative_config is None else - executor_config.speculative_config.max_draft_tokens) + executor_config.speculative_config.max_draft_len) return TorchSampler.Args( max_seq_len=max_seq_len, max_draft_tokens=max_draft_tokens, diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index be243c5e8e75..e4d5cc078816 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/tensorrt_llm/_torch/pyexecutor/config.py @@ -8,7 +8,6 @@ from ...logger import logger from ...mapping import Mapping from ..model_config import MoeLoadBalancerConfig -from ..speculative import SpecConfig from .resource_manager import BaseResourceManager @@ -110,7 +109,7 @@ def update_executor_config( pytorch_backend_config: Optional[PyTorchConfig] = None, mapping: Optional[Mapping] = None, build_config: Optional[BuildConfig] = None, - speculative_config: Optional[SpecConfig] = None, + speculative_config: Optional["DecodingBaseConfig"] = None, hf_model_dir: Optional[str] = None, max_input_len: Optional[int] = None, max_seq_len: Optional[int] = None): diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 1aa0849ec2ea..07f0faf1853e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -21,6 +21,8 @@ import tensorrt_llm.bindings.internal.userbuffers as ub from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors +from tensorrt_llm._torch.speculative import ( + get_num_extra_kv_tokens, update_spec_config_from_model_config) from tensorrt_llm._torch.speculative.mtp import SampleStateTensorsMTP from tensorrt_llm._utils import (is_trace_enabled, local_mpi_rank, local_mpi_size, nvtx_range, release_gc, @@ -51,7 +53,7 @@ run_concurrently, timing) from ..modules.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, maybe_create_moe_load_balancer) -from ..speculative import SpecConfig, SpecMetadata, get_spec_metadata +from ..speculative import SpecMetadata, get_spec_metadata from ..utils import (get_model_extra_attrs, set_torch_compiling, with_model_extra_attrs) from .config import LoadFormat, PyTorchConfig @@ -353,7 +355,7 @@ def __init__( mapping: Optional[Mapping] = None, attn_runtime_features: Optional[AttentionRuntimeFeatures] = None, dist: Optional[MPIDist] = None, - spec_config: Optional[SpecConfig] = None, + spec_config: Optional["DecodingBaseConfig"] = None, guided_decoding_config: Optional[GuidedDecodingConfig] = None, lora_config: Optional[LoraConfig] = None, is_draft_model: bool = False, @@ -455,8 +457,9 @@ def __init__( if self.is_spec_decode: self.spec_metadata = None - self.spec_config.update_from_model_config(self.model.config) - max_num_draft_tokens = self.spec_config.max_draft_tokens * batch_size + update_spec_config_from_model_config(self.spec_config, + self.model.config) + max_num_draft_tokens = self.spec_config.max_draft_len * batch_size self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, device='cuda') @@ -472,7 +475,7 @@ def __init__( device='cuda') self.without_logits = self.spec_config.spec_dec_mode.without_logits( ) - self.max_draft_len = spec_config.max_draft_tokens + self.max_draft_len = spec_config.max_draft_len else: self.without_logits = False self.max_draft_len = 0 @@ -858,6 +861,7 @@ def _set_up_spec_metadata( if no_cache: return get_spec_metadata( self.spec_config, + self.model.config, self.batch_size, max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, @@ -867,6 +871,7 @@ def _set_up_spec_metadata( return self.spec_metadata self.spec_metadata = get_spec_metadata( self.spec_config, + self.model.config, self.batch_size, max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, @@ -951,7 +956,7 @@ def _round_up_batch_size(self, batch_size: int) -> int: def _maybe_get_cuda_graph( self, batch: ScheduledRequests, - spec_config: Optional[SpecConfig] = None + spec_config: Optional["DecodingBaseConfig"] = None ) -> Optional[DecodingCUDAGraphRunner]: """ Get a CUDA graph runner or return None (e.g. if CUDA graphs are disabled @@ -961,7 +966,7 @@ def _maybe_get_cuda_graph( if ExpertStatistic.set_iter(self.iter_counter): return None - spec_max_draft_tokens = spec_config.max_draft_tokens if self.is_spec_decode else 0 + spec_max_draft_tokens = spec_config.max_draft_len if self.is_spec_decode else 0 can_run_cuda_graph = batch.can_run_cuda_graph batch_size = len(batch.generation_requests) if self._run_cuda_graphs and self.enable_attention_dp and self.mapping.tp_size > 1: @@ -1078,7 +1083,7 @@ def init_meta_tensor(t: torch.Tensor): if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): - weights = load_weights(self.spec_config.draft_model_path) + weights = load_weights(self.spec_config.speculative_model) model.load_draft_weights(weights) elif load_format == LoadFormat.DUMMY: @@ -1263,7 +1268,7 @@ def _prepare_tp_inputs( if not self._disable_overlap_scheduler and self.is_spec_decode: spec_dec_mode = self.spec_config.spec_dec_mode assert spec_dec_mode.support_overlap_scheduler( - ), f"{self.spec_config.spec_dec_name} does not support overlap scheduler" + ), f"{self.spec_config.spec_dec_mode} does not support overlap scheduler" # will contain previous batch indices of generation requests previous_batch_indices = [] @@ -1491,8 +1496,7 @@ def previous_seq_slots_device(): attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq, - num_extra_kv_tokens=0 if self.spec_config is None else - self.spec_config.num_extra_kv_tokens) + num_extra_kv_tokens=get_num_extra_kv_tokens(self.spec_config)) attn_metadata.kv_cache_manager = kv_cache_manager attn_metadata.prepare() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 7ea20581f992..ee338971216b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -17,6 +17,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager +from tensorrt_llm._torch.speculative import get_draft_model_prompt from tensorrt_llm._utils import (customized_gc_thresholds, global_mpi_rank, is_trace_enabled, nvtx_range, trace_func) from tensorrt_llm.bindings.executor import (DisServingRequestStats, @@ -979,7 +980,7 @@ def _prepare_draft_requests(self): LlmRequestState.DISAGG_GENERATION_INIT): continue req.py_last_draft_tokens = req.py_draft_tokens - max_draft_len = self.model_engine.spec_config.max_draft_tokens + max_draft_len = self.model_engine.spec_config.max_draft_len if max_draft_len > 0: req.py_draft_tokens = [0] * max_draft_len @@ -1767,9 +1768,10 @@ def _prepare_draft_batch( num_rejected_tokens = num_draft_tokens - num_accepted_tokens assert num_rejected_tokens >= 0 - spec_config = self.model_engine.spec_config + spec_mode = self.model_engine.spec_config.spec_dec_mode beam_idx = 0 - input_tokens = spec_config.get_draft_model_prompt( + input_tokens = get_draft_model_prompt( + spec_mode, request.get_tokens()[beam_idx]) def create_new_request(input_tokens): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 91441e50f594..c24e6bc7a379 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -19,7 +19,8 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist -from ..speculative import get_spec_drafter, get_spec_resource_manager +from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, + get_spec_resource_manager) from ._util import (KvCacheCreator, create_py_executor_instance, instantiate_sampler, is_mla) from .config import PyTorchConfig @@ -247,10 +248,10 @@ def create_py_executor( draft_spec_config = copy.copy(spec_config) # The draft model won't have any draft tokens attached to # generation requests when we invoke it autoregressively - draft_spec_config.max_draft_tokens = 0 + draft_spec_config.max_draft_len = 0 draft_model_engine = PyTorchModelEngine( - model_path=spec_config.draft_model_path, + model_path=spec_config.speculative_model, pytorch_backend_config=pytorch_backend_config, batch_size=executor_config.max_batch_size, max_beam_width=executor_config.max_beam_width, @@ -276,11 +277,11 @@ def create_py_executor( if not pytorch_backend_config.disable_overlap_scheduler: max_seq_len = model_engine.max_seq_len + 1 if spec_config is not None: - max_seq_len += spec_config.max_draft_tokens + max_seq_len += spec_config.max_draft_len if spec_config is not None: - max_seq_len += spec_config.num_extra_kv_tokens - max_seq_len += spec_config.max_draft_tokens + max_seq_len += get_num_extra_kv_tokens(spec_config) + max_seq_len += spec_config.max_draft_len executor_config.max_seq_len = max_seq_len executor_config.max_num_tokens = model_engine.max_num_tokens diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 360ec6679d53..68a82c1ca3a7 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2,7 +2,7 @@ import math from abc import ABC, abstractmethod from collections import OrderedDict, defaultdict -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import torch @@ -22,9 +22,6 @@ from tensorrt_llm._utils import mpi_comm -if TYPE_CHECKING: - from ..speculative.interface import SpecConfig - BufferManagerCpp = tensorrt_llm.bindings.internal.runtime.BufferManager KVCacheManagerCpp = tensorrt_llm.bindings.internal.batch_manager.KVCacheManager KvCacheConfigCpp = tensorrt_llm.bindings.executor.KvCacheConfig @@ -83,7 +80,7 @@ def shutdown(self): def get_pp_layers( num_layers: int, mapping: Mapping, - spec_config: Optional["SpecConfig"] = None, + spec_config: Optional["DecodingBaseConfig"] = None, layer_mask: Optional[List[bool]] = None, ) -> Tuple[List[int], int]: from ..speculative.utils import get_num_spec_layers @@ -127,7 +124,7 @@ def __init__( max_batch_size: int, mapping: Mapping, dtype: DataType = DataType.HALF, - spec_config: Optional["SpecConfig"] = None, + spec_config: Optional["DecodingBaseConfig"] = None, layer_mask: Optional[List[bool]] = None, max_num_tokens: int = 8192, model_config: Optional[ModelConfig] = None, @@ -178,7 +175,9 @@ def __init__( self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 # Some speculative decoding methods need to use different kv lengths for the # draft/target layers. Add extra tokens to handle this issue. - self.num_extra_kv_tokens = 0 if spec_config is None else spec_config.num_extra_kv_tokens + # Import here to avoid circular imports + from ..speculative import get_num_extra_kv_tokens + self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) self.event_buffer_max_size = kv_cache_config.event_buffer_max_size self.max_num_tokens = max_num_tokens @@ -902,7 +901,7 @@ def __init__( max_batch_size: int, mapping: Mapping, dtype: DataType = DataType.HALF, - spec_config: Optional["SpecConfig"] = None, + spec_config: Optional["DecodingBaseConfig"] = None, ) -> None: # mamba hybrid cache requires block reuse to be disabled in KV cache config diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 1ec11d22be6c..5feb0760feb5 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -1,31 +1,27 @@ -from .draft_target import DraftTargetConfig -from .eagle3 import Eagle3Config, Eagle3SpecMetadata -from .interface import SpecConfig, SpecMetadata -from .mtp import MTPConfig, MTPEagleWorker, MTPSpecMetadata, MTPWorker -from .ngram import NGramConfig, NGramDrafter, NGramPoolManager -from .user_provided import UserProvidedConfig -from .utils import (get_num_spec_layers, get_spec_decoder, get_spec_drafter, +from .eagle3 import Eagle3SpecMetadata +from .interface import SpecMetadata +from .mtp import MTPEagleWorker, MTPSpecMetadata, MTPWorker +from .ngram import NGramDrafter, NGramPoolManager +from .utils import (get_draft_model_prompt, get_num_extra_kv_tokens, + get_num_spec_layers, get_spec_decoder, get_spec_drafter, get_spec_metadata, get_spec_resource_manager, - get_spec_worker) + get_spec_worker, update_spec_config_from_model_config) __all__ = [ - "DraftTargetConfig", - "Eagle3Config", "Eagle3SpecMetadata", - "MTPConfig", "MTPEagleWorker", "MTPSpecMetadata", "MTPWorker", - "NGramConfig", "NGramDrafter", "NGramPoolManager", - "SpecConfig", "SpecMetadata", - "UserProvidedConfig", + "get_draft_model_prompt", + "get_num_extra_kv_tokens", "get_num_spec_layers", "get_spec_decoder", "get_spec_drafter", "get_spec_metadata", "get_spec_resource_manager", "get_spec_worker", + "update_spec_config_from_model_config", ] diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index 690bb3bc98a3..3fdf4865c9b6 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -1,28 +1,6 @@ from dataclasses import dataclass -import torch - -from .interface import SpecConfig, SpecMetadata, SpeculativeDecodingMode - - -@dataclass -class DraftTargetConfig(SpecConfig): - spec_dec_name: str = "DRAFT_TARGET" - - def __post_init__(self): - if self.draft_model_path is None: - raise ValueError("Path to Draft weights must be specified.") - - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - self.num_extra_kv_tokens = 0 - - def update_from_model_config(self, model_config): - pass - - def get_draft_model_prompt(self, - input_tokens: torch.Tensor) -> torch.Tensor: - return input_tokens +from .interface import SpecMetadata @dataclass diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 78a18100c73c..22f3f645e43d 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -4,7 +4,6 @@ import torch from torch import nn -from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata @@ -12,42 +11,10 @@ from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.sampler import TorchSampler from ..pyexecutor.scheduler import ScheduledRequests -from .interface import SpecConfig, SpecMetadata, SpeculativeDecodingMode +from .interface import SpecMetadata from .mtp import MTPSampler -@dataclass -class Eagle3Config(SpecConfig): - spec_dec_name: str = "EAGLE3" - num_layers: int = 0 - hidden_size: int = 0 - eagle3_one_model: bool = True - - def __post_init__(self): - if self.draft_model_path is None: - raise ValueError("Path to EAGLE3 weights must be specified.") - - if self.eagle3_one_model: - self.spec_dec_mode = SpeculativeDecodingMode.EAGLE3_ONE_MODEL - self.num_extra_kv_tokens = self.max_draft_tokens - 1 - else: - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - logger.info(f"EAGLE3 Config: {self}") - - def update_from_model_config(self, model_config): - self.num_layers = model_config.num_hidden_layers - self.hidden_size = model_config.hidden_size - self.dtype = model_config.torch_dtype - - def get_draft_model_prompt(self, - input_tokens: torch.Tensor) -> torch.Tensor: - """ - Eagle3 always throws away the first token when processing draft inputs - """ - return input_tokens[1:] - - class Eagle3ResourceManager(BaseResourceManager): """ Eagle3 needs to save the hidden states for the draft model. When using @@ -55,11 +22,11 @@ class Eagle3ResourceManager(BaseResourceManager): and one for the draft model. Use this class to manage the hidden states. """ - def __init__(self, config: Eagle3Config, dtype: torch.dtype, + def __init__(self, config: "EagleDecodingConfig", dtype: torch.dtype, hidden_size: int, max_num_requests: int, max_seq_len: int, max_num_tokens: int): self.dtype = dtype - self.max_draft_tokens = config.max_draft_tokens + self.max_draft_tokens = config.max_draft_len self.hidden_size = hidden_size self.max_num_requests = max_num_requests self.max_seq_len = max_seq_len @@ -293,10 +260,10 @@ def __init__(self, args: TorchSampler.Args): class Eagle3OneModelWorker(nn.Module): - def __init__(self, spec_config: Eagle3Config, mapping: Mapping): + def __init__(self, spec_config: "EagleDecodingConfig", mapping: Mapping): super().__init__() self.spec_config = spec_config - self.max_draft_tokens = self.spec_config.max_draft_tokens + self.max_draft_tokens = self.spec_config.max_draft_len self.mapping = mapping @torch.compile(mode="max-autotune-no-cudagraphs") diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index b369e6bb0042..b7ceed8cebec 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -7,7 +7,6 @@ from ..._utils import get_sm_version from ..attention_backend.trtllm import AttentionBackend, TrtllmAttention -from ..model_config import TConfig class SpeculativeDecodingMode(IntEnum): @@ -16,7 +15,7 @@ class SpeculativeDecodingMode(IntEnum): EAGLE3 = auto() EAGLE3_ONE_MODEL = auto() NGRAM = auto() - DRAFT_TARGET = auto() + DRAFTTARGET = auto() USER_PROVIDED = auto() NONE = auto() @@ -45,7 +44,7 @@ def is_none(self): return self == SpeculativeDecodingMode.NONE def is_draft_target(self): - return self == SpeculativeDecodingMode.DRAFT_TARGET + return self == SpeculativeDecodingMode.DRAFTTARGET def without_logits(self): return self.is_mtp() or self.is_eagle3_one_model() @@ -106,38 +105,6 @@ def from_string(name: Optional[str]) -> "SpeculativeDecodingMode": return SpeculativeDecodingMode[name.upper()] -@dataclass -class SpecConfig: - """ - Configuration for speculative decoding. - """ - # The name of speculative decoding. - spec_dec_name = None - # The mode of speculative decoding. - spec_dec_mode: SpeculativeDecodingMode = SpeculativeDecodingMode.NONE - # The max number of draft tokens - max_draft_tokens: int = 1024 - # The path to the draft model - draft_model_path: Optional[str] = None - # The number of extra kv tokens - num_extra_kv_tokens: int = 0 - - def __post_init__(self) -> None: - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - - def update_from_model_config(self, model_config: TConfig): - pass - - def get_draft_model_prompt(self, - input_tokens: torch.Tensor) -> torch.Tensor: - """ - Override for spec dec modes that need to preprocess prompt - tokens before passing them to the draft model. - """ - return input_tokens - - @dataclass class SpecMetadata: """ @@ -180,7 +147,7 @@ class SpecMetadata: # Some speculative decoding methods need to use different kv lengths for the # draft/target layers. But KVCacheManager can only support kv caches with the # same kv lengths for different layers. Add extra kv token in kv cache manager - # to haddle this issue. + # to handle this issue. num_extra_kv_tokens: Optional[int] = 0 # Number of layers in target model # The number of layers num_layers: int = 0 diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 93e9787d3f86..86e6901053a6 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -10,7 +10,7 @@ from ..pyexecutor.sampler import (SampleState, SampleStateTensors, TorchSampler, add_token, int_tensor) from ..pyexecutor.scheduler import ScheduledRequests -from .interface import SpecConfig, SpecMetadata, SpeculativeDecodingMode +from .interface import SpecMetadata @dataclass(kw_only=True) @@ -25,52 +25,10 @@ class SampleStateMTP(SampleState): host: SampleStateTensorsMTP -@dataclass -class MTPConfig(SpecConfig): - """ - Configuration for MTP. - """ - # The name of speculative decoding. - spec_dec_name = "MTP" - # The number of MTP modules - num_nextn_predict_layers: int = 1 - # The number of max batch size - max_batch_size: int = 8 - - # Whether to use relaxed acceptance during thinking phase for reasoning model - use_relaxed_acceptance_for_thinking: bool = False - # The top-N tokens are sampled from logits to obtain a candidate set. - relaxed_topk: int = 1 - # The threshold to further filter the candidate set. - # Filter out tokens with a large probability gap between the top-1 token's log probability. - relaxed_delta: float = 0. - - # Whether to use vanilla MTP - use_mtp_vanilla: bool = False - - # TODO: Hard code for DeepSeek R1 - # When encounter , start thinking phase. - # When encounter , end thinking phase. - # [thinking phase] [real output] - BEGIN_THINKING_PHASE_TOKEN: int = 128798 - END_THINKING_PHASE_TOKEN: int = 128799 - - def __post_init__(self) -> None: - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - self.max_draft_tokens = self.num_nextn_predict_layers - - def update_from_model_config(self, model_config): - assert self.num_nextn_predict_layers > 0 - if model_config.num_nextn_predict_layers == 1 and not self.use_mtp_vanilla: - self.spec_dec_mode = SpeculativeDecodingMode.MTP_EAGLE - self.num_extra_kv_tokens = self.num_nextn_predict_layers - 1 - - class MTPHiddenStatesManager(BaseResourceManager): - def __init__(self, config: MTPConfig, dtype: torch.dtype, hidden_size: int, - max_num_requests: int): + def __init__(self, config: "MTPDecodingConfig", dtype: torch.dtype, + hidden_size: int, max_num_requests: int): self.dtype = dtype self.num_nextn_predict_layers = config.num_nextn_predict_layers self.hidden_size = hidden_size @@ -353,7 +311,7 @@ def sample_async(self, scheduled_requests: ScheduledRequests, class MTPWorker(nn.Module): - def __init__(self, spec_config: MTPConfig): + def __init__(self, spec_config: "MTPDecodingConfig"): super().__init__() self.spec_config = spec_config self.is_thop = False @@ -1090,7 +1048,7 @@ def draft_sampler( class MTPEagleWorker(MTPWorker): - def __init__(self, spec_config: MTPConfig): + def __init__(self, spec_config: "MTPDecodingConfig"): super().__init__(spec_config) self.mtp_num_modules = spec_config.num_nextn_predict_layers diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py index dceb04c53d30..47f0315c3f33 100644 --- a/tensorrt_llm/_torch/speculative/ngram.py +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from itertools import chain from ordered_set import OrderedSet @@ -9,34 +8,6 @@ from ..pyexecutor.resource_manager import BaseResourceManager from ..pyexecutor.scheduler import ScheduledRequests from .drafter import Drafter -from .interface import SpecConfig, SpeculativeDecodingMode - - -@dataclass -class NGramConfig(SpecConfig): - """ - Configuration for NGram drafter. - """ - # The name of speculative decoding. - spec_dec_name = "NGRAM" - - num_extra_kv_tokens: int = 0 - max_draft_tokens: int = 0 - - prompt_lookup_num_tokens: int = 5 - max_matching_ngram_size: int = 5 - end_id: int = -1 - is_keep_all: bool = True - is_use_oldest: bool = True - is_public_pool: bool = True - - def __post_init__(self) -> None: - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - self.max_draft_tokens = self.prompt_lookup_num_tokens - - def update_from_model_config(self, model_config): - pass class NGramPoolManager(BaseResourceManager): @@ -76,7 +47,8 @@ class NGramPoolManager(BaseResourceManager): It maps from request ID to the index of the prompt to update the pool in the next step. """ - def __init__(self, spec_config: SpecConfig, max_num_requests: int): + def __init__(self, spec_config: "NGramDecodingConfig", + max_num_requests: int): self.prompt_lookup_num_tokens = spec_config.prompt_lookup_num_tokens self.max_matching_ngram_size = spec_config.max_matching_ngram_size self.is_keep_all = spec_config.is_keep_all @@ -191,12 +163,12 @@ class NGramDrafter(Drafter): def __init__( self, - spec_config: SpecConfig, + spec_config: "NGramDecodingConfig", ngram_pool_manager: NGramPoolManager = None, ): assert ngram_pool_manager is not None, "NGram needs a resource manager to maintain the pool." super().__init__(spec_resource_manager=ngram_pool_manager) - self.max_num_draft_tokens = spec_config.max_draft_tokens + self.max_draft_len = spec_config.max_draft_len def prepare_draft_tokens( self, @@ -220,8 +192,8 @@ def prepare_draft_tokens( request.py_end_id, request.py_orig_prompt_len + request.py_max_new_tokens, ) - # Pad length to `self.max_num_draft_tokens` + # Pad length to `self.max_draft_len` if len(draft_tokens) > 0: - pad_length = self.max_num_draft_tokens - len(draft_tokens) + pad_length = self.max_draft_len - len(draft_tokens) draft_tokens.extend([request.py_end_id] * pad_length) request.py_draft_tokens = draft_tokens diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 087e3578bfc5..78505aacb755 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,23 +1,27 @@ +import torch + from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler -from tensorrt_llm._torch.speculative.interface import SpecConfig, SpecMetadata +from tensorrt_llm._torch.speculative.interface import SpecMetadata from .draft_target import DraftTargetSpecMetadata from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata) +from .interface import SpeculativeDecodingMode from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker) from .ngram import NGramDrafter, NGramPoolManager def get_spec_metadata(spec_config, + model_config, max_num_requests, max_num_tokens, spec_resource_manager=None, is_draft_model=False): if spec_config.spec_dec_mode.is_mtp(): return MTPSpecMetadata( - max_draft_tokens=spec_config.max_draft_tokens, + max_draft_tokens=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, mtp_num_modules=spec_config.num_nextn_predict_layers, max_num_requests=max_num_requests, @@ -25,28 +29,28 @@ def get_spec_metadata(spec_config, ) if spec_config.spec_dec_mode.is_eagle3(): return Eagle3SpecMetadata( - max_draft_tokens=spec_config.max_draft_tokens, + max_draft_tokens=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, - num_layers=spec_config.num_layers, - hidden_size=spec_config.hidden_size, + num_layers=model_config.num_hidden_layers, + hidden_size=model_config.hidden_size, max_num_tokens=max_num_tokens, - dtype=spec_config.dtype, + dtype=model_config.torch_dtype, is_draft_model=is_draft_model, eagle3_resource_manager=spec_resource_manager, ) if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelSpecMetadata( - max_draft_tokens=spec_config.max_draft_tokens, + max_draft_tokens=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, - num_layers=spec_config.num_layers, - hidden_size=spec_config.hidden_size, + num_layers=model_config.num_hidden_layers, + hidden_size=model_config.hidden_size, max_num_tokens=max_num_tokens, ) if spec_config.spec_dec_mode.is_draft_target(): return DraftTargetSpecMetadata( - max_draft_tokens=spec_config.max_draft_tokens, + max_draft_tokens=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, ) @@ -104,7 +108,8 @@ def get_spec_resource_manager(model_engine, return None -def get_spec_decoder(sampler_args: TorchSampler.Args, spec_config: SpecConfig): +def get_spec_decoder(sampler_args: TorchSampler.Args, + spec_config: "DecodingBaseConfig"): if spec_config.spec_dec_mode.is_mtp(): return MTPSampler(sampler_args, nextn=spec_config.num_nextn_predict_layers) @@ -133,18 +138,49 @@ def get_spec_drafter(model_engine): def get_num_spec_layers(spec_config): if spec_config.spec_dec_mode.is_mtp(): return spec_config.num_nextn_predict_layers - elif spec_config.spec_dec_mode.is_eagle3_one_model(): + if spec_config.spec_dec_mode.is_eagle3_one_model(): return 1 - else: - return 0 + return 0 def get_spec_worker(spec_config, mapping): if spec_config.spec_dec_mode.is_mtp(): return MTPWorker(spec_config) - elif spec_config.spec_dec_mode.is_mtp_eagle(): + if spec_config.spec_dec_mode.is_mtp_eagle(): return MTPEagleWorker(spec_config) - elif spec_config.spec_dec_mode.is_eagle3_one_model(): + if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelWorker(spec_config, mapping) - else: - return None + return None + + +def get_draft_model_prompt(spec_dec_mode: SpeculativeDecodingMode, + input_tokens: torch.Tensor) -> torch.Tensor: + """ + Can be used to modify prompts for speculative algorithms that need to update tokens + before drafting. + """ + if spec_dec_mode.is_eagle3(): + # EAGLE3 always throws away the first token when processing draft inputs + return input_tokens[1:] + return input_tokens + + +def get_num_extra_kv_tokens(spec_config): + """ + Implementation detail for one model implementations of speculative decoding. Extra + KV cache tokens are required. + """ + if spec_config is None: + return 0 + if spec_config.spec_dec_mode.is_eagle3_one_model( + ) or spec_config.spec_dec_mode.is_mtp_eagle(): + return spec_config.max_draft_len - 1 + return 0 + + +def update_spec_config_from_model_config(spec_config, model_config): + if spec_config.spec_dec_mode.is_mtp(): + # Use `max_draft_len` for several low-level APIs. TODO: Remove this after distinguishing them, + spec_config.max_draft_len = spec_config.num_nextn_predict_layers + # Use `num_nextn_predict_layers_from_model_config` to decide decoding mode MTP / MTP_EAGLE. + spec_config.num_nextn_predict_layers_from_model_config = model_config.num_nextn_predict_layers diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6bb1a4a65a2c..1a111a15bd40 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1,4 +1,5 @@ import copy +import functools import json import math import os @@ -247,6 +248,28 @@ def from_dict(cls, data: dict): def _check_fields(self): pass + def supports_backend(self, backend: str) -> bool: + """ + Override if the speculation algorithm does not support + a subset of the possible backends. + """ + return True + + def validate(self) -> None: + """ + Do any additional error checking here. + """ + + @functools.cached_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 + # TRT backend. + from tensorrt_llm._torch.speculative.interface import \ + SpeculativeDecodingMode as TorchSpeculativeDecodingMode + return TorchSpeculativeDecodingMode.from_string( + self.decoding_type.upper()) + class MedusaDecodingConfig(DecodingBaseConfig): medusa_choices: Optional[List[List[int]]] = None @@ -258,6 +281,9 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "Medusa" + def supports_backend(self, backend: str) -> bool: + return backend not in ("pytorch", "_autodeploy") + class EagleDecodingConfig(DecodingBaseConfig): eagle_choices: Optional[List[List[int]]] = None @@ -267,7 +293,6 @@ class EagleDecodingConfig(DecodingBaseConfig): dynamic_tree_max_topK: Optional[int] = None num_eagle_layers: Optional[int] = None max_non_leaves_per_layer: Optional[int] = None - pytorch_weights_path: Optional[str] = None eagle3_one_model: Optional[bool] = True @classmethod @@ -276,6 +301,18 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "Eagle" + def validate(self) -> None: + if self.speculative_model is None: + raise ValueError("Draft model must be provided for EAGLE") + + @functools.cached_property + def spec_dec_mode(self): + from tensorrt_llm._torch.speculative.interface import \ + SpeculativeDecodingMode as TorchSpeculativeDecodingMode + if self.eagle3_one_model: + return TorchSpeculativeDecodingMode.EAGLE3_ONE_MODEL + return TorchSpeculativeDecodingMode.EAGLE3 + class UserProvidedDecodingConfig(DecodingBaseConfig): # Type should be Drafter, but it leads to circular import @@ -321,9 +358,11 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "NGram" + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + class DraftTargetDecodingConfig(DecodingBaseConfig): - pytorch_weights_path: Optional[str] = None @classmethod def from_dict(cls, data: dict): @@ -331,6 +370,9 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "DraftTarget" + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + class MTPDecodingConfig(DecodingBaseConfig): num_nextn_predict_layers: Optional[int] = 1 @@ -339,12 +381,34 @@ class MTPDecodingConfig(DecodingBaseConfig): relaxed_delta: Optional[float] = 0. use_mtp_vanilla: Optional[bool] = False + # TODO: remove this after distinguishing `max_draft_len` and `num_nextn_predict_layers` + # Now we need a flag when MTPDecodingConfig is updated by PyTorchModelEngine. + num_nextn_predict_layers_from_model_config: Optional[int] = 1 + + # TODO: Hard code for DeepSeek R1 + # When encounter , start thinking phase. + # When encounter , end thinking phase. + # [thinking phase] [real output] + BEGIN_THINKING_PHASE_TOKEN: int = 128798 + END_THINKING_PHASE_TOKEN: int = 128799 + @classmethod def from_dict(cls, data: dict): return cls(**data) decoding_type: ClassVar[str] = "MTP" + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + + @functools.cached_property + def spec_dec_mode(self): + from tensorrt_llm._torch.speculative.interface import \ + SpeculativeDecodingMode as TorchSpeculativeDecodingMode + if self.num_nextn_predict_layers_from_model_config == 1 and not self.use_mtp_vanilla: + return TorchSpeculativeDecodingMode.MTP_EAGLE + return TorchSpeculativeDecodingMode.MTP + class PybindMirror(ABC): ''' A class containing the utilities for mirroring Python classes to @@ -635,6 +699,9 @@ def _to_pybind(self): self.max_ngram_size, self.max_verification_set_size) + def supports_backend(self, backend: str) -> bool: + return backend not in ("pytorch", "_autodeploy") + decoding_type: ClassVar[str] = "Lookahead" @@ -1314,6 +1381,13 @@ def validate_build_config_remaining(self): @model_validator(mode="after") def validate_speculative_config(self): if self.speculative_config: + if not self.speculative_config.supports_backend(self.backend): + raise ValueError( + f"Speculation type {self.speculative_config.decoding_type} does not " + f"support backend {self.backend}") + + # Below, we only need to set speculative_decoding_mode/decoding_config for speculation + # on the TRT backend. if isinstance(self.speculative_config, LookaheadDecodingConfig): lookahead_config = self.speculative_config # Update the build config @@ -1327,6 +1401,7 @@ def validate_speculative_config(self): decoding_mode=DecodingMode.Lookahead(), lookahead_decoding_config=PybindMirror.maybe_to_pybind( lookahead_config)) + elif isinstance(self.speculative_config, MedusaDecodingConfig): self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.MEDUSA @@ -1335,6 +1410,16 @@ def validate_speculative_config(self): self.decoding_config = DecodingConfig( decoding_mode=DecodingMode.Medusa(), medusa_choices=self.speculative_config.medusa_choices) + + elif isinstance(self.speculative_config, DraftTargetDecodingConfig): + self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL + assert self.speculative_config.max_draft_len > 0 + self.build_config.max_draft_len = self.speculative_config.max_draft_len + + elif isinstance(self.speculative_config, MTPDecodingConfig): + assert self.speculative_config.num_nextn_predict_layers > 0 + self.build_config.max_draft_len = self.speculative_config.num_nextn_predict_layers + elif isinstance(self.speculative_config, EagleDecodingConfig): self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.EAGLE assert self.speculative_config.max_draft_len > 0 @@ -1351,62 +1436,21 @@ def validate_speculative_config(self): self.decoding_config = DecodingConfig( decoding_mode=DecodingMode.Eagle(), eagle_config=eagle_config) - else: - from tensorrt_llm._torch.speculative import Eagle3Config - self.speculative_config = Eagle3Config( - max_draft_tokens=self.speculative_config.max_draft_len, - draft_model_path=self.speculative_config. - pytorch_weights_path, - eagle3_one_model=self.speculative_config. - eagle3_one_model) + elif isinstance(self.speculative_config, NGramDecodingConfig): - assert self.backend in ['pytorch', '_autodeploy'] assert self.speculative_config.prompt_lookup_num_tokens > 0 and self.speculative_config.max_matching_ngram_size > 0 self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.NGRAM self.build_config.max_draft_len = self.speculative_config.max_draft_len - from tensorrt_llm._torch.speculative import NGramConfig - self.speculative_config = NGramConfig( - prompt_lookup_num_tokens=self.speculative_config. - prompt_lookup_num_tokens, - max_matching_ngram_size=self.speculative_config. - max_matching_ngram_size, - is_keep_all=self.speculative_config.is_keep_all, - is_use_oldest=self.speculative_config.is_use_oldest, - is_public_pool=self.speculative_config.is_public_pool, - ) - elif isinstance(self.speculative_config, DraftTargetDecodingConfig): - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL - assert self.backend == 'pytorch' - assert self.speculative_config.max_draft_len > 0 - self.build_config.max_draft_len = self.speculative_config.max_draft_len - from tensorrt_llm._torch.speculative import DraftTargetConfig - self.speculative_config = DraftTargetConfig( - max_draft_tokens=self.speculative_config.max_draft_len, - draft_model_path=self.speculative_config. - pytorch_weights_path) - elif isinstance(self.speculative_config, MTPDecodingConfig): - from tensorrt_llm._torch.speculative import MTPConfig - self.speculative_config = MTPConfig( - num_nextn_predict_layers=self.speculative_config. - num_nextn_predict_layers, - max_batch_size=self.build_config.max_batch_size, - use_relaxed_acceptance_for_thinking=self.speculative_config. - use_relaxed_acceptance_for_thinking, - relaxed_topk=self.speculative_config.relaxed_topk, - relaxed_delta=self.speculative_config.relaxed_delta, - use_mtp_vanilla=self.speculative_config.use_mtp_vanilla) + elif isinstance(self.speculative_config, UserProvidedDecodingConfig): - assert self.backend in ['pytorch', '_autodeploy'] - from tensorrt_llm._torch.speculative import UserProvidedConfig - self.speculative_config = UserProvidedConfig( - max_draft_tokens=self.speculative_config.max_draft_len, - drafter=self.speculative_config.drafter) + assert self.speculative_config.max_draft_len > 0 self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.USER_PROVIDED - self.build_config.max_draft_len = self.speculative_config.max_draft_tokens + self.build_config.max_draft_len = self.speculative_config.max_draft_len + else: raise ValueError( - f"Speculative config type not recognized: {self.speculative_config}" + f"Unrecognized speculative config type {type(self.speculative_config)}" ) else: self.decoding_config = None diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index 811bb80a109e..71057092f97d 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -25,7 +25,6 @@ import tensorrt_llm.evaluate from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm._torch.speculative import SpecConfig from tensorrt_llm.builder import BuildConfig from tensorrt_llm.llmapi import SamplingParams from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig @@ -156,8 +155,6 @@ def evaluate(self, spec_dec_algo = None elif isinstance(llm.args.speculative_config, DecodingBaseConfig): spec_dec_algo = llm.args.speculative_config.decoding_type - elif isinstance(llm.args.speculative_config, SpecConfig): - spec_dec_algo = llm.args.speculative_config.spec_dec_name else: raise ValueError( f"Not recognized speculative_config: {llm.args.speculative_config}." diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8afa75ed43ad..8111ffd26cb5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -243,7 +243,7 @@ def test_eagle3(self): draft_len = 4 spec_config = EagleDecodingConfig(max_draft_len=draft_len, - pytorch_weights_path=eagle_model_dir) + speculative_model=eagle_model_dir) llm = LLM(model=target_model_dir, **pytorch_config, @@ -262,6 +262,7 @@ def test_ngram(self): draft_len = 4 spec_config = NGramDecodingConfig( + max_draft_len=draft_len, prompt_lookup_num_tokens=draft_len, max_matching_ngram_size=draft_len, is_keep_all=True, diff --git a/tests/unittest/_torch/speculative/test_draft_target.py b/tests/unittest/_torch/speculative/test_draft_target.py index 88ff1ac92b92..75b93b3a29c0 100644 --- a/tests/unittest/_torch/speculative/test_draft_target.py +++ b/tests/unittest/_torch/speculative/test_draft_target.py @@ -35,8 +35,8 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str): draft_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" draft_len = 4 - spec_config = DraftTargetDecodingConfig( - max_draft_len=draft_len, pytorch_weights_path=draft_model_dir) + spec_config = DraftTargetDecodingConfig(max_draft_len=draft_len, + speculative_model=draft_model_dir) llm_spec = LLM(model=target_model_dir, max_batch_size=max_batch_size, disable_overlap_scheduler=True, diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 3630edee9e0e..07f1035531b1 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -49,7 +49,7 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, draft_len = 4 spec_config = EagleDecodingConfig( max_draft_len=draft_len, - pytorch_weights_path=eagle_model_dir, + speculative_model=eagle_model_dir, # Llama 3 does not support one model eagle. eagle3_one_model=use_one_model) diff --git a/tests/unittest/_torch/speculative/test_mtp.py b/tests/unittest/_torch/speculative/test_mtp.py index 61dcf2a3ce62..3d6438dfd638 100644 --- a/tests/unittest/_torch/speculative/test_mtp.py +++ b/tests/unittest/_torch/speculative/test_mtp.py @@ -6,9 +6,9 @@ import tensorrt_llm from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams -from tensorrt_llm._torch.speculative.mtp import (MTPConfig, - MTPHiddenStatesManager, +from tensorrt_llm._torch.speculative.mtp import (MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker) +from tensorrt_llm.llmapi import MTPDecodingConfig def unittest_name_func(testcase_func, param_num, param): @@ -297,7 +297,8 @@ def test_sample_and_accept_draft_tokens(self, test_case_name, ref_accepted_tokens, ref_num_accepted_tokens): batch_size = len(draft_len) - spec_config = MTPConfig(num_nextn_predict_layers=mtp_num_modules) + spec_config = MTPDecodingConfig( + num_nextn_predict_layers=mtp_num_modules) # attention metedata attn_metadata = TrtllmAttentionMetadata(max_num_requests=batch_size, @@ -871,7 +872,7 @@ def test_mtp_update_mtp_hidden_states( batch_size = len(request_ids) batch_size - num_context_request hidden_size = hidden_states.shape[1] - spec_config = MTPConfig( + spec_config = MTPDecodingConfig( num_nextn_predict_layers=num_nextn_predict_layers) attn_metadata = TrtllmAttentionMetadata(max_num_requests=batch_size, @@ -1362,7 +1363,7 @@ def test_prepare_drafter_inputs( hidden_size = previous_layer_hidden_states.shape[1] else: hidden_size = 10 - spec_config = MTPConfig( + spec_config = MTPDecodingConfig( num_nextn_predict_layers=num_nextn_predict_layers) if attn_metadata is None: diff --git a/tests/unittest/_torch/speculative/test_ngram.py b/tests/unittest/_torch/speculative/test_ngram.py index 918ead95ec9c..39e5dd058c85 100644 --- a/tests/unittest/_torch/speculative/test_ngram.py +++ b/tests/unittest/_torch/speculative/test_ngram.py @@ -27,19 +27,23 @@ def test_llama_ngram(disable_overlap_scheduler: bool, use_cuda_graph: bool, cuda_graph_config = CudaGraphConfig( batch_sizes=[1]) if use_cuda_graph else None + max_batch_size = 4 + max_draft_len = 4 + llm_common_config = dict( \ model=llm_models_root() / "llama-3.1-model" /"Meta-Llama-3.1-8B", backend='pytorch', attn_backend=attn_backend, disable_overlap_scheduler=disable_overlap_scheduler, cuda_graph_config=cuda_graph_config, - max_batch_size=4, + max_batch_size=max_batch_size, kv_cache_config=KvCacheConfig(enable_block_reuse=False), max_num_tokens=2048, ) spec_config = NGramDecodingConfig( - prompt_lookup_num_tokens=4, + max_draft_len=max_draft_len, + prompt_lookup_num_tokens=max_draft_len, max_matching_ngram_size=2, is_keep_all=True, is_use_oldest=True, diff --git a/tests/unittest/_torch/speculative/test_user_provided.py b/tests/unittest/_torch/speculative/test_user_provided.py index 89f256b15fd9..609bcd0f8fb6 100644 --- a/tests/unittest/_torch/speculative/test_user_provided.py +++ b/tests/unittest/_torch/speculative/test_user_provided.py @@ -6,9 +6,9 @@ import torch from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm._torch.speculative.ngram import (NGramConfig, NGramDrafter, - NGramPoolManager) +from tensorrt_llm._torch.speculative.ngram import NGramDrafter, NGramPoolManager from tensorrt_llm.llmapi import (CudaGraphConfig, KvCacheConfig, + NGramDecodingConfig, UserProvidedDecodingConfig) sys.path.append(os.path.join(os.path.dirname(__file__), '..')) @@ -29,6 +29,7 @@ def test_llama_user_provided(disable_overlap_scheduler: bool, batch_sizes=[1]) if use_cuda_graph else None max_batch_size = 4 + max_draft_len = 4 llm_common_config = dict( \ model=llm_models_root() / "llama-3.1-model" /"Meta-Llama-3.1-8B", @@ -41,23 +42,29 @@ def test_llama_user_provided(disable_overlap_scheduler: bool, max_num_tokens=2048, ) - draft_len = 4 - - ngram_config = NGramConfig( - prompt_lookup_num_tokens=draft_len, + ngram_config = NGramDecodingConfig( + max_draft_len=max_draft_len, + prompt_lookup_num_tokens=max_draft_len, max_matching_ngram_size=2, is_keep_all=True, is_use_oldest=True, is_public_pool=True, ) - drafter = NGramDrafter(spec_config=ngram_config, - ngram_pool_manager=NGramPoolManager( - spec_config=ngram_config, - max_num_requests=max_batch_size)) + ngram_pool_manager = NGramPoolManager( + spec_config=ngram_config, + max_num_requests=max_batch_size, + ) - spec_config = UserProvidedDecodingConfig(max_draft_len=draft_len, - drafter=drafter) + drafter = NGramDrafter( + spec_config=ngram_config, + ngram_pool_manager=ngram_pool_manager, + ) + + spec_config = UserProvidedDecodingConfig( + max_draft_len=draft_len, + drafter=drafter, + ) prompts = [ "The capital of France is", From 04d2a23a0e537c01ba326dd2022795dfc8d5d3b8 Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Tue, 8 Jul 2025 17:07:25 +0800 Subject: [PATCH 2/3] v0.2: rebase after UserProvided Signed-off-by: wili-65535 --- .../_torch/speculative/user_provided.py | 26 ------------------- tensorrt_llm/_torch/speculative/utils.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 2 +- .../_torch/speculative/test_user_provided.py | 2 +- 4 files changed, 3 insertions(+), 29 deletions(-) delete mode 100644 tensorrt_llm/_torch/speculative/user_provided.py diff --git a/tensorrt_llm/_torch/speculative/user_provided.py b/tensorrt_llm/_torch/speculative/user_provided.py deleted file mode 100644 index d8514cdab18d..000000000000 --- a/tensorrt_llm/_torch/speculative/user_provided.py +++ /dev/null @@ -1,26 +0,0 @@ -from dataclasses import dataclass -from typing import Optional - -from tensorrt_llm._torch.speculative.drafter import Drafter - -from .interface import SpecConfig, SpeculativeDecodingMode - - -@dataclass -class UserProvidedConfig(SpecConfig): - """ - Configuration for user provided speculative decoding. - """ - # The name of speculative decoding. - spec_dec_name = "USER_PROVIDED" - - num_extra_kv_tokens: int = 0 - max_draft_tokens: int = 0 - drafter: Optional[Drafter] = None - - def __post_init__(self) -> None: - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - - def update_from_model_config(self, model_config): - pass diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 78505aacb755..bfb719f19d35 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -57,7 +57,7 @@ def get_spec_metadata(spec_config, if spec_config.spec_dec_mode.is_ngram( ) or spec_config.spec_dec_mode.is_user_provided(): return SpecMetadata( - max_draft_tokens=spec_config.max_draft_tokens, + max_draft_tokens=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1a111a15bd40..1ba53c0b7760 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -322,7 +322,7 @@ class UserProvidedDecodingConfig(DecodingBaseConfig): def from_dict(cls, data: dict): return cls(**data) - decoding_type: ClassVar[str] = "UserProvided" + decoding_type: ClassVar[str] = "User_Provided" class NGramDecodingConfig(DecodingBaseConfig): diff --git a/tests/unittest/_torch/speculative/test_user_provided.py b/tests/unittest/_torch/speculative/test_user_provided.py index 609bcd0f8fb6..178165400d04 100644 --- a/tests/unittest/_torch/speculative/test_user_provided.py +++ b/tests/unittest/_torch/speculative/test_user_provided.py @@ -62,7 +62,7 @@ def test_llama_user_provided(disable_overlap_scheduler: bool, ) spec_config = UserProvidedDecodingConfig( - max_draft_len=draft_len, + max_draft_len=max_draft_len, drafter=drafter, ) From 6064bb416bd196a93528177e3ccc37f7f4590987 Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Thu, 10 Jul 2025 11:55:04 +0800 Subject: [PATCH 3/3] v0.3: safe change v1 (no flaky error, no tool function) Signed-off-by: wili-65535 --- .../blog6_Llama4_maverick_eagle_guide.md | 2 +- examples/eagle/README.md | 2 +- .../_tensorrt_engine/llm_eagle2_decoding.py | 4 +- .../_tensorrt_engine/llm_eagle_decoding.py | 4 +- .../_tensorrt_engine/llm_medusa_decoding.py | 4 +- examples/llm-api/llm_speculative_decoding.py | 4 +- examples/llm-api/quickstart_advanced.py | 5 +- .../_torch/auto_deploy/shim/ad_executor.py | 8 +- .../_torch/models/modeling_speculative.py | 2 +- tensorrt_llm/_torch/pyexecutor/_util.py | 12 +-- .../_torch/pyexecutor/cuda_graph_runner.py | 2 +- .../_torch/pyexecutor/model_engine.py | 18 ++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 18 ++-- .../_torch/pyexecutor/py_executor_creator.py | 7 +- .../_torch/pyexecutor/resource_manager.py | 4 +- tensorrt_llm/_torch/pyexecutor/sampler.py | 4 +- tensorrt_llm/_torch/speculative/__init__.py | 11 +-- .../_torch/speculative/draft_target.py | 13 --- tensorrt_llm/_torch/speculative/eagle3.py | 29 ++++--- tensorrt_llm/_torch/speculative/interface.py | 17 +++- tensorrt_llm/_torch/speculative/mtp.py | 10 +-- tensorrt_llm/_torch/speculative/ngram.py | 8 +- tensorrt_llm/_torch/speculative/utils.py | 56 ++----------- tensorrt_llm/llmapi/llm_args.py | 82 ++++++++++++------- tensorrt_llm/llmapi/llm_utils.py | 10 +-- tensorrt_llm/models/eagle/config.py | 2 +- tensorrt_llm/models/eagle/model.py | 6 +- tensorrt_llm/models/medusa/config.py | 2 +- tensorrt_llm/models/medusa/model.py | 2 +- tensorrt_llm/models/redrafter/model.py | 6 +- .../defs/accuracy/accuracy_core.py | 5 ++ .../accuracy/test_disaggregated_serving.py | 2 +- .../integration/defs/accuracy/test_llm_api.py | 4 +- .../defs/accuracy/test_llm_api_pytorch.py | 3 +- .../test_configs/disagg_config_ngram.yaml | 2 +- .../_torch/speculative/test_draft_target.py | 56 ++++++------- .../_torch/speculative/test_eagle3.py | 77 ++++++++--------- tests/unittest/_torch/speculative/test_mtp.py | 37 ++++----- .../unittest/_torch/speculative/test_ngram.py | 10 +-- .../_torch/speculative/test_user_provided.py | 10 +-- tests/unittest/llmapi/test_llm.py | 8 +- 41 files changed, 260 insertions(+), 308 deletions(-) delete mode 100644 tensorrt_llm/_torch/speculative/draft_target.py diff --git a/docs/source/blogs/tech_blog/blog6_Llama4_maverick_eagle_guide.md b/docs/source/blogs/tech_blog/blog6_Llama4_maverick_eagle_guide.md index 93727adc5182..888898664703 100644 --- a/docs/source/blogs/tech_blog/blog6_Llama4_maverick_eagle_guide.md +++ b/docs/source/blogs/tech_blog/blog6_Llama4_maverick_eagle_guide.md @@ -68,7 +68,7 @@ docker run -d --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \ -p 8000:8000 --gpus=all -e "TRTLLM_ENABLE_PDL=1" \ -v /path/to/maverick:/config/models/maverick -v /path/to/eagle:/config/models/eagle \ docker.io//tensorrt_llm:main sh \ - -c "echo -e 'enable_attention_dp: false\nenable_min_latency: true\nenable_autotuner: false\ncuda_graph_config:\n max_batch_size: 8\nspeculative_config:\n decoding_type: Eagle\n max_draft_len: 3\n pytorch_weights_path: /config/models/eagle\nkv_cache_config:\n enable_block_reuse: false' > c.yaml && \ + -c "echo -e 'enable_attention_dp: false\nenable_min_latency: true\nenable_autotuner: false\ncuda_graph_config:\n max_batch_size: 8\nspeculative_config:\n decoding_type: Eagle\n max_draft_len: 3\n speculative_model_dir: /config/models/eagle\nkv_cache_config:\n enable_block_reuse: false' > c.yaml && \ TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL=True \ trtllm-serve /config/models/maverick \ --host 0.0.0.0 --port 8000 \ diff --git a/examples/eagle/README.md b/examples/eagle/README.md index 79756f045c5f..637223afb912 100644 --- a/examples/eagle/README.md +++ b/examples/eagle/README.md @@ -108,7 +108,7 @@ When using EAGLE-2, please enable `--eagle_use_dynamic_tree`, which indicates wh - In EagleNet2, the `N` output nodes of EagleNet1 are expanded, and each node expands `N` new draft tokens. Therefore, this layer also has a total of `N * N` draft tokens. And select the top `N` as the output of this layer. - Etc. -Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_tokens)` nodes. +Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_len)` nodes. diff --git a/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py b/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py index 86b5ca28af4a..a1343cc5757b 100755 --- a/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py +++ b/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py @@ -23,12 +23,12 @@ def main(): model = "lmsys/vicuna-7b-v1.3" # The end user can customize the eagle decoding configuration by specifying the - # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices + # speculative_model_dir, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK # with the EagleDecodingConfig class speculative_config = EagleDecodingConfig( - speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3", + speculative_model_dir="yuhuili/EAGLE-Vicuna-7B-v1.3", max_draft_len=63, num_eagle_layers=4, max_non_leaves_per_layer=10, diff --git a/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py b/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py index e6e89a622eeb..c66e15f66468 100644 --- a/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py +++ b/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py @@ -23,12 +23,12 @@ def main(): model = "lmsys/vicuna-7b-v1.3" # The end user can customize the eagle decoding configuration by specifying the - # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices + # speculative_model_dir, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK # with the EagleDecodingConfig class speculative_config = EagleDecodingConfig( - speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3", + speculative_model_dir="yuhuili/EAGLE-Vicuna-7B-v1.3", max_draft_len=63, num_eagle_layers=4, max_non_leaves_per_layer=10, diff --git a/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py b/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py index c683e04dbf94..b6d7f90c0f5f 100644 --- a/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py +++ b/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py @@ -48,10 +48,10 @@ def run_medusa_decoding(use_modelopt_ckpt=False, model_dir=None): model = "lmsys/vicuna-7b-v1.3" # The end user can customize the medusa decoding configuration by specifying the - # speculative_model, max_draft_len, medusa heads num and medusa choices + # speculative_model_dir, max_draft_len, medusa heads num and medusa choices # with the MedusaDecodingConfig class speculative_config = MedusaDecodingConfig( - speculative_model="FasterDecoding/medusa-vicuna-7b-v1.3", + speculative_model_dir="FasterDecoding/medusa-vicuna-7b-v1.3", max_draft_len=63, num_medusa_heads=4, medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ diff --git a/examples/llm-api/llm_speculative_decoding.py b/examples/llm-api/llm_speculative_decoding.py index bcba5d3c8c0b..8048bede8ee3 100644 --- a/examples/llm-api/llm_speculative_decoding.py +++ b/examples/llm-api/llm_speculative_decoding.py @@ -35,7 +35,7 @@ def run_MTP(model: Optional[str] = None): def run_Eagle3(): spec_config = EagleDecodingConfig( max_draft_len=3, - pytorch_weights_path="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", + speculative_model_dir="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", eagle3_one_model=True) llm = LLM( @@ -50,7 +50,7 @@ def run_Eagle3(): def run_ngram(): spec_config = NGramDecodingConfig( - prompt_lookup_num_tokens=3, + max_draft_len=3, max_matching_ngram_size=3, is_keep_all=True, is_use_oldest=True, diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 22a79a7eb6e6..352a23893ca0 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -169,16 +169,15 @@ def setup_llm(args): elif spec_decode_algo == "EAGLE3": spec_config = EagleDecodingConfig( max_draft_len=args.spec_decode_nextn, - speculative_model=args.draft_model_dir, + speculative_model_dir=args.draft_model_dir, eagle3_one_model=args.use_one_model) elif spec_decode_algo == "DRAFT_TARGET": spec_config = DraftTargetDecodingConfig( max_draft_len=args.spec_decode_nextn, - speculative_model=args.draft_model_dir) + speculative_model_dir=args.draft_model_dir) elif spec_decode_algo == "NGRAM": spec_config = NGramDecodingConfig( max_draft_len=args.spec_decode_nextn, - prompt_lookup_num_tokens=args.spec_decode_nextn, max_matching_ngram_size=args.max_matching_ngram_size, is_keep_all=True, is_use_oldest=True, diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 6aca50ba844a..c1a0fb151d47 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -261,8 +261,8 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir: max_num_sequences = ad_config.max_batch_size * dist_mapping.pp_size # some derivative properties - max_draft_tokens = ( - 0 if ad_config.speculative_config is None else ad_config.speculative_config.max_draft_tokens + max_draft_len = ( + 0 if ad_config.speculative_config is None else ad_config.speculative_config.max_draft_len ) # initialize model engine @@ -299,7 +299,7 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir: # correctly for models as needed. sampler_args = TorchSampler.Args( max_seq_len=ad_config.max_seq_len, - max_draft_tokens=max_draft_tokens, + max_draft_len=max_draft_len, max_num_sequences=max_num_sequences, max_beam_width=executor_config.max_beam_width, enable_mixed_sampler=ad_config.enable_mixed_sampler, @@ -317,7 +317,7 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir: disable_overlap_scheduler=ad_config.disable_overlap_scheduler, max_input_len=ad_config.max_input_len, max_batch_size=ad_config.max_batch_size, - max_draft_tokens=max_draft_tokens, + max_draft_len=max_draft_len, max_beam_width=ad_config.max_beam_width, ) return py_executor diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 11f360965e26..c6f05219775a 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -340,7 +340,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): model_config, 'spec_config', None ) and model_config.spec_config.spec_dec_mode.use_one_engine(): draft_config = ModelConfig.from_pretrained( - model_config.spec_config.speculative_model, + model_config.spec_config.speculative_model_dir, trust_remote_code=True, attn_backend=model_config.attn_backend, moe_backend=model_config.moe_backend, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 39d253a96d86..6e969a8d1de7 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -17,7 +17,7 @@ from tensorrt_llm.mapping import Mapping from ..model_config import ModelConfig -from ..speculative import get_num_extra_kv_tokens, get_spec_decoder +from ..speculative import get_spec_decoder from .config_utils import is_mla, is_nemotron_hybrid from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver from .llm_request import ExecutorResponse @@ -161,7 +161,7 @@ def _get_token_num_for_estimation(self) -> int: if spec_cfg is not None: num_extra_tokens_per_seq += spec_cfg.max_draft_len - num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg) + num_extra_tokens_per_seq += spec_cfg.num_extra_kv_tokens for req in self._dummy_reqs: num_req_tokens = len(req.input_token_ids) + num_extra_tokens_per_seq # Requests cannot share KV cache blocks. Round up to nearest integer multiple of block size. @@ -538,7 +538,7 @@ def create_py_executor_instance( disable_overlap_scheduler, max_batch_size=executor_config.max_batch_size, max_beam_width=executor_config.max_beam_width, - max_draft_tokens=spec_config.max_draft_len + max_draft_len=spec_config.max_draft_len if spec_config is not None else 0, kv_cache_transceiver=kv_cache_transceiver, draft_model_engine=draft_model_engine, @@ -549,11 +549,11 @@ def create_py_executor_instance( def create_torch_sampler_args(executor_config: ExecutorConfig, mapping: Mapping, *, max_seq_len: int, enable_mixed_sampler: bool): max_num_sequences = executor_config.max_batch_size * mapping.pp_size - max_draft_tokens = (0 if executor_config.speculative_config is None else - executor_config.speculative_config.max_draft_len) + max_draft_len = (0 if executor_config.speculative_config is None else + executor_config.speculative_config.max_draft_len) return TorchSampler.Args( max_seq_len=max_seq_len, - max_draft_tokens=max_draft_tokens, + max_draft_len=max_draft_len, max_num_sequences=max_num_sequences, max_beam_width=executor_config.max_beam_width, enable_mixed_sampler=enable_mixed_sampler, diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index f63fad0d9bb4..241fc1447cee 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -53,7 +53,7 @@ def __init__( # [CUDA graph spec decode padding] # We pad input IDs/position IDs to the maximum draft length (token per request). # We're forced to do this because we cannot reallocate inputs over many graph runs. - token_per_request = spec_metadata.max_draft_tokens + 1 if spec_metadata is not None else 1 + token_per_request = spec_metadata.max_draft_len + 1 if spec_metadata is not None else 1 # Using ones instead of zeros prevents NaNs in e.g. Deepseek self.input_ids = torch.ones((batch_size * token_per_request, ), diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 07f0faf1853e..7f99bab3e7fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -21,8 +21,6 @@ import tensorrt_llm.bindings.internal.userbuffers as ub from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors -from tensorrt_llm._torch.speculative import ( - get_num_extra_kv_tokens, update_spec_config_from_model_config) from tensorrt_llm._torch.speculative.mtp import SampleStateTensorsMTP from tensorrt_llm._utils import (is_trace_enabled, local_mpi_rank, local_mpi_size, nvtx_range, release_gc, @@ -457,8 +455,7 @@ def __init__( if self.is_spec_decode: self.spec_metadata = None - update_spec_config_from_model_config(self.spec_config, - self.model.config) + self.spec_config.update_from_model_config(self.model.config) max_num_draft_tokens = self.spec_config.max_draft_len * batch_size self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, @@ -1083,7 +1080,8 @@ def init_meta_tensor(t: torch.Tensor): if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): - weights = load_weights(self.spec_config.speculative_model) + weights = load_weights( + self.spec_config.speculative_model_dir) model.load_draft_weights(weights) elif load_format == LoadFormat.DUMMY: @@ -1266,9 +1264,8 @@ def _prepare_tp_inputs( extend_requests += extend_dummy_requests if not self._disable_overlap_scheduler and self.is_spec_decode: - spec_dec_mode = self.spec_config.spec_dec_mode - assert spec_dec_mode.support_overlap_scheduler( - ), f"{self.spec_config.spec_dec_mode} does not support overlap scheduler" + assert self.spec_config.spec_dec_mode.support_overlap_scheduler( + ), f"{self.spec_config.decoding_type} does not support overlap scheduler" # will contain previous batch indices of generation requests previous_batch_indices = [] @@ -1496,7 +1493,8 @@ def previous_seq_slots_device(): attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq, - num_extra_kv_tokens=get_num_extra_kv_tokens(self.spec_config)) + num_extra_kv_tokens=0 if self.spec_config is None else + self.spec_config.num_extra_kv_tokens) attn_metadata.kv_cache_manager = kv_cache_manager attn_metadata.prepare() @@ -2082,7 +2080,7 @@ def forward( spec_metadata.spec_dec_mode.attention_need_spec_dec_mode(), spec_metadata.is_spec_dec_tree, spec_metadata.is_spec_dec_dynamic_tree, - spec_metadata.max_draft_tokens) + spec_metadata.max_draft_len) else: spec_metadata = None diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ee338971216b..c42b999eced7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -17,7 +17,6 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager -from tensorrt_llm._torch.speculative import get_draft_model_prompt from tensorrt_llm._utils import (customized_gc_thresholds, global_mpi_rank, is_trace_enabled, nvtx_range, trace_func) from tensorrt_llm.bindings.executor import (DisServingRequestStats, @@ -174,7 +173,7 @@ def __init__(self, max_input_len: int = 2048, max_batch_size: int = 8, max_beam_width: int = 1, - max_draft_tokens: int = 0, + max_draft_len: int = 0, kv_cache_transceiver: Optional[KvCacheTransceiver] = None, draft_model_engine: Optional[ModelEngine] = None, garbage_collection_gen0_threshold: Optional[int] = None, @@ -208,7 +207,7 @@ def __init__(self, self.active = True self.next_req_id = max_batch_size # The first max_batch_size request IDs are reserved for dummy requests self.max_beam_width = max_beam_width - self.max_draft_tokens = max_draft_tokens + self.max_draft_len = max_draft_len self.print_log = model_engine.pytorch_backend_config.print_iter_log self.enable_iter_perf_stats = model_engine.pytorch_backend_config.enable_iter_perf_stats self.enable_iter_req_stats = model_engine.pytorch_backend_config.enable_iter_req_stats @@ -1524,7 +1523,7 @@ def _pad_attention_dp_dummy_request(self): request_ids=[0], is_gen=not self.has_context_request, prepare_resource=not self.has_context_request, - max_num_draft_tokens=self.max_draft_tokens, + max_num_draft_tokens=self.max_draft_len, )[0] llm_request.is_attention_dp_dummy = True spec_resource_manager = self.resource_manager.get_resource_manager( @@ -1768,10 +1767,9 @@ def _prepare_draft_batch( num_rejected_tokens = num_draft_tokens - num_accepted_tokens assert num_rejected_tokens >= 0 - spec_mode = self.model_engine.spec_config.spec_dec_mode + spec_config = self.model_engine.spec_config beam_idx = 0 - input_tokens = get_draft_model_prompt( - spec_mode, + input_tokens = spec_config.get_draft_model_prompt( request.get_tokens()[beam_idx]) def create_new_request(input_tokens): @@ -1873,15 +1871,15 @@ def _process_decoded_tokens(draft_batch): # this? Just needs proper kernel support. def _pad_to_max_draft_tokens(): for req in scheduled_requests.generation_requests: - max_draft_tokens = self.max_draft_tokens + max_draft_len = self.max_draft_len num_draft_tokens = len(req.py_draft_tokens) req.py_draft_tokens.extend( - 0 for _ in range(max_draft_tokens - num_draft_tokens)) + 0 for _ in range(max_draft_len - num_draft_tokens)) draft_batch.generation_requests = draft_batch.context_requests + draft_batch.generation_requests draft_batch.context_requests = [] - for i in range(self.max_draft_tokens - 1): + for i in range(self.max_draft_len - 1): if len(draft_batch.generation_requests) == 0: break diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index c24e6bc7a379..a72f6a58b127 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -19,8 +19,7 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist -from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, - get_spec_resource_manager) +from ..speculative import get_spec_drafter, get_spec_resource_manager from ._util import (KvCacheCreator, create_py_executor_instance, instantiate_sampler, is_mla) from .config import PyTorchConfig @@ -251,7 +250,7 @@ def create_py_executor( draft_spec_config.max_draft_len = 0 draft_model_engine = PyTorchModelEngine( - model_path=spec_config.speculative_model, + model_path=spec_config.speculative_model_dir, pytorch_backend_config=pytorch_backend_config, batch_size=executor_config.max_batch_size, max_beam_width=executor_config.max_beam_width, @@ -280,7 +279,7 @@ def create_py_executor( max_seq_len += spec_config.max_draft_len if spec_config is not None: - max_seq_len += get_num_extra_kv_tokens(spec_config) + max_seq_len += spec_config.num_extra_kv_tokens max_seq_len += spec_config.max_draft_len executor_config.max_seq_len = max_seq_len diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 68a82c1ca3a7..e52096727c6b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -175,9 +175,7 @@ def __init__( self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 # Some speculative decoding methods need to use different kv lengths for the # draft/target layers. Add extra tokens to handle this issue. - # Import here to avoid circular imports - from ..speculative import get_num_extra_kv_tokens - self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + self.num_extra_kv_tokens = 0 if spec_config is None else spec_config.num_extra_kv_tokens self.event_buffer_max_size = kv_cache_config.event_buffer_max_size self.max_num_tokens = max_num_tokens diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 1dcf3185c137..8d2bacc7b86d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -220,7 +220,7 @@ def create_store(self) -> Store: @dataclass(frozen=True, kw_only=True) class Args: max_seq_len: int - max_draft_tokens: int + max_draft_len: int max_num_sequences: int max_beam_width: int enable_mixed_sampler: bool @@ -228,7 +228,7 @@ class Args: def __init__(self, args: Args): self.max_seq_len = args.max_seq_len self.enable_mixed_sampler = args.enable_mixed_sampler - self.max_tokens = args.max_draft_tokens + 1 + self.max_tokens = args.max_draft_len + 1 assert args.max_beam_width == self.MAX_BEAM_WIDTH, "TorchSampler only supports beam_width = 1" self.num_seq_slots = args.max_num_sequences diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 5feb0760feb5..4667398e5908 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -1,11 +1,10 @@ from .eagle3 import Eagle3SpecMetadata -from .interface import SpecMetadata +from .interface import SpecConfig, SpecMetadata from .mtp import MTPEagleWorker, MTPSpecMetadata, MTPWorker from .ngram import NGramDrafter, NGramPoolManager -from .utils import (get_draft_model_prompt, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, get_spec_drafter, +from .utils import (get_num_spec_layers, get_spec_decoder, get_spec_drafter, get_spec_metadata, get_spec_resource_manager, - get_spec_worker, update_spec_config_from_model_config) + get_spec_worker) __all__ = [ "Eagle3SpecMetadata", @@ -14,14 +13,12 @@ "MTPWorker", "NGramDrafter", "NGramPoolManager", + "SpecConfig", "SpecMetadata", - "get_draft_model_prompt", - "get_num_extra_kv_tokens", "get_num_spec_layers", "get_spec_decoder", "get_spec_drafter", "get_spec_metadata", "get_spec_resource_manager", "get_spec_worker", - "update_spec_config_from_model_config", ] diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py deleted file mode 100644 index 3fdf4865c9b6..000000000000 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ /dev/null @@ -1,13 +0,0 @@ -from dataclasses import dataclass - -from .interface import SpecMetadata - - -@dataclass -class DraftTargetSpecMetadata(SpecMetadata): - - def __post_init__(self): - pass - - def prepare(self): - pass diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 22f3f645e43d..e6059abfc51f 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -26,7 +26,7 @@ def __init__(self, config: "EagleDecodingConfig", dtype: torch.dtype, hidden_size: int, max_num_requests: int, max_seq_len: int, max_num_tokens: int): self.dtype = dtype - self.max_draft_tokens = config.max_draft_len + self.max_draft_len = config.max_draft_len self.hidden_size = hidden_size self.max_num_requests = max_num_requests self.max_seq_len = max_seq_len @@ -235,7 +235,7 @@ def prepare(self): pin_memory=True) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - self.num_tokens -= (self.num_generations) * self.max_draft_tokens + self.num_tokens -= (self.num_generations) * self.max_draft_len def maybe_capture_hidden_states( self, @@ -255,7 +255,7 @@ def maybe_capture_hidden_states( class Eagle3OneModelSampler(MTPSampler): def __init__(self, args: TorchSampler.Args): - super().__init__(args, nextn=args.max_draft_tokens) + super().__init__(args, nextn=args.max_draft_len) class Eagle3OneModelWorker(nn.Module): @@ -263,7 +263,7 @@ class Eagle3OneModelWorker(nn.Module): def __init__(self, spec_config: "EagleDecodingConfig", mapping: Mapping): super().__init__() self.spec_config = spec_config - self.max_draft_tokens = self.spec_config.max_draft_len + self.max_draft_len = self.spec_config.max_draft_len self.mapping = mapping @torch.compile(mode="max-autotune-no-cudagraphs") @@ -300,7 +300,7 @@ def forward(self, input_ids, position_ids, hidden_states, logits, # Predict draft tokens next_draft_tokens = [] - for i in range(self.max_draft_tokens): + for i in range(self.max_draft_len): hidden_states, hidden_states_to_save = draft_model.model(**inputs) # FIXME (jhaotingc): Currently we disable use_spec_decoding mode for Eagle engine nth steps except 1st step. @@ -311,7 +311,7 @@ def forward(self, input_ids, position_ids, hidden_states, logits, attn_metadata.use_spec_decoding = False if i == 0: start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] * - (self.max_draft_tokens + 1)).long() + (self.max_draft_len + 1)).long() gather_ids_gen = (start_ids_gen + num_accepted_tokens[num_contexts:] - 1 + attn_metadata.num_ctx_tokens) @@ -341,8 +341,7 @@ def forward(self, input_ids, position_ids, hidden_states, logits, # update kv_lens_cuda if hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( - self.max_draft_tokens - - num_accepted_tokens[num_contexts:]) + 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 @@ -395,7 +394,7 @@ def sample_and_accept_draft_tokens( logits = logits.unsqueeze(0) # The return buffer - accepted_tokens = torch.empty((batch_size, (self.max_draft_tokens + 1)), + accepted_tokens = torch.empty((batch_size, (self.max_draft_len + 1)), dtype=torch.int, device=logits.device) num_accepted_tokens = torch.ones(batch_size, @@ -409,13 +408,13 @@ def sample_and_accept_draft_tokens( # generation gen_target_tokens = target_tokens[num_contexts:].reshape( - num_gens, self.max_draft_tokens + 1) + num_gens, self.max_draft_len + 1) accepted_tokens[num_contexts:, :] = gen_target_tokens draft_tokens = spec_metadata.draft_tokens.reshape( - num_gens, self.max_draft_tokens) - num_accepted_tokens[num_contexts:] += torch.cumprod(( - draft_tokens == gen_target_tokens[:, :self.max_draft_tokens]).int(), - dim=-1).sum(1) + num_gens, self.max_draft_len) + num_accepted_tokens[num_contexts:] += torch.cumprod( + (draft_tokens == gen_target_tokens[:, :self.max_draft_len]).int(), + dim=-1).sum(1) return accepted_tokens, num_accepted_tokens def draft_decoder( @@ -435,7 +434,7 @@ def draft_decoder( Returns: draft_tokens: torch.Tensor - [batch_size * max_draft_tokens] + [batch_size * max_draft_len] Draft token ids. Flattened. ''' diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index b7ceed8cebec..e374592a600e 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -15,7 +15,7 @@ class SpeculativeDecodingMode(IntEnum): EAGLE3 = auto() EAGLE3_ONE_MODEL = auto() NGRAM = auto() - DRAFTTARGET = auto() + DRAFT_TARGET = auto() USER_PROVIDED = auto() NONE = auto() @@ -44,7 +44,7 @@ def is_none(self): return self == SpeculativeDecodingMode.NONE def is_draft_target(self): - return self == SpeculativeDecodingMode.DRAFTTARGET + return self == SpeculativeDecodingMode.DRAFT_TARGET def without_logits(self): return self.is_mtp() or self.is_eagle3_one_model() @@ -105,6 +105,17 @@ def from_string(name: Optional[str]) -> "SpeculativeDecodingMode": return SpeculativeDecodingMode[name.upper()] +@dataclass +class SpecConfig: + """ + Configuration for speculative decoding. + This class is deprecated, but thread-leak of pytest raises flaky error if removing it. + TODO: remove this class safely. + """ + # The name of speculative decoding. + spec_dec_name = None + + @dataclass class SpecMetadata: """ @@ -113,7 +124,7 @@ class SpecMetadata: # The max number of requests in a single batch. max_num_requests: int # The max number of draft tokens. - max_draft_tokens: int + max_draft_len: int # The number of gen-phase sequences in the batch. num_generations: int = 0 # Whether CUDA graph is enabled. diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 86e6901053a6..72316a2e474a 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -157,8 +157,8 @@ def prepare(self): pin_memory=True) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - # MTP vanilla worker uses total max_draft_tokens input tokens in generation phase, - # while MTP Eagle worker uses (max_draft_tokens + 1) input tokens in the 1st draft + # MTP vanilla worker uses total max_draft_len input tokens in generation phase, + # while MTP Eagle worker uses (max_draft_len + 1) input tokens in the 1st draft # forward and only one input token in the following draft forward. # This num_tokens is used to set the all_rank_num_tokens for attention dp. if not self.spec_dec_mode.is_mtp_eagle(): @@ -700,7 +700,7 @@ def sample_and_accept_draft_tokens( Returns: accepted_tokens: torch.Tensor - [batch_size, (max_draft_tokens + 1)] + [batch_size, (max_draft_len + 1)] Accepted token ids. Flattened. num_accepted_tokens: torch.Tensor @@ -900,7 +900,7 @@ def prepare_drafter_inputs( Target model's hidden states. accepted_tokens: torch.Tensor - [batch_size, max_draft_tokens + 1] + [batch_size, max_draft_len + 1] Accepted token ids. Flattened. num_accepted_tokens: torch.Tensor @@ -1038,7 +1038,7 @@ def draft_sampler( Returns: draft_tokens: torch.Tensor - [batch_size * max_draft_tokens] + [batch_size * max_draft_len] Draft token ids. Flattened. ''' diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py index 47f0315c3f33..1d015a58b9b9 100644 --- a/tensorrt_llm/_torch/speculative/ngram.py +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -23,7 +23,7 @@ class NGramPoolManager(BaseResourceManager): `matches` is a list of candidate draft token ids attaching to a pattern. Arguments: - prompt_lookup_num_tokens: int + max_draft_len: int The length maximum of draft tokens (can be understood as length maximum of output draft tokens). max_matching_ngram_size: int @@ -49,7 +49,7 @@ class NGramPoolManager(BaseResourceManager): def __init__(self, spec_config: "NGramDecodingConfig", max_num_requests: int): - self.prompt_lookup_num_tokens = spec_config.prompt_lookup_num_tokens + self.max_draft_len = spec_config.max_draft_len self.max_matching_ngram_size = spec_config.max_matching_ngram_size self.is_keep_all = spec_config.is_keep_all self.is_use_oldest = spec_config.is_use_oldest # TODO: remove this if updating strategy is supported @@ -105,7 +105,7 @@ def get_draft_tokens( -1): # Find each possible pattern-match combination, and use tuple for hash for l in range(len(sequence) - size): - r = min(l + size + self.prompt_lookup_num_tokens, len(sequence)) + r = min(l + size + self.max_draft_len, len(sequence)) pattern = tuple(sequence[l:l + size]) new_match = tuple(sequence[l + size:r]) if pattern not in pool or \ @@ -137,7 +137,7 @@ def get_draft_tokens( # Update start_index self.start_index[request_id] = max( 0, prefix_len - - (self.prompt_lookup_num_tokens + self.max_matching_ngram_size - 1)) + (self.max_draft_len + self.max_matching_ngram_size - 1)) return draft_tokens diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index bfb719f19d35..882dfdf924d7 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,13 +1,9 @@ -import torch - from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler from tensorrt_llm._torch.speculative.interface import SpecMetadata -from .draft_target import DraftTargetSpecMetadata from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata) -from .interface import SpeculativeDecodingMode from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker) from .ngram import NGramDrafter, NGramPoolManager @@ -21,7 +17,7 @@ def get_spec_metadata(spec_config, is_draft_model=False): if spec_config.spec_dec_mode.is_mtp(): return MTPSpecMetadata( - max_draft_tokens=spec_config.max_draft_len, + max_draft_len=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, mtp_num_modules=spec_config.num_nextn_predict_layers, max_num_requests=max_num_requests, @@ -29,7 +25,7 @@ def get_spec_metadata(spec_config, ) if spec_config.spec_dec_mode.is_eagle3(): return Eagle3SpecMetadata( - max_draft_tokens=spec_config.max_draft_len, + max_draft_len=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, num_layers=model_config.num_hidden_layers, @@ -41,23 +37,18 @@ def get_spec_metadata(spec_config, ) if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelSpecMetadata( - max_draft_tokens=spec_config.max_draft_len, + max_draft_len=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, num_layers=model_config.num_hidden_layers, hidden_size=model_config.hidden_size, max_num_tokens=max_num_tokens, ) - if spec_config.spec_dec_mode.is_draft_target(): - return DraftTargetSpecMetadata( - max_draft_tokens=spec_config.max_draft_len, - spec_dec_mode=spec_config.spec_dec_mode, - max_num_requests=max_num_requests, - ) - if spec_config.spec_dec_mode.is_ngram( - ) or spec_config.spec_dec_mode.is_user_provided(): + if spec_config.spec_dec_mode.is_draft_target() or \ + spec_config.spec_dec_mode.is_ngram() or \ + spec_config.spec_dec_mode.is_user_provided(): return SpecMetadata( - max_draft_tokens=spec_config.max_draft_len, + max_draft_len=spec_config.max_draft_len, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, ) @@ -151,36 +142,3 @@ def get_spec_worker(spec_config, mapping): if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelWorker(spec_config, mapping) return None - - -def get_draft_model_prompt(spec_dec_mode: SpeculativeDecodingMode, - input_tokens: torch.Tensor) -> torch.Tensor: - """ - Can be used to modify prompts for speculative algorithms that need to update tokens - before drafting. - """ - if spec_dec_mode.is_eagle3(): - # EAGLE3 always throws away the first token when processing draft inputs - return input_tokens[1:] - return input_tokens - - -def get_num_extra_kv_tokens(spec_config): - """ - Implementation detail for one model implementations of speculative decoding. Extra - KV cache tokens are required. - """ - if spec_config is None: - return 0 - if spec_config.spec_dec_mode.is_eagle3_one_model( - ) or spec_config.spec_dec_mode.is_mtp_eagle(): - return spec_config.max_draft_len - 1 - return 0 - - -def update_spec_config_from_model_config(spec_config, model_config): - if spec_config.spec_dec_mode.is_mtp(): - # Use `max_draft_len` for several low-level APIs. TODO: Remove this after distinguishing them, - spec_config.max_draft_len = spec_config.num_nextn_predict_layers - # Use `num_nextn_predict_layers_from_model_config` to decide decoding mode MTP / MTP_EAGLE. - spec_config.num_nextn_predict_layers_from_model_config = model_config.num_nextn_predict_layers diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1ba53c0b7760..4d0ac968e7d9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -223,7 +223,8 @@ class _ModelFormatKind(Enum): class DecodingBaseConfig(BaseModel): max_draft_len: Optional[int] = None - speculative_model: Optional[Union[str, Path]] = None + speculative_model_dir: Optional[Union[str, Path]] = None + num_extra_kv_tokens: int = 0 @classmethod def from_dict(cls, data: dict): @@ -270,6 +271,13 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.from_string( self.decoding_type.upper()) + def update_from_model_config(self, model_config): + pass + + def get_draft_model_prompt(self, + input_tokens: torch.Tensor) -> torch.Tensor: + return input_tokens + class MedusaDecodingConfig(DecodingBaseConfig): medusa_choices: Optional[List[List[int]]] = None @@ -302,7 +310,7 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "Eagle" def validate(self) -> None: - if self.speculative_model is None: + if self.speculative_model_dir is None: raise ValueError("Draft model must be provided for EAGLE") @functools.cached_property @@ -313,6 +321,13 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.EAGLE3_ONE_MODEL return TorchSpeculativeDecodingMode.EAGLE3 + def get_draft_model_prompt(self, + input_tokens: torch.Tensor) -> torch.Tensor: + """ + Eagle3 always throws away the first token when processing draft inputs + """ + return input_tokens[1:] + class UserProvidedDecodingConfig(DecodingBaseConfig): # Type should be Drafter, but it leads to circular import @@ -330,7 +345,7 @@ class NGramDecodingConfig(DecodingBaseConfig): Configuration for NGram drafter speculative decoding. Arguments: - prompt_lookup_num_tokens: int + max_draft_len: int The length maximum of draft tokens (can be understood as length maximum of output draft tokens). max_matching_ngram_size: int @@ -346,7 +361,6 @@ class NGramDecodingConfig(DecodingBaseConfig): Whether to use a common pool for all requests, or the pool is private for each request if False. """ - prompt_lookup_num_tokens: int = 2 max_matching_ngram_size: int = 4 is_keep_all: bool = True is_use_oldest: bool = True @@ -368,7 +382,7 @@ class DraftTargetDecodingConfig(DecodingBaseConfig): def from_dict(cls, data: dict): return cls(**data) - decoding_type: ClassVar[str] = "DraftTarget" + decoding_type: ClassVar[str] = "Draft_Target" def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -409,6 +423,11 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.MTP_EAGLE return TorchSpeculativeDecodingMode.MTP + def update_from_model_config(self, model_config): + assert self.num_nextn_predict_layers > 0 + if model_config.num_nextn_predict_layers == 1 and not self.use_mtp_vanilla: + self.num_extra_kv_tokens = self.num_nextn_predict_layers - 1 + class PybindMirror(ABC): ''' A class containing the utilities for mirroring Python classes to @@ -1104,7 +1123,7 @@ def model_format(self) -> _ModelFormatKind: return self._model_format @property - def speculative_model(self) -> Optional[_ModelFormatKind]: + def speculative_model_dir(self) -> Optional[_ModelFormatKind]: return self._speculative_model @property @@ -1389,43 +1408,32 @@ def validate_speculative_config(self): # Below, we only need to set speculative_decoding_mode/decoding_config for speculation # on the TRT backend. if isinstance(self.speculative_config, LookaheadDecodingConfig): - lookahead_config = self.speculative_config - # Update the build config - _, _, max_draft_tokens, _ = lookahead_config.calculate_speculative_resource( - ) + max_draft_len = self.speculative_config.calculate_speculative_resource( + )[2] + assert max_draft_len > 0 self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.LOOKAHEAD_DECODING - if max_draft_tokens > self.build_config.max_draft_len: - self.build_config.max_draft_len = max_draft_tokens - + self.build_config.max_draft_len = max( + self.build_config.max_draft_len, max_draft_len) self.decoding_config = DecodingConfig( decoding_mode=DecodingMode.Lookahead(), lookahead_decoding_config=PybindMirror.maybe_to_pybind( - lookahead_config)) + self.speculative_config)) elif isinstance(self.speculative_config, MedusaDecodingConfig): - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.MEDUSA - assert self.speculative_config.max_draft_len > 0 + self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.MEDUSA self.build_config.max_draft_len = self.speculative_config.max_draft_len self.decoding_config = DecodingConfig( decoding_mode=DecodingMode.Medusa(), medusa_choices=self.speculative_config.medusa_choices) - elif isinstance(self.speculative_config, DraftTargetDecodingConfig): - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL - assert self.speculative_config.max_draft_len > 0 - self.build_config.max_draft_len = self.speculative_config.max_draft_len - - elif isinstance(self.speculative_config, MTPDecodingConfig): - assert self.speculative_config.num_nextn_predict_layers > 0 - self.build_config.max_draft_len = self.speculative_config.num_nextn_predict_layers - elif isinstance(self.speculative_config, EagleDecodingConfig): - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.EAGLE assert self.speculative_config.max_draft_len > 0 - + assert self.speculative_config.speculative_model_dir is not None, "Path to EAGLE3 weights must be specified." self.build_config.max_draft_len = self.speculative_config.max_draft_len - + self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.EAGLE + if self.speculative_config.eagle3_one_model: + self.num_extra_kv_tokens = self.max_draft_len - 1 if self.backend not in ['pytorch', '_autodeploy']: eagle_config = _EagleConfig( self.speculative_config.eagle_choices, @@ -1438,13 +1446,24 @@ def validate_speculative_config(self): eagle_config=eagle_config) elif isinstance(self.speculative_config, NGramDecodingConfig): - assert self.speculative_config.prompt_lookup_num_tokens > 0 and self.speculative_config.max_matching_ngram_size > 0 + assert self.backend in ['pytorch', '_autodeploy'] + assert self.speculative_config.max_draft_len > 0 and self.speculative_config.max_matching_ngram_size > 0 self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.NGRAM self.build_config.max_draft_len = self.speculative_config.max_draft_len + elif isinstance(self.speculative_config, DraftTargetDecodingConfig): + assert self.backend in ['pytorch'] + assert self.speculative_config.max_draft_len > 0 + self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL + self.build_config.max_draft_len = self.speculative_config.max_draft_len + + elif isinstance(self.speculative_config, MTPDecodingConfig): + assert self.speculative_config.num_nextn_predict_layers > 0 + self.speculative_config.max_draft_len = self.speculative_config.num_nextn_predict_layers + elif isinstance(self.speculative_config, UserProvidedDecodingConfig): - assert self.speculative_config.max_draft_len > 0 + assert self.backend in ['pytorch', '_autodeploy'] self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.USER_PROVIDED self.build_config.max_draft_len = self.speculative_config.max_draft_len @@ -1452,11 +1471,12 @@ def validate_speculative_config(self): raise ValueError( f"Unrecognized speculative config type {type(self.speculative_config)}" ) + else: self.decoding_config = None self._speculative_model = getattr(self.speculative_config, - "speculative_model", None) + "speculative_model_dir", None) speculative_model_obj = _ModelWrapper( self._speculative_model ) if self._speculative_model is not None else None diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index e5bbe61fb0ab..cf2bdb26c143 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -110,8 +110,8 @@ def __init__(self, self.model_obj = _ModelWrapper(self.llm_args.model) self.speculative_model_obj = _ModelWrapper( - self.llm_args.speculative_model - ) if self.llm_args.speculative_model is not None else None + self.llm_args.speculative_model_dir + ) if self.llm_args.speculative_model_dir is not None else None if isinstance(self.llm_args, TrtLlmArgs): self.convert_checkpoint_options = self.llm_args._convert_checkpoint_options @@ -440,8 +440,8 @@ def _load_model_from_hf(self): model_cls = AutoModelForCausalLM.get_trtllm_model_class( self._model_dir, self.llm_args.trust_remote_code, self.llm_args.decoding_config.decoding_mode - if hasattr(self.llm_args, "speculative_model") - and self.llm_args.speculative_model else None) + if hasattr(self.llm_args, "speculative_model_dir") + and self.llm_args.speculative_model_dir else None) prequantized = self._update_from_hf_quant_config() @@ -484,7 +484,7 @@ def _load_model_from_hf(self): load_model_on_cpu= True, # TODO:TRTLLM-195 to enhance the weights loading memory usage and chose best location trust_remote_code=self.llm_args.trust_remote_code, - speculative_model=self._speculative_model_dir, + speculative_model_dir=self._speculative_model_dir, speculative_config=self.llm_args.speculative_config if not isinstance(self.llm_args.speculative_config, LookaheadDecodingConfig) else None, diff --git a/tensorrt_llm/models/eagle/config.py b/tensorrt_llm/models/eagle/config.py index 6124884f723b..f81e43bb03ff 100644 --- a/tensorrt_llm/models/eagle/config.py +++ b/tensorrt_llm/models/eagle/config.py @@ -59,7 +59,7 @@ def from_hugging_face( **kwargs): import transformers trust_remote_code = kwargs.pop('trust_remote_code', True) - speculative_config_or_dir = kwargs.pop('speculative_model', None) + speculative_config_or_dir = kwargs.pop('speculative_model_dir', None) if isinstance(hf_config_or_dir, transformers.PretrainedConfig): hf_config = hf_config_or_dir diff --git a/tensorrt_llm/models/eagle/model.py b/tensorrt_llm/models/eagle/model.py index 01b3905ba798..07a9d97843f2 100644 --- a/tensorrt_llm/models/eagle/model.py +++ b/tensorrt_llm/models/eagle/model.py @@ -948,11 +948,11 @@ def prepare_inputs(self, *args, **kwargs): spec_decoding_position_offsets: [bs, max_gen_tokens] spec_decoding_packed_mask: [bs, max_draft_len, packed_length] ** eagle_temperature: [bs] - rand_data_validation: [bs, max_draft_tokens] + rand_data_validation: [bs, max_draft_len] ** The mask is tricky since the boolean mask will need to be packed in runtime. So, the last dim will be: - packed_length = ceil((max_draft_tokens+1)/32) + packed_length = ceil((max_draft_len+1)/32) """ default_range = GenerationMixin.default_range remove_input_padding = default_net().plugin_config.remove_input_padding @@ -1228,7 +1228,7 @@ def from_hugging_face( quant_config: Optional[QuantConfig] = None, **kwargs): assert hf_model_or_dir is not None - speculative_model_dir = kwargs.get('speculative_model', None) + speculative_model_dir = kwargs.get('speculative_model_dir', None) tllm_config = EagleConfig.from_hugging_face(hf_model_or_dir, dtype=dtype, mapping=mapping, diff --git a/tensorrt_llm/models/medusa/config.py b/tensorrt_llm/models/medusa/config.py index 361cb523068c..b686cf77143d 100644 --- a/tensorrt_llm/models/medusa/config.py +++ b/tensorrt_llm/models/medusa/config.py @@ -70,7 +70,7 @@ def from_hugging_face( import transformers trust_remote_code = kwargs.pop('trust_remote_code', True) - speculative_config_or_dir = kwargs.pop('speculative_model', None) + speculative_config_or_dir = kwargs.pop('speculative_model_dir', None) speculative_config = kwargs.pop("speculative_config", None) if isinstance(hf_config_or_dir, transformers.PretrainedConfig): diff --git a/tensorrt_llm/models/medusa/model.py b/tensorrt_llm/models/medusa/model.py index de044de6fa19..7eafbbafc760 100644 --- a/tensorrt_llm/models/medusa/model.py +++ b/tensorrt_llm/models/medusa/model.py @@ -191,7 +191,7 @@ def from_hugging_face( import transformers assert hf_model_or_dir is not None - speculative_model_dir = kwargs.get('speculative_model', None) + speculative_model_dir = kwargs.get('speculative_model_dir', None) use_preloading = isinstance(hf_model_or_dir, transformers.PreTrainedModel) diff --git a/tensorrt_llm/models/redrafter/model.py b/tensorrt_llm/models/redrafter/model.py index aa909a22b9dd..4f9f139f053b 100644 --- a/tensorrt_llm/models/redrafter/model.py +++ b/tensorrt_llm/models/redrafter/model.py @@ -179,15 +179,15 @@ def prepare_inputs(self, *args, **kwargs): bb_range = default_range(max_batch_size) bb0_range = default_range(max_batch_size, min_range=0, opt_offset=1) num_beam_tokens = self.num_beams * self.beam_length - max_draft_tokens = num_beam_tokens - self.num_beams # ignore the true token - max_gen_token_len = 1 + max_draft_tokens # for the true token + max_draft_len = num_beam_tokens - self.num_beams # ignore the true token + max_gen_token_len = 1 + max_draft_len # for the true token max_gen_token_len_range = default_range(max_gen_token_len) bb_max_gen_token_len_range = default_range(max_gen_token_len * max_batch_size, min_range=0) kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = max_draft_tokens + kwargs['max_draft_len'] = max_draft_len kwargs['spec_decoding_is_generation_length_variable'] = True inputs = super().prepare_inputs(*args, **kwargs) assert inputs['spec_decoding_params'] is not None diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index 71057092f97d..7a8a8a888100 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -25,6 +25,7 @@ import tensorrt_llm.evaluate from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm._torch.speculative import SpecConfig from tensorrt_llm.builder import BuildConfig from tensorrt_llm.llmapi import SamplingParams from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig @@ -155,6 +156,10 @@ def evaluate(self, spec_dec_algo = None elif isinstance(llm.args.speculative_config, DecodingBaseConfig): spec_dec_algo = llm.args.speculative_config.decoding_type + elif isinstance(llm.args.speculative_config, SpecConfig): + # This branch is deprecated, but thread-leak of pytest raises flaky error if removing it. + # TODO: remove this branch safely. + spec_dec_algo = llm.args.speculative_config.spec_dec_name else: raise ValueError( f"Not recognized speculative_config: {llm.args.speculative_config}." diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 46006d0545c7..f8f503fffaed 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -214,7 +214,7 @@ def test_auto_dtype(self, disable_overlap_scheduler): def test_ngram(self): speculative_decoding_config = { "decoding_type": "NGram", - "prompt_lookup_num_tokens": 4, + "max_draft_len": 4, "max_matching_ngram_size": 4, "is_keep_all": True, "is_use_oldest": True, diff --git a/tests/integration/defs/accuracy/test_llm_api.py b/tests/integration/defs/accuracy/test_llm_api.py index 872685ef22fb..6033eae3b6ad 100644 --- a/tests/integration/defs/accuracy/test_llm_api.py +++ b/tests/integration/defs/accuracy/test_llm_api.py @@ -393,7 +393,7 @@ class TestEagleVicuna_7B_v1_3(LlmapiAccuracyTestHarness): speculative_config = EagleDecodingConfig( max_draft_len=63, - speculative_model=f"{llm_models_root()}/EAGLE-Vicuna-7B-v1.3", + speculative_model_dir=f"{llm_models_root()}/EAGLE-Vicuna-7B-v1.3", num_eagle_layers=4, max_non_leaves_per_layer=10, eagle_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ @@ -419,7 +419,7 @@ class TestEagle2Vicuna_7B_v1_3(LlmapiAccuracyTestHarness): speculative_config = EagleDecodingConfig( max_draft_len=63, - speculative_model=f"{llm_models_root()}/EAGLE-Vicuna-7B-v1.3", + speculative_model_dir=f"{llm_models_root()}/EAGLE-Vicuna-7B-v1.3", num_eagle_layers=4, max_non_leaves_per_layer=10, use_dynamic_tree=True, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8111ffd26cb5..8eb4377bdf7a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -243,7 +243,7 @@ def test_eagle3(self): draft_len = 4 spec_config = EagleDecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir) + speculative_model_dir=eagle_model_dir) llm = LLM(model=target_model_dir, **pytorch_config, @@ -263,7 +263,6 @@ def test_ngram(self): draft_len = 4 spec_config = NGramDecodingConfig( max_draft_len=draft_len, - prompt_lookup_num_tokens=draft_len, max_matching_ngram_size=draft_len, is_keep_all=True, is_use_oldest=True, diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml index 9c9c9e0c8b4c..667262df4a3e 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml @@ -18,7 +18,7 @@ generation_servers: - "localhost:8002" speculative_config: decoding_type: NGram - prompt_lookup_num_tokens: 4 + max_draft_len: 4 max_matching_ngram_size: 4 is_keep_all: True is_use_oldest: True diff --git a/tests/unittest/_torch/speculative/test_draft_target.py b/tests/unittest/_torch/speculative/test_draft_target.py index 75b93b3a29c0..23d26126f897 100644 --- a/tests/unittest/_torch/speculative/test_draft_target.py +++ b/tests/unittest/_torch/speculative/test_draft_target.py @@ -22,45 +22,43 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str): pytest.skip("Not enough memory to load target model") models_path = llm_models_root() + draft_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=2080) - - sampling_params = SamplingParams( - max_tokens=32, - temperature=0, + max_batch_size = 2 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=False) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + backend='pytorch', + attn_backend=attn_backend, + disable_overlap_scheduler=True, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_num_tokens=2048, ) - max_batch_size = 1 - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - draft_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - draft_len = 4 - spec_config = DraftTargetDecodingConfig(max_draft_len=draft_len, - speculative_model=draft_model_dir) - llm_spec = LLM(model=target_model_dir, - max_batch_size=max_batch_size, - disable_overlap_scheduler=True, - cuda_graph_config=CudaGraphConfig( - cuda_graph_batch_sizes=[1]) if use_cuda_graph else None, - attn_backend=attn_backend, - kv_cache_config=kv_cache_config, - speculative_config=spec_config) + spec_config = DraftTargetDecodingConfig( + max_draft_len=max_draft_len, + speculative_model_dir=draft_model_dir, + ) prompts = [ - "The capital of France is", "The president of the United States is" + "The capital of France is", + "The president of the United States is", ] + sampling_params = SamplingParams(max_tokens=32) + + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) results_spec = llm_spec.generate(prompts, sampling_params) generated_text_spec = [result.outputs[0].text for result in results_spec] llm_spec.shutdown() - llm_ref = LLM(model=target_model_dir, - max_batch_size=max_batch_size, - disable_overlap_scheduler=True, - cuda_graph_config=CudaGraphConfig( - cuda_graph_batch_sizes=[1]) if use_cuda_graph else None, - attn_backend=attn_backend, - kv_cache_config=kv_cache_config) - + llm_ref = LLM(**llm_common_config) results_ref = llm_ref.generate(prompts, sampling_params) generated_text_ref = [result.outputs[0].text for result in results_ref] llm_ref.shutdown() diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 07f1035531b1..a228c5429130 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -32,81 +32,70 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, pytest.skip("Not enough memory to load target + draft model") models_path = llm_models_root() - - pytorch_config = dict( - disable_overlap_scheduler=disable_overlap_scheduler, - # Only create a single CUDA graph to prevent OOM in CI - attn_backend=attn_backend, - ) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1]) if use_cuda_graph else None - - kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, ) - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - draft_len = 4 - spec_config = EagleDecodingConfig( - max_draft_len=draft_len, - speculative_model=eagle_model_dir, - # Llama 3 does not support one model eagle. - eagle3_one_model=use_one_model) + # bs > 1 gives non-deterministic when doing IFB. There are slight chances + # that ref and spec does not match 100% + max_batch_size = 1 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, + free_gpu_memory_fraction=0.5) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None - llm_spec = LLM( + llm_common_config = dict( model=target_model_dir, - **pytorch_config, - # bs > 1 gives non-deterministic when doing IFB. There are slight chances - # that ref and spec does not match 100% - max_batch_size=1, + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, # This max_seq_len is larger than the one specified # in the llama 3 8B eagle's config. We want to make sure # that the draft model won't go above its max in warmup # in this test. max_seq_len=8192, - kv_cache_config=kv_cache_config, - cuda_graph_config=cuda_graph_config, - speculative_config=spec_config) + ) - sampling_params = SamplingParams( - max_tokens=10, - temperature=0, + spec_config = EagleDecodingConfig( + max_draft_len=max_draft_len, + speculative_model_dir=eagle_model_dir, + # Llama 3 does not support one model eagle. + eagle3_one_model=use_one_model, ) - # First make sure the acceptance rate is reasonable. + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + # Acceptance rate tests tok_ids = llm_spec.tokenizer.encode("The future of AI is") num_tokens = 0 - num_drafted = 0 num_accepted = 0 - + sampling_params = SamplingParams(max_tokens=128, temperature=0) for output in llm_spec.generate_async(tok_ids, - SamplingParams(max_tokens=128, - temperature=0), + sampling_params, streaming=True): - beam = output.outputs[0] - new_tokens = beam.token_ids - - num_drafted += draft_len + new_tokens = output.outputs[0].token_ids + num_drafted += max_draft_len num_accepted += len(new_tokens) - num_tokens - 1 - num_tokens = len(new_tokens) accept_rate = num_accepted / num_drafted assert accept_rate > 0.15 + # Output tests prompts = [ - "The capital of France is", "The president of the United States is" + "The capital of France is", + "The president of the United States is", ] + sampling_params = SamplingParams(max_tokens=10, temperature=0) + results_spec = llm_spec.generate(prompts, sampling_params) generated_text_spec = [result.outputs[0].text for result in results_spec] llm_spec.shutdown() - llm_ref = LLM(model=target_model_dir, - **pytorch_config, - kv_cache_config=kv_cache_config, - cuda_graph_config=cuda_graph_config) - + llm_ref = LLM(**llm_common_config) results_ref = llm_ref.generate(prompts, sampling_params) generated_text_ref = [result.outputs[0].text for result in results_ref] llm_ref.shutdown() diff --git a/tests/unittest/_torch/speculative/test_mtp.py b/tests/unittest/_torch/speculative/test_mtp.py index 3d6438dfd638..c4a9783e792f 100644 --- a/tests/unittest/_torch/speculative/test_mtp.py +++ b/tests/unittest/_torch/speculative/test_mtp.py @@ -40,15 +40,14 @@ def load_sample_and_accept_draft_tokens_test_cases(): device="cuda") # [num_tokens, vocab_size] draft_tokens = torch.tensor( - [], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + [], dtype=torch.int, device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([0], dtype=torch.int, device="cuda") # [batch_size] ref_accepted_tokens = torch.tensor( [[1, 0]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([1], dtype=torch.int, @@ -74,15 +73,14 @@ def load_sample_and_accept_draft_tokens_test_cases(): device="cuda") # [num_tokens, vocab_size] draft_tokens = torch.tensor( - [], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + [], dtype=torch.int, device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([0, 0, 0, 0], dtype=torch.int, device="cuda") # [batch_size] ref_accepted_tokens = torch.tensor( [[1, 0], [3, 0], [3, 0], [6, 0]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([1, 1, 1, 1], dtype=torch.int, @@ -111,14 +109,14 @@ def load_sample_and_accept_draft_tokens_test_cases(): draft_tokens = torch.tensor( [1, 3, 4], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([3], dtype=torch.int, device="cuda") # [batch_size] ref_accepted_tokens = torch.tensor( [[1, 3, 2, 0]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([3], dtype=torch.int, @@ -147,14 +145,14 @@ def load_sample_and_accept_draft_tokens_test_cases(): draft_tokens = torch.tensor( [1, 5], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([1, 1], dtype=torch.int, device="cuda") # [batch_size] ref_accepted_tokens = torch.tensor( [[1, 3], [4, 0]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([2, 1], dtype=torch.int, @@ -187,14 +185,14 @@ def load_sample_and_accept_draft_tokens_test_cases(): draft_tokens = torch.tensor( [1, 3, 4, 4, 7, 3], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([3, 3], dtype=torch.int, device="cuda") # [batch_size] ref_accepted_tokens = torch.tensor( [[1, 3, 2, 0], [4, 6, 0, 0]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([3, 2], dtype=torch.int, @@ -231,7 +229,7 @@ def load_sample_and_accept_draft_tokens_test_cases(): draft_tokens = torch.tensor( [1, 3, 5, 4, 6, 5, 5, 7, 4], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([3, 3, 3], dtype=torch.int, device="cuda") # [batch_size] @@ -239,7 +237,7 @@ def load_sample_and_accept_draft_tokens_test_cases(): ref_accepted_tokens = torch.tensor( [[1, 3, 2, 0], [4, 6, 5, 2], [4, 0, 0, 0]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([3, 4, 1], dtype=torch.int, @@ -267,15 +265,14 @@ def load_sample_and_accept_draft_tokens_test_cases(): device="cuda") # [num_tokens, vocab_size] draft_tokens = torch.tensor( - [4], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + [4], dtype=torch.int, device="cuda") # [batch_size * max_draft_len] draft_len = torch.tensor([0, 1], dtype=torch.int, device="cuda") # [batch_size] ref_accepted_tokens = torch.tensor( [[1, 0], [4, 6]], dtype=torch.int, - device="cuda") # [batch_size * max_draft_tokens] + device="cuda") # [batch_size * max_draft_len] ref_num_accepted_tokens = torch.tensor([1, 2], dtype=torch.int, @@ -311,7 +308,7 @@ def test_sample_and_accept_draft_tokens(self, test_case_name, # speculative decoding metadata spec_metadata = MTPSpecMetadata(max_num_requests=32, spec_dec_mode=spec_config.spec_dec_mode, - max_draft_tokens=mtp_num_modules, + max_draft_len=mtp_num_modules, mtp_num_modules=mtp_num_modules) spec_metadata.draft_tokens = draft_tokens @@ -893,7 +890,7 @@ def test_mtp_update_mtp_hidden_states( spec_metadata = MTPSpecMetadata( max_num_requests=32, spec_dec_mode=spec_config.spec_dec_mode, - max_draft_tokens=num_nextn_predict_layers, + max_draft_len=num_nextn_predict_layers, mtp_num_modules=num_nextn_predict_layers, mtp_hidden_states_manager=spec_manager) spec_metadata.request_ids = request_ids @@ -1388,7 +1385,7 @@ def test_prepare_drafter_inputs( spec_metadata = MTPSpecMetadata( max_num_requests=32, spec_dec_mode=spec_config.spec_dec_mode, - max_draft_tokens=num_nextn_predict_layers, + max_draft_len=num_nextn_predict_layers, mtp_num_modules=num_nextn_predict_layers, mtp_hidden_states_manager=spec_manager) spec_metadata.request_ids = request_ids diff --git a/tests/unittest/_torch/speculative/test_ngram.py b/tests/unittest/_torch/speculative/test_ngram.py index 39e5dd058c85..90bd64715611 100644 --- a/tests/unittest/_torch/speculative/test_ngram.py +++ b/tests/unittest/_torch/speculative/test_ngram.py @@ -24,11 +24,12 @@ def test_llama_ngram(disable_overlap_scheduler: bool, use_cuda_graph: bool, total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if total_mem_gb < 20: pytest.skip("Not enough memory to load target model") - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1]) if use_cuda_graph else None - max_batch_size = 4 + max_batch_size = 2 max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=False) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None llm_common_config = dict( \ model=llm_models_root() / "llama-3.1-model" /"Meta-Llama-3.1-8B", @@ -37,13 +38,12 @@ def test_llama_ngram(disable_overlap_scheduler: bool, use_cuda_graph: bool, disable_overlap_scheduler=disable_overlap_scheduler, cuda_graph_config=cuda_graph_config, max_batch_size=max_batch_size, - kv_cache_config=KvCacheConfig(enable_block_reuse=False), + kv_cache_config=kv_cache_config, max_num_tokens=2048, ) spec_config = NGramDecodingConfig( max_draft_len=max_draft_len, - prompt_lookup_num_tokens=max_draft_len, max_matching_ngram_size=2, is_keep_all=True, is_use_oldest=True, diff --git a/tests/unittest/_torch/speculative/test_user_provided.py b/tests/unittest/_torch/speculative/test_user_provided.py index 178165400d04..580c700e5aab 100644 --- a/tests/unittest/_torch/speculative/test_user_provided.py +++ b/tests/unittest/_torch/speculative/test_user_provided.py @@ -25,11 +25,12 @@ def test_llama_user_provided(disable_overlap_scheduler: bool, total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if total_mem_gb < 20: pytest.skip("Not enough memory to load target model") - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1]) if use_cuda_graph else None - max_batch_size = 4 + max_batch_size = 2 max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=False) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None llm_common_config = dict( \ model=llm_models_root() / "llama-3.1-model" /"Meta-Llama-3.1-8B", @@ -38,13 +39,12 @@ def test_llama_user_provided(disable_overlap_scheduler: bool, disable_overlap_scheduler=disable_overlap_scheduler, cuda_graph_config=cuda_graph_config, max_batch_size=max_batch_size, - kv_cache_config=KvCacheConfig(enable_block_reuse=False), + kv_cache_config=kv_cache_config, max_num_tokens=2048, ) ngram_config = NGramDecodingConfig( max_draft_len=max_draft_len, - prompt_lookup_num_tokens=max_draft_len, max_matching_ngram_size=2, is_keep_all=True, is_use_oldest=True, diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index ad87291aadc5..b17ada143948 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -1119,7 +1119,7 @@ def test_llm_api_medusa(): speculative_config = MedusaDecodingConfig(num_medusa_heads=4, max_draft_len=63, - speculative_model=get_model_path("medusa-vicuna-7b-v1.3"), + speculative_model_dir=get_model_path("medusa-vicuna-7b-v1.3"), medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], \ [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], \ @@ -1158,7 +1158,7 @@ def test_llm_api_medusa_tp2(): speculative_config = MedusaDecodingConfig(num_medusa_heads=4, max_draft_len=63, - speculative_model=get_model_path("medusa-vicuna-7b-v1.3"), + speculative_model_dir=get_model_path("medusa-vicuna-7b-v1.3"), medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], \ [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], \ @@ -1196,7 +1196,7 @@ def test_llm_api_eagle(**llm_kwargs): speculative_config = EagleDecodingConfig( max_draft_len=63, - speculative_model=get_model_path("EAGLE-Vicuna-7B-v1.3"), + speculative_model_dir=get_model_path("EAGLE-Vicuna-7B-v1.3"), num_eagle_layers=4, max_non_leaves_per_layer=10, eagle_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ @@ -1243,7 +1243,7 @@ def test_llm_api_eagle2(**llm_kwargs): speculative_config = EagleDecodingConfig( max_draft_len=63, - speculative_model=get_model_path("EAGLE-Vicuna-7B-v1.3"), + speculative_model_dir=get_model_path("EAGLE-Vicuna-7B-v1.3"), num_eagle_layers=4, max_non_leaves_per_layer=10, use_dynamic_tree=True,