From b7b994d5c4e35fe452978cb73dc0b64f5c31dac8 Mon Sep 17 00:00:00 2001 From: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:20:06 -0400 Subject: [PATCH 1/2] [None][chore] Remove dead code in model engine Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 8 +- .../_torch/pyexecutor/cuda_graph_runner.py | 79 ++-- .../_torch/pyexecutor/engine/metadata.py | 8 +- .../pyexecutor/engine/runners/encoder.py | 10 +- .../pyexecutor/engine/runners/no_kv_cache.py | 3 - .../_torch/pyexecutor/model_engine.py | 364 ++---------------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 30 +- tensorrt_llm/_torch/route_capture.py | 4 +- tensorrt_llm/_torch/speculative/interface.py | 9 +- tensorrt_llm/_torch/speculative/utils.py | 3 - .../_torch/executor/engine/test_encoder.py | 1 - .../_torch/executor/engine/test_metadata.py | 50 --- .../executor/engine/test_no_kv_cache.py | 1 - .../_torch/executor/engine/test_runners.py | 2 - .../kv_cache/test_dual_pool_kv_cache.py | 1 - .../executor/test_distributed_warmup_oom.py | 1 - .../executor/test_pytorch_model_engine.py | 15 +- .../test_pytorch_model_engine_warmup.py | 3 - .../_torch/executor/test_seq_slot_sizing.py | 1 - tests/unittest/_torch/helpers.py | 1 - tests/unittest/_torch/test_route_capture.py | 26 +- 21 files changed, 92 insertions(+), 528 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 660ffe11b601..1c222c791e1f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1768,8 +1768,7 @@ def _create_kv_cache_manager( layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, disable_overlap_scheduler=self._disable_overlap_scheduler, - kv_events_config=None - if estimating_kv_cache or model_engine.is_draft_model else + kv_events_config=None if estimating_kv_cache else self._llm_args.kv_cache_config.kv_events_config, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, @@ -2754,7 +2753,7 @@ def _create_kv_cache_manager( # Optional overrides for one-model draft case (when model_engine is None) model_config: Optional[ModelConfig] = None, dtype: Optional[torch.dtype] = None, - is_draft: Optional[bool] = None, + is_draft: bool = False, layer_mask: Optional[List[bool]] = None, num_layers: Optional[int] = None, num_kv_heads: Optional[Union[int, List[int]]] = None, @@ -2802,9 +2801,6 @@ def _create_kv_cache_manager( if dtype is None: dtype = model_engine.dtype - if is_draft is None: - is_draft = model_engine.is_draft_model - if kv_cache_type is None: kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 29c5640664b9..4614477976e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -18,7 +18,6 @@ from ..memory_buffer_utils import Buffers, get_memory_buffers from ..modules.multi_stream_utils import with_multi_stream from ..moe.expert_statistic import ExpertStatistic -from ..speculative.eagle3 import Eagle3ResourceManager from ..speculative.interface import SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_draft_kv_cache_manager @@ -127,7 +126,6 @@ class CUDAGraphRunnerConfig: use_mrope: bool original_max_draft_len: int original_max_total_draft_tokens: int - is_draft_model: bool enable_attention_dp: bool is_encoder_decoder: bool batch_size: int @@ -379,52 +377,37 @@ def get_graph_key( else self._resolve_sample_type( batch, promoted_context_request_ids)) - if self.config.is_draft_model and spec_resource_manager is not None and isinstance( - spec_resource_manager, Eagle3ResourceManager): - # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'. - # Because we will pad the input to 'max_draft_len' length for the first draft layer. - draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0 - key = KeyType(batch_size=batch_size, - draft_len=draft_len, - is_first_draft=spec_resource_manager.is_first_draft, - short_seq_len_mode=short_seq_len_mode, - is_all_greedy_sample=is_all_greedy_sample, - sample_type=sample_type, - peft_cache_data_type=peft_cache_data_type, - use_lora_graph=use_lora_graph) - else: - # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, - # so we need to get the draft length from the batch instead of using enable_spec_decode. - draft_len_list = [] - for request in batch.generation_requests: - draft_len_list.append(len(request.py_draft_tokens)) - draft_len = max(draft_len_list) - assert len( - set(draft_len_list)) == 1, "All draft lengths must be the same" - context_requests = batch.context_requests - num_contexts = len(context_requests) - context_query_len = 0 - if num_contexts: - context_query_len = int(context_requests[0].context_chunk_size) - if any( - int(request.context_chunk_size) != context_query_len - for request in context_requests[1:]): - return None - num_encoder_tokens = sum( - int(request.encoder_output_len) for request in context_requests - if not request.py_skip_cross_kv_projection) - key = KeyType(batch_size=batch_size, - draft_len=draft_len, - is_first_draft=False, - short_seq_len_mode=short_seq_len_mode, - is_all_greedy_sample=is_all_greedy_sample, - sample_type=sample_type, - num_contexts=num_contexts, - context_query_len=context_query_len, - num_encoder_tokens=num_encoder_tokens, - peft_cache_data_type=peft_cache_data_type, - use_lora_graph=use_lora_graph) - return key + # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, + # so we need to get the draft length from the batch instead of using enable_spec_decode. + draft_len_list = [] + for request in batch.generation_requests: + draft_len_list.append(len(request.py_draft_tokens)) + draft_len = max(draft_len_list) + assert len( + set(draft_len_list)) == 1, "All draft lengths must be the same" + context_requests = batch.context_requests + num_contexts = len(context_requests) + context_query_len = 0 + if num_contexts: + context_query_len = int(context_requests[0].context_chunk_size) + if any( + int(request.context_chunk_size) != context_query_len + for request in context_requests[1:]): + return None + num_encoder_tokens = sum( + int(request.encoder_output_len) for request in context_requests + if not request.py_skip_cross_kv_projection) + return KeyType(batch_size=batch_size, + draft_len=draft_len, + is_first_draft=False, + short_seq_len_mode=short_seq_len_mode, + is_all_greedy_sample=is_all_greedy_sample, + sample_type=sample_type, + num_contexts=num_contexts, + context_query_len=context_query_len, + num_encoder_tokens=num_encoder_tokens, + peft_cache_data_type=peft_cache_data_type, + use_lora_graph=use_lora_graph) def set_capture_sample_type(self, sample_type: Optional[SampleType]) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index 088db863fbfe..d5abb50e1cdf 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -99,7 +99,6 @@ def update_spec_metadata( *, runtime_draft_len: int, runtime_tokens_per_gen_step: int, - is_draft_model: bool, attention_backend: type[AttentionBackend], original_max_draft_len: int, original_max_total_draft_tokens: int, @@ -109,10 +108,7 @@ def update_spec_metadata( spec_metadata.runtime_draft_len = runtime_draft_len spec_metadata.runtime_tokens_per_gen_step = runtime_tokens_per_gen_step - is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode( - is_draft_model, - attention_backend, - ) + is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode(attention_backend) # Parallel-draft modes advertise their full generation width rather than a # conventional draft length, so attention needs the total-token capacity. if spec_metadata.spec_dec_mode.is_parallel_draft(): @@ -122,7 +118,7 @@ def update_spec_metadata( max_draft_len = original_max_draft_len max_total_draft_tokens = spec_dec_max_total_draft_tokens - if spec_tree_manager is not None and not is_draft_model: + if spec_tree_manager is not None: spec_tree_manager.slot_storage.fill_all_slot_ids( scheduled_requests.context_requests, scheduled_requests.generation_requests, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py index 67fd8d935ec3..ae436e4503af 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py @@ -80,7 +80,6 @@ def create( attention_runtime_features: AttentionRuntimeFeatures, enable_autotuner: bool, is_encoder_decoder: bool, - draft_model: bool, ) -> Self: """Resolve encoder settings and construct the concrete runner config.""" batch_sizes = list(graph_config.batch_sizes or []) if graph_config is not None else [] @@ -107,12 +106,7 @@ def create( feature_shape = None feature_dtype = None fixed_seq_len = None - if ( - graph_config is not None - and not draft_model - and is_encoder_decoder - and model_graph_spec is not None - ): + if graph_config is not None and is_encoder_decoder and model_graph_spec is not None: if mapping.tp_size > 1: logger.warning( "Feature-mode encoder CUDA graphs require TP=1; the encoder phase stays eager." @@ -226,7 +220,6 @@ def create( attention_backend: type[AttentionBackend], attention_runtime_features: AttentionRuntimeFeatures, enable_autotuner: bool, - draft_model: bool, ) -> Self: return super().create( model=model, @@ -241,7 +234,6 @@ def create( attention_runtime_features=attention_runtime_features, enable_autotuner=enable_autotuner, is_encoder_decoder=False, - draft_model=draft_model, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py index 363034b43221..e38688e83d0b 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py @@ -45,7 +45,6 @@ class NoKVCacheRunnerConfig(RunnerConfig): prefill_cuda_graph_num_tokens: list[int] mm_encoder_cache_enabled: bool spec_config: DecodingBaseConfig | None - is_draft_model: bool num_seq_slots: int | None original_max_draft_len: int original_max_total_draft_tokens: int @@ -107,7 +106,6 @@ def setup_spec_metadata( runner_config.max_batch_size, max_num_tokens=runner_config.max_num_tokens, spec_resource_manager=spec_resource_manager, - is_draft_model=runner_config.is_draft_model, max_seq_len=runner_config.max_seq_len, num_seq_slots=runner_config.num_seq_slots, ) @@ -121,7 +119,6 @@ def setup_spec_metadata( runtime_tokens_per_gen_step=spec_config.get_runtime_tokens_per_gen_step( runtime_draft_len ), - is_draft_model=runner_config.is_draft_model, attention_backend=runner_config.attention_backend, original_max_draft_len=runner_config.original_max_draft_len, original_max_total_draft_tokens=(runner_config.original_max_total_draft_tokens), diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4fc771e1ce64..edca33e053a3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -69,7 +69,6 @@ prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) -from ..speculative.eagle3 import Eagle3ResourceManager from ..speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_static_draft_len, update_draft_len @@ -255,8 +254,7 @@ def forward(self, resource_manager: Optional[ResourceManager], new_tensors_device: Optional[SampleStateTensors], gather_context_logits: bool = False, - cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None): + cache_indirection_buffer: Optional[torch.Tensor] = None): raise NotImplementedError def warmup(self, resource_manager: Optional[ResourceManager]) -> None: @@ -376,7 +374,6 @@ def __init__( attn_runtime_features: Optional[AttentionRuntimeFeatures] = None, dist: Optional[Distributed] = None, spec_config: Optional[DecodingBaseConfig] = None, - is_draft_model: bool = False, model: Optional[torch.nn.Module] = None, checkpoint_loader: Optional[BaseCheckpointLoader] = None, model_weights_memory_tag: Optional[str] = None, @@ -475,16 +472,10 @@ def __init__( else: self.max_draft_loop_tokens = self.original_max_total_draft_tokens - # The draft model won't have any draft tokens attached to - # generation requests when we invoke it autoregressively - if spec_config is not None and is_draft_model: - spec_config.max_draft_len = 0 - spec_config.max_total_draft_tokens = 0 self.spec_config = spec_config self.is_spec_decode = spec_config is not None - self.sparse_attention_config = None if is_draft_model else llm_args.sparse_attention_config + self.sparse_attention_config = llm_args.sparse_attention_config self.enable_spec_decode = self.is_spec_decode - self.is_draft_model = is_draft_model self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures( ) @@ -503,8 +494,7 @@ def __init__( self.moe_load_balancer: Optional[MoeLoadBalancer] = None self.model_loader: Optional[ModelLoader] = None if model is None: - lora_config: Optional[ - LoraConfig] = None if is_draft_model else llm_args.lora_config + lora_config: Optional[LoraConfig] = llm_args.lora_config # Keep the model_loader to support reloading the model weights later self.model_loader = ModelLoader( llm_args=llm_args, @@ -572,8 +562,7 @@ def __init__( model_engine=self, enabled=self.llm_args.enable_return_routed_experts, pp_size=self.mapping.pp_size, - is_spec_decode=self.is_spec_decode, - is_draft_model=self.is_draft_model) + is_spec_decode=self.is_spec_decode) if self.route_capture is not None: self.model.extra_attrs[ROUTE_CAPTURE_ATTR] = self.route_capture # Every MM item-scheduling decision -- policy, capability, feature @@ -754,9 +743,7 @@ def __init__( self.spec_metadata = None if self.is_spec_decode: - if not self.is_draft_model: - update_spec_config_from_loaded_model(self.spec_config, - self.model) + update_spec_config_from_loaded_model(self.spec_config, self.model) max_num_draft_tokens = self.max_draft_loop_tokens * self.batch_size self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, @@ -841,49 +828,10 @@ def __init__( self.mrope_position_ids_cuda = torch.empty( (3, 1, self.max_num_tokens), dtype=torch.int, device='cuda') - # Pre-allocated buffers for draft model to avoid implicit synchronization - # These are used to build index tensors without creating tensors from Python lists - max_first_draft_tokens = self.batch_size * ( - self.original_max_total_draft_tokens + - 1) if spec_config else self.batch_size - tokens_per_draft = self.original_max_total_draft_tokens + 1 - self.idx_accepted_tokens_cache = None - self.draft_token_positions_cache = None - if spec_config: - # Cache for idx_accepted_tokens (pattern: 0,0,0...1,1,1...2,2,2...) - self.idx_accepted_tokens_cache = torch.arange( - max_first_draft_tokens, dtype=torch.long, - device='cuda') // tokens_per_draft - - if self.is_draft_model: - self.draft_ctx_token_indices_cuda = torch.empty((self.batch_size, ), - dtype=torch.long, - device='cuda') - self.draft_ctx_seq_slots_cuda = torch.empty((self.batch_size, ), - dtype=torch.long, - device='cuda') - # Buffers for first_draft requests (max_draft_len+1 tokens per request) - self.draft_first_draft_indices_cuda = torch.empty( - (max_first_draft_tokens, ), dtype=torch.long, device='cuda') - self.draft_first_draft_seq_slots_cuda = torch.empty( - (max_first_draft_tokens, ), dtype=torch.long, device='cuda') - # Buffers for seq_slots and request indices - self.draft_seq_slots_buffer_cuda = torch.empty((self.batch_size, ), - dtype=torch.int, - device='cuda') - self.draft_request_indices_buffer_cuda = torch.empty( - (self.batch_size, ), dtype=torch.int, device='cuda') - - # Pre-computed constant tensors for incremental update optimization - # Cache for token_positions (pattern: 0,1,2...N repeated) - self.draft_token_positions_cache = torch.arange(tokens_per_draft, - dtype=torch.long, - device='cuda') - # We look up this key in resource_manager during forward to find the # kv cache manager. Can be changed to support multiple model engines # with different KV cache managers. - self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER + self.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False self._warmup_timer = _WarmupTimer(self.mapping.rank) @@ -948,7 +896,6 @@ def _initialize_cuda_graph_runner(self) -> Optional[CUDAGraphRunner]: original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=self. original_max_total_draft_tokens, - is_draft_model=self.is_draft_model, enable_attention_dp=self.enable_attention_dp, is_encoder_decoder=is_encoder_decoder, batch_size=self.batch_size, @@ -1005,7 +952,6 @@ def _initialize_encoder_runner( attention_backend=self.attn_backend, attention_runtime_features=self.attn_runtime_features, enable_autotuner=self.llm_args.enable_autotuner, - draft_model=self.is_draft_model, ) return runner_cls( self.model, @@ -1029,7 +975,6 @@ def _initialize_encoder_decoder_runner( attention_runtime_features=self.attn_runtime_features, enable_autotuner=self.llm_args.enable_autotuner, is_encoder_decoder=True, - draft_model=self.is_draft_model, ) runner = runner_cls( self.model, @@ -1058,7 +1003,6 @@ def _initialize_no_kv_cache_runner( attention_runtime_features=self.attn_runtime_features, mm_encoder_cache_enabled=self._mm_encoder_cache_enabled, spec_config=self.spec_config, - is_draft_model=self.is_draft_model, num_seq_slots=self.max_num_seq_slots, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=( @@ -1534,8 +1478,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, # this predicate is not rank-uniform on its own. Agree it before it # gates either the attention or the general phase. can_run_general_warmup = self._agree_warmup_flag( - not is_enc_dec and not self.is_draft_model - and not self.mapping.has_cp_helix() and self.guided_decoder is None + not is_enc_dec and not self.mapping.has_cp_helix() + and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager)) log_mem_snapshot("warmup/before_warmup") @@ -2136,8 +2080,8 @@ def trtllm_gen_fmha_jit_warmup(): logger.info("Running TRTLLM-Gen FMHA JIT warmup") warmup_requests_configs = [] - if not self.is_draft_model and self.guided_decoder is None: - # doesn't support 2-model speculative draft and guided decoding + if self.guided_decoder is None: + # doesn't support guided decoding warmup_requests_configs.append( (1 + self.max_total_draft_tokens, 1)) # one generation request else: @@ -2159,8 +2103,7 @@ def trtllm_gen_fmha_jit_warmup(): f"{_KIMI_KDA_PREFILL_WARMUP_TOKENS} context tokens") warmup_requests_configs.append((_KIMI_KDA_PREFILL_WARMUP_TOKENS, 0)) - if (not self.is_draft_model and self.guided_decoder is None - and can_run_general_warmup): + if self.guided_decoder is None and can_run_general_warmup: # The cute_dsl_mla FMHA lib now only support the generation-only batch, we need to warmup the TRTLLM-Gen FMHA lib for the mixed context+generation batch. # One MIXED context+generation batch (1 ctx token + 1 gen request). warmup_requests_configs.append( @@ -2311,8 +2254,7 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: max_num_draft_tokens=self.original_max_draft_len) warmup_configs = [(curr_max_num_tokens, 0)] - if (not self.is_draft_model and self.guided_decoder is None - and not self.mapping.has_pp()): + if self.guided_decoder is None and not self.mapping.has_pp(): # Add generation request to warmup the autotuner cache. warmup_configs.append((1 + self.max_total_draft_tokens, 1)) @@ -2331,14 +2273,6 @@ def run_autotuner_pass(autotune_context: Any, f"autotuner, num_tokens={num_tokens}, " f"num_gen_requests={num_gen_requests}"): continue - # Reset the flag is_first_draft for the draft model. - # This is necessary for overlap scheduler. - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance( - spec_resource_manager, Eagle3ResourceManager): - spec_resource_manager.is_first_draft = True - with self._warmup_timer.phase( f"autotuner shape num_tokens={num_tokens}, " f"num_gen_requests={num_gen_requests}", @@ -2538,13 +2472,6 @@ def _run_mamba_hybrid_warmup(self, if not self._should_run_warmup_batch( batch, num_tokens_i, shape): continue - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance( - spec_resource_manager, - Eagle3ResourceManager): - spec_resource_manager.is_first_draft = True - self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) @@ -2617,13 +2544,7 @@ def _get_graphs_to_capture( Returns: List of (batch_size, draft_len) tuples for CUDA graph capture. """ - # Case 1: Draft model (two-model speculative decoding) - # Two-model path is deprecated and will be removed in the near future - if self.is_draft_model: - draft_len = self.max_total_draft_tokens - return [(bs, draft_len) for bs in cuda_graph_batch_sizes] - - # Case 2: One-model with dynamic draft length + # Case 1: One-model with dynamic draft length if self.spec_config is not None and self.spec_config.draft_len_schedule is not None and self.spec_config.spec_dec_mode.support_dynamic_draft_len( ): graphs = [(graph_bs, draft_len) for graph_bs, draft_len in @@ -2651,7 +2572,7 @@ def _get_graphs_to_capture( f"Capturing {len(graphs)} graphs: {graphs}") return graphs - # Case 3: Target model (two-model) or one-model without dynamic draft + # Case 2: One-model without dynamic draft # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. # spec_config is None for non-spec models — fall back to max_draft_len (= 0). @@ -2850,11 +2771,9 @@ def _run_capture_pass(force_non_greedy: bool, f"Run generation-only CUDA graph {operation} ({label}) " f"for batch size={bs}, draft_len={draft_len}, " f"max_seq_len={max_seq_len}") - self.enable_spec_decode = draft_len > 0 or self.is_draft_model or ( + self.enable_spec_decode = draft_len > 0 or ( self.spec_config is not None and self.spec_config.spec_dec_mode.use_one_engine()) - self._update_draft_inference_state_for_warmup( - batch, draft_len > 0, resource_manager) if self._is_encoder_decoder_model(): prepare_cross_batch(batch, resource_manager) self.forward(batch, @@ -3609,20 +3528,6 @@ def _get_enc_dec_hidden_size(self) -> int: "hidden size from the model config.") return int(hidden_size) - def _update_draft_inference_state_for_warmup( - self, batch: ScheduledRequests, is_first_draft: bool, - resource_manager: ResourceManager): - """Updates request states for specific draft model warmups like Eagle3.""" - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance(spec_resource_manager, - Eagle3ResourceManager): - spec_resource_manager.is_first_draft = is_first_draft - if is_first_draft: - for req in batch.generation_requests: - req.py_is_first_draft = True - req.py_draft_tokens = [] - def _set_up_attn_metadata( self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], @@ -3702,7 +3607,6 @@ def _set_up_spec_metadata( self.batch_size, max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, - is_draft_model=self.is_draft_model, max_seq_len=self.max_seq_len, num_seq_slots=self.max_num_seq_slots) return self.spec_metadata @@ -4307,7 +4211,7 @@ def _can_use_encoder_decoder_input_fast_path( static_eligible = ( hasattr(batch_manager_bindings, "prepare_encoder_decoder_inputs") - and self._is_encoder_decoder_model() and not self.is_draft_model + and self._is_encoder_decoder_model() and self.max_beam_width == 1 and self.sparse_attention_config is None and not self.use_mrope and not self.enable_attention_dp @@ -4743,8 +4647,6 @@ def _prepare_tp_inputs( spec_metadata: Optional[SpecMetadata] = None, new_tensors_device: Optional[SampleStateTensors] = None, cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, resource_manager: Optional[ResourceManager] = None, maybe_graph: bool = False, promoted_context_request_ids: frozenset[int] = frozenset(), @@ -4760,7 +4662,7 @@ def _prepare_tp_inputs( new_tokens_device = new_tensors_device.new_tokens # When using overlap scheduler with speculative decoding, the target model's inputs would be SampleStateTensorsSpec. if isinstance(new_tensors_device, SampleStateTensorsSpec): - assert self.enable_spec_decode and not self.is_draft_model + assert self.enable_spec_decode new_tokens_lens_device = new_tensors_device.new_tokens_lens # [batch] next_draft_tokens_device = new_tensors_device.next_draft_tokens # [batch, draft_len] @@ -4833,18 +4735,6 @@ def _prepare_tp_inputs( cross_encoder_seq_lens: List[int] = [ ] # new encoder K/V tokens per decoder sequence cross_encoder_cached_tokens_per_seq: List[int] = [] - # Variables for updating the inputs of draft model - # Base values for gather_ids computation - first_draft_base_gather_ids = [] - # seq_slots to index into num_accepted_tokens_device - first_draft_seq_slots = [] - # Indices in the num_accepted_draft_tokens list - first_draft_request_indices = [] - - # (start_idx, end_idx, seq_slot) for context requests - context_input_ids_positions = [] - # (start_idx, end_idx, seq_slot) for first_draft requests - first_draft_input_ids_positions = [] context_prompt_lookahead = None if (spec_metadata is not None @@ -4905,17 +4795,7 @@ def append_cross_attention_state(request: LlmRequest, # that rewrite token IDs in place write into the request's own span # rather than assuming a contiguous multimodal prefix. context_start_idx = len(input_ids) - # Track position for updating the inputs of draft model - if self.is_draft_model and num_accepted_tokens_device is not None: - input_ids.extend(prompt_tokens) - end_idx = len(input_ids) - slot_idx = req_id_to_old_request[ - request.py_request_id].py_seq_slot - context_input_ids_positions.append( - (context_start_idx, end_idx - 1, - slot_idx)) # end_idx-1 is the last token position - else: - input_ids.extend(prompt_tokens) + input_ids.extend(prompt_tokens) gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) @@ -5268,37 +5148,12 @@ def _helix_pack_extend(request, group: int) -> int: end_compute) position_ids.extend( range(begin_compute, begin_compute + len(prompt_tokens))) - - # Track position for updating the inputs of draft model - if self.is_draft_model and num_accepted_tokens_device is not None: - start_idx = len(input_ids) - input_ids.extend(prompt_tokens) - end_idx = len(input_ids) - # For first_draft, we need to replace the last original_max_draft_len+1 tokens - slot_idx = req_id_to_old_request[ - request.py_request_id].py_seq_slot - first_draft_input_ids_positions.append( - (start_idx, end_idx, slot_idx)) - - # Store info for GPU computation of gather_ids and num_accepted_draft_tokens - base_gather_id = len( - input_ids) - 1 - self.original_max_draft_len - # Placeholder, will be corrected on GPU - gather_ids.append(base_gather_id) - first_draft_base_gather_ids.append(base_gather_id) - first_draft_seq_slots.append(slot_idx) - first_draft_request_indices.append( - len(num_accepted_draft_tokens)) - - # Placeholder, will be corrected on GPU - num_accepted_draft_tokens.append(0) - else: - input_ids.extend(prompt_tokens) - gather_ids.append( - len(input_ids) - 1 - (self.original_max_draft_len - - request.py_num_accepted_draft_tokens)) - num_accepted_draft_tokens.append( - request.py_num_accepted_draft_tokens) + input_ids.extend(prompt_tokens) + gather_ids.append( + len(input_ids) - 1 - (self.original_max_draft_len - + request.py_num_accepted_draft_tokens)) + num_accepted_draft_tokens.append( + request.py_num_accepted_draft_tokens) sequence_lengths.append(1 + self.original_max_draft_len) prompt_lengths.append(request.py_prompt_len) @@ -5388,17 +5243,7 @@ def _helix_pack_extend(request, group: int) -> int: # can be aligned to the correct positions. if not request.is_cuda_graph_dummy: for beam in range(beam_width): - # Track position for GPU update (draft model only) - if self.is_draft_model and num_accepted_tokens_device is not None: - start_idx = len(input_ids) - input_ids.append(request.get_last_tokens(beam)) - end_idx = len(input_ids) - slot_idx = req_id_to_old_request[ - request.py_request_id].py_seq_slot - first_draft_input_ids_positions.append( - (start_idx, end_idx, slot_idx)) - else: - input_ids.append(request.get_last_tokens(beam)) + input_ids.append(request.get_last_tokens(beam)) past_seen_token_num = request.max_beam_num_tokens - 1 request_has_previous_tensor = False else: @@ -5573,80 +5418,6 @@ def previous_seq_slots_device(): pin_memory=prefer_pinned()) self.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True) - # Update input_ids_cuda with new tokens from new_tensors_device (draft model only) - if self.is_draft_model and num_accepted_tokens_device is not None: - # For context requests: replace the last token with new_tensors_device[0, seq_slot, 0] - if len(context_input_ids_positions) > 0: - # Build tensors on CPU first, then copy to GPU to avoid implicit sync - num_ctx_positions = len(context_input_ids_positions) - ctx_token_indices_cpu = torch.tensor( - [ - last_token_idx for _, last_token_idx, _ in - context_input_ids_positions - ], - dtype=torch.long, - pin_memory=prefer_pinned()) - ctx_seq_slots_cpu = torch.tensor([ - seq_slot - for _, _, seq_slot in context_input_ids_positions - ], - dtype=torch.long, - pin_memory=prefer_pinned()) - # Copy to pre-allocated GPU buffers - self.draft_ctx_token_indices_cuda[:num_ctx_positions].copy_( - ctx_token_indices_cpu, non_blocking=True) - self.draft_ctx_seq_slots_cuda[:num_ctx_positions].copy_( - ctx_seq_slots_cpu, non_blocking=True) - self.input_ids_cuda[ - self. - draft_ctx_token_indices_cuda[:num_ctx_positions]] = new_tensors_device.new_tokens[ - 0, - self.draft_ctx_seq_slots_cuda[:num_ctx_positions], - 0] - - # For first_draft requests: replace the last (original_max_draft_len+1) tokens - # with new_tensors_device[:, seq_slot, 0] - if len(first_draft_input_ids_positions) > 0: - # All first_draft requests have same token length (original_max_draft_len + 1) - # Build index tensors on CPU first, then copy to GPU to avoid implicit sync - num_requests = len(first_draft_input_ids_positions) - tokens_per_request = first_draft_input_ids_positions[0][ - 1] - first_draft_input_ids_positions[0][0] - - # Create flat index array for all tokens to update on CPU - all_indices = [] - all_seq_slots = [] - for start_idx, end_idx, seq_slot in first_draft_input_ids_positions: - all_indices.extend(range(start_idx, end_idx)) - all_seq_slots.extend([seq_slot] * (end_idx - start_idx)) - - # Create CPU tensors with pinned memory - total_tokens = len(all_indices) - idx_tensor_cpu = torch.tensor(all_indices, - dtype=torch.long, - pin_memory=prefer_pinned()) - seq_slots_tensor_cpu = torch.tensor( - all_seq_slots, - dtype=torch.long, - pin_memory=prefer_pinned()) - - # Copy to pre-allocated GPU buffers - self.draft_first_draft_indices_cuda[:total_tokens].copy_( - idx_tensor_cpu, non_blocking=True) - self.draft_first_draft_seq_slots_cuda[:total_tokens].copy_( - seq_slots_tensor_cpu, non_blocking=True) - - # Create token position indices (repeating 0..tokens_per_request for each request) - token_positions = torch.arange( - tokens_per_request, dtype=torch.long, - device='cuda').repeat(num_requests) - - self.input_ids_cuda[ - self. - draft_first_draft_indices_cuda[:total_tokens]] = new_tensors_device.new_tokens[ - token_positions, self. - draft_first_draft_seq_slots_cuda[:total_tokens], 0] - if num_draft_tokens > 0: draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, @@ -5660,34 +5431,6 @@ def previous_seq_slots_device(): self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)].copy_(num_accepted_draft_tokens, non_blocking=True) - - # Update num_accepted_draft_tokens_cuda for first_draft_requests directly from num_accepted_tokens_device (draft model only) - if self.is_draft_model and len(first_draft_seq_slots) > 0: - # Build tensors on CPU first, then copy to GPU to avoid implicit sync - num_first_draft = len(first_draft_seq_slots) - first_draft_seq_slots_cpu = torch.tensor( - first_draft_seq_slots, - dtype=torch.int, - pin_memory=prefer_pinned()) - first_draft_indices_cpu = torch.tensor( - first_draft_request_indices, - dtype=torch.int, - pin_memory=prefer_pinned()) - - # Copy to pre-allocated GPU buffers - self.draft_seq_slots_buffer_cuda[:num_first_draft].copy_( - first_draft_seq_slots_cpu, non_blocking=True) - self.draft_request_indices_buffer_cuda[:num_first_draft].copy_( - first_draft_indices_cpu, non_blocking=True) - - # Extract accepted tokens for first_draft requests from device tensor - accepted_tokens = num_accepted_tokens_device[ - self.draft_seq_slots_buffer_cuda[:num_first_draft]] - # Update the correct positions in num_accepted_draft_tokens_cuda - self.num_accepted_draft_tokens_cuda[ - self. - draft_request_indices_buffer_cuda[: - num_first_draft]] = accepted_tokens if next_draft_tokens_device is not None: # Initialize these two values to zeros self.previous_pos_id_offsets_cuda *= 0 @@ -5833,35 +5576,6 @@ def previous_seq_slots_device(): gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), non_blocking=True) - # Update gather_ids for first_draft_requests on GPU (draft model only) - if self.is_draft_model and len(first_draft_seq_slots) > 0: - # Build tensors on CPU first, then copy to GPU to avoid implicit sync - num_first_draft = len(first_draft_seq_slots) - first_draft_seq_slots_cpu = torch.tensor( - first_draft_seq_slots, - dtype=torch.int, - pin_memory=prefer_pinned()) - first_draft_indices_cpu = torch.tensor( - first_draft_request_indices, - dtype=torch.int, - pin_memory=prefer_pinned()) - - # Copy to pre-allocated GPU buffers - self.draft_seq_slots_buffer_cuda[:num_first_draft].copy_( - first_draft_seq_slots_cpu, non_blocking=True) - self.draft_request_indices_buffer_cuda[:num_first_draft].copy_( - first_draft_indices_cpu, non_blocking=True) - - # Extract accepted tokens for first_draft requests from device tensor - accepted_tokens = num_accepted_tokens_device[ - self.draft_seq_slots_buffer_cuda[:num_first_draft]] - # Update gather_ids: gather_id = base_gather_id + num_accepted_tokens - # (since gather_id = len(input_ids) - 1 - (max_draft_len - num_accepted)) - self.gather_ids_cuda[ - self. - draft_request_indices_buffer_cuda[: - num_first_draft]] += accepted_tokens - if self.mapping.has_cp_helix(): # A non-None owned-count list is what arms # _helix_spec_tokens_valid, and the per-token slots/bounds that @@ -6115,8 +5829,8 @@ def previous_seq_slots_device(): # guards stable), with all three axes equal to the scalar # positions, so the fast path advances that buffer in place and # returns the same layout (see _apply_steady_gen_fast_prepare). - if (self.spec_config is None and not self.is_draft_model - and spec_metadata is None and new_tokens_device is not None + if (self.spec_config is None and spec_metadata is None + and new_tokens_device is not None and self.guided_decoder is None and not self.enable_attention_dp and not mrope_position_ids and not mrope_delta_write_seq_slots @@ -6160,8 +5874,6 @@ def _prepare_inputs( spec_metadata: Optional[SpecMetadata] = None, new_tensors_device: Optional[SampleStateTensors] = None, cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, resource_manager: Optional[ResourceManager] = None, maybe_graph: bool = False, promoted_context_request_ids: frozenset[int] = frozenset(), @@ -6206,8 +5918,6 @@ def _prepare_inputs( spec_metadata, new_tensors_device, cache_indirection_buffer, - num_accepted_tokens_device, - req_id_to_old_request, resource_manager, maybe_graph, promoted_context_request_ids, @@ -6220,9 +5930,7 @@ def forward(self, resource_manager: Optional[ResourceManager] = None, new_tensors_device: Optional[SampleStateTensors] = None, gather_context_logits: bool = False, - cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None): + cache_indirection_buffer: Optional[torch.Tensor] = None): if isinstance(batch, PackedEncoderBatch): assert self._is_packed_runner, ( "a packed batch requires a packed-batch runner") @@ -6236,18 +5944,13 @@ def forward(self, new_tensors_device=new_tensors_device, gather_context_logits=gather_context_logits, cache_indirection_buffer=cache_indirection_buffer, - num_accepted_tokens_device=num_accepted_tokens_device, - req_id_to_old_request=req_id_to_old_request, ) def _forward_scheduled(self, scheduled_requests: ScheduledRequests, resource_manager: ResourceManager, *, new_tensors_device: Optional[SampleStateTensors], gather_context_logits: bool, - cache_indirection_buffer: Optional[torch.Tensor], - num_accepted_tokens_device: Optional[torch.Tensor], - req_id_to_old_request: Optional[Dict[int, - LlmRequest]]): + cache_indirection_buffer: Optional[torch.Tensor]): assert not self._is_packed_runner, ( "a packed-batch runner cannot execute scheduled requests") kv_cache_manager = resource_manager.get_resource_manager( @@ -6291,7 +5994,6 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, runtime_tokens_per_gen_step=( self.get_runtime_tokens_per_gen_step( self.runtime_draft_len)), - is_draft_model=self.is_draft_model, attention_backend=self.attn_backend, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=( @@ -6309,8 +6011,7 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, # total tree width after graph selection. Only linear-tree zero-draft # iterations can therefore safely reuse a zero-draft graph. can_promote_spec_decode = (not self.enable_spec_decode - or (not self.is_draft_model - and self.runtime_draft_len == 0 + or (self.runtime_draft_len == 0 and self.spec_config is not None and self.spec_config.is_linear_tree)) # TODO: Generalize these conservative gates as actual-draft, beam, and @@ -6407,8 +6108,7 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, staged_sample_type) # Fill slot-ID buffer for scatter inside draft loop - if (self.enable_spec_decode and spec_tree_manager is not None - and not self.is_draft_model): + if self.enable_spec_decode and spec_tree_manager is not None: spec_tree_manager.slot_storage.fill_all_slot_ids( execution_requests.context_requests, execution_requests.generation_requests, @@ -6421,8 +6121,6 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, spec_metadata, new_tensors_device, cache_indirection_buffer, - num_accepted_tokens_device, - req_id_to_old_request, resource_manager, can_run_graph, execution_promoted_context_ids, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fc4c549f63ef..4c46b24ee23b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5238,16 +5238,14 @@ def _executor_loop_overlap(self): # When there's any accepted tokens, we can't directly use the previous batch's outputs in this iteration for the target model, # so we'll set the target model's input to None and skip updating the target requests after target model forward. use_previous_draft_tokens = self.has_previous_draft_tokens - num_accepted_tokens_device = None target_inputs = None - num_accepted_tokens_device = None if has_draft_batch: self.execution_stream.wait_stream( torch.cuda.current_stream()) with torch.cuda.stream(self.execution_stream): - target_inputs, num_accepted_tokens_device = self._handle_speculative_decoding( + target_inputs = self._handle_speculative_decoding( scheduled_batch, previous_tensors, previous_tensors_device) torch.cuda.current_stream().wait_stream( @@ -5281,8 +5279,7 @@ def _executor_loop_overlap(self): gpu_forward_start, gpu_forward_end) as fwd_timing: with self._step_scope(scheduled_batch): batch_outputs = self._forward_step( - scheduled_batch, previous_tensors_device, - num_accepted_tokens_device) + scheduled_batch, previous_tensors_device) self._maybe_prefetch_next_iter_mm_encoders(scheduled_batch) @@ -7499,11 +7496,9 @@ def _maybe_prefetch_next_iter_mm_encoders( f"Cross-iter MM encoder prefetch failed; falling back to " f"in-iter encode.\n{traceback.format_exc()}") - def _forward_step( - self, - scheduled_requests: ScheduledRequests, - new_tensors_device: Optional[SampleStateTensors] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None): + def _forward_step(self, + scheduled_requests: ScheduledRequests, + new_tensors_device: Optional[SampleStateTensors] = None): self._maybe_record_hang_diagnostic_phase( "forward_call", scheduled_requests, @@ -7541,15 +7536,13 @@ def _forward_step( f"[Executor] _forward_step {self.iter_counter}: {scheduled_requests.num_context_requests} ctx reqs, {num_ctx_tokens} ctx tokens, {scheduled_requests.num_generation_requests} gen reqs" ) def forward(scheduled_requests, resource_manager, new_tensors_device, - gather_context_logits, cache_indirection_buffer, - num_accepted_tokens_device): + gather_context_logits, cache_indirection_buffer): return self.model_engine.forward( scheduled_requests, resource_manager, new_tensors_device, gather_context_logits=gather_context_logits, - cache_indirection_buffer=cache_indirection_buffer, - num_accepted_tokens_device=num_accepted_tokens_device) + cache_indirection_buffer=cache_indirection_buffer) try: gather_context_logits = any( @@ -7564,8 +7557,7 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, with torch.cuda.stream(self.execution_stream): outputs = forward(scheduled_requests, self.resource_manager, new_tensors_device, gather_context_logits, - cache_indirection_buffer, - num_accepted_tokens_device) + cache_indirection_buffer) self._maybe_record_hang_diagnostic_phase( "forward_returned", scheduled_requests, @@ -8509,8 +8501,8 @@ def _remove_inflight_ids(self, scheduled_requests: ScheduledRequests): scheduled_requests.added_inflight_req_ids = [] def _handle_speculative_decoding( - self, scheduled_batch, previous_tensors, target_inputs - ) -> Tuple[Optional[SampleStateTensorsSpec], Optional[torch.Tensor]]: + self, scheduled_batch, previous_tensors, + target_inputs) -> Optional[SampleStateTensorsSpec]: with request_context(is_draft=self.draft_model_engine is not None, scheduled_requests=scheduled_batch): target_outputs = self.previous_batch.sample_state and self.previous_batch.sample_state.device @@ -8528,7 +8520,7 @@ def _handle_speculative_decoding( # Pad draft tokens to the max draft length for CUDA graph compatibility self.has_previous_draft_tokens = new_target_inputs is not None and new_target_inputs.next_draft_tokens is not None - return new_target_inputs, num_accepted_tokens_device + return new_target_inputs def reset_prefix_cache(self): self.kv_cache_manager.reset_reuse_state() diff --git a/tensorrt_llm/_torch/route_capture.py b/tensorrt_llm/_torch/route_capture.py index 50dd7d9d245a..e7c16b6b371c 100644 --- a/tensorrt_llm/_torch/route_capture.py +++ b/tensorrt_llm/_torch/route_capture.py @@ -196,18 +196,16 @@ def create( enabled: bool, pp_size: int, is_spec_decode: bool, - is_draft_model: bool, ) -> Optional["RouteCapture"]: """Build the engine's capturer, or ``None`` when Router Replay is off. - Draft engines never capture: their MoE layers are not target routing. Fails closed on paths this capture cannot attribute correctly: pipeline parallelism (routes live on the last PP stage only) and speculative decoding / MTP (accepted-token remap not handled). All flags are passed explicitly so this does not depend on engine attribute initialization order. """ - if not enabled or is_draft_model: + if not enabled: return None if pp_size > 1 or is_spec_decode: raise RuntimeError( diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 5a0715f2ce0f..37992e2f8835 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -436,22 +436,21 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]): def attention_need_spec_dec_mode( self, - is_draft_model: bool, attention_backend: Type[AttentionBackend], ): """ If true, the attention backend kernel needs to run in spec-dec mode (multi-token query mode). Args: - is_draft_model: whether the model is a draft model. attention_backend: the attention backend. """ is_trtllm_attention = issubclass(attention_backend, TrtllmAttention) # Always use the multi-token query mode for 1-model if the kernels are available. use_case_1 = self.use_one_engine() - # For 2-model, only the target model (verification) processes multiple tokens at once. - use_case_2 = (not self.use_one_engine() and not is_draft_model - and is_trtllm_attention) + # For modes that do not run in one engine (NGram, user-provided drafts), + # the target model verifies multiple draft tokens per step and needs the + # multi-token query kernel. + use_case_2 = not self.use_one_engine() and is_trtllm_attention return use_case_1 or use_case_2 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 318d8fb14339..053723280738 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -338,7 +338,6 @@ def get_spec_metadata(spec_config, max_num_requests, max_num_tokens, spec_resource_manager=None, - is_draft_model=False, max_seq_len=262144, num_seq_slots=None): metadata = _build_spec_metadata(spec_config, @@ -346,7 +345,6 @@ def get_spec_metadata(spec_config, max_num_requests, max_num_tokens, spec_resource_manager=spec_resource_manager, - is_draft_model=is_draft_model, max_seq_len=max_seq_len) # Set here rather than in each branch below: every one-model mode needs it and # the per-mode constructors are easy to miss one of. @@ -370,7 +368,6 @@ def _build_spec_metadata(spec_config, max_num_requests, max_num_tokens, spec_resource_manager=None, - is_draft_model=False, max_seq_len=262144): use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) diff --git a/tests/unittest/_torch/executor/engine/test_encoder.py b/tests/unittest/_torch/executor/engine/test_encoder.py index f3848d93feba..2a8e8b7c6366 100644 --- a/tests/unittest/_torch/executor/engine/test_encoder.py +++ b/tests/unittest/_torch/executor/engine/test_encoder.py @@ -59,7 +59,6 @@ def encoder_graph_spec(self) -> tuple[tuple[int, ...], torch.dtype, int]: attention_backend=TrtllmAttention, attention_runtime_features=AttentionRuntimeFeatures(), enable_autotuner=False, - draft_model=False, ) if encoder_decoder: kwargs["is_encoder_decoder"] = True diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index 089317ef907c..6ccc97670cf1 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -128,7 +128,6 @@ def test_update_spec_metadata_handles_parallel_draft_and_dynamic_tree() -> None: spec_tree_manager, runtime_draft_len=3, runtime_tokens_per_gen_step=4, - is_draft_model=False, attention_backend=_AttentionBackend, original_max_draft_len=2, original_max_total_draft_tokens=6, @@ -152,52 +151,3 @@ def test_update_spec_metadata_handles_parallel_draft_and_dynamic_tree() -> None: spec_tree_manager=spec_tree_manager, num_contexts=1, ) - - -def test_update_spec_metadata_uses_non_parallel_limits_for_draft_model() -> None: - spec_mode = SimpleNamespace( - attention_need_spec_dec_mode=Mock(return_value=False), - is_parallel_draft=Mock(return_value=False), - ) - spec_metadata = SimpleNamespace( - spec_dec_mode=spec_mode, - is_spec_dec_tree=False, - is_spec_dec_dynamic_tree=False, - ) - scheduled_requests = SimpleNamespace( - batch_size=1, - num_context_requests=0, - context_requests=[], - generation_requests=[object()], - ) - attn_metadata = SimpleNamespace(update_spec_dec_param=Mock()) - spec_tree_manager = SimpleNamespace( - slot_storage=SimpleNamespace(fill_all_slot_ids=Mock()), - ) - - update_spec_metadata( - spec_metadata, - scheduled_requests, - attn_metadata, - spec_tree_manager, - runtime_draft_len=2, - runtime_tokens_per_gen_step=3, - is_draft_model=True, - attention_backend=_AttentionBackend, - original_max_draft_len=4, - original_max_total_draft_tokens=8, - spec_dec_max_total_draft_tokens=6, - ) - - spec_tree_manager.slot_storage.fill_all_slot_ids.assert_not_called() - attn_metadata.update_spec_dec_param.assert_called_once_with( - batch_size=1, - is_spec_decoding_enabled=False, - is_spec_dec_tree=False, - is_spec_dec_dynamic_tree=False, - max_draft_len=4, - max_total_draft_tokens=6, - spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager, - num_contexts=0, - ) diff --git a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py b/tests/unittest/_torch/executor/engine/test_no_kv_cache.py index f6ef18999409..d28e58926287 100644 --- a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py +++ b/tests/unittest/_torch/executor/engine/test_no_kv_cache.py @@ -107,7 +107,6 @@ def _prepare( prefill_cuda_graph_num_tokens=[], mm_encoder_cache_enabled=True, spec_config=object() if enable_spec_decode else None, - is_draft_model=False, num_seq_slots=None, original_max_draft_len=0, original_max_total_draft_tokens=0, diff --git a/tests/unittest/_torch/executor/engine/test_runners.py b/tests/unittest/_torch/executor/engine/test_runners.py index 223d58925884..2cd0eaa329f0 100644 --- a/tests/unittest/_torch/executor/engine/test_runners.py +++ b/tests/unittest/_torch/executor/engine/test_runners.py @@ -89,7 +89,6 @@ def _config() -> NoKVCacheRunnerConfig: prefill_cuda_graph_num_tokens=[], mm_encoder_cache_enabled=False, spec_config=None, - is_draft_model=False, num_seq_slots=None, original_max_draft_len=0, original_max_total_draft_tokens=0, @@ -211,7 +210,6 @@ def test_encoder_runner_graph_config_comes_only_from_cuda_graph_config( engine.without_logits = False engine.attn_backend = _AttentionBackend engine.attn_runtime_features = AttentionRuntimeFeatures() - engine.is_draft_model = False deps = object() engine._create_runner_deps = Mock(return_value=deps) runner_config = object() diff --git a/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py index 53e49a976521..444586c0c5ed 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py @@ -139,7 +139,6 @@ def _make_mock_model_engine(model_config): engine = Mock() engine.model.model_config = model_config engine.dtype = torch.bfloat16 - engine.is_draft_model = False engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER return engine diff --git a/tests/unittest/_torch/executor/test_distributed_warmup_oom.py b/tests/unittest/_torch/executor/test_distributed_warmup_oom.py index 3808183d7c2b..560c6dab6718 100644 --- a/tests/unittest/_torch/executor/test_distributed_warmup_oom.py +++ b/tests/unittest/_torch/executor/test_distributed_warmup_oom.py @@ -342,7 +342,6 @@ def _mamba_engine(*, world_size: int = 1, dwdp_size: int = 0) -> tuple[PyTorchMo engine.batch_size = 4 engine.max_seq_len = 8 engine.original_max_draft_len = 0 - engine.is_draft_model = False engine.llm_args = SimpleNamespace(enable_autotuner=False) engine.no_cuda_graph = contextlib.nullcontext batch = object() diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index b5b303461610..04480291cfe7 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -293,7 +293,6 @@ def _make_forward_only_engine( engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER engine.enable_spec_decode = False engine.is_spec_decode = False - engine.is_draft_model = False engine.guided_decoder = None engine.max_beam_width = 1 engine._is_encode_only = False @@ -734,7 +733,7 @@ def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: def test_graph_key_forwards_promoted_context_ids(self) -> None: runner = Mock() - runner.config = SimpleNamespace(is_draft_model=False) + runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = True request = _make_request_stub(7) @@ -761,7 +760,7 @@ def test_graph_key_forwards_promoted_context_ids(self) -> None: def test_graph_key_aggregates_encoder_tokens(self) -> None: runner = Mock() - runner.config = SimpleNamespace(is_draft_model=False) + runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner.max_beam_width = 1 runner._get_seq_len_mode.return_value = False @@ -786,7 +785,7 @@ def test_graph_key_aggregates_encoder_tokens(self) -> None: def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: runner = Mock() - runner.config = SimpleNamespace(is_draft_model=False) + runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False first_context = _make_request_stub(1) @@ -833,7 +832,7 @@ def test_graph_key_rounds_encoder_tokens_up_to_captured_extent( def test_graph_key_includes_peft_cache_dtype(self) -> None: runner = Mock() - runner.config = SimpleNamespace(is_draft_model=False) + runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False request = _make_request_stub(7) @@ -891,7 +890,7 @@ def test_graph_dtype_change_falls_back_to_eager(self) -> None: def test_graph_key_includes_lora_variant(self) -> None: runner = Mock() - runner.config = SimpleNamespace(is_draft_model=False) + runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False request = _make_request_stub(7) @@ -1268,7 +1267,6 @@ def test_global_incompatibilities_bypass_candidate_selection(self) -> None: cases = ( "graphs_disabled", "speculative_nonzero_runtime_draft", - "speculative_draft_model", "beam", "encoder_decoder", "ple_recurrent_state", @@ -1285,9 +1283,6 @@ def test_global_incompatibilities_bypass_candidate_selection(self) -> None: elif case == "speculative_nonzero_runtime_draft": engine.enable_spec_decode = True engine.runtime_draft_len = 1 - elif case == "speculative_draft_model": - engine.enable_spec_decode = True - engine.is_draft_model = True elif case == "beam": engine.max_beam_width = 2 elif case == "encoder_decoder": diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 4e6f70734077..828af15a2dc4 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -889,7 +889,6 @@ def flashinfer_autotune(): original_max_draft_len=0, mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False), dist=object(), - is_draft_model=False, guided_decoder=None, max_total_draft_tokens=0, no_cuda_graph=lambda: contextlib.nullcontext(), @@ -1005,7 +1004,6 @@ def flashinfer_autotune(): original_max_draft_len=0, mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False), dist=object(), - is_draft_model=False, guided_decoder=None, max_total_draft_tokens=0, no_cuda_graph=lambda: contextlib.nullcontext(), @@ -1106,7 +1104,6 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): original_max_draft_len=0, mapping=SimpleNamespace(tp_size=2, has_pp=lambda: True), dist=dist, - is_draft_model=False, guided_decoder=None, max_total_draft_tokens=0, no_cuda_graph=lambda: contextlib.nullcontext(), diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index a147c3610d10..2da9d0cac649 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -736,7 +736,6 @@ def test_kv_cache_manager_receives_the_overlap_scheduler_flag(disable_overlap_sc creator = _make_kv_cache_creator(disable_overlap_scheduler) model_engine = SimpleNamespace( model=SimpleNamespace(model_config=SimpleNamespace(is_generation=True)), - is_draft_model=False, ) with patch( diff --git a/tests/unittest/_torch/helpers.py b/tests/unittest/_torch/helpers.py index 301ac1a88964..c2adb2b9b68d 100644 --- a/tests/unittest/_torch/helpers.py +++ b/tests/unittest/_torch/helpers.py @@ -253,7 +253,6 @@ def create_mock_cuda_graph_runner(batch_size: int, enable_attention_dp=False, original_max_draft_len=0, original_max_total_draft_tokens=0, - is_draft_model=False, is_encoder_decoder=False, mapping=Mapping(), dist=None, diff --git a/tests/unittest/_torch/test_route_capture.py b/tests/unittest/_torch/test_route_capture.py index 3476d2563cca..a3c40189f9d1 100644 --- a/tests/unittest/_torch/test_route_capture.py +++ b/tests/unittest/_torch/test_route_capture.py @@ -150,33 +150,15 @@ def test_prefix_store_and_readback_roundtrip(): def test_create_gating_and_fail_closed(): common = dict(rank=0, model_engine=None) # Feature off -> no capturer at all. - assert ( - RouteCapture.create( - **common, enabled=False, pp_size=1, is_spec_decode=False, is_draft_model=False - ) - is None - ) - # Draft engines never capture, even with the feature on. - assert ( - RouteCapture.create( - **common, enabled=True, pp_size=1, is_spec_decode=False, is_draft_model=True - ) - is None - ) + assert RouteCapture.create(**common, enabled=False, pp_size=1, is_spec_decode=False) is None # Supported path -> a capturer instance. - rc = RouteCapture.create( - **common, enabled=True, pp_size=1, is_spec_decode=False, is_draft_model=False - ) + rc = RouteCapture.create(**common, enabled=True, pp_size=1, is_spec_decode=False) assert isinstance(rc, RouteCapture) # Unsupported paths fail closed instead of returning wrong routes. with pytest.raises(RuntimeError): - RouteCapture.create( - **common, enabled=True, pp_size=2, is_spec_decode=False, is_draft_model=False - ) + RouteCapture.create(**common, enabled=True, pp_size=2, is_spec_decode=False) with pytest.raises(RuntimeError): - RouteCapture.create( - **common, enabled=True, pp_size=1, is_spec_decode=True, is_draft_model=False - ) + RouteCapture.create(**common, enabled=True, pp_size=1, is_spec_decode=True) def test_active_capturer_is_looked_up_per_engine_forward(): From b76890a85b19be8415d07859a60925664bf82322 Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Sat, 26 Sep 2026 12:06:08 -0400 Subject: [PATCH 2/2] Update comment Signed-off-by: Mike Iovine --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index edca33e053a3..2e5ce22d0296 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2572,7 +2572,7 @@ def _get_graphs_to_capture( f"Capturing {len(graphs)} graphs: {graphs}") return graphs - # Case 2: One-model without dynamic draft + # Case 2: Static draft length # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. # spec_config is None for non-spec models — fall back to max_draft_len (= 0).