diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 62633fc7e00f..4dcede1c4751 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -319,7 +319,6 @@ common-files: &common_files | tensorrt_llm/_torch/speculative/__init__.py | tensorrt_llm/_torch/speculative/auto_heuristic.py | tensorrt_llm/_torch/speculative/drafter.py | - tensorrt_llm/_torch/speculative/drafting_loops.py | tensorrt_llm/_torch/speculative/eagle3.py | tensorrt_llm/_torch/speculative/interface.py | tensorrt_llm/_torch/speculative/model_drafter.py | @@ -1086,7 +1085,6 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/speculative/__init__.py | tensorrt_llm/_torch/speculative/auto_heuristic.py | tensorrt_llm/_torch/speculative/drafter.py | - tensorrt_llm/_torch/speculative/drafting_loops.py | tensorrt_llm/_torch/speculative/eagle3.py | tensorrt_llm/_torch/speculative/interface.py | tensorrt_llm/_torch/speculative/model_drafter.py | diff --git a/legacy-files.txt b/legacy-files.txt index 9bfbd437f244..73fe3fee5899 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -311,7 +311,6 @@ tensorrt_llm/_torch/shared_tensor/shared_tensor.py tensorrt_llm/_torch/speculative/__init__.py tensorrt_llm/_torch/speculative/auto_heuristic.py tensorrt_llm/_torch/speculative/drafter.py -tensorrt_llm/_torch/speculative/drafting_loops.py tensorrt_llm/_torch/speculative/eagle3.py tensorrt_llm/_torch/speculative/interface.py tensorrt_llm/_torch/speculative/model_drafter.py diff --git a/pyproject.toml b/pyproject.toml index 3ba5e950b89b..f8e1447ca03b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -368,7 +368,6 @@ exclude = [ "tensorrt_llm/_torch/speculative/__init__.py", "tensorrt_llm/_torch/speculative/auto_heuristic.py", "tensorrt_llm/_torch/speculative/drafter.py", - "tensorrt_llm/_torch/speculative/drafting_loops.py", "tensorrt_llm/_torch/speculative/eagle3.py", "tensorrt_llm/_torch/speculative/interface.py", "tensorrt_llm/_torch/speculative/model_drafter.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 386e23c9adfb..612b315564ca 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -328,7 +328,6 @@ include = [ "tensorrt_llm/_torch/speculative/__init__.py", "tensorrt_llm/_torch/speculative/auto_heuristic.py", "tensorrt_llm/_torch/speculative/drafter.py", - "tensorrt_llm/_torch/speculative/drafting_loops.py", "tensorrt_llm/_torch/speculative/eagle3.py", "tensorrt_llm/_torch/speculative/interface.py", "tensorrt_llm/_torch/speculative/model_drafter.py", diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 42708a4f33ac..8de88be14dd8 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -445,7 +445,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, num_contexts: int = 0): diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index bb900fa90bbd..b0e15e281917 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -851,7 +851,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped: bool = False, spec_metadata: Optional["SpecMetadata"] = None, spec_tree_manager: Optional["SpecTreeManager"] = None, num_contexts: int = 0, @@ -864,7 +863,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped, spec_metadata, spec_tree_manager, num_contexts=num_contexts, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py index a711bbe8b53e..454164115ff9 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py @@ -123,7 +123,6 @@ def add_dummy_requests( use_mrope: bool = False, max_beam_width: int = 1, encoder_output_lens: Optional[List[int]] = None, - num_extra_decoding_steps: int = 0, draft_kv_cache_manager=None, ): requests = super().add_dummy_requests( @@ -136,7 +135,6 @@ def add_dummy_requests( use_mrope=use_mrope, max_beam_width=max_beam_width, encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager, ) if prepare_resource: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 03fb9e36df7a..eb02be80f82a 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1187,7 +1187,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, num_contexts: int = 0, @@ -1201,7 +1200,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree: bool, whether using dynamic tree. max_draft_len: int, the number of the draft layers. max_total_draft_tokens: int, the number of all nodes in the tree (except the root). - model_is_wrapped: Optional[bool] = False, whether the drafter model is wrapped (i.e, CDL). spec_metadata: Optional['SpecMetadata'] = None, the metadata of the spec-dec. spec_tree_manager: Optional['SpecTreeManager'] = None, the spec_tree_manager for draft token tree. num_contexts: int = 0, the number of context (prefill) requests in the @@ -1329,61 +1327,26 @@ def update_spec_dec_param( self.spec_decoding_generation_lengths[:batch_size].fill_(n_dt) cpp_query_len = n_dt - # Case 2/3: static tree + # Case 2: static tree (target model only) elif self.is_spec_dec_tree and not self.is_spec_dec_dynamic_tree and spec_metadata is not None: assert (spec_metadata.spec_dec_mode.is_eagle3() or spec_metadata.spec_dec_mode.is_eagle3_one_model() ), "Tree decoding is only supported for Eagle3 now" + assert not getattr(spec_metadata, 'is_draft_model', False), ( + "Static tree spec-dec params are only prepared for the target model" + ) - is_target_model = not getattr(spec_metadata, 'is_draft_model', - False) - - # Case 2: static tree and target model - if is_target_model: - # For the target model, we update the spec-dec parameters with the spec_tree_manager, which is prepared in advance. - self.spec_decoding_position_offsets[:batch_size, :].copy_( - spec_tree_manager.spec_dec_position_offsets[0, :], - non_blocking=True) - self.spec_decoding_packed_mask[:batch_size, :, :].copy_( - spec_tree_manager.spec_dec_packed_mask[0, :, :], - non_blocking=True) - self.spec_decoding_generation_lengths[:batch_size].fill_( - spec_tree_manager.max_total_draft_tokens + 1) - - # Case 3: static tree and the first drafter layer - else: - assert model_is_wrapped == True, "The drafter model should be wrapped" - # The first drafter layer will take the padded tokens as input (padding to the max_draft_len + 1) - # But the spec-dec parameters are still in the shape of max_total_draft_tokens + 1. - # Considering that these spec-dec params are accessed consecutively (without padding) in the attention Op, - # we need to write them consecutively when setting them. - # For the next drafter layers, we will prepare these spec-dec params in the drafting loops. - - # position_offsets - position_offset = torch.arange( - max_draft_len + 1, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()).repeat(batch_size) - self.spec_decoding_position_offsets.reshape( - -1)[:(max_draft_len + 1) * batch_size].copy_( - position_offset, non_blocking=True) - # packed_mask - dummy_idx = torch.arange(max_draft_len + 1) - spec_decoding_packed_mask = torch.pow( - 2, dummy_idx + 1) - 1 # [max_draft_len + 1] - spec_decoding_packed_mask = spec_decoding_packed_mask.repeat( - batch_size) # [batch_size * (max_draft_len + 1)] - self.spec_decoding_packed_mask.reshape( - -1)[:(max_draft_len + 1) * batch_size].copy_( - spec_decoding_packed_mask, non_blocking=True) - self.generate_spec_decoding_generation_length( - runtime_draft_len=max_draft_len) - if (self.spec_decoding_position_offsets is not None - and self.spec_decoding_position_offsets.dim() == 1): - cpp_query_len = max_draft_len + 1 - - # Case 4: linear tree + # For the target model, we update the spec-dec parameters with the spec_tree_manager, which is prepared in advance. + self.spec_decoding_position_offsets[:batch_size, :].copy_( + spec_tree_manager.spec_dec_position_offsets[0, :], + non_blocking=True) + self.spec_decoding_packed_mask[:batch_size, :, :].copy_( + spec_tree_manager.spec_dec_packed_mask[0, :, :], + non_blocking=True) + self.spec_decoding_generation_lengths[:batch_size].fill_( + spec_tree_manager.max_total_draft_tokens + 1) + + # Case 3: linear tree else: # Currently dynamic draft length is only supported for linear tree # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index b1a965d20e32..43f7d9773876 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -3357,7 +3357,6 @@ def add_dummy_requests( use_mrope: bool = False, max_beam_width: int = 1, encoder_output_lens: Optional[List[int]] = None, - num_extra_decoding_steps: int = 0, draft_kv_cache_manager: Optional["BaseResourceManager"] = None, ): _kv_draft = ( @@ -3430,7 +3429,7 @@ def release_resources( release_resources(req) return None kv_cache.stop_committing() - dummy_capacity = token_num + self.num_extra_kv_tokens + num_extra_decoding_steps + dummy_capacity = token_num + self.num_extra_kv_tokens if is_gen and not materialize_history: kv_cache.enable_swa_scratch_reuse = False # Need to hint the committed history to activate stale-block diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 041c4933542a..696eceed0e57 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1029,7 +1029,6 @@ def __init__( # Whether the request is for the first forward of the draft model. self.py_is_first_draft = is_first_draft self.d2t = None - self.py_draft_use_greedy_sampling = False self.py_disable_speculative_decoding = False # Chunked logits parameters diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 12035139b2ec..61adda1a378d 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -2488,11 +2488,6 @@ def add_dummy_requests( use_mrope: bool = False, max_beam_width: int = 1, encoder_output_lens: Optional[List[int]] = None, - # For capturable drafting loops. During normal inference, the draft model always - # has enough KV cache space to fit all of our draft tokens. During warmup, however, - # we need to make the KV cache manager aware that multiple autoregressive steps will - # occur. - num_extra_decoding_steps: int = 0, draft_kv_cache_manager: Optional[KVCacheManager] = None, ) -> List[LlmRequest]: requests = super().add_dummy_requests( @@ -2505,7 +2500,6 @@ def add_dummy_requests( use_mrope=use_mrope, max_beam_width=max_beam_width, encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager, ) if requests: @@ -3551,7 +3545,6 @@ def add_dummy_requests( use_mrope: bool = False, max_beam_width: int = 1, encoder_output_lens: Optional[List[int]] = None, - num_extra_decoding_steps: int = 0, draft_kv_cache_manager: Optional[BaseResourceManager] = None, ) -> List[LlmRequest]: requests = super().add_dummy_requests( @@ -3564,7 +3557,6 @@ def add_dummy_requests( use_mrope=use_mrope, max_beam_width=max_beam_width, encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager, ) if requests and prepare_resource: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ae97f89bb386..7efe960af8ac 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -75,7 +75,6 @@ prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) -from ..speculative.drafting_loops import BaseDraftingLoopWrapper from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..tensor_lru_cache import TensorLRUCache @@ -217,10 +216,9 @@ def _filter_piecewise_capture_num_tokens( max_num_tokens: int, max_batch_size: int, max_seq_len: int, - num_extra_decoding_steps: int = 0, ) -> Tuple[list[int], list[int]]: """Cap piecewise CUDA graph capture candidates at the engine's reachable - `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` + `num_tokens` ceiling `max_batch_size * (max_seq_len - 1)` clamping user-requested sizes above it down to the ceiling. Each in-flight request must leave room for at least one decode token, @@ -240,8 +238,7 @@ def _filter_piecewise_capture_num_tokens( ceiling but within `max_num_tokens` (the clamped ones, reported so the caller's warning fires). """ - max_capturable_num_tokens = max( - 0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)) + max_capturable_num_tokens = max(0, max_batch_size * (max_seq_len - 1)) piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) if piecewise_capacity_limit > 0: kept = sorted({ @@ -390,8 +387,6 @@ def __init__( dist: Optional[Distributed] = None, spec_config: Optional[DecodingBaseConfig] = None, is_draft_model: bool = False, - drafting_loop_wrapper: Optional[Callable[[torch.nn.Module], - torch.nn.Module]] = None, model: Optional[torch.nn.Module] = None, checkpoint_loader: Optional[BaseCheckpointLoader] = None, model_weights_memory_tag: Optional[str] = None, @@ -480,14 +475,9 @@ def __init__( else: self.max_draft_loop_tokens = self.original_max_total_draft_tokens - preserve_wrapped_eagle3_widths = (spec_config is not None - and is_draft_model - and drafting_loop_wrapper is not None - and - spec_config.spec_dec_mode.is_eagle3()) # 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 and not preserve_wrapped_eagle3_widths: + 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 @@ -546,11 +536,6 @@ def __init__( self._enable_non_overlap_adp_forward_intent = ( should_enable_non_overlap_adp_forward_intent( mapping, llm_args.disable_overlap_scheduler)) - if drafting_loop_wrapper is not None: - self.model = drafting_loop_wrapper(self.model) - self.model_is_wrapped = True - else: - self.model_is_wrapped = False self.sparse_attention_config = self.model.model_config.sparse_attention_config # In case that some tests use stub models and override `_load_model`. if not hasattr(self.model, 'extra_attrs'): @@ -776,21 +761,19 @@ def __init__( if prefill_cuda_graph_num_tokens is None: prefill_cuda_graph_num_tokens = cuda_graph_batch_sizes or [] - num_extra_decoding_steps = self._get_num_extra_decoding_steps() self._prefill_cuda_graph_num_tokens, unrecordable = ( _filter_prefill_capture_num_tokens( prefill_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_batch_size=self.batch_size, max_seq_len=self.max_seq_len, - num_extra_decoding_steps=num_extra_decoding_steps, )) if unrecordable: logger.warning( f"Skipping prefill CUDA graph capture for num_tokens=" f"{unrecordable}: exceeds reachable ceiling " - f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)=" - f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. " + f"max_batch_size*(max_seq_len-1)=" + f"{max(0, self.batch_size * (self.max_seq_len - 1))}. " f"Clamping them to the ceiling; raise max_seq_len for larger graphs." ) @@ -878,7 +861,7 @@ def __init__( self.previous_kv_lens_offsets_cuda = torch.zeros( (self.batch_size, ), dtype=torch.int, device='cuda') self.without_logits = self.spec_config.spec_dec_mode.without_logits( - ) or self.model_is_wrapped + ) self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 self.max_draft_len = spec_config.max_draft_len # Mutable per-iteration draft length (updated each iteration when @@ -1999,8 +1982,8 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): # even when ``curr_max_num_tokens`` is close to the block ceiling. # ``least_requests=False`` instead spreads the token budget across # ``batch_size`` short sequences; when each sequence's length lands - # exactly on a block boundary AND the KV cache has ``num_extra_kv_tokens`` - # or ``num_extra_decoding_steps`` > 0 (e.g. spec decoding cases), + # exactly on a block boundary AND the KV cache has + # ``num_extra_kv_tokens`` > 0 (e.g. spec decoding cases), # ``add_token`` needs to allocate one extra block per sequence, which # ``_create_warmup_request``'s ``blocks_to_use`` estimate doesn't # account for. On a small KV pool (e.g. Qwen3.5 hybrid with DFlash spec @@ -2119,9 +2102,7 @@ def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: return mapping def _get_graphs_to_capture( - self, cuda_graph_batch_sizes: list[int], - spec_resource_manager: Optional[BaseResourceManager] - ) -> list[tuple[int, int]]: + self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, int]]: """Determine which (batch_size, draft_len) graphs to capture. Returns: @@ -2130,12 +2111,7 @@ def _get_graphs_to_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: - if self.model_is_wrapped and self.is_spec_decode and spec_resource_manager is not None and isinstance( - spec_resource_manager, Eagle3ResourceManager): - # The CDL path uses draft_len > 0 for the number of iterations in the drafting loop. - draft_len = self.original_max_total_draft_tokens - else: - draft_len = self.max_total_draft_tokens + 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 @@ -2278,16 +2254,13 @@ def _capture_generation_cuda_graphs(self, if self.cuda_graph_runner.is_warmup_only else "capture") logger.info(f"Running CUDA graph {operation} for " f"{len(self._cuda_graph_batch_sizes)} batch sizes.") - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) # Reverse order so smaller graphs can reuse memory from larger ones cuda_graph_batch_sizes = sorted(self._cuda_graph_batch_sizes, reverse=True) # Determine which graph shapes to process. - graphs_to_capture = self._get_graphs_to_capture(cuda_graph_batch_sizes, - spec_resource_manager) + graphs_to_capture = self._get_graphs_to_capture(cuda_graph_batch_sizes) graphs_to_capture = sorted(graphs_to_capture, reverse=True) # Create CUDA graphs for short and long sequences separately for sparse attention. # self.max_seq_len is the global max sequence length. For Helix CP each @@ -2698,16 +2671,6 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests], if spec_resource_manager is not None: spec_resource_manager.free_resources(req) - def _get_num_extra_decoding_steps(self) -> int: - """Determines extra decoding steps needed for fused drafting loops.""" - if isinstance(self.model, BaseDraftingLoopWrapper): - return self.model.max_total_draft_tokens - else: - assert not self.model_is_wrapped, ( - f"Please add logic to determine num_extra_decoding_steps for drafting loop {type(self.model)}" - ) - return 0 - def _create_warmup_request( self, resource_manager: ResourceManager, @@ -2730,8 +2693,6 @@ def _create_warmup_request( if num_tokens > self.max_num_tokens or num_tokens > available_tokens: return None - num_extra_decoding_steps = self._get_num_extra_decoding_steps() - if num_gen_requests > self.batch_size: return None num_gen_tokens = num_gen_requests * (1 + self.max_total_draft_tokens) @@ -2743,10 +2704,10 @@ def _create_warmup_request( ctx_requests = [] gen_requests = [] - # For drafting loops, reduce max_seq_len to leave room for extra decoding steps - max_seq_len = self.max_seq_len - 1 - num_extra_decoding_steps + # Leave room for at least one decode token per request. + max_seq_len = self.max_seq_len - 1 if max_seq_len < 1: - return None # Not enough sequence length for drafting loop + return None num_full_seqs = 0 num_left_over_tokens = 0 @@ -2774,9 +2735,9 @@ def _create_warmup_request( return None # Not enough batch size to fill the request # Mirror add_dummy_requests' actual allocation: on top of the raw - # token count, every sequence gets num_extra_kv_tokens + - # num_extra_decoding_steps add_token calls, and generation dummies - # additionally reserve max_draft_loop_tokens for the draft loop. + # token count, every sequence gets num_extra_kv_tokens add_token + # calls, and generation dummies additionally reserve + # max_draft_loop_tokens for the draft loop. # In one-engine spec modes that is (max_draft_len - 1) extra KV # tokens plus max_draft_len draft-loop tokens per gen dummy, i.e. # 2 * max_draft_len - 1 on top of the single prompt token. @@ -2788,7 +2749,7 @@ def blocks_for_seq(num_tokens: int) -> int: return math.ceil(num_tokens / kv_cache_manager.tokens_per_block) extra_ctx_tokens = (getattr(kv_cache_manager, "num_extra_kv_tokens", 0) - or 0) + num_extra_decoding_steps + or 0) extra_gen_tokens = extra_ctx_tokens + self.max_draft_loop_tokens blocks_to_use = num_full_seqs * blocks_for_seq(max_seq_len + extra_ctx_tokens) @@ -2814,7 +2775,6 @@ def blocks_for_seq(num_tokens: int) -> int: max_num_draft_tokens=self.max_total_draft_tokens, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if ctx_requests is None: @@ -2835,7 +2795,6 @@ def blocks_for_seq(num_tokens: int) -> int: kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if gen_requests is None: @@ -2878,7 +2837,6 @@ def _create_cuda_graph_warmup_request( return None result = ScheduledRequests() - num_extra_decoding_steps = self._get_num_extra_decoding_steps() runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 @@ -2906,7 +2864,6 @@ def _create_cuda_graph_warmup_request( use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=list(mixed_context_encoder_output_lens), - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if context_requests is None: return None @@ -2926,7 +2883,6 @@ def _create_cuda_graph_warmup_request( max_beam_width=self.max_beam_width, encoder_output_lens=[max_encoder_output_len] * len(generation_request_ids), - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if generation_requests is None: for request in context_requests: @@ -2949,7 +2905,6 @@ def _create_cuda_graph_warmup_request( use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if requests is None: return None @@ -3002,10 +2957,6 @@ def free_warmup_requests() -> None: if max_position_embeddings is not None: token_num = min(token_num, max_position_embeddings - _kv_draft) - assert token_num > num_extra_decoding_steps, ( - "Cannot fuse drafting loop. Not enough KV cache space for all draft tokens." - ) - token_num -= num_extra_decoding_steps token_num = int( token_num) # Ensure int for range() in add_dummy_requests @@ -3019,7 +2970,6 @@ def free_warmup_requests() -> None: max_beam_width=self.max_beam_width, encoder_output_lens=[max_encoder_output_len] if is_enc_dec else None, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if max_seq_len_request is None: @@ -4528,8 +4478,7 @@ def _can_use_incremental_update( return False has_current_device_draft = next_draft_tokens_device is not None - return (self.is_draft_model and self.model_is_wrapped) or ( - has_current_device_draft and self.has_previous_device_draft) + return has_current_device_draft and self.has_previous_device_draft @nvtx_range("_apply_incremental_update") def _apply_incremental_update( @@ -7269,8 +7218,7 @@ def forward(self, is None) # attn_metadata now depends on spec_metadata since it determines the shape/content of spec_dec parameter Tensors is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode( - spec_resource_manager, self.is_draft_model, self.attn_backend, - self.model_is_wrapped) + self.is_draft_model, self.attn_backend) # Propagate runtime_draft_len (already set on self by py_executor) # to spec_metadata so downstream code (eagle3, interface, trtllm) can read it. spec_metadata.runtime_draft_len = self.runtime_draft_len @@ -7304,7 +7252,6 @@ def forward(self, is_spec_dec_dynamic_tree=spec_metadata.is_spec_dec_dynamic_tree, max_draft_len=sd_max_draft_len, max_total_draft_tokens=sd_max_total, - model_is_wrapped=self.model_is_wrapped, spec_metadata=spec_metadata, spec_tree_manager=spec_tree_manager, num_contexts=scheduled_requests.num_context_requests) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7f04fb9742f3..c4bc59ec28fe 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -32,7 +32,6 @@ from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from ..attention_backend.interface import AttentionRuntimeFeatures -from ..attention_backend.trtllm import TrtllmAttention from ..distributed import Distributed from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, get_spec_resource_manager) @@ -606,38 +605,6 @@ def allocation_scope(current_stage: ExecutorMemoryType): with allocation_scope(ExecutorMemoryType.MODEL_ENGINE_DRAFT): draft_spec_config = copy.copy(spec_config) - use_chain_drafter = ( - guided_decoding_config is None - and draft_spec_config._allow_chain_drafter - and draft_spec_config._allow_greedy_draft_tokens - and llm_args.attn_backend == "TRTLLM" - and draft_spec_config.draft_len_schedule is None) - - logger.debug(f"USE CHAIN DRAFTER: {use_chain_drafter}") - if use_chain_drafter: - - def drafting_loop_wrapper(model): - from tensorrt_llm._torch.speculative.drafting_loops import ( - LinearDraftingLoopWrapper, - StaticTreeDraftingLoopWrapper) - from tensorrt_llm.llmapi import EagleDecodingConfig - - static_tree_drafter = isinstance( - draft_spec_config, EagleDecodingConfig - ) and draft_spec_config.eagle_choices is not None - - if static_tree_drafter: - return StaticTreeDraftingLoopWrapper( - spec_config.max_draft_len, - spec_config.tokens_per_gen_step - 1, max_batch_size, - model) - else: - return LinearDraftingLoopWrapper( - spec_config.max_draft_len, - spec_config.tokens_per_gen_step - 1, model) - else: - drafting_loop_wrapper = None - draft_llm_args = copy.copy(llm_args) if spec_config.load_format == "dummy": draft_llm_args.load_format = LoadFormat.DUMMY @@ -657,7 +624,6 @@ def drafting_loop_wrapper(model): dist=dist, spec_config=draft_spec_config, is_draft_model=True, - drafting_loop_wrapper=drafting_loop_wrapper, model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, ) @@ -669,13 +635,10 @@ def drafting_loop_wrapper(model): else: draft_model_engine = None - # TODO: Overlap scheduler is not supported for below cases: - # 1. non-CDL is used - # 2. non-TrtllmAttention attention backend is used - if has_draft_model_engine and (not use_chain_drafter or not issubclass( - draft_model_engine.attn_backend, TrtllmAttention)): + # TODO: Overlap scheduler is not supported for two-model speculative decoding. + if has_draft_model_engine: logger.warning( - "Overlap scheduler is not supported for non-CDL or non-TrtllmAttention backend." + "Overlap scheduler is not supported for two-model speculative decoding." ) llm_args.disable_overlap_scheduler = True diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 7e4fed2a0eb2..81ba164f9693 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1140,11 +1140,6 @@ def add_dummy_requests( use_mrope: bool = False, max_beam_width: int = 1, encoder_output_lens: Optional[List[int]] = None, - # For capturable drafting loops. During normal inference, the draft model always - # has enough KV cache space to fit all of our draft tokens. During warmup, however, - # we need to make the KV cache manager aware that multiple autoregressive steps will - # occur. - num_extra_decoding_steps: int = 0, draft_kv_cache_manager: Optional[BaseResourceManager] = None, ): _kv_draft = kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens @@ -1210,8 +1205,6 @@ def add_dummy_requests( for req_id, token_num, _ in batch_request_infos: for _ in range(self.num_extra_kv_tokens): self.impl.add_token(req_id) - for _ in range(num_extra_decoding_steps): - self.impl.add_token(req_id) if draft_batch_request_infos and draft_kv_cache_manager is not None: draft_kv_cache_manager.impl.add_sequence_batch( diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/two_model_spec_dec.py b/tensorrt_llm/_torch/pyexecutor/sampler/two_model_spec_dec.py index e5ef3a39c334..92db168539f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/two_model_spec_dec.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/two_model_spec_dec.py @@ -248,11 +248,7 @@ def _process_draft_tokens_rejection_sampling( # filtering of vocab_size logits, out of vocab_size in # total. The 'sample' below should generally be avoided # by retaining the draft_probs during drafting (TRTLLM-7772). - draft_sampling_strategy = ( - GREEDY - if request.py_draft_use_greedy_sampling - else _request_strategy(request, vocab_size=2**31) - ) + draft_sampling_strategy = _request_strategy(request, vocab_size=2**31) generator = self._sampler.get_generator(request.py_draft_logits.device) _, draft_probs, _ = sample( draft_sampling_strategy, diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py deleted file mode 100644 index 133814bb2d22..000000000000 --- a/tensorrt_llm/_torch/speculative/drafting_loops.py +++ /dev/null @@ -1,522 +0,0 @@ -""" -This module contains capturable drafting loops for speculative decoding. - -These are torch modules wrap another draft model. The wrapped module -is supposed to invoke the draft model autoregressively and invoke -a sampling algorithm to obtain draft tokens. By structuring the code -like this, we are able to avoid host overhead: the entire drafting process -for speculation can be launched as a single CUDA graph. -""" - -from abc import ABC, abstractmethod -from contextlib import contextmanager -from typing import Optional, final - -import torch - -from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import \ - greedy_search_sampling_batch -from tensorrt_llm._torch.speculative.eagle3 import Eagle3SpecMetadata -from tensorrt_llm._torch.speculative.interface import SpecMetadata -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager - -# Enable capture_scalar_outputs to avoid graph breaks from Tensor.item() calls -torch._dynamo.config.capture_scalar_outputs = True - - -class BaseDraftingLoopWrapper(ABC, torch.nn.Module): - - @abstractmethod - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, - attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, - **kwargs) -> dict[str, torch.Tensor]: - raise NotImplementedError - - @abstractmethod - def sample(self, - logits: torch.Tensor, - max_top_k: Optional[int] = None) -> torch.Tensor: - raise NotImplementedError - - @abstractmethod - def prepare_for_generation( - self, - attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata, - position_ids: torch.Tensor, - spec_tree_manager: Optional[SpecTreeManager] = None - ) -> torch.Tensor | None: - raise NotImplementedError - - @final - def load_weights_from_target_model(self, target_model) -> None: - loader = getattr(self.draft_model, "load_weights_from_target_model", - None) - if callable(loader): - self.draft_model.load_weights_from_target_model(target_model) - - -@contextmanager -def save_metadata_state(attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata) -> None: - attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") - batch_size = attn_metadata.num_seqs - # Do not use prepare_for_spec_dec for this special field. - # TRTLLM attention uses views of this tensor internally and prepare_for_spec_dec - # creates a copy. If you write to the copy, TRTLLM attention won't see the updates. - kv_lens = attn_metadata.kv_lens_cuda[:batch_size].clone() - - if attn_metadata.is_cuda_graph: - assert spec_metadata.is_cuda_graph - num_tokens = spec_metadata.num_tokens - if isinstance(spec_metadata, Eagle3SpecMetadata): - read_indices = spec_metadata.hidden_states_read_indices[: - batch_size].clone( - ) - write_indices = spec_metadata.hidden_states_write_indices[: - batch_size].clone( - ) - - try: - yield - finally: - attn_metadata.restore_from_spec_dec() - attn_metadata.kv_lens_cuda[:batch_size].copy_(kv_lens) - attn_metadata.on_update() - if attn_metadata.is_cuda_graph: - spec_metadata.num_tokens = num_tokens - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.hidden_states_read_indices[:batch_size].copy_( - read_indices) - spec_metadata.hidden_states_write_indices[:batch_size].copy_( - write_indices) - - # This restore has to happen even if the spec_metadata is not being used - # for CUDA graphs. It won't be reset by spec_metadata.prepare(). - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.is_first_draft = True - spec_metadata.eagle3_resource_manager.is_first_draft = True - - -class LinearDraftingLoopWrapper(BaseDraftingLoopWrapper): - - def __init__(self, max_draft_len: int, max_total_draft_tokens: int, - draft_model: torch.nn.Module): - super().__init__() - self.draft_model = draft_model - self.config = self.draft_model.config - self.model_config = self.draft_model.model_config - self.max_draft_len = max_draft_len - self.max_total_draft_tokens = max_total_draft_tokens - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, - attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, - **kwargs) -> dict[str, torch.Tensor]: - logits = self.draft_model.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - return_context_logits=True) - logits = logits[spec_metadata.gather_ids] - - new_draft_tokens = [self.sample(logits)] - draft_logits = [logits] - if self.max_draft_len > 1: - is_eagle3 = isinstance(spec_metadata, Eagle3SpecMetadata) - with save_metadata_state(attn_metadata, spec_metadata): - batch_size = attn_metadata.num_seqs - - new_position_ids = self.prepare_for_generation( - attn_metadata, spec_metadata, position_ids) - for i in range(self.max_draft_len - 1): - logits = self.draft_model.forward( - input_ids=new_draft_tokens[-1], - position_ids=new_position_ids, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata) - new_draft_tokens.append(self.sample(logits)) - draft_logits.append(logits) - new_position_ids += 1 - attn_metadata.kv_lens_cuda[:batch_size] += 1 - if i == 0 and is_eagle3: - spec_metadata.hidden_states_read_indices[:batch_size].copy_( - spec_metadata. - hidden_states_write_indices[:batch_size]) - - return { - "new_draft_tokens": torch.stack(new_draft_tokens), - "draft_logits": torch.stack(draft_logits) - } - - def sample(self, logits: torch.Tensor) -> torch.Tensor: - # TODO: inject the sampler here so we can support non-greedy - tokens, _ = greedy_search_sampling_batch(logits, return_probs=False) - if hasattr(self.draft_model.model, "d2t"): - d2t = self.draft_model.model.d2t.data - return tokens + d2t[tokens] - - return tokens - - def prepare_for_generation(self, attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata, - position_ids: torch.Tensor) -> torch.Tensor: - batch_size = attn_metadata.num_seqs - num_accepted_draft_tokens = spec_metadata.num_accepted_draft_tokens[: - batch_size] - seq_lens = attn_metadata.seq_lens_cuda[:batch_size] - attn_metadata.kv_lens_cuda[: - batch_size] -= seq_lens - num_accepted_draft_tokens - 1 - # Calculate last accepted token indices - last_tokens_idx = torch.cumsum( - seq_lens, dim=0, - dtype=torch.long) - seq_lens + num_accepted_draft_tokens - new_position_ids = position_ids[0, last_tokens_idx] + 1 - - attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[:batch_size].fill_(1) - attn_metadata.on_update() - attn_metadata.kv_lens_cuda[:batch_size] += 1 - - attn_metadata.host_request_types[:attn_metadata.num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - # The next inference of draft model will not use spec decoding and the number of input tokens is 1 - attn_metadata.use_spec_decoding = False - - spec_metadata.num_tokens = batch_size - - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.eagle3_resource_manager.is_first_draft = False - spec_metadata.is_first_draft = False - - old_write_indices = spec_metadata.hidden_states_write_indices - - spec_metadata.hidden_states_read_indices[:batch_size].copy_( - old_write_indices[last_tokens_idx]) - spec_metadata.hidden_states_write_indices[:batch_size].copy_( - torch.arange( - batch_size, - dtype=spec_metadata.hidden_states_write_indices.dtype, - device=spec_metadata.hidden_states_write_indices.device)) - - return new_position_ids - - -class StaticTreeDraftingLoopWrapper(BaseDraftingLoopWrapper): - - def __init__(self, max_draft_len: int, max_total_draft_tokens: int, - max_batch_size: int, draft_model: torch.nn.Module): - super().__init__() - self.draft_model = draft_model - self.config = self.draft_model.config - self.model_config = self.draft_model.model_config - self.max_draft_len = max_draft_len - self.max_total_draft_tokens = max_total_draft_tokens - self.max_batch_size = max_batch_size - - self.draft_tokens_buffer = torch.zeros( - (max_batch_size, max_total_draft_tokens + 1), - dtype=torch.int64, - device='cuda') - self.position_ids_buffer = torch.zeros( - (max_batch_size, max_total_draft_tokens + 1), - dtype=torch.int64, - device='cuda') - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, - attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, - **kwargs) -> dict[str, torch.Tensor]: - assert isinstance(spec_metadata, Eagle3SpecMetadata) - spec_tree_manager = spec_metadata.eagle3_resource_manager.spec_tree_manager - - logits = self.draft_model.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - return_context_logits=True) - batch_size = attn_metadata.num_seqs - vocab_size = logits.shape[-1] - logits = logits[spec_metadata.gather_ids] # [batch_size, vocab_size] - - # new_draft_tokens: [batch_size * max_top_k] - new_draft_tokens = self.sample(logits=logits, - max_top_k=spec_tree_manager.max_top_k) - - self.extract_real_draft_tokens( - cur_draft_idx=0, - batch_size=batch_size, - new_draft_tokens=new_draft_tokens, - use_cuda_graph=attn_metadata.is_cuda_graph, - spec_tree_manager=spec_tree_manager) - return_draft_logits = None - with save_metadata_state(attn_metadata, spec_metadata): - batch_size = attn_metadata.num_seqs - - self.prepare_for_generation(attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager, - position_ids=position_ids) - - for layer_idx in range(1, self.max_draft_len): - # input_ids: [batch_size * (max_total_draft_tokens + 1)] - # position_ids: [batch_size * (max_total_draft_tokens + 1)] - # logits: [batch_size * (max_total_draft_tokens + 1), vocab_size] - logits = self.draft_model.forward( - input_ids=self.draft_tokens_buffer[:batch_size, :self. - max_total_draft_tokens + - 1].reshape(-1), - position_ids=self. - position_ids_buffer[:batch_size, :self. - max_total_draft_tokens + 1].reshape(-1), - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - return_context_logits=True) - - # new_draft_tokens: [batch_size * (max_total_draft_tokens + 1) * max_top_k] - new_draft_tokens = self.sample( - logits=logits, max_top_k=spec_tree_manager.max_top_k) - # Keep updating - self.extract_real_draft_tokens( - cur_draft_idx=layer_idx, - batch_size=batch_size, - new_draft_tokens=new_draft_tokens, - use_cuda_graph=attn_metadata.is_cuda_graph, - spec_tree_manager=spec_tree_manager) - - if layer_idx == self.max_draft_len - 1: - return_draft_logits = logits - - # self.draft_tokens_buffer[:batch_size, :]: [batch_size, max_total_draft_tokens + 1] - # return_new_draft_tokens: [max_total_draft_tokens, batch_size] - return_new_draft_tokens = torch.transpose( - self.draft_tokens_buffer[:batch_size, :-1], 0, 1) - - # return_draft_logits: [batch_size, max_total_draft_tokens + 1, vocab_size] -> [max_total_draft_tokens, batch_size, vocab_size] - if return_draft_logits is None: - # When max_draft_len == 1, the loop doesn't execute. - # Expand the initial logits to match the expected shape. - return_draft_logits = logits.unsqueeze(1).expand( - batch_size, self.max_total_draft_tokens + 1, - vocab_size).reshape(-1, vocab_size) - - return_draft_logits = return_draft_logits.reshape( - batch_size, self.max_total_draft_tokens + 1, vocab_size) - return_draft_logits = torch.transpose(return_draft_logits[:, :-1, :], 0, - 1) - - assert return_new_draft_tokens.shape == (self.max_total_draft_tokens, - batch_size) - assert return_draft_logits.shape == (self.max_total_draft_tokens, - batch_size, vocab_size) - - return { - "new_draft_tokens": return_new_draft_tokens, - "draft_logits": return_draft_logits - } - - def sample(self, logits: torch.Tensor, max_top_k: int) -> torch.Tensor: - # TODO: inject the sampler here so we can support non-greedy - - # for draft_layer_idx == 0, logits is of shape [batch_size, vocab_size] - # for draft_layer_idx > 0, logits is of shape [batch_size * (max_total_draft_tokens + 1), vocab_size] - indices = torch.topk( - logits, k=max_top_k, dim=-1 - ).indices # [batch_size, max_top_k] or [batch_size * max_total_draft_tokens, max_top_k] - tokens = indices.reshape(-1) - - if hasattr(self.draft_model.model, "d2t"): - d2t = self.draft_model.model.d2t.data - tokens = tokens + d2t[tokens] - - return tokens - - def extract_real_draft_tokens(self, cur_draft_idx: int, batch_size: int, - new_draft_tokens: torch.Tensor, - use_cuda_graph: bool, - spec_tree_manager: SpecTreeManager): - ''' - Extract the real draft tokens from the new draft tokens to self.draft_tokens_buffer. - ''' - # After the first drafter layer, new_draft_tokens: [batch_size * max_top_k] - # For other drafter layers, new_draft_tokens: [batch_size * (max_total_draft_tokens + 1) * max_top_k] - if cur_draft_idx == 0: - assert new_draft_tokens.shape[0] == (batch_size * - spec_tree_manager.max_top_k) - else: - assert new_draft_tokens.shape[0] == ( - batch_size * (self.max_total_draft_tokens + 1) * - spec_tree_manager.max_top_k) - - # reshape the new_draft_tokens to [batch_size, -1, spec_tree_manager.max_top_k] - new_draft_tokens = new_draft_tokens.reshape(batch_size, -1, - spec_tree_manager.max_top_k) - - # If using cuda graph, we need to use a torch op to implement this logic - if use_cuda_graph: - torch.ops.trtllm.extract_real_draft_tokens_op( - new_draft_tokens, self.draft_tokens_buffer, spec_tree_manager. - tokens_gather_idx_for_drafter_model[cur_draft_idx], - spec_tree_manager.top_k_list_cuda[cur_draft_idx], - spec_tree_manager.draft_tokens_indices_cumsum, cur_draft_idx, - batch_size, self.max_draft_len, self.max_total_draft_tokens, - spec_tree_manager.max_top_k) - else: - # 1) Gather the real tokens processed by this layer - process_tokens = new_draft_tokens[:, spec_tree_manager. - tokens_gather_idx_for_drafter_model[ - cur_draft_idx], :] # [batch_size, num_tokens_process_this_layer, max_top_k] - process_tokens = process_tokens.reshape( - -1, spec_tree_manager.max_top_k - ) # [batch_size * num_tokens_process_this_layer, max_top_k] - - # 2) Gather the real draft tokens samples by these processed tokens' logits - top_k_list = spec_tree_manager.top_k_list_cuda[ - cur_draft_idx].repeat( - batch_size) # [batch_size * num_tokens_process_this_layer] - assert top_k_list.shape[0] == process_tokens.shape[0] - - # [batch_size * num_tokens_process_this_layer, spec_tree_manager.max_top_k] - col_indices = torch.arange( - spec_tree_manager.max_top_k, - device=new_draft_tokens.device).unsqueeze(0).repeat( - top_k_list.shape[0], 1) - - mask = col_indices < top_k_list.unsqueeze( - 1 - ) # [batch_size * num_tokens_process_this_layer, spec_tree_manager.max_top_k] - - real_new_draft_tokens = process_tokens[ - mask] # [batch_size * sum(spec_tree_manager.top_k_list_cuda[cur_draft_idx])] - real_new_draft_tokens = real_new_draft_tokens.reshape( - batch_size, -1 - ) # [batch_size, sum(spec_tree_manager.top_k_list_cuda[cur_draft_idx])] - - self.draft_tokens_buffer[:batch_size, spec_tree_manager. - draft_tokens_indices_cumsum[cur_draft_idx]: - spec_tree_manager. - draft_tokens_indices_cumsum[ - cur_draft_idx + - 1]] = real_new_draft_tokens[:, :] - - def prepare_for_generation(self, attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata, - spec_tree_manager: SpecTreeManager, - position_ids: torch.Tensor): - ''' - Prepare the inputs for the subsequent draft layers. - Note: Except for the 0th drafter layer, in each subsequent drafter layer, - we take 'max_total_drafter_tokens + 1' draft tokens as input. - Only the first part of the draft tokens is meaningful, and the later tokens can be regarded as padding - until we continuously write the correct value. - - This introduces additional redundant computation, but it makes it compatible with cuda graphs. - - What we need to prepare are: - 1) position_ids - 2) attn_metadata - 2.1) kv_lens_cuda - 2.2) _seq_lens, _seq_lens_cuda - 2.3) host_request_types - 2.4) num_contexts - 2.5) use_spec_decoding - 2.6) spec_decoding_position_offsets - 2.7) spec_decoding_packed_mask - 2.8) spec_decoding_generation_lengths - 3) spec_metadata - 3.1) num_tokens - 3.2) hidden_states_read_indices, hidden_states_write_indices - 3.3) is_first_draft - ''' - batch_size = attn_metadata.num_seqs - - # 1) Prepare the position_ids - num_accepted_draft_tokens = spec_metadata.num_accepted_draft_tokens[: - batch_size] - seq_lens = attn_metadata.seq_lens_cuda[:batch_size] - # Calculate last accepted token indices - last_tokens_idx = torch.cumsum( - seq_lens, dim=0, - dtype=torch.long) - seq_lens + num_accepted_draft_tokens - position_start_idx = position_ids[0, - last_tokens_idx] + 1 # [batch_size] - self.position_ids_buffer[:batch_size, :-1] = position_start_idx.unsqueeze( - 1) + spec_tree_manager.spec_dec_position_offsets[0, 1:].unsqueeze( - 0) - 1 # exclude the root node - # 2) Prepare the attn_metadata - ## 2.1) kv_lens_cuda - attn_metadata.kv_lens_cuda[: - batch_size] -= seq_lens - num_accepted_draft_tokens - 1 - attn_metadata.kv_lens_cuda[:batch_size] += ( - self.max_total_draft_tokens + 1) - - ## 2.2) _seq_lens, _seq_lens_cuda - attn_metadata._seq_lens[:batch_size].fill_(self.max_total_draft_tokens + - 1) - attn_metadata._seq_lens_cuda[:batch_size].fill_( - self.max_total_draft_tokens + 1) - attn_metadata.on_update() - - ## 2.3) host_request_types - attn_metadata.host_request_types[:attn_metadata.num_contexts].fill_(1) - - ## 2.4) num_contexts - attn_metadata.num_contexts = 0 - - ## 2.5) use_spec_decoding - attn_metadata.use_spec_decoding = True - - ## 2.6) spec_decoding_position_offsets - ### attn_metadata.spec_decoding_position_offsets: [max_num_requests, max_total_draft_tokens + 1] - attn_metadata.spec_decoding_position_offsets[:batch_size, :self. - max_total_draft_tokens] = spec_tree_manager.spec_dec_position_offsets[ - 0, 1:self. - max_total_draft_tokens + - 1].unsqueeze( - 0 - ) - 1 # exclude the root node - attn_metadata.spec_decoding_position_offsets[:batch_size, self. - max_total_draft_tokens] = 0 # padding - - ## 2.7) spec_decoding_packed_mask - ### attn_metadata.spec_decoding_packed_mask: [max_num_requests, max_total_draft_tokens + 1, math.ceil((max_total_draft_tokens + 1) / 32)] - attn_metadata.spec_decoding_packed_mask[: - batch_size, :, :] = spec_tree_manager.spec_dec_packed_mask_for_drafter_model - - ## 2.8) spec_decoding_generation_lengths - ### attn_metadata.spec_decoding_generation_lengths: [max_num_requests] - attn_metadata.spec_decoding_generation_lengths[: - batch_size] = self.max_total_draft_tokens + 1 - - # 3) Update spec_metadata - ## 3.1) num_tokens - spec_metadata.num_tokens = batch_size * (self.max_total_draft_tokens + - 1) - ## 3.2) hidden_states_read_indices, hidden_states_write_indices - ### spec_metadata.hidden_states_read_indices: [self.max_num_tokens] - ### spec_metadata.hidden_states_write_indices: [self.max_num_tokens] - old_write_indices = spec_metadata.hidden_states_write_indices - start_idx = old_write_indices[ - last_tokens_idx] # [batch_size], already take the accepted tokens into account. - - ### shape: [batch_size, self.max_total_draft_tokens + 1] - hidden_states_read_indices_offset = spec_tree_manager.hidden_states_read_indices_offset_for_drafter_model[:self.max_total_draft_tokens + 1].repeat( - batch_size).reshape(batch_size, self.max_total_draft_tokens + 1) - hidden_states_read_indices_offset = hidden_states_read_indices_offset + start_idx.unsqueeze( - 1) - spec_metadata.hidden_states_read_indices[:batch_size * ( - self.max_total_draft_tokens + - 1)] = hidden_states_read_indices_offset.reshape(-1) - - hidden_states_write_offset = torch.arange( - 1, self.max_total_draft_tokens + 1 + 1, - device=position_ids.device).unsqueeze(0).repeat( - batch_size, 1) + start_idx.unsqueeze(1) - spec_metadata.hidden_states_write_indices[:batch_size * ( - self.max_total_draft_tokens + - 1)] = hidden_states_write_offset.reshape(-1) - - ## 3.3) is_first_draft - spec_metadata.eagle3_resource_manager.is_first_draft = False - spec_metadata.is_first_draft = False - - return diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 55c9b703ac30..79498baab45f 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -32,8 +32,7 @@ from ..attention_backend.trtllm import (AttentionBackend, TrtllmAttention, TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE -from ..pyexecutor.resource_manager import (BaseResourceManager, - ResourceManagerType) +from ..pyexecutor.resource_manager import ResourceManagerType if TYPE_CHECKING: from ..pyexecutor.guided_decoder import CapturableGuidedDecoder @@ -396,31 +395,23 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]): TrtllmAttention) or not xqa_supported def attention_need_spec_dec_mode( - self, - spec_resource_manager: Optional[BaseResourceManager], - is_draft_model: bool, - attention_backend: Type[AttentionBackend], - use_chain_drafter: bool, # CDL + 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: - spec_resource_manager: the resource manager for the spec-dec mode. is_draft_model: whether the model is a draft model. attention_backend: the attention backend. - use_chain_drafter: whether to use capturable drafting loops (CDL). For the target model, it is always False. """ 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, we need to enable it when we process multiple tokens at once. This occurs with - # the target model (verification) or on the first draft for CDL based speculation. - use_case_2 = not self.use_one_engine() and ( - not is_draft_model or - (spec_resource_manager is not None - and spec_resource_manager.is_first_draft - and use_chain_drafter)) and is_trtllm_attention + # 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) return use_case_1 or use_case_2 diff --git a/tensorrt_llm/_torch/speculative/model_drafter.py b/tensorrt_llm/_torch/speculative/model_drafter.py index 5eae9b7e44cc..33a497ab64f2 100644 --- a/tensorrt_llm/_torch/speculative/model_drafter.py +++ b/tensorrt_llm/_torch/speculative/model_drafter.py @@ -8,7 +8,6 @@ from tensorrt_llm._utils import nvtx_range, prefer_pinned from tensorrt_llm.logger import logger -from ..attention_backend.trtllm import TrtllmAttention from ..pyexecutor.guided_decoder import GuidedDecoder from ..pyexecutor.handle_logits import HandleLogits from ..pyexecutor.llm_request import LlmRequest, LlmRequestState @@ -92,14 +91,7 @@ def __init__( self.sampler = sampler self.guided_decoder = guided_decoder - self.use_static_draft_loop = draft_model_engine.model_is_wrapped - if self.use_static_draft_loop: - # TODO: enable sampling/guided decoding on static draft loop - assert guided_decoder is None - assert spec_config._allow_greedy_draft_tokens - assert spec_config.draft_len_schedule is None - - # Create accumulator for draft tokens in non-CDL mode + # Create accumulator for draft tokens self.draft_tokens_accumulator: Dict[int, List[int]] = {} # Initialize draft latency tracking for specDecodingStats @@ -176,47 +168,6 @@ def _create_accepted_tokens_request(self, request: LlmRequest, input_tokens) - num_accepted_tokens - 1 return new_request - def _get_previous_draft_request( - self, request: LlmRequest) -> Optional[LlmRequest]: - """Get the previous draft request for the given request.""" - if self.previous_draft_batch is None: - return None - for req in self.previous_draft_batch.all_requests(): - if req.py_request_id == request.py_request_id: - return req - return None - - def _create_accepted_tokens_request_for_trtllm_attn( - self, request: LlmRequest, input_tokens: Any, - num_accepted_tokens: int) -> LlmRequest: - """ - Create a chunked context request for accepted tokens. - Only applicable if the draft model needs to recompute KV cache for accepted tokens (e.g. eagle 3) - """ - # Pad input_tokens to max_draft_len - # We use max_draft_len instead of max_total_draft_tokens here, - # because at most max_draft_len draft tokens are accepted. - input_tokens.extend( - 0 for _ in range(self.max_draft_len - num_accepted_tokens)) - - # Reuse the previous draft request if it exists. - # This can reduce host overhead significantly. - draft_request = self._get_previous_draft_request(request) - if draft_request is not None: - generated_tokens = input_tokens[draft_request.py_prompt_len:] - draft_request.set_generated_tokens([generated_tokens]) - else: - draft_request = self._create_draft_request(request, input_tokens) - - draft_request.state = LlmRequestState.GENERATION_IN_PROGRESS - draft_request.py_num_accepted_draft_tokens = request.py_num_accepted_draft_tokens - draft_request.py_is_first_draft = True - # For tree decoding, we need to store the accepted tokens indices for these requests, - # which will be used to update the hidden_states_read_indices. - draft_request.py_num_accepted_draft_tokens_indices = request.py_num_accepted_draft_tokens_indices - - return draft_request - def _create_draft_request_for_request( self, request: LlmRequest) -> Optional[LlmRequest]: """Create a draft request based on the original request state.""" @@ -227,9 +178,6 @@ def _create_draft_request_for_request( request, self.disable_overlap_scheduler) - is_eagle_style = self.spec_config.spec_dec_mode.is_eagle3( - ) or self.spec_config.spec_dec_mode.is_mtp_eagle() - # First time seeing this request - context request num_overlap_tokens = 0 if self.disable_overlap_scheduler else 1 if request.max_beam_num_tokens - 1 + num_overlap_tokens == request.py_prompt_len: @@ -240,12 +188,6 @@ def _create_draft_request_for_request( assert num_draft_tokens == 0 return self._create_context_request(request, input_tokens) - # For TRTLLM attention backend, we need to create a generation request for both no tokens accepted and tokens accepted - elif issubclass(self.draft_model_engine.attn_backend, TrtllmAttention - ) and self.use_static_draft_loop and is_eagle_style: - return self._create_accepted_tokens_request_for_trtllm_attn( - request, input_tokens, num_accepted_tokens) - # No tokens accepted - generation request. This only applies to speculation algorithms # that need to recompute KV cache for accepted tokens like eagle3. elif num_accepted_tokens == 0 or not self.spec_config.spec_dec_mode.needs_kv_cache_recompute( @@ -290,11 +232,10 @@ def _prepare_draft_batch( ScheduledRequests: The prepared draft batch """ try: - for req in scheduled_requests.all_requests(): - draft_model = self.draft_model_engine.model.draft_model if self.use_static_draft_loop else self.draft_model_engine.model - if hasattr(draft_model.model, "d2t"): + draft_model = self.draft_model_engine.model + if hasattr(draft_model.model, "d2t"): + for req in scheduled_requests.all_requests(): req.d2t = draft_model.model.d2t.data - req.py_draft_use_greedy_sampling = self.use_static_draft_loop draft_batch = ScheduledRequests() @@ -353,8 +294,6 @@ def _should_disable_cuda_graph(self, is_first_draft_token: bool) -> bool: """Check if CUDA graph should be disabled for the current forward pass.""" if not is_first_draft_token: return False - if self.use_static_draft_loop: - return False return self.spec_config.spec_dec_mode.needs_kv_cache_recompute() @nvtx_range("forward_draft_model") @@ -383,10 +322,7 @@ def forward_draft_model( num_accepted_tokens_device=num_accepted_tokens_device, req_id_to_old_request=self.req_id_to_old_request) - # Handle d2t data if available. Static drafting loops should incorporate d2t - # in their implementations. - if not self.use_static_draft_loop and hasattr( - self.draft_model_engine.model.model, 'd2t'): + if hasattr(self.draft_model_engine.model.model, 'd2t'): outputs['d2t'] = self.draft_model_engine.model.model.d2t.data return outputs @@ -585,13 +521,13 @@ def _update_draft_tokens_for_target_inputs( draft_indices = [] target_indices = [] - for req_idx, request in enumerate(draft_batch.all_requests()): + for request in draft_batch.all_requests(): target_req = self.req_id_to_old_request[request.py_request_id] if target_req.state != LlmRequestState.GENERATION_IN_PROGRESS: # Skip prefill requests continue # Get the index of the draft/target tokens in the device tensor - draft_idx = req_idx if self.use_static_draft_loop else request.py_seq_slot + draft_idx = request.py_seq_slot target_idx = target_req.py_seq_slot draft_indices.append(draft_idx) target_indices.append(target_idx) @@ -662,41 +598,6 @@ def _setup_draft_batch_and_resources( self.draft_seq_slot_manager.prepare_resources(draft_batch) return draft_batch - def process_static_draft_outputs(self, outputs: dict[str, torch.Tensor] - | tuple[torch.Tensor, SampleState], - draft_batch: ScheduledRequests) -> None: - """ - Process outputs from static draft loop, update target requests, and clean up resources. - - Args: - outputs: The outputs from the draft model - draft_batch: The draft batch that was processed - """ - - if isinstance(outputs, dict): - draft_tokens_host = outputs["new_draft_tokens"].cpu() - draft_logits = outputs["draft_logits"] - else: - draft_logits = outputs[0] - draft_tokens_host = outputs[1].host.new_tokens - outputs[1].sampler_event.synchronize() - - for req_idx, req in enumerate(draft_batch.all_requests()): - target_model_req = self.req_id_to_old_request[req.py_request_id] - if target_model_req.state != LlmRequestState.GENERATION_IN_PROGRESS: - # Chunked prefill request in progress; no need to append draft tokens - continue - target_model_req.py_draft_tokens = [] - py_draft_logits = [] - for token_idx in range(self.max_total_draft_tokens): - target_model_req.py_draft_tokens.append( - draft_tokens_host[token_idx][req_idx]) - py_draft_logits.append(draft_logits[token_idx][req_idx]) - - # The overlap scheduler doesn't support rejection sampling yet, so we don't update the py_draft_logits to get it fallback to greedy sampling. - if self.disable_overlap_scheduler: - target_model_req.py_draft_logits = torch.stack(py_draft_logits) - def process_dynamic_draft_outputs( self, outputs: Any, @@ -766,18 +667,14 @@ def _process_previous_draft_results( current_req_id_to_old_request = self.req_id_to_old_request # Set req_id_to_old_request for the previous batch, - # this will be used in process_static_draft_outputs and process_dynamic_draft_outputs + # this will be used in process_dynamic_draft_outputs self.req_id_to_old_request = { req.py_request_id: req for req in self.previous_scheduled_batch.all_requests() } - if self.use_static_draft_loop: - self.process_static_draft_outputs(self.previous_draft_outputs, - self.previous_draft_batch) - elif self.previous_draft_outputs is not None: - self.process_dynamic_draft_outputs(self.previous_draft_outputs, - cleanup_resources=False) + self.process_dynamic_draft_outputs(self.previous_draft_outputs, + cleanup_resources=False) self.req_id_to_old_request = current_req_id_to_old_request @@ -905,36 +802,7 @@ def generate_draft_tokens_with_overlap( self._process_previous_draft_results(resource_manager) num_draft_reqs = len(draft_batch.all_requests()) - if self.use_static_draft_loop: - # Only update target inputs, cleanup will be done in executor loop - self._update_draft_tokens_for_target_inputs( - target_inputs, - outputs["new_draft_tokens"], - draft_position=0, - draft_length=self.max_draft_len, - draft_batch=draft_batch) - - new_tokens_host = outputs["new_draft_tokens"].to(device='cpu', - non_blocking=True) - sampler_event = torch.cuda.Event() - sampler_event.record() - - sample_state = SampleState( - requests=draft_batch.all_requests(), - device=SampleStateTensors( - new_tokens=outputs["new_draft_tokens"]), - host=SampleStateTensors(new_tokens=new_tokens_host), - sampler_event=sampler_event) - - # Store current batch for processing in next iteration - self.previous_draft_batch = draft_batch - self.previous_draft_outputs = (outputs["draft_logits"], - sample_state) - self.previous_scheduled_batch = scheduled_batch - return - - # Handle guided decoder and sampling for non-static loop if self.guided_decoder is not None: self.guided_decoder.add_batch(draft_batch) self.guided_decoder.execute(outputs['logits'], @@ -1000,23 +868,11 @@ def prepare_draft_tokens( self.update_cur_draft_layer_idx( 0, resource_manager ) # Update the current draft layer index in the resource manager. - # Initial forward pass. May do the complete drafting loop - # if use_static_draft_loop is set. + # Initial forward pass. outputs = self.forward_draft_model(draft_batch, resource_manager, is_first_draft_token=True) - if self.use_static_draft_loop: - self.process_static_draft_outputs(outputs, draft_batch) - # Clean up draft_seq_slot_manager resources - for req in draft_batch.all_requests(): - self.draft_seq_slot_manager.free_resources(req) - # Record draft latency before returning - draft_end_time = time.time() - self.last_draft_latency_ms = (draft_end_time - - draft_start_time) * 1e3 - return - if self.guided_decoder is not None: self.guided_decoder.add_batch(draft_batch) self.guided_decoder.execute(outputs['logits'], diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6a1b0ec24ce3..e42a52744f62 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1873,8 +1873,6 @@ class DecodingBaseConfig(StrictBaseModel): "any of these penalties is rejected at admission rather than silently decoded " "without them.") - # If set, drafting is allowed to use chain drafter. - _allow_chain_drafter: bool = PrivateAttr(True) # If set, drafting uses greedy sampling, irrespective of sampling parameters. _allow_greedy_draft_tokens: bool = PrivateAttr(True) # Internal: record decoding_type alias used during parsing (for warnings). diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 7d44477e520e..d3199c1e42a0 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -216,7 +216,6 @@ def _make_forward_only_engine( ) engine.runtime_draft_len = 0 engine.attn_backend = None - engine.model_is_wrapped = False engine.original_max_draft_len = 0 engine.original_max_total_draft_tokens = 0 engine._spec_dec_max_total_draft_tokens = 0 diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index 089ef80645f7..6ee84e35f199 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -559,7 +559,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_tree=is_spec_dec_tree, max_draft_len=max_total_draft_tokens, max_total_draft_tokens=max_total_draft_tokens, - model_is_wrapped=False, spec_tree_manager=spec_tree_mgr, ) @@ -633,7 +632,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_dynamic_tree=is_tree_phase1, max_draft_len=gen_input_ids_1.size(-1) - 1, max_total_draft_tokens=gen_input_ids_1.size(-1) - 1, - model_is_wrapped=False, spec_metadata=spec_metadata_phase1, spec_tree_manager=spec_tree_mgr_phase1) @@ -708,7 +706,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_dynamic_tree=is_tree_ref, max_draft_len=gen_input_ids_ref.size(-1) - 1, max_total_draft_tokens=gen_input_ids_ref.size(-1) - 1, - model_is_wrapped=False, spec_metadata=spec_metadata_ref, spec_tree_manager=spec_tree_mgr_ref) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py deleted file mode 100644 index 78acbde58cb9..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py +++ /dev/null @@ -1,455 +0,0 @@ -import math -import os -import sys - -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.speculative.drafting_loops import StaticTreeDraftingLoopWrapper -from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig - -sys.path.append(os.path.join(os.path.dirname(__file__), "..")) - - -class DummyModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.model_config = None - self.config = None - self.model = {} - self.model_is_wrapped = True - - def forward(self, *args, **kwargs) -> torch.Tensor: - pass - - -def test_draft_token_static_tree_prepare_for_generation(): - # Fix parameters - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" # It will not actually be used. - use_dynamic_tree = False - max_new_tokens = 128 - kv_cache_manager = None - - # Create related object and run test - def run_test( - max_batch_size, - prepare_for_layer_idx, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - input_seq_lens_cuda, - input_kv_lens_cuda, - input_num_accepted_draft_tokens, - input_hidden_states_write_indices, - input_hidden_states_read_indices, - input_position_ids, - ref_position_ids, - ref_attn_metadata, - ref_spec_metadata, - ): - # 1) Create attention metadata - attn_metadata = TrtllmAttentionMetadata( - max_num_requests=max_batch_size, - max_num_tokens=max_new_tokens, - kv_cache_manager=kv_cache_manager, - ) - - # Set initial values - attn_metadata._seq_lens_cuda = input_seq_lens_cuda # set from input - attn_metadata.kv_lens_cuda = input_kv_lens_cuda # set from input - attn_metadata._seq_lens = torch.zeros([max_batch_size], device="cpu") - attn_metadata.host_request_types = torch.zeros([max_batch_size], device="cuda") - attn_metadata.spec_decoding_position_offsets = torch.zeros( - [max_batch_size, max_total_draft_tokens + 1], - dtype=torch.int, - device="cuda", - ) - attn_metadata.spec_decoding_packed_mask = torch.zeros( - [ - max_batch_size, - max_total_draft_tokens + 1, - math.ceil((max_total_draft_tokens + 1) / 32), - ], - dtype=torch.int, - device="cuda", - ) - attn_metadata.spec_decoding_generation_lengths = torch.zeros( - [max_batch_size], - dtype=torch.int, - device="cuda", - ) - - # 2) Create spec metadata - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - speculative_model=eagle_model_dir, - eagle3_one_model=False, - eagle_choices=eagle_choices, - use_dynamic_tree=use_dynamic_tree, - ) - - eagle3_resource_manager = Eagle3ResourceManager( - spec_config, - torch.bfloat16, - 1024, - max_batch_size, - max_new_tokens, - max_new_tokens, - ) - - spec_tree_manager = SpecTreeManager( - max_num_requests=max_batch_size, - use_dynamic_tree=spec_config.use_dynamic_tree, - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - eagle_choices=spec_config.eagle_choices, - dynamic_tree_max_topK=spec_config.dynamic_tree_max_topK, - ) - - spec_metadata = Eagle3SpecMetadata( - max_draft_len=spec_config.max_draft_len, - spec_dec_mode=spec_config.spec_dec_mode, - max_num_requests=max_batch_size, - num_layers=32, - hidden_size=1024, - max_num_tokens=max_new_tokens, - dtype=torch.bfloat16, - is_draft_model=True, - eagle3_resource_manager=eagle3_resource_manager, - layers_to_capture=spec_config.eagle3_layers_to_capture, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - eagle_choices=spec_config.eagle_choices, - is_spec_dec_tree=spec_config.eagle_choices is not None or spec_config.use_dynamic_tree, - is_spec_dec_dynamic_tree=spec_config.use_dynamic_tree, - ) - - # Set initial values - spec_metadata.num_accepted_draft_tokens = input_num_accepted_draft_tokens # set from input - spec_metadata.num_tokens = 0 - spec_metadata.hidden_states_write_indices = ( - input_hidden_states_write_indices # set from input - ) - spec_metadata.hidden_states_read_indices = ( - input_hidden_states_read_indices # set from input - ) - - # 3) Create StaticTreeDraftingLoopWrapper - static_tree_drafting_loop_wrapper = StaticTreeDraftingLoopWrapper( - max_batch_size=max_batch_size, - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - draft_model=DummyModel(), - ) - - # 3) Run the function - static_tree_drafting_loop_wrapper.prepare_for_generation( - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager, - position_ids=input_position_ids, - ) - - # Compare input_ids and position_ids - print( - f"static_tree_drafting_loop_wrapper.position_ids_buffer: \ - {static_tree_drafting_loop_wrapper.position_ids_buffer}, \ - ref_output_position_ids: {ref_position_ids}" - ) - - # Compare the attention metadata - print( - f"attn_metadata.kv_lens_cuda: {attn_metadata.kv_lens_cuda}, \ - ref_attn_metadata.kv_lens_cuda: {ref_attn_metadata['kv_lens_cuda']}" - ) - print( - f"attn_metadata._seq_lens: {attn_metadata._seq_lens}, \ - ref_attn_metadata._seq_lens: {ref_attn_metadata['_seq_lens']}" - ) - print( - f"attn_metadata._seq_lens_cuda: {attn_metadata._seq_lens_cuda}, \ - ref_attn_metadata._seq_lens_cuda: {ref_attn_metadata['_seq_lens_cuda']}" - ) - print( - f"attn_metadata.host_request_types: {attn_metadata.host_request_types}, \ - ref_attn_metadata.host_request_types: {ref_attn_metadata['host_request_types']}" - ) - print( - f"attn_metadata.num_contexts: {attn_metadata.num_contexts}, \ - ref_attn_metadata.num_contexts: {ref_attn_metadata['num_contexts']}" - ) - print( - f"attn_metadata.spec_decoding_position_offsets: {attn_metadata.spec_decoding_position_offsets}, \ - ref_attn_metadata.spec_decoding_position_offsets: {ref_attn_metadata['spec_decoding_position_offsets']}" - ) - print( - f"attn_metadata.spec_decoding_packed_mask: {attn_metadata.spec_decoding_packed_mask}, \ - ref_attn_metadata.spec_decoding_packed_mask: {ref_attn_metadata['spec_decoding_packed_mask']}" - ) - print( - f"attn_metadata.spec_decoding_generation_lengths: {attn_metadata.spec_decoding_generation_lengths}, \ - ref_attn_metadata.spec_decoding_generation_lengths: {ref_attn_metadata['spec_decoding_generation_lengths']}" - ) - - # Compare the spec metadata - print( - f"spec_metadata.num_tokens: {spec_metadata.num_tokens}, \ - ref_spec_metadata.num_tokens: {ref_spec_metadata['num_tokens']}" - ) - print( - f"spec_metadata.hidden_states_read_indices: {spec_metadata.hidden_states_read_indices}, \ - ref_spec_metadata.hidden_states_read_indices: {ref_spec_metadata['hidden_states_read_indices']}" - ) - print( - f"spec_metadata.hidden_states_write_indices: {spec_metadata.hidden_states_write_indices}, \ - ref_spec_metadata.hidden_states_write_indices: {ref_spec_metadata['hidden_states_write_indices']}" - ) - - assert torch.all(static_tree_drafting_loop_wrapper.position_ids_buffer == ref_position_ids) - assert torch.all(attn_metadata.kv_lens_cuda == ref_attn_metadata["kv_lens_cuda"]) - assert torch.all(attn_metadata._seq_lens == ref_attn_metadata["_seq_lens"]) - assert torch.all(attn_metadata._seq_lens_cuda == ref_attn_metadata["_seq_lens_cuda"]) - assert torch.all( - attn_metadata.host_request_types == ref_attn_metadata["host_request_types"] - ) - assert torch.all( - torch.tensor(attn_metadata.num_contexts) - == torch.tensor(ref_attn_metadata["num_contexts"]) - ) - assert torch.all( - attn_metadata.spec_decoding_generation_lengths - == ref_attn_metadata["spec_decoding_generation_lengths"] - ) - assert torch.all( - attn_metadata.spec_decoding_position_offsets[:max_batch_size, :].reshape(-1) - == ref_attn_metadata["spec_decoding_position_offsets"].reshape(-1) - ) - assert torch.all( - attn_metadata.spec_decoding_packed_mask[:max_batch_size, :, :].reshape(-1) - == ref_attn_metadata["spec_decoding_packed_mask"].reshape(-1) - ) - assert torch.all( - torch.tensor(spec_metadata.num_tokens) == torch.tensor(ref_spec_metadata["num_tokens"]) - ) - - output_hidden_states_read_indices = spec_metadata.hidden_states_read_indices[ - : max_batch_size * (max_total_draft_tokens + 1) - ].reshape(max_batch_size, max_total_draft_tokens + 1) - assert torch.all( - # We do not compare the last element of the hidden_states_read_indices, because it is padding. - output_hidden_states_read_indices[:, :-1] - == ref_spec_metadata["hidden_states_read_indices"][:, :-1] - ) - - output_hidden_states_write_indices = spec_metadata.hidden_states_write_indices[ - : max_batch_size * (max_total_draft_tokens + 1) - ].reshape(max_batch_size, max_total_draft_tokens + 1) - assert torch.all( - output_hidden_states_write_indices == ref_spec_metadata["hidden_states_write_indices"] - ) - - ##### CASE 1 static tree, batch size = 1, prefill, prepare_for_layer_idx = 1 ############# - max_batch_size = 1 - prepare_for_layer_idx = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - prompt_len_1 = 15 - input_position_ids = torch.arange(prompt_len_1, dtype=torch.int32, device="cuda").reshape( - 1, prompt_len_1 - ) - input_seq_lens_cuda = torch.tensor([prompt_len_1], dtype=torch.int32, device="cuda") - input_kv_lens_cuda = torch.tensor([prompt_len_1], dtype=torch.int32, device="cuda") - input_num_accepted_draft_tokens = torch.tensor( - [prompt_len_1 - 1], dtype=torch.int32, device="cuda" - ) - input_hidden_states_write_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - input_hidden_states_write_indices[:prompt_len_1] = torch.arange( - prompt_len_1, dtype=torch.long, device="cuda" - ) - input_hidden_states_read_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - - ref_position_ids = torch.tensor( - [[15, 15, 15, 16, 16, 16, 16, 16, 16, 17, 17, 17, 0]], dtype=torch.int32, device="cuda" - ) - - ref_attn_metadata = {} - # prompt_len_1 + max_total_draft_tokens + 1 - ref_attn_metadata["kv_lens_cuda"] = torch.tensor([28], dtype=torch.int32, device="cuda") - - # max_total_draft_tokens + 1 - ref_attn_metadata["_seq_lens"] = torch.tensor([13], dtype=torch.int32, device="cpu") - - # max_total_draft_tokens + 1 - ref_attn_metadata["_seq_lens_cuda"] = torch.tensor([13], dtype=torch.int32, device="cuda") - - ref_attn_metadata["host_request_types"] = torch.tensor([0], dtype=torch.int32, device="cuda") - ref_attn_metadata["num_contexts"] = 0 - ref_attn_metadata["spec_decoding_position_offsets"] = torch.tensor( - [0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) # [max_batch_size * (max_total_draft_tokens + 1)] - ref_attn_metadata["spec_decoding_packed_mask"] = torch.tensor( - [1, 2, 4, 9, 17, 33, 66, 130, 260, 521, 1041, 2114, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) # [max_batch_size * (max_total_draft_tokens + 1) * 1] - ref_attn_metadata["spec_decoding_generation_lengths"] = torch.tensor( - [13], dtype=torch.int32, device="cuda" - ) - - ref_spec_metadata = {} - ref_spec_metadata["num_tokens"] = 13 - ref_spec_metadata["hidden_states_read_indices"] = torch.tensor( - [[14, 14, 14, 15, 15, 15, 16, 16, 17, 18, 19, 21, 0]], dtype=torch.int32, device="cuda" - ) # [max_batch_size, max_total_draft_tokens + 1] - ref_spec_metadata["hidden_states_write_indices"] = torch.tensor( - [[15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]], dtype=torch.int32, device="cuda" - ) # [max_batch_size, max_total_draft_tokens + 1] - - run_test( - max_batch_size, - prepare_for_layer_idx, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - input_seq_lens_cuda, - input_kv_lens_cuda, - input_num_accepted_draft_tokens, - input_hidden_states_write_indices, - input_hidden_states_read_indices, - input_position_ids, - ref_position_ids, - ref_attn_metadata, - ref_spec_metadata, - ) - - ##### CASE 2 static tree, batch size = 2, one prefill, one decode, prepare_for_layer_idx = 1 ##### - max_batch_size = 2 - prepare_for_layer_idx = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - prompt_len_1 = 15 # prefill - prompt_len_2 = 18 - seq_len_2 = ( - 3 + 1 - ) # accepted 2 draft tokens. For the 0-th drafter layer, the sequence length will be pad to max_draft_len + 1 - - input_position_ids_1 = torch.arange(prompt_len_1, dtype=torch.int32, device="cuda").reshape( - 1, prompt_len_1 - ) - input_position_ids_2 = torch.tensor([18, 19, 20, 21], dtype=torch.int32, device="cuda").reshape( - 1, max_draft_len + 1 - ) # for target model - input_position_ids = torch.cat([input_position_ids_1, input_position_ids_2], dim=1) - - input_seq_lens_cuda = torch.tensor([prompt_len_1, seq_len_2], dtype=torch.int32, device="cuda") - input_kv_lens_cuda = torch.tensor( - [prompt_len_1, prompt_len_2 + seq_len_2], dtype=torch.int32, device="cuda" - ) - input_num_accepted_draft_tokens = torch.tensor( - [prompt_len_1 - 1, 2], dtype=torch.int32, device="cuda" - ) # Suppose 2 are received. - input_hidden_states_write_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - input_hidden_states_write_indices[: prompt_len_1 + seq_len_2] = torch.arange( - prompt_len_1 + seq_len_2, dtype=torch.long, device="cuda" - ) - input_hidden_states_read_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - ref_position_ids = torch.tensor( - [ - [15, 15, 15, 16, 16, 16, 16, 16, 16, 17, 17, 17, 0], - [21, 21, 21, 22, 22, 22, 22, 22, 22, 23, 23, 23, 0], - ], - dtype=torch.int32, - device="cuda", - ) - - ref_attn_metadata = {} - ref_attn_metadata["kv_lens_cuda"] = torch.tensor([28, 34], dtype=torch.int32, device="cuda") - ref_attn_metadata["_seq_lens"] = torch.tensor([13, 13], dtype=torch.int32, device="cpu") - ref_attn_metadata["_seq_lens_cuda"] = torch.tensor([13, 13], dtype=torch.int32, device="cuda") - ref_attn_metadata["host_request_types"] = torch.tensor([0, 0], dtype=torch.int32, device="cuda") - ref_attn_metadata["num_contexts"] = 0 - ref_attn_metadata["spec_decoding_position_offsets"] = torch.tensor( - [0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) - ref_attn_metadata["spec_decoding_packed_mask"] = torch.tensor( - [1, 2, 4, 9, 17, 33, 66, 130, 260, 521, 1041, 2114, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) - ref_attn_metadata["spec_decoding_generation_lengths"] = torch.tensor( - [13, 13], dtype=torch.int32, device="cuda" - ) - - ref_spec_metadata = {} - ref_spec_metadata["num_tokens"] = 26 - ref_spec_metadata["hidden_states_read_indices"] = torch.tensor( - [ - [14, 14, 14, 15, 15, 15, 16, 16, 17, 18, 19, 21, 0], - [17, 17, 17, 18, 18, 18, 19, 19, 20, 21, 22, 24, 0], - ], - dtype=torch.int32, - device="cuda", - ) - ref_spec_metadata["hidden_states_write_indices"] = torch.tensor( - [ - [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27], - [18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30], - ], - dtype=torch.int32, - device="cuda", - ) - - run_test( - max_batch_size, - prepare_for_layer_idx, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - input_seq_lens_cuda, - input_kv_lens_cuda, - input_num_accepted_draft_tokens, - input_hidden_states_write_indices, - input_hidden_states_read_indices, - input_position_ids, - ref_position_ids, - ref_attn_metadata, - ref_spec_metadata, - ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py deleted file mode 100644 index fd20a445d035..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py +++ /dev/null @@ -1,939 +0,0 @@ -import os -import sys -import unittest - -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm._torch.speculative.drafting_loops import StaticTreeDraftingLoopWrapper -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm.llmapi import Eagle3DecodingConfig - -sys.path.append(os.path.join(os.path.dirname(__file__), "..")) - - -class DummyModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.model_config = None - self.config = None - self.model = {} - - def forward(self, *args, **kwargs) -> torch.Tensor: - pass - - -def test_draft_token_static_tree_sampling(): - # Fix parameters - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" # It will not actually be used. - use_dynamic_tree = False - - # Create related object and run test - def run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - use_cuda_graph, - ref_new_tokens, - ): - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - speculative_model=eagle_model_dir, - eagle3_one_model=False, - eagle_choices=eagle_choices, - use_dynamic_tree=use_dynamic_tree, - ) - spec_tree_manager = SpecTreeManager( - max_num_requests=max_batch_size, - use_dynamic_tree=spec_config.use_dynamic_tree, - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - eagle_choices=spec_config.eagle_choices, - dynamic_tree_max_topK=spec_config.dynamic_tree_max_topK, - ) - - # Create the chain drafter - tree_drafter = StaticTreeDraftingLoopWrapper( - max_batch_size=max_batch_size, - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - draft_model=DummyModel(), - ) - - sampled_tokens = tree_drafter.sample( - logits=logits, - max_top_k=spec_tree_manager.max_top_k, - ) - - tree_drafter.extract_real_draft_tokens( - cur_draft_idx=draft_layer_id, - batch_size=max_batch_size, - new_draft_tokens=sampled_tokens, - use_cuda_graph=use_cuda_graph, # use torch op or not - spec_tree_manager=spec_tree_manager, - ) - - real_new_draft_tokens = tree_drafter.draft_tokens_buffer[:max_batch_size, :] - - print(f"ref_new_tokens.shape: {ref_new_tokens.shape}, ref_new_tokens: {ref_new_tokens}") - print( - f"real_new_draft_tokens.shape: {real_new_draft_tokens.shape}, output_tokens: {real_new_draft_tokens}" - ) - assert torch.all(real_new_draft_tokens == ref_new_tokens) - - ################## CASE 1 static tree, batch size = 1, draft_layer_id = 0 ########################## - max_batch_size = 1 - draft_layer_id = 0 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - logits = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [4, 1, 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 2 static tree, batch size = 1, draft_layer_id = 1 ########################## - max_batch_size = 1 - draft_layer_id = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([0, 1, 2], device="cuda") - logits[:, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top2 indices = [7, 1] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 4, 1, 9, 7, 1, 9, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 3 static tree, batch size = 1, draft_layer_id = 2 ########################## - max_batch_size = 1 - draft_layer_id = 2 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([3, 4, 6], device="cuda") - logits[:, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top1 indices = [4] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top1 indices = [7] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 7, 9, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 4 static tree, batch size = 2, draft_layer_id = 0 ########################## - max_batch_size = 2 - draft_layer_id = 0 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - [0.1, 0.3, 1.1, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 2, 9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [4, 1, 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [4, 2, 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 5 static tree, batch size = 2, draft_layer_id = 1 ########################## - max_batch_size = 2 - draft_layer_id = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([0, 1, 2], device="cuda") - - logits[0, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top2 indices = [7, 1] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - - logits[1, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 0.6, 1.2, 0.7, 0.8, 1.0, 0.9], # top3 indices = [5, 1, 8] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 1.2, 0.8, 0.9, 1.0], # top2 indices = [6, 1] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.2, 1.0], # top1 indices = [8] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 4, 1, 9, 7, 1, 9, 0, 0, 0, 0], - [0, 0, 0, 5, 1, 8, 6, 1, 8, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 6 static tree, batch size = 2, draft_layer_id = 2 ########################## - max_batch_size = 2 - draft_layer_id = 2 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([3, 4, 6], device="cuda") - - logits[0, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top1 indices = [4] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top1 indices = [7] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - - logits[1, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 0.6, 1.2, 0.7, 0.8, 0.9, 1.0], # top1 indices = [5] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.2, 1.0], # top1 indices = [8] - [1.2, 0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0], # top1 indices = [0] - ], - device="cuda", - ) - - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 7, 9, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 8, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 7 static tree, batch size = 1, draft_layer_id = 0, bigger tree ########################## - max_batch_size = 1 - draft_layer_id = 0 - max_total_draft_tokens = 63 - max_draft_len = 4 - eagle_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], - [0, 0, 0, 0], - [0, 1, 1], - [0, 0, 6], - [0, 3, 0], - [5, 0], - [1, 3], - [0, 0, 7], - [0, 0, 8], - [0, 0, 9], - [6, 0], - [0, 4, 0], - [1, 4], - [7, 0], - [0, 1, 2], - [2, 0, 0], - [3, 1], - [2, 2], - [8, 0], - [0, 5, 0], - [1, 5], - [1, 0, 1], - [0, 2, 1], - [9, 0], - [0, 6, 0], - [0, 0, 0, 1], - [1, 6], - [0, 7, 0], - ] # mc_sim_7b_63 - - logits = torch.tensor( - [ - [ - 1.1, - 1.9, - 0.7, - 0.3, - 4.4, - 4.3, - 4.9, - 1.2, - 3.9, - 4.7, - 1.4, - 2.5, - 3.2, - 0.8, - 2.0, - 3.0, - 1.8, - 2.3, - 4.2, - 1.3, - ] - ], - device="cuda", - ) - - ref_new_tokens = torch.zeros((max_batch_size, max_total_draft_tokens + 1), device="cuda") - - ref_new_tokens[0, :10] = torch.tensor( - [ - [6, 9, 4, 5, 18, 8, 12, 15, 11, 17], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 8 static tree, batch size = 1, draft_layer_id = 1, bigger tree ########################## - max_batch_size = 1 - draft_layer_id = 1 - max_total_draft_tokens = 63 - max_draft_len = 4 - eagle_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], - [0, 0, 0, 0], - [0, 1, 1], - [0, 0, 6], - [0, 3, 0], - [5, 0], - [1, 3], - [0, 0, 7], - [0, 0, 8], - [0, 0, 9], - [6, 0], - [0, 4, 0], - [1, 4], - [7, 0], - [0, 1, 2], - [2, 0, 0], - [3, 1], - [2, 2], - [8, 0], - [0, 5, 0], - [1, 5], - [1, 0, 1], - [0, 2, 1], - [9, 0], - [0, 6, 0], - [0, 0, 0, 1], - [1, 6], - [0, 7, 0], - ] # mc_sim_7b_63 - - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 20), device="cuda") - set_indices = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], device="cuda") - - logits[0, set_indices, :] = torch.tensor( - [ - [ - 3.8, - 2.5, - 1.9, - 0.2, - 4.5, - 0.9, - 0.5, - 3.9, - 4.0, - 2.9, - 2.7, - 0.7, - 2.8, - 2.1, - 1.2, - 1.4, - 3.3, - 3.7, - 2.4, - 5.4, - ], - [ - 3.2, - 5.4, - 0.1, - 4.6, - 1.9, - 4.8, - 3.8, - 0.9, - 5.3, - 3.7, - 4.4, - 0.4, - 4.3, - 1.3, - 0.3, - 3.6, - 2.5, - 6.0, - 3.1, - 3.9, - ], - [ - 4.1, - 2.8, - 3.9, - 4.2, - 5.1, - 5.8, - 3.7, - 1.5, - 0.8, - 1.4, - 2.2, - 0.0, - 1.9, - 2.5, - 3.2, - 0.4, - 3.5, - 5.0, - 2.9, - 4.3, - ], - [ - 0.7, - 1.4, - 0.6, - 2.0, - 3.0, - 0.9, - 1.7, - 0.0, - 3.4, - 2.8, - 4.7, - 3.6, - 5.6, - 4.2, - 1.1, - 4.5, - 1.2, - 1.3, - 3.9, - 3.2, - ], - [ - 3.6, - 5.5, - 4.0, - 6.0, - 4.7, - 2.3, - 5.8, - 0.8, - 3.2, - 0.3, - 4.8, - 2.1, - 0.0, - 3.1, - 0.1, - 2.5, - 2.2, - 5.0, - 3.4, - 3.9, - ], - [ - 3.3, - 3.8, - 2.5, - 2.4, - 6.0, - 5.5, - 3.6, - 0.6, - 1.8, - 2.6, - 4.7, - 1.2, - 3.9, - 4.0, - 4.8, - 5.0, - 1.1, - 1.9, - 0.4, - 0.3, - ], - [ - 3.1, - 1.7, - 0.0, - 4.6, - 3.7, - 3.2, - 2.8, - 4.7, - 2.7, - 3.8, - 6.0, - 4.2, - 1.1, - 3.4, - 5.5, - 2.1, - 2.6, - 1.8, - 0.1, - 1.3, - ], - [ - 3.6, - 5.1, - 2.3, - 5.3, - 3.5, - 1.8, - 2.7, - 2.4, - 3.2, - 3.0, - 4.8, - 5.9, - 2.9, - 1.4, - 5.7, - 0.9, - 1.0, - 1.3, - 2.0, - 3.9, - ], - [ - 3.9, - 0.2, - 4.5, - 1.5, - 1.4, - 0.1, - 4.3, - 0.3, - 3.2, - 3.4, - 0.4, - 3.8, - 4.8, - 3.5, - 5.6, - 1.9, - 2.8, - 4.1, - 1.0, - 5.0, - ], - [ - 5.9, - 1.8, - 3.5, - 0.3, - 1.7, - 2.5, - 3.0, - 3.3, - 2.1, - 4.3, - 6.0, - 0.1, - 0.5, - 0.9, - 4.0, - 5.4, - 1.1, - 2.7, - 5.3, - 3.4, - ], - ], - device="cuda", - ) - - ref_new_tokens = torch.zeros((max_batch_size, max_total_draft_tokens + 1), device="cuda") - ref_new_tokens[0, 10:38] = torch.tensor( - [ - [ - 19, - 4, - 8, - 7, - 0, - 17, - 16, - 9, - 12, - 10, # top-10 - 17, - 1, - 8, - 5, - 3, - 10, - 12, # top-7 - 5, - 4, - 17, # top-3 - 12, - 10, # top-2 - 3, # top-1 - 4, # top-1 - 10, # top-1 - 11, # top-1 - 14, # top-1 - 10, # top-1 - ] - ], - device="cuda", - ) # shape: [max_batch_size, num_new_draft_tokens] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index fb4406b13a20..e76cf655b01c 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -472,69 +472,35 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): @pytest.mark.parametrize( - "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp,use_hf_speculative_model", + "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,multi_batch,attention_dp,use_hf_speculative_model", [ - [True, "TRTLLM", True, False, False, False, True, False, False, False], - [True, "TRTLLM", True, False, False, False, False, False, False, False], - [False, "TRTLLM", True, False, False, False, True, False, False, False], - [ - False, "TRTLLM", True, False, False, False, False, False, False, - False - ], - [ - True, "FLASHINFER", True, False, False, False, True, False, False, - False - ], - [ - False, "FLASHINFER", True, False, False, False, True, False, False, - False - ], - [False, "TRTLLM", False, True, True, False, True, False, False, False], - [True, "TRTLLM", False, True, True, False, True, False, False, False], - [True, "TRTLLM", True, False, True, True, True, False, False, False], - [True, "TRTLLM", True, False, True, False, True, False, False, False], - [True, "TRTLLM", True, False, False, True, True, False, False, False], - [True, "TRTLLM", False, False, False, False, True, False, False, False], - [ - False, "TRTLLM", False, False, False, False, True, False, False, - False - ], - [True, "TRTLLM", False, False, False, False, False, True, False, False], - [True, "TRTLLM", False, False, False, False, False, True, True, False], - [ - False, "TRTLLM", False, False, False, False, False, True, False, - False - ], - [True, "TRTLLM", False, False, False, False, True, True, False, False], - [False, "TRTLLM", False, False, False, False, True, True, False, False], - [ - True, "TRTLLM", False, False, False, False, False, False, False, - False - ], - [ - False, "TRTLLM", False, False, False, False, False, False, False, - False - ], - [True, "TRTLLM", False, False, False, True, True, False, False, False], - [True, "TRTLLM", False, False, False, True, False, False, False, False], - [ - True, "FLASHINFER", False, False, False, False, True, False, False, - False - ], - [ - False, "FLASHINFER", False, False, False, False, True, False, False, - False - ], + [True, "TRTLLM", True, False, False, False, False, False, False], + [False, "TRTLLM", True, False, False, False, False, False, False], + [True, "FLASHINFER", True, False, False, False, False, False, False], + [False, "FLASHINFER", True, False, False, False, False, False, False], + [False, "TRTLLM", False, True, True, False, False, False, False], + [True, "TRTLLM", False, True, True, False, False, False, False], + [True, "TRTLLM", True, False, True, True, False, False, False], + [True, "TRTLLM", True, False, True, False, False, False, False], + [True, "TRTLLM", True, False, False, True, False, False, False], + [True, "TRTLLM", False, False, False, False, False, False, False], + [False, "TRTLLM", False, False, False, False, False, False, False], + [True, "TRTLLM", False, False, False, False, True, False, False], + [True, "TRTLLM", False, False, False, False, True, True, False], + [False, "TRTLLM", False, False, False, False, True, False, False], + [True, "TRTLLM", False, False, False, True, False, False, False], + [True, "FLASHINFER", False, False, False, False, False, False, False], + [False, "FLASHINFER", False, False, False, False, False, False, False], # Tests (mocked) speculative model auto-download from HuggingFace - [False, "TRTLLM", True, False, False, False, True, False, False, True], + [False, "TRTLLM", True, False, False, False, False, False, True], ]) @pytest.mark.high_cuda_memory @with_mocked_hf_download_for_single_gpu def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, disable_overlap_scheduler: bool, enable_block_reuse: bool, use_one_model: bool, enable_chunked_prefill: bool, - use_chain_drafter: bool, multi_batch: bool, - attention_dp: bool, use_hf_speculative_model: bool): + multi_batch: bool, attention_dp: bool, + use_hf_speculative_model: bool): if not use_one_model: pytest.skip("Two model Eagle3 is deprecated") @@ -586,7 +552,6 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, # Llama 3 does not support one model eagle. eagle3_one_model=use_one_model, ) - spec_config._allow_chain_drafter = use_chain_drafter # Create the LLM instance llm_spec = LLM(**llm_common_config, speculative_config=spec_config) @@ -986,56 +951,6 @@ def test_multi_eagle3(use_one_model: bool): pass -@pytest.mark.parametrize("disable_overlap_scheduler", [True, False]) -def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): - """Test CDL sampling with 2 requests and max_batch_size=2.""" - attn_backend = "TRTLLM" - enable_block_reuse = False - use_one_model = False - enable_chunked_prefill = False - - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - 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" - - max_batch_size = 1 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, - max_tokens=8192) - cuda_graph_config = CudaGraphConfig(batch_sizes=[1, 2, 4], - enable_padding=True) - - llm_common_config = dict( - model=target_model_dir, - 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, - max_seq_len=8192, - enable_chunked_prefill=enable_chunked_prefill, - ) - - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, - ) - - # Create the LLM instance - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - prompts = ["The president of the United States is"] - - sampling_params = SamplingParams(max_tokens=20, temperature=1.0, top_p=0.9) - llm_spec.generate(prompts, sampling_params) - llm_spec.shutdown() - - @pytest.mark.parametrize("use_dynamic_tree", [False, True], ids=["no_dynamic_tree", "dynamic_tree"]) @pytest.mark.parametrize("use_cuda_graph", [False, True]) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 529851947843..f1713d6c40df 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2024,7 +2024,7 @@ class TestPiecewiseCudaGraphCaptureDefaults: (invariants 2 and 3) clamps out-of-range entries to the reachable ceiling and never invents sizes beyond this list. 2. `_filter_piecewise_capture_num_tokens` caps the candidate list at - `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` -- + `max_batch_size * (max_seq_len - 1)` -- the largest forward-pass `num_tokens` the warmup builder can construct, since every in-flight request must leave room for at least one decode token. @@ -2315,42 +2315,6 @@ def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self): assert kept.count(128) == 1 assert unrecordable == [] - def test_piecewise_filter_subtracts_extra_decoding_steps(self): - """Subtract `num_extra_decoding_steps` from the ceiling. - - Drafting loops consume extra decode steps; the filter must mirror - the `max_seq_len - 1 - num_extra_decoding_steps` constraint - applied when warmup requests are built. Candidates above the - reduced ceiling are clamped down to it; nothing is appended. - """ - from tensorrt_llm._torch.pyexecutor.model_engine import \ - _filter_piecewise_capture_num_tokens - - candidates = [1, 2, 4, 8, 16, 32, 64, 100, 120] - # max_seq_len=128, batch=1, 5 extra decoding steps -> ceiling 122. - kept, unrecordable = _filter_piecewise_capture_num_tokens( - candidates, - max_num_tokens=128, - max_batch_size=1, - max_seq_len=128, - num_extra_decoding_steps=5, - ) - assert kept[-1] == 120 # nothing above the 122 ceiling to clamp - assert 120 in kept - assert unrecordable == [] - # Same setup with 9 extra decoding steps -> ceiling 118; 120 drops. - kept, unrecordable = _filter_piecewise_capture_num_tokens( - candidates, - max_num_tokens=128, - max_batch_size=1, - max_seq_len=128, - num_extra_decoding_steps=9, - ) - assert kept[-1] == 118 - assert 100 in kept - assert 120 not in kept - assert unrecordable == [120] - def test_piecewise_filter_does_not_double_append_ceiling(self): """Ceiling already present in candidates -> not duplicated.""" from tensorrt_llm._torch.pyexecutor.model_engine import \