diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py index de6de657009a..346cbcf0b490 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py @@ -775,7 +775,7 @@ def get_output_embeddings(self): @dataclass class EagleWrapperOutput(ModelOutput): - """Output format compatible with Eagle3OneModelSampler/MTPSampler. + """Output format compatible with SpecSampler. This output format allows the one-model speculative decoding flow to bypass logits-based sampling in the sampler. The EagleWrapper performs all sampling diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 86488bc36cda..1ed345a1914f 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -48,7 +48,7 @@ SimpleScheduler, ) from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager -from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSampler +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler from tensorrt_llm._utils import get_free_port, mpi_rank, mpi_world_size, nvtx_range from tensorrt_llm.inputs.multimodal import MultimodalRuntimeData, check_mm_embed_cumsum_if_needed from tensorrt_llm.llmapi.llm_args import ContextChunkingPolicy, MultimodalConfig, SamplerType @@ -1113,7 +1113,7 @@ def instantiate_sampler( max_beam_width=ad_config.max_beam_width, disable_overlap_scheduler=ad_config.disable_overlap_scheduler, ) - return Eagle3OneModelSampler(sampler_args) + return SpecSampler(sampler_args) sampler_type = ad_config.sampler_type if sampler_type == SamplerType.auto: diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/top_p_decay.py b/tensorrt_llm/_torch/pyexecutor/sampler/top_p_decay.py index 853b687fe3a7..3343209ff3ec 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/top_p_decay.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/top_p_decay.py @@ -153,7 +153,7 @@ def validate_request(request: LlmRequest) -> None: # tokens and produces multiple tokens per step (req_num_steps = # 1 + draft_token_length). One-model speculation (vanilla MTP, one-model # Eagle3 / MTP-Eagle, SA, draft-target-one-model) uses its own - # SpecSamplerBase-derived sampler and never reaches TorchSampler; the + # SpecSampler and never reaches TorchSampler; the # drafter-based modes that DO flow draft tokens through TorchSampler # (two-model draft-target, NGram, user-provided, two-model Eagle3 / # MTP-Eagle) are what can make this length non-zero. top-p decay does not diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 347f3aae5dc2..9151d1ac0589 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -7,15 +7,15 @@ prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, should_use_separate_draft_kv_cache) -from .mtp import MTPSampler, MTPSpecMetadata, MTPWorker +from .mtp import MTPSpecMetadata, MTPWorker from .ngram import NGramDrafter, NGramPoolManager from .pard import PARDSpecMetadata, PARDWorker from .sa_enhancer import SADraftEnhancer -from .sa_worker import SASampler, SASpecMetadata, SAWorker +from .sa_worker import SASpecMetadata, SAWorker from .save_hidden_state import (SaveHiddenStatesResourceManager, SaveHiddenStatesSpecMetadata) from .spec_sampler_base import (SampleStateSpec, SampleStateTensorsSpec, - SpecSamplerBase) + SpecSampler) from .spec_tree_manager import SpecTreeManager from .suffix_automaton import SuffixAutomatonManager from .utils import (get_draft_kv_cache_manager, get_num_extra_kv_tokens, @@ -32,7 +32,6 @@ "DraftTargetOneModelWorker", "Eagle3SpecMetadata", "MTPEagleWorker", - "MTPSampler", "MTPSpecMetadata", "MTPWorker", "NGramDrafter", @@ -40,7 +39,6 @@ "PARDSpecMetadata", "PARDWorker", "SADraftEnhancer", - "SASampler", "SASpecMetadata", "SAWorker", "SuffixAutomatonManager", @@ -49,7 +47,7 @@ "SaveHiddenStatesResourceManager", "SaveHiddenStatesSpecMetadata", "SpecMetadata", - "SpecSamplerBase", + "SpecSampler", "SpecWorkerBase", "get_draft_kv_cache_manager", "get_num_extra_kv_tokens", diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index 6edeac334fdd..8d25eeb68c9f 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -30,9 +30,7 @@ from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata -from ..pyexecutor.sampler import TorchSampler from .interface import SpecMetadata, SpecWorkerBase -from .mtp import MTPSampler if TYPE_CHECKING: from ...llmapi.llm_args import DraftTargetDecodingConfig @@ -75,17 +73,6 @@ def prepare(self): self.is_spec_dec_dynamic_tree = False -class DraftTargetOneModelSampler(MTPSampler): - """ - Sampler for DraftTarget one-model speculative decoding. - - Inherits from MTPSampler to reuse the speculative decoding sampling logic. - """ - - def __init__(self, args: TorchSampler.Args): - super().__init__(args, nextn=args.max_draft_len) - - class DraftTargetOneModelWorker(SpecWorkerBase): def __init__( self, diff --git a/tensorrt_llm/_torch/speculative/drafter.py b/tensorrt_llm/_torch/speculative/drafter.py index 0396bbeb0dd7..a5adba45881e 100644 --- a/tensorrt_llm/_torch/speculative/drafter.py +++ b/tensorrt_llm/_torch/speculative/drafter.py @@ -67,7 +67,7 @@ def should_use_spec_decode(self, requests: List[LlmRequest], # Drafters that use TorchSampler (NGram, two-model) compute py_rewind_len # from len(py_draft_tokens), which includes padding. They must set this # to True so that extend_capacity_for_tokens is called after padding. - # One-model drafters (MTP / Eagle3 / SA) use SpecSamplerBase which + # One-model drafters (MTP / Eagle3 / SA) use SpecSampler which # computes rewind from runtime_draft_len, so padding is harmless. _needs_padding_kv_extension: bool = False diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index a1f4f82869bd..bb9252801536 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -14,10 +14,9 @@ from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager -from ..pyexecutor.sampler import TorchSampler from ..pyexecutor.scheduler import ScheduledRequests from .interface import SpecMetadata, SpecWorkerBase -from .mtp import MTPSampler, _select_mtp_position_ids +from .mtp import _select_mtp_position_ids from .sa_enhancer import SADraftEnhancer from .spec_tree_manager import SpecTreeManager @@ -591,22 +590,6 @@ def maybe_capture_hidden_states( break -class Eagle3OneModelSampler(MTPSampler): - """Sampler for one-model EAGLE3 (linear and dynamic tree modes).""" - - def __init__(self, args: TorchSampler.Args, spec_config=None): - self._spec_config = spec_config - super().__init__(args, nextn=args.max_total_draft_tokens) - - def _get_max_new_tokens(self, args: TorchSampler.Args, - draft_len: int) -> int: - """Dynamic tree: accepted path depth <= max_draft_len + 1.""" - if (self._spec_config is not None - and getattr(self._spec_config, 'use_dynamic_tree', False)): - return self._spec_config.max_draft_len + 1 - return self._get_max_tokens(args, draft_len) - - class Eagle3OneModelWorker(SpecWorkerBase): """Unified one-model worker for Eagle3 and MTP Eagle speculative decoding. diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 721b3942ed04..05839454f997 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -750,7 +750,7 @@ def _normalize_request_sampling_params( # decoding does not support min_p (there is no request_min_p buffer # nor min_p wiring in the sampling_batch_spec_dec_one_model* # kernels); a min_p request is rejected at admission by - # SpecSamplerBase.validate_request, so nothing reaching this scan + # SpecSampler.validate_request, so nothing reaching this scan # carries a min_p that would change its classification. The # two-model draft/target path honors min_p via _request_strategy. is_greedy = SamplingParams.params_imply_greedy_decoding( diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 3dc12a669aa6..044c2c679e8a 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1,4 +1,3 @@ -import sys from dataclasses import dataclass from typing import TYPE_CHECKING, List, Optional @@ -9,22 +8,13 @@ from ..attention_backend import AttentionMetadata from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager -from ..pyexecutor.sampler import TorchSampler from ..pyexecutor.scheduler import ScheduledRequests from .interface import SpecMetadata, SpecWorkerBase from .sa_enhancer import SADraftEnhancer -from .spec_sampler_base import SampleStateSpec, SpecSamplerBase if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig -if sys.version_info[:2] >= (3, 12): - from typing import override -else: - from typing_extensions import override - -SampleStateMTP = SampleStateSpec - def _normalize_mtp_position_ids(position_ids: torch.Tensor) -> torch.Tensor: """Collapse plain [1, N] position IDs while preserving MRoPE axes.""" @@ -239,38 +229,6 @@ def prepare(self): sa_manager.prepare(gen_request_ids, self.runtime_draft_len) -class MTPSampler(SpecSamplerBase): - """ - MTP sampler. - - Inherits from SpecSamplerBase with overrides for tree-based speculation - using max_total_draft_tokens instead of draft_len. - """ - - SampleState = SampleStateMTP - - @override - def is_generation_model(self) -> bool: - return True - - def setup_sampler_step(self, scheduled_requests: ScheduledRequests): - pass - - def __init__(self, args: TorchSampler.Args, *, nextn: int): - super().__init__(args, draft_len=nextn) - - @override - def _get_max_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: - """MTP uses max_total_draft_tokens + 1 for tree-based speculation.""" - return args.max_total_draft_tokens + 1 - - @override - def _get_draft_tokens_storage_size(self, args: TorchSampler.Args, - draft_len: int) -> int: - """MTP uses max_total_draft_tokens for draft token storage.""" - return args.max_total_draft_tokens - - class MTPWorker(SpecWorkerBase): def __init__(self, diff --git a/tensorrt_llm/_torch/speculative/sa_worker.py b/tensorrt_llm/_torch/speculative/sa_worker.py index 9f97c8d69404..d2c6f0ec0400 100644 --- a/tensorrt_llm/_torch/speculative/sa_worker.py +++ b/tensorrt_llm/_torch/speculative/sa_worker.py @@ -21,7 +21,6 @@ Key components: - SASpecMetadata: Metadata for SA speculative decoding - SAWorker: Spec worker that uses suffix automaton for draft generation -- SASampler: Sampler that handles GPU->CPU result extraction """ from dataclasses import dataclass, field @@ -32,9 +31,7 @@ from tensorrt_llm._utils import prefer_pinned from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager -from ..pyexecutor.sampler import TorchSampler from .interface import SpecMetadata, SpecWorkerBase -from .spec_sampler_base import SampleStateSpec, SpecSamplerBase from .suffix_automaton import SuffixAutomatonManager if TYPE_CHECKING: @@ -334,20 +331,3 @@ def _generate_draft_tokens( draft_tokens = draft_tokens * mask return draft_tokens # [batch_size, max_draft_len] GPU tensor - - -class SASampler(SpecSamplerBase): - """ - Sampler for SA that extracts GPU results to CPU after graph replay. - - Uses SpecSamplerBase with default behavior (draft_len + 1 storage, - adds dummy draft tokens for context requests). - """ - - SampleState = SampleStateSpec - - def __init__(self, args: TorchSampler.Args, *, max_draft_len: int): - super().__init__(args, draft_len=max_draft_len) - - def is_generation_model(self) -> bool: - return True diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index f6b2a60569bb..f006a6ae1d10 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -13,10 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Base class for speculative decoding samplers. +Sampler for one-model speculative decoding. -This module provides a common base class for MTPSampler, SASampler, and -Eagle3OneModelSampler. +Every one-model speculative mode (MTP, Eagle3, SA, DraftTarget, PARD, DFlash, +DSpark) shares a single sampler: the worker's fused kernel already performs +drafting, target verification and acceptance, so the sampler only scatters that +output into slot-indexed buffers, starts the async D2H copy, and updates +requests host-side. Buffer shapes derive entirely from ``TorchSampler.Args``. """ from dataclasses import dataclass @@ -57,19 +60,19 @@ class SampleStateSpec(SampleState): host: SampleStateTensorsSpec -class SpecSamplerBase(Sampler[SampleStateSpec], AsyncWorkerMixin): +class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): """ - Base class for speculative decoding samplers (MTP, NGram, Eagle3, SA). + Sampler for all one-model speculative decoding modes. - Provides common functionality: - - Pre-allocated GPU storage buffers + Provides: + - Pre-allocated, slot-indexed GPU storage buffers - Async GPU->CPU copy in sample_async - Request state updates in update_requests - Subclasses can customize behavior by overriding: - - _get_max_tokens(): How to calculate max_tokens for storage - - _get_draft_tokens_storage_size(): Size of next_draft_tokens tensor - - _add_dummy_draft_tokens(): Whether to add dummy drafts for context requests + This class carries no per-mode behavior. ``args.max_total_draft_tokens`` is + ``spec_config.tokens_per_gen_step - 1`` (see ``create_torch_sampler_args``), + i.e. the target's per-step input width minus one, which is exactly the + draft length every mode used to compute for itself. """ SampleState = SampleStateSpec @@ -108,68 +111,57 @@ class Store: next_draft_tokens: torch.Tensor new_tokens_lens: torch.Tensor - def __init__(self, args: TorchSampler.Args, *, draft_len: int): + def __init__(self, args: TorchSampler.Args, *, accepted_path_len: Optional[int] = None): """ Initialize the speculative sampler. Args: args: TorchSampler.Args with max_num_sequences, max_seq_len, etc. - draft_len: Maximum number of draft tokens per iteration. + accepted_path_len: Upper bound on the number of tokens a single step + can accept, used to size new_tokens. Defaults to + ``args.max_draft_len + 1``; see the store comment below for the + one mode that has to override it. """ self._async_worker_init(args.enable_async_worker) self.mapping = None - self.draft_len = draft_len self.max_seq_len = args.max_seq_len + # Wire width minus one: the number of draft slots the target verifies + # per step. Linear modes set max_total_draft_tokens == max_draft_len; + # tree modes set it to the total node count; PARD sets it to 2K-1 + # because it also feeds mask tokens through the target. + self.draft_len = args.max_total_draft_tokens seq_slots = args.max_num_sequences - max_tokens = self._get_max_tokens(args, draft_len) - max_new_tokens = self._get_max_new_tokens(args, draft_len) - draft_tokens_size = self._get_draft_tokens_storage_size(args, draft_len) self.max_beam_width = args.max_beam_width assert self.max_beam_width == 1, "beam width must be 1 for speculative decoding" + # new_tokens holds the accepted tokens only, so it is sized to how many + # a step can accept rather than to the wire width. Normally that is + # max_draft_len + 1: the drafter advances max_draft_len times, and the + # golden token the target always accepts adds one. Verified against + # Eagle3 dynamic tree (K=6, T=60), MTP dynamic tree, PARD (T=2K-1) and + # the linear modes -- none exceed it. + # + # The exception is the deprecated eagle_choices static tree. There the + # one-model drafter ignores the tree and runs _forward_draft_loop, a + # linear loop over runtime_draft_len == max_total_draft_tokens, so a + # step can accept up to max_total_draft_tokens + 1 tokens even though + # max_draft_len only describes the depth of a tree that is never built. + # (Tree-aware acceptance lives in TorchSampler, i.e. the two-model + # path.) get_spec_decoder passes the wire width for that mode; both it + # and this workaround go away with the feature in release 1.4. + self.max_accepted_path_len = ( + accepted_path_len if accepted_path_len is not None else args.max_draft_len + 1 + ) self.store = self.Store( - new_tokens=int_tensor((max_new_tokens, seq_slots, self.max_beam_width)), - next_new_tokens=int_tensor((max_tokens, seq_slots, self.max_beam_width)), - next_draft_tokens=int_tensor((seq_slots, draft_tokens_size)), + new_tokens=int_tensor((self.max_accepted_path_len, seq_slots, self.max_beam_width)), + next_new_tokens=int_tensor( + (args.max_total_draft_tokens + 1, seq_slots, self.max_beam_width) + ), + next_draft_tokens=int_tensor((seq_slots, args.max_total_draft_tokens)), new_tokens_lens=int_tensor((seq_slots,)), ) - def _get_max_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: - """ - Calculate max_tokens for storage allocation. - - Override in subclasses if needed. Default: draft_len + 1. - MTP uses args.max_total_draft_tokens + 1 for tree-based speculation. - """ - return draft_len + 1 - - def _get_max_new_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: - """Max depth of accepted token path for new_tokens buffer. - - Defaults to _get_max_tokens (same size as next_new_tokens). - Override when accepted path depth differs from total draft tokens, - e.g. dynamic tree where max_draft_len < max_total_draft_tokens. - """ - return self._get_max_tokens(args, draft_len) - - def _get_draft_tokens_storage_size(self, args: TorchSampler.Args, draft_len: int) -> int: - """ - Calculate storage size for next_draft_tokens tensor. - - Override in subclasses if needed. Default: draft_len. - MTP uses args.max_total_draft_tokens for tree-based speculation. - """ - return draft_len - - def _add_dummy_draft_tokens(self) -> bool: - """ - Whether to add dummy draft tokens for context requests. - - Override in subclasses. Default: True (needed for KV cache preparation). - """ - return True - def _request_common_handling( self, request: LlmRequest, @@ -223,6 +215,13 @@ def update_requests( if req.state == LlmRequestState.GENERATION_COMPLETE: continue num_new_tokens = new_tokens_lens_list[req.py_seq_slot] + # new_tokens is sized to this bound, and add_token indexes a plain + # host-side list, so a violation would otherwise surface as an + # opaque IndexError. + assert num_new_tokens <= self.max_accepted_path_len, ( + f"accepted {num_new_tokens} tokens in one step, but new_tokens is " + f"sized for {self.max_accepted_path_len}" + ) for i in range(num_new_tokens): new_token = add_token(req, new_tokens, beam_idx=beam_idx, step=i) if TorchSampler._handle_stop_criteria( @@ -272,8 +271,8 @@ def sample_async( runtime_draft_len = o_next_draft_tokens.shape[1] # Pad or truncate to match fixed-size store buffers for index_copy_. - # Use actual store buffer dimensions (which may differ from draft_len - # when _get_max_new_tokens is overridden, e.g. dynamic tree mode). + # The worker output width tracks runtime_draft_len, which dynamic draft + # length shrinks below the statically allocated store width. new_tokens_width = self.store.new_tokens.shape[0] next_new_tokens_width = self.store.next_new_tokens.shape[0] draft_tokens_width = self.store.next_draft_tokens.shape[1] @@ -317,9 +316,8 @@ def sample_async( sampler_event = self._record_sampler_event() # Add dummy draft tokens to context requests for KV cache preparation - if self._add_dummy_draft_tokens(): - for request in finished_context_requests: - request.py_draft_tokens = [1] * self.draft_len + for request in finished_context_requests: + request.py_draft_tokens = [1] * self.draft_len return SampleStateSpec( requests=sampling_requests, diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index bffa8833058c..96c36b1856f9 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -17,24 +17,23 @@ from ..pyexecutor.seq_slot_manager import SeqSlotManager from ..speculative.interface import SpecMetadata from .dflash import DFlashSpecMetadata, DFlashWorker -from .draft_target import (DraftTargetOneModelSampler, - DraftTargetOneModelSpecMetadata, +from .draft_target import (DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) from .dspark import DSparkSpecMetadata, DSparkWorker from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, - Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, - Eagle3OneModelWorker, Eagle3ResourceManager, - Eagle3SpecMetadata, MTPEagleWorker) + Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, + Eagle3ResourceManager, Eagle3SpecMetadata, MTPEagleWorker) from .eagle3_dynamic_tree import Eagle3OneModelDynamicTreeWorker from .model_drafter import ModelDrafter -from .mtp import MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker +from .mtp import MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker from .mtp_dynamic_tree import (MTPEagleDynamicTreeResourceManager, MTPEagleDynamicTreeWorker) from .ngram import NGramDrafter, NGramPoolManager from .pard import PARDSpecMetadata, PARDWorker -from .sa_worker import SASampler, SASpecMetadata, SAWorker +from .sa_worker import SASpecMetadata, SAWorker from .save_hidden_state import (SaveHiddenStatesResourceManager, SaveHiddenStatesSpecMetadata) +from .spec_sampler_base import SpecSampler from .suffix_automaton import SuffixAutomatonManager @@ -386,27 +385,33 @@ def get_spec_decoder( sampler_args: TorchSampler.Args, spec_config: "DecodingBaseConfig", ): - if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): - # MTP Eagle one-model now uses the same sampler as Eagle3 one-model. - return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) - if spec_config.spec_dec_mode.is_mtp_vanilla(): - nextn = spec_config.max_draft_len - if getattr(spec_config, "use_dynamic_tree", False): - nextn = spec_config.max_total_draft_tokens - return MTPSampler(sampler_args, nextn=nextn) - if spec_config.spec_dec_mode.is_eagle3( - ) or spec_config.spec_dec_mode.is_mtp_eagle(): - # TorchSampler handles Eagle3 gracefully, by integrating d2t into the sampling process + spec_dec_mode = spec_config.spec_dec_mode + if spec_dec_mode.is_eagle3() or spec_dec_mode.is_mtp_eagle(): + # Two-model path: the target model emits logits, so the general-purpose + # TorchSampler does the actual sampling (and folds in the d2t vocab + # mapping). One-model modes below sample inside the worker kernel. return TorchSampler(sampler_args) - if spec_config.spec_dec_mode.is_eagle3_one_model(): - return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) - if spec_config.spec_dec_mode.is_parallel_draft(): - return MTPSampler(sampler_args, - nextn=spec_config.tokens_per_gen_step - 1) - if spec_config.spec_dec_mode.is_sa(): - return SASampler(sampler_args, max_draft_len=spec_config.max_draft_len) - if spec_config.spec_dec_mode.is_draft_target_one_model(): - return DraftTargetOneModelSampler(sampler_args) + if (spec_dec_mode.is_mtp_eagle_one_model() + or spec_dec_mode.is_mtp_vanilla() + or spec_dec_mode.is_eagle3_one_model() + or spec_dec_mode.is_parallel_draft() or spec_dec_mode.is_sa() + or spec_dec_mode.is_draft_target_one_model()): + # One sampler for every one-model mode: it only moves the worker's + # pre-sampled output around, and its buffer shapes derive from + # sampler_args alone. + # + # WORKAROUND (remove with eagle_choices in release 1.4): the static + # tree is the one mode where a step can accept more than + # max_draft_len + 1 tokens. The one-model drafter never builds the tree + # -- _forward_draft_loop is linear over runtime_draft_len, which for a + # non-linear tree is max_total_draft_tokens -- so max_draft_len only + # describes a tree depth that is never used, and acceptance is bounded + # by the wire width instead. Tree-aware acceptance only exists in the + # two-model TorchSampler path, which is deprecated alongside this. + accepted_path_len = None + if getattr(spec_config, "eagle_choices", None): + accepted_path_len = sampler_args.max_total_draft_tokens + 1 + return SpecSampler(sampler_args, accepted_path_len=accepted_path_len) raise ValueError( f"Unsupported speculative decoding mode: {spec_config.spec_dec_mode}") diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 245cc79ec1f2..132bdec6e0e2 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -104,7 +104,7 @@ def test_super_mtp_ssm_replay_smoke(): Verifies that the full pipeline — transforms, cache manager init with replay buffers, and MTP inference — completes without error. The AD SSM custom ops are not directly - invoked at runtime in this configuration (Eagle3OneModelSampler drives its own forward + invoked at runtime in this configuration (SpecSampler drives its own forward loop); the replay kernel path is covered by test_flashinfer_extend_replay_calls_replay_kernel. Uses mamba_head_dim=64 and ssm_state_size=64 to satisfy FlashInfer constraints on the decode path (which IS called in this config).