# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import contextlib import functools import gc import inspect import math import os import weakref from abc import ABC, abstractmethod from contextlib import contextmanager from typing import (Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Type, Union, cast) import torch import torch._dynamo.config import tensorrt_llm.bindings.internal.userbuffers as ub from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.manager import LoraModelConfig from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._utils import (global_mpi_rank, is_trace_enabled, maybe_pin_memory, nvtx_range, prefer_pinned, release_gc, trace_func) from tensorrt_llm.bindings.internal import \ batch_manager as batch_manager_bindings from tensorrt_llm.inputs.multimodal import (MultimodalParams, MultimodalRuntimeData, _has_mm_payload_keys, check_mm_embed_cumsum_if_needed, strip_mm_data_for_generation) from tensorrt_llm.inputs.registry import (BaseMultimodalInputProcessor, create_input_processor) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) # isort: split from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping from tensorrt_llm.sampling_params import SamplingParams from ..attention.backends.interface import (AttentionMetadata, AttentionRuntimeFeatures) from ..attention.backends.trtllm import TrtllmAttentionMetadata from ..attention.backends.utils import get_attention_backend from ..autotuner import AutoTuner, autotune from ..compilation.backend import Backend from ..compilation.utils import capture_piecewise_cuda_graph from ..distributed import Distributed from ..distributed.communicator import init_pp_comm from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool from ..metadata import KVCacheParams from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.modeling_multimodal_mixin import (MultimodalModelMixin, _build_request_multimodal_input) from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.mamba.mamba2_metadata import Mamba2Metadata from ..moe.expert_statistic import ExpertStatistic from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer, MoeLoadBalancerIterContext) from ..peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture from ..speculative import (SpecMetadata, get_draft_kv_cache_manager, get_num_extra_kv_tokens, get_spec_metadata, prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) from ..speculative.eagle3 import Eagle3ResourceManager from ..speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_static_draft_len, update_draft_len from ..utils import (get_model_extra_attrs, get_per_request_prefill_cuda_graph_flag, helix_local_len, set_per_request_prefill_cuda_graph_flag, set_torch_compiling, torch_compiling, with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner from .config_utils import is_hybrid_linear from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, get_mrope_dummy_seq_slot) from .engine.cuda_graph import (filter_cuda_graph_batch_sizes, resolve_cuda_graph_batch_sizes) from .engine.lora import (LoraParamBuilder, make_cuda_graph_lora_manager, make_lora_model_config) from .engine.metadata import build_attention_metadata, update_spec_metadata from .engine.multimodal import (MultimodalItemScheduler, is_multimodal, mm_encoder_cache_enabled, setup_mm_encoder_attn_metadata) from .engine.runners import (apply_position_id_offset, get_all_rank_num_tokens, get_padding_params, get_position_id_offset, get_top_level_model, prepare_multimodal_indices, resolve_runner_type, set_spec_metadata_all_rank_num_tokens, ship_multimodal_indices) from .engine.runners.encoder import EncoderRunner, EncoderRunnerConfig from .engine.runners.encoder_decoder import (EncoderDecoderRunner, EncoderDecoderRunnerConfig) from .engine.runners.interface import (ModelRunner, PackedEncoderBatch, PackedModelRunner, RunnerDeps) from .engine.runners.no_kv_cache import NoKVCacheRunner, NoKVCacheRunnerConfig from .guided_decoder import CapturableGuidedDecoder from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, MambaHybridCacheManager) from .layerwise_nvtx_marker import LayerwiseNvtxMarker from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length from .model_loader import ModelLoader, _construct_checkpoint_loader from .resource_manager import (BaseResourceManager, KVCacheManager, ResourceManager, ResourceManagerType) from .sampler import SampleStateTensors from .sampler.ops.flashinfer import (warmup_sample_from_logits_op, warmup_sampling_module) from .sampler.sampler_common import SampleType from .scheduler import ScheduledRequests from .trace_log_utils import log_mem_snapshot from .workspace import EagerWorkspaceReclaimer def _get_context_prompt_lookahead_token(request: LlmRequest, chunk_end: int) -> int: """Prompt token immediately following a context chunk. Uses the live C++ ``mPromptLen``; ``py_prompt_len`` goes stale after a non-recompute preemption. """ if request.is_last_context_chunk: return INVALID_PROMPT_LOOKAHEAD_TOKEN return request.get_token(0, chunk_end) def resolve_mamba_metadata_cls(model: torch.nn.Module) -> Type[Mamba2Metadata]: """Resolve the model-specific Mamba metadata class with a default.""" return getattr(model, 'mamba_metadata_cls', None) or Mamba2Metadata def resolve_mrope_position_deltas_cache( model: Optional[torch.nn.Module]) -> Optional[torch.Tensor]: """The MRoPE delta cache held by ``model`` or by its draft model. ``None`` for every model that does not keep one, which is also how ``should_enable_overlap_headroom`` learns that the seat pool may be widened: the cache is sized from ``max_num_tokens`` rather than from the seat pool. """ cache = getattr(model, "mrope_position_deltas_cache", None) if cache is None: cache = getattr(getattr(model, "draft_model", None), "mrope_position_deltas_cache", None) return cache def _make_single_token_context_graph_batch( scheduled_requests: ScheduledRequests, is_multimodal_decode_compatible: Optional[Callable[[LlmRequest], bool]] = None, ) -> tuple[ScheduledRequests, frozenset[int]]: """Build a decode-shaped graph candidate for final one-token contexts. Multimodal rows remain fail-closed unless the engine proves that their one remaining prompt token is representable by the existing decode provider. """ if scheduled_requests.num_context_requests == 0: return scheduled_requests, frozenset() context_requests = scheduled_requests.context_requests_last_chunk if (scheduled_requests.encoder_requests or scheduled_requests.context_requests_chunking): return scheduled_requests, frozenset() for request in context_requests: if (request.context_chunk_size != 1 or request.context_remaining_length != 1 or request.context_current_position + 1 != request.py_prompt_len or request.py_beam_width != 1 or get_draft_token_length(request) > 0 or request.py_is_first_draft or request.is_context_only_request or request.is_generation_only_request or request.py_disaggregated_params is not None or request.py_mm_encoder_event is not None or (request.py_multimodal_data is not None and (is_multimodal_decode_compatible is None or not is_multimodal_decode_compatible(request)))): return scheduled_requests, frozenset() for request in scheduled_requests.generation_requests: if (request.py_beam_width != 1 or get_draft_token_length(request) > 0 or request.py_is_first_draft or request.py_disaggregated_params is not None): return scheduled_requests, frozenset() graph_batch = ScheduledRequests() graph_batch.generation_requests = list(context_requests) + list( scheduled_requests.generation_requests) graph_batch.paused_requests = list(scheduled_requests.paused_requests) promoted_context_request_ids = frozenset(request.py_request_id for request in context_requests) return graph_batch, promoted_context_request_ids class _PrefillCompiledModel(torch.nn.Module): """Share weights between eligible prefill/mixed and original eager paths. The prefill flag includes the all-rank attention-DP decision and capture ceiling. A decode-only rank must still compile when another rank prefills. """ def __init__(self, eager_model: torch.nn.Module, compiled_model: torch.nn.Module) -> None: """Keep eager and compiled entry points sharing the same model weights.""" super().__init__() self.eager_model = eager_model # The compiled callable references the same weights. Register only the # eager tree so state_dict(), children() and _apply() visit it once. object.__setattr__(self, "compiled_model", compiled_model) def named_modules( self, memo: Optional[set[torch.nn.Module]] = None, prefix: str = "", remove_duplicate: bool = True, ) -> Iterator[Tuple[str, torch.nn.Module]]: """Expose checkpoint-compatible module names for partial weight reloads.""" # Weight reloads match checkpoint prefixes against this traversal. # Hide the eager_model prefix even with remove_duplicate=False. yield from self.eager_model.named_modules(memo, prefix, remove_duplicate) def forward(self, *args: Any, **kwargs: Any) -> Any: """Use the compiled path only for globally eligible prefill batches.""" model = (self.compiled_model if get_per_request_prefill_cuda_graph_flag() else self.eager_model) return model(*args, **kwargs) def __getattr__(self, name: str) -> Any: """Delegate model-specific attributes to the original eager model.""" # Epilogues can access transformer attributes after forward returns. try: return super().__getattr__(name) except AttributeError: return getattr(super().__getattr__("eager_model"), name) class ModelEngine(ABC): @abstractmethod def get_max_num_sequences(self) -> int: raise NotImplementedError @abstractmethod def forward(self, scheduled_requests: ScheduledRequests, resource_manager: Optional[ResourceManager], new_tensors_device: Optional[SampleStateTensors], gather_context_logits: bool = False, cache_indirection_buffer: Optional[torch.Tensor] = None, num_accepted_tokens_device: Optional[torch.Tensor] = None): raise NotImplementedError def warmup(self, resource_manager: Optional[ResourceManager]) -> None: """ This method is called after runtime resources are initialized. The resource manager is absent for drivers that allocate none. Override to perform warmup actions: instantiating CUDA graphs, torch.compile, etc. """ return def _filter_piecewise_capture_num_tokens( candidate_num_tokens: list[int], max_num_tokens: int, max_batch_size: int, max_seq_len: int, ) -> 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)` clamping user-requested sizes above it down to the ceiling. Each in-flight request must leave room for at least one decode token, so the ceiling is the largest forward-pass `num_tokens` the warmup builder can construct. Candidates above the ceiling cannot be recorded; clamping them down to the ceiling preserves the user's intent (a requested 128 becomes 127 when only 127 is recordable) without inventing capture sizes the user never asked for. Appending sizes beyond the user's list is harmful: runtime padding rounds iterations up to the nearest captured size, so a far appended ceiling (e.g. 65536 over a list topping at 13914) would make every iteration in the gap execute the full ceiling shape. Returns `(kept, unrecordable)` where `kept` is sorted ascending and deduped, with above-ceiling candidates clamped to the ceiling. `unrecordable` is the sorted unique set of input entries above the 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)) piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) if piecewise_capacity_limit > 0: kept = sorted({ min(i, piecewise_capacity_limit) for i in candidate_num_tokens if 0 < i <= max_num_tokens }) else: kept = [] unrecordable = sorted({ i for i in candidate_num_tokens if max_capturable_num_tokens < i <= max_num_tokens }) return kept, unrecordable # BCG uses the same capture-bucket filtering semantics as PCG. _filter_prefill_capture_num_tokens = _filter_piecewise_capture_num_tokens _DEEP_GEMM_PDL_CONFIGURED = False # Arbitrary non-greedy params used to force the advanced-sampling CUDA graph # warmup capture path. NON_GREEDY_CAPTURE_SAMPLING_PARAMS = SamplingParams(temperature=0.7, top_k=50, top_p=0.9, min_p=0.05) def _configure_deep_gemm_pdl() -> None: global _DEEP_GEMM_PDL_CONFIGURED if _DEEP_GEMM_PDL_CONFIGURED: return from tensorrt_llm import deep_gemm deep_gemm.set_pdl(os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1") _DEEP_GEMM_PDL_CONFIGURED = True @contextlib.contextmanager def _moe_a2a_steady_state_budget_for_capture(): """Force the steady-state MoE all-to-all budget across CUDA-graph capture. The budget is a kernel launch argument, so it is frozen into each captured graph. Capture happens inside the warmup window, so without this a replay would keep warmup's relaxed deadline for the life of the process. """ _set_moe_a2a_warmup(False) try: yield finally: _set_moe_a2a_warmup(True) def _set_moe_a2a_warmup(in_warmup: bool) -> None: """Select the MoE all-to-all completion-flag budget for the current phase. No-op when the op is unavailable (older bindings). """ try: torch.ops.trtllm.moe_a2a_set_warmup(in_warmup) logger.info(f"moe_a2a completion-flag budget: in_warmup={in_warmup}") except (AttributeError, RuntimeError) as e: logger.warning( f"moe_a2a_set_warmup unavailable, the all-to-all timeout " f"budget was not switched: {type(e).__name__}: {e}") class PyTorchModelEngine(ModelEngine): def __init__( self, *, model_path: str, llm_args: TorchLlmArgs, mapping: Optional[Mapping] = None, attn_runtime_features: Optional[AttentionRuntimeFeatures] = None, dist: Optional[Distributed] = None, spec_config: Optional[DecodingBaseConfig] = None, is_draft_model: bool = False, model: Optional[torch.nn.Module] = None, checkpoint_loader: Optional[BaseCheckpointLoader] = None, model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): _configure_deep_gemm_pdl() self.forward_pass_callable = None self._cleanup_done = False self._runner: Optional[Union[ModelRunner, PackedModelRunner]] = None # Transitional snapshot for decoder capture and encoder scheduling. # Encoder graph resources and lifecycle remain entirely runner-owned. self._encoder_graph_shapes: frozenset[tuple[int, int]] = frozenset() self._encoder_graph_batch_sizes: tuple[int, ...] = () self._encoder_graph_pad_to_limit = False # Optional in-graph sampling hook, registered by PyExecutor. Called at # the tail of _forward_step -- i.e. right after the LM head, and inside # capture_forward_fn -- so that for graph-capturable sampling tiers the # sampling lands at the end of the captured forward graph instead of # costing a separate launch. Left None when the sampler does not opt in. self.sample_in_graph_callable = None # Stages (or clears) in-graph sampling state once the batch is settled. self._stage_in_graph_sampling = None if llm_args.encode_only and llm_args.mm_encoder_only: raise ValueError( "encode_only and mm_encoder_only are mutually exclusive.") ( max_beam_width, max_num_tokens, max_seq_len, max_batch_size, ) = llm_args.get_runtime_sizes() self.batch_size = max_batch_size self.max_num_tokens = max_num_tokens self.max_seq_len = max_seq_len self.max_beam_width = max_beam_width self.encoder_batch_size = (llm_args.encoder_max_batch_size if llm_args.encoder_max_batch_size is not None else self.batch_size) # The multimodal encoder token budget falls back to the LLM-side value # when unset. It may be raised after model load because atomic MM items # cannot be split. self.encoder_max_num_tokens = (llm_args.encoder_max_num_tokens if llm_args.encoder_max_num_tokens is not None else self.max_num_tokens) if checkpoint_loader is None: checkpoint_loader = _construct_checkpoint_loader( llm_args.checkpoint_loader, llm_args.checkpoint_format, mx_config=llm_args.mx_config, checkpoint_io_policy=llm_args.checkpoint_io_policy, load_format=llm_args.load_format, partial_model_loading=llm_args.is_partial_model_loading, ) self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) from ._util import (compute_max_num_sequences, resolved_kv_cache_manager_is_v2, should_enable_adp_dummy_fixes, should_enable_non_overlap_adp_forward_intent, should_enable_overlap_headroom, should_enable_scheduler_aware_adp_dummy) self._enable_adp_dummy_fixes = should_enable_adp_dummy_fixes(mapping) self.dist = dist self.llm_args = llm_args if dist is not None: ExpertStatistic.create(self.dist.rank) # Opt-in tiered sampling captured into the forward graph. Off by # default: it captures one extra graph per enabled tier, which costs # startup time and memory that deployments not bound by sampling # overhead should not pay. self.enable_in_graph_sampling = bool( getattr(llm_args, "enable_in_graph_sampling", False)) self.original_max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 self.original_max_total_draft_tokens = ( spec_config.tokens_per_gen_step - 1) if spec_config is not None else 0 # Saved before zeroing for draft models; used by update_spec_dec_param. self._spec_dec_max_total_draft_tokens = ( spec_config.max_total_draft_tokens if spec_config is not None else 0) # Dynamic tree draft loop produces up to K * max_draft_len tokens, # which may exceed max_total_draft_tokens. Use the larger value for # KV cache reservation only; verify/tree output stays at max_total_draft_tokens. if (spec_config is not None and getattr(spec_config, 'use_dynamic_tree', False) and getattr(spec_config, 'dynamic_tree_max_topK', 0) > 0): self.max_draft_loop_tokens = max( self.original_max_total_draft_tokens, spec_config.dynamic_tree_max_topK * spec_config.max_draft_len) else: self.max_draft_loop_tokens = self.original_max_total_draft_tokens # The draft model won't have any draft tokens attached to # generation requests when we invoke it autoregressively if spec_config is not None and is_draft_model: spec_config.max_draft_len = 0 spec_config.max_total_draft_tokens = 0 self.spec_config = spec_config self.is_spec_decode = spec_config is not None self.sparse_attention_config = None if is_draft_model else llm_args.sparse_attention_config self.enable_spec_decode = self.is_spec_decode self.is_draft_model = is_draft_model self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures( ) input_processor_kwargs = {} video_pruning_rate = llm_args.multimodal_config.video_pruning_rate if video_pruning_rate is not None: input_processor_kwargs['video_pruning_rate'] = video_pruning_rate self.input_processor = create_input_processor( model_path, tokenizer=None, checkpoint_format=llm_args.checkpoint_format, trust_remote_code=llm_args.trust_remote_code, **input_processor_kwargs) self.moe_load_balancer: Optional[MoeLoadBalancer] = None self.model_loader: Optional[ModelLoader] = None if model is None: lora_config: Optional[ LoraConfig] = None if is_draft_model else llm_args.lora_config # Keep the model_loader to support reloading the model weights later self.model_loader = ModelLoader( llm_args=llm_args, mapping=self.mapping, spec_config=self.spec_config, sparse_attention_config=self.sparse_attention_config, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, lora_config=lora_config, model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, ) # Open checkpoint and load the LLM module object. self.model, moe_load_balancer = self.model_loader.load( checkpoint_dir=model_path, checkpoint_loader=checkpoint_loader) if isinstance(moe_load_balancer, MoeLoadBalancer): self.moe_load_balancer = moe_load_balancer else: self.model = model self._validate_breakable_cuda_graph_compatibility() # In-graph sampling needs the full vocabulary: top-k / top-p over a # tensor-parallel shard would rank against a slice of the logits and # emit the wrong token, silently. The LM head gathers by default, but # the draft models of some speculation modes turn that off, so refuse # the fast path rather than sample a sharded row. if self.enable_in_graph_sampling and not getattr( self.model.model_config, "lm_head_gather_output", True): logger.warning( "Disabling enable_in_graph_sampling: this model's LM head does not " "gather its output, so the logits are sharded across tensor " "parallel ranks and cannot be sampled in-graph.") self.enable_in_graph_sampling = False pretrained_config = self.model.model_config.pretrained_config model_type = getattr(pretrained_config, "model_type", None) self._enable_scheduler_aware_adp_dummy = ( should_enable_scheduler_aware_adp_dummy( model_type, mapping, llm_args.disable_overlap_scheduler)) self._enable_non_overlap_adp_forward_intent = ( should_enable_non_overlap_adp_forward_intent( mapping, llm_args.disable_overlap_scheduler)) self._enable_overlap_headroom = should_enable_overlap_headroom( mapping, llm_args.disable_overlap_scheduler, kv_cache_manager_is_v2=resolved_kv_cache_manager_is_v2( llm_args.kv_cache_config, self.max_beam_width), is_hybrid=is_hybrid_linear(pretrained_config), has_mrope_delta_cache=resolve_mrope_position_deltas_cache( self.model) is not None) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, llm_args.disable_overlap_scheduler, enable_overlap_headroom=self._enable_overlap_headroom, ) 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'): self.model.extra_attrs = {} # Router Replay (R3) capturer owned by this engine (None when disabled or # for draft engines). Registered in the model extra attrs so the MoE # routing hook reaches this engine's capturer during its forward, the # same way ``moe_layers`` is looked up -- no process-wide state. self.route_capture = RouteCapture.create( rank=self.dist.rank if self.dist is not None else 0, model_engine=self, enabled=self.llm_args.enable_return_routed_experts, pp_size=self.mapping.pp_size, is_spec_decode=self.is_spec_decode, is_draft_model=self.is_draft_model) if self.route_capture is not None: self.model.extra_attrs[ROUTE_CAPTURE_ATTR] = self.route_capture # Every MM item-scheduling decision -- policy, capability, feature # validation, budget resolution -- lives in engine/multimodal.py; the # engine only copies back the three budgets that are external contract. mm_item_scheduler = MultimodalItemScheduler.maybe_create( llm_args=self.llm_args, model=self.model, input_processor=self.input_processor, encoder_max_num_tokens=self.encoder_max_num_tokens) self._mm_item_scheduler = mm_item_scheduler # `getattr`-read by py_executor.py and _util.py. self.mm_encoder_item_scheduling_enabled = mm_item_scheduler is not None self.mm_encoder_output_budget_bytes: Optional[int] = None if mm_item_scheduler is not None: # The raised encoder token budget is read back off the engine by # `_util.py`, and sizes the encoder metadata set up below. self.encoder_max_num_tokens = mm_item_scheduler.encoder_max_num_tokens self.mm_encoder_output_budget_bytes = mm_item_scheduler.output_budget_bytes # Absent, not None, when item scheduling is off: external readers # rely on the `getattr` default. self.bytes_per_mm_encoder_embedding = mm_item_scheduler.bytes_per_embedding setup_mm_encoder_attn_metadata( self.model, self.input_processor, self.encoder_max_num_tokens, mm_item_scheduler.attention_metadata_capacity if mm_item_scheduler is not None else None) if self.llm_args.enable_layerwise_nvtx_marker: layerwise_nvtx_marker = LayerwiseNvtxMarker() module_prefix = 'Model' if self.model.model_config and self.model.model_config.pretrained_config and self.model.model_config.pretrained_config.architectures: module_prefix = '|'.join( self.model.model_config.pretrained_config.architectures) layerwise_nvtx_marker.register_hooks(self.model, module_prefix) self.enable_attention_dp = self.model.model_config.mapping.enable_attention_dp self._disable_overlap_scheduler = self.llm_args.disable_overlap_scheduler self._torch_compile_backend = None self.dtype = self.model.config.torch_dtype self._init_model_capacity() self.cuda_graph_config = self.llm_args.cuda_graph_config self._is_encode_only = self.llm_args.encode_only if (isinstance(self.cuda_graph_config, EncodeCudaGraphConfig) and self._is_encoder_decoder_model()): logger.warning( "EncodeCudaGraphConfig is not supported for encoder-decoder " "models through cuda_graph_config. Use DecodeCudaGraphConfig " "for cuda_graph_config and configure encoder graphs through " "encoder_cuda_graph_config. Decoder CUDA graphs will be " "disabled.") self.cuda_graph_config = None cuda_graph_batch_sizes = self.cuda_graph_config.batch_sizes if self.cuda_graph_config else CudaGraphConfig.model_fields[ 'batch_sizes'].default cuda_graph_padding_enabled = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[ 'enable_padding'].default self._cuda_graph_padding_enabled = cuda_graph_padding_enabled decode_tokens_per_request = 1 + self.original_max_total_draft_tokens self._cuda_graph_batch_sizes = filter_cuda_graph_batch_sizes( cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens, decode_tokens_per_request, self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else [] self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if self._cuda_graph_batch_sizes else 0) self.torch_compile_config = self.llm_args.torch_compile_config self.prefill_cuda_graph_backend = self.llm_args.prefill_cuda_graph_backend torch_compile_enabled = bool(self.torch_compile_config is not None) torch_compile_fullgraph = self.torch_compile_config.enable_fullgraph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_fullgraph'].default torch_compile_inductor_enabled = self.torch_compile_config.enable_inductor if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_inductor'].default torch_compile_piecewise_cuda_graph = (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE) torch_compile_enable_userbuffers = self.torch_compile_config.enable_userbuffers if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'enable_userbuffers'].default torch_compile_max_num_streams = self.torch_compile_config.max_num_streams if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ 'max_num_streams'].default self._torch_compile_enabled = torch_compile_enabled self._torch_compile_piecewise_cuda_graph = torch_compile_piecewise_cuda_graph self._torch_compile_prefill_only = False prefill_cuda_graph_num_tokens = self.llm_args.prefill_capture_num_tokens if prefill_cuda_graph_num_tokens is None: prefill_cuda_graph_num_tokens = cuda_graph_batch_sizes or [] 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, )) 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)=" f"{max(0, self.batch_size * (self.max_seq_len - 1))}. " f"Clamping them to the ceiling; raise max_seq_len for larger graphs." ) try: use_ub_for_nccl = ( self.llm_args.allreduce_strategy == "NCCL_SYMMETRIC" and self._init_userbuffers(self.model.config.hidden_size)) if self._torch_compile_enabled: set_torch_compiling(True) use_ub = not use_ub_for_nccl and ( torch_compile_enable_userbuffers and self._init_userbuffers(self.model.config.hidden_size)) self.backend_num_streams = Backend.Streams([ torch.cuda.Stream() for _ in range(torch_compile_max_num_streams - 1) ]) self._torch_compile_backend = Backend( torch_compile_inductor_enabled, enable_userbuffers=use_ub, enable_piecewise_cuda_graph=self. _torch_compile_piecewise_cuda_graph, capture_num_tokens=self._prefill_cuda_graph_num_tokens, max_num_streams=torch_compile_max_num_streams, mapping=self.mapping) apply_llm_torch_compile = getattr(self.model, "apply_llm_torch_compile", None) if isinstance(self.model, DecoderModelForCausalLM): eager_model = self.model.model compiled_model = torch.compile( eager_model, backend=self._torch_compile_backend, fullgraph=torch_compile_fullgraph) self._torch_compile_prefill_only = ( self._torch_compile_piecewise_cuda_graph and not self.model.use_fx_for_pcg_fallback) self.model.model = ( _PrefillCompiledModel(eager_model, compiled_model) if self._torch_compile_prefill_only else compiled_model) elif callable(apply_llm_torch_compile): # TODO: Move this contract to MultimodalModelMixin once # multimodal models consistently expose their LLM compile # scope through the mixin. apply_llm_torch_compile(backend=self._torch_compile_backend, fullgraph=torch_compile_fullgraph) else: self.model = torch.compile( self.model, backend=self._torch_compile_backend, fullgraph=torch_compile_fullgraph) torch._dynamo.config.cache_size_limit = 16 else: set_torch_compiling(False) except Exception as e: import traceback traceback.print_exception(Exception, e, e.__traceback__) raise e self.is_warmup = False self.previous_request_ids = [] self._encoder_decoder_host_buffer_pool: List[Dict[str, Any]] = [] self._encoder_decoder_input_fast_path_static_eligible: Optional[ bool] = None self._encoder_decoder_position_id_offset: Optional[int] = None sparse_params = (self.sparse_attention_config.to_sparse_params( pretrained_config=self.model.model_config.pretrained_config) if self.sparse_attention_config is not None else None) self.attn_backend = get_attention_backend(self.llm_args.attn_backend, sparse_params=sparse_params) self.get_runtime_tokens_per_gen_step = spec_config.get_runtime_tokens_per_gen_step if spec_config is not None else lambda runtime_draft_len: 1 self.spec_metadata = None if self.is_spec_decode: if not self.is_draft_model: update_spec_config_from_loaded_model(self.spec_config, self.model) max_num_draft_tokens = self.max_draft_loop_tokens * self.batch_size self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, device='cuda') self.gather_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') self.num_accepted_draft_tokens_cuda = torch.empty( (self.batch_size, ), dtype=torch.int, device='cuda') self.previous_pos_indices_cuda = torch.empty( (self.max_num_tokens, ), dtype=torch.int, device='cuda') self.previous_pos_id_offsets_cuda = torch.zeros( (self.max_num_tokens, ), dtype=torch.int, device='cuda') 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( ) 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 # dynamic draft length is enabled; otherwise stays fixed). Tree # modes verify all tree nodes per step, which can be wider than the # tree depth used by the drafter loop. self.runtime_draft_len = get_static_draft_len(self) else: self.without_logits = False self.max_draft_len = 0 self.runtime_draft_len = 0 self.max_total_draft_tokens = 0 self.guided_decoder: Optional[CapturableGuidedDecoder] = None # This field is initialized lazily on the first forward pass. # This is convenient because: # 1) The attention metadata depends on the KV cache manager. # 2) The KV cache manager depends on the model configuration. # 3) The model configuration is not loaded until the model engine # is initialized. # # NOTE: This can be simplified by decoupling the model config loading and # the model engine. self.attn_metadata = None self._eager_workspace_reclaimer: Optional[ EagerWorkspaceReclaimer] = None self.spec_metadata = None self.iter_states = {} # Log cached prefixes in model-input sequence order when enabled. self._log_cached_kv_tokens_per_req = os.getenv( 'TLLM_LOG_CACHED_KV_TOKENS_PER_REQ', '0') == '1' # Let the first CUDA graph capture create its private pool. Piecewise # CUDA graphs use a separate pool owned by their runners, so sharing a # pre-created pool handle with the outer graph runner is unnecessary. self._cuda_graph_mem_pool = None self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping( ) self.previous_batch_indices_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') self._encoder_decoder_staged_request_ids: Optional[List[int]] = None self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') self.position_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') # Steady-state generation-only prepare cache (non-speculative overlap # decode). Holds the per-request lists that are invariant while the # scheduled generation batch keeps the same composition, plus a pinned # cached-token counter advanced by one per step (host-side bookkeeping # only; the device position buffer is advanced in place and this # buffer is never the source of an async H2D). Invalidated (set to # None) by every full _prepare_tp_inputs pass. self._steady_gen_cache: Optional[Dict[str, Any]] = None self._steady_gen_positions_pinned = torch.empty( (self.max_num_tokens, ), dtype=torch.int, pin_memory=prefer_pinned()) if self.use_mrope: self.mrope_position_ids_cuda = torch.empty( (3, 1, self.max_num_tokens), dtype=torch.int, device='cuda') # Pre-allocated buffers for draft model to avoid implicit synchronization # These are used to build index tensors without creating tensors from Python lists max_first_draft_tokens = self.batch_size * ( self.original_max_total_draft_tokens + 1) if spec_config else self.batch_size tokens_per_draft = self.original_max_total_draft_tokens + 1 self.idx_accepted_tokens_cache = None self.draft_token_positions_cache = None if spec_config: # Cache for idx_accepted_tokens (pattern: 0,0,0...1,1,1...2,2,2...) self.idx_accepted_tokens_cache = torch.arange( max_first_draft_tokens, dtype=torch.long, device='cuda') // tokens_per_draft if self.is_draft_model: self.draft_ctx_token_indices_cuda = torch.empty((self.batch_size, ), dtype=torch.long, device='cuda') self.draft_ctx_seq_slots_cuda = torch.empty((self.batch_size, ), dtype=torch.long, device='cuda') # Buffers for first_draft requests (max_draft_len+1 tokens per request) self.draft_first_draft_indices_cuda = torch.empty( (max_first_draft_tokens, ), dtype=torch.long, device='cuda') self.draft_first_draft_seq_slots_cuda = torch.empty( (max_first_draft_tokens, ), dtype=torch.long, device='cuda') # Buffers for seq_slots and request indices self.draft_seq_slots_buffer_cuda = torch.empty((self.batch_size, ), dtype=torch.int, device='cuda') self.draft_request_indices_buffer_cuda = torch.empty( (self.batch_size, ), dtype=torch.int, device='cuda') # Pre-computed constant tensors for incremental update optimization # Cache for token_positions (pattern: 0,1,2...N repeated) self.draft_token_positions_cache = torch.arange(tokens_per_draft, dtype=torch.long, device='cuda') # We look up this key in resource_manager during forward to find the # kv cache manager. Can be changed to support multiple model engines # with different KV cache managers. self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False self._warmup_timer = _WarmupTimer(self.mapping.rank) self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None self._force_lora_graph_for_capture: Optional[bool] = None self._lora = LoraParamBuilder(spec_config=self.spec_config, attn_backend=self.attn_backend) # Sampling tier pinned during an in-graph sampling capture pass; None outside # capture, where the tier comes from the batch instead. self._capture_sample_type: Optional[SampleType] = None # Setup the local cache indirection buffer only once and reuse it. # This way it can also be used for CUDA graphs. if self.use_beam_search: self.cache_indirection_attention = torch.zeros( (self.batch_size, self.max_beam_width, self.max_seq_len), device="cuda", dtype=torch.int32) else: self.cache_indirection_attention = None runner_cls = resolve_runner_type(self.model, self.llm_args) self._runner = self._initialize_runner(runner_cls) self.cuda_graph_runner = self._initialize_cuda_graph_runner() self.breakable_cuda_graph_runner = \ self._initialize_breakable_cuda_graph_runner() self.kv_cache_dtype_byte_size = self.get_kv_cache_dtype_byte_size() self._prepare_inputs_event: Optional[torch.cuda.Event] = None # Cache for enc-dec cross-attention stable generation steps. # Populated on the first CUDA-graph generation step; cleared whenever # the batch composition changes (new encoder request arrives). self._cross_attn_stable_cached_tokens: Optional[List[int]] = None self._cross_attn_stable_request_ids: Optional[List[int]] = None def _initialize_cuda_graph_runner(self) -> Optional[CUDAGraphRunner]: is_encoder_decoder = self._is_encoder_decoder_model() if self._runner is not None and not is_encoder_decoder: return None enable_encoder_decoder_mixed_cuda_graph = ( is_encoder_decoder and bool(self._encoder_graph_shapes) and self.cuda_graph_config is not None and self.llm_args.enable_encoder_decoder_mixed_cuda_graph) config = CUDAGraphRunnerConfig( use_cuda_graph=(not self._is_encode_only and self.cuda_graph_config is not None), cuda_graph_padding_enabled=self._cuda_graph_padding_enabled, cuda_graph_batch_sizes=self._cuda_graph_batch_sizes, max_cuda_graph_batch_size=self._max_cuda_graph_batch_size, max_beam_width=self.max_beam_width, spec_config=self.spec_config, cuda_graph_mem_pool=self._cuda_graph_mem_pool, dynamic_draft_len_mapping=self._dynamic_draft_len_mapping, max_num_tokens=self.max_num_tokens, use_mrope=self.use_mrope, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=self. original_max_total_draft_tokens, is_draft_model=self.is_draft_model, enable_attention_dp=self.enable_attention_dp, is_encoder_decoder=is_encoder_decoder, batch_size=self.batch_size, mapping=self.mapping, dist=self.dist, kv_cache_manager_key=self.kv_cache_manager_key, sparse_attention_config=self.sparse_attention_config, enable_encoder_decoder_mixed_cuda_graph=( enable_encoder_decoder_mixed_cuda_graph), enable_in_graph_sampling=self.enable_in_graph_sampling, ) return CUDAGraphRunner(config) def _initialize_breakable_cuda_graph_runner( self) -> Optional[BreakableCUDAGraphRunner]: if (self.cuda_graph_runner is None or self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.BREAKABLE): return None decoder_model = (self.model if isinstance( self.model, DecoderModelForCausalLM) else getattr( self.model, "llm", None)) if not isinstance(decoder_model, DecoderModelForCausalLM): raise ValueError( "breakable prefill CUDA graph requires a decoder model body") return BreakableCUDAGraphRunner(decoder_model.model) def _initialize_runner( self, runner_cls: Optional[Type[Union[ModelRunner, PackedModelRunner]]] ) -> Optional[Union[ModelRunner, PackedModelRunner]]: if runner_cls is None: return None if issubclass(runner_cls, EncoderRunner): return self._initialize_encoder_runner(runner_cls) if issubclass(runner_cls, EncoderDecoderRunner): return self._initialize_encoder_decoder_runner(runner_cls) if issubclass(runner_cls, NoKVCacheRunner): return self._initialize_no_kv_cache_runner(runner_cls) raise TypeError(f"No runner initializer registered for " f"{runner_cls.__module__}.{runner_cls.__qualname__}") def _initialize_encoder_runner( self, runner_cls: Type[EncoderRunner]) -> EncoderRunner: runner_config = EncoderRunnerConfig.create( model=self.model, mapping=self.mapping, graph_config=(self.cuda_graph_config if isinstance( self.cuda_graph_config, EncodeCudaGraphConfig) else None), max_batch_size=self.batch_size, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, max_beam_width=self.max_beam_width, without_logits=self.without_logits, attention_backend=self.attn_backend, attention_runtime_features=self.attn_runtime_features, enable_autotuner=self.llm_args.enable_autotuner, draft_model=self.is_draft_model, ) return runner_cls( self.model, self._create_runner_deps(), runner_config, ) def _initialize_encoder_decoder_runner( self, runner_cls: Type[EncoderDecoderRunner]) -> EncoderDecoderRunner: runner_config = EncoderDecoderRunnerConfig.create( model=self.model, mapping=self.mapping, graph_config=self.llm_args.encoder_cuda_graph_config, max_batch_size=self.encoder_batch_size, max_num_tokens=self.encoder_max_num_tokens, max_seq_len=self.max_seq_len, max_beam_width=self.max_beam_width, without_logits=self.without_logits, attention_backend=self.attn_backend, attention_runtime_features=self.attn_runtime_features, enable_autotuner=self.llm_args.enable_autotuner, is_encoder_decoder=True, draft_model=self.is_draft_model, ) runner = runner_cls( self.model, self._create_runner_deps(), runner_config, ) # Remove this bridge when decoder graph ownership moves into the runner. # Only immutable planning data is shared; graph resources stay private. self._encoder_graph_shapes = runner._encoder_graph_shapes self._encoder_graph_batch_sizes = runner._encoder_graph_batch_sizes self._encoder_graph_pad_to_limit = runner._encoder_graph_pad_to_limit return runner def _initialize_no_kv_cache_runner( self, runner_cls: Type[NoKVCacheRunner]) -> NoKVCacheRunner: runner_config = NoKVCacheRunnerConfig( max_batch_size=self.batch_size, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, max_beam_width=self.max_beam_width, without_logits=self.without_logits, enable_attention_dp=self.enable_attention_dp, prefill_cuda_graph_backend=self.prefill_cuda_graph_backend, prefill_cuda_graph_num_tokens=self._prefill_cuda_graph_num_tokens, attention_backend=self.attn_backend, attention_runtime_features=self.attn_runtime_features, mm_encoder_cache_enabled=self._mm_encoder_cache_enabled, spec_config=self.spec_config, is_draft_model=self.is_draft_model, num_seq_slots=self.max_num_seq_slots, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=( self.original_max_total_draft_tokens), spec_dec_max_total_draft_tokens=( self._spec_dec_max_total_draft_tokens), ) return runner_cls(self.model, self._create_runner_deps(), runner_config) def _create_runner_deps(self) -> RunnerDeps: return RunnerDeps( dist=self.dist, mapping=self.mapping, input_ids_cuda=self.input_ids_cuda, position_ids_cuda=self.position_ids_cuda, gather_ids_cuda=getattr(self, "gather_ids_cuda", None), draft_tokens_cuda=getattr(self, "draft_tokens_cuda", None), cache_indirection=(self.cache_indirection_attention if self.attn_backend.Metadata is TrtllmAttentionMetadata else None), lora=self._lora, moe_load_balancer=self.moe_load_balancer, # Do not retain the engine through a bound method. model_forward=functools.partial( type(self).model_forward, weakref.proxy(self)), ) def register_forward_pass_callable(self, callable: Callable): self.forward_pass_callable = callable def register_sample_in_graph_callable(self, callable: Optional[Callable]): """Register the hook that samples at the tail of the forward graph.""" self.sample_in_graph_callable = callable def register_sample_type_resolver(self, resolver: Optional[Callable], stage: Optional[Callable] = None): """Register how a batch maps to its sampling tier, for the graph key.""" self.cuda_graph_runner.register_sample_type_resolver(resolver) self._stage_in_graph_sampling = stage def get_kv_cache_dtype_byte_size(self) -> float: """ Returns the size (in bytes) occupied by kv cache type. """ layer_quant_mode = self.model.model_config.quant_config.layer_quant_mode if layer_quant_mode.has_fp4_kv_cache(): return 1 / 2 elif layer_quant_mode.has_fp8_kv_cache( ) or layer_quant_mode.has_int8_kv_cache(): return 1 else: return 2 def set_lora_model_config(self, lora_target_modules: list[str], trtllm_modules_to_hf_modules: dict[str, str], swap_gate_up_proj_lora_b_weight: bool = True): # Called by `_util.py` after the engine exists. Both LoRA handles stay # engine state: warmup, capture and the enc-dec fast path read them. self.lora_model_config = make_lora_model_config( self.model, lora_target_modules, trtllm_modules_to_hf_modules, swap_gate_up_proj_lora_b_weight) def _init_cuda_graph_lora_manager(self, lora_config: LoraConfig): """Initialize CUDA Graph LoRA manager with model configuration.""" if (self.cuda_graph_runner is not None and self.cuda_graph_runner.enabled): # For spec decode, each generation request contributes # max_draft_len + 1 tokens per forward pass. max_tokens_per_seq = (self.original_max_draft_len + 1) if self.is_spec_decode else 1 self.cuda_graph_lora_manager = make_cuda_graph_lora_manager( self.model, lora_config, self.lora_model_config, self.batch_size, # Use engine's max batch size max_tokens_per_seq, self.max_num_tokens) def _use_lora_cuda_graph(self, scheduled_requests: ScheduledRequests) -> bool: """ Determines whether a non-LoRA or LoRA CUDA graph should be used, if both are available (cuda_graph_specialize_lora==True). """ if self.cuda_graph_lora_manager is None: return False # Needed during graph capture to enforce a given mode if self._force_lora_graph_for_capture is not None: return self._force_lora_graph_for_capture if not self.llm_args.lora_config.cuda_graph_specialize_lora: return True return any(request.lora_task_id is not None for request in scheduled_requests.generation_requests) def set_guided_decoder(self, guided_decoder: CapturableGuidedDecoder) -> bool: if hasattr(self.model, "set_guided_decoder"): success = self.model.set_guided_decoder(guided_decoder) if success: self.guided_decoder = guided_decoder return success return False @property def use_mrope(self): use_mrope = False try: use_mrope = self.model.model_config.pretrained_config.rope_scaling[ 'type'] == 'mrope' except Exception: pass logger.debug(f"Detected use_mrope: {use_mrope}") return use_mrope @functools.cached_property def _mm_encoder_cache_enabled(self) -> bool: """Whether the multimodal encoder cache is active for this model.""" return mm_encoder_cache_enabled(self.model) @property def is_warmup(self): return getattr(self, "_is_warmup", False) @is_warmup.setter def is_warmup(self, value: bool): self._is_warmup = value # This setter is the one choke point every warmup transition passes # through, including PyExecutor's, so select the MoE all-to-all budget # here rather than in set_warmup_flag(). _set_moe_a2a_warmup(value) self.moe_load_balancer_iter_info = (not value, not value) @property def moe_load_balancer_iter_info(self): moe_load_balancer = self.moe_load_balancer if moe_load_balancer is not None: return moe_load_balancer.enable_statistic, moe_load_balancer.enable_update_weights return False, False @moe_load_balancer_iter_info.setter def moe_load_balancer_iter_info(self, value: Tuple[bool, bool]): moe_load_balancer = self.moe_load_balancer if moe_load_balancer is not None: moe_load_balancer.set_iter_info(enable_statistic=value[0], enable_update_weights=value[1]) @property def use_beam_search(self): return self.max_beam_width > 1 @property def _is_packed_runner(self) -> bool: """Which of the two runner contracts `_runner` implements.""" return isinstance(self._runner, EncoderRunner) def _get_draft_kv_cache_manager( self, resource_manager: ResourceManager ) -> Optional[Union[KVCacheManager, KVCacheManagerV2]]: """ Returns the draft KV cache manager only in one-model speculative decoding mode where the target model manages a separate draft KV cache. """ return get_draft_kv_cache_manager(self.spec_config, resource_manager) @contextmanager def set_warmup_flag(self): prev_is_warmup = self.is_warmup self.is_warmup = True try: yield finally: self.is_warmup = prev_is_warmup @staticmethod def with_warmup_flag(method): @functools.wraps(method) def wrapper(self, *args, **kwargs): with self.set_warmup_flag(): return method(self, *args, **kwargs) return wrapper @contextlib.contextmanager def no_cuda_graph(self): if self.cuda_graph_runner is None: yield return _run_cuda_graphs = self.cuda_graph_runner.enabled self.cuda_graph_runner.enabled = False try: yield finally: self.cuda_graph_runner.enabled = _run_cuda_graphs def _pad_batch_seed_mrope_delta_cache( self, padded_requests: ScheduledRequests) -> None: if not self.use_mrope or padded_requests.num_generation_requests == 0: return mrope_position_deltas_cache = resolve_mrope_position_deltas_cache( self.model) if mrope_position_deltas_cache is None: return mrope_seed_seq_slots = [] mrope_seed_deltas = [] mrope_seed_requests = [] for request in padded_requests.generation_requests: if (request.py_seq_slot is None or request.is_dummy or getattr(request, "py_mrope_delta_cache_slot", None) == request.py_seq_slot): continue mrope_position_delta = getattr(request, "py_mrope_position_delta", None) if mrope_position_delta is None and request.py_multimodal_data: mrope_config = request.py_multimodal_data.get('mrope_config') if mrope_config is not None: mrope_position_delta = mrope_config.get( 'mrope_position_deltas') if mrope_position_delta is None: continue if mrope_position_delta.device.type == "cpu": mrope_position_delta = maybe_pin_memory( mrope_position_delta).to(device='cuda', dtype=torch.int32, non_blocking=True) elif mrope_position_delta.dtype != torch.int32: mrope_position_delta = mrope_position_delta.to( dtype=torch.int32) request.py_mrope_position_delta = mrope_position_delta mrope_seed_seq_slots.append(request.py_seq_slot) mrope_seed_deltas.append(mrope_position_delta.reshape(1)) mrope_seed_requests.append(request) if not mrope_seed_seq_slots: return mrope_seed_seq_slots_tensor = torch.tensor( mrope_seed_seq_slots, dtype=torch.long, pin_memory=prefer_pinned()).to(device='cuda', non_blocking=True) mrope_seed_deltas_tensor = torch.cat(mrope_seed_deltas, dim=0) mrope_position_deltas_cache.index_copy_( 0, mrope_seed_seq_slots_tensor, mrope_seed_deltas_tensor.to( dtype=mrope_position_deltas_cache.dtype)) for request in mrope_seed_requests: request.py_mrope_delta_cache_slot = request.py_seq_slot @staticmethod def warmup_with_kv_cache_cleanup(method): """ Decorator for warmup methods that cleans up NaNs/Infs in KV Cache after warmup execution. Why this is needed: - Our attention kernel uses multiplication by zero to mask out invalid tokens within the same page. Since NaN/Inf * 0 = NaN, any NaNs/Infs in these invalid KV areas will persist after masking. - These NaNs/Infs propagate to outputs and subsequent KV Cache entries, corrupting future computations with higher probability. - During warmup, we execute with placeholder data rather than actual valid inputs, which can introduce NaNs/Infs into KV Cache pages and cause random, hard-to-debug accuracy issues. """ @functools.wraps(method) def wrapper(self, resource_manager: Optional[ResourceManager] = None, *args, **kwargs): result = method(self, resource_manager, *args, **kwargs) kv_cache_manager = (resource_manager.get_resource_manager( self.kv_cache_manager_key) if resource_manager is not None else None) if kv_cache_manager is not None: has_invalid_values = kv_cache_manager.check_invalid_values_in_kv_cache( fill_with_zero=True) if has_invalid_values: logger.warning( "NaNs/Infs have been introduced to KVCache during warmup, KVCache was filled with zeros to avoid potential issues" ) return result return wrapper def _get_max_shape_warmup_requests( self, resource_manager: ResourceManager) -> List[Tuple[int, int]]: """ Returns warmup configs covering the maximum context and generation shapes. """ kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) token_num_upper_bound = min(self.max_num_tokens, self.batch_size * (self.max_seq_len - 1)) curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( token_num_upper_bound=token_num_upper_bound, max_num_draft_tokens=self.original_max_draft_len) max_batch_size = min( self.batch_size, curr_max_num_tokens // (1 + self.max_draft_loop_tokens) // self.max_beam_width) warmup_requests_configs = [ (curr_max_num_tokens, 0), # max_num_tokens, pure context (max_batch_size, max_batch_size), # max_batch_size, pure generation ] return warmup_requests_configs def _get_full_general_warmup_requests( self, resource_manager: ResourceManager) -> List[Tuple[int, int]]: """ Returns the ordered warmup configs for torch.compile specialization. Covers 1-token (0-1 graph specialization), max-shape (best triton autotuning), and small-context (2-token path) cases. """ max_configs = self._get_max_shape_warmup_requests(resource_manager) # Specialize for 1 token pure ctx and pure gen one_token_configs = [(1, 0), (1, 1)] # Small ctx specialization small_ctx_configs = [(2, 0)] # Ordering matters for torch.compile graph specialization: # 1-token first to capture the 0→1 transition graph; max-shape next to seed # triton autotuning with the largest inputs; 2-token last for the small-ctx path. warmup_configs = one_token_configs + max_configs + small_ctx_configs # Deduplicate the warmup_configs while keeping the order. return list(dict.fromkeys(warmup_configs)) @contextmanager def maybe_autotune_lora(self): """Enable autotuning while warming up CUDA-graph LoRA kernels.""" if not (self.llm_args.enable_autotuner and self.cuda_graph_lora_manager is not None): yield return cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) with autotune(cache_path=cache_path): try: yield finally: # Complete the PP cache hand-off even on ranks without a # CUDA-graph-only tunable op. autotuner = AutoTuner.get() autotuner.cache_pp_recv() autotuner.cache_pp_send() autotuner.clean_pp_flag() @with_warmup_flag def _warmup_encoder_cuda_graphs_enc_dec( self, resource_manager: ResourceManager, ) -> None: if not self._is_encoder_decoder_model(): return if self._runner is None or self._is_packed_runner: raise RuntimeError( "Encoder-decoder model did not initialize a model runner.") runner = cast(ModelRunner, self._runner) runner.warmup(resource_manager) runner.capture_graphs(resource_manager) def _get_encoder_cuda_graph_batch_sizes( self, max_batch_size: int) -> tuple[int, ...]: """Use startup settings while encoder scheduling remains in PyExecutor.""" return resolve_cuda_graph_batch_sizes( self._encoder_graph_batch_sizes, max_batch_size, pad_to_limit=self._encoder_graph_pad_to_limit) def forward_encoder( self, encoder_requests: List[LlmRequest], resource_manager: Optional[ResourceManager] = None, ) -> Tuple[torch.Tensor, List[int]]: if (not self._is_encoder_decoder_model() or self._runner is None or self._is_packed_runner): raise RuntimeError( "Encoder phase requires an initialized encoder-decoder model runner." ) assert resource_manager is not None, ( "the encoder phase requires a resource manager") scheduled_requests = ScheduledRequests() scheduled_requests.encoder_requests = list(encoder_requests) outputs = cast(ModelRunner, self._runner).forward( scheduled_requests, resource_manager=resource_manager, cuda_graph_lora_manager=None, runtime_draft_len=0, gather_context_logits=False, ) return ( outputs["encoder_hidden_states"], outputs["encoder_seq_lens"], ) @with_warmup_flag @warmup_with_kv_cache_cleanup def warmup(self, resource_manager: Optional[ResourceManager] = None) -> None: """ Orchestrates the warmup process by calling specialized warmup methods for torch.compile, the autotuner, and CUDA graphs. """ if self._is_packed_runner: packed_runner = cast(PackedModelRunner, self._runner) packed_runner.warmup() packed_runner.capture_graphs() return assert resource_manager is not None, ( "scheduled warmup requires a resource manager") kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) if self._runner is not None and not self._is_encoder_decoder_model(): assert kv_cache_manager is None, ( "a no-KV-cache runner was initialized, but a KV cache manager was allocated" ) runner = cast(ModelRunner, self._runner) runner.warmup(resource_manager) runner.capture_graphs(resource_manager) return # Retain a partial summary when a phase raises. with self._warmup_timer: self._warmup_scheduled(resource_manager, kv_cache_manager) def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) -> None: """Body of ``warmup`` for the scheduled (KV-cache-backed) path.""" # Ahead of the legacy early returns below: only the advanced-sampling # CUDA graph capture pass exercises the non-greedy sampler, so with # cuda_graph_config=None flashinfer's sampling kernels would be # JIT-built mid-serving. self._eager_workspace_reclaimer = None with self._warmup_timer.phase("sampling_module_prewarm"): warmup_sampling_module() if self.enable_in_graph_sampling: # The fast tier samples inside the captured graph via a # torch.compile'd op; compile it now so capture does not. warmup_sample_from_logits_op(self.model.config.vocab_size, torch.device('cuda'), self.dtype, self._cuda_graph_batch_sizes or []) if kv_cache_manager is None: logger.info("Skipping warm up as no KV Cache manager allocated.") return # The lifetime of model engine and kv cache manager can be different. # Reset the global cuda graph dummy requests in warmup. self.cuda_graph_runner.padding_dummy_requests = {} is_enc_dec = self._is_encoder_decoder_model() if self.mapping.cp_size > 1: cp_type = self.mapping.cp_config.get("cp_type", None) if cp_type != CpType.HELIX: logger.info( f"[ModelEngine::warmup] Skipping warmup for cp_type: {None if cp_type is None else cp_type.name}." ) return # Create AutoTuner singleton in eager context before any compiled forward. # Otherwise the first get() can happen inside torch.compile tracing and # trigger non-traceable code (time.time(), torch.cuda.*) in the cache. AutoTuner.get() # ``guided_decoder`` is installed only on the last pipeline rank, so # this predicate is not rank-uniform on its own. Agree it before it # gates either the attention or the general phase. can_run_general_warmup = self._agree_warmup_flag( not is_enc_dec and not self.is_draft_model and not self.mapping.has_cp_helix() and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager)) log_mem_snapshot("warmup/before_warmup") # Compile the DSv4 indexer-Q CuTe DSL kernels before the first # collective-bearing forward, so their JIT cost is not charged against the # MoE all-to-all completion-flag deadline. with self._warmup_timer.phase("cute_dsl_indexer_q"): self._prewarm_cute_dsl_indexer_q() log_mem_snapshot("warmup/after_cute_dsl_indexer_q") if not is_enc_dec: with self._warmup_timer.phase("attention_jit"): self._run_attention_warmup(resource_manager, can_run_general_warmup) if can_run_general_warmup: # Specialize torch.compile graphs across the key input shapes before CUDA graph capture. with self._warmup_timer.phase("general"): warmup_requests_configs = self._agree_warmup_shapes( self._get_full_general_warmup_requests(resource_manager)) # Currently graph has not been captured, disable cuda graph for this warmup. with self.no_cuda_graph(): self._general_warmup(resource_manager, warmup_requests_configs) # Release C++ MoE workspace buffers so the autotuner can # reclaim the memory. They will be re-allocated on next use. from ..custom_ops.torch_custom_ops import MoERunner MoERunner.clear_all_workspaces() # Clear Cache now as autotuner may use additional memory. # Memory pool will be warmed up later. gc.collect() torch.cuda.empty_cache() # Helix CP is decode-only and runs into issues with the # autotuner warmup's context requests. if not is_enc_dec and not self.mapping.has_cp_helix(): with self._warmup_timer.phase("autotuner"): self._run_autotuner_warmup(resource_manager) log_mem_snapshot("warmup/after_autotuner") # Pre-JIT Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels # for Mamba hybrid models. Runs regardless of enable_autotuner, # since MambaHybridCacheManager skips _general_warmup and the # default autotuner shape is single-seq / no-initstates. Safe # no-op for non-Mamba models. with self._warmup_timer.phase("mamba_hybrid"): self._run_mamba_hybrid_warmup(resource_manager) log_mem_snapshot("warmup/after_mamba_hybrid") # Release the autotuner's exploration-mode intermediates. The # exploration leftovers are pure waste that hide tens of GiB from # non-torch allocators (cuBLAS handle workspace, UCX/NIXL, # NVSHMEM). gc.collect() torch.cuda.empty_cache() # Warm up every graph shape before capturing any graph. Attention # kernels can switch implementations at smaller batch sizes and require # a larger workspace, so the first pass grows the workspace to its # maximum size. The second pass runs the final per-shape warmup and # captures without resizing the workspace. # Capture with the steady-state MoE all-to-all budget: the timeout is a # launch argument and is baked into every later replay. with self._warmup_timer.phase("cuda_graph_capture"), \ _moe_a2a_steady_state_budget_for_capture(): with self.cuda_graph_runner.allow_capture(): self.cuda_graph_runner.is_warmup_only = True try: with self.maybe_autotune_lora(): self._run_cuda_graph_warmup(resource_manager) finally: self.cuda_graph_runner.is_warmup_only = False self.cuda_graph_runner.padding_dummy_requests = {} self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned # batch bucket the runtime can produce (max_batch_size scaled by the # MTP / DSL expansion factor when applicable). CUDA-graph warmup only # exercises the batch sizes in cuda_graph_batch_sizes, which round # up to a subset of buckets; any inference iter whose # context_lens.size(0) lands on an uncovered bucket triggers an # nvcc-driven JIT compile (~3s stall inside _prepare_inputs) on # first touch. Pre-touching every bucket funnels that cost into # warmup. No-op on non-DSA models. # Both DSA hooks read attn_metadata, which only a warmup forward # creates; build it when every forward above was skipped. with self._warmup_timer.phase("dsa_prewarm"): self._ensure_dsa_attn_metadata_for_warmup(resource_manager) self._warmup_dg_paged_mqa_logits_metadata() log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") self._warmup_cute_dsl_radix_topk() log_mem_snapshot("warmup/after_cute_dsl_radix_topk") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. with self._warmup_timer.phase("memory_pool_prepop"): warmup_requests_configs = self._get_max_shape_warmup_requests( resource_manager) self._general_warmup(resource_manager, warmup_requests_configs) log_mem_snapshot("warmup/after_memory_pool_prepop") # Allocate the CUDA graph padding dummies now, while the KV cache is # empty. Waiting for the first padded step can race KV saturation: # once the cache is full, the lazy allocation in _get_padded_batch # fails every step and padded batches silently run eager. with self._warmup_timer.phase("preallocate_padding_dummies"): self.cuda_graph_runner.preallocate_padding_dummies(resource_manager) log_mem_snapshot("warmup/after_preallocate_padding_dummies") # If this is a BOLT-instrumented build (the profile-gen job sets # TLLM_BOLT_CLEAR_COUNTERS=1), reset the instrumentation counters now # that all startup JIT/autotune/graph-capture is done, so the emitted # .fdata reflects steady-state serving only. No-op on normal builds. from ..bolt_profiling import maybe_bolt_clear_counters maybe_bolt_clear_counters() self._freeze_eager_workspace_floor() def _freeze_eager_workspace_floor(self) -> None: if os.environ.get("TRTLLM_RECLAIM_WORKSPACE", "1") == "0": return metadata = self.attn_metadata if (self.is_spec_decode or self.mapping.cp_size != 1 or self._is_encoder_decoder_model() or self.sparse_attention_config is not None or self._torch_compile_backend is not None or self.breakable_cuda_graph_runner is not None): logger.info( "Eager workspace reclamation is disabled for speculative, " "CP, encoder-decoder, sparse, compiled, or breakable-graph models" ) return if (type(metadata) is not TrtllmAttentionMetadata or not metadata.workspace_reclaimable or metadata.workspace is None or metadata.workspace.untyped_storage().nbytes() == 0): logger.info( "Eager workspace reclamation requires warmup of pure fallback scratch" ) return self._eager_workspace_reclaimer = EagerWorkspaceReclaimer(metadata) def _warmup_dg_paged_mqa_logits_metadata(self) -> None: """Pre-compile DeepGEMM's `get_paged_mqa_logits_metadata` helper for every 32-aligned batch bucket the runtime can produce. DSA's `Indexer.prepare_scheduler_metadata` calls `deep_gemm.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms)` inside `_prepare_inputs` every iteration. The underlying kernel is templated on `` where `kAlignedBatchSize = align(context_lens.size(0), 32)` and `split_kv` / `num_sms` are fixed for a given (block_kv, device). deep_gemm's Python-side JIT compiles a fresh cubin (spawning nvcc/cicc/ptxas, ~3s on GB300) the first time each `aligned_bs` is requested. CUDA-graph warmup exercises only the batch sizes in `cuda_graph_batch_sizes`, which round up to a subset of the 32- aligned buckets; every uncovered bucket that the inference workload later touches produces a 3s stall on that iteration. Pre-touching every bucket here funnels those compiles into the deterministic warmup phase. `context_lens.size(0)` is not always `num_generations`. For MTP with `use_expanded_buffers_for_mtp=True` the expanded call passes `num_generations * (1 + max_draft_tokens)`. For DSL expansion the call passes `num_generations * dsl_expand_factor`, where `dsl_expand_factor = next_n // eff` (`eff in kernel_atoms`, see `_pick_dsl_expand` in `dsa.py`); its worst case is `next_n = 1 + max_draft_tokens` when `eff == 1`. Reading the current `dsl_expand_factor` off the metadata would under-estimate the eventual max (it defaults to 1 before any prepare() has run, and per-iter picks can differ across iters when CUDA graph is off), so we use the static upper bound `1 + max_draft_tokens` for both expansion paths. Bucket range is also scaled by `max_beam_width` as a defense-in-depth ceiling for future beam support (no-op today — DSA does not use beam). No-op on non-DSA models. Best-effort: per-bucket JIT failures are logged and skipped so a single broken bucket does not abort PyExecutor startup. """ attn_meta = getattr(self, "attn_metadata", None) if attn_meta is None: return try: from tensorrt_llm._torch.attention.backends.sparse.dsa import ( _DG_SCHEDULE_BLOCK_KV, DSAtrtllmAttentionMetadata) except ImportError: return if not isinstance(attn_meta, DSAtrtllmAttentionMetadata): return try: from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata except ImportError: logger.info( "[DG warmup] deep_gemm.get_paged_mqa_logits_metadata not " "available; skipping paged_mqa_logits_metadata prewarm.") return num_sms = attn_meta.num_sms max_bs = max(1, int(self.batch_size)) beam_width = max(1, int(getattr(self, "max_beam_width", 1) or 1)) # Static upper bound on the row-count multiplier applied to # `context_lens`. Both MTP-expanded and DSL-expanded call sites # are bounded above by `(1 + max_draft_tokens)`; see the # docstring for why we don't read the runtime `dsl_expand_factor` # here. max_draft_tokens = int(getattr(attn_meta, "max_draft_tokens", 0) or 0) expands_batch = (getattr(attn_meta, "use_expanded_buffers_for_mtp", False) or getattr(attn_meta, "expand_for_dsl", False)) expand_factor = 1 + max_draft_tokens if expands_batch else 1 max_aligned = ((max_bs * beam_width * expand_factor + 31) // 32) * 32 buckets = list(range(32, max_aligned + 32, 32)) logger.info(f"[DG warmup] Pre-compiling paged_mqa_logits_metadata for " f"{len(buckets)} aligned batch buckets up to {max_aligned} " f"(block_kv={_DG_SCHEDULE_BLOCK_KV}, num_sms={num_sms}, " f"max_bs={max_bs}, beam_width={beam_width}, " f"expand_factor={expand_factor})") for aligned_bs in buckets: # Kernel scans `context_lens` and prefix-sums schedules; a # zero-filled 2D tensor of shape (aligned_bs, 1) is enough to # trigger dispatch and compile — the metadata output is # discarded. dummy = torch.zeros(aligned_bs, 1, dtype=torch.int32, device="cuda") try: _ = get_paged_mqa_logits_metadata(dummy, _DG_SCHEDULE_BLOCK_KV, num_sms) except RuntimeError as e: # Narrow to RuntimeError so signature drifts in # get_paged_mqa_logits_metadata (TypeError / ValueError) # surface loudly instead of silently degrading perf. logger.warning( f"[DG warmup] paged_mqa_logits_metadata prewarm failed " f"for aligned_bs={aligned_bs} " f"(block_kv={_DG_SCHEDULE_BLOCK_KV}, num_sms={num_sms}); " f"skipping bucket. {type(e).__name__}: {e}") torch.cuda.synchronize() def _prewarm_cute_dsl_indexer_q(self) -> None: """Pre-compile the DSv4 indexer-Q CuTe DSL kernels, then barrier. Runs before any collective-bearing forward so this op's first-touch ``cute.compile`` is not charged against the MoE all-to-all completion-flag deadline. It is a partial mitigation only: other first-touch compiles remain inside collective-bearing forwards, and some sit on the all-to-all path itself and cannot be pre-compiled this way. The runtime budget (``moeA2AGetTimeoutCycles``) covers the general case. Only the fallback tactics are compiled -- what an eager, cache-miss forward selects. The runner's kernel cache key excludes m/n/k, so one compile per tactic covers every shape. Uses the real module and weights, so it cannot drift from what the model runs. No-op on non-DSA models. See nvbugs/6482566. """ try: from ..attention.backends.sparse.deepseek_v4.deepseek_v4 import \ DeepseekV4Indexer except ImportError: return indexer = next( (m for m in self.model.modules() if isinstance(m, DeepseekV4Indexer) and getattr(m, "wq_b", None) is not None), None) if indexer is None: return weight = indexer.wq_b.weight # _fallback_tactic() branches on m at 4 and 8, so these three token # counts cover every fallback tactic it can return. with torch.inference_mode(): for num_tokens in (4, 8, 16): try: qr = torch.zeros((num_tokens, weight.shape[1]), dtype=torch.bfloat16, device=weight.device) position_ids = torch.zeros((num_tokens, ), dtype=torch.int32, device=weight.device) indexer._project_and_quantize_q(qr, position_ids) except Exception as e: # Never fail startup for a prewarm miss; the kernel would # simply be compiled later, as it is today. logger.warning( f"indexer-Q CuTe DSL prewarm skipped for {num_tokens} " f"tokens. {type(e).__name__}: {e}") torch.cuda.synchronize() # Hold every rank here until the slowest has finished compiling, so the # first MoE all-to-all dispatch is entered without JIT skew. if self.mapping.tp_size > 1 and self.dist is not None: self.dist.tp_allgather(1) logger.info("indexer-Q CuTe DSL prewarm complete") def _ensure_dsa_attn_metadata_for_warmup( self, resource_manager: ResourceManager) -> None: """Build the DSA attention metadata if no warmup forward created it, so the top-K pre-compile hooks still run (draft engine, guided decoder or context-only server without general warmup). No-op unless DSA.""" if getattr(self, "attn_metadata", None) is not None: return try: from ..attention.backends.sparse.dsa import \ DSAtrtllmAttentionMetadata except ImportError: return metadata_cls = getattr(self.attn_backend, "Metadata", None) if metadata_cls is None or not issubclass(metadata_cls, DSAtrtllmAttentionMetadata): return kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) if kv_cache_manager is None: return self._set_up_attn_metadata( kv_cache_manager, self._get_draft_kv_cache_manager(resource_manager)) def _warmup_cute_dsl_radix_topk(self) -> None: """Pre-compile the DSA radix-filter CuTe DSL decode top-k for every cluster_size band during warmup, before serving. Captured geometries are already compiled by the warmup-step forwards; this fills in the bands the eager (non-captured) decode path can still hit (mixed prefill+decode batch, or cuda_graph disabled) so they do not pay a first-touch JIT stall on a live request. DSA-specific params live on the metadata, so delegate to it. No-op on non-DSA models. """ attn_meta = getattr(self, "attn_metadata", None) if attn_meta is None: return try: from ..attention.backends.sparse.dsa import \ DSAtrtllmAttentionMetadata except ImportError: return if isinstance(attn_meta, DSAtrtllmAttentionMetadata): next_n = 1 + self.original_max_draft_len attn_meta.warmup_cute_dsl_radix_topk(next_n) if hasattr(attn_meta, "warmup_selfsampling_topk"): attn_meta.warmup_selfsampling_topk( next_n, batch_sizes=self._cuda_graph_batch_sizes) def _general_warmup(self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]): """ Runs forward passes for each config in warmup_requests_configs. Serves both torch.compile graph specialization and memory pool pre-population. """ # Disable CUDA graph replay during general warmup to avoid replaying # graphs with stale KV cache block offsets from capture time. with self.no_cuda_graph(): self._general_warmup_impl(resource_manager, warmup_requests_configs) def _is_distributed_forward(self) -> bool: """Return whether model forward can communicate with peer workers. ``dist`` is optional. An engine built without a communicator cannot enter a collective at all, so it has no peers to strand and every warmup failure stays rank-local. """ if self.dist is None: return False return self.dist.world_size > 1 or self.mapping.dwdp_enabled def _warmup_agreement_allgather( self) -> Optional[Callable[[int], List[int]]]: """Return an allgather over the ranks this warmup forward synchronizes with. ``None`` means agreement cannot be established here, so a missing batch stays fatal rather than being skipped unilaterally. DWDP is the case that cannot be answered: its peers are reached through a ``COMM_WORLD``-derived subgroup built in ``dwdp.py``, not through ``self.dist``, so a ``self.dist`` allgather would report a unanimity it never observed. """ if self.dist is None or self.mapping.dwdp_enabled: return None if self.dist.world_size <= 1: return None return self.dist.allgather def _agree_warmup_flag(self, flag: bool) -> bool: """Reduce a phase-entry decision to one the whole forward group shares. Several predicates that gate a warmup phase are rank-local. The capturable guided decoder is installed only on the last pipeline rank, so ``guided_decoder is None`` -- and through it ``can_run_general_warmup`` -- differs across pipeline stages. Mamba's entry test reads this rank's free KV capacity. Letting one rank enter a phase its peers skip strands whoever enters the collective, and it also unbalances the per-shape agreement below. Any rank opting out takes the whole group out with it. """ allgather = self._warmup_agreement_allgather() if allgather is None: return flag return all(allgather(int(flag))) def _agree_warmup_shapes( self, configs: List[Tuple[int, int]]) -> List[Tuple[int, int]]: """Reduce warmup shapes to the ones every rank in the group proposed. ``_get_max_shape_warmup_requests`` derives shapes from this rank's free KV capacity, so under attention-DP the values -- and, once ``dict.fromkeys`` drops a collision, the length -- differ per rank. Ranks would then walk different loops and meet in different forwards. The intersection keeps rank 0's ordering, which matters because the general warmup list is ordered for torch.compile specialization. """ allgather = self._warmup_agreement_allgather() if allgather is None: return configs per_rank = allgather([list(config) for config in configs]) shared = set.intersection(*({tuple(config) for config in rank_configs} for rank_configs in per_rank)) agreed = [config for config in configs if config in shared] dropped = [config for config in configs if config not in shared] if dropped: logger.warning( f"Dropping warmup shapes {dropped} that not every rank could " f"propose; per-rank KV capacity differs. Remaining: {agreed}.") return agreed def _should_run_warmup_batch(self, batch: Optional[ScheduledRequests], num_tokens: int, shape: str) -> bool: """Decide whether this warmup shape runs, is skipped, or fails the rank. A rank that skips a shape its peers run leaves them blocked in that forward's collectives for the rest of the job. Skipping is therefore safe exactly when every rank in the forward group skips too, and an allgather establishes that before any rank enters the forward. The plan agreed by ``_agree_warmup_plan`` is what makes that allgather safe: every rank walks the same shape list, so this runs the same number of times everywhere. """ if not self._is_distributed_forward(): if batch is not None: return True # Safe to skip, but never silently: a skip during KV cache # estimation makes the profiling peak unrepresentative of # this shape. logger.warning(f"Skipping warmup shape ({shape}): not enough KV " f"cache space.") return False allgather = self._warmup_agreement_allgather() if allgather is None: # No reachable group to agree with. Keep the TP-only check, which # still catches the attention-DP asymmetry it was written for. self._assert_all_tp_ranks_have_warmup_batch(batch, num_tokens) if batch is None: raise RuntimeError( f"Warmup batch creation failed for shape ({shape}) on " f"global_rank={global_mpi_rank()}, " f"model_rank={self.dist.rank}, and this topology offers no " f"way to confirm that peers are skipping it too. They may " f"already be inside the matching forward, so this rank " f"cannot skip the shape without stranding them.") return True flags = list(allgather(int(batch is not None))) if all(flags): return True if not any(flags): # Every rank in the forward group is skipping, so none of them is # left inside a collective. This is the ordinary outcome for a # shape that does not fit the configuration at all, such as a # mixed context+generation shape under ``max_batch_size=1``. logger.info(f"Skipping warmup shape ({shape}) on all " f"{len(flags)} ranks: not enough KV cache space.") return False all_tokens = list(allgather(num_tokens)) failed_ranks = [i for i, flag in enumerate(flags) if not flag] raise RuntimeError( f"Warmup batch creation failed for shape ({shape}) on rank(s) " f"{failed_ranks} but succeeded on others, so entering this forward " f"would deadlock the ranks that still hold a batch. Per-rank " f"curr_max_num_tokens: {all_tokens}. This indicates asymmetric KV " f"cache capacity across ranks. Consider increasing " f"--kv_cache_free_gpu_mem_fraction.") def _assert_all_tp_ranks_have_warmup_batch(self, batch, num_tokens: int) -> None: """Assert every TP rank has a valid warmup batch, or raise with diagnostics. Under attention-DP, each rank's KV cache available capacity can differ at runtime, causing _create_warmup_request to return None on some ranks while others proceed into forward() with tp_comm collectives — deadlocking the job. This check prevents the deadlock by failing early with diagnostic info. ``tp_size`` alone does not establish that peers are reachable: ``dist`` is optional, and without a communicator there is no tp_comm collective to deadlock in. """ if self.mapping.tp_size <= 1 or self.dist is None: return has_batch = int(batch is not None) all_flags = list(self.dist.tp_allgather(has_batch)) if any(all_flags) and not all(all_flags): # Gather token counts for diagnostics all_tokens = list(self.dist.tp_allgather(num_tokens)) failed_ranks = [i for i, f in enumerate(all_flags) if not f] raise RuntimeError( f"Warmup batch creation failed on TP rank(s) {failed_ranks} " f"but succeeded on others. This would cause a collective " f"deadlock. Per-rank curr_max_num_tokens: {all_tokens}. " f"This indicates asymmetric KV cache capacity across TP ranks. " f"Consider increasing --kv_cache_free_gpu_mem_fraction.") def _general_warmup_impl( self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]) -> None: for num_tokens, num_gen_tokens in warmup_requests_configs: # Helix CP does not support warmup with context requests. if self.mapping.has_cp_helix() and num_tokens != num_gen_tokens: continue try: with self._release_batch_context( self._create_warmup_request(resource_manager, num_tokens, num_gen_tokens), resource_manager) as batch: if not self._should_run_warmup_batch( batch, num_tokens, f"general, num_tokens={num_tokens}, " f"num_gen_tokens={num_gen_tokens}"): continue logger.info( f"Run warmup with {num_tokens} tokens, include {num_gen_tokens} generation tokens" ) with self._warmup_timer.phase( f"general shape num_tokens={num_tokens}, " f"num_gen_tokens={num_gen_tokens}", record=False, log_start=False): self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) torch.cuda.synchronize() except torch.OutOfMemoryError: if self._is_distributed_forward(): # Peers are inside the same forward's collectives and # cannot follow a rank-local skip. raise logger.warning( f"OOM during general warmup with {num_tokens} tokens, " f"{num_gen_tokens} generation tokens. Skipping.") # If the OOM aborted the forward between dispatch() and # combine(), the MoE A2A state machines are stuck in # ``dispatched`` and the next warmup will hit # ``dispatch called twice``. Reset them before retrying a # smaller shape. self._reset_moe_alltoall_state() torch.cuda.empty_cache() def _reset_moe_alltoall_state(self) -> None: """Reset all MoE all-to-all state machines reachable from ``self.model``. Each MoE backend keeps a small dispatch/combine phase state per layer (``MoeAlltoAll`` or ``NVLinkOneSided``). A forward that calls ``dispatch`` but raises before reaching ``combine`` (e.g., a warmup OOM mid-MoE) leaves that state in ``dispatched``, which fails the invariant on the next ``dispatch`` call. This helper walks the model and resets any A2A state found, so subsequent forwards start clean. """ for module in self.model.modules(): for attr_name in ("moe_a2a", "comm"): obj = getattr(module, attr_name, None) reset = getattr(obj, "reset_state", None) if callable(reset): try: reset() except Exception as e: # noqa: BLE001 logger.warning( f"Failed to reset MoE A2A state on {type(module).__name__}.{attr_name}: {e}" ) def _run_attention_warmup(self, resource_manager: ResourceManager, can_run_general_warmup: bool = True) -> None: if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata): return @contextlib.contextmanager def trtllm_gen_fmha_jit_warmup(): previous = self._trtllm_gen_jit_warmup self._trtllm_gen_jit_warmup = True try: yield finally: self._trtllm_gen_jit_warmup = previous logger.info("Running TRTLLM-Gen FMHA JIT warmup") warmup_requests_configs = [] if not self.is_draft_model and self.guided_decoder is None: # doesn't support 2-model speculative draft and guided decoding warmup_requests_configs.append( (1 + self.max_total_draft_tokens, 1)) # one generation request else: logger.debug("Skipped TRTLLM-Gen FMHA JIT warmup for Gen kernels") if can_run_general_warmup: warmup_requests_configs.append((1, 0)) # one context token else: logger.debug("Skipped TRTLLM-Gen FMHA JIT warmup for Ctx kernels") model_type = getattr(self.model.model_config.pretrained_config, "model_type", None) if can_run_general_warmup and model_type in ("kimi_k3", "kimi_linear"): # Kimi's one-token context takes the NT < 4 FLA fallback and does # not compile the optimized single-sequence K123 variant. A # non-aligned five-chunk context enters the pure K123 path. _KIMI_KDA_PREFILL_WARMUP_TOKENS = 257 logger.info("Adding Kimi KDA pure-prefill warmup with " f"{_KIMI_KDA_PREFILL_WARMUP_TOKENS} context tokens") warmup_requests_configs.append((_KIMI_KDA_PREFILL_WARMUP_TOKENS, 0)) if (not self.is_draft_model and self.guided_decoder is None and can_run_general_warmup): # The cute_dsl_mla FMHA lib now only support the generation-only batch, we need to warmup the TRTLLM-Gen FMHA lib for the mixed context+generation batch. # One MIXED context+generation batch (1 ctx token + 1 gen request). warmup_requests_configs.append( (1 + self.max_total_draft_tokens + 1, 1)) else: logger.debug( "Skipped TRTLLM-Gen flashinfer_trtllm_gen FMHA lib JIT warmup When enable cute_dsl_mla FMHA lib" ) for num_tokens, num_gen_requests in warmup_requests_configs: warmup_request = self._create_warmup_request( resource_manager, num_tokens=num_tokens, num_gen_requests=num_gen_requests) with self.no_cuda_graph(), self._release_batch_context( warmup_request, resource_manager) as batch: if not self._should_run_warmup_batch( batch, num_tokens, f"attention, num_tokens={num_tokens}, " f"num_gen_requests={num_gen_requests}"): continue # The first forward of the process lands here, so this shape # also absorbs every first-touch JIT (CuTe DSL GEMM, FMHA # NVRTC grid, ...). Time it per shape so a slow startup can be # attributed without a debugger. with self._warmup_timer.phase( f"attention shape num_tokens={num_tokens}, " f"num_gen_requests={num_gen_requests}", record=False, log_start=True): with trtllm_gen_fmha_jit_warmup(): self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) torch.cuda.synchronize() @staticmethod def _release_megamoe_profiling_scratch(): # MegaMoE tuning resources are shared across layers, so only the engine # can release them after its full autotune warmup and before graph # capture. Later eviction could invalidate a captured workspace pointer. from ..moe.custom_ops import cute_dsl_megamoe_custom_op as _megamoe_op release_megamoe_scratch = getattr(_megamoe_op, "release_megamoe_profiling_scratch", None) if release_megamoe_scratch is not None: release_megamoe_scratch() def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: """Runs forward passes to populate the autotuner cache.""" from ..custom_ops.torch_custom_ops import ( IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, MXFP8GemmRunner) from ..modules.linear import (MXFP8LinearMethod, flashinfer_mxfp8_autotune) enable_trtllm_autotuner = self.llm_args.enable_autotuner if not enable_trtllm_autotuner: return mxfp8_methods = [] for module in self.model.modules(): quant_method = getattr(module, "quant_method", None) if isinstance(quant_method, MXFP8LinearMethod): mxfp8_methods.append(quant_method) # This engine owns startup warmup, so it explicitly opts its MXFP8 # methods into native tuning. Standalone modules and engine paths that # skip this warmup remain on the direct native op. for method in mxfp8_methods: method.enable_native_autotune() # Native and FlashInfer tuning are independent. Capture native # eligibility before enabling graph-only FlashInfer dispatch. native_mxfp8_methods = [ method for method in mxfp8_methods if method.needs_native_autotune ] compile_all_batches = (self._torch_compile_enabled and not self._torch_compile_prefill_only) if (compile_all_batches and "TRTLLM_MXFP8_GEMM_BACKEND" not in os.environ): # Compiled auto dispatch uses native; do not tune unused backends. # Prefill-only compile retains the eager generation-graph policy. for method in mxfp8_methods: method.disable_flashinfer_auto() use_mxfp8_flashinfer_graph_default = ( self.cuda_graph_runner.enabled and not compile_all_batches and "TRTLLM_MXFP8_GEMM_BACKEND" not in os.environ and any( getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) for module in self.model.modules())) if use_mxfp8_flashinfer_graph_default: # CuTeDSL is the alternative backend; PP has no graph-pass handoff. tune_with_cute_dsl = (IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE and not self.mapping.has_pp()) for quant_method in mxfp8_methods: quant_method.enable_flashinfer_auto() quant_method.tune_decode_graph_backends = ( tune_with_cute_dsl and quant_method.uses_flashinfer) # An explicit auto setting stays intact, but compiled auto dispatch # uses native GEMM. Do not run or mark an unused FlashInfer tuning pass. flashinfer_mxfp8_methods = [ method for method in mxfp8_methods if method.needs_flashinfer_autotune and ( not compile_all_batches or method.backend == "flashinfer") ] # Every TP and PP rank must make the same backend decision before any # rank returns or enters a tuning forward with model collectives. if self.mapping.tp_size > 1 or self.mapping.has_pp(): local_flashinfer_enabled = int(bool(flashinfer_mxfp8_methods)) all_flashinfer_enabled = [local_flashinfer_enabled] if self.mapping.tp_size > 1: all_flashinfer_enabled = list( self.dist.tp_allgather(local_flashinfer_enabled)) if self.mapping.has_pp(): all_flashinfer_enabled = [ enabled for stage_flags in self.dist.pp_allgather( all_flashinfer_enabled) for enabled in stage_flags ] if any(all_flashinfer_enabled) and not all(all_flashinfer_enabled): forced_flashinfer = any(method.backend == "flashinfer" for method in mxfp8_methods) for method in mxfp8_methods: method.disable_flashinfer_auto() flashinfer_mxfp8_methods = [] if forced_flashinfer: raise RuntimeError( "FlashInfer MXFP8 was explicitly requested but is not " "available on every TP/PP rank") logger.warning( "FlashInfer MXFP8 availability differs across TP/PP ranks; " "using the native TensorRT-LLM GEMM backend on every rank.") enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods) enable_native_mxfp8_autotuner = bool(native_mxfp8_methods) AutoTuner.get().setup_distributed_state(self.mapping, self.dist) logger.info( f"Running autotuner warmup (TRT-LLM={enable_trtllm_autotuner}, " f"native MXFP8={enable_native_mxfp8_autotuner}, " f"FlashInfer MXFP8={enable_flashinfer_mxfp8_autotuner})...") kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) token_num_upper_bound = min(self.max_num_tokens, self.batch_size * (self.max_seq_len - 1)) curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( token_num_upper_bound=token_num_upper_bound, max_num_draft_tokens=self.original_max_draft_len) warmup_configs = [(curr_max_num_tokens, 0)] if (not self.is_draft_model and self.guided_decoder is None and not self.mapping.has_pp()): # Add generation request to warmup the autotuner cache. warmup_configs.append((1 + self.max_total_draft_tokens, 1)) def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> bool: """Run one isolated tuning pass with fresh synthetic batches.""" ran_forward = False with self.no_cuda_graph(), autotune_context: for num_tokens, num_gen_requests in warmup_configs: warmup_request = self._create_warmup_request( resource_manager, num_tokens, num_gen_requests) with self._release_batch_context(warmup_request, resource_manager) as batch: if not self._should_run_warmup_batch( batch, num_tokens, f"autotuner, num_tokens={num_tokens}, " f"num_gen_requests={num_gen_requests}"): continue # Reset the flag is_first_draft for the draft model. # This is necessary for overlap scheduler. spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) if self.is_draft_model and isinstance( spec_resource_manager, Eagle3ResourceManager): spec_resource_manager.is_first_draft = True with self._warmup_timer.phase( f"autotuner shape num_tokens={num_tokens}, " f"num_gen_requests={num_gen_requests}", record=False, log_start=True): self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) ran_forward = True torch.cuda.synchronize() if ran_forward and synchronize_trtllm_cache: # pp_recv in AutoTuner choose_one will never be called if there is no tuning op during the forward pass. # So we need to make an extra call to consume the previous rank's pp_send to guarantee that the previous rank's pp_send is released. AutoTuner.get().cache_pp_recv() # Send the cache after the tuning process to the next PP rank AutoTuner.get().cache_pp_send() # Clean the pp flag to avoid deadlock with synchronous send/recv AutoTuner.get().clean_pp_flag() return ran_forward cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) ran_native_forward = run_autotuner_pass(autotune(cache_path=cache_path), synchronize_trtllm_cache=True) ran_flashinfer_forward = False if enable_flashinfer_mxfp8_autotuner: ran_flashinfer_forward = run_autotuner_pass( flashinfer_mxfp8_autotune(), synchronize_trtllm_cache=False) if enable_flashinfer_mxfp8_autotuner: if ran_flashinfer_forward: for method in flashinfer_mxfp8_methods: method.mark_flashinfer_autotuned() else: forced_flashinfer = any(method.backend == "flashinfer" for method in flashinfer_mxfp8_methods) for method in flashinfer_mxfp8_methods: method.disable_flashinfer_auto() if forced_flashinfer: raise RuntimeError( "FlashInfer MXFP8 was explicitly requested but its autotuner " "warmup forward could not run") logger.warning( "FlashInfer MXFP8 autotuning could not run; using the native " "TensorRT-LLM GEMM backend.") if enable_native_mxfp8_autotuner: if ran_native_forward: MXFP8GemmRunner.sync_all_tactic_caches(AutoTuner.get()) for method in native_mxfp8_methods: method.mark_native_autotuned() else: for method in native_mxfp8_methods: method.disable_native_autotune() logger.warning( "Native MXFP8 autotuning had no runnable warmup batch; " "using the default native GEMM tactic.") logger.info( f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}" ) AutoTuner.get().print_profiling_cache() self._release_megamoe_profiling_scratch() # Clear workspace buffers allocated during the autotuner forward pass. # The autotuner runs a context-only forward with max_num_tokens, which # causes the global Buffers pool to cache large MoE/GEMM workspaces. # If not cleared, these inflate the memory baseline seen by the KV cache # profiler, reducing memory available for activations during inference. clear_memory_buffers() torch.cuda.empty_cache() def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager) -> None: """Pre-JIT the Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels. Mamba hybrid models (e.g. Nemotron 3 Super 120B, Nemotron-Nano-12B-v2) skip ``_general_warmup`` because ``can_run_general_warmup`` is False when the KV cache manager is a ``MambaHybridCacheManager``. The default ``_run_autotuner_warmup`` then issues a single ``least_requests=True`` prefill = 1 sequence with ``num_cached_tokens_per_seq = 0``, which only compiles the ``num_seqs == 1`` / ``HAS_INITSTATES=False`` variants of the SSD kernels. The first real serve iteration with chunked prefill and multiple context requests then triggers autotune of the missing variants mid-inference, producing a ~30 s stall / large P99 spike. This method runs two extra forward passes to compile those variants during warmup: 1. ``least_requests=False`` — splits ``curr_max_num_tokens`` into many short sequences, forcing the multi-seq path of ``cu_seqlens_to_chunk_indices_offsets_triton`` and its ``_cu_seqlens_triton_kernel``. 2. ``least_requests=False`` inside ``Mamba2Metadata.force_initial_states_for_warmup()`` — same as (1) plus the ``HAS_INITSTATES=True`` variants of ``_state_passing_fwd_kernel``, ``_chunk_scan_fwd_kernel``, and ``_chunk_state_varlen_kernel``. Runs regardless of ``enable_autotuner``. Wraps in ``autotune()`` when the autotuner is enabled so op-level (M,N,K) caches also get primed for these shapes. Set ``TLLM_MAMBA_MULTISEQ_WARMUP=0`` to disable. """ if os.environ.get("TLLM_MAMBA_MULTISEQ_WARMUP", "1") != "1": return kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) if kv_cache_manager is None or not isinstance(kv_cache_manager, MambaHybridCacheManager): return token_num_upper_bound = min(self.max_num_tokens, self.batch_size * (self.max_seq_len - 1)) curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( token_num_upper_bound=token_num_upper_bound, max_num_draft_tokens=self.original_max_draft_len) # Rank-local capacity, so peers can disagree. Leaving the phase alone # would unbalance the per-shape agreement inside it. if not self._agree_warmup_flag(curr_max_num_tokens >= 4): return # Cap the multi-seq warmup token count so we don't fill the KV cache # to the brim. The autotuner warmup that ran just before this uses # ``least_requests=True`` (few long sequences) which fits comfortably # 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`` > 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 # decoding on a single H100: 259 blocks total, ``max_num_tokens=8192`` # nearly saturates it), that extra per-sequence block overflows the # pool and crashes with "Can't allocate new blocks for window size N". # The point of this warmup is only to trigger ``num_seqs > 1`` + # ``HAS_INITSTATES=True`` kernel variants — a modest token budget # achieves that with plenty of block headroom. WARMUP_TOKEN_CAP = 4096 capped_num_tokens = min(curr_max_num_tokens, WARMUP_TOKEN_CAP) logger.info( "Running Mamba hybrid warmup (multi-seq + HAS_INITSTATES=True)...") # (num_tokens, num_gen_requests, least_requests, force_initstates) mamba_warmup_shapes = [ (capped_num_tokens, 0, False, False), (capped_num_tokens, 0, False, True), ] autotuner_enabled = self.llm_args.enable_autotuner cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) autotune_ctx = (autotune(cache_path=cache_path) if autotuner_enabled else contextlib.nullcontext()) with self.no_cuda_graph(), autotune_ctx: for (num_tokens_i, num_gen_requests_i, least_req_i, force_init_i) in mamba_warmup_shapes: init_ctx = (Mamba2Metadata.force_initial_states_for_warmup() if force_init_i else contextlib.nullcontext()) shape = (f"Mamba hybrid, num_tokens={num_tokens_i}, " f"num_gen_requests={num_gen_requests_i}, " f"force_initstates={force_init_i}") with init_ctx: try: warmup_request = self._create_warmup_request( resource_manager, num_tokens_i, num_gen_requests_i, least_requests=least_req_i) except torch.OutOfMemoryError as e: if self._is_distributed_forward(): raise logger.warning(f"Warmup skipped for shape ({shape}): " f"{type(e).__name__}: {e}") torch.cuda.empty_cache() continue except RuntimeError as e: # The known KV allocation failure happens before # forward and is recoverable only when no peer worker # can advance independently. Any other RuntimeError is # a defect, not a capacity limit, and is fatal. if self._is_distributed_forward(): raise if "Can't allocate new blocks for window size" not in str( e): raise logger.warning(f"Warmup skipped for shape ({shape}): " f"{type(e).__name__}: {e}") torch.cuda.empty_cache() continue try: with self._release_batch_context( warmup_request, resource_manager) as batch: if not self._should_run_warmup_batch( batch, num_tokens_i, shape): continue spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) if self.is_draft_model and isinstance( spec_resource_manager, Eagle3ResourceManager): spec_resource_manager.is_first_draft = True self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) if autotuner_enabled: AutoTuner.get().cache_pp_recv() AutoTuner.get().cache_pp_send() AutoTuner.get().clean_pp_flag() torch.cuda.synchronize() # Once peers can enter warmup synchronization or forward, # any rank-local exception can strand them in a collective. except Exception as e: # noqa: BLE001 if self._is_distributed_forward(): raise # ``torch.OutOfMemoryError`` is a ``RuntimeError`` # subclass; anything outside that hierarchy is a defect # rather than a capacity limit. if not isinstance(e, RuntimeError): raise # A single-rank warmup is a pure perf optimization. If # a forward shape does not fit, it can be compiled # lazily on the first real request. logger.warning(f"Warmup skipped for shape ({shape}): " f"{type(e).__name__}: {e}") # An OOM between dispatch() and combine() leaves the # local MoE A2A state in ``dispatched``. self._reset_moe_alltoall_state() torch.cuda.empty_cache() clear_memory_buffers() torch.cuda.empty_cache() def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: """Compute graph_bs → draft_len mapping for dynamic draft length feature. Example: draft_len_schedule = {4:4, 8:2, 32:1}, cuda_graph_batch_sizes = [1,2,3,4,5,6,7,8,16,24,32,64] - Batch sizes 1-4: use draft_len=4 (up to key 4) - Batch sizes 5-8: use draft_len=2 (up to key 8) - Batch sizes 9-32: use draft_len=1 (up to key 32) - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled) Returns: {1:4, 2:4, 3:4, 4:4, 5:2, 6:2, 7:2, 8:2, 16:1, 24:1, 32:1, 64:0} """ # Dynamic draft length for CUDA graphs is only supported for one-model path if (not self.spec_config or not self.spec_config.draft_len_schedule or not self.spec_config.spec_dec_mode.support_dynamic_draft_len()): return None schedule = self.spec_config.draft_len_schedule schedule_keys = list(schedule.keys()) mapping = {} key_idx = 0 for graph_bs in self._cuda_graph_batch_sizes: while key_idx < len( schedule_keys) and schedule_keys[key_idx] < graph_bs: key_idx += 1 if key_idx < len(schedule_keys): draft_len = schedule[schedule_keys[key_idx]] else: draft_len = 0 mapping[graph_bs] = draft_len return mapping def _get_graphs_to_capture( self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, int]]: """Determine which (batch_size, draft_len) graphs to capture. Returns: List of (batch_size, draft_len) tuples for CUDA graph capture. """ # Case 1: Draft model (two-model speculative decoding) # Two-model path is deprecated and will be removed in the near future if self.is_draft_model: draft_len = self.max_total_draft_tokens return [(bs, draft_len) for bs in cuda_graph_batch_sizes] # Case 2: One-model with dynamic draft length if self.spec_config is not None and self.spec_config.draft_len_schedule is not None and self.spec_config.spec_dec_mode.support_dynamic_draft_len( ): graphs = [(graph_bs, draft_len) for graph_bs, draft_len in self._dynamic_draft_len_mapping.items()] # Workaround for dynamic draft length: # capture the maximum speculative graph shape up front. Dynamic draft length # breaks the previous assumption that attention workspace demand can be safely # ordered by batch size alone; a later graph shape may require a larger shared # graph workspace, and resizing that workspace can change its data_ptr and # invalidate pointers captured by earlier graphs, causing illegal memory access # on replay. # # This adds the overhead of one extra captured graph, and that graph is not # expected to be used by the normal schedule-driven dynamic draft-length path. # # Follow-up first-principles fix: # query or precompute the exact attention workspace requirement for all # reachable graph shapes, pre-size the shared graph workspace once without # capturing an extra graph, and avoid resizing it in graph mode afterward. max_spec_graph = (max(cuda_graph_batch_sizes), self.original_max_draft_len) if max_spec_graph not in graphs: graphs.append(max_spec_graph) logger.info(f"Dynamic draft length enabled for one-model path. " f"Capturing {len(graphs)} graphs: {graphs}") return graphs # Case 3: Target model (two-model) or one-model without dynamic draft # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. # spec_config is None for non-spec models — fall back to max_draft_len (= 0). draft_lengths = [get_static_draft_len(self)] should_capture_no_spec = ( self.max_total_draft_tokens > 0 and not self.spec_config.spec_dec_mode.use_one_engine() # Assume speculation is always on if no max_concurrency set (saves memory) and self.spec_config.max_concurrency is not None) if should_capture_no_spec: draft_lengths.append(0) return [(bs, draft_len) for bs in cuda_graph_batch_sizes for draft_len in draft_lengths] def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" if not (self.cuda_graph_runner.enabled or self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED): return from ..modules.linear import (MXFP8LinearMethod, flashinfer_mxfp8_autotune, flashinfer_mxfp8_decode_graph_capture) # The automatic MiniMax-M3 MXFP8 selection is decode-graph-only. # Tune every generation graph shape during the warmup-only pass. Keep # piecewise context/prefill graph capture on the native backend. flashinfer_methods = [ quant_method for module in self.model.modules() if isinstance((quant_method := getattr(module, "quant_method", None) ), MXFP8LinearMethod) and quant_method.needs_flashinfer_autotune ] flashinfer_autotune_context = ( flashinfer_mxfp8_autotune() if self.cuda_graph_runner.is_warmup_only and flashinfer_methods else contextlib.nullcontext()) with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture( ): self._capture_generation_cuda_graphs(resource_manager) self._capture_mixed_encoder_decoder_cuda_graphs(resource_manager) # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. if not self.cuda_graph_runner.is_warmup_only: self._capture_prefill_cuda_graphs(resource_manager) def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): """Warm up or capture pure-generation CUDA graph shapes.""" if not self.cuda_graph_runner.enabled: return operation = ("warmup" 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.") # 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) 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 # rank only holds max_seq_len / cp_size tokens, so scale accordingly to # avoid creating warmup requests whose position_ids exceed the RoPE # table (max_position_embeddings). effective_max_seq_len = self.max_seq_len if self.mapping is not None and self.mapping.has_cp_helix(): effective_max_seq_len = self.max_seq_len // self.mapping.cp_size sparse_config = self.sparse_attention_config if (isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) and sparse_config.needs_separate_short_long_cuda_graphs()): # For short sequences, subtract the maximum runtime tokens consumed # by a generation step so all current-step tokens stay within the # sequence length threshold. PARD uses 2K tokens here, not K+1. max_runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.max_draft_len) # For long sequences, use the default maximum sequence length. max_seq_len = (sparse_config.seq_len_threshold - max_runtime_tokens_per_gen_step) if max_seq_len < effective_max_seq_len: max_seq_len_list = [effective_max_seq_len, max_seq_len] else: max_seq_len_list = [effective_max_seq_len] else: max_seq_len_list = [effective_max_seq_len] def prepare_cross_batch(batch: ScheduledRequests, resource_manager: ResourceManager) -> None: """Populate dummy gen requests' cross-KV cache before capture. Dummy generation requests used for graph capture never ran a context step, so their cross-KV cache blocks are uninitialized and captured kernels would read garbage. Temporarily switch each request to a one-token context chunk with a fake encoder output to run just the cross-KV projection (via _populate_cross_kv_cache), then restore generation state for the actual capture. """ if not batch.generation_requests: return max_encoder_output_len = self._get_max_encoder_output_len( resource_manager) hidden_size = self._get_enc_dec_hidden_size() saved_request_state = [] for request in batch.generation_requests: saved_request_state.append( (request, request.py_encoder_output, request.py_skip_cross_kv_projection, request.state, request.py_batch_idx, request._cached_tokens, request._cached_tokens_set)) request.py_encoder_output = torch.ones( (max_encoder_output_len, hidden_size), device="cuda", dtype=self.dtype) request.py_skip_cross_kv_projection = False request.state = LlmRequestState.CONTEXT_INIT request.context_current_position = 0 request.context_chunk_size = 1 projection_batch = ScheduledRequests() projection_batch.reset_context_requests(batch.generation_requests) kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) attn_metadata = self._set_up_attn_metadata(kv_cache_manager, draft_kv_cache_manager) with self.no_cuda_graph(): projection_inputs, _ = self._prepare_inputs( projection_batch, kv_cache_manager, attn_metadata, spec_metadata=None, new_tensors_device=None, resource_manager=resource_manager, maybe_graph=False) self._populate_cross_kv_cache(projection_inputs) torch.cuda.synchronize() for (request, encoder_output, skip_cross_kv_projection, state, batch_idx, cached_tokens, cached_tokens_set) in saved_request_state: request.py_encoder_output = encoder_output request.py_skip_cross_kv_projection = skip_cross_kv_projection request.state = state if state == LlmRequestState.GENERATION_IN_PROGRESS: request.context_current_position = request.prompt_len request.py_batch_idx = batch_idx request._cached_tokens = cached_tokens request._cached_tokens_set = cached_tokens_set def _run_capture_pass(force_non_greedy: bool, label: str, force_lora_graph: bool, sample_type: Optional[SampleType] = None) -> None: assert self._force_lora_graph_for_capture is None self._force_lora_graph_for_capture = force_lora_graph # Pin the sampling tier for this pass. maybe_get_cuda_graph reads # it to build the graph key, so every graph captured below records # the kernels of this tier and nothing else. Passes that do not name # a tier capture FULL graphs -- ones carrying no sampling at all -- # which is what a batch resolving to FULL replays. pinned_tier = sample_type or SampleType.FULL self._capture_sample_type = pinned_tier self.cuda_graph_runner.set_capture_sample_type(pinned_tier) try: for bs, draft_len in graphs_to_capture: if bs > self.batch_size: continue for max_seq_len in max_seq_len_list: warmup_request = self._create_cuda_graph_warmup_request( resource_manager, bs, draft_len, max_seq_len, force_non_greedy=force_non_greedy) with self._release_batch_context( warmup_request, resource_manager) as batch: if batch is None: # No KV cache space for this batch size. During KV # cache estimation this makes the profiling peak # unrepresentative (the final executor still # captures this graph), so don't skip silently. logger.warning( f"Skipping CUDA graph warmup ({label}) for " f"batch size={bs}, draft_len={draft_len}: " f"not enough KV cache space.") continue logger.info( f"Run generation-only CUDA graph {operation} ({label}) " f"for batch size={bs}, draft_len={draft_len}, " f"max_seq_len={max_seq_len}") self.enable_spec_decode = draft_len > 0 or self.is_draft_model or ( self.spec_config is not None and self.spec_config.spec_dec_mode.use_one_engine()) self._update_draft_inference_state_for_warmup( batch, draft_len > 0, resource_manager) if self._is_encoder_decoder_model(): prepare_cross_batch(batch, resource_manager) self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) torch.cuda.synchronize() finally: self._force_lora_graph_for_capture = None self._capture_sample_type = None self.cuda_graph_runner.set_capture_sample_type(None) if self.cuda_graph_lora_manager is None: lora_graph_cases = [False] elif self.llm_args.lora_config.cuda_graph_specialize_lora: # Capture the larger LoRA graph first so the base-only graph can # reuse its CUDA graph memory-pool allocations. lora_graph_cases = [True, False] else: lora_graph_cases = [True] # Which variants to capture depends on which sampler this engine will # actually run, and the two are mutually exclusive: a spec-decoding mode # in a one-engine mode samples inside its worker and gets a dedicated # sampler from get_spec_decoder, so TorchSampler -- the only sampler # implementing in-graph sampling -- never runs. Capturing the other # branch's variants would spend warmup on graphs whose keys no batch can # ever produce. # # The predicate is use_one_engine() rather than has_spec_decoder(): the # two-model eagle3 / mtp_eagle modes also "have a spec decoder", but # get_spec_decoder hands them a plain TorchSampler, so they belong on # the TorchSampler branch and do need the FAST graphs. uses_own_spec_decoder = ( self.spec_config is not None and self.spec_config.spec_dec_mode.use_one_engine()) def _capture_variant(label: str, force_non_greedy: bool = False, sample_type: Optional[SampleType] = None) -> None: for use_lora_graph in lora_graph_cases: variant_label = label if self.cuda_graph_lora_manager is not None: variant_label += (", LoRA" if use_lora_graph else ", base-only") _run_capture_pass(force_non_greedy=force_non_greedy, label=variant_label, force_lora_graph=use_lora_graph, sample_type=sample_type) if uses_own_spec_decoder: # One-engine spec branch: the greedy argmax fast path plus the # advanced-sampling variant. The latter is needed because on-the-fly # capture is disabled outside warmup, so a batch containing a # non-greedy request would otherwise fall back to eager. # # Dummy warmup requests carry no sampling params, so the greedy pass # needs no override while the advanced one has to force the flag. _capture_variant("greedy") _capture_variant("advanced sampling", force_non_greedy=True) else: # TorchSampler branch: FULL carries no sampling in the graph and is # what every batch replays unless in-graph sampling is enabled, so it # is always captured. FAST adds the in-graph sampling kernels and is # captured only when opted in, since it costs extra warmup time and # memory. _capture_variant("full", sample_type=SampleType.FULL) if self.enable_in_graph_sampling: # Give the warmup requests real non-greedy sampling params: # dummies otherwise carry none, resolve to greedy, and the # capture would record argmax rather than the fast tier's # top-k/top-p kernels. Same substitution the advanced-sampling # pass relies on. _capture_variant("fast", force_non_greedy=True, sample_type=SampleType.FAST) # Set the value back to the original value after cuda graph warmups are complete self.enable_spec_decode = self.is_spec_decode # update_is_all_greedy_sample inside each forward call during the # non-greedy capture pass leaves is_all_greedy_sample=False on # spec_metadata. Reset it so the first real iteration starts clean; # update_is_all_greedy_sample will refresh it on every iteration anyway. # This is a defensive guard. if self.spec_metadata is not None: self.spec_metadata.is_all_greedy_sample = True def _capture_mixed_encoder_decoder_cuda_graphs( self, resource_manager: ResourceManager) -> None: """Warm and capture reachable mixed encoder-decoder graph shapes. The first global CUDA-graph pass warms every shape so shared attention workspace reaches its final size. The second pass captures the same shapes. Runtime capture is deliberately disabled because graph capture executes KV-cache writes and must never run against live requests. """ runner = self.cuda_graph_runner if not runner.enable_encoder_decoder_mixed_cuda_graph: return max_encoder_output_len = self._get_max_encoder_output_len( resource_manager) context_shapes = set(self._encoder_graph_shapes) if not context_shapes: logger.warning("Skipping mixed encoder-decoder CUDA graph capture: " "no encoder CUDA graph shapes are configured.") return max_encoder_batch_size = max(batch_size for batch_size, _ in context_shapes) max_batch_token_counts = { total_tokens for batch_size, total_tokens in context_shapes if batch_size == max_encoder_batch_size } paired_context_count = 2 * max_encoder_batch_size if runner.max_supported_batch_size > paired_context_count: paired_token_counts = { first + second for first in max_batch_token_counts for second in max_batch_token_counts } context_shapes.update((paired_context_count, token_count) for token_count in paired_token_counts) operation = ("warmup" if runner.is_warmup_only else "capture") hidden_size = self._get_enc_dec_hidden_size() max_num_encoder_tokens = max( (total_encoder_tokens for num_contexts, total_encoder_tokens in context_shapes if total_encoder_tokens <= num_contexts * max_encoder_output_len and any(batch_size > num_contexts for batch_size in runner.supported_batch_sizes)), default=0) if max_num_encoder_tokens == 0: return model_config = self.model.model_config.pretrained_config # The capture query length must equal the runtime decoder prefix or # every mixed batch misses its graph, silently and with no counter to # show it. Prefer the input processor's actual prefix (Whisper forces # [decoder_start, lang, task, no_timestamps] = 4); fall back to the # token-model heuristic: BART/mBART prepend a forced BOS token after # decoder_start, T5 uses decoder_start alone. prefix_fn = getattr(self.input_processor, "get_decoder_prefix_len", None) mixed_context_query_len = prefix_fn() if prefix_fn is not None else None if not mixed_context_query_len: mixed_context_query_len = (2 if getattr( model_config, "model_type", None) in ("bart", "mbart") else 1) logger.info("Mixed encoder/decoder graph capture using decoder prefix " f"length {mixed_context_query_len}.") for num_contexts, total_encoder_tokens in sorted( context_shapes, key=lambda shape: shape[1], reverse=True): if total_encoder_tokens > num_contexts * max_encoder_output_len: continue base_encoder_len, remainder = divmod(total_encoder_tokens, num_contexts) encoder_output_lens = ([base_encoder_len + 1] * remainder + [base_encoder_len] * (num_contexts - remainder)) if not encoder_output_lens or encoder_output_lens[-1] <= 0: continue for batch_size in runner.supported_batch_sizes: if batch_size <= num_contexts: continue warmup_request = self._create_cuda_graph_warmup_request( resource_manager, batch_size, draft_len=0, mixed_context_encoder_output_lens=encoder_output_lens, mixed_context_query_len=mixed_context_query_len) with self._release_batch_context(warmup_request, resource_manager) as batch: if batch is None: logger.warning( "Skipping mixed encoder-decoder CUDA graph " f"{operation}: not enough KV cache space for " f"batch size={batch_size}.") continue context_requests = batch.context_requests for request, encoder_output_len in zip( context_requests, encoder_output_lens): request.state = LlmRequestState.CONTEXT_INIT request.context_current_position = 0 request.context_chunk_size = mixed_context_query_len request.cached_tokens = 0 request.py_batch_idx = None request.py_encoder_output = torch.ones( (encoder_output_len, hidden_size), device="cuda", dtype=self.dtype, ) request.py_skip_cross_kv_projection = False runner._get_static_encoder_hidden_states( context_requests[0].py_encoder_output, max_num_encoder_tokens, allow_allocate=True, ) logger.info("Run mixed encoder-decoder CUDA graph " f"{operation} for batch size={batch_size}, " f"context requests={num_contexts}, " f"packed encoder tokens={total_encoder_tokens}") saved_enable_spec_decode = self.enable_spec_decode # The builder above has already set runtime_draft_len to 0, # so this save/restore cannot recover the pre-builder value. # This is harmless while encoder-decoder models use # max_draft_len == 0. Before supporting nonzero drafting, # save the length before building the batch and extend the # try/finally to cover batch creation as well. saved_runtime_draft_len = self.runtime_draft_len try: self.enable_spec_decode = False self.runtime_draft_len = 0 self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) torch.cuda.synchronize() finally: self.enable_spec_decode = saved_enable_spec_decode self.runtime_draft_len = saved_runtime_draft_len def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): """Capture configured CUDA graphs for context/prefill steps.""" if (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.DISABLED or (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE and not self._torch_compile_enabled)): return logger.info("Running prefill CUDA graph warmup...") prefill_cuda_graph_num_tokens = sorted( self._prefill_cuda_graph_num_tokens, reverse=True) capture_context = (capture_piecewise_cuda_graph(True) if self._torch_compile_piecewise_cuda_graph else contextlib.nullcontext()) with capture_context, self.no_cuda_graph(): for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request( resource_manager, num_tokens, 0) with self._release_batch_context(warmup_request, resource_manager) as batch: self._assert_all_tp_ranks_have_warmup_batch( batch, num_tokens) if batch is None: continue logger.info( f"Run prefill CUDA graph capture for num tokens={num_tokens}" ) if self.breakable_cuda_graph_runner is not None: self.breakable_cuda_graph_runner.capture( num_tokens, lambda: self.forward( batch, new_tensors_device=None, resource_manager=resource_manager)) else: # Run a few times to ensure torch.compile capture. for _ in range(4): self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) # The logits allocations grow with the number of requests and are not # part of the captured model body. Warm up the largest request count so # those allocations can be reused during stable inference. for num_tokens in prefill_cuda_graph_num_tokens: warmup_request = self._create_warmup_request(resource_manager, num_tokens, 0, least_requests=False) with self._release_batch_context(warmup_request, resource_manager) as batch: self._assert_all_tp_ranks_have_warmup_batch(batch, num_tokens) if batch is None: continue logger.info( f"Run prefill CUDA graph warmup for num tokens={num_tokens} with most requests" ) if self.breakable_cuda_graph_runner is not None: with self.no_cuda_graph(): self.breakable_cuda_graph_runner.warmup( lambda: self.forward(batch, new_tensors_device=None, resource_manager= resource_manager), steps=1) else: self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) torch.cuda.synchronize() ### Helper methods promoted from the original warmup method ### @contextlib.contextmanager def _release_batch_context(self, batch: Optional[ScheduledRequests], resource_manager: ResourceManager): """A context manager to automatically free resources of a dummy batch.""" kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) cross_kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.CROSS_KV_CACHE_MANAGER) spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) try: yield batch finally: if batch is not None and kv_cache_manager is not None: for req in batch.all_requests(): kv_cache_manager.free_resources(req) if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(req) if cross_kv_cache_manager is not None: cross_kv_cache_manager.free_resources(req) if spec_resource_manager is not None: spec_resource_manager.free_resources(req) def _create_warmup_request( self, resource_manager: ResourceManager, num_tokens: int, num_gen_requests: int, least_requests: bool = True) -> Optional[ScheduledRequests]: """Creates a generic dummy ScheduledRequests object for warmup.""" kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) available_tokens = kv_cache_manager.get_num_available_tokens( token_num_upper_bound=num_tokens, max_num_draft_tokens=self.max_total_draft_tokens) available_blocks = kv_cache_manager.get_num_free_blocks() if num_tokens > self.max_num_tokens or num_tokens > available_tokens: return None if num_gen_requests > self.batch_size: return None num_gen_tokens = num_gen_requests * (1 + self.max_total_draft_tokens) if num_gen_tokens > self.max_num_tokens: return None num_ctx_tokens = num_tokens - num_gen_tokens num_ctx_requests = 0 ctx_requests = [] gen_requests = [] # 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 num_full_seqs = 0 num_left_over_tokens = 0 max_context_requests = self.batch_size - num_gen_requests if max_context_requests * max_seq_len < num_ctx_tokens: return None if num_ctx_tokens > 0: if least_requests: num_full_seqs = num_ctx_tokens // max_seq_len num_left_over_tokens = num_ctx_tokens - num_full_seqs * max_seq_len else: max_bs = min(num_ctx_tokens, max_context_requests) if num_ctx_tokens % max_bs == 0: num_full_seqs = max_bs else: num_full_seqs = max_bs - 1 max_seq_len = num_ctx_tokens // num_full_seqs num_left_over_tokens = num_ctx_tokens - max_seq_len * num_full_seqs num_ctx_requests = num_full_seqs + (1 if num_left_over_tokens > 0 else 0) if num_ctx_requests + num_gen_requests > self.batch_size: 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 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. # Under-counting these let warmup start an allocation that fails # midway and, before the partial-allocation cleanup existed, # permanently leaked most of the estimation-sized KV pool # (TRTLLM-14903). 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) 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) if num_left_over_tokens > 0: blocks_to_use += blocks_for_seq(num_left_over_tokens + extra_ctx_tokens) blocks_to_use += (num_gen_requests * self.max_beam_width * blocks_for_seq(1 + extra_gen_tokens)) if blocks_to_use > available_blocks and isinstance( kv_cache_manager, KVCacheManager): return None if num_ctx_tokens > 0: ctx_token_nums = [max_seq_len] * num_full_seqs if num_left_over_tokens > 0: ctx_token_nums.append(num_left_over_tokens) ctx_requests = kv_cache_manager.add_dummy_requests( list(range(num_ctx_requests)), token_nums=ctx_token_nums, is_gen=False, max_num_draft_tokens=self.max_total_draft_tokens, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, draft_kv_cache_manager=draft_kv_cache_manager) if ctx_requests is None: return None if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests( request_ids=list(range(num_ctx_requests))) if num_gen_requests > 0: gen_requests = kv_cache_manager.add_dummy_requests( list( range(num_ctx_requests, num_ctx_requests + num_gen_requests)), token_nums=[1] * num_gen_requests, is_gen=True, max_num_draft_tokens=self.max_total_draft_tokens, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, draft_kv_cache_manager=draft_kv_cache_manager) if gen_requests is None: for r in ctx_requests: kv_cache_manager.free_resources(r) if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(r) return None if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests(request_ids=list( range(num_ctx_requests, num_ctx_requests + num_gen_requests))) result = ScheduledRequests() result.reset_context_requests(ctx_requests) result.generation_requests = gen_requests update_draft_len(self, result, draft_len=get_static_draft_len(self)) return result def _create_cuda_graph_warmup_request( self, resource_manager: ResourceManager, batch_size: int, draft_len: int, max_seq_len: int = None, mixed_context_encoder_output_lens: Optional[Sequence[int]] = None, mixed_context_query_len: int = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, force_non_greedy: bool = False, ) -> Optional[ScheduledRequests]: """Creates a dummy ScheduledRequests tailored for CUDA graph capture.""" capture_sampling_params = (NON_GREEDY_CAPTURE_SAMPLING_PARAMS if force_non_greedy else None) kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) available_blocks = kv_cache_manager.get_num_free_blocks( ) // self.max_beam_width if available_blocks < batch_size: return None result = ScheduledRequests() 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 is_enc_dec = self._is_encoder_decoder_model() max_encoder_output_len = ( self._get_max_encoder_output_len(resource_manager) if is_enc_dec else None) num_mixed_contexts = len(mixed_context_encoder_output_lens or ()) if is_enc_dec else 0 if num_mixed_contexts >= batch_size: return None # Add (batch_size - 1) dummy requests with the minimal sequence # length. Mixed capture must create its context rows as real context # requests; converting generation dummies afterward leaves their # native prompt/context bookkeeping at one token. if mixed_context_encoder_output_lens: context_request_ids = list(range(num_mixed_contexts)) context_requests = kv_cache_manager.add_dummy_requests( context_request_ids, token_nums=[mixed_context_query_len] * num_mixed_contexts, is_gen=False, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=list(mixed_context_encoder_output_lens), draft_kv_cache_manager=draft_kv_cache_manager, capture_sampling_params=capture_sampling_params) if context_requests is None: return None generation_request_ids = list( range(num_mixed_contexts, batch_size - 1)) generation_requests = [] if generation_request_ids: generation_requests = kv_cache_manager.add_dummy_requests( generation_request_ids, token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * len(generation_request_ids), is_gen=True, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=[max_encoder_output_len] * len(generation_request_ids), draft_kv_cache_manager=draft_kv_cache_manager, capture_sampling_params=capture_sampling_params) if generation_requests is None: for request in context_requests: kv_cache_manager.free_resources(request) if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(request) return None requests = context_requests + generation_requests else: token_nums = ([ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * (batch_size - 1)) if is_enc_dec else None encoder_output_lens = ([max_encoder_output_len] * (batch_size - 1)) if is_enc_dec else None requests = kv_cache_manager.add_dummy_requests( list(range(batch_size - 1)), token_nums=token_nums, is_gen=True, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=encoder_output_lens, draft_kv_cache_manager=draft_kv_cache_manager, capture_sampling_params=capture_sampling_params) if requests is None: return None def free_warmup_requests() -> None: for r in requests: kv_cache_manager.free_resources(r) if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(r) # Add one dummy request with the maximum possible sequence length. max_seq_len = min( self.max_seq_len if max_seq_len is None else max_seq_len, kv_cache_manager.max_seq_len) # Use max_draft_loop_tokens for capacity estimation to account # for the actual KV reservation per request. _kv_draft = self.max_draft_loop_tokens available_tokens = kv_cache_manager.get_num_available_tokens( token_num_upper_bound=max_seq_len, batch_size=batch_size, max_num_draft_tokens=_kv_draft) # Also consider draft KV cache capacity when it exists if draft_kv_cache_manager is not None: draft_available_tokens = draft_kv_cache_manager.get_num_available_tokens( batch_size=batch_size, token_num_upper_bound=max_seq_len, max_num_draft_tokens=_kv_draft) available_tokens = min(available_tokens, draft_available_tokens) token_num = max( ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1, min( available_tokens, max_seq_len - 1 - get_num_extra_kv_tokens(self.spec_config) - _kv_draft)) model_config = self.model.model_config.pretrained_config max_position_embeddings = getattr(model_config, 'max_position_embeddings', None) if is_enc_dec: # For enc-dec models the engine max_seq_len covers the encoder # sequence, which may exceed the decoder's position table (e.g. # Whisper: 1500 encoder positions vs max_target_positions=448). decoder_position_limit = getattr(model_config, 'max_target_positions', None) if decoder_position_limit is not None: max_position_embeddings = ( decoder_position_limit if max_position_embeddings is None else min(max_position_embeddings, decoder_position_limit)) if max_position_embeddings is not None: token_num = min(token_num, max_position_embeddings - _kv_draft) token_num = int( token_num) # Ensure int for range() in add_dummy_requests max_seq_len_request = kv_cache_manager.add_dummy_requests( request_ids=[batch_size - 1], token_nums=[token_num], is_gen=True, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=[max_encoder_output_len] if is_enc_dec else None, draft_kv_cache_manager=draft_kv_cache_manager, capture_sampling_params=capture_sampling_params) if max_seq_len_request is None: free_warmup_requests() return None else: max_seq_len_request = max_seq_len_request[0] if mixed_context_encoder_output_lens: requests.append(max_seq_len_request) for request in requests[:num_mixed_contexts]: request.state = LlmRequestState.CONTEXT_INIT request.context_current_position = 0 request.context_chunk_size = mixed_context_query_len request.cached_tokens = 0 request.py_batch_idx = None result.context_requests_last_chunk = requests[:num_mixed_contexts] result.generation_requests = requests[num_mixed_contexts:] else: # Insert the longest request first to simulate padding for the CUDA # graph. requests.insert(0, max_seq_len_request) result.generation_requests = requests if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests( request_ids=list(range(batch_size))) if self._is_encoder_decoder_model(): if not self._add_cross_dummy_requests(result.all_requests(), resource_manager): return None update_draft_len(self, result, draft_len=draft_len) return result def _get_max_encoder_output_len(self, resource_manager: ResourceManager) -> int: cross_kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.CROSS_KV_CACHE_MANAGER) max_encoder_output_len = int(self.max_seq_len) if cross_kv_cache_manager is not None: max_encoder_output_len = min( max_encoder_output_len, int( getattr(cross_kv_cache_manager, "max_seq_len", max_encoder_output_len))) return max(1, max_encoder_output_len) def _add_cross_dummy_requests(self, requests: List[LlmRequest], resource_manager: ResourceManager) -> bool: if not requests: return True cross_kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.CROSS_KV_CACHE_MANAGER) if cross_kv_cache_manager is None: raise RuntimeError("Encoder-decoder CUDA graph warmup requires " "ResourceManagerType.CROSS_KV_CACHE_MANAGER.") max_encoder_output_len = self._get_max_encoder_output_len( resource_manager) for request in requests: request.py_encoder_output = None request.py_skip_cross_kv_projection = True encoder_output_lens = [max_encoder_output_len] * len(requests) cross_dummy_requests = cross_kv_cache_manager.add_dummy_requests( request_ids=[request.py_request_id for request in requests], token_nums=encoder_output_lens, is_gen=True, max_beam_width=1, encoder_output_lens=encoder_output_lens) if cross_dummy_requests is not None: return True kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) for request in requests: kv_cache_manager.free_resources(request) if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(request) if spec_resource_manager is not None: spec_resource_manager.free_resources(request) return False def _populate_cross_kv_cache(self, inputs: Dict[str, Any]) -> None: encoder_hidden_states = inputs.get("encoder_hidden_states") cross_attn_metadata = inputs.get("cross_attn_metadata") if encoder_hidden_states is None or cross_attn_metadata is None: return decoder = getattr(get_top_level_model(self.model), "decoder", None) layers = getattr(decoder, "layers", None) if layers is None: raise RuntimeError("Encoder-decoder CUDA graph warmup requires a " "decoder with cross-attention layers.") attn_metadata = inputs["attn_metadata"] hidden_states = torch.ones( (attn_metadata.num_tokens, self._get_enc_dec_hidden_size()), device=encoder_hidden_states.device, dtype=encoder_hidden_states.dtype) for layer in layers: cross_attn = getattr(layer, "cross_attn", None) if cross_attn is None: raise RuntimeError( "Encoder-decoder CUDA graph warmup requires every decoder " "layer to expose a cross_attn module.") cross_attn(hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, attn_metadata=attn_metadata, cross_attn_metadata=cross_attn_metadata, skip_cross_kv_projection=False) def _get_enc_dec_hidden_size(self) -> int: config = self.model.model_config.pretrained_config hidden_size = getattr(config, "hidden_size", None) if hidden_size is None: hidden_size = getattr(config, "d_model", None) if hidden_size is None: raise RuntimeError( "Encoder-decoder CUDA graph warmup could not infer encoder " "hidden size from the model config.") return int(hidden_size) def _update_draft_inference_state_for_warmup( self, batch: ScheduledRequests, is_first_draft: bool, resource_manager: ResourceManager): """Updates request states for specific draft model warmups like Eagle3.""" spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) if self.is_draft_model and isinstance(spec_resource_manager, Eagle3ResourceManager): spec_resource_manager.is_first_draft = is_first_draft if is_first_draft: for req in batch.generation_requests: req.py_is_first_draft = True req.py_draft_tokens = [] def _set_up_attn_metadata( self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], draft_kv_cache_manager: Optional[Union[KVCacheManager, KVCacheManagerV2]] = None): if self.attn_metadata is not None: # This assertion can be relaxed if needed: just create a new metadata # object if it changes. assert self.attn_metadata.kv_cache_manager is kv_cache_manager return self.attn_metadata config = self.model.model_config.pretrained_config self.attn_metadata = build_attention_metadata( self.model.model_config, max_batch_size=self.batch_size, max_num_tokens=self.max_num_tokens, max_beam_width=self.max_beam_width, attention_backend=self.attn_backend, attention_runtime_features=self.attn_runtime_features, mapping=self.mapping, cache_indirection=self.cache_indirection_attention if self.attn_backend.Metadata is TrtllmAttentionMetadata else None, kv_cache_manager=kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, ) if isinstance(kv_cache_manager, BaseMambaCacheManager): self.attn_metadata.mamba_chunk_size = getattr( config, 'chunk_size', self.attn_metadata.mamba_chunk_size) self.attn_metadata.mamba_metadata_cls = resolve_mamba_metadata_cls( self.model) return self.attn_metadata @property def is_multimodal(self) -> bool: """True iff this engine drives a multimodal model.""" return is_multimodal(self.model, self.input_processor) def _validate_breakable_cuda_graph_compatibility(self) -> None: if self.llm_args.prefill_cuda_graph_backend != PrefillCudaGraphBackend.BREAKABLE: return if isinstance(self.model, DecoderModelForCausalLM): return decoder_model = getattr(self.model, "llm", None) if (self.llm_args.disable_mm_encoder and isinstance(decoder_model, DecoderModelForCausalLM) and getattr(self.model, "mm_encoder", None) is None): return if (isinstance(self.model, MultimodalModelMixin) or isinstance( self.input_processor, BaseMultimodalInputProcessor)): raise ValueError( "breakable prefill CUDA graph does not support multimodal models" ) def forward_multimodal_encoder_items( self, requests: List[LlmRequest], scheduled_items: Dict[int, List[int]], ) -> None: """Forward selected MM encoder items and commit request-local outputs.""" if not scheduled_items: return if self._mm_item_scheduler is None: raise TypeError( "Item-level MM scheduling requires MultimodalModelMixin") self._mm_item_scheduler.forward_items(requests, scheduled_items) def _set_up_spec_metadata( self, spec_resource_manager: Optional[BaseResourceManager]): spec_config = self.spec_config if self.enable_spec_decode else None if self.spec_metadata is not None: return self.spec_metadata self.spec_metadata = get_spec_metadata( spec_config, self.model.config, self.batch_size, max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, is_draft_model=self.is_draft_model, max_seq_len=self.max_seq_len, num_seq_slots=self.max_num_seq_slots) return self.spec_metadata def cleanup(self) -> None: """Release resources owned by this model engine. Tears down, in order: 1. The optional ``ModelLoader`` (which in turn releases any GMS client; see :meth:`ModelLoader.cleanup`). 2. CUDA Graph captures (via :meth:`_release_cuda_graphs`). 3. The runner, MM item scheduler, and model module reference, which hold references to the model. 4. Input processors. Idempotency: Subsequent calls are no-ops (guarded by ``_cleanup_done``). The flag is set only at the end, so a partial cleanup that raises mid-way will be retried on the next call. Called from: :meth:`__del__`, and only from there. ``PyExecutor.shutdown`` deliberately does *not* call this: it is also invoked mid-init by ``configure_kv_cache_capacity``, which reads ``model`` right afterwards, so clearing ``model`` here would break it. That path calls :meth:`_release_cuda_graphs` and then drops its reference instead. """ if self._cleanup_done: return # Cleanup is not truly atomic: released CUDA/GMS resources cannot be # rolled back. Keep each handle live until its own release succeeds, # so a failed cleanup can be retried without double-freeing resources # that were already released. model_loader = self.model_loader if model_loader is not None: model_loader.cleanup() self.model_loader = None # Release runner-owned graphs before dropping the runner. Keep the # handle available if graph release fails and cleanup is retried. self._release_cuda_graphs() # The runner and scheduler keep their own references to the model, so # clearing the engine's attribute alone would leave the weights # reachable past `release_gc()` below. self._runner = None self._mm_item_scheduler = None self.model = None self.input_processor = None # Release model weights. release_gc() self._cleanup_done = True def __del__(self) -> None: """Best-effort cleanup during garbage collection. Delegates to :meth:`cleanup`. Catches ``RuntimeError`` (which a release step such as :meth:`_release_cuda_graphs` or ``ModelLoader.cleanup`` may raise) and ``AttributeError`` (typical on partially-initialized engines torn down during interpreter shutdown when module references have already been cleared); both are logged and swallowed because destructors cannot reliably surface exceptions. This is the only production caller of :meth:`cleanup` -- see the note there on why ``PyExecutor.shutdown`` must not call it. """ try: self.cleanup() except (RuntimeError, AttributeError) as e: logger.warning( "PyTorchModelEngine cleanup failed during destruction: %s", e) def _init_max_seq_len(self): # Allow user to override the inferred max_seq_len with a warning. allow_long_max_model_len = os.getenv( "TLLM_ALLOW_LONG_MAX_MODEL_LEN", "0").lower() in ["1", "true", "yes", "y"] # Vision encoders may expose their own sequence-length inference. if hasattr(self.model, 'infer_max_seq_len'): inferred_max_seq_len = self.model.infer_max_seq_len() else: inferred_max_seq_len = self._infer_max_seq_len_from_config() if self.max_seq_len is None: logger.info( f"max_seq_len is not specified, using inferred value {inferred_max_seq_len}" ) self.max_seq_len = inferred_max_seq_len elif inferred_max_seq_len < self.max_seq_len: if allow_long_max_model_len: logger.warning( f"User specified max_seq_len is larger than the config in the model config file " f"({inferred_max_seq_len}). Setting max_seq_len to user's specified value {self.max_seq_len}. " ) else: # NOTE: py_executor_creator makes sure that the executor uses this # smaller value as its max_seq_len too. logger.warning( f"Specified {self.max_seq_len=} is larger than what the model can support " f"({inferred_max_seq_len}). Setting max_seq_len to {inferred_max_seq_len}. " ) self.max_seq_len = inferred_max_seq_len def _infer_max_seq_len_from_config(self) -> int: if hasattr(self.model, 'model_config') and self.model.model_config: model_config = self.model.model_config.pretrained_config rope_scaling = getattr(model_config, 'rope_scaling', None) rope_factor = 1 if rope_scaling is not None: rope_type = rope_scaling.get('type', rope_scaling.get('rope_type')) if rope_type not in ("su", "longrope", "llama3", "yarn"): rope_factor = rope_scaling.get('factor', 1.0) # Step 1: Find the upper bound of max_seq_len inferred_max_seq_len = 2048 max_position_embeddings = getattr(model_config, 'max_position_embeddings', None) if max_position_embeddings is None and hasattr( model_config, 'text_config'): max_position_embeddings = getattr(model_config.text_config, 'max_position_embeddings', None) if max_position_embeddings is not None: inferred_max_seq_len = max_position_embeddings # Step 2: Scale max_seq_len with rotary scaling if rope_factor != 1: inferred_max_seq_len = int( math.ceil(inferred_max_seq_len * rope_factor)) logger.warning( f'max_seq_len is scaled to {inferred_max_seq_len} by rope scaling {rope_factor}' ) return inferred_max_seq_len default_max_seq_len = 8192 logger.warning( f"Could not infer max_seq_len from model config, using default value: {default_max_seq_len}" ) return default_max_seq_len def _init_max_num_tokens(self): # Modified from tensorrt_llm/_bootstrap.py check_max_num_tokens if self.max_num_tokens is None: self.max_num_tokens = self.max_seq_len * self.batch_size if self.max_num_tokens > self.max_seq_len * self.batch_size: logger.warning( f"max_num_tokens ({self.max_num_tokens}) shouldn't be greater than " f"max_seq_len * max_batch_size ({self.max_seq_len * self.batch_size}), " f"specifying to max_seq_len * max_batch_size ({self.max_seq_len * self.batch_size})." ) self.max_num_tokens = self.max_seq_len * self.batch_size def _init_model_capacity(self): self._init_max_seq_len() self._init_max_num_tokens() def _release_cuda_graphs(self): if self._runner is not None: self._runner.cleanup() if self._torch_compile_backend is not None: self._torch_compile_backend.clear_piecewise_cuda_graphs() if hasattr(self, 'cuda_graph_runner') and self.cuda_graph_runner is not None: self.cuda_graph_runner.clear() if (hasattr(self, 'breakable_cuda_graph_runner') and self.breakable_cuda_graph_runner is not None): self.breakable_cuda_graph_runner.clear() def get_max_num_sequences(self) -> int: """ Return the maximum number of sequences that the model supports. PyExecutor needs this to compute max_num_active_requests """ num_batches = self.mapping.pp_size return num_batches * self.batch_size def _should_use_full_generation_page_table( self, spec_config: Optional[DecodingBaseConfig], attn_metadata: AttentionMetadata) -> bool: """Return whether overlap decode needs every reserved generation page.""" # FlashInfer metadata owns the optional device-side KV-length correction used with this # wider page table. return (self.enable_spec_decode and not self._disable_overlap_scheduler and getattr(spec_config, '_use_shared_kv_cache', False) and hasattr(attn_metadata, 'apply_spec_decode_kv_lens_offsets')) def _preprocess_inputs(self, inputs: Dict[str, Any]): """ Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get('attn_metadata') # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. if attn_meta is not None: attn_meta.on_update_kv_lens() if self.enable_spec_decode and not self._disable_overlap_scheduler: # When enabling overlap scheduler, the kv cache for draft tokens will # be prepared in advance by using the max_total_draft_tokens. But we need to use # new_tokens_lens_device to get the real past kv lengths and the # correct position ids. And to avoid blocking the async data transfer, # we need to preprocess the inputs in forward to update the position_ids and # kv cache length. if inputs['attn_metadata'].kv_cache_manager is not None: num_seqs = inputs['attn_metadata'].num_seqs num_ctx_requests = inputs['attn_metadata'].num_contexts num_gen_requests = inputs['attn_metadata'].num_generations num_ctx_tokens = inputs['attn_metadata'].num_ctx_tokens num_chunked_ctx_requests = inputs[ 'attn_metadata'].num_chunked_ctx_requests previous_batch_tokens = inputs['input_ids'].shape[ 0] - num_ctx_tokens if inputs['position_ids'].ndim == 3: # mrope: [3, 1, N] inputs['position_ids'][:, :, num_ctx_tokens:] += ( self. previous_pos_id_offsets_cuda[:previous_batch_tokens]) else: inputs['position_ids'][0, num_ctx_tokens:] += ( self. previous_pos_id_offsets_cuda[:previous_batch_tokens]) if hasattr(inputs['attn_metadata'], 'kv_lens_cuda'): if num_ctx_requests >= num_chunked_ctx_requests and num_chunked_ctx_requests > 0: # The generation requests with draft_tokens are treated as chunked context requests when extend_ctx returns True. inputs['attn_metadata'].kv_lens_cuda[ num_ctx_requests - num_chunked_ctx_requests:num_ctx_requests] += ( self. previous_kv_lens_offsets_cuda[: num_chunked_ctx_requests] ) else: inputs['attn_metadata'].kv_lens_cuda[ num_ctx_requests:num_seqs] += ( self. previous_kv_lens_offsets_cuda[:num_gen_requests] ) inputs['attn_metadata'].on_update_kv_lens() # TRTLLM uses `kv_lens_cuda` above; FlashInfer exposes this backend-specific # correction without coupling the engine to its metadata type. elif hasattr(inputs['attn_metadata'], 'apply_spec_decode_kv_lens_offsets'): inputs['attn_metadata'].apply_spec_decode_kv_lens_offsets( self.previous_kv_lens_offsets_cuda, num_gen_requests, self.get_runtime_tokens_per_gen_step( self.runtime_draft_len), num_chunked_contexts=num_chunked_ctx_requests, ) if self.enable_spec_decode and self.mapping.has_cp_helix(): # Helix verify groups: the per-token device buffers (write slots, # attention bounds, rank-local kv lens) must be derived on EVERY # spec step, overlap or not -- the append/mask kernels consume # them whenever _helix_spec_tokens_valid is armed. Under overlap # the host packed provisional positions from a stale base, so # first apply the same accepted-count correction position_ids # got above; without overlap the host values are already exact. md = inputs.get('attn_metadata') if (md is not None and md.kv_cache_manager is not None and getattr(md, '_helix_spec_tokens_valid', False)): helix_gen_tokens = (inputs['input_ids'].shape[0] - md.num_ctx_tokens) if not self._disable_overlap_scheduler: # The kv_lens override in the recompute supersedes the # generic previous_kv_lens_offsets adjustment above, # which is not ownership-aware. md.helix_position_offsets[:helix_gen_tokens] += ( self.previous_pos_id_offsets_cuda[:helix_gen_tokens]) md.recompute_helix_spec_buffers( helix_gen_tokens, self.get_runtime_tokens_per_gen_step( self.runtime_draft_len)) md.on_update_kv_lens() if self.guided_decoder is not None: self.guided_decoder.token_event.record() return inputs def _postprocess_inputs(self, inputs: Dict[str, Any]): """ Postprocess to make sure model forward doesn't change the inputs. It is only used in cuda graph capture, because other cases will prepare new inputs before the model forward. """ if self.enable_spec_decode and not self._disable_overlap_scheduler: if inputs['attn_metadata'].kv_cache_manager is not None: num_seqs = inputs['attn_metadata'].num_seqs num_ctx_requests = inputs['attn_metadata'].num_contexts num_gen_requests = inputs['attn_metadata'].num_generations num_ctx_tokens = inputs['attn_metadata'].num_ctx_tokens num_chunked_ctx_requests = inputs[ 'attn_metadata'].num_chunked_ctx_requests previous_batch_tokens = inputs['input_ids'].shape[ 0] - num_ctx_tokens if inputs['position_ids'].ndim == 3: # mrope: [3, 1, N] inputs['position_ids'][:, :, num_ctx_tokens:] -= ( self. previous_pos_id_offsets_cuda[:previous_batch_tokens]) else: inputs['position_ids'][0, num_ctx_tokens:] -= ( self. previous_pos_id_offsets_cuda[:previous_batch_tokens]) # Only TrtllmAttentionMetadata has kv_lens_cuda. if isinstance(inputs['attn_metadata'], TrtllmAttentionMetadata): if num_ctx_requests >= num_chunked_ctx_requests and num_chunked_ctx_requests > 0: inputs['attn_metadata'].kv_lens_cuda[ num_ctx_requests - num_chunked_ctx_requests:num_ctx_requests] -= ( self. previous_kv_lens_offsets_cuda[: num_chunked_ctx_requests] ) else: inputs['attn_metadata'].kv_lens_cuda[ num_ctx_requests:num_seqs] -= ( self. previous_kv_lens_offsets_cuda[:num_gen_requests] ) # Restore the FlashInfer-specific logical KV lengths through the same optional hook # used by `_preprocess_inputs`. elif hasattr(inputs['attn_metadata'], 'apply_spec_decode_kv_lens_offsets'): inputs['attn_metadata'].apply_spec_decode_kv_lens_offsets( self.previous_kv_lens_offsets_cuda, num_gen_requests, self.get_runtime_tokens_per_gen_step( self.runtime_draft_len), num_chunked_contexts=num_chunked_ctx_requests, restore=True, ) if (self.mapping.has_cp_helix() and getattr(inputs['attn_metadata'], '_helix_spec_tokens_valid', False)): # Mirror of the helix position correction in # _preprocess_inputs (capture symmetry, like position_ids # above). The recompute's OVERWRITES (slots/bounds/ # kv_lens) need no reversal: every consumer buffer is # rewritten from host state at the next step's prepare. inputs[ 'attn_metadata'].helix_position_offsets[:previous_batch_tokens] -= ( self. previous_pos_id_offsets_cuda[:previous_batch_tokens] ) def _get_all_rank_num_tokens_and_spec_counts( self, attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata ) -> Tuple[Optional[List[int]], Optional[List[List[int]]]]: """Exchange the attention and speculative per-rank counts in a single collective instead of one collective each. The spec token count is derived via ``SpecMetadata.dp_num_tokens`` because this collective runs *before* ``spec_metadata.prepare()`` rewrites ``num_tokens`` into the same shape (the attention count it shares the collective with is consumed earlier, by padding selection). Both read the same fields, so callers must have set ``spec_metadata.num_tokens``, ``num_generations`` and ``seq_lens`` for this batch before calling this. ``num_generations`` counts every generation request, CUDA-graph dummies included. """ if not self.enable_attention_dp: return None, None spec_counts = (spec_metadata.dp_num_tokens(), len(spec_metadata.seq_lens), spec_metadata.num_generations) if self.mapping.cp_size > 1 and not self.mapping.has_cp_helix(): # attn counts span TP only while spec counts span TP*CP; keep the # two exchanges separate. gathered = self.dist.tp_cp_allgather_int64(list(spec_counts)) return (get_all_rank_num_tokens( attn_metadata, enable_attention_dp=self.enable_attention_dp, mapping=self.mapping, dist=self.dist), gathered.T.tolist()) num_tokens = attn_metadata.num_tokens if self.mapping.has_cp_helix(): num_tokens = math.ceil(num_tokens / self.mapping.cp_size) gathered = self.dist.tp_cp_allgather_int64([num_tokens, *spec_counts]) cols = gathered.T.tolist() return cols[0], cols[1:] def _sync_group_all_greedy_sample(self, spec_metadata) -> None: """All-gather the per-rank greedy flags and store the group AND. Why the sampling-path choice must be group-uniform under ADP + LM-head TP is documented on the anchor, ``SpecMetadata.group_all_greedy_sample``. Local contract: called once per iteration, right after ``update_is_all_greedy_sample`` and BEFORE the CUDA graph key is built. The gate is pure config (identical on every rank), so ranks also agree on whether the exchange happens; the gather spans the whole TP group, a superset of any LM-head-TP subgroup. A dedicated host all-gather rather than a piggyback on the ``all_rank_num_tokens`` exchange, which runs in ``_prepare_inputs`` -- after the graph key, too late for the key to see the synced value. """ # enable_lm_head_tp_in_adp implies enable_attention_dp (asserted in # Mapping.__init__), so ADP needs no separate check here. if not (self.mapping.enable_lm_head_tp_in_adp and spec_metadata.use_rejection_sampling): return local_flag = bool(spec_metadata.is_all_greedy_sample) all_flags = self.dist.tp_allgather_int64([local_flag])[:, 0] spec_metadata.group_all_greedy_sample = bool(all_flags.all()) # Also overwrite the live flag directly: this iteration's scan already # ran (update_is_all_greedy_sample just returned) and the CUDA graph # key reads the flag next -- the stored override only takes effect on # the NEXT rescan (populate), which is after key selection. spec_metadata.is_all_greedy_sample = spec_metadata.group_all_greedy_sample def _is_final_multimodal_context_decode_compatible( self, request: LlmRequest) -> bool: """Return whether the final prompt token uses the decode input path. KV reuse has already materialized every preceding prompt token. A multimodal final-context row therefore needs its prepared embedding only when the one remaining token is itself an MM placeholder. Text tokens can use the existing decode provider; MRoPE deltas are seeded into the per-sequence cache before graph lookup. An MRoPE request with real MM payload remains eager until its delta is available. """ final_prompt_token = request.get_tokens(0)[ request.context_current_position] _, mm_token_indices = prepare_multimodal_indices([final_prompt_token], model=self.model) if mm_token_indices.numel() != 0: return False multimodal_data = request.py_multimodal_data if not self.use_mrope or not _has_mm_payload_keys(multimodal_data): return True return CUDAGraphRunner._get_mrope_position_delta(request) is not None def _is_encoder_decoder_model(self) -> bool: return bool( getattr(getattr(self.model, "model_config", None), "is_encoder_decoder", False)) @functools.cached_property def _model_uses_ple_recurrent_state(self) -> bool: """Detect PLE on text-only and multimodal model wrappers. The answer is fixed once the model is loaded, and the CUDA-graph gate below consults it on every forward that has context requests. """ top_level_model = get_top_level_model(self.model) if getattr(top_level_model, "has_ple", False): return True llm = getattr(top_level_model, "llm", None) text_model = getattr(llm, "model", llm) return bool(getattr(text_model, "has_ple", False)) def _prepare_enc_dec_cross_attn_inputs( self, encoder_hidden_states: List[torch.Tensor], encoder_seq_lens: List[int], encoder_num_cached_tokens_per_seq: List[int], attn_metadata: AttentionMetadata, resource_manager: Optional[ResourceManager], encoder_kv_lens: Optional[torch.Tensor] = None, context_encoder_kv_tokens: int = 0, generation_encoder_kv_tokens: int = 0, max_encoder_kv_len: int = 0, ) -> Dict[str, Any]: if not encoder_seq_lens: return {} if len(encoder_seq_lens) != attn_metadata.num_seqs: raise RuntimeError( "Cross-attention encoder lengths must align with decoder " f"sequences: got {len(encoder_seq_lens)} encoder lengths for " f"{attn_metadata.num_seqs} decoder sequences.") if resource_manager is None: raise RuntimeError( "Encoder-decoder decoder forward requires a resource manager " "with a cross-KV cache manager.") cross_kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.CROSS_KV_CACHE_MANAGER) if cross_kv_cache_manager is None: raise RuntimeError("Encoder-decoder decoder forward requires " "ResourceManagerType.CROSS_KV_CACHE_MANAGER.") new_encoder_tokens = sum(encoder_seq_lens) if encoder_hidden_states: packed_encoder_hidden_states = ( encoder_hidden_states[0] if len(encoder_hidden_states) == 1 else torch.cat(encoder_hidden_states, dim=0)) if packed_encoder_hidden_states.shape[0] != new_encoder_tokens: raise RuntimeError( "Packed encoder hidden states do not match cross-attention " "metadata: got " f"{packed_encoder_hidden_states.shape[0]} rows for " f"{new_encoder_tokens} new encoder KV tokens.") skip_cross_kv_projection = False else: if new_encoder_tokens != 0: raise RuntimeError( "Cross-attention metadata asks to project encoder K/V, " "but no encoder hidden states were supplied.") packed_encoder_hidden_states = None skip_cross_kv_projection = True def prepare_cross_metadata( cross_attn_metadata: AttentionMetadata) -> None: if encoder_kv_lens is None: cross_attn_metadata.prepare() return assert isinstance(cross_attn_metadata, TrtllmAttentionMetadata) cross_attn_metadata.prepare_encoder_decoder_from_precomputed_lengths( prompt_lens=attn_metadata.prompt_lens, kv_lens=encoder_kv_lens, context_kv_tokens=context_encoder_kv_tokens, generation_kv_tokens=generation_encoder_kv_tokens, max_kv_len=max_encoder_kv_len) if attn_metadata.is_cuda_graph and attn_metadata.has_cross_sub_metadata: # Fast path for stable CUDA-graph generation steps: the encoder # KV lengths (kv_lens_cuda) and the frozen prompt lengths # (prompt_lens_cuda) are identical across all generation steps # for a fixed batch. Skip the expensive torch.tensor() allocations # and H2D copies inside prepare() when nothing has changed. is_stable_gen_step = ( new_encoder_tokens == 0 # pure generation, no new cross-KV and self._cross_attn_stable_cached_tokens == encoder_num_cached_tokens_per_seq and self._cross_attn_stable_request_ids == attn_metadata.request_ids # same batch and row order ) if is_stable_gen_step: cross_attn_metadata = attn_metadata.cross # Only refresh the decoder-side Python references that the # kernel reads; these are pointer-level updates with no alloc. cross_attn_metadata._seq_lens = attn_metadata.seq_lens cross_attn_metadata._seq_lens_cuda = attn_metadata.seq_lens_cuda cross_attn_metadata.prompt_lens = attn_metadata.prompt_lens cross_attn_metadata.request_ids = attn_metadata.request_ids cross_attn_metadata.num_contexts = attn_metadata.num_contexts else: cross_attn_metadata = attn_metadata.update_cross_metadata( encoder_seq_lens=encoder_seq_lens, cross_kv_cache_manager=cross_kv_cache_manager, encoder_num_cached_tokens_per_seq= encoder_num_cached_tokens_per_seq, ) prepare_cross_metadata(cross_attn_metadata) if new_encoder_tokens == 0: # Record this stable state for future fast-path use. self._cross_attn_stable_cached_tokens = list( encoder_num_cached_tokens_per_seq) self._cross_attn_stable_request_ids = list( attn_metadata.request_ids) else: # Batch changed (new encoder request); reset cache. self._cross_attn_stable_cached_tokens = None self._cross_attn_stable_request_ids = None else: cross_attn_metadata = attn_metadata.create_cross_metadata( cross_kv_cache_manager=cross_kv_cache_manager, encoder_seq_lens=encoder_seq_lens, encoder_num_cached_tokens_per_seq= encoder_num_cached_tokens_per_seq, ) if attn_metadata.is_cuda_graph: attn_metadata.cross = cross_attn_metadata if new_encoder_tokens == 0: self._cross_attn_stable_cached_tokens = list( encoder_num_cached_tokens_per_seq) self._cross_attn_stable_request_ids = list( attn_metadata.request_ids) else: self._cross_attn_stable_cached_tokens = None self._cross_attn_stable_request_ids = None else: self._cross_attn_stable_cached_tokens = None self._cross_attn_stable_request_ids = None prepare_cross_metadata(cross_attn_metadata) return { "encoder_hidden_states": packed_encoder_hidden_states, "cross_attn_metadata": cross_attn_metadata, "skip_cross_kv_projection": skip_cross_kv_projection, } def _can_use_encoder_decoder_input_fast_path( self, scheduled_requests: ScheduledRequests, new_tokens_device: Optional[torch.Tensor], next_draft_tokens_device: Optional[torch.Tensor]) -> bool: """Return whether the TRT-like persistent input path is sufficient.""" static_eligible = self._encoder_decoder_input_fast_path_static_eligible if static_eligible is None: static_eligible = ( hasattr(batch_manager_bindings, "prepare_encoder_decoder_inputs") and self._is_encoder_decoder_model() and not self.is_draft_model and self.max_beam_width == 1 and self.sparse_attention_config is None and not self.use_mrope and not self.enable_attention_dp and not self.mapping.has_cp_helix() and not self.is_multimodal and not self.attn_runtime_features.chunked_prefill and not self.attn_runtime_features.cache_reuse and not self.attn_runtime_features.has_speculative_draft_tokens) self._encoder_decoder_input_fast_path_static_eligible = \ static_eligible if (not static_eligible or self.enable_spec_decode or self.lora_model_config is not None or new_tokens_device is None or next_draft_tokens_device is not None or self.guided_decoder is not None): return False if scheduled_requests.batch_size == 0: return False for request in scheduled_requests.generation_requests: if request.py_batch_idx is None and not request.is_dummy: return False return True def _acquire_encoder_decoder_host_buffers(self) -> Dict[str, Any]: """Acquire pinned staging whose preceding asynchronous copies finished.""" pool = self._encoder_decoder_host_buffer_pool for buffers in pool: event = buffers['event'] if event is None or event.query(): return buffers buffers = { 'input_ids': torch.empty(self.max_num_tokens, dtype=torch.int, pin_memory=prefer_pinned()), 'position_ids': torch.empty(self.max_num_tokens, dtype=torch.int, pin_memory=prefer_pinned()), 'sequence_lengths': torch.empty(self.batch_size, dtype=torch.int, pin_memory=prefer_pinned()), 'prompt_lengths': torch.empty(self.batch_size, dtype=torch.int, pin_memory=prefer_pinned()), 'cached_token_lengths': torch.empty(self.batch_size, dtype=torch.int, pin_memory=prefer_pinned()), 'kv_lengths': torch.empty(self.batch_size, dtype=torch.int, pin_memory=prefer_pinned()), 'encoder_kv_lengths': torch.empty(self.batch_size, dtype=torch.int, pin_memory=prefer_pinned()), 'previous_batch_indices': torch.empty(self.batch_size, dtype=torch.int, pin_memory=prefer_pinned()), 'event': None, } pool.append(buffers) return buffers def _record_cached_kv_tokens_per_req( self, num_cached_tokens_per_seq: Sequence[int], request_groups: Sequence[tuple[Sequence[LlmRequest], int]] = (), ) -> None: """Log post-prepare (backend-adjusted) counts in packed order, retaining ADP dummies to match num_scheduled_requests at beam width one and separating CUDA graph padding. """ counts = [] if request_groups: offset = 0 for requests, rows_per_request in request_groups: for request in requests: end = offset + rows_per_request if not request.is_cuda_graph_dummy: counts.extend(num_cached_tokens_per_seq[offset:end]) offset = end assert offset == len(num_cached_tokens_per_seq) else: counts = list(num_cached_tokens_per_seq) self.iter_states['cached_kv_tokens_per_req'] = counts self.iter_states['cached_kv_tokens_cuda_graph_padding'] = ( sum(num_cached_tokens_per_seq) - sum(counts)) @nvtx_range("_prepare_encoder_decoder_inputs_fast") def _prepare_encoder_decoder_inputs_fast( self, scheduled_requests: ScheduledRequests, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], attn_metadata: AttentionMetadata, new_tokens_device: torch.Tensor, resource_manager: Optional[ResourceManager]): """Prepare a simple BART batch with native collation and reused buffers.""" buffers = self._acquire_encoder_decoder_host_buffers() position_id_offset = self._encoder_decoder_position_id_offset if position_id_offset is None: position_id_offset = get_position_id_offset(self.model) self._encoder_decoder_position_id_offset = position_id_offset (request_ids, encoder_seq_lens, encoder_cached_token_lengths, total_num_tokens, num_context_tokens, num_previous_batch_requests, cached_kv_tokens, context_kv_tokens, generation_kv_tokens, max_kv_len, context_encoder_kv_tokens, generation_encoder_kv_tokens, max_encoder_kv_len ) = batch_manager_bindings.prepare_encoder_decoder_inputs( scheduled_requests.context_requests, scheduled_requests.generation_requests, buffers['input_ids'], buffers['position_ids'], buffers['sequence_lengths'], buffers['prompt_lengths'], buffers['cached_token_lengths'], buffers['kv_lengths'], buffers['encoder_kv_lengths'], buffers['previous_batch_indices'], position_id_offset, ) num_sequences = scheduled_requests.batch_size num_context_requests = scheduled_requests.num_context_requests num_generation_requests = scheduled_requests.num_generation_requests generation_request_ids = request_ids[num_context_requests:] if num_context_tokens: self.input_ids_cuda[:num_context_tokens].copy_( buffers['input_ids'][:num_context_tokens], non_blocking=True) if num_previous_batch_requests: previous_slots = self.previous_batch_indices_cuda[: num_previous_batch_requests] staged_request_ids = generation_request_ids[: num_previous_batch_requests] # Sequence slots are stable for a request's lifetime, so the # device indices remain valid while this ordered batch does. if self._encoder_decoder_staged_request_ids != staged_request_ids: previous_slots.copy_(buffers['previous_batch_indices'] [:num_previous_batch_requests], non_blocking=True) self._encoder_decoder_staged_request_ids = staged_request_ids generation_begin = num_context_tokens generation_end = generation_begin + num_previous_batch_requests torch.index_select( new_tokens_device[0, :, 0], dim=0, index=previous_slots, out=self.input_ids_cuda[generation_begin:generation_end]) else: self._encoder_decoder_staged_request_ids = None dummy_begin = num_context_tokens + num_previous_batch_requests if dummy_begin < total_num_tokens: self.input_ids_cuda[dummy_begin:total_num_tokens].fill_(0) self.position_ids_cuda[:total_num_tokens].copy_( buffers['position_ids'][:total_num_tokens], non_blocking=True) final_position_ids = self.position_ids_cuda[: total_num_tokens].unsqueeze( 0) sequence_lengths = buffers['sequence_lengths'][:num_sequences] attn_metadata._seq_lens = sequence_lengths if (attn_metadata.is_cuda_graph and attn_metadata._seq_lens_cuda is not None): attn_metadata._seq_lens_cuda.copy_(sequence_lengths, non_blocking=True) else: attn_metadata._seq_lens_cuda = sequence_lengths.cuda( non_blocking=True) attn_metadata._num_contexts = scheduled_requests.num_context_requests attn_metadata._num_ctx_tokens = num_context_tokens attn_metadata._num_generations = num_generation_requests attn_metadata._num_tokens = total_num_tokens attn_metadata.beam_width = 1 attn_metadata.request_ids = request_ids attn_metadata.prompt_lens = buffers['prompt_lengths'][:num_sequences] attn_metadata.num_chunked_ctx_requests = 0 attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=buffers['cached_token_lengths'] [:num_sequences], num_extra_kv_tokens=0) attn_metadata.kv_cache_manager = kv_cache_manager assert isinstance(attn_metadata, TrtllmAttentionMetadata) attn_metadata.prepare_encoder_decoder_from_precomputed_lengths( prompt_lens=buffers['prompt_lengths'][:num_sequences], kv_lens=buffers['kv_lengths'][:num_sequences], context_kv_tokens=context_kv_tokens, generation_kv_tokens=generation_kv_tokens, max_kv_len=max_kv_len) encoder_hidden_states = [] for request in scheduled_requests.context_requests: encoder_output = request.py_encoder_output if encoder_output is None: raise RuntimeError( f"Decoder context request {request.py_request_id} has no " "encoder output.") encoder_hidden_states.append(encoder_output) request.py_batch_idx = request.py_seq_slot cross_attention_inputs = self._prepare_enc_dec_cross_attn_inputs( encoder_hidden_states, encoder_seq_lens, encoder_cached_token_lengths, attn_metadata, resource_manager, encoder_kv_lens=buffers['encoder_kv_lengths'][:num_sequences], context_encoder_kv_tokens=context_encoder_kv_tokens, generation_encoder_kv_tokens=generation_encoder_kv_tokens, max_encoder_kv_len=max_encoder_kv_len, ) attn_all_rank_num_tokens = get_all_rank_num_tokens( attn_metadata, enable_attention_dp=self.enable_attention_dp, mapping=self.mapping, dist=self.dist) (padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens) = get_padding_params( total_num_tokens, scheduled_requests.num_context_requests, attn_all_rank_num_tokens, dist=self.dist, enable_attention_dp=self.enable_attention_dp, prefill_cuda_graph_backend=self.prefill_cuda_graph_backend, prefill_cuda_graph_num_tokens=self._prefill_cuda_graph_num_tokens) set_per_request_prefill_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = (padded_num_tokens if padded_num_tokens != total_num_tokens else None) virtual_num_tokens = total_num_tokens if attn_metadata.padded_num_tokens is not None: self.input_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0) self.position_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0) virtual_num_tokens = padded_num_tokens final_position_ids = self.position_ids_cuda[: virtual_num_tokens].unsqueeze( 0) inputs = { 'attn_metadata': attn_metadata, 'input_ids': self.input_ids_cuda[:virtual_num_tokens], 'position_ids': final_position_ids, 'inputs_embeds': None, 'multimodal_params': [], 'resource_manager': resource_manager, } inputs.update(cross_attention_inputs) self.iter_states[ 'num_ctx_requests'] = scheduled_requests.num_context_requests self.iter_states['num_ctx_tokens'] = num_context_tokens self.iter_states['num_generation_tokens'] = num_generation_requests self.iter_states['cached_kv_tokens'] = cached_kv_tokens if self._log_cached_kv_tokens_per_req: self._record_cached_kv_tokens_per_req( buffers['cached_token_lengths'][:num_sequences].tolist(), ((scheduled_requests.context_requests, 1), (scheduled_requests.generation_requests, 1))) if not self.is_warmup: self.previous_request_ids = generation_request_ids event = torch.cuda.Event() event.record(torch.cuda.current_stream()) buffers['event'] = event return inputs, None def _can_use_steady_gen_fast_prepare( self, scheduled_requests: ScheduledRequests, new_tokens_device: Optional[torch.Tensor], next_draft_tokens_device: Optional[torch.Tensor], spec_metadata: Optional[SpecMetadata]) -> bool: """Check whether the cached steady-state generation prepare applies. The cache is only recorded by a full _prepare_tp_inputs pass whose batch consisted purely of non-dummy generation requests that all had a previous overlap-scheduler tensor (see the recording site), so the per-step check only needs to confirm the dynamic conditions: still a generation-only batch with the exact same requests in the same order. """ cache = self._steady_gen_cache if cache is None or self.is_warmup: return False if new_tokens_device is None or next_draft_tokens_device is not None \ or spec_metadata is not None: return False if scheduled_requests.num_context_requests > 0: return False generation_requests = scheduled_requests.generation_requests if len(generation_requests) != cache['num_requests']: return False return cache['request_ids'] == [ request.py_request_id for request in generation_requests ] @nvtx_range("_apply_steady_gen_fast_prepare") def _apply_steady_gen_fast_prepare( self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], attn_metadata: AttentionMetadata, new_tensors_device: SampleStateTensors, resource_manager: Optional[ResourceManager]): """Prepare inputs for an unchanged generation-only batch. Every request advanced by exactly one committed token since the last prepare, so instead of re-walking the batch in Python this advances the cached positions in place (device position buffer plus a pinned host counter), reuses the seq-slot buffer already on device, and refreshes only the per-step metadata. For mrope models (recorded only for batches with no actual mrope work) the (3,1,N) broadcast buffer the model reads is the one advanced. """ cache = self._steady_gen_cache num_requests = cache['num_requests'] # Positions and cached-token counts are the same values in this # regime; advance both by one. The device-side position buffer is # advanced in place: it still holds the previous step's positions # because only _prepare_tp_inputs writes it and the cache validity # invariant guarantees the previous pass wrote these same rows. This # avoids reusing a mutated pinned buffer as the source of an async # H2D whose previous-step copy may still be pending under the overlap # scheduler (the nvbug 6293536 hazard class; see # KVCacheManager._stage_block_offsets_for_copy). The pinned buffer is # host-side bookkeeping only. use_mrope = cache['use_mrope'] positions = self._steady_gen_positions_pinned[:num_requests] positions.add_(1) if use_mrope: # Text-only batch on an mrope model: the recording pass broadcast # the scalar positions onto all three axes of the (3,1,N) buffer, # which is what the model (and any captured CUDA graph) reads, so # advance it in place. position_ids_cuda is reseeded by the next # full pass. self.mrope_position_ids_cuda[:, :, :num_requests].add_(1) else: self.position_ids_cuda[:num_requests].add_(1) num_cached_tokens_per_seq = positions.tolist() # Gather this step's input tokens from the previous iteration's device # sample buffer; the seq-slot indices in previous_batch_indices_cuda # are unchanged since the last full pass. previous_slots = self.previous_batch_indices_cuda[:num_requests] torch.index_select( new_tensors_device.new_tokens[0, :, :self.max_beam_width], 0, previous_slots, out=self.input_ids_cuda[:num_requests * self.max_beam_width].view( num_requests, self.max_beam_width), ) if not attn_metadata.is_cuda_graph: attn_metadata.seq_lens = cache['seq_lens_ones'] attn_metadata.beam_width = 1 attn_metadata.request_ids = cache['request_ids'] attn_metadata.prompt_lens = cache['prompt_lens'] attn_metadata.num_contexts = 0 attn_metadata.num_chunked_ctx_requests = 0 attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq, num_extra_kv_tokens=get_num_extra_kv_tokens(None)) attn_metadata.kv_cache_manager = kv_cache_manager if hasattr(self.model.model_config.pretrained_config, 'chunk_size'): attn_metadata.mamba_chunk_size = \ self.model.model_config.pretrained_config.chunk_size with nvtx_range("steady_gen_metadata_prepare"): attn_metadata.prepare() attn_all_rank_num_tokens = get_all_rank_num_tokens( attn_metadata, enable_attention_dp=self.enable_attention_dp, mapping=self.mapping, dist=self.dist) padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = \ get_padding_params( num_requests, 0, attn_all_rank_num_tokens, dist=self.dist, enable_attention_dp=self.enable_attention_dp, prefill_cuda_graph_backend=self.prefill_cuda_graph_backend, prefill_cuda_graph_num_tokens=self._prefill_cuda_graph_num_tokens) set_per_request_prefill_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = ( padded_num_tokens if padded_num_tokens != num_requests else None) virtual_num_tokens = num_requests if attn_metadata.padded_num_tokens is not None: self.input_ids_cuda[num_requests:padded_num_tokens].fill_(0) # Zero-fill the padding tail of whichever position layout the # model consumes, matching the full pass. if use_mrope: self.mrope_position_ids_cuda[:, :, num_requests: padded_num_tokens].fill_(0) else: self.position_ids_cuda[num_requests:padded_num_tokens].fill_(0) virtual_num_tokens = padded_num_tokens self.iter_states['num_ctx_requests'] = 0 self.iter_states['num_ctx_tokens'] = 0 self.iter_states['num_generation_tokens'] = num_requests self.iter_states['cached_kv_tokens'] = sum(num_cached_tokens_per_seq) if self._log_cached_kv_tokens_per_req: self._record_cached_kv_tokens_per_req(num_cached_tokens_per_seq) if use_mrope: final_position_ids = \ self.mrope_position_ids_cuda[:, :, :virtual_num_tokens] else: final_position_ids = \ self.position_ids_cuda[:virtual_num_tokens].unsqueeze(0) inputs = { 'attn_metadata': attn_metadata, 'input_ids': self.input_ids_cuda[:virtual_num_tokens], 'position_ids': final_position_ids, 'inputs_embeds': None, 'multimodal_params': [], 'resource_manager': resource_manager, } return inputs, None def _prepare_tp_inputs( self, scheduled_requests: ScheduledRequests, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], attn_metadata: AttentionMetadata, spec_metadata: Optional[SpecMetadata] = None, new_tensors_device: Optional[SampleStateTensors] = None, cache_indirection_buffer: Optional[torch.Tensor] = None, num_accepted_tokens_device: Optional[torch.Tensor] = None, req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, resource_manager: Optional[ResourceManager] = None, maybe_graph: bool = False, promoted_context_request_ids: frozenset[int] = frozenset(), use_lora_graph: bool = False, ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]: """ Prepare inputs for Pytorch Model. """ new_tokens_device, new_tokens_lens_device, next_draft_tokens_device = None, None, None if new_tensors_device is not None: # speculative decoding cases: [batch, 1 + draft_len], others: [batch] new_tokens_device = new_tensors_device.new_tokens # When using overlap scheduler with speculative decoding, the target model's inputs would be SampleStateTensorsSpec. if isinstance(new_tensors_device, SampleStateTensorsSpec): assert self.enable_spec_decode and not self.is_draft_model new_tokens_lens_device = new_tensors_device.new_tokens_lens # [batch] next_draft_tokens_device = new_tensors_device.next_draft_tokens # [batch, draft_len] # Must be before the update of py_batch_idx if self.guided_decoder is not None: self.guided_decoder.add_batch( scheduled_requests, new_tokens=new_tokens_device, runtime_draft_len=self.runtime_draft_len) if (not promoted_context_request_ids and type(attn_metadata) is TrtllmAttentionMetadata and self._can_use_encoder_decoder_input_fast_path( scheduled_requests, new_tokens_device, next_draft_tokens_device)): return self._prepare_encoder_decoder_inputs_fast( scheduled_requests, kv_cache_manager, attn_metadata, new_tokens_device, resource_manager) self._encoder_decoder_staged_request_ids = None if (not promoted_context_request_ids and self._can_use_steady_gen_fast_prepare( scheduled_requests, new_tokens_device, next_draft_tokens_device, spec_metadata)): return self._apply_steady_gen_fast_prepare(kv_cache_manager, attn_metadata, new_tensors_device, resource_manager) # Any full pass invalidates the steady-state cache; it is re-recorded # at the end of this pass when the batch qualifies. self._steady_gen_cache = None # Hoist self.use_mrope to a function-scope local so the per-request / # per-context-request mrope branches use LOAD_FAST instead of LOAD_ATTR. _use_mrope = self.use_mrope # if new_tensors_device exist, input_ids will only contain new context tokens input_ids = [] # per sequence sequence_lengths = [] # per sequence prompt_lengths = [] # per sequence request_ids = [] # per request gather_ids = [] position_ids = [] # per sequence num_cached_tokens_per_seq = [] # per sequence draft_tokens = [] draft_lens = [] gen_request_seq_slots = [] # per generation request # One-model rejection: slots of gen requests that produced 0 real draft # tokens this step (marked in _handle_dynamic_draft_len); their stale # draft_probs rows are one-hot'd after spec_metadata.prepare(). padding_gen_slots = [] multimodal_params_list = [] mrope_position_ids = [ ] # (start_idx, end_idx, (3,1,L) mrope_pos_ids) per multimodal request mrope_delta_write_seq_slots = [] mrope_delta_read_seq_slots = [] # Whether any generation request in this batch carries real MRoPE # metadata; see the post-loop cleanup below. has_gen_mrope_delta = False # Extra model-side cache slot reserved for CUDA graph / warmup dummy # requests, whose outputs are discarded, and for generation requests # that carry no MRoPE metadata at all. The cache is zero-initialized and # the write path only ever targets real ``py_seq_slot``s, so this slot # permanently reads back a zero delta. mrope_dummy_seq_slot = get_mrope_dummy_seq_slot(self.max_num_tokens, self.mapping.pp_size) num_accepted_draft_tokens = [] # per request is_enc_dec = self._is_encoder_decoder_model() cross_encoder_hidden_states: List[torch.Tensor] = [] cross_encoder_seq_lens: List[int] = [ ] # new encoder K/V tokens per decoder sequence cross_encoder_cached_tokens_per_seq: List[int] = [] # Variables for updating the inputs of draft model # Base values for gather_ids computation first_draft_base_gather_ids = [] # seq_slots to index into num_accepted_tokens_device first_draft_seq_slots = [] # Indices in the num_accepted_draft_tokens list first_draft_request_indices = [] # (start_idx, end_idx, seq_slot) for context requests context_input_ids_positions = [] # (start_idx, end_idx, seq_slot) for first_draft requests first_draft_input_ids_positions = [] context_prompt_lookahead = None if (spec_metadata is not None and spec_metadata.context_prompt_lookahead_tokens is not None): context_prompt_lookahead = [] def append_cross_attention_state(request: LlmRequest, project_encoder_output: bool, repeat: int = 1) -> None: if not is_enc_dec: return encoder_output_len = int(request.encoder_output_len) if project_encoder_output: encoder_output = getattr(request, "py_encoder_output", None) if encoder_output is None: raise RuntimeError( "Decoder context request " f"{request.py_request_id} has no encoder output. " "The encoder iteration must populate " "req.py_encoder_output before the first decoder " "context step.") if encoder_output.shape[0] != encoder_output_len: raise RuntimeError( "Decoder context request " f"{request.py_request_id} encoder output length " f"({encoder_output.shape[0]}) does not match " f"encoder_output_len ({encoder_output_len}).") cross_encoder_hidden_states.append(encoder_output) cross_encoder_seq_lens.append(encoder_output_len) cross_encoder_cached_tokens_per_seq.append(0) return for _ in range(repeat): cross_encoder_seq_lens.append(0) cross_encoder_cached_tokens_per_seq.append(encoder_output_len) for request in scheduled_requests.context_requests: request_ids.append(request.py_request_id) draft_lens.append(0) begin_compute = request.context_current_position end_compute = begin_compute + request.context_chunk_size if context_prompt_lookahead is not None: context_prompt_lookahead.append( _get_context_prompt_lookahead_token(request, end_compute)) # Fetch only the current chunk. get_tokens(0) marshals the whole # O(seq_len) VecTokens into a Python list of boxed ints; chunked # prefill re-enters this loop for every chunk of the same prompt, so # that is O(L) per chunk = O(L^2/chunk) over the prefill. # get_tokens_range copies only [begin, end) -> O(chunk). prompt_tokens = request.get_tokens_range(0, begin_compute, end_compute) position_ids.extend( range(begin_compute, begin_compute + len(prompt_tokens))) # Start offset of this request's (current-chunk) tokens within the # flattened input_ids. Recorded on multimodal_params below so models # that rewrite token IDs in place write into the request's own span # rather than assuming a contiguous multimodal prefix. context_start_idx = len(input_ids) # Track position for updating the inputs of draft model if self.is_draft_model and num_accepted_tokens_device is not None: input_ids.extend(prompt_tokens) end_idx = len(input_ids) slot_idx = req_id_to_old_request[ request.py_request_id].py_seq_slot context_input_ids_positions.append( (context_start_idx, end_idx - 1, slot_idx)) # end_idx-1 is the last token position else: input_ids.extend(prompt_tokens) gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) num_accepted_draft_tokens.append(len(prompt_tokens) - 1) prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - request.py_num_compressed_tokens) request.cached_tokens = past_seen_token_num append_cross_attention_state( request, project_encoder_output=not request.py_skip_cross_kv_projection and (not getattr(request, "is_dummy", False) or getattr(request, "py_encoder_output", None) is not None)) # Embed mask is required only for partial iterations (chunked # prefill or KV-cache reuse); full-prefill degrades gracefully. check_mm_embed_cumsum_if_needed( request.py_multimodal_data, begin_compute=past_seen_token_num, end_compute=end_compute, prompt_len=request.get_num_tokens(0), ) mm_data = request.py_multimodal_data or {} cumsum = mm_data.get('multimodal_embed_mask_cumsum') py_multimodal_runtime = None if cumsum is not None: py_multimodal_runtime = MultimodalRuntimeData( embed_mask_cumsum=cumsum, past_seen_token_num=past_seen_token_num, chunk_end_pos=end_compute, ) multimodal_params = MultimodalParams( multimodal_input=_build_request_multimodal_input( request, self._mm_encoder_cache_enabled), multimodal_data=request.py_multimodal_data, multimodal_runtime=py_multimodal_runtime, mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx) # Transfer any cross-iter MM encoder prefetch event stamped on the request onto the # freshly-built MultimodalParams. The downstream consume site reads it from the wrapper, # not from the request. # NOTE: the prefetch producer always writes the cached embedding into # `py_multimodal_data` before stamping the event, so whenever the event is present, # `has_content()` below is `True` and the wrapper reaches the consume site that waits on # it. mm_encoder_event = request.py_mm_encoder_event if mm_encoder_event is not None: multimodal_params.encoder_event = mm_encoder_event request.py_mm_encoder_event = None if multimodal_params.has_content(): # TODO(TRTLLM-14726): Check the persistent MM encoder cache before H2D and avoid # transferring raw encoder inputs for full hits in both regular and # side-stream-prefetched paths. multimodal_params.to_device("multimodal_data", "cuda", pin_memory=prefer_pinned(), target_keywords=getattr( self.model, "multimodal_data_device_paths", None)) if _use_mrope: # A request may carry multimodal content but no MRoPE # metadata (a text-only prompt whose input processor skips # ``mrope_config``, or a model that does not consume it). # Its per-axis positions are just the scalar positions, # which the (3,1,N) seeding further below already # broadcasts, so leave that span alone. mrope_config = multimodal_params.multimodal_data.get( 'mrope_config') or {} mrope_pos_ids = mrope_config.get('mrope_position_ids') if mrope_pos_ids is not None: ctx_mrope_position_ids = mrope_pos_ids[:, :, begin_compute: begin_compute + len(prompt_tokens )] # Record as (start_idx, end_idx, (3,1,L) mrope_pos_ids) mrope_position_ids.append( (len(position_ids) - len(prompt_tokens), len(position_ids), ctx_mrope_position_ids)) mrope_position_delta = mrope_config.get( 'mrope_position_deltas') if mrope_position_delta is not None: request.py_mrope_position_delta = mrope_position_delta if (mrope_position_delta is not None and request.py_seq_slot is not None): mrope_delta_write_seq_slots.append(request.py_seq_slot) request.py_mrope_delta_cache_slot = request.py_seq_slot #re-assign the multimodal_data to the request after to_device for generation requests request.py_multimodal_data = multimodal_params.multimodal_data multimodal_params_list.append(multimodal_params) # Re-register mrope tensors for context-only requests (EPD disaggregated serving). # This creates new IPC handles owned by the prefill worker, so the decode worker # can access them even after the encode worker's GC deallocates the original memory. # Without this, the decode worker would receive handles pointing to freed memory. if (request.is_context_only_request and _use_mrope and "mrope_config" in multimodal_params.multimodal_data): mrope_config = multimodal_params.multimodal_data[ "mrope_config"] _mrope_position_ids = mrope_config.get("mrope_position_ids") _mrope_position_deltas = mrope_config.get( "mrope_position_deltas") if _mrope_position_ids is not None and _mrope_position_deltas is not None: # Clone to allocate new memory owned by this (prefill) worker. request.py_result.set_mrope_position( _mrope_position_ids.clone(), _mrope_position_deltas.clone()) request.py_batch_idx = request.py_seq_slot num_ctx_requests = scheduled_requests.num_context_requests num_ctx_tokens = len(input_ids) if len(multimodal_params_list) > 0: # input_ids holds only context tokens here; extend/draft tokens are # appended below and are by construction text, so we reuse the # CPU-side text_token_indices and just extend it with the # post-context arange instead of recomputing via a bool mask + # torch.where over the full range. text_token_indices_ctx, mm_token_indices = \ prepare_multimodal_indices(input_ids, model=self.model) else: text_token_indices_ctx = None mm_token_indices = None # Requests with draft tokens are treated like extend requests. Dummy extend requests should be # at the end of extend_requests. extend_requests = [] extend_dummy_requests = [] generation_requests = [] first_draft_requests = [] # Collect generation request IDs during categorization to avoid # a separate iteration over scheduled_requests.generation_requests later. all_gen_request_ids = [] for request in scheduled_requests.generation_requests: is_promoted_context = (request.py_request_id in promoted_context_request_ids) if not is_promoted_context: all_gen_request_ids.append(request.py_request_id) # In speculative iterations, keep promoted rows ahead of existing # generation rows in the extend-request packing order. Although # their q_len is one, this category provides the "no previous # speculative tensor" branch needed to source their prompt token # without disturbing the overlap offsets of ordinary generation # siblings. Non-speculative promoted rows retain the established # ordinary generation path below. if is_promoted_context and self.enable_spec_decode: extend_requests.append(request) elif is_promoted_context: generation_requests.append(request) elif (get_draft_token_length(request) > 0 or next_draft_tokens_device is not None): if request.is_dummy: extend_dummy_requests.append(request) else: extend_requests.append(request) elif request.py_is_first_draft: first_draft_requests.append(request) else: generation_requests.append(request) extend_requests += extend_dummy_requests # Helix bookkeeping is needed by BOTH the extend (speculative verify # group) and the plain generation packing loops below, so initialize # it ahead of them. Positions are global; KV ownership follows the # round-robin ledger (page b -> rank b % cp); the host-side # provisional packing values come from the one shared definition in # _torch.utils.helix_local_len. helix_is_inactive_rank, helix_position_offsets = [], [] helix_owned_new_tokens = [] _has_cp_helix = self.mapping.has_cp_helix() if _has_cp_helix and kv_cache_manager is not None: _helix_phys = kv_cache_manager.tokens_per_block _helix_cp_size = self.mapping.cp_size _helix_cp_rank = self.mapping.cp_rank def _helix_local_len_host(global_len: int) -> int: return helix_local_len(global_len, _helix_phys, _helix_cp_size, _helix_cp_rank) def _helix_pack_extend(request, group: int) -> int: # A helix gen worker's token list is the rank-LOCAL # round-robin subset, so max_beam_num_tokens is not a global # base; rebuild it from the global prompt length plus the # rank-invariant generated count. Also repacks position_ids, # which the caller filled from the local base. generated_len = (request.max_beam_num_tokens - request.py_prompt_len) base = request.total_input_len_cp + generated_len - 1 helix_position_offsets.extend(range(base, base + group)) position_ids[-group:] = range(base, base + group) helix_is_inactive_rank.append(False) return base spec_config = self.spec_config if self.enable_spec_decode else None if not self._disable_overlap_scheduler and spec_config is not None: assert spec_config.spec_dec_mode.support_overlap_scheduler( ), f"{spec_config.decoding_type} does not support overlap scheduler" # For tree decoding, runtime_draft_len should match total tree # tokens (not tree depth). py_executor resets it every iteration. if spec_config is not None and not spec_config.is_linear_tree: self.runtime_draft_len = get_static_draft_len(self) # will contain previous batch indices of generation requests previous_batch_indices = [] previous_pos_indices = [] runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: is_promoted_context = (request.py_request_id in promoted_context_request_ids) if getattr(request, "py_needs_onehot_draft_probs", False): if request.py_seq_slot is not None: padding_gen_slots.append(request.py_seq_slot) request.py_needs_onehot_draft_probs = False # consume once request_ids.append(request.py_request_id) # the request has no previous tensor: # (1) next_draft_tokens_device is None, which means overlap scheduler is disabled; or # (2) a dummy request; or # (3) the first step in the generation server of disaggregated serving if (is_promoted_context or next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None): # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): if is_promoted_context: input_ids.append( request.get_tokens(0)[ request.context_current_position]) else: input_ids.append(request.get_last_tokens(0)) input_ids.extend(request.py_draft_tokens) draft_tokens.extend(request.py_draft_tokens) # get other ids and lengths num_draft_tokens = get_draft_token_length(request) past_seen_token_num = (request.context_current_position if is_promoted_context else request.max_beam_num_tokens - 1) draft_lens.append(num_draft_tokens) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: # We're treating the prompt lengths as context requests here, so # the the prompt lens should not include the cached tokens. prompt_lengths.append(1 + num_draft_tokens) else: prompt_lengths.append(request.py_prompt_len) sequence_lengths.append(1 + num_draft_tokens) num_accepted_draft_tokens.append(num_draft_tokens) gather_ids.extend( list( range(len(position_ids), len(position_ids) + 1 + num_draft_tokens))) position_ids.extend( list( range(past_seen_token_num, past_seen_token_num + 1 + num_draft_tokens))) num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens) request.cached_tokens = past_seen_token_num if _has_cp_helix: # Verify group [base, base+group) in GLOBAL positions. # On a helix gen worker the request's token list is the # rank-LOCAL round-robin subset, so max_beam_num_tokens # (= local_prompt + generated) must NOT be used as a # global base; reconstruct it from the global prompt # length plus the (rank-invariant) generated count. This # branch has no in-flight predecessor, so every value is # exact (no device correction needed). group = 1 + num_draft_tokens base = _helix_pack_extend(request, group) local_cached = _helix_local_len_host(base) helix_owned_new_tokens.append( _helix_local_len_host(base + group) - local_cached) num_cached_tokens_per_seq[-1] = ( local_cached - request.py_num_compressed_tokens) request.cached_tokens = local_cached # update batch index request.py_batch_idx = request.py_seq_slot else: # update batch index previous_batch_idx = request.py_batch_idx request.py_batch_idx = request.py_seq_slot sequence_lengths.append(runtime_tokens_per_gen_step) num_accepted_draft_tokens.append( request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 draft_lens.append(runtime_draft_token_buffer_width) gather_ids.extend( list( range(len(position_ids), len(position_ids) + runtime_tokens_per_gen_step))) position_ids.extend( list( range(past_seen_token_num, past_seen_token_num + runtime_tokens_per_gen_step))) # previous tensor previous_batch_indices.append(previous_batch_idx) previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) num_cached_tokens_per_seq.append( past_seen_token_num + runtime_tokens_per_gen_step - request.py_num_compressed_tokens) request.cached_tokens = (past_seen_token_num + runtime_tokens_per_gen_step) if _has_cp_helix: # In-flight predecessor: mirror the non-helix convention # above -- positions are packed from the stale base (the # overlap device correction adds the accepted count) and # KV numbers assume full acceptance (the device recompute # in recompute_helix_spec_buffers overrides them). The # base is reconstructed GLOBALLY (see the no-previous # branch: the token list is rank-local under helix). group = runtime_tokens_per_gen_step base = _helix_pack_extend(request, group) local_full = _helix_local_len_host(base + group) helix_owned_new_tokens.append(0) num_cached_tokens_per_seq[-1] = ( local_full - request.py_num_compressed_tokens) request.cached_tokens = local_full if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: prompt_lengths.append(runtime_tokens_per_gen_step) else: prompt_lengths.append(request.py_prompt_len) append_cross_attention_state(request, project_encoder_output=False) for request in first_draft_requests: request_ids.append(request.py_request_id) draft_lens.append(0) # Only the length and the last (original_max_draft_len+1) tokens are # needed here; get_num_tokens is O(1) and get_tokens_range copies only # the requested window, whereas get_tokens(0) marshals the whole # O(seq_len) VecTokens into a Python list. _num_tokens = request.get_num_tokens(0) begin_compute = _num_tokens - self.original_max_draft_len - 1 end_compute = begin_compute + self.original_max_draft_len + 1 prompt_tokens = request.get_tokens_range(0, begin_compute, end_compute) position_ids.extend( range(begin_compute, begin_compute + len(prompt_tokens))) # Track position for updating the inputs of draft model if self.is_draft_model and num_accepted_tokens_device is not None: start_idx = len(input_ids) input_ids.extend(prompt_tokens) end_idx = len(input_ids) # For first_draft, we need to replace the last original_max_draft_len+1 tokens slot_idx = req_id_to_old_request[ request.py_request_id].py_seq_slot first_draft_input_ids_positions.append( (start_idx, end_idx, slot_idx)) # Store info for GPU computation of gather_ids and num_accepted_draft_tokens base_gather_id = len( input_ids) - 1 - self.original_max_draft_len # Placeholder, will be corrected on GPU gather_ids.append(base_gather_id) first_draft_base_gather_ids.append(base_gather_id) first_draft_seq_slots.append(slot_idx) first_draft_request_indices.append( len(num_accepted_draft_tokens)) # Placeholder, will be corrected on GPU num_accepted_draft_tokens.append(0) else: input_ids.extend(prompt_tokens) gather_ids.append( len(input_ids) - 1 - (self.original_max_draft_len - request.py_num_accepted_draft_tokens)) num_accepted_draft_tokens.append( request.py_num_accepted_draft_tokens) sequence_lengths.append(1 + self.original_max_draft_len) prompt_lengths.append(request.py_prompt_len) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - request.py_num_compressed_tokens) append_cross_attention_state(request, project_encoder_output=False) # update batch index request.py_batch_idx = request.py_seq_slot # Cache invariant method result to avoid repeated calls per-request _n_gen = len(generation_requests) # One-shot batch-level flag — True iff any generation request actually # carries multimodal payload. Lets the strip_mm_data branch below # short-circuit on a LOAD_FAST rather than a per-request LOAD_ATTR # of py_multimodal_data for non-multimodal models (the gpt-oss-120b # GEN case). _has_any_multimodal_request = any(r.py_multimodal_data is not None for r in generation_requests) if _n_gen > 0: # The whole batch is laid out with request 0's beam width: every # generation request contributes exactly this many rows to # input_ids / position_ids / sequence_lengths and to the logits the # model returns. The sampler, in turn, locates a request's logits by # accumulating the *per-request* beam widths # (TorchSampler._select_generated_logits -> # calculate_request_offsets). Both agree only while every request in # the batch has the same beam width. # # Mixing widths would desynchronize the two: the sampler would read # a request's rows at the wrong offset, and `logits.view(batch, # beam_width_in, vocab)` succeeds for any shape whose element count # divides, so the result is silently wrong rather than an error. # Supporting mixed widths needs the forward path to emit a fixed # max_beam_width stride and the sampler offsets to match; until # then, fail loudly. beam_width = generation_requests[0].py_beam_width # Admission pins every request to max_beam_width, but a # variable-beam-width request narrows or widens per iteration, so # the widths can still diverge mid-batch. Compare the # *per-iteration* width: py_beam_width is fixed at admission and # would be identical across those requests. Dummy requests are # excluded -- they carry no user request and are built at their own # width (CUDA-graph padding at the engine width, attention-DP and # warmup dummies at width one), so they would otherwise trip this # on an ordinary padded batch. real_requests = [ req for req in generation_requests if not req.is_dummy ] iter_widths = { req.get_beam_width_by_iter() for req in real_requests } if len(iter_widths) > 1: # NB: this aborts the whole batch, not just the offending # requests -- ModelEngine has no per-request failure channel, # and by this point the batch is already scheduled. Scoping the # failure needs the scheduler to group by beam width in the # first place, so that no such batch is formed; TRTLLM-14792. raise ValueError( "Generation requests in one batch must all have the same " f"beam width; got {sorted(iter_widths)}. Mixed beam widths " "within a batch are not supported yet (TRTLLM-14792).") # Pre-extend constant-value lists to avoid per-request append # overhead (saves ~3 append calls per request). draft_lens.extend([0] * (_n_gen * beam_width)) sequence_lengths.extend([1] * (_n_gen * beam_width)) num_accepted_draft_tokens.extend([0] * (_n_gen * beam_width)) for request in generation_requests: request_ids.append(request.py_request_id) is_promoted_context = (request.py_request_id in promoted_context_request_ids) if is_promoted_context: input_ids.append( request.get_tokens(0)[request.context_current_position]) past_seen_token_num = request.context_current_position request_has_previous_tensor = False # The request has no previous tensor: # (1) new_tokens_device is None, which means overlap scheduler is disabled; or # (2) a dummy request; or # (3) the first step in the generation server of disaggregated serving. elif new_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # skip adding input_ids of CUDA graph dummy requests so that new_tokens_device # can be aligned to the correct positions. if not request.is_cuda_graph_dummy: for beam in range(beam_width): # Track position for GPU update (draft model only) if self.is_draft_model and num_accepted_tokens_device is not None: start_idx = len(input_ids) input_ids.append(request.get_last_tokens(beam)) end_idx = len(input_ids) slot_idx = req_id_to_old_request[ request.py_request_id].py_seq_slot first_draft_input_ids_positions.append( (start_idx, end_idx, slot_idx)) else: input_ids.append(request.get_last_tokens(beam)) past_seen_token_num = request.max_beam_num_tokens - 1 request_has_previous_tensor = False else: # the request has previous tensor # previous_batch_indices is per-request, not per-beam previous_batch_indices.append(request.py_batch_idx) past_seen_token_num = request.max_beam_num_tokens request_has_previous_tensor = True position_id = past_seen_token_num if _has_cp_helix: # We compute a global position_id because each helix rank has only a subset of # tokens for a sequence. position_id = request.total_input_len_cp + request.py_decoding_iter - 1 if request_has_previous_tensor: # With the overlap scheduler this batch is prepared # before the previous iteration's _update_requests has # advanced py_decoding_iter, so the counter is one # behind. Compensate exactly like the non-helix path # above, which uses max_beam_num_tokens *without* the # -1 in this case. Without this, the position repeats # once (L, L, L+1, ...) and the new token's K is roped # at the wrong position before being written to the KV # cache, corrupting every later step. # TODO: revisit for helix x speculative decoding - # the base formula and this +1 both assume exactly # one new token per step (draft-token modes are # currently rejected under helix). position_id += 1 if request.py_helix_is_inactive_rank: past_seen_token_num = request.seqlen_this_rank_cp else: # Discount the token added to active rank in resource manager as it hasn't # been previously seen. past_seen_token_num = request.seqlen_this_rank_cp - 1 for beam in range(beam_width): # Update helix-specific parameters. helix_is_inactive_rank.append( request.py_helix_is_inactive_rank) helix_position_offsets.append(position_id) # Keep the per-seq owned-count list aligned when the # spec path is active in the same batch. Whether the # list arms the spec path at all is decided once, at # the update_helix_param call below. helix_owned_new_tokens.append( 0 if request.py_helix_is_inactive_rank else 1) request.cached_tokens = past_seen_token_num for beam in range(beam_width): position_ids.append(position_id) num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens) prompt_lengths.append(request.py_prompt_len) gather_ids.append(len(position_ids) - 1) if _use_mrope: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) if mrope_position_delta is None and request.py_multimodal_data: mrope_config = request.py_multimodal_data.get( 'mrope_config') or {} mrope_position_delta = mrope_config.get( 'mrope_position_deltas') if mrope_position_delta is not None: if mrope_position_delta.device.type == "cpu": mrope_position_delta = maybe_pin_memory( mrope_position_delta).to(device='cuda', dtype=torch.int32, non_blocking=True) mrope_config[ 'mrope_position_deltas'] = mrope_position_delta request.py_mrope_position_delta = mrope_position_delta if mrope_position_delta is not None: has_gen_mrope_delta = True # NOTE: Expanding position_ids to 3D tensor who is using mrope gen_mrope_position_ids = (past_seen_token_num + mrope_position_delta).expand( 3, 1, 1) update_mrope_delta = ( request.py_seq_slot is not None and not request.is_dummy and getattr(request, "py_mrope_delta_cache_slot", None) != request.py_seq_slot) delta_read_seq_slot = (mrope_dummy_seq_slot if request.is_dummy or request.py_seq_slot is None else request.py_seq_slot) if update_mrope_delta: multimodal_params = MultimodalParams( multimodal_data={ 'mrope_config': { 'mrope_position_deltas': mrope_position_delta } }) mrope_delta_write_seq_slots.append( request.py_seq_slot) multimodal_params_list.append(multimodal_params) request.py_mrope_delta_cache_slot = request.py_seq_slot for beam in range(beam_width): # Locate this beam's single token in the flat array. token_start = len(position_ids) - beam_width + beam mrope_position_ids.append( (token_start, token_start + 1, gen_mrope_position_ids)) mrope_delta_read_seq_slots.append( delta_read_seq_slot) else: # No MRoPE metadata for this request (text-only prompt # on an MRoPE model): its delta is zero by construction, # so read the reserved zero slot instead of skipping the # append. The kernel indexes ``mrope_position_deltas`` # by *generation batch index* # (decoderMaskedMultiheadAttentionTemplate.h), so a list # that is sparse w.r.t. the generation batch would # silently shift every later request onto another # request's delta. No ``mrope_position_ids`` span is # recorded: the broadcast scalar position is already # this request's answer on all three axes. for _ in range(beam_width): mrope_delta_read_seq_slots.append( mrope_dummy_seq_slot) # Equivalent to the original `is_generation_admission and # request.py_multimodal_data`. The batch-level flag is checked # first so non-multimodal models pay one LOAD_FAST per request # instead of LOAD_ATTR(py_multimodal_data) + LOAD_ATTR(py_batch_idx). if (_has_any_multimodal_request and request.py_multimodal_data and request.py_batch_idx is None): strip_mm_data_for_generation(request.py_multimodal_data) request.py_batch_idx = request.py_seq_slot append_cross_attention_state(request, project_encoder_output=False, repeat=beam_width) # Do not add a gen_request_seq_slot for CUDA graph dummy requests # to prevent access errors due to None values if not request.is_cuda_graph_dummy: gen_request_seq_slots.append(request.py_seq_slot) if _use_mrope and not has_gen_mrope_delta: # Every generation request in this batch resolved to the zero slot, # so the gathered deltas would be an all-zero vector -- identical to # passing no deltas at all. Dropping the list keeps the steady-state # generation fast path (which requires the mrope lists to be empty) # reachable for text-only batches on MRoPE models. mrope_delta_read_seq_slots.clear() previous_batch_len = len(previous_batch_indices) def previous_seq_slots_device(): previous_batch_indices_host = torch.tensor( previous_batch_indices, dtype=torch.int, pin_memory=prefer_pinned()) previous_slots = self.previous_batch_indices_cuda[: previous_batch_len] previous_slots.copy_(previous_batch_indices_host, non_blocking=True) return previous_slots num_tokens = len(input_ids) num_draft_tokens = len(draft_tokens) total_num_tokens = len(position_ids) assert total_num_tokens <= self.max_num_tokens, ( f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})" ) # if exist requests that do not have previous batch, copy input_ids and draft_tokens if num_tokens > 0: input_ids = torch.tensor(input_ids, dtype=torch.int, pin_memory=prefer_pinned()) self.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True) # Update input_ids_cuda with new tokens from new_tensors_device (draft model only) if self.is_draft_model and num_accepted_tokens_device is not None: # For context requests: replace the last token with new_tensors_device[0, seq_slot, 0] if len(context_input_ids_positions) > 0: # Build tensors on CPU first, then copy to GPU to avoid implicit sync num_ctx_positions = len(context_input_ids_positions) ctx_token_indices_cpu = torch.tensor( [ last_token_idx for _, last_token_idx, _ in context_input_ids_positions ], dtype=torch.long, pin_memory=prefer_pinned()) ctx_seq_slots_cpu = torch.tensor([ seq_slot for _, _, seq_slot in context_input_ids_positions ], dtype=torch.long, pin_memory=prefer_pinned()) # Copy to pre-allocated GPU buffers self.draft_ctx_token_indices_cuda[:num_ctx_positions].copy_( ctx_token_indices_cpu, non_blocking=True) self.draft_ctx_seq_slots_cuda[:num_ctx_positions].copy_( ctx_seq_slots_cpu, non_blocking=True) self.input_ids_cuda[ self. draft_ctx_token_indices_cuda[:num_ctx_positions]] = new_tensors_device.new_tokens[ 0, self.draft_ctx_seq_slots_cuda[:num_ctx_positions], 0] # For first_draft requests: replace the last (original_max_draft_len+1) tokens # with new_tensors_device[:, seq_slot, 0] if len(first_draft_input_ids_positions) > 0: # All first_draft requests have same token length (original_max_draft_len + 1) # Build index tensors on CPU first, then copy to GPU to avoid implicit sync num_requests = len(first_draft_input_ids_positions) tokens_per_request = first_draft_input_ids_positions[0][ 1] - first_draft_input_ids_positions[0][0] # Create flat index array for all tokens to update on CPU all_indices = [] all_seq_slots = [] for start_idx, end_idx, seq_slot in first_draft_input_ids_positions: all_indices.extend(range(start_idx, end_idx)) all_seq_slots.extend([seq_slot] * (end_idx - start_idx)) # Create CPU tensors with pinned memory total_tokens = len(all_indices) idx_tensor_cpu = torch.tensor(all_indices, dtype=torch.long, pin_memory=prefer_pinned()) seq_slots_tensor_cpu = torch.tensor( all_seq_slots, dtype=torch.long, pin_memory=prefer_pinned()) # Copy to pre-allocated GPU buffers self.draft_first_draft_indices_cuda[:total_tokens].copy_( idx_tensor_cpu, non_blocking=True) self.draft_first_draft_seq_slots_cuda[:total_tokens].copy_( seq_slots_tensor_cpu, non_blocking=True) # Create token position indices (repeating 0..tokens_per_request for each request) token_positions = torch.arange( tokens_per_request, dtype=torch.long, device='cuda').repeat(num_requests) self.input_ids_cuda[ self. draft_first_draft_indices_cuda[:total_tokens]] = new_tensors_device.new_tokens[ token_positions, self. draft_first_draft_seq_slots_cuda[:total_tokens], 0] if num_draft_tokens > 0: draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()) self.draft_tokens_cuda[:len(draft_tokens)].copy_(draft_tokens, non_blocking=True) if self.is_spec_decode and len(num_accepted_draft_tokens) > 0: num_accepted_draft_tokens = torch.tensor(num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned()) self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)].copy_(num_accepted_draft_tokens, non_blocking=True) # Update num_accepted_draft_tokens_cuda for first_draft_requests directly from num_accepted_tokens_device (draft model only) if self.is_draft_model and len(first_draft_seq_slots) > 0: # Build tensors on CPU first, then copy to GPU to avoid implicit sync num_first_draft = len(first_draft_seq_slots) first_draft_seq_slots_cpu = torch.tensor( first_draft_seq_slots, dtype=torch.int, pin_memory=prefer_pinned()) first_draft_indices_cpu = torch.tensor( first_draft_request_indices, dtype=torch.int, pin_memory=prefer_pinned()) # Copy to pre-allocated GPU buffers self.draft_seq_slots_buffer_cuda[:num_first_draft].copy_( first_draft_seq_slots_cpu, non_blocking=True) self.draft_request_indices_buffer_cuda[:num_first_draft].copy_( first_draft_indices_cpu, non_blocking=True) # Extract accepted tokens for first_draft requests from device tensor accepted_tokens = num_accepted_tokens_device[ self.draft_seq_slots_buffer_cuda[:num_first_draft]] # Update the correct positions in num_accepted_draft_tokens_cuda self.num_accepted_draft_tokens_cuda[ self. draft_request_indices_buffer_cuda[: num_first_draft]] = accepted_tokens if next_draft_tokens_device is not None: # Initialize these two values to zeros self.previous_pos_id_offsets_cuda *= 0 self.previous_kv_lens_offsets_cuda *= 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 if previous_batch_len > 0: previous_slots = previous_seq_slots_device() # previous input ids previous_batch_tokens = (previous_batch_len * runtime_tokens_per_gen_step) new_tokens = new_tokens_device.transpose( 0, 1)[previous_slots, :runtime_tokens_per_gen_step].flatten() self.input_ids_cuda[num_tokens:num_tokens + previous_batch_tokens].copy_( new_tokens, non_blocking=True) # previous draft tokens previous_batch_draft_tokens = (previous_batch_len * runtime_draft_token_buffer_width) if runtime_draft_token_buffer_width > 0: self.draft_tokens_cuda[ num_draft_tokens:num_draft_tokens + previous_batch_draft_tokens].copy_( next_draft_tokens_device[ previous_slots, : runtime_draft_token_buffer_width].flatten(), non_blocking=True) # prepare data for the preprocess inputs kv_len_offsets_device = (new_tokens_lens_device - runtime_tokens_per_gen_step) previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned()) self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( previous_pos_indices_host, non_blocking=True) # The order of requests in a batch: [context requests, generation requests] # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] # 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. # 2) 'requests that already have previous batch': previous iteration's requests. # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments. # For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs. # Already set to '0' during initialization. # For 2) 'requests that already have previous batch': enable overlap scheduler. # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device. # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. # Already set to '0' during initialization. num_extend_reqeust_wo_dummy = len(extend_requests) - len( extend_dummy_requests) self.previous_pos_id_offsets_cuda[ (num_extend_reqeust_wo_dummy - previous_batch_len) * runtime_tokens_per_gen_step:num_extend_reqeust_wo_dummy * runtime_tokens_per_gen_step].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ 0:previous_batch_tokens]], non_blocking=True) self.previous_kv_lens_offsets_cuda[ num_extend_reqeust_wo_dummy - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( kv_len_offsets_device[previous_slots], non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() max_draft_len = max(draft_lens) new_tokens = new_tokens_device[:max_draft_len + 1, seq_slots_device, :self. max_beam_width] self.input_ids_cuda[num_tokens:num_tokens + previous_batch_len * self.max_beam_width].copy_( new_tokens.flatten(), non_blocking=True) if (not self._disable_overlap_scheduler and next_draft_tokens_device is None and len(extend_requests) > 0): # During warmup, for those generation requests, we don't have previous tensors, # so we need to set the previous_pos_id_offsets and previous_kv_lens_offsets to zeros # to skip the value changes in _preprocess_inputs. Otherwise, there will be illegal memory access # when writing key/values to the KV cache. self.previous_pos_id_offsets_cuda *= 0 self.previous_kv_lens_offsets_cuda *= 0 position_ids = apply_position_id_offset(position_ids, model=self.model) host_position_ids = torch.tensor(position_ids, dtype=torch.int, pin_memory=prefer_pinned()) # Use the (3,1,N) MRoPE layout whenever the model declares MRoPE, even # for text-only batches: keeping position_ids rank-consistent between # warmup and serving keeps torch.compile guards stable, so piecewise # CUDA graphs captured at warmup remain usable at runtime. if self.use_mrope: # Mixed batches may have only some requests with multimodal MRoPE # data. Seed the full (3,1,N) buffer from scalar position_ids # (text-only tokens get the same value on all 3 axes), then # overwrite only the multimodal spans with their real MRoPE coords. self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, non_blocking=True) # Broadcast [N] to [3,1,N]: default for text-only tokens. self.mrope_position_ids_cuda[:, :, :total_num_tokens].copy_( self.position_ids_cuda[:total_num_tokens].view(1, 1, -1).expand( 3, 1, -1), non_blocking=True) # Overwrite multimodal spans with per-axis MRoPE positions. for start_idx, end_idx, segment in mrope_position_ids: if segment.ndim != 3: raise RuntimeError( f"Expected 3D mrope_position_ids, got shape {tuple(segment.shape)}" ) if segment.shape[0] != 3 and segment.shape[-1] == 3: logger.warning( "Transposing unexpected mrope_position_ids shape from " f"{tuple(segment.shape)}") segment = segment.transpose(0, 2).contiguous() if segment.shape[:2] != (3, 1): raise RuntimeError( f"Unexpected mrope_position_ids shape {tuple(segment.shape)} for span {start_idx}:{end_idx}" ) segment = segment.contiguous() if segment.device.type == "cpu": segment = maybe_pin_memory(segment) self.mrope_position_ids_cuda[:, :, start_idx:end_idx].copy_( segment[:, :, :end_idx - start_idx], non_blocking=True) final_position_ids = self.mrope_position_ids_cuda[:, :, : total_num_tokens] else: self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, non_blocking=True) final_position_ids = self.position_ids_cuda[: total_num_tokens].unsqueeze( 0) if self.enable_spec_decode: self.gather_ids_cuda[:len(gather_ids)].copy_(torch.tensor( gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), non_blocking=True) # Update gather_ids for first_draft_requests on GPU (draft model only) if self.is_draft_model and len(first_draft_seq_slots) > 0: # Build tensors on CPU first, then copy to GPU to avoid implicit sync num_first_draft = len(first_draft_seq_slots) first_draft_seq_slots_cpu = torch.tensor( first_draft_seq_slots, dtype=torch.int, pin_memory=prefer_pinned()) first_draft_indices_cpu = torch.tensor( first_draft_request_indices, dtype=torch.int, pin_memory=prefer_pinned()) # Copy to pre-allocated GPU buffers self.draft_seq_slots_buffer_cuda[:num_first_draft].copy_( first_draft_seq_slots_cpu, non_blocking=True) self.draft_request_indices_buffer_cuda[:num_first_draft].copy_( first_draft_indices_cpu, non_blocking=True) # Extract accepted tokens for first_draft requests from device tensor accepted_tokens = num_accepted_tokens_device[ self.draft_seq_slots_buffer_cuda[:num_first_draft]] # Update gather_ids: gather_id = base_gather_id + num_accepted_tokens # (since gather_id = len(input_ids) - 1 - (max_draft_len - num_accepted)) self.gather_ids_cuda[ self. draft_request_indices_buffer_cuda[: num_first_draft]] += accepted_tokens if self.mapping.has_cp_helix(): # A non-None owned-count list is what arms # _helix_spec_tokens_valid, and the per-token slots/bounds that # flag gates are only ever filled by recompute_helix_spec_buffers, # which _preprocess_inputs runs under enable_spec_decode. Gate the # hand-off here, at the single choke point, so no packing loop can # arm the spec path for ordinary helix generation and send its # consumers to uninitialized buffers. helix_spec_active = bool(self.enable_spec_decode and helix_owned_new_tokens) attn_metadata.update_helix_param( helix_position_offsets=helix_position_offsets, helix_is_inactive_rank=helix_is_inactive_rank, helix_owned_new_tokens=(helix_owned_new_tokens if helix_spec_active else None), ) if not attn_metadata.is_cuda_graph: # Assumes seq lens do not change between CUDA graph invocations. This applies # to draft sequences too. This means that all draft sequences must be padded. attn_metadata.seq_lens = torch.tensor( sequence_lengths, dtype=torch.int, pin_memory=prefer_pinned(), ) num_generation_requests = len(gen_request_seq_slots) # Cache indirection is only used for beam search on generation requests if self.use_beam_search and num_generation_requests > 0: if cache_indirection_buffer is not None: #Copy cache indirection to local buffer with offsets changing: seq_slots[i] -> i # Convert to GPU tensor to avoid implicit sync gen_request_seq_slots_tensor = torch.tensor( gen_request_seq_slots, dtype=torch.long, pin_memory=prefer_pinned()).to(device='cuda', non_blocking=True) self.cache_indirection_attention[:num_generation_requests].copy_( cache_indirection_buffer[gen_request_seq_slots_tensor]) if cache_indirection_buffer is not None or self.is_warmup: attn_metadata.beam_width = self.max_beam_width else: attn_metadata.beam_width = 1 attn_metadata.request_ids = request_ids attn_metadata.prompt_lens = prompt_lengths attn_metadata.num_contexts = scheduled_requests.num_context_requests # Use num_chunked_ctx_requests to record the number of extend context requests, # so that we can update the kv_lens_cuda correctly in _preprocess_inputs. attn_metadata.num_chunked_ctx_requests = 0 if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: # For the tree decoding, we want to use XQA to process the draft tokens for the target model. # Therefore, we do not treat them as the chunked context requests. attn_metadata.num_contexts += len(extend_requests) attn_metadata.num_chunked_ctx_requests = len(extend_requests) attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq, num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config), use_full_generation_page_table=( self._should_use_full_generation_page_table( spec_config, attn_metadata))) attn_metadata.kv_cache_manager = kv_cache_manager if hasattr(self.model.model_config.pretrained_config, 'chunk_size'): attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size # Some sparse backends (RocketKV) clamp # kv_cache_params.num_cached_tokens_per_seq in place during prepare(), # and KVCacheParams holds the list by reference. Snapshot the true # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, cross_encoder_seq_lens, cross_encoder_cached_tokens_per_seq, attn_metadata, resource_manager, ) if is_enc_dec else {}) peft_cache_manager = resource_manager and resource_manager.get_resource_manager( ResourceManagerType.PEFT_CACHE_MANAGER) lora_params = self._lora.build( scheduled_requests, attn_metadata, cuda_graph_lora_manager=self.cuda_graph_lora_manager, enable_spec_decode=self.enable_spec_decode, runtime_draft_len=self.runtime_draft_len, peft_cache_manager=peft_cache_manager, maybe_graph=maybe_graph, use_lora_graph=use_lora_graph) if spec_metadata is not None: # Set the per-batch counts here, before the attention-DP allgather # below: the allgather and prepare() must derive the DP token count # from the same fields (see SpecMetadata.dp_num_tokens). Use # scheduled_requests.num_generation_requests -- the same-named # local above excludes CUDA-graph dummies and would not match the # count prepare() uses. spec_metadata.num_tokens = total_num_tokens spec_metadata.num_generations = \ scheduled_requests.num_generation_requests spec_metadata.seq_lens = sequence_lengths spec_all_rank_counts = None if spec_metadata is not None and self.enable_attention_dp: (attn_all_rank_num_tokens, spec_all_rank_counts ) = self._get_all_rank_num_tokens_and_spec_counts( attn_metadata, spec_metadata) else: attn_all_rank_num_tokens = get_all_rank_num_tokens( attn_metadata, enable_attention_dp=self.enable_attention_dp, mapping=self.mapping, dist=self.dist) (padded_num_tokens, can_run_prefill_cuda_graph, attn_all_rank_num_tokens) = get_padding_params( total_num_tokens, num_ctx_requests, attn_all_rank_num_tokens, dist=self.dist, enable_attention_dp=self.enable_attention_dp, prefill_cuda_graph_backend=self.prefill_cuda_graph_backend, prefill_cuda_graph_num_tokens=self._prefill_cuda_graph_num_tokens) set_per_request_prefill_cuda_graph_flag(can_run_prefill_cuda_graph) attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != total_num_tokens else None virtual_num_tokens = total_num_tokens if attn_metadata.padded_num_tokens is not None: self.input_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0) virtual_num_tokens = padded_num_tokens # Match the rank of the unpadded branch: MRoPE models always use # the (3,1,N) layout (see the seeding block above), so the padded # view must stay 3D as well to keep torch.compile guards stable. if self.use_mrope: # Zero-fill padding on dim 2 (token dim) of (3,1,N) buffer. self.mrope_position_ids_cuda[:, :, total_num_tokens: padded_num_tokens].fill_(0) final_position_ids = self.mrope_position_ids_cuda[:, :, : virtual_num_tokens] else: self.position_ids_cuda[ total_num_tokens:padded_num_tokens].fill_(0) final_position_ids = self.position_ids_cuda[: virtual_num_tokens].unsqueeze( 0) if self.enable_attention_dp: attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens # Prepare inputs inputs = { 'attn_metadata': attn_metadata, 'input_ids': self.input_ids_cuda[:virtual_num_tokens], 'position_ids': final_position_ids, 'inputs_embeds': None, "multimodal_params": multimodal_params_list, 'resource_manager': resource_manager, } inputs.update(cross_attention_inputs) if self.use_mrope: if mrope_delta_write_seq_slots: delta_write_seq_slots = torch.tensor( mrope_delta_write_seq_slots, dtype=torch.long, pin_memory=prefer_pinned()) inputs[ 'mrope_delta_write_seq_slots'] = delta_write_seq_slots.to( device='cuda', non_blocking=True) if mrope_delta_read_seq_slots: delta_read_seq_slots = torch.tensor(mrope_delta_read_seq_slots, dtype=torch.long, pin_memory=prefer_pinned()) inputs['mrope_delta_read_seq_slots'] = delta_read_seq_slots.to( device='cuda', non_blocking=True) if bool(lora_params): inputs['lora_params'] = lora_params if spec_metadata is not None: total_draft_lens = sum(draft_lens) spec_metadata.draft_tokens = self.draft_tokens_cuda[: total_draft_lens] spec_metadata.request_ids = request_ids spec_metadata.gather_ids = self.gather_ids_cuda[:len(gather_ids)] # num_generations / num_tokens / seq_lens are set above, before the # attention-DP allgather that must agree with prepare(). spec_metadata.host_position_ids = host_position_ids spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead( context_prompt_lookahead) # No-op for non 1-model spec_metadata.populate_sampling_params_for_one_model( scheduled_requests.all_requests()) spec_metadata.prepare() # One-model rejection: one-hot the stale draft_probs rows of gen # requests that produced no draft tokens this step, so the (possibly # captured) rejection kernel reads a legal placeholder distribution. spec_metadata.write_padding_onehot_draft_probs( padding_gen_slots, self.runtime_draft_len) inputs['spec_metadata'] = spec_metadata if self.enable_attention_dp: set_spec_metadata_all_rank_num_tokens(spec_metadata, *spec_all_rank_counts) if mm_token_indices is not None: ship_multimodal_indices( inputs, mm_token_indices_cpu=mm_token_indices, text_token_indices_cpu=text_token_indices_ctx, num_ctx_tokens=num_ctx_tokens, total_num_tokens=total_num_tokens, ) num_generation_tokens = len(generation_requests) + len( extend_requests) + sum(draft_lens) + len(first_draft_requests) self.iter_states['num_ctx_requests'] = num_ctx_requests self.iter_states['num_ctx_tokens'] = num_ctx_tokens self.iter_states['num_generation_tokens'] = num_generation_tokens # Count the already-cached prefix for the sequences scheduled this iteration. self.iter_states['cached_kv_tokens'] = sum(num_cached_tokens_per_seq) if self._log_cached_kv_tokens_per_req: self._record_cached_kv_tokens_per_req( num_cached_tokens_per_seq, ((scheduled_requests.context_requests, 1), (extend_requests, 1), (first_draft_requests, 1), (generation_requests, beam_width if generation_requests else 1))) if not self.is_warmup: self.previous_request_ids = all_gen_request_ids # Record the steady-state generation cache when this pass handled # purely non-dummy generation requests that all carried a previous # overlap-scheduler tensor (previous_batch_len == _n_gen implies # every request took that branch and none appended input_ids). # While the batch composition holds, the next passes only need to # advance positions by one and refresh per-step metadata. # MRoPE models are supported only for batches with no actual mrope # work (text-only requests, empty mrope lists below): the full # pass routes use_mrope models through the (3,1,N) # mrope_position_ids_cuda layout even then (to keep torch.compile # guards stable), with all three axes equal to the scalar # positions, so the fast path advances that buffer in place and # returns the same layout (see _apply_steady_gen_fast_prepare). if (self.spec_config is None and not self.is_draft_model and spec_metadata is None and new_tokens_device is not None and self.guided_decoder is None and not self.enable_attention_dp and not mrope_position_ids and not mrope_delta_write_seq_slots and not mrope_delta_read_seq_slots and not self.use_beam_search and self.max_beam_width == 1 and not is_enc_dec and not _has_cp_helix and num_ctx_requests == 0 and not extend_requests and not first_draft_requests and _n_gen > 0 and previous_batch_len == _n_gen and num_tokens == 0 and not _has_any_multimodal_request and not multimodal_params_list and not lora_params and attn_metadata.padded_num_tokens is None and get_position_id_offset(self.model) == 0 and not getattr(kv_cache_manager, "kv_compression_manages_history", False)): self._steady_gen_positions_pinned[:_n_gen].copy_( torch.as_tensor(num_cached_tokens_snapshot, dtype=torch.int)) self._steady_gen_cache = { 'num_requests': _n_gen, 'request_ids': all_gen_request_ids, 'prompt_lens': prompt_lengths, 'seq_lens_ones': maybe_pin_memory(torch.ones(_n_gen, dtype=torch.int)), 'use_mrope': _use_mrope, } return inputs, self.gather_ids_cuda[:len( gather_ids)] if self.enable_spec_decode else None @nvtx_range("_prepare_inputs") def _prepare_inputs( self, scheduled_requests: ScheduledRequests, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], attn_metadata: AttentionMetadata, spec_metadata: Optional[SpecMetadata] = None, new_tensors_device: Optional[SampleStateTensors] = None, cache_indirection_buffer: Optional[torch.Tensor] = None, num_accepted_tokens_device: Optional[torch.Tensor] = None, req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, resource_manager: Optional[ResourceManager] = None, maybe_graph: bool = False, promoted_context_request_ids: frozenset[int] = frozenset(), use_lora_graph: bool = False, ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]: set_per_request_prefill_cuda_graph_flag(False) if self.mapping is not None and 'cp_type' in self.mapping.cp_config: cp_type = self.mapping.cp_config['cp_type'] if cp_type in (CpType.HELIX, CpType.ULYSSES): # Take the usual route of _prepare_tp_inputs. pass else: raise NotImplementedError( f"Unsupported cp_type {getattr(cp_type, 'name', cp_type)}.") # Initialize SA state for new requests (MTP+SA, EAGLE3+SA, PARD+SA, etc.) has_sa_enhancer = (self.spec_config is not None and getattr( self.spec_config, 'sa_config', None) is not None) if has_sa_enhancer and resource_manager is not None and self.mapping.is_last_pp_rank( ): from tensorrt_llm._torch.speculative.suffix_automaton import \ SuffixAutomatonManager spec_rm = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) sa_manager = None if spec_rm is not None: if isinstance(spec_rm, SuffixAutomatonManager): sa_manager = spec_rm else: sa_manager = getattr(spec_rm, 'sa_manager', None) if sa_manager is not None: for request in scheduled_requests.all_requests(): if request.py_request_id not in sa_manager._initialized_requests: sa_manager.add_request(request.py_request_id, request.get_tokens(0)) sa_manager._initialized_requests.add( request.py_request_id) return self._prepare_tp_inputs(scheduled_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, num_accepted_tokens_device, req_id_to_old_request, resource_manager, maybe_graph, promoted_context_request_ids, use_lora_graph=use_lora_graph) @torch.inference_mode() @with_model_extra_attrs(lambda self: self.model.extra_attrs) def forward(self, batch: Union[ScheduledRequests, PackedEncoderBatch], resource_manager: Optional[ResourceManager] = None, new_tensors_device: Optional[SampleStateTensors] = None, gather_context_logits: bool = False, cache_indirection_buffer: Optional[torch.Tensor] = None, num_accepted_tokens_device: Optional[torch.Tensor] = None, req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None): if isinstance(batch, PackedEncoderBatch): assert self._is_packed_runner, ( "a packed batch requires a packed-batch runner") return cast(PackedModelRunner, self._runner).forward( batch, gather_context_logits=gather_context_logits) assert resource_manager is not None, ( "scheduled execution requires a resource manager") return self._forward_scheduled( batch, resource_manager, new_tensors_device=new_tensors_device, gather_context_logits=gather_context_logits, cache_indirection_buffer=cache_indirection_buffer, num_accepted_tokens_device=num_accepted_tokens_device, req_id_to_old_request=req_id_to_old_request, ) def _forward_scheduled(self, scheduled_requests: ScheduledRequests, resource_manager: ResourceManager, *, new_tensors_device: Optional[SampleStateTensors], gather_context_logits: bool, cache_indirection_buffer: Optional[torch.Tensor], num_accepted_tokens_device: Optional[torch.Tensor], req_id_to_old_request: Optional[Dict[int, LlmRequest]]): assert not self._is_packed_runner, ( "a packed-batch runner cannot execute scheduled requests") kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) if self._runner is not None and not self._is_encoder_decoder_model(): assert kv_cache_manager is None, ( "a no-KV-cache runner was initialized, but a KV cache manager was allocated" ) runner = cast(ModelRunner, self._runner) return runner.forward( scheduled_requests, resource_manager=resource_manager, cuda_graph_lora_manager=self.cuda_graph_lora_manager, runtime_draft_len=self.runtime_draft_len, gather_context_logits=gather_context_logits, ) assert kv_cache_manager is not None, ( "the legacy runner requires a KV cache manager") draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) attn_metadata = self._set_up_attn_metadata(kv_cache_manager, draft_kv_cache_manager) if isinstance(attn_metadata, TrtllmAttentionMetadata): attn_metadata.trtllm_gen_jit_warmup = self._trtllm_gen_jit_warmup if self.enable_spec_decode: spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) spec_tree_manager = None if spec_resource_manager is not None and hasattr( spec_resource_manager, 'spec_tree_manager'): spec_tree_manager = spec_resource_manager.spec_tree_manager spec_metadata = self._set_up_spec_metadata(spec_resource_manager) assert spec_metadata is not None update_spec_metadata( spec_metadata, scheduled_requests, attn_metadata, spec_tree_manager=spec_tree_manager, runtime_draft_len=self.runtime_draft_len, runtime_tokens_per_gen_step=( self.get_runtime_tokens_per_gen_step( self.runtime_draft_len)), is_draft_model=self.is_draft_model, attention_backend=self.attn_backend, original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=( self.original_max_total_draft_tokens), spec_dec_max_total_draft_tokens=( self._spec_dec_max_total_draft_tokens)) else: spec_resource_manager = None spec_metadata = None moe_load_balancer = self.moe_load_balancer graph_requests = scheduled_requests promoted_context_request_ids: frozenset[int] = frozenset() # Non-linear tree input preparation expands runtime_draft_len to the # total tree width after graph selection. Only linear-tree zero-draft # iterations can therefore safely reuse a zero-draft graph. can_promote_spec_decode = (not self.enable_spec_decode or (not self.is_draft_model and self.runtime_draft_len == 0 and self.spec_config is not None and self.spec_config.is_linear_tree)) # TODO: Generalize these conservative gates as actual-draft, beam, and # context-parallel providers for decoder-only LLMs gain support for # promoted final-context rows. Each relaxation must preserve whole-batch # fallback on graph miss and prove parity with the provider's native # q_len=1 path. Encoder-decoder and non-LLM engines remain out of scope. if (scheduled_requests.num_context_requests > 0 and self.cuda_graph_runner.enabled and can_promote_spec_decode and not self.use_beam_search and not self._is_encoder_decoder_model() # PLE owns recurrent n-gram and convolution state. Promoting a # fresh final-context row would skip its cache-slot reset. and not self._model_uses_ple_recurrent_state and self.mapping.cp_size == 1): graph_requests, promoted_context_request_ids = \ _make_single_token_context_graph_batch( scheduled_requests, self._is_final_multimodal_context_decode_compatible) with self.cuda_graph_runner.pad_batch( graph_requests, resource_manager, self.runtime_draft_len) as padded_graph_requests: # Callee already no-ops when use_mrope=False, but the Python call / # frame setup itself is non-trivial under high concurrency. Gating # at the caller avoids that overhead for non-mrope models. if self.use_mrope: self._pad_batch_seed_mrope_delta_cache(padded_graph_requests) # Refresh is_all_greedy_sample for the *current* batch BEFORE the # CUDA graph key is built below. The key includes this flag to pick # the argmax vs advanced-sampling graph variant; populate (inside # _prepare_inputs) runs later and fills the matching GPU buffers. # Without this pre-scan the key would use the previous iteration's # stale value and could replay the advanced graph against # unpopulated (greedy) buffers, hanging the run (e.g. MTP nextn>=2). if spec_metadata is not None: spec_metadata.update_is_all_greedy_sample( padded_graph_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) peft_cache_data_type = None if getattr(self, "cuda_graph_lora_manager", None) is not None: peft_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.PEFT_CACHE_MANAGER) peft_cache_data_type = peft_cache_manager.data_type use_lora_graph = self._use_lora_cuda_graph(padded_graph_requests) maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( padded_graph_requests, enable_spec_decode=self.enable_spec_decode, attn_metadata=attn_metadata, spec_metadata=spec_metadata, draft_tokens_cuda=self.draft_tokens_cuda if self.is_spec_decode else None, new_tensors_device=new_tensors_device, spec_resource_manager=spec_resource_manager, promoted_context_request_ids=promoted_context_request_ids, peft_cache_data_type=peft_cache_data_type, use_lora_graph=use_lora_graph, ) can_run_graph = key is not None if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata execution_requests = padded_graph_requests execution_promoted_context_ids = promoted_context_request_ids else: attn_metadata = self.attn_metadata if self.enable_spec_decode: spec_metadata = self.spec_metadata else: spec_metadata = None execution_requests = scheduled_requests execution_promoted_context_ids = frozenset() # Stage in-graph sampling now that the batch is settled: the staged # scatter width has to match the batch the forward actually runs on, # which differs between the graph (padded) and eager (unpadded) # branches above. Falling back to eager also means no graph replays, # so the tier is dropped and the sampler runs after the forward. # # Promoted one-token contexts join the graph batch's generation # requests, but sample_async is handed the original scheduled batch, # which excludes them. Staging a tier here would sample rows that # are then discarded and advance those requests' Philox streams an # extra time, so keep those steps on the eager path. if self._stage_in_graph_sampling is not None: staged_sample_type = (key.sample_type if can_run_graph and not execution_promoted_context_ids else SampleType.FULL) self._stage_in_graph_sampling(execution_requests, staged_sample_type) # Fill slot-ID buffer for scatter inside draft loop if (self.enable_spec_decode and spec_tree_manager is not None and not self.is_draft_model): spec_tree_manager.slot_storage.fill_all_slot_ids( execution_requests.context_requests, execution_requests.generation_requests, ) inputs, gather_ids = self._prepare_inputs( execution_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, num_accepted_tokens_device, req_id_to_old_request, resource_manager, can_run_graph, execution_promoted_context_ids, use_lora_graph=use_lora_graph) if execution_promoted_context_ids: self.iter_states[ 'num_ctx_requests'] = scheduled_requests.num_context_requests self.iter_states['num_ctx_tokens'] = sum( request.context_chunk_size for request in scheduled_requests.context_requests) self.iter_states[ 'num_generation_tokens'] = scheduled_requests.num_generation_requests self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() breakable_runner = self.breakable_cuda_graph_runner with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): def forward_step(): with MoeLoadBalancerIterContext(moe_load_balancer): return self._forward_step( inputs, gather_ids=gather_ids, gather_context_logits=gather_context_logits) if not can_run_graph: if (breakable_runner is not None and breakable_runner.is_capturing): return breakable_runner.capture_model_body(forward_step) num_tokens = inputs['input_ids'].shape[0] can_run_breakable_graph = ( breakable_runner is not None and get_per_request_prefill_cuda_graph_flag() and not gather_context_logits and breakable_runner.has_graph(num_tokens)) if can_run_breakable_graph and not breakable_runner.is_warming_up: outputs = breakable_runner.execute( num_tokens, forward_step) else: # real eager or BCG warmup or PCG outputs = forward_step() else: needs_capture = self.cuda_graph_runner.needs_capture(key) if needs_capture: def capture_forward_fn(inputs: Dict[str, Any]): with MoeLoadBalancerIterContext(moe_load_balancer): return self._forward_step( inputs, gather_ids=gather_ids, gather_context_logits=gather_context_logits) def capture_postprocess_fn(inputs: Dict[str, Any]): self._postprocess_inputs(inputs) capture_outputs = self.cuda_graph_runner.capture( key, capture_forward_fn, inputs, enable_spec_decode=self.enable_spec_decode, postprocess_fn=capture_postprocess_fn) if self.cuda_graph_runner.is_warmup_only: outputs = capture_outputs elif needs_capture: # Refresh attention metadata for the current batch's # draft cache before replaying the captured graph. saved_draft = prepare_attn_metadata_for_draft_replay( attn_metadata, draft_kv_cache_manager) try: outputs = self.cuda_graph_runner.replay(key, inputs) finally: restore_attn_metadata_after_draft_replay( attn_metadata, saved_draft) else: saved_draft = prepare_attn_metadata_for_draft_replay( attn_metadata, draft_kv_cache_manager) try: with MoeLoadBalancerIterContext(moe_load_balancer): outputs = self.cuda_graph_runner.replay( key, inputs) finally: restore_attn_metadata_after_draft_replay( attn_metadata, saved_draft) if self.forward_pass_callable is not None: self.forward_pass_callable() self._execute_logit_post_processors(scheduled_requests, outputs) return outputs def model_forward(self, **kwargs): attrs = get_model_extra_attrs() assert attrs is not None, "Model extra attrs is not set" attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) attrs.update(self.model.model_config.extra_attrs) attrs["spec_metadata"] = kwargs.get('spec_metadata', None) if self._torch_compile_backend is not None: # Register aux streams and events to model extra attrs. # The streams and events are list which could be updated during compilation. attrs["aux_streams"] = weakref.ref(self.backend_num_streams) attrs["events"] = weakref.ref(self._torch_compile_backend.events) attrs["global_stream"] = torch.cuda.current_stream() reclaimer = self._eager_workspace_reclaimer metadata = kwargs['attn_metadata'] reclaim_scope = (reclaimer.forward(metadata) if reclaimer is not None and not self.is_warmup and isinstance(metadata, TrtllmAttentionMetadata) else contextlib.nullcontext()) # Scope the entire top-level forward, including Eagle3's epilogue, so # eager decode and over-ceiling prefill do not select compile-only ops. compile_scope = ( torch_compiling(get_per_request_prefill_cuda_graph_flag()) if self._torch_compile_prefill_only else contextlib.nullcontext()) with reclaim_scope, compile_scope: if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): return trace_func(self.model.forward)(**kwargs) else: return self.model.forward(**kwargs) @nvtx_range("_forward_step") def _forward_step(self, inputs: Dict[str, Any], *, gather_ids: Optional[torch.Tensor] = None, gather_context_logits: bool = False) -> Dict[str, Any]: inputs = self._preprocess_inputs(inputs) if inputs.get('spec_metadata', None): gather_ids = inputs['spec_metadata'].gather_ids # For simplicity, just return all the the logits if we have special gather_ids # from speculative decoding. outputs = self.model_forward( **inputs, return_context_logits=gather_ids is not None or gather_context_logits, ) if self.without_logits: return outputs if isinstance(outputs, dict): # If the model returns a dict, get the logits from it. All other keys are kept. logits = outputs.get('logits', None) # If the logits are not found, no further processing is needed. if logits is None: return outputs else: # If the model returns a single tensor, assume it is the logits and wrap it in a dict. logits = outputs outputs = {'logits': logits} # If we have special gather_ids, gather the logits if gather_ids is not None: outputs['logits'] = logits[gather_ids] # Sample at the tail of the forward pass, so that under CUDA graph # capture the sampling kernels are recorded as part of this graph. The # hook is a no-op unless the sampler staged a graph-capturable tier for # this batch. # # Only the last PP rank runs the LM head. The others get a placeholder # that is the right shape but never filled, so a shape check waves it # through and they would sample garbage every step -- discarded, but # not free. _execute_logit_post_processors skips them for this reason. if (self.sample_in_graph_callable is not None and self.mapping.is_last_pp_rank()): self.sample_in_graph_callable(outputs) return outputs def _init_userbuffers(self, hidden_size): if self.mapping.tp_size <= 1 or self.mapping.pp_size > 1: return False # Disable UB for unsupported platforms if not ub.ub_supported(): return False # NCCL_SYMMETRIC strategy no longer requires UserBuffer allocator initialization. # It uses NCCLWindowAllocator from ncclUtils directly. if self.llm_args.allreduce_strategy == "NCCL_SYMMETRIC": # Skip UB initialization for NCCL_SYMMETRIC - it uses NCCLWindowAllocator directly return False ub.initialize_userbuffers_manager(self.mapping.tp_size, self.mapping.pp_size, self.mapping.cp_size, self.mapping.rank, self.mapping.gpus_per_node, hidden_size * self.max_num_tokens * 2) return True def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: """ When doing spec decode, sometimes draft models need to share certain weights with their target models. Here, we set up such weights by invoking self.model.load_weights_from_target_model if such a method exists. """ loader = getattr(self.model, "load_weights_from_target_model", None) if callable(loader): loader(target_model) @staticmethod def _apply_logits_processors(request, logits_processors, logits_tensor, beam_width, token_ids, logits_row_offset): logits_rows = logits_tensor[logits_row_offset:logits_row_offset + beam_width] # Reshape to align w/ the shape used in the TRT backend, # so the same logit processors can be used across both backends. logits_rows = logits_rows.view(beam_width, 1, -1) for lp in logits_processors: lp_params = inspect.signature(lp).parameters assert 4 <= len(lp_params) <= 5, ( "Logit post processor signature must match the `LogitsProcessor` interface " "defined in `tensorrtllm.sampling_params`.") lp(request.py_request_id, logits_rows, token_ids, None, None) # logits_rows is a view into logits_tensor (narrow + view never # copy), so the processors already mutated it in place. Writing it # back would be a self-assignment, which torch rejects for the # non-contiguous slices a TP-padded vocab produces. def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, outputs: dict): """Apply logit post processors (in-place modify outputs Tensors) if any.""" if not (self.mapping.is_last_pp_rank()): return if not isinstance(outputs, dict) or "logits" not in outputs: # TODO: support models that don't return outputs as dict return logits_tensor = outputs["logits"] logits_row_offset = 0 request_groups = ( (scheduled_requests.context_requests, True), (scheduled_requests.generation_requests, False), ) for requests, is_context_request in request_groups: for request in requests: if is_context_request: beam_width = 1 row_stride = 1 else: # Generation rows are laid out at the static admission # width, so that is the stride between requests, while # only the leading beam_width rows hold live beams under # a variable beam width array. Advancing the offset by the # narrower width would make every request after the first # rewrite another request's logits rows in place. beam_width = request.get_beam_width_by_iter( for_next_iteration=False) row_stride = request.py_beam_width logits_processors = getattr(request, "py_logits_post_processors", None) if logits_processors: token_ids = ([request.get_tokens(0)] if is_context_request else [ request.get_tokens(beam_idx) for beam_idx in range(beam_width) ]) if (is_context_request and request.py_orig_prompt_len < len(token_ids[0])): # Skip as we only need to apply logit processor on the last context request logits_row_offset += row_stride continue self._apply_logits_processors(request, logits_processors, logits_tensor, beam_width, token_ids, logits_row_offset) logits_row_offset += row_stride def wait_for_input_copy(self): """ Wait for input preparation and H2D copy of previous iteration before modifying host input, otherwise the input of previous iteration will be overwritten. """ if self._prepare_inputs_event is not None: self._prepare_inputs_event.synchronize()