From 4529767370b39dd32afcd784c9cecdbdc33e8cfb Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:40:21 -0700 Subject: [PATCH 1/8] [TRTLLM-14778][feat] Add a feature-encoder CUDA graph config An encoder that consumes fixed-shape per-request features emits the same number of positions for every request, so its graph key is the batch size alone and the token-shaped num_tokens / seq_lens buckets do not apply. encoder_cuda_graph_config becomes a discriminated union on mode so each encoder kind accepts only the buckets it has. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/llmapi/__init__.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 73 ++++++++++++++++++- .../usage/llm_args_golden_manifest.json | 3 +- .../api_stability/references/llm.yaml | 2 +- 4 files changed, 76 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 2ddc301eb01a..d707f397e265 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -15,7 +15,8 @@ DraftTargetDecodingConfig, DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, + ExtendedRuntimePerfKnobConfig, + FeatureEncoderCudaGraphConfig, KvCacheConfig, LlmArgs, LookaheadDecodingConfig, MambaStateConfig, MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, @@ -51,6 +52,7 @@ 'CudaGraphConfig', 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', + 'FeatureEncoderCudaGraphConfig', 'MoeConfig', 'LookaheadDecodingConfig', 'MedusaDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1afb2cc25994..bb6b4b91734a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -446,6 +446,47 @@ def _generate_cuda_graph_seq_lens(max_seq_len: int, return sizes +class FeatureEncoderCudaGraphConfig(StrictBaseModel): + """CUDA graph configuration for fixed-shape feature encoder requests. + + Applies to encoder-decoder models whose encoder consumes a fixed-shape + per-request feature tensor rather than packed tokens, e.g. Whisper's + 30 s-padded waveform. Such an encoder emits the same number of positions + for every request, so the graph key is the batch size alone and the + token-shaped `num_tokens` / `seq_lens` buckets of + :class:`EncodeCudaGraphConfig` do not apply. + """ + + mode: Literal["feature_encode"] = Field( + default="feature_encode", description="CUDA graph configuration mode.") + + batch_sizes: Optional[List[PositiveInt]] = Field( + default=None, + min_length=1, + description=( + "Encoder batch sizes to capture. None derives them from " + "`encoder_max_batch_size`, capped by the scheduler's encoder-batch " + "bound (max_num_tokens // encoder output length)."), + status="prototype", + ) + + enable_padding: bool = Field( + default=True, + description=( + "Pad an encoder batch up to the next captured batch size. Each pad " + "slot costs a full encoder forward, so padding is skipped when it " + "would add disproportionate encoder work."), + status="prototype", + ) + + @model_validator(mode='after') + def validate_feature_encoder_cuda_graph_config( + self) -> 'FeatureEncoderCudaGraphConfig': + if self.batch_sizes is not None: + self.batch_sizes = sorted(set(self.batch_sizes)) + return self + + # For CudaGraphConfig's backward compatibility CudaGraphConfig = DecodeCudaGraphConfig @@ -454,6 +495,11 @@ def _generate_cuda_graph_seq_lens(max_seq_len: int, Field(discriminator="mode"), ] +EncoderCudaGraphConfigType: TypeAlias = Annotated[ + Union[EncodeCudaGraphConfig, FeatureEncoderCudaGraphConfig], + Field(discriminator="mode"), +] + class MultimodalEncoderCudaGraphConfig(StrictBaseModel): """CUDA graph capture for multimodal vision / audio encoders. @@ -5083,13 +5129,16 @@ class TorchLlmArgs(BaseLlmArgs): Note that each CUDA graph can use up to 200 MB of extra memory.", status="beta") - encoder_cuda_graph_config: Optional[EncodeCudaGraphConfig] = Field( + encoder_cuda_graph_config: Optional[EncoderCudaGraphConfigType] = Field( default=None, description=( "CUDA graph configuration for the encoder forward pass of an " "encoder-decoder model. Use `cuda_graph_config` for the decoder " - "and this field for the encoder. Encoder CUDA graphs require " - "`encoder_max_batch_size` to be set."), + "and this field for the encoder. Pass an `EncodeCudaGraphConfig` " + "for a token encoder (T5/BART) or a " + "`FeatureEncoderCudaGraphConfig` for a fixed-shape feature encoder " + "(Whisper). Encoder CUDA graphs require `encoder_max_batch_size` " + "to be set."), status="prototype") enable_encoder_decoder_mixed_cuda_graph: bool = Field( @@ -5179,6 +5228,19 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: raise ValueError("must be a positive integer when set") return v + @field_validator('encoder_cuda_graph_config', mode='before') + @classmethod + def infer_encoder_cuda_graph_config_mode(cls, v): + if isinstance(v, dict) and "mode" not in v: + token_keys = { + "num_tokens", "max_num_token", "seq_lens", "max_seq_len" + } + v = dict(v) + v["mode"] = "encode" if any( + k in v and v[k] not in (None, 0) + for k in token_keys) else "feature_encode" + return v + @model_validator(mode="after") def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': if self.encoder_cuda_graph_config is None: @@ -5191,6 +5253,11 @@ def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': if self.encoder_max_batch_size is None: raise ValueError( "encoder_cuda_graph_config requires encoder_max_batch_size.") + if isinstance(self.encoder_cuda_graph_config, + FeatureEncoderCudaGraphConfig): + # A feature encoder's token counts and sequence lengths follow from + # the model, so batch_sizes is the only bucket dimension to require. + return self missing = [] if not self.encoder_cuda_graph_config.num_tokens: missing.append("num_tokens/max_num_token") diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 38436f22514a..143a2c3cf5f3 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -528,7 +528,8 @@ }, { "allowed_values": [ - "encode" + "encode", + "feature_encode" ], "annotation": "Literal['encode']", "converter": "", diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 409eafe66895..d1a3832fdcc4 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -100,7 +100,7 @@ methods: default: null status: beta encoder_cuda_graph_config: - annotation: Optional[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig] + annotation: Union[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig, tensorrt_llm.llmapi.llm_args.FeatureEncoderCudaGraphConfig, NoneType] default: null status: prototype enable_encoder_decoder_mixed_cuda_graph: From 411cbe346f44ab56d9e3f01acede5370c63515cf Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:40:21 -0700 Subject: [PATCH 2/8] [TRTLLM-14778][perf] Capture the encoder step in a CUDA graph for feature encoders Whisper's encoder takes a 30 s-padded waveform per request, so the runner swaps its packed-token static tensors for an input_features buffer keyed on batch size. Capture goes through the shared two-pass warmup helper and runs on a dedicated stream, because encoder replay is device-concurrent with decoder replay. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_whisper.py | 32 ++ .../_torch/pyexecutor/cuda_graph_runner.py | 257 +++++++++++- .../_torch/pyexecutor/model_engine.py | 370 ++++++++++++++++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 27 +- 4 files changed, 653 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_whisper.py b/tensorrt_llm/_torch/models/modeling_whisper.py index 420b7ea6f9de..edc3ae28f832 100644 --- a/tensorrt_llm/_torch/models/modeling_whisper.py +++ b/tensorrt_llm/_torch/models/modeling_whisper.py @@ -654,6 +654,17 @@ def _build_decoder_prompt(self) -> List[int]: forced = self.processor.get_decoder_prompt_ids(no_timestamps=True) return [int(start_id)] + [int(tok) for _, tok in sorted(forced)] + def get_decoder_prefix_len(self) -> int: + """Tokens every request's decoder prompt starts with. + + Mixed encoder/decoder CUDA graphs capture at this query length, and a + mismatch makes every mixed batch miss its graph silently. This reports + the checkpoint default, so a request carrying a text prompt of a + different length (see `_resolve_decoder_prompt`) misses the mixed graph + and runs that batch eagerly. + """ + return len(self._build_decoder_prompt()) + def _resolve_decoder_prompt(self, prompt_text: Optional[str]) -> List[int]: """Checkpoint-default forced prompt, or the user's decoder prompt. @@ -880,6 +891,27 @@ def __pp_init__(self): def config(self): return self.model_config.pretrained_config + def encoder_graph_spec(self): + """Fixed-shape encoder contract for enc-dec encoder CUDA graphs. + + Every Whisper encoder request is an fp32 waveform zero-padded by + `WhisperInputProcessor` to the fixed window that yields exactly + ``max_source_positions`` encoder positions — so the encoder graph key + degenerates to the batch size. + + The window is derived rather than hardcoded: the conv stem halves the + STFT frame count, so ``n_samples = max_source_positions * 2 * + hop_length``. ``hop_length`` comes from the checkpoint's feature + extractor, the same source the input processor validates its own + window against. + + Returns ``(per_request_feature_shape, dtype, fixed_seq_len)``. + """ + fixed_seq_len = int(self.config.max_source_positions) + hop_length = int(self.model.encoder.log_mel.hop_length) + n_samples = fixed_seq_len * 2 * hop_length + return ((n_samples,), torch.float32, fixed_seq_len) + def forward( self, attn_metadata: AttentionMetadata, diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 4019c037f513..d23d5284623b 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1047,6 +1047,15 @@ class EncoderCUDAGraphRunnerConfig: is_encoder_decoder: bool = False use_fixed_sequence_slots: bool = False + # Feature mode (encoders taking fixed-shape per-request feature tensors, + # e.g. Whisper's [480000] fp32 waveform). When feature_shape is set, the + # runner replaces the input_ids/position_ids static tensors with an + # input_features buffer and the graph key degenerates to + # (bs, bs * fixed_seq_len, fixed_seq_len). + feature_shape: Optional[Tuple[int, ...]] = None + feature_dtype: Optional[torch.dtype] = None + fixed_seq_len: Optional[int] = None + class EncoderCUDAGraphRunner: """CUDA graph runner for no-cache encoder forward passes. @@ -1061,6 +1070,10 @@ class EncoderCUDAGraphRunner: """ WARMUP_STEPS = 1 + MAX_FEATURE_PADDING_RATIO = 9 / 8 + # Host-side feature mirrors. Two is enough to hide the host fill behind + # one in-flight H2D; more would only add pinned memory. + FEATURE_MIRROR_SLOTS = 2 def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.config = config @@ -1069,14 +1082,50 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.padding_enabled = config.cuda_graph_padding_enabled self.supported_batch_sizes = sorted(config.cuda_graph_batch_sizes) self.max_supported_batch_size = config.max_cuda_graph_batch_size - self.supported_num_tokens = sorted(config.cuda_graph_num_tokens) - self.max_supported_num_tokens = config.max_cuda_graph_num_tokens - self.supported_seq_lens = sorted(config.cuda_graph_seq_lens) + self.feature_mode = config.feature_shape is not None self.is_encoder_decoder = config.is_encoder_decoder self.use_fixed_sequence_slots = config.use_fixed_sequence_slots + + if self.feature_mode and not self.is_encoder_decoder: + # Nothing structural forbids this - a standalone feature encoder + # (audio/vision embedding tower) would land here - but no in-tree + # model exercises it, so fail loudly rather than capture untested + # shapes. + raise NotImplementedError( + "Feature-mode encoder CUDA graphs are only supported for " + "encoder-decoder models today.") + + if self.feature_mode: + # A feature encoder produces a fixed number of positions per + # request, so the configured token/seq-len buckets are not free + # parameters: they degenerate to multiples of fixed_seq_len over + # the batch sizes. Any user-supplied values are ignored. + fixed = config.fixed_seq_len + self.supported_num_tokens = sorted( + bs * fixed for bs in self.supported_batch_sizes) + self.max_supported_num_tokens = (self.max_supported_batch_size * + fixed) + self.supported_seq_lens = [fixed] + else: + self.supported_num_tokens = sorted(config.cuda_graph_num_tokens) + self.max_supported_num_tokens = config.max_cuda_graph_num_tokens + self.supported_seq_lens = sorted(config.cuda_graph_seq_lens) + self.capture_keys: frozenset[EncoderKeyType] = frozenset() self._capture_sequence_lengths: Dict[EncoderKeyType, List[int]] = {} - if self.is_encoder_decoder: + if self.feature_mode: + # Every request contributes exactly fixed_seq_len positions, so one + # key per batch size is the complete reachable set. Taking the + # token path's cross product of batch sizes and token counts would + # instead yield keys no batch can ever match, and `capture_keys` + # also drives mixed encoder/decoder decoder-graph warmup. + self._capture_sequence_lengths = { + (bs, bs * config.fixed_seq_len, config.fixed_seq_len): + [config.fixed_seq_len] * bs + for bs in self.supported_batch_sizes + } + self.capture_keys = frozenset(self._capture_sequence_lengths) + elif self.is_encoder_decoder: self._capture_sequence_lengths = ( self._build_encoder_decoder_capture_layouts()) self.capture_keys = frozenset(self._capture_sequence_lengths) @@ -1101,11 +1150,27 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.is_warmup_only = False self._staging_retirement_event: Optional[torch.cuda.Event] = None + # `torch.cuda.graph` falls back to a process-wide singleton capture + # stream when `stream=` is omitted, so the encoder graphs would capture + # on the same stream as the decoder graphs. Sharing it couples the two + # graph sets through stream-keyed cuBLAS/cuBLASLt scratch: a serial + # split-K GEMM captured into one graph spins forever in + # cutlass::Semaphore::wait() once the other graph's matmuls leave that + # region non-zero, because the captured graph has no node that re-zeroes + # it. Capture on our own stream instead. + self._capture_stream: Optional[torch.cuda.Stream] = None + # CUDA graph H2D memcpy nodes require pinned host sources. In CC mode # prefer_pinned() is false: pageable host buffers are preferred, so the # H2D copies must be issued before graph replay instead of captured. self._capture_h2d_copy = prefer_pinned() + def _get_capture_stream(self) -> torch.cuda.Stream: + """Return this runner's dedicated capture stream, creating it lazily.""" + if self._capture_stream is None: + self._capture_stream = torch.cuda.Stream() + return self._capture_stream + def _create_shared_static_tensors(self): """Allocates static tensors sized for the largest supported num_tokens.""" max_total_tokens = ( @@ -1113,6 +1178,47 @@ def _create_shared_static_tensors(self): self.max_supported_num_tokens, self.config.max_num_tokens)) max_batch_size = self.max_supported_batch_size + if self.feature_mode: + feature_shape = (max_batch_size, *self.config.feature_shape) + self.shared_static_tensors = { + "input_features": + torch.zeros(feature_shape, + device="cuda", + dtype=self.config.feature_dtype), + } + self.shared_static_tensors_cpu = { + "seq_lens": + torch.full((max_batch_size, ), + self.config.fixed_seq_len, + device="cpu", + dtype=torch.int32, + pin_memory=prefer_pinned()), + } + # Host mirrors are double-buffered; the device buffer is not. The + # device buffer is captured into every graph, so its refill must + # stay ordered behind the previous replay that reads it - that is + # a real data dependency, not an artifact of stream choice. The + # *host* fill has no such constraint, so filling mirror B while + # the device still drains mirror A takes it off the critical path. + # One event per mirror records the H2D that read it; a mirror is + # refilled only once its own H2D has completed, which with two + # slots is normally already true. + self._feature_mirrors = [ + torch.zeros(feature_shape, + device="cpu", + dtype=self.config.feature_dtype, + pin_memory=prefer_pinned()) + for _ in range(self.FEATURE_MIRROR_SLOTS) + ] + self._feature_h2d_events = [ + torch.cuda.Event() for _ in range(self.FEATURE_MIRROR_SLOTS) + ] + # Record once so the first use of each slot does not block. + for event in self._feature_h2d_events: + event.record() + self._feature_mirror_slot = 0 + return + self.shared_static_tensors = { "input_ids": torch.ones((max_total_tokens, ), device="cuda", dtype=torch.int32), @@ -1400,6 +1506,23 @@ def pad_batch(self, inputs: Dict[str, Any], yield inputs return + if self.feature_mode: + # A feature pad slot is a full fixed_seq_len request of compute + # (zero-filled input rows, outputs discarded at scatter), unlike + # the 1-token pads of the token path. Fall back to eager across + # large bucket gaps so graph replay cannot add more than 12.5% + # encoder work. + if (batch_size == 0 or padded_batch_size + > batch_size * self.MAX_FEATURE_PADDING_RATIO): + yield inputs + return + padded_inputs = dict(inputs) + padded_inputs['seq_lens'] = (list(inputs['seq_lens']) + + [self.config.fixed_seq_len] * + (padded_batch_size - batch_size)) + yield padded_inputs + return + padding_size = padded_batch_size - batch_size # Should not pad inputs if it would exceed the max supported number of tokens # maybe_get_cuda_graph will check this and fall back to eager if batch size is not in the supported list @@ -1495,6 +1618,12 @@ def maybe_get_cuda_graph( if padded_batch_size not in self.supported_batch_sizes: return None, None + if self.feature_mode and any(s != self.config.fixed_seq_len + for s in seq_lens): + # Fixed-shape contract violated (should not happen for feature + # encoders); fall back to eager rather than replay a wrong shape. + return None, None + key, is_padding_performed, is_padding_successful = self.get_graph_key( inputs) if self.is_encoder_decoder and key not in self.capture_keys: @@ -1717,6 +1846,9 @@ def capture( """Warm up and/or capture the forward pass for a graph key.""" padded_num_tokens = key[1] + if self.feature_mode: + return self._capture_features(key, forward_fn, inputs) + sliced_static_tensors = { "input_ids": self.shared_static_tensors["input_ids"][:padded_num_tokens], @@ -1767,6 +1899,7 @@ def capture( graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, pool=self.memory_pool, + stream=self._get_capture_stream(), capture_error_mode="thread_local"): if self._capture_h2d_copy: # H2D copies for captured inside the graph: at replay @@ -1798,12 +1931,128 @@ def retire_staging(self) -> None: self._staging_retirement_event.synchronize() self._staging_retirement_event = None + def _capture_features( + self, + key: EncoderKeyType, + forward_fn: Callable[[Dict[str, Any]], Any], + inputs: Dict[str, Any], + ) -> Any: + """Capture path for the fixed-shape feature mode (enc-dec encoders). + + The capture region receives the static device feature buffer sliced + to the padded batch size; in pinned mode the H2D from the pinned CPU + mirror is captured inside the graph so replay re-issues it without an + eager driver call. + """ + padded_batch_size, _, _ = key + + static_features = ( + self.shared_static_tensors["input_features"][:padded_batch_size]) + + capture_inputs = dict(inputs) + capture_inputs["input_features"] = static_features + + attn_md = capture_inputs["attn_metadata"] + self.graph_metadata[key] = { + "attn_metadata": attn_md, + } + + # Feature-mode seq_lens never change for this key: populate the + # metadata's device seq_lens once, eagerly, instead of capturing the + # H2D like the token path does per replay. + attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) + + # NOTE: unlike the token path, the input H2D is NOT captured inside + # the graph. All buckets share one pinned mirror, and consecutive + # encoder batches (different buckets) can be enqueued back-to-back — + # a captured H2D would read the mirror at replay-execution time, + # after the host has already refilled it for the next batch. The + # eager H2D in `_replay_features` is stream-ordered and guarded by + # per-mirror events instead. + output = None + with with_multi_stream(True), piecewise_cuda_graph(False): + for _ in range(self.WARMUP_STEPS): + output = forward_fn(capture_inputs) + + # The warmup pass runs these shapes eagerly to settle PyTorch and + # attention state; it must not build a graph, and its caller + # consumes the eager output directly. + if self.is_warmup_only: + return output + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, + pool=self.memory_pool, + stream=self._get_capture_stream()): + output = forward_fn(capture_inputs) + + if self._contains_nested_tensor(output): + raise TypeError( + "Encoder CUDA graph does not support nested tensor outputs.") + self.graphs[key] = graph + graph_output = make_weak_ref(output) + self.graph_outputs[key] = graph_output + self.memory_pool = graph.pool() + return graph_output + + def _replay_features( + self, + key: EncoderKeyType, + inputs: Dict[str, Any], + ) -> Any: + """Replay path for the fixed-shape feature mode.""" + stored_meta = self.graph_metadata[key] + assert inputs["attn_metadata"] is stored_meta["attn_metadata"] + + padded_batch_size, _, _ = key + features = inputs["input_features"] + + slot = self._feature_mirror_slot + self._feature_mirror_slot = (slot + 1) % self.FEATURE_MIRROR_SLOTS + + # Wait only for the H2D that last read *this* mirror. With two slots + # that copy was issued two batches ago, so this is normally already + # satisfied and the host proceeds straight to the fill. + self._feature_h2d_events[slot].synchronize() + + mirror = self._feature_mirrors[slot] + if isinstance(features, list): + # Per-request CPU tensors straight from the requests — one copy + # into the host mirror, no intermediate packing. + rows = 0 + for f in features: + n = int(f.shape[0]) + mirror[rows:rows + n].copy_(f) + rows += n + else: + rows = int(features.shape[0]) + mirror[:rows].copy_(features) + if rows < padded_batch_size: + mirror[rows:padded_batch_size].zero_() + + # Eager, stream-ordered H2D: runs after any previously enqueued replay + # on this stream, so it cannot race an in-flight graph reading the + # single device buffer. Keeping it on this stream is load-bearing. + self.shared_static_tensors["input_features"][:padded_batch_size].copy_( + mirror[:padded_batch_size], non_blocking=True) + self._feature_h2d_events[slot].record() + + self.graphs[key].replay() + return self.graph_outputs[key] + def replay( self, key: EncoderKeyType, inputs: Dict[str, Any], ) -> Any: """Replay a captured graph with current inputs.""" + if self.feature_mode: + # Feature mode stages through its own double-buffered pinned + # mirrors and guards them with per-mirror events; `retire_staging` + # covers the token path's shared host staging buffers, which + # feature mode does not touch. + return self._replay_features(key, inputs) + self.retire_staging() stored_meta = self.graph_metadata[key] diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index bd832625d042..159b954198ac 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -35,6 +35,7 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, + FeatureEncoderCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -557,17 +558,19 @@ def __init__( encoder_cuda_graph_batch_sizes = ( self.encoder_cuda_graph_config.batch_sizes if self.encoder_cuda_graph_config is not None else []) - encoder_cuda_graph_num_tokens = ( - self.encoder_cuda_graph_config.num_tokens - if self.encoder_cuda_graph_config is not None else []) - encoder_cuda_graph_seq_lens = (self.encoder_cuda_graph_config.seq_lens - if self.encoder_cuda_graph_config - is not None else []) + # A feature encoder config carries batch sizes only: its token counts + # and sequence lengths follow from the model's fixed encoder output + # length, which the encoder graph runner derives. + encoder_cuda_graph_num_tokens = getattr(self.encoder_cuda_graph_config, + 'num_tokens', None) or [] + encoder_cuda_graph_seq_lens = getattr(self.encoder_cuda_graph_config, + 'seq_lens', None) or [] encoder_cuda_graph_padding_enabled = ( self.encoder_cuda_graph_config.enable_padding if self.encoder_cuda_graph_config is not None else False) if (self.encoder_cuda_graph_config is not None + and not self._is_feature_encoder_cuda_graph_config() and (not encoder_cuda_graph_num_tokens or not encoder_cuda_graph_seq_lens)): missing = [] @@ -618,11 +621,12 @@ def __init__( self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) - use_encoder_cuda_graph = ((self._is_encoder_decoder_model() - or self._is_encode_only) - and self.encoder_cuda_graph_config is not None - and bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens)) + use_encoder_cuda_graph = ( + (self._is_encoder_decoder_model() or self._is_encode_only) + and self.encoder_cuda_graph_config is not None + and (self._is_feature_encoder_cuda_graph_config() or + (bool(self._cuda_graph_num_tokens) + and bool(self._cuda_graph_seq_lens)))) self.torch_compile_config = self.llm_args.torch_compile_config torch_compile_enabled = bool(self.torch_compile_config is not None) @@ -866,6 +870,48 @@ def __init__( encoder_graph_max_batch_size = (encoder_graph_batch_sizes[-1] if encoder_graph_batch_sizes else 0) encoder_graph_max_num_tokens = self._max_cuda_graph_num_tokens + + # A feature-driven encoder (Whisper) declares a fixed-shape per-request + # contract instead of packed tokens, so its graph shapes follow from + # the batch sizes alone. Enablement still goes through the same + # `encoder_cuda_graph_config` opt-in as the token path. + feature_shape, feature_dtype, fixed_seq_len = self._encoder_graph_spec() + if feature_shape is not None: + bs_cap = max( + 1, + min(self.batch_size, + self.encoder_max_num_tokens // fixed_seq_len)) + if not encoder_graph_batch_sizes: + # batch_sizes left unset: derive them, bounded by the encoder + # batch limit and the scheduler's encoder-batch bound. + encoder_graph_batch_sizes = ( + self._derive_feature_encoder_batch_sizes( + min(self.encoder_batch_size, bs_cap))) + encoder_graph_batch_sizes = sorted( + bs for bs in encoder_graph_batch_sizes if bs <= bs_cap) + if not encoder_graph_batch_sizes: + logger.warning( + "Feature-mode encoder CUDA graphs: no configured batch " + f"size fits within max_num_tokens // encoder_output_len = " + f"{bs_cap}; the encoder step stays eager.") + feature_shape = feature_dtype = fixed_seq_len = None + else: + encoder_graph_max_batch_size = encoder_graph_batch_sizes[-1] + encoder_graph_max_num_tokens = (encoder_graph_max_batch_size * + fixed_seq_len) + elif (use_encoder_cuda_graph + and self._model_encoder_graph_spec() is not None): + # A feature encoder cannot consume the packed token inputs the + # token-shaped capture path synthesizes, so keep it eager rather + # than let warmup drive tokens into it. + logger.warning( + "This model's encoder consumes fixed-shape features, but " + "encoder_cuda_graph_config is an EncodeCudaGraphConfig, whose " + "token buckets do not apply; the encoder step stays eager. " + "Pass FeatureEncoderCudaGraphConfig(batch_sizes=[...]) to " + "capture it.") + use_encoder_cuda_graph = False + encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( use_cuda_graph=use_encoder_cuda_graph, cuda_graph_padding_enabled=( @@ -877,15 +923,27 @@ def __init__( max_cuda_graph_num_tokens=encoder_graph_max_num_tokens, max_num_tokens=self.encoder_max_num_tokens, max_seq_len=self.max_seq_len, - cuda_graph_mem_pool=self._cuda_graph_mem_pool, + # The encoder runner takes its own graph pool. Encoder replay runs + # on `encoder_stream`, device-concurrent with decoder replay, and + # torch's pool-sharing contract assumes replays from a shared pool + # are not concurrent. + cuda_graph_mem_pool=None, is_encoder_decoder=self._is_encoder_decoder_model(), use_fixed_sequence_slots=(self._is_encoder_decoder_model() and hasattr( pretrained_config, "relative_attention_num_buckets")), + feature_shape=feature_shape, + feature_dtype=feature_dtype, + fixed_seq_len=fixed_seq_len, ) self.encoder_cuda_graph_runner = EncoderCUDAGraphRunner( encoder_cuda_graph_runner_config) + if feature_shape is not None: + logger.info( + f"Feature-mode encoder CUDA graphs enabled for batch sizes " + f"{encoder_graph_batch_sizes} (fixed_seq_len={fixed_seq_len}, " + f"feature_shape={tuple(feature_shape)}).") # Once encoder CUDA graphs are usable, enable mixed decoder graphs by # default unless the user explicitly opts out. @@ -1353,6 +1411,11 @@ def warmup(self, resource_manager: ResourceManager) -> None: 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 self.encoder_cuda_graph_runner.feature_mode: + # After decoder-graph capture, so the decoder pool's high-water + # mark is set before the encoder runner allocates its own pool. + self._capture_enc_dec_encoder_graphs() + log_mem_snapshot("warmup/after_enc_dec_encoder_graph_capture") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -2064,6 +2127,11 @@ def _warmup_encoder_cuda_graphs_enc_dec( runner = self.encoder_cuda_graph_runner if not runner.is_encoder_decoder: return + if runner.feature_mode: + # This driver synthesizes packed token inputs, which a feature + # encoder cannot consume. Feature mode captures from + # `_capture_enc_dec_encoder_graphs` during engine warmup instead. + return capture = functools.partial( self._capture_encoder_cuda_graphs_enc_dec, @@ -2379,10 +2447,20 @@ def _capture_mixed_encoder_decoder_cuda_graphs( if max_num_encoder_tokens == 0: return model_config = self.model.model_config.pretrained_config - # BART/mBART prepend a forced BOS token after decoder_start; T5 uses - # decoder_start alone. Match the LLM API's decoder-prefix construction. - mixed_context_query_len = (2 if getattr( - model_config, "model_type", None) in ("bart", "mbart") else 1) + # 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: @@ -3626,6 +3704,71 @@ def _is_encoder_decoder_model(self) -> bool: getattr(getattr(self.model, "model_config", None), "is_encoder_decoder", False)) + def _is_feature_encoder_cuda_graph_config(self) -> bool: + """True when the encoder graph config selects the feature-shaped path.""" + return isinstance(self.encoder_cuda_graph_config, + FeatureEncoderCudaGraphConfig) + + def _model_encoder_graph_spec(self): + """The model's fixed-shape encoder contract, or None. Queried once.""" + if not hasattr(self, "_cached_model_encoder_graph_spec"): + # torch.compile wraps the model; the spec is on the original. + model = getattr(self.model, "_orig_mod", self.model) + spec_fn = getattr(model, "encoder_graph_spec", None) + self._cached_model_encoder_graph_spec = (spec_fn() if spec_fn + is not None else None) + return self._cached_model_encoder_graph_spec + + @staticmethod + def _derive_feature_encoder_batch_sizes(max_batch_size: int) -> List[int]: + """Encoder graph batch sizes for a fixed-shape feature encoder. + + Dense up to 8, then multiples of 8. A feature pad slot costs a full + encoder forward, so padding is refused once it would add more than + ``MAX_FEATURE_PADDING_RATIO`` of encoder work; sparse buckets at the + low end would leave small batches permanently eager. + """ + sizes = set(range(1, min(max_batch_size, 8) + 1)) + sizes.update(range(8, max_batch_size + 1, 8)) + sizes.add(max_batch_size) + return sorted(bs for bs in sizes if 0 < bs <= max_batch_size) + + def _encoder_graph_spec(self): + """Fixed-shape encoder contract, or (None, None, None) if unavailable. + + Returns ``(feature_shape, feature_dtype, fixed_seq_len)`` when the user + selected ``FeatureEncoderCudaGraphConfig``, the model declares + ``encoder_graph_spec()`` and feature-mode encoder CUDA graphs are + viable. Gated to TP=1 (allreduce inside encoder capture is unverified) + and to non-draft models. + """ + none = (None, None, None) + if (self.encoder_cuda_graph_config is None or self.is_draft_model + or not self._is_encoder_decoder_model()): + return none + + spec = self._model_encoder_graph_spec() + + if not self._is_feature_encoder_cuda_graph_config(): + return none + + if spec is None: + raise ValueError( + "FeatureEncoderCudaGraphConfig requires a model whose encoder " + "declares a fixed-shape encoder_graph_spec(); this model does " + "not. Token-driven encoders such as T5 and BART take " + "EncodeCudaGraphConfig with num_tokens and seq_lens set. Note " + "that an encoder_cuda_graph_config supplied without those two " + "fields is read as a feature-encoder config.") + + if self.mapping.tp_size > 1: + logger.warning( + "Feature-mode encoder CUDA graphs are gated to TP=1 in this " + "phase; the encoder step stays eager.") + return none + + return spec + def _get_top_level_model(self) -> Any: model = getattr(self.model, "_orig_mod", self.model) top_level_model = getattr(model, "model", model) @@ -7506,17 +7649,64 @@ def _prepare_tp_inputs_encoder_features( sequence_lengths, request_ids) inputs = { - 'input_features': - torch.cat(features, dim=0).to('cuda', non_blocking=True), - 'encoder_attn_metadata': - encoder_attn_metadata, - 'encoder_seq_lens': - sequence_lengths, - 'resource_manager': - resource_manager, + 'input_features': self._pack_encoder_features(features), + 'encoder_attn_metadata': encoder_attn_metadata, + 'encoder_seq_lens': sequence_lengths, + 'resource_manager': resource_manager, } return inputs + def _pack_encoder_features(self, + features: List[torch.Tensor]) -> torch.Tensor: + """Pack per-request feature tensors into one device tensor. + + Copies through a lazily-grown pinned staging buffer so the H2D + transfer is a single async DMA. ``torch.cat(...).to('cuda')`` from + pageable request tensors forces a synchronous driver-staged copy per + batch, which dominates encoder host time at large batch sizes + (measured 51.7 ms/call at bs32 on a Xeon 8570 host). + """ + first = features[0] + uniform = first.device.type == 'cpu' and all( + f.shape[1:] == first.shape[1:] and f.dtype == first.dtype + and f.device.type == 'cpu' for f in features) + if not uniform: + return torch.cat(features, dim=0).to('cuda', non_blocking=True) + + rows = sum(f.shape[0] for f in features) + staging = getattr(self, '_encoder_feature_staging', None) + if (staging is None or staging.dtype != first.dtype + or staging.shape[1:] != first.shape[1:] + or staging.shape[0] < rows): + staging = torch.empty((rows, *first.shape[1:]), + dtype=first.dtype, + pin_memory=prefer_pinned()) + self._encoder_feature_staging = staging + self._encoder_feature_staging_event = torch.cuda.Event() + # Dedicated copy stream: enqueued on the encoder stream the H2D + # would queue behind the previous encoder forward, and the next + # batch's staging reuse would host-block on that forward. + self._encoder_feature_copy_stream = torch.cuda.Stream() + else: + # The previous batch's H2D from this buffer must be complete + # before its rows are overwritten. It ran on the copy stream, + # concurrent with the previous forward, so this is ~always done. + self._encoder_feature_staging_event.synchronize() + + offset = 0 + for f in features: + staging[offset:offset + f.shape[0]].copy_(f) + offset += f.shape[0] + consumer_stream = torch.cuda.current_stream() + with torch.cuda.stream(self._encoder_feature_copy_stream): + packed = staging[:rows].to('cuda', non_blocking=True) + self._encoder_feature_staging_event.record() + consumer_stream.wait_event(self._encoder_feature_staging_event) + # The device tensor was allocated on the copy stream; mark it used by + # the consumer stream so the allocator does not recycle it early. + packed.record_stream(consumer_stream) + return packed + @nvtx_range("_prepare_tp_inputs_encoder") def _prepare_tp_inputs_encoder( self, @@ -7749,12 +7939,144 @@ def forward_encoder( raise ValueError("forward_encoder called with no encoder requests") with torch.inference_mode(): + graph_result = self._maybe_forward_encoder_graph(encoder_requests) + if graph_result is not None: + return graph_result + inputs = self._prepare_tp_inputs_encoder( encoder_requests, resource_manager=resource_manager) encoder_hidden_states = self._encoder_forward_enc_dec(inputs) return encoder_hidden_states, inputs['encoder_seq_lens'] + def _maybe_forward_encoder_graph( + self, + encoder_requests: List[LlmRequest], + ) -> Optional[Tuple[torch.Tensor, List[int]]]: + """Try to serve the encoder batch from a captured CUDA graph. + + Returns ``(encoder_hidden_states, encoder_seq_lens)`` on a graph hit + (the hidden states are CLONED from the graph's static output buffer — + the executor stores views of the result across scheduler iterations, + and a later replay of the same bucket would clobber them), or None to + fall back to the eager path. + """ + runner = self.encoder_cuda_graph_runner + if runner is None or not runner.enabled or not runner.feature_mode: + return None + + fixed = runner.config.fixed_seq_len + features: List[torch.Tensor] = [] + for request in encoder_requests: + f = request.py_encoder_input_features + if (f is None or int(request.encoder_output_len) != fixed + or f.shape[1:] != runner.config.feature_shape + or f.dtype != runner.config.feature_dtype): + return None + features.append(f) + + seq_lens = [fixed] * len(encoder_requests) + graph_inputs = { + 'seq_lens': seq_lens, + 'input_features': features, + } + with runner.pad_batch(graph_inputs, + len(encoder_requests)) as padded_inputs: + # `pad_batch` extends seq_lens to the captured bucket, and the + # metadata takes one request id per sequence. Pad slots carry no + # request; the encoder pass runs without a KV cache, so their ids + # are never looked up and only have to exist and stay distinct. + request_ids = [r.py_request_id for r in encoder_requests] + request_ids += [ + -(i + 1) for i in range( + len(padded_inputs['seq_lens']) - len(request_ids)) + ] + eager_attn_metadata = self._make_encoder_attn_metadata( + padded_inputs['seq_lens'], request_ids) + graph_attn_metadata, key = runner.maybe_get_cuda_graph( + padded_inputs, eager_attn_metadata) + if key is None: + return None + padded_inputs['attn_metadata'] = graph_attn_metadata + + if runner.needs_capture(key): + padded_batch_size, padded_num_tokens, _ = key + # Feature-mode seq_lens are constant per bucket: initialize + # the graph-resident metadata once at capture. + graph_attn_metadata.prepare_encoder_cuda_graph_replay( + [fixed] * padded_batch_size, padded_num_tokens) + runner.capture(key, self._enc_dec_encoder_graph_forward_fn, + padded_inputs) + + output = runner.replay(key, padded_inputs) + + real_tokens = fixed * len(encoder_requests) + return output[:real_tokens].clone(), seq_lens + + def _enc_dec_encoder_graph_forward_fn( + self, capture_inputs: Dict[str, Any]) -> torch.Tensor: + return self._forward_step_encoder({ + 'input_features': + capture_inputs['input_features'], + 'encoder_attn_metadata': + capture_inputs['attn_metadata'], + 'encoder_seq_lens': + capture_inputs['seq_lens'], + }) + + def _capture_enc_dec_encoder_graphs(self) -> None: + """Capture enc-dec encoder graphs for every configured batch size. + + Runs at engine warmup, and goes through the shared two-pass helper: + every shape must be warmed before any graph is captured, because a + smaller batch can select a different attention kernel and grow the + shared workspace, which would move a buffer a larger batch's graph + had already captured the address of. + """ + with torch.inference_mode(): + self._warmup_and_capture_encoder_cuda_graphs( + self._capture_feature_encoder_graphs_once) + + def _capture_feature_encoder_graphs_once(self) -> None: + """One pass over every feature encoder batch size. + + Called twice by `_warmup_and_capture_encoder_cuda_graphs`: once with + the runner in warmup-only mode, then once to capture. Largest bucket + first so the runner's graph pool high-water mark is set on the first + capture. Inputs are synthesized without LlmRequests — a zero waveform + is a valid fixed-shape feature — and no KV/cross-pool resources are + involved (the encoder step writes no KV cache). + """ + runner = self.encoder_cuda_graph_runner + fixed = runner.config.fixed_seq_len + for bs in sorted(runner.supported_batch_sizes, reverse=True): + seq_lens = [fixed] * bs + features = [ + torch.zeros((1, *runner.config.feature_shape), + dtype=runner.config.feature_dtype) + for _ in range(bs) + ] + graph_inputs = { + 'seq_lens': seq_lens, + 'input_features': features, + } + eager_md = self._make_encoder_attn_metadata(seq_lens, + list(range(bs))) + graph_md, key = runner.maybe_get_cuda_graph(graph_inputs, eager_md) + if key is None: + logger.warning( + "Enc-dec encoder CUDA graph capture skipped for " + f"batch size {bs} (unsupported metadata/backend).") + continue + if not runner.needs_capture(key): + continue + logger.info( + f"Capturing enc-dec encoder CUDA graph for batch size {bs}.") + graph_inputs['attn_metadata'] = graph_md + graph_md.prepare_encoder_cuda_graph_replay(seq_lens, key[1]) + runner.capture(key, self._enc_dec_encoder_graph_forward_fn, + graph_inputs) + def _init_userbuffers(self, hidden_size): if self.mapping.tp_size <= 1 or self.mapping.pp_size > 1: return False diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 62b7c6488ee9..006c89606089 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -41,7 +41,8 @@ ReqIdsSet) from tensorrt_llm.executor.request import TruncateKVCacheRequest from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation -from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy +from tensorrt_llm.llmapi.llm_args import (FeatureEncoderCudaGraphConfig, + PeftCacheConfig, WaitingQueuePolicy) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError @@ -5384,14 +5385,30 @@ def _waiting_encoder_requests( encoder_max_batch_size = self.llm_args.encoder_max_batch_size encoder_cuda_graph_config = self.llm_args.encoder_cuda_graph_config + # A feature encoder has no token/seq-len buckets to gate on. + is_feature_encoder_config = isinstance(encoder_cuda_graph_config, + FeatureEncoderCudaGraphConfig) if (encoder_max_batch_size is not None and encoder_cuda_graph_config is not None - and bool(encoder_cuda_graph_config.num_tokens) - and bool(encoder_cuda_graph_config.seq_lens)): + and (is_feature_encoder_config or + (bool(encoder_cuda_graph_config.num_tokens) + and bool(encoder_cuda_graph_config.seq_lens)))): encoder_batch_size_limit = min(encoder_max_batch_size, self.max_batch_size) - configured_batch_sizes = (encoder_cuda_graph_config.batch_sizes - or []) + if is_feature_encoder_config: + # Feature batch sizes may have been derived rather than + # configured, so take the ones the runner actually resolved. + # They are populated from the config even when capture was + # declined (TP > 1, no bucket fits), so waiting on them would + # delay a batch that can only ever run eager. + runner = getattr(self.model_engine, 'encoder_cuda_graph_runner', + None) + configured_batch_sizes = (list(runner.supported_batch_sizes) + if runner is not None + and runner.enabled else []) + else: + configured_batch_sizes = (encoder_cuda_graph_config.batch_sizes + or []) supported_batch_sizes = [ batch_size for batch_size in configured_batch_sizes if batch_size <= encoder_batch_size_limit From 7618a7134bdfab7225a365125cb9f8148ee9498e Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:40:21 -0700 Subject: [PATCH 3/8] [TRTLLM-14778][test] Cover the Whisper encoder CUDA-graph path Covers capture and replay across the configured encoder batch sizes, the eager fallback for an uncaptured size, the config/model mismatch branches, and encoder microbatch admission with the feature config enabled and declined. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../llmapi/test_llm_api_pytorch_whisper.py | 93 +++++++++++--- .../test_lists/test-db/l0_h100.yml | 1 + .../test_lists/test-db/l0_l40s.yml | 1 + .../_torch/executor/test_py_executor.py | 76 +++++++++++ .../executor/test_pytorch_model_engine.py | 119 +++++++++++++++++- .../test_pytorch_model_engine_warmup.py | 1 + tests/unittest/llmapi/test_llm_args.py | 60 +++++++++ 7 files changed, 329 insertions(+), 22 deletions(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py index bb72ccb7121a..5c9e82f65c5d 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py @@ -28,7 +28,14 @@ import pytest import soundfile -from tensorrt_llm.llmapi import LLM, CudaGraphConfig, KvCacheConfig, SamplingParams, SchedulerConfig +from tensorrt_llm.llmapi import ( + LLM, + CudaGraphConfig, + FeatureEncoderCudaGraphConfig, + KvCacheConfig, + SamplingParams, + SchedulerConfig, +) from ..conftest import llm_models_root @@ -36,6 +43,9 @@ _MIN_GPU_MEMORY_MB = 16_000 _FREE_GPU_MEMORY_FRACTION = 0.2 _CROSS_KV_CACHE_FRACTION = 0.5 +# Every Whisper request produces this many encoder positions, whatever the audio +# length. It sets the cross-KV pool capacity and the encoder graph shapes. +_ENCODER_OUTPUT_LEN = 1500 # whisper-tiny fp32 greedy on 1221-135766-0002.wav (matches HF transformers). _EXPECTED_GREEDY_OUTPUT_TOKEN_IDS = [ 1939, @@ -104,15 +114,32 @@ def _make_llm( torch_dtype: str | None = None, cuda_graph_batch_sizes: list[int] | None = None, tensor_parallel_size: int = 1, + encoder_graphs: bool = False, ) -> LLM: - # CudaGraphConfig captures the decode step only; fp32 enc-dec declines + # CudaGraphConfig captures the decode step; the enc-dec encoder step opts in + # separately through `encoder_cuda_graph_config`. fp32 enc-dec declines # graphs at engine init (workspace-sizing guard), so requesting them must # still work for every dtype. cuda_graph_config = ( - CudaGraphConfig(batch_sizes=cuda_graph_batch_sizes, enable_padding=True) + CudaGraphConfig( + batch_sizes=cuda_graph_batch_sizes, + enable_padding=True, + ) if cuda_graph_batch_sizes is not None else None ) + encoder_kwargs = {} + if encoder_graphs: + # Whisper's encoder emits a fixed `_ENCODER_OUTPUT_LEN` positions per + # request, so the graph key is the batch size alone. + encoder_batch_sizes = list(cuda_graph_batch_sizes or [1]) + encoder_kwargs = { + "encoder_max_batch_size": max(encoder_batch_sizes), + "encoder_cuda_graph_config": FeatureEncoderCudaGraphConfig( + batch_sizes=encoder_batch_sizes, + enable_padding=True, + ), + } dtype_kwargs = {} if torch_dtype is not None: # The checkpoint's torch_dtype wins over `dtype` in the PyTorch @@ -134,10 +161,11 @@ def _make_llm( max_beam_width=max_beam_width, # Cross-KV pool capacity; the default (1024) is smaller than the # 1500 encoder positions every Whisper request produces. - max_input_len=1500, - max_num_tokens=3000, + max_input_len=_ENCODER_OUTPUT_LEN, + max_num_tokens=2 * _ENCODER_OUTPUT_LEN, scheduler_config=SchedulerConfig(use_python_scheduler=True), tensor_parallel_size=tensor_parallel_size, + **encoder_kwargs, **dtype_kwargs, ) @@ -240,38 +268,60 @@ def test_whisper_pytorch_beam_search( ) outputs = llm.generate([_audio_prompt(wave, sample_rate)], beam_params) assert _EXPECTED_TRANSCRIPT_FRAGMENT in outputs[0].outputs[0].text.lower() - _assert_decoder_cuda_graph_state(llm, captured=graphs_captured) + _assert_cuda_graph_state(llm, captured=graphs_captured) -def _assert_decoder_cuda_graph_state(llm: LLM, captured: bool) -> None: +def _assert_cuda_graph_state(llm: LLM, captured: bool, encoder_captured: bool = False) -> None: """Introspect the in-process engine (single-process mode only). - Decoder graphs captured (or not); the enc-dec encoder step stays eager. + The enc-dec encoder step shares `encoder_cuda_graph_runner` with the + `llm.encode()` path; feature mode is a mode of that one runner, selected by + `encoder_cuda_graph_config`, not a second runner. """ model_engine = llm._executor.engine.model_engine - assert not model_engine.encoder_cuda_graph_runner.enabled - assert not model_engine.encoder_cuda_graph_runner.graphs assert model_engine.cuda_graph_runner.enabled == captured assert bool(model_engine.cuda_graph_runner.graphs) == captured + encoder_runner = model_engine.encoder_cuda_graph_runner + if not encoder_captured: + assert not encoder_runner.enabled + assert not encoder_runner.graphs + return + # Capture must actually have happened: a silent fallback to the eager + # encoder path would otherwise pass every output assertion above. + assert encoder_runner.enabled + assert encoder_runner.graphs + assert encoder_runner.feature_mode + assert encoder_runner.is_encoder_decoder + # Feature-combination matrix mirroring the T5/BART enc-dec coverage. Cases: # (torch_dtype override or None for checkpoint fp32, kv manager v2, decoder -# cuda-graph batch sizes, graphs must capture, TP size). KVCacheManagerV2 -# requires beam width 1, so v2 rides greedy; the fp32+graphs-requested case -# asserts the engine declines graphs (fp32 enc-dec guard) yet stays exact. +# cuda-graph batch sizes, graphs must capture, TP size, encoder graphs). +# KVCacheManagerV2 requires beam width 1, so v2 rides greedy; the +# fp32+graphs-requested case asserts the engine declines graphs (fp32 enc-dec +# guard) yet stays exact. The encoder-graphs case additionally captures the +# encoder step, which must not change a single token. _FEATURE_COMBINATION_CASES = [ - pytest.param(None, True, None, False, 1, id="fp32-kv-v2-graphs-off-greedy"), - pytest.param(None, False, [1, 2], False, 1, id="fp32-kv-v1-graphs-requested-greedy"), - pytest.param("bfloat16", False, [1, 2], True, 1, id="bf16-kv-v1-decoder-graphs-on-greedy"), - pytest.param("bfloat16", True, [1, 2], True, 1, id="bf16-kv-v2-decoder-graphs-on-greedy"), - pytest.param("float16", False, None, False, 1, id="fp16-kv-v1-graphs-off-greedy"), + pytest.param(None, True, None, False, 1, False, id="fp32-kv-v2-graphs-off-greedy"), + pytest.param(None, False, [1, 2], False, 1, False, id="fp32-kv-v1-graphs-requested-greedy"), + pytest.param( + "bfloat16", False, [1, 2], True, 1, False, id="bf16-kv-v1-decoder-graphs-on-greedy" + ), + pytest.param( + "bfloat16", True, [1, 2], True, 1, False, id="bf16-kv-v2-decoder-graphs-on-greedy" + ), + pytest.param( + "bfloat16", False, [1, 2], True, 1, True, id="bf16-kv-v1-encoder-graphs-on-greedy" + ), + pytest.param("float16", False, None, False, 1, False, id="fp16-kv-v1-graphs-off-greedy"), pytest.param( None, False, None, False, 2, + False, id="fp32-kv-v1-graphs-off-greedy-tp2", marks=pytest.mark.skip_less_device(2), ), @@ -279,7 +329,8 @@ def _assert_decoder_cuda_graph_state(llm: LLM, captured: bool) -> None: @pytest.mark.parametrize( - "torch_dtype,use_kv_cache_manager_v2,cuda_graph_batch_sizes,graphs_captured,tp_size", + "torch_dtype,use_kv_cache_manager_v2,cuda_graph_batch_sizes,graphs_captured,tp_size," + "encoder_graphs", _FEATURE_COMBINATION_CASES, ) def test_whisper_pytorch_feature_combinations( @@ -289,6 +340,7 @@ def test_whisper_pytorch_feature_combinations( cuda_graph_batch_sizes, graphs_captured, tp_size, + encoder_graphs, ): """Greedy transcription across dtype/kv-cache-manager/CUDA-graph/TP combos. @@ -308,6 +360,7 @@ def test_whisper_pytorch_feature_combinations( torch_dtype=torch_dtype, cuda_graph_batch_sizes=cuda_graph_batch_sizes, tensor_parallel_size=tp_size, + encoder_graphs=encoder_graphs, ) with llm: for batch_size in (1, 2): @@ -326,4 +379,4 @@ def test_whisper_pytorch_feature_combinations( assert _EXPECTED_TRANSCRIPT_FRAGMENT in completion.text.lower() if tp_size == 1: - _assert_decoder_cuda_graph_state(llm, captured=graphs_captured) + _assert_cuda_graph_state(llm, captured=graphs_captured, encoder_captured=encoder_graphs) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 88130605c153..4ef714fdaf8e 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -346,6 +346,7 @@ l0_h100: - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-requested-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-decoder-graphs-on-greedy] + - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp16-kv-v1-graphs-off-greedy] - examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b] - unittest/bindings # 8 mins on H100 diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index ff242644533f..854439219d7c 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -76,6 +76,7 @@ l0_l40s: - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-requested-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-decoder-graphs-on-greedy] + - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp16-kv-v1-graphs-off-greedy] - condition: ranges: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index fa8441426dbf..5f708387b6ac 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -38,6 +38,7 @@ ScheduledRequests, SerializableSchedulerOutput, ) +from tensorrt_llm.llmapi.llm_args import FeatureEncoderCudaGraphConfig pytestmark = pytest.mark.cpu_only @@ -169,6 +170,31 @@ def _make_encoder_batch_wait_executor(batch_sizes=None, encoder_max_batch_size=8 return executor +def _make_feature_encoder_batch_wait_executor( + runner_batch_sizes, encoder_max_batch_size=8, runner_enabled=True +): + """Batch-wait executor whose encoder graph config is the feature variant. + + `FeatureEncoderCudaGraphConfig` has no `num_tokens` / `seq_lens`, and its + `batch_sizes` may have been derived rather than configured, so the resolved + sizes come from the engine's encoder graph runner rather than the config. + """ + executor = object.__new__(PyExecutor) + executor.max_batch_size = 32 + executor.llm_args = types.SimpleNamespace( + encoder_cuda_graph_config=FeatureEncoderCudaGraphConfig(enable_padding=True), + encoder_max_batch_size=encoder_max_batch_size, + ) + executor.model_engine = types.SimpleNamespace( + encoder_cuda_graph_runner=types.SimpleNamespace( + supported_batch_sizes=runner_batch_sizes, enabled=runner_enabled + ) + ) + executor.batch_wait_timeout_iters = 48 + executor.encoder_batch_wait_iters_count = 0 + return executor + + def _make_encoder_fallback_batch_wait_executor(): executor = object.__new__(PyExecutor) executor.llm_args = types.SimpleNamespace( @@ -213,6 +239,56 @@ def test_encoder_graph_warmup_uses_runtime_encoder_stream(): ) +def test_encoder_microbatch_admission_supports_feature_encoder_config(): + # A feature config carries no num_tokens / seq_lens, so reading them to + # gate this path raises AttributeError on the first encoder batch. + executor = _make_feature_encoder_batch_wait_executor([1, 2, 4, 8]) + encoder_requests = [object() for _ in range(12)] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object()] * 20, + ) + + assert scheduled == encoder_requests[:8] + assert executor.encoder_batch_wait_iters_count == 0 + + +def test_encoder_microbatch_admission_uses_derived_feature_batch_sizes(): + # batch_sizes left unset on the config: the engine derived them, so the + # runner is the only place the resolved list exists. + executor = _make_feature_encoder_batch_wait_executor([1, 2, 3, 4]) + assert executor.llm_args.encoder_cuda_graph_config.batch_sizes is None + encoder_requests = [object() for _ in range(6)] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object()] * 20, + ) + + assert scheduled == encoder_requests[:4] + + +def test_encoder_microbatch_admission_ignores_disabled_feature_runner(): + # supported_batch_sizes stays populated from the config even when capture + # was declined (TP > 1, or no bucket fits), so waiting on those shapes + # would stall a batch that can only ever run eager. With no decoder work + # the request must be released immediately instead. + executor = _make_feature_encoder_batch_wait_executor([1, 2, 4, 8], runner_enabled=False) + executor.batch_wait_max_tokens_ratio = 0.5 + executor.max_num_tokens = 32 + executor.active_requests = [] + executor.inflight_req_ids = _InflightRequestIds() + encoder_requests = [_make_encoder_request(0)] + + scheduled = executor._waiting_encoder_requests(encoder_requests, [], []) + + assert scheduled == encoder_requests + assert executor.encoder_batch_wait_iters_count == 0 + + def test_encoder_microbatch_graph_admission_boundaries(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object()] * 7 diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 03d76b437123..04f89d7ed270 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -18,13 +18,16 @@ from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - CUDAGraphRunner, EncoderCUDAGraphRunner, KeyType, - _restore_spec_decode_capture_state, _save_spec_decode_capture_state) + CUDAGraphRunner, EncoderCUDAGraphRunner, EncoderCUDAGraphRunnerConfig, + KeyType, _restore_spec_decode_capture_state, + _save_spec_decode_capture_state) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import ( PyTorchModelEngine, _build_request_multimodal_input, _make_single_token_context_graph_batch) from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + EncodeCudaGraphConfig, + FeatureEncoderCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchLlmArgs) @@ -1019,6 +1022,118 @@ def test_global_incompatibilities_bypass_candidate_selection(self) -> None: class PyTorchModelEngineTestCase(unittest.TestCase): + @staticmethod + def _feature_encoder_runner(batch_sizes, fixed_seq_len=1500): + """A feature-mode runner with capture disabled, so no CUDA is touched.""" + config = EncoderCUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=True, + cuda_graph_batch_sizes=batch_sizes, + cuda_graph_num_tokens=[], + cuda_graph_seq_lens=[], + max_cuda_graph_batch_size=max(batch_sizes), + max_cuda_graph_num_tokens=max(batch_sizes) * fixed_seq_len, + max_num_tokens=max(batch_sizes) * fixed_seq_len, + max_seq_len=fixed_seq_len, + cuda_graph_mem_pool=None, + is_encoder_decoder=True, + use_fixed_sequence_slots=False, + feature_shape=(480000, ), + feature_dtype=torch.float32, + fixed_seq_len=fixed_seq_len, + ) + return EncoderCUDAGraphRunner(config) + + def test_feature_encoder_capture_keys_are_all_reachable(self) -> None: + # Every request contributes exactly fixed_seq_len positions, so the + # only reachable key per batch size is (bs, bs * fixed, fixed). The + # token path's cross product would also emit keys whose token count no + # batch can produce, and capture_keys drives mixed encoder/decoder + # decoder-graph warmup. + fixed = 1500 + batch_sizes = [1, 2, 4, 8] + runner = self._feature_encoder_runner(batch_sizes, fixed) + + self.assertEqual( + sorted(runner.capture_keys), + [(bs, bs * fixed, fixed) for bs in batch_sizes], + ) + + def test_feature_encoder_capture_layout_is_uniform(self) -> None: + fixed = 1500 + runner = self._feature_encoder_runner([2], fixed) + self.assertEqual( + runner._capture_sequence_lengths[(2, 2 * fixed, fixed)], + [fixed, fixed]) + + def test_derived_feature_encoder_batch_sizes_are_dense_below_eight( + self) -> None: + # The padding guard refuses a graph once padding would add more than + # MAX_FEATURE_PADDING_RATIO of encoder work, so sparse low buckets + # would leave batch sizes 5-7 permanently eager. + derived = PyTorchModelEngine._derive_feature_encoder_batch_sizes(32) + self.assertEqual(derived, [1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32]) + + def test_derived_feature_encoder_batch_sizes_respect_cap(self) -> None: + self.assertEqual( + PyTorchModelEngine._derive_feature_encoder_batch_sizes(6), + [1, 2, 3, 4, 5, 6]) + self.assertEqual( + PyTorchModelEngine._derive_feature_encoder_batch_sizes(1), [1]) + + @staticmethod + def _encoder_spec_engine(encoder_cuda_graph_config, + declares_spec: bool, + tp_size: int = 1): + """A bare engine carrying only what `_encoder_graph_spec` reads.""" + spec = ((480000, ), torch.float32, 1500) + + class _Model: + model_config = SimpleNamespace(is_encoder_decoder=True) + + if declares_spec: + + def encoder_graph_spec(self): + return spec + + engine = PyTorchModelEngine.__new__(PyTorchModelEngine) + engine.encoder_cuda_graph_config = encoder_cuda_graph_config + engine.is_draft_model = False + engine.model = _Model() + engine.mapping = SimpleNamespace(tp_size=tp_size) + return engine, spec + + def test_encoder_graph_spec_returns_spec_for_feature_config(self) -> None: + engine, spec = self._encoder_spec_engine( + FeatureEncoderCudaGraphConfig(batch_sizes=[1, 2]), + declares_spec=True) + self.assertEqual(engine._encoder_graph_spec(), spec) + + def test_encoder_graph_spec_rejects_feature_config_on_token_model( + self) -> None: + # A token encoder cannot satisfy the fixed-shape contract, and a bare + # `batch_sizes` config is read as feature mode, so the message has to + # point at both possibilities. + engine, _ = self._encoder_spec_engine( + FeatureEncoderCudaGraphConfig(batch_sizes=[1]), declares_spec=False) + with self.assertRaises(ValueError) as ctx: + engine._encoder_graph_spec() + self.assertIn("EncodeCudaGraphConfig", str(ctx.exception)) + + def test_encoder_graph_spec_declines_token_config_on_feature_model( + self) -> None: + engine, _ = self._encoder_spec_engine(EncodeCudaGraphConfig( + batch_sizes=[1], num_tokens=[1500], seq_lens=[1500]), + declares_spec=True) + self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) + + def test_encoder_graph_spec_declines_tensor_parallel(self) -> None: + engine, _ = self._encoder_spec_engine( + FeatureEncoderCudaGraphConfig(batch_sizes=[1]), + declares_spec=True, + tp_size=2) + self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) + def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( self) -> None: runner = EncoderCUDAGraphRunner.__new__(EncoderCUDAGraphRunner) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 8d95fb9f1a53..e3fff9008be2 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -195,6 +195,7 @@ def allow_capture(): enabled=True, is_encoder_decoder=True, is_warmup_only=False, + feature_mode=False, allow_capture=allow_capture, ) model_engine.encoder_cuda_graph_runner = runner diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2fa8772fd073..f288a0f54455 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -43,6 +43,7 @@ EncodeCudaGraphConfig, ExecutorMemoryType, ExtendedRuntimePerfKnobConfig, + FeatureEncoderCudaGraphConfig, KvCacheConfig, LookaheadDecodingConfig, MambaStateConfig, MoeConfig, @@ -319,6 +320,65 @@ def test_rejects_non_positive(self, llm_args_cls, field_name, llm_args_cls(model=llama_model_path, **{field_name: invalid_value}) +@pytest.mark.cpu_only +class TestFeatureEncoderCudaGraphConfig: + """Cover the fixed-shape feature encoder branch of encoder_cuda_graph_config. + + A feature encoder emits the same number of positions for every request, so + `batch_sizes` is its only bucket dimension. The token-shaped + `EncodeCudaGraphConfig` fields must not be expressible here, or a caller + could supply values the encoder graph runner silently overrides. + """ + + @pytest.mark.parametrize("batch_sizes", [[], [0], [-1], [1, -2]]) + def test_rejects_non_positive_or_empty_batch_sizes(self, batch_sizes): + with pytest.raises(ValidationError): + FeatureEncoderCudaGraphConfig(batch_sizes=batch_sizes) + + def test_sorts_and_deduplicates_batch_sizes(self): + cfg = FeatureEncoderCudaGraphConfig(batch_sizes=[4, 2, 2, 1]) + assert cfg.batch_sizes == [1, 2, 4] + + @pytest.mark.parametrize("field", ["num_tokens", "seq_lens"]) + def test_rejects_token_shaped_buckets(self, field): + # These belong to EncodeCudaGraphConfig; accepting them here would be + # accepting a value the runner discards. + with pytest.raises(ValidationError): + FeatureEncoderCudaGraphConfig(batch_sizes=[1], **{field: [64]}) + + @pytest.mark.parametrize( + "config_dict, expected_type", + [ + (dict(batch_sizes=[1, 2]), FeatureEncoderCudaGraphConfig), + (dict(batch_sizes=[1], seq_lens=[64], + num_tokens=[64]), EncodeCudaGraphConfig), + (dict(mode="feature_encode", + batch_sizes=[1]), FeatureEncoderCudaGraphConfig), + ], + ids=["bare_batch_sizes", "token_buckets", "explicit_mode"], + ) + def test_encoder_config_mode_is_inferred(self, config_dict, expected_type): + llm_args = TorchLlmArgs(model=llama_model_path, + encoder_max_batch_size=8, + encoder_cuda_graph_config=config_dict) + assert isinstance(llm_args.encoder_cuda_graph_config, expected_type) + + def test_feature_config_does_not_require_token_buckets(self): + llm_args = TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=8, + encoder_cuda_graph_config=FeatureEncoderCudaGraphConfig( + batch_sizes=[1, 2])) + assert llm_args.encoder_cuda_graph_config.batch_sizes == [1, 2] + + def test_token_config_still_requires_token_buckets(self): + with pytest.raises(ValueError): + TorchLlmArgs(model=llama_model_path, + encoder_max_batch_size=8, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 2])) + + @pytest.mark.cpu_only def test_decoding_type_eagle3_parses_to_eagle3_decoding_config(): adapter = TypeAdapter(SpeculativeConfig) From 9f8cf21b889ce8d925824cc3d38f8177ba37125e Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:20:20 -0700 Subject: [PATCH 4/8] [TRTLLM-14778][fix] Select feature-mode encoder graphs from the model, not a new config Encoder-graph capture keyed on batch size alone applies to an encoder whose input is a fixed-shape per-request feature tensor, which is a property of the model rather than a choice the caller makes. Detect it from encoder_graph_spec() and drop the separate config type, so encoder_cuda_graph_config keeps its existing shape and the token buckets a feature encoder derives become optional. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 90 ++++++------------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 23 ++--- tensorrt_llm/llmapi/__init__.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 85 ++---------------- .../usage/llm_args_golden_manifest.json | 3 +- .../llmapi/test_llm_api_pytorch_whisper.py | 7 +- .../_torch/executor/test_py_executor.py | 33 ++++--- .../executor/test_pytorch_model_engine.py | 45 +++------- .../api_stability/references/llm.yaml | 2 +- tests/unittest/llmapi/test_llm_args.py | 77 ++++------------ 10 files changed, 102 insertions(+), 267 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 159b954198ac..f6b7d045e381 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -35,7 +35,6 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, - FeatureEncoderCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -558,19 +557,21 @@ def __init__( encoder_cuda_graph_batch_sizes = ( self.encoder_cuda_graph_config.batch_sizes if self.encoder_cuda_graph_config is not None else []) - # A feature encoder config carries batch sizes only: its token counts - # and sequence lengths follow from the model's fixed encoder output - # length, which the encoder graph runner derives. - encoder_cuda_graph_num_tokens = getattr(self.encoder_cuda_graph_config, - 'num_tokens', None) or [] - encoder_cuda_graph_seq_lens = getattr(self.encoder_cuda_graph_config, - 'seq_lens', None) or [] + encoder_cuda_graph_num_tokens = ( + self.encoder_cuda_graph_config.num_tokens + if self.encoder_cuda_graph_config is not None else []) + encoder_cuda_graph_seq_lens = (self.encoder_cuda_graph_config.seq_lens + if self.encoder_cuda_graph_config + is not None else []) encoder_cuda_graph_padding_enabled = ( self.encoder_cuda_graph_config.enable_padding if self.encoder_cuda_graph_config is not None else False) + # A fixed-shape feature encoder derives both bucket lists from the + # model's encoder output length, so only a token-driven encoder needs + # the user to supply them. if (self.encoder_cuda_graph_config is not None - and not self._is_feature_encoder_cuda_graph_config() + and self._model_encoder_graph_spec() is None and (not encoder_cuda_graph_num_tokens or not encoder_cuda_graph_seq_lens)): missing = [] @@ -624,7 +625,7 @@ def __init__( use_encoder_cuda_graph = ( (self._is_encoder_decoder_model() or self._is_encode_only) and self.encoder_cuda_graph_config is not None - and (self._is_feature_encoder_cuda_graph_config() or + and (self._model_encoder_graph_spec() is not None or (bool(self._cuda_graph_num_tokens) and bool(self._cuda_graph_seq_lens)))) @@ -872,21 +873,15 @@ def __init__( encoder_graph_max_num_tokens = self._max_cuda_graph_num_tokens # A feature-driven encoder (Whisper) declares a fixed-shape per-request - # contract instead of packed tokens, so its graph shapes follow from - # the batch sizes alone. Enablement still goes through the same - # `encoder_cuda_graph_config` opt-in as the token path. + # contract instead of packed tokens, so its graph shapes follow from the + # batch sizes alone and `num_tokens` / `seq_lens` are ignored if set. + # Enablement is the same `encoder_cuda_graph_config` opt-in either way. feature_shape, feature_dtype, fixed_seq_len = self._encoder_graph_spec() if feature_shape is not None: bs_cap = max( 1, min(self.batch_size, self.encoder_max_num_tokens // fixed_seq_len)) - if not encoder_graph_batch_sizes: - # batch_sizes left unset: derive them, bounded by the encoder - # batch limit and the scheduler's encoder-batch bound. - encoder_graph_batch_sizes = ( - self._derive_feature_encoder_batch_sizes( - min(self.encoder_batch_size, bs_cap))) encoder_graph_batch_sizes = sorted( bs for bs in encoder_graph_batch_sizes if bs <= bs_cap) if not encoder_graph_batch_sizes: @@ -901,15 +896,14 @@ def __init__( fixed_seq_len) elif (use_encoder_cuda_graph and self._model_encoder_graph_spec() is not None): - # A feature encoder cannot consume the packed token inputs the - # token-shaped capture path synthesizes, so keep it eager rather - # than let warmup drive tokens into it. + # Feature mode was declined above (TP > 1). This model's encoder + # cannot consume the packed token inputs the token-shaped capture + # path synthesizes, so keep it eager rather than let warmup drive + # tokens into it. logger.warning( - "This model's encoder consumes fixed-shape features, but " - "encoder_cuda_graph_config is an EncodeCudaGraphConfig, whose " - "token buckets do not apply; the encoder step stays eager. " - "Pass FeatureEncoderCudaGraphConfig(batch_sizes=[...]) to " - "capture it.") + "This model's encoder consumes fixed-shape features and " + "feature-mode encoder CUDA graphs are unavailable; the encoder " + "step stays eager.") use_encoder_cuda_graph = False encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( @@ -3704,11 +3698,6 @@ def _is_encoder_decoder_model(self) -> bool: getattr(getattr(self.model, "model_config", None), "is_encoder_decoder", False)) - def _is_feature_encoder_cuda_graph_config(self) -> bool: - """True when the encoder graph config selects the feature-shaped path.""" - return isinstance(self.encoder_cuda_graph_config, - FeatureEncoderCudaGraphConfig) - def _model_encoder_graph_spec(self): """The model's fixed-shape encoder contract, or None. Queried once.""" if not hasattr(self, "_cached_model_encoder_graph_spec"): @@ -3719,28 +3708,15 @@ def _model_encoder_graph_spec(self): is not None else None) return self._cached_model_encoder_graph_spec - @staticmethod - def _derive_feature_encoder_batch_sizes(max_batch_size: int) -> List[int]: - """Encoder graph batch sizes for a fixed-shape feature encoder. - - Dense up to 8, then multiples of 8. A feature pad slot costs a full - encoder forward, so padding is refused once it would add more than - ``MAX_FEATURE_PADDING_RATIO`` of encoder work; sparse buckets at the - low end would leave small batches permanently eager. - """ - sizes = set(range(1, min(max_batch_size, 8) + 1)) - sizes.update(range(8, max_batch_size + 1, 8)) - sizes.add(max_batch_size) - return sorted(bs for bs in sizes if 0 < bs <= max_batch_size) - def _encoder_graph_spec(self): """Fixed-shape encoder contract, or (None, None, None) if unavailable. - Returns ``(feature_shape, feature_dtype, fixed_seq_len)`` when the user - selected ``FeatureEncoderCudaGraphConfig``, the model declares - ``encoder_graph_spec()`` and feature-mode encoder CUDA graphs are - viable. Gated to TP=1 (allreduce inside encoder capture is unverified) - and to non-draft models. + Returns ``(feature_shape, feature_dtype, fixed_seq_len)`` when the model + declares ``encoder_graph_spec()`` and feature-mode encoder CUDA graphs + are viable. The model selects the mode, not the config: an encoder + either takes fixed-shape features or it does not. Gated to TP=1 + (allreduce inside encoder capture is unverified) and to non-draft + models. """ none = (None, None, None) if (self.encoder_cuda_graph_config is None or self.is_draft_model @@ -3748,18 +3724,8 @@ def _encoder_graph_spec(self): return none spec = self._model_encoder_graph_spec() - - if not self._is_feature_encoder_cuda_graph_config(): - return none - if spec is None: - raise ValueError( - "FeatureEncoderCudaGraphConfig requires a model whose encoder " - "declares a fixed-shape encoder_graph_spec(); this model does " - "not. Token-driven encoders such as T5 and BART take " - "EncodeCudaGraphConfig with num_tokens and seq_lens set. Note " - "that an encoder_cuda_graph_config supplied without those two " - "fields is read as a feature-encoder config.") + return none if self.mapping.tp_size > 1: logger.warning( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 006c89606089..f97a097e008a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -41,8 +41,7 @@ ReqIdsSet) from tensorrt_llm.executor.request import TruncateKVCacheRequest from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation -from tensorrt_llm.llmapi.llm_args import (FeatureEncoderCudaGraphConfig, - PeftCacheConfig, WaitingQueuePolicy) +from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError @@ -5385,27 +5384,23 @@ def _waiting_encoder_requests( encoder_max_batch_size = self.llm_args.encoder_max_batch_size encoder_cuda_graph_config = self.llm_args.encoder_cuda_graph_config - # A feature encoder has no token/seq-len buckets to gate on. - is_feature_encoder_config = isinstance(encoder_cuda_graph_config, - FeatureEncoderCudaGraphConfig) + # A fixed-shape feature encoder has no token/seq-len buckets to gate on. + runner = getattr(self.model_engine, 'encoder_cuda_graph_runner', None) + is_feature_encoder = bool(getattr(runner, 'feature_mode', False)) if (encoder_max_batch_size is not None and encoder_cuda_graph_config is not None - and (is_feature_encoder_config or + and (is_feature_encoder or (bool(encoder_cuda_graph_config.num_tokens) and bool(encoder_cuda_graph_config.seq_lens)))): encoder_batch_size_limit = min(encoder_max_batch_size, self.max_batch_size) - if is_feature_encoder_config: + if is_feature_encoder: # Feature batch sizes may have been derived rather than # configured, so take the ones the runner actually resolved. - # They are populated from the config even when capture was - # declined (TP > 1, no bucket fits), so waiting on them would - # delay a batch that can only ever run eager. - runner = getattr(self.model_engine, 'encoder_cuda_graph_runner', - None) + # They stay populated even when capture was declined, so + # waiting on them would delay a batch that can only run eager. configured_batch_sizes = (list(runner.supported_batch_sizes) - if runner is not None - and runner.enabled else []) + if runner.enabled else []) else: configured_batch_sizes = (encoder_cuda_graph_config.batch_sizes or []) diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index d707f397e265..2ddc301eb01a 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -15,8 +15,7 @@ DraftTargetDecodingConfig, DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - ExtendedRuntimePerfKnobConfig, - FeatureEncoderCudaGraphConfig, KvCacheConfig, LlmArgs, + ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, LookaheadDecodingConfig, MambaStateConfig, MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, @@ -52,7 +51,6 @@ 'CudaGraphConfig', 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', - 'FeatureEncoderCudaGraphConfig', 'MoeConfig', 'LookaheadDecodingConfig', 'MedusaDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index bb6b4b91734a..64be4738f0b0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -446,47 +446,6 @@ def _generate_cuda_graph_seq_lens(max_seq_len: int, return sizes -class FeatureEncoderCudaGraphConfig(StrictBaseModel): - """CUDA graph configuration for fixed-shape feature encoder requests. - - Applies to encoder-decoder models whose encoder consumes a fixed-shape - per-request feature tensor rather than packed tokens, e.g. Whisper's - 30 s-padded waveform. Such an encoder emits the same number of positions - for every request, so the graph key is the batch size alone and the - token-shaped `num_tokens` / `seq_lens` buckets of - :class:`EncodeCudaGraphConfig` do not apply. - """ - - mode: Literal["feature_encode"] = Field( - default="feature_encode", description="CUDA graph configuration mode.") - - batch_sizes: Optional[List[PositiveInt]] = Field( - default=None, - min_length=1, - description=( - "Encoder batch sizes to capture. None derives them from " - "`encoder_max_batch_size`, capped by the scheduler's encoder-batch " - "bound (max_num_tokens // encoder output length)."), - status="prototype", - ) - - enable_padding: bool = Field( - default=True, - description=( - "Pad an encoder batch up to the next captured batch size. Each pad " - "slot costs a full encoder forward, so padding is skipped when it " - "would add disproportionate encoder work."), - status="prototype", - ) - - @model_validator(mode='after') - def validate_feature_encoder_cuda_graph_config( - self) -> 'FeatureEncoderCudaGraphConfig': - if self.batch_sizes is not None: - self.batch_sizes = sorted(set(self.batch_sizes)) - return self - - # For CudaGraphConfig's backward compatibility CudaGraphConfig = DecodeCudaGraphConfig @@ -495,11 +454,6 @@ def validate_feature_encoder_cuda_graph_config( Field(discriminator="mode"), ] -EncoderCudaGraphConfigType: TypeAlias = Annotated[ - Union[EncodeCudaGraphConfig, FeatureEncoderCudaGraphConfig], - Field(discriminator="mode"), -] - class MultimodalEncoderCudaGraphConfig(StrictBaseModel): """CUDA graph capture for multimodal vision / audio encoders. @@ -5129,16 +5083,13 @@ class TorchLlmArgs(BaseLlmArgs): Note that each CUDA graph can use up to 200 MB of extra memory.", status="beta") - encoder_cuda_graph_config: Optional[EncoderCudaGraphConfigType] = Field( + encoder_cuda_graph_config: Optional[EncodeCudaGraphConfig] = Field( default=None, description=( "CUDA graph configuration for the encoder forward pass of an " "encoder-decoder model. Use `cuda_graph_config` for the decoder " - "and this field for the encoder. Pass an `EncodeCudaGraphConfig` " - "for a token encoder (T5/BART) or a " - "`FeatureEncoderCudaGraphConfig` for a fixed-shape feature encoder " - "(Whisper). Encoder CUDA graphs require `encoder_max_batch_size` " - "to be set."), + "and this field for the encoder. Encoder CUDA graphs require " + "`encoder_max_batch_size` to be set."), status="prototype") enable_encoder_decoder_mixed_cuda_graph: bool = Field( @@ -5228,19 +5179,6 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: raise ValueError("must be a positive integer when set") return v - @field_validator('encoder_cuda_graph_config', mode='before') - @classmethod - def infer_encoder_cuda_graph_config_mode(cls, v): - if isinstance(v, dict) and "mode" not in v: - token_keys = { - "num_tokens", "max_num_token", "seq_lens", "max_seq_len" - } - v = dict(v) - v["mode"] = "encode" if any( - k in v and v[k] not in (None, 0) - for k in token_keys) else "feature_encode" - return v - @model_validator(mode="after") def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': if self.encoder_cuda_graph_config is None: @@ -5253,19 +5191,10 @@ def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': if self.encoder_max_batch_size is None: raise ValueError( "encoder_cuda_graph_config requires encoder_max_batch_size.") - if isinstance(self.encoder_cuda_graph_config, - FeatureEncoderCudaGraphConfig): - # A feature encoder's token counts and sequence lengths follow from - # the model, so batch_sizes is the only bucket dimension to require. - return self - missing = [] - if not self.encoder_cuda_graph_config.num_tokens: - missing.append("num_tokens/max_num_token") - if not self.encoder_cuda_graph_config.seq_lens: - missing.append("seq_lens/max_seq_len") - if missing: - raise ValueError("encoder_cuda_graph_config requires " - f"{' and '.join(missing)}.") + # `num_tokens` / `seq_lens` are checked by the model engine rather than + # here: an encoder whose input is a fixed-shape per-request feature + # tensor derives both from the model, and only the engine knows which + # kind of encoder the model has. return self attn_backend: str = Field( diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 143a2c3cf5f3..38436f22514a 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -528,8 +528,7 @@ }, { "allowed_values": [ - "encode", - "feature_encode" + "encode" ], "annotation": "Literal['encode']", "converter": "", diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py index 5c9e82f65c5d..0572ffa83d5c 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py @@ -31,7 +31,7 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, - FeatureEncoderCudaGraphConfig, + EncodeCudaGraphConfig, KvCacheConfig, SamplingParams, SchedulerConfig, @@ -131,11 +131,12 @@ def _make_llm( encoder_kwargs = {} if encoder_graphs: # Whisper's encoder emits a fixed `_ENCODER_OUTPUT_LEN` positions per - # request, so the graph key is the batch size alone. + # request, so the graph key is the batch size alone; num_tokens and + # seq_lens are derived from the model and are not supplied here. encoder_batch_sizes = list(cuda_graph_batch_sizes or [1]) encoder_kwargs = { "encoder_max_batch_size": max(encoder_batch_sizes), - "encoder_cuda_graph_config": FeatureEncoderCudaGraphConfig( + "encoder_cuda_graph_config": EncodeCudaGraphConfig( batch_sizes=encoder_batch_sizes, enable_padding=True, ), diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 5f708387b6ac..7ce9a773a7e1 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -38,7 +38,7 @@ ScheduledRequests, SerializableSchedulerOutput, ) -from tensorrt_llm.llmapi.llm_args import FeatureEncoderCudaGraphConfig +from tensorrt_llm.llmapi.llm_args import EncodeCudaGraphConfig pytestmark = pytest.mark.cpu_only @@ -165,6 +165,11 @@ def _make_encoder_batch_wait_executor(batch_sizes=None, encoder_max_batch_size=8 ), encoder_max_batch_size=encoder_max_batch_size, ) + executor.model_engine = types.SimpleNamespace( + encoder_cuda_graph_runner=types.SimpleNamespace( + feature_mode=False, enabled=True, supported_batch_sizes=batch_sizes + ) + ) executor.batch_wait_timeout_iters = 48 executor.encoder_batch_wait_iters_count = 0 return executor @@ -175,19 +180,21 @@ def _make_feature_encoder_batch_wait_executor( ): """Batch-wait executor whose encoder graph config is the feature variant. - `FeatureEncoderCudaGraphConfig` has no `num_tokens` / `seq_lens`, and its - `batch_sizes` may have been derived rather than configured, so the resolved - sizes come from the engine's encoder graph runner rather than the config. + A feature encoder leaves `num_tokens` / `seq_lens` unset and may have had + its `batch_sizes` derived rather than configured, so the resolved sizes come + from the engine's encoder graph runner rather than the config. """ executor = object.__new__(PyExecutor) executor.max_batch_size = 32 executor.llm_args = types.SimpleNamespace( - encoder_cuda_graph_config=FeatureEncoderCudaGraphConfig(enable_padding=True), + encoder_cuda_graph_config=EncodeCudaGraphConfig(enable_padding=True), encoder_max_batch_size=encoder_max_batch_size, ) executor.model_engine = types.SimpleNamespace( encoder_cuda_graph_runner=types.SimpleNamespace( - supported_batch_sizes=runner_batch_sizes, enabled=runner_enabled + supported_batch_sizes=runner_batch_sizes, + enabled=runner_enabled, + feature_mode=True, ) ) executor.batch_wait_timeout_iters = 48 @@ -201,6 +208,7 @@ def _make_encoder_fallback_batch_wait_executor(): encoder_cuda_graph_config=None, encoder_max_batch_size=None, ) + executor.model_engine = types.SimpleNamespace(encoder_cuda_graph_runner=None) executor.batch_wait_timeout_iters = 48 executor.encoder_batch_wait_iters_count = 0 executor.batch_wait_max_tokens_ratio = 0.5 @@ -240,8 +248,8 @@ def test_encoder_graph_warmup_uses_runtime_encoder_stream(): def test_encoder_microbatch_admission_supports_feature_encoder_config(): - # A feature config carries no num_tokens / seq_lens, so reading them to - # gate this path raises AttributeError on the first encoder batch. + # A feature encoder leaves num_tokens / seq_lens unset, so gating this + # path on them would skip microbatch admission entirely. executor = _make_feature_encoder_batch_wait_executor([1, 2, 4, 8]) encoder_requests = [object() for _ in range(12)] @@ -255,11 +263,12 @@ def test_encoder_microbatch_admission_supports_feature_encoder_config(): assert executor.encoder_batch_wait_iters_count == 0 -def test_encoder_microbatch_admission_uses_derived_feature_batch_sizes(): - # batch_sizes left unset on the config: the engine derived them, so the - # runner is the only place the resolved list exists. +def test_encoder_microbatch_admission_uses_resolved_feature_batch_sizes(): + # The runner's list is authoritative: the engine filters the configured + # sizes by the scheduler's encoder-batch bound, so the config alone can + # name sizes that were never captured. executor = _make_feature_encoder_batch_wait_executor([1, 2, 3, 4]) - assert executor.llm_args.encoder_cuda_graph_config.batch_sizes is None + executor.llm_args.encoder_cuda_graph_config.batch_sizes = [1, 2, 3, 4, 8] encoder_requests = [object() for _ in range(6)] scheduled = executor._waiting_encoder_requests( diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 04f89d7ed270..71fee50804fb 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -27,7 +27,6 @@ _make_single_token_context_graph_batch) from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, EncodeCudaGraphConfig, - FeatureEncoderCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchLlmArgs) @@ -1066,21 +1065,6 @@ def test_feature_encoder_capture_layout_is_uniform(self) -> None: runner._capture_sequence_lengths[(2, 2 * fixed, fixed)], [fixed, fixed]) - def test_derived_feature_encoder_batch_sizes_are_dense_below_eight( - self) -> None: - # The padding guard refuses a graph once padding would add more than - # MAX_FEATURE_PADDING_RATIO of encoder work, so sparse low buckets - # would leave batch sizes 5-7 permanently eager. - derived = PyTorchModelEngine._derive_feature_encoder_batch_sizes(32) - self.assertEqual(derived, [1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32]) - - def test_derived_feature_encoder_batch_sizes_respect_cap(self) -> None: - self.assertEqual( - PyTorchModelEngine._derive_feature_encoder_batch_sizes(6), - [1, 2, 3, 4, 5, 6]) - self.assertEqual( - PyTorchModelEngine._derive_feature_encoder_batch_sizes(1), [1]) - @staticmethod def _encoder_spec_engine(encoder_cuda_graph_config, declares_spec: bool, @@ -1103,33 +1087,26 @@ def encoder_graph_spec(self): engine.mapping = SimpleNamespace(tp_size=tp_size) return engine, spec - def test_encoder_graph_spec_returns_spec_for_feature_config(self) -> None: + def test_encoder_graph_spec_returns_spec_for_feature_model(self) -> None: + # The model selects feature mode, not the config: an encoder either + # takes fixed-shape features or it does not. engine, spec = self._encoder_spec_engine( - FeatureEncoderCudaGraphConfig(batch_sizes=[1, 2]), - declares_spec=True) + EncodeCudaGraphConfig(batch_sizes=[1, 2]), declares_spec=True) self.assertEqual(engine._encoder_graph_spec(), spec) - def test_encoder_graph_spec_rejects_feature_config_on_token_model( - self) -> None: - # A token encoder cannot satisfy the fixed-shape contract, and a bare - # `batch_sizes` config is read as feature mode, so the message has to - # point at both possibilities. - engine, _ = self._encoder_spec_engine( - FeatureEncoderCudaGraphConfig(batch_sizes=[1]), declares_spec=False) - with self.assertRaises(ValueError) as ctx: - engine._encoder_graph_spec() - self.assertIn("EncodeCudaGraphConfig", str(ctx.exception)) - - def test_encoder_graph_spec_declines_token_config_on_feature_model( - self) -> None: + def test_encoder_graph_spec_declines_token_model(self) -> None: engine, _ = self._encoder_spec_engine(EncodeCudaGraphConfig( batch_sizes=[1], num_tokens=[1500], seq_lens=[1500]), - declares_spec=True) + declares_spec=False) + self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) + + def test_encoder_graph_spec_declines_without_config(self) -> None: + engine, _ = self._encoder_spec_engine(None, declares_spec=True) self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) def test_encoder_graph_spec_declines_tensor_parallel(self) -> None: engine, _ = self._encoder_spec_engine( - FeatureEncoderCudaGraphConfig(batch_sizes=[1]), + EncodeCudaGraphConfig(batch_sizes=[1]), declares_spec=True, tp_size=2) self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d1a3832fdcc4..409eafe66895 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -100,7 +100,7 @@ methods: default: null status: beta encoder_cuda_graph_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig, tensorrt_llm.llmapi.llm_args.FeatureEncoderCudaGraphConfig, NoneType] + annotation: Optional[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig] default: null status: prototype enable_encoder_decoder_mixed_cuda_graph: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index f288a0f54455..94a4c61f0bdd 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -43,7 +43,6 @@ EncodeCudaGraphConfig, ExecutorMemoryType, ExtendedRuntimePerfKnobConfig, - FeatureEncoderCudaGraphConfig, KvCacheConfig, LookaheadDecodingConfig, MambaStateConfig, MoeConfig, @@ -321,59 +320,34 @@ def test_rejects_non_positive(self, llm_args_cls, field_name, @pytest.mark.cpu_only -class TestFeatureEncoderCudaGraphConfig: - """Cover the fixed-shape feature encoder branch of encoder_cuda_graph_config. +class TestEncoderCudaGraphConfigValidation: + """`encoder_cuda_graph_config` accepts batch sizes on their own. - A feature encoder emits the same number of positions for every request, so - `batch_sizes` is its only bucket dimension. The token-shaped - `EncodeCudaGraphConfig` fields must not be expressible here, or a caller - could supply values the encoder graph runner silently overrides. + An encoder whose input is a fixed-shape per-request feature tensor + (Whisper) derives `num_tokens` / `seq_lens` from the model, and only the + model engine knows which kind of encoder a model has, so those buckets are + checked there rather than at config validation. """ - @pytest.mark.parametrize("batch_sizes", [[], [0], [-1], [1, -2]]) - def test_rejects_non_positive_or_empty_batch_sizes(self, batch_sizes): - with pytest.raises(ValidationError): - FeatureEncoderCudaGraphConfig(batch_sizes=batch_sizes) - - def test_sorts_and_deduplicates_batch_sizes(self): - cfg = FeatureEncoderCudaGraphConfig(batch_sizes=[4, 2, 2, 1]) - assert cfg.batch_sizes == [1, 2, 4] - - @pytest.mark.parametrize("field", ["num_tokens", "seq_lens"]) - def test_rejects_token_shaped_buckets(self, field): - # These belong to EncodeCudaGraphConfig; accepting them here would be - # accepting a value the runner discards. - with pytest.raises(ValidationError): - FeatureEncoderCudaGraphConfig(batch_sizes=[1], **{field: [64]}) - - @pytest.mark.parametrize( - "config_dict, expected_type", - [ - (dict(batch_sizes=[1, 2]), FeatureEncoderCudaGraphConfig), - (dict(batch_sizes=[1], seq_lens=[64], - num_tokens=[64]), EncodeCudaGraphConfig), - (dict(mode="feature_encode", - batch_sizes=[1]), FeatureEncoderCudaGraphConfig), - ], - ids=["bare_batch_sizes", "token_buckets", "explicit_mode"], - ) - def test_encoder_config_mode_is_inferred(self, config_dict, expected_type): + def test_accepts_batch_sizes_without_token_buckets(self): llm_args = TorchLlmArgs(model=llama_model_path, encoder_max_batch_size=8, - encoder_cuda_graph_config=config_dict) - assert isinstance(llm_args.encoder_cuda_graph_config, expected_type) - - def test_feature_config_does_not_require_token_buckets(self): - llm_args = TorchLlmArgs( - model=llama_model_path, - encoder_max_batch_size=8, - encoder_cuda_graph_config=FeatureEncoderCudaGraphConfig( - batch_sizes=[1, 2])) + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 2], enable_padding=True)) assert llm_args.encoder_cuda_graph_config.batch_sizes == [1, 2] + assert not llm_args.encoder_cuda_graph_config.num_tokens + assert not llm_args.encoder_cuda_graph_config.seq_lens + + def test_still_requires_encoder_max_batch_size(self): + with pytest.raises(ValueError): + TorchLlmArgs(model=llama_model_path, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 2])) - def test_token_config_still_requires_token_buckets(self): + def test_still_rejects_encode_only(self): with pytest.raises(ValueError): TorchLlmArgs(model=llama_model_path, + encode_only=True, encoder_max_batch_size=8, encoder_cuda_graph_config=EncodeCudaGraphConfig( batch_sizes=[1, 2])) @@ -1930,19 +1904,6 @@ def test_encoder_cuda_graph_config_validation(self): }, "encoder_cuda_graph_config requires encoder_max_batch_size", ), - ( - { - "encoder_max_batch_size": - 4, - "encoder_cuda_graph_config": - EncodeCudaGraphConfig( - batch_sizes=[1, 4], - enable_padding=True, - ), - }, - ("encoder_cuda_graph_config requires " - "num_tokens/max_num_token and seq_lens/max_seq_len"), - ), ] for kwargs, error_match in invalid_cases: From 71799171b339e0bd2d7df7f620b11409df24c2f7 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:20:38 -0700 Subject: [PATCH 5/8] [TRTLLM-14778][chore] Consolidate encoder graph paths and drop dead branches Delete five unreachable branches, fold the feature and token capture setups into one parameterized capture epilogue, and route feature-mode warmup through the existing enc-dec driver, which captures on the worker owning runtime replay. Feature capture now uses capture_error_mode="thread_local" like the token path, and a feature model with no fitting batch size disables the runner outright. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 219 ++++++++---------- .../_torch/pyexecutor/model_engine.py | 169 ++++++-------- .../test_lists/test-db/l0_h100.yml | 1 - .../test_lists/test-db/l0_l40s.yml | 7 +- .../_torch/executor/test_py_executor.py | 43 ++-- .../executor/test_pytorch_model_engine.py | 65 +++--- tests/unittest/llmapi/test_llm_args.py | 83 +++---- 7 files changed, 248 insertions(+), 339 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index d23d5284623b..7028ca1fe398 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1086,15 +1086,6 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.is_encoder_decoder = config.is_encoder_decoder self.use_fixed_sequence_slots = config.use_fixed_sequence_slots - if self.feature_mode and not self.is_encoder_decoder: - # Nothing structural forbids this - a standalone feature encoder - # (audio/vision embedding tower) would land here - but no in-tree - # model exercises it, so fail loudly rather than capture untested - # shapes. - raise NotImplementedError( - "Feature-mode encoder CUDA graphs are only supported for " - "encoder-decoder models today.") - if self.feature_mode: # A feature encoder produces a fixed number of positions per # request, so the configured token/seq-len buckets are not free @@ -1618,12 +1609,6 @@ def maybe_get_cuda_graph( if padded_batch_size not in self.supported_batch_sizes: return None, None - if self.feature_mode and any(s != self.config.fixed_seq_len - for s in seq_lens): - # Fixed-shape contract violated (should not happen for feature - # encoders); fall back to eager rather than replay a wrong shape. - return None, None - key, is_padding_performed, is_padding_successful = self.get_graph_key( inputs) if self.is_encoder_decoder and key not in self.capture_keys: @@ -1844,52 +1829,22 @@ def capture( inputs: Dict[str, Any], ) -> Any: """Warm up and/or capture the forward pass for a graph key.""" - padded_num_tokens = key[1] + capture_inputs, capture_h2d = (self._prepare_feature_capture( + key, inputs) if self.feature_mode else self._prepare_token_capture( + key, inputs)) - if self.feature_mode: - return self._capture_features(key, forward_fn, inputs) - - sliced_static_tensors = { - "input_ids": - self.shared_static_tensors["input_ids"][:padded_num_tokens], - "position_ids": - self.shared_static_tensors["position_ids"][:, :padded_num_tokens], - } - sliced_static_tensors_cpu = { - "input_ids": - self.shared_static_tensors_cpu["input_ids"][:padded_num_tokens], - "position_ids": - self.shared_static_tensors_cpu["position_ids"] - [:, :padded_num_tokens], + self.graph_metadata[key] = { + "attn_metadata": capture_inputs["attn_metadata"] } - capture_inputs = dict(inputs) - capture_inputs.update(sliced_static_tensors) - - attn_md = capture_inputs["attn_metadata"] - - self.graph_metadata[key] = {"attn_metadata": attn_md} - - # Warmup must see the same runtime data as capture. In particular, - # graph metadata initializes _seq_lens_cuda to ones, while - # prepare_encoder_cuda_graph_replay updates its stable host buffer. - # Populate every device input before warmup so packed-token counts and - # sequence boundaries are consistent. - self._stage_inputs(key, inputs) - if self._capture_h2d_copy: - capture_inputs["input_ids"].copy_( - sliced_static_tensors_cpu["input_ids"], non_blocking=True) - capture_inputs["position_ids"].copy_( - sliced_static_tensors_cpu["position_ids"], non_blocking=True) - attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) - torch.cuda.current_stream().synchronize() - output = None with with_multi_stream(True), piecewise_cuda_graph(False): # Warmup runs required by CUDA graph semantics. See # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graph-semantics # Warmups initialize PyTorch and attention metadata state, and # resize the shared attention workspace before any graph is captured. + # The warmup pass must not build a graph; its caller consumes the + # eager output directly. for _ in range(self.WARMUP_STEPS): output = forward_fn(capture_inputs) @@ -1901,18 +1856,8 @@ def capture( pool=self.memory_pool, stream=self._get_capture_stream(), capture_error_mode="thread_local"): - if self._capture_h2d_copy: - # H2D copies for captured inside the graph: at replay - # time it re-issues from the pinned static buffer without - # an eager driver call. - capture_inputs["input_ids"].copy_( - sliced_static_tensors_cpu["input_ids"], - non_blocking=True) - capture_inputs["position_ids"].copy_( - sliced_static_tensors_cpu["position_ids"], - non_blocking=True) - attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, - non_blocking=True) + if capture_h2d is not None: + capture_h2d() output = forward_fn(capture_inputs) if self._contains_nested_tensor(output): @@ -1925,75 +1870,97 @@ def capture( self.memory_pool = graph.pool() return graph_output - def retire_staging(self) -> None: - """Wait until a prior replay no longer reads shared staging buffers.""" - if self._staging_retirement_event is not None: - self._staging_retirement_event.synchronize() - self._staging_retirement_event = None - - def _capture_features( + def _prepare_token_capture( self, key: EncoderKeyType, - forward_fn: Callable[[Dict[str, Any]], Any], inputs: Dict[str, Any], - ) -> Any: - """Capture path for the fixed-shape feature mode (enc-dec encoders). + ) -> Tuple[Dict[str, Any], Optional[Callable[[], None]]]: + """Capture setup for the packed-token mode. - The capture region receives the static device feature buffer sliced - to the padded batch size; in pinned mode the H2D from the pinned CPU - mirror is captured inside the graph so replay re-issues it without an - eager driver call. + Returns the capture inputs and, in pinned mode, a callable replaying + the input H2D inside the capture region so that graph replay re-issues + it from the pinned static buffer without an eager driver call. """ - padded_batch_size, _, _ = key + padded_num_tokens = key[1] - static_features = ( - self.shared_static_tensors["input_features"][:padded_batch_size]) + sliced_static_tensors = { + "input_ids": + self.shared_static_tensors["input_ids"][:padded_num_tokens], + "position_ids": + self.shared_static_tensors["position_ids"][:, :padded_num_tokens], + } + sliced_static_tensors_cpu = { + "input_ids": + self.shared_static_tensors_cpu["input_ids"][:padded_num_tokens], + "position_ids": + self.shared_static_tensors_cpu["position_ids"] + [:, :padded_num_tokens], + } capture_inputs = dict(inputs) - capture_inputs["input_features"] = static_features - + capture_inputs.update(sliced_static_tensors) attn_md = capture_inputs["attn_metadata"] - self.graph_metadata[key] = { - "attn_metadata": attn_md, - } + + def copy_inputs() -> None: + capture_inputs["input_ids"].copy_( + sliced_static_tensors_cpu["input_ids"], non_blocking=True) + capture_inputs["position_ids"].copy_( + sliced_static_tensors_cpu["position_ids"], non_blocking=True) + + # Warmup must see the same runtime data as capture. In particular, + # graph metadata initializes _seq_lens_cuda to ones, while + # prepare_encoder_cuda_graph_replay updates its stable host buffer. + # Populate every device input before warmup so packed-token counts and + # sequence boundaries are consistent. + self._stage_inputs(key, inputs) + if self._capture_h2d_copy: + copy_inputs() + attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) + torch.cuda.current_stream().synchronize() + + if not self._capture_h2d_copy: + return capture_inputs, None + + def capture_h2d() -> None: + copy_inputs() + attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) + + return capture_inputs, capture_h2d + + def _prepare_feature_capture( + self, + key: EncoderKeyType, + inputs: Dict[str, Any], + ) -> Tuple[Dict[str, Any], Optional[Callable[[], None]]]: + """Capture setup for the fixed-shape feature mode. + + The capture region receives the static device feature buffer sliced to + the padded batch size, and never captures the input H2D: all buckets + share one pinned mirror, and consecutive encoder batches (different + buckets) can be enqueued back-to-back, so a captured H2D would read the + mirror at replay-execution time, after the host has already refilled it + for the next batch. The eager H2D in `_replay_features` is + stream-ordered and guarded by per-mirror events instead. + """ + padded_batch_size, _, _ = key + + capture_inputs = dict(inputs) + capture_inputs["input_features"] = ( + self.shared_static_tensors["input_features"][:padded_batch_size]) # Feature-mode seq_lens never change for this key: populate the # metadata's device seq_lens once, eagerly, instead of capturing the # H2D like the token path does per replay. + attn_md = capture_inputs["attn_metadata"] attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) - # NOTE: unlike the token path, the input H2D is NOT captured inside - # the graph. All buckets share one pinned mirror, and consecutive - # encoder batches (different buckets) can be enqueued back-to-back — - # a captured H2D would read the mirror at replay-execution time, - # after the host has already refilled it for the next batch. The - # eager H2D in `_replay_features` is stream-ordered and guarded by - # per-mirror events instead. - output = None - with with_multi_stream(True), piecewise_cuda_graph(False): - for _ in range(self.WARMUP_STEPS): - output = forward_fn(capture_inputs) - - # The warmup pass runs these shapes eagerly to settle PyTorch and - # attention state; it must not build a graph, and its caller - # consumes the eager output directly. - if self.is_warmup_only: - return output - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, - pool=self.memory_pool, - stream=self._get_capture_stream()): - output = forward_fn(capture_inputs) + return capture_inputs, None - if self._contains_nested_tensor(output): - raise TypeError( - "Encoder CUDA graph does not support nested tensor outputs.") - self.graphs[key] = graph - graph_output = make_weak_ref(output) - self.graph_outputs[key] = graph_output - self.memory_pool = graph.pool() - return graph_output + def retire_staging(self) -> None: + """Wait until a prior replay no longer reads shared staging buffers.""" + if self._staging_retirement_event is not None: + self._staging_retirement_event.synchronize() + self._staging_retirement_event = None def _replay_features( self, @@ -2015,18 +1982,14 @@ def _replay_features( # satisfied and the host proceeds straight to the fill. self._feature_h2d_events[slot].synchronize() + # Per-request CPU tensors straight from the requests — one copy into + # the host mirror, no intermediate packing. mirror = self._feature_mirrors[slot] - if isinstance(features, list): - # Per-request CPU tensors straight from the requests — one copy - # into the host mirror, no intermediate packing. - rows = 0 - for f in features: - n = int(f.shape[0]) - mirror[rows:rows + n].copy_(f) - rows += n - else: - rows = int(features.shape[0]) - mirror[:rows].copy_(features) + rows = 0 + for f in features: + n = int(f.shape[0]) + mirror[rows:rows + n].copy_(f) + rows += n if rows < padded_batch_size: mirror[rows:padded_batch_size].zero_() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f6b7d045e381..87ee8454bc35 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -890,6 +890,10 @@ def __init__( f"size fits within max_num_tokens // encoder_output_len = " f"{bs_cap}; the encoder step stays eager.") feature_shape = feature_dtype = fixed_seq_len = None + # Without this the runner would be built in *token* mode with + # empty bucket lists: no capture keys, but `enabled` true and + # token-mode static buffers allocated for nothing. + use_encoder_cuda_graph = False else: encoder_graph_max_batch_size = encoder_graph_batch_sizes[-1] encoder_graph_max_num_tokens = (encoder_graph_max_batch_size * @@ -1405,11 +1409,6 @@ def warmup(self, resource_manager: ResourceManager) -> None: 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 self.encoder_cuda_graph_runner.feature_mode: - # After decoder-graph capture, so the decoder pool's high-water - # mark is set before the encoder runner allocates its own pool. - self._capture_enc_dec_encoder_graphs() - log_mem_snapshot("warmup/after_enc_dec_encoder_graph_capture") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -2121,11 +2120,6 @@ def _warmup_encoder_cuda_graphs_enc_dec( runner = self.encoder_cuda_graph_runner if not runner.is_encoder_decoder: return - if runner.feature_mode: - # This driver synthesizes packed token inputs, which a feature - # encoder cannot consume. Feature mode captures from - # `_capture_enc_dec_encoder_graphs` during engine warmup instead. - return capture = functools.partial( self._capture_encoder_cuda_graphs_enc_dec, @@ -2164,23 +2158,36 @@ def _capture_encoder_cuda_graphs_enc_dec( if sequence_lengths is None: continue - encoder_input_ids = [0] * sum(sequence_lengths) - encoder_position_ids = [] - for sequence_length in sequence_lengths: - encoder_position_ids.extend( - self._apply_position_id_offset(list( - range(sequence_length)))) - inputs = self._prepare_encoder_decoder_encoder_inputs( - encoder_input_ids=encoder_input_ids, - encoder_position_ids=encoder_position_ids, - sequence_lengths=sequence_lengths, - request_ids=list(range(len(sequence_lengths))), - resource_manager=resource_manager, - ) - logger.info("Encoder-decoder encoder CUDA graph " f"{operation}: key={key}") - self._encoder_forward_enc_dec(inputs) + if runner.feature_mode: + # A zero waveform is a valid fixed-shape feature, and the + # encoder step writes no KV cache, so no LlmRequests and no + # KV/cross-pool resources are involved. + self._feature_encoder_graph_forward( + features=[ + torch.zeros((1, *runner.config.feature_shape), + dtype=runner.config.feature_dtype) + for _ in sequence_lengths + ], + seq_lens=list(sequence_lengths), + request_ids=list(range(len(sequence_lengths))), + ) + else: + encoder_input_ids = [0] * sum(sequence_lengths) + encoder_position_ids = [] + for sequence_length in sequence_lengths: + encoder_position_ids.extend( + self._apply_position_id_offset( + list(range(sequence_length)))) + inputs = self._prepare_encoder_decoder_encoder_inputs( + encoder_input_ids=encoder_input_ids, + encoder_position_ids=encoder_position_ids, + sequence_lengths=sequence_lengths, + request_ids=list(range(len(sequence_lengths))), + resource_manager=resource_manager, + ) + self._encoder_forward_enc_dec(inputs) torch.cuda.synchronize() num_processed += 1 @@ -7928,7 +7935,7 @@ def _maybe_forward_encoder_graph( fall back to the eager path. """ runner = self.encoder_cuda_graph_runner - if runner is None or not runner.enabled or not runner.feature_mode: + if not runner.enabled or not runner.feature_mode: return None fixed = runner.config.fixed_seq_len @@ -7942,42 +7949,65 @@ def _maybe_forward_encoder_graph( features.append(f) seq_lens = [fixed] * len(encoder_requests) - graph_inputs = { - 'seq_lens': seq_lens, - 'input_features': features, - } - with runner.pad_batch(graph_inputs, - len(encoder_requests)) as padded_inputs: + output = self._feature_encoder_graph_forward( + features=features, + seq_lens=seq_lens, + request_ids=[r.py_request_id for r in encoder_requests], + ) + if output is None: + return None + + real_tokens = fixed * len(encoder_requests) + return output[:real_tokens].clone(), seq_lens + + def _feature_encoder_graph_forward( + self, + features: List[torch.Tensor], + seq_lens: List[int], + request_ids: List[int], + ) -> Optional[torch.Tensor]: + """Run one feature encoder batch through its CUDA graph. + + Shared by the runtime path and by warmup/capture. Returns the packed + hidden states for the *padded* batch (the caller slices back to the + real rows), or None when no captured graph fits and the caller must + fall back to eager. + """ + runner = self.encoder_cuda_graph_runner + fixed = runner.config.fixed_seq_len + graph_inputs = {'seq_lens': seq_lens, 'input_features': features} + + with runner.pad_batch(graph_inputs, len(seq_lens)) as padded_inputs: # `pad_batch` extends seq_lens to the captured bucket, and the # metadata takes one request id per sequence. Pad slots carry no # request; the encoder pass runs without a KV cache, so their ids # are never looked up and only have to exist and stay distinct. - request_ids = [r.py_request_id for r in encoder_requests] - request_ids += [ - -(i + 1) for i in range( - len(padded_inputs['seq_lens']) - len(request_ids)) + padded_seq_lens = padded_inputs['seq_lens'] + padded_request_ids = list(request_ids) + [ + -(i + 1) + for i in range(len(padded_seq_lens) - len(request_ids)) ] eager_attn_metadata = self._make_encoder_attn_metadata( - padded_inputs['seq_lens'], request_ids) + padded_seq_lens, padded_request_ids) graph_attn_metadata, key = runner.maybe_get_cuda_graph( padded_inputs, eager_attn_metadata) if key is None: return None padded_inputs['attn_metadata'] = graph_attn_metadata + capture_output = None if runner.needs_capture(key): padded_batch_size, padded_num_tokens, _ = key # Feature-mode seq_lens are constant per bucket: initialize # the graph-resident metadata once at capture. graph_attn_metadata.prepare_encoder_cuda_graph_replay( [fixed] * padded_batch_size, padded_num_tokens) - runner.capture(key, self._enc_dec_encoder_graph_forward_fn, - padded_inputs) + capture_output = runner.capture( + key, self._enc_dec_encoder_graph_forward_fn, padded_inputs) - output = runner.replay(key, padded_inputs) - - real_tokens = fixed * len(encoder_requests) - return output[:real_tokens].clone(), seq_lens + if runner.is_warmup_only: + return capture_output + return runner.replay(key, padded_inputs) def _enc_dec_encoder_graph_forward_fn( self, capture_inputs: Dict[str, Any]) -> torch.Tensor: @@ -7990,59 +8020,6 @@ def _enc_dec_encoder_graph_forward_fn( capture_inputs['seq_lens'], }) - def _capture_enc_dec_encoder_graphs(self) -> None: - """Capture enc-dec encoder graphs for every configured batch size. - - Runs at engine warmup, and goes through the shared two-pass helper: - every shape must be warmed before any graph is captured, because a - smaller batch can select a different attention kernel and grow the - shared workspace, which would move a buffer a larger batch's graph - had already captured the address of. - """ - with torch.inference_mode(): - self._warmup_and_capture_encoder_cuda_graphs( - self._capture_feature_encoder_graphs_once) - - def _capture_feature_encoder_graphs_once(self) -> None: - """One pass over every feature encoder batch size. - - Called twice by `_warmup_and_capture_encoder_cuda_graphs`: once with - the runner in warmup-only mode, then once to capture. Largest bucket - first so the runner's graph pool high-water mark is set on the first - capture. Inputs are synthesized without LlmRequests — a zero waveform - is a valid fixed-shape feature — and no KV/cross-pool resources are - involved (the encoder step writes no KV cache). - """ - runner = self.encoder_cuda_graph_runner - fixed = runner.config.fixed_seq_len - for bs in sorted(runner.supported_batch_sizes, reverse=True): - seq_lens = [fixed] * bs - features = [ - torch.zeros((1, *runner.config.feature_shape), - dtype=runner.config.feature_dtype) - for _ in range(bs) - ] - graph_inputs = { - 'seq_lens': seq_lens, - 'input_features': features, - } - eager_md = self._make_encoder_attn_metadata(seq_lens, - list(range(bs))) - graph_md, key = runner.maybe_get_cuda_graph(graph_inputs, eager_md) - if key is None: - logger.warning( - "Enc-dec encoder CUDA graph capture skipped for " - f"batch size {bs} (unsupported metadata/backend).") - continue - if not runner.needs_capture(key): - continue - logger.info( - f"Capturing enc-dec encoder CUDA graph for batch size {bs}.") - graph_inputs['attn_metadata'] = graph_md - graph_md.prepare_encoder_cuda_graph_replay(seq_lens, key[1]) - runner.capture(key, self._enc_dec_encoder_graph_forward_fn, - graph_inputs) - def _init_userbuffers(self, hidden_size): if self.mapping.tp_size <= 1 or self.mapping.pp_size > 1: return False diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 4ef714fdaf8e..88130605c153 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -346,7 +346,6 @@ l0_h100: - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-requested-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-decoder-graphs-on-greedy] - - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp16-kv-v1-graphs-off-greedy] - examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b] - unittest/bindings # 8 mins on H100 diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index 854439219d7c..2e9e5bd661b9 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -54,9 +54,11 @@ l0_l40s: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-on-greedy-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-overlap-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-t5-small] - # Whisper (encoder-decoder) — customer-side deployment targets L40S/H200 + # Whisper (encoder-decoder) — customer-side deployment targets L40S/H200. + # The encoder-graphs case also exercises decoder graphs, so it stands in for + # a decoder-only case rather than adding to it; KV-v2 stays covered on H100. - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_transcribe_end_to_end - - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v2-decoder-graphs-on-greedy] + - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy] - condition: ranges: system_gpu_count: @@ -76,7 +78,6 @@ l0_l40s: - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-requested-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-decoder-graphs-on-greedy] - - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp16-kv-v1-graphs-off-greedy] - condition: ranges: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 7ce9a773a7e1..da52d489b242 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -247,11 +247,25 @@ def test_encoder_graph_warmup_uses_runtime_encoder_stream(): ) -def test_encoder_microbatch_admission_supports_feature_encoder_config(): - # A feature encoder leaves num_tokens / seq_lens unset, so gating this - # path on them would skip microbatch admission entirely. - executor = _make_feature_encoder_batch_wait_executor([1, 2, 4, 8]) - encoder_requests = [object() for _ in range(12)] +@pytest.mark.parametrize( + "runner_batch_sizes,config_batch_sizes,num_requests,expected", + [ + # A feature encoder leaves num_tokens / seq_lens unset, so gating this + # path on them would skip microbatch admission entirely. + ([1, 2, 4, 8], None, 12, 8), + # The runner's list is authoritative: the engine filters the configured + # sizes by the scheduler's encoder-batch bound, so the config alone can + # name sizes that were never captured. + ([1, 2, 3, 4], [1, 2, 3, 4, 8], 6, 4), + ], +) +def test_encoder_microbatch_admission_uses_resolved_feature_batch_sizes( + runner_batch_sizes, config_batch_sizes, num_requests, expected +): + executor = _make_feature_encoder_batch_wait_executor(runner_batch_sizes) + if config_batch_sizes is not None: + executor.llm_args.encoder_cuda_graph_config.batch_sizes = config_batch_sizes + encoder_requests = [object() for _ in range(num_requests)] scheduled = executor._waiting_encoder_requests( encoder_requests, @@ -259,27 +273,10 @@ def test_encoder_microbatch_admission_supports_feature_encoder_config(): [object()] * 20, ) - assert scheduled == encoder_requests[:8] + assert scheduled == encoder_requests[:expected] assert executor.encoder_batch_wait_iters_count == 0 -def test_encoder_microbatch_admission_uses_resolved_feature_batch_sizes(): - # The runner's list is authoritative: the engine filters the configured - # sizes by the scheduler's encoder-batch bound, so the config alone can - # name sizes that were never captured. - executor = _make_feature_encoder_batch_wait_executor([1, 2, 3, 4]) - executor.llm_args.encoder_cuda_graph_config.batch_sizes = [1, 2, 3, 4, 8] - encoder_requests = [object() for _ in range(6)] - - scheduled = executor._waiting_encoder_requests( - encoder_requests, - [], - [object()] * 20, - ) - - assert scheduled == encoder_requests[:4] - - def test_encoder_microbatch_admission_ignores_disabled_feature_runner(): # supported_batch_sizes stays populated from the config even when capture # was declined (TP > 1, or no bucket fits), so waiting on those shapes diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 71fee50804fb..2b85cb76c202 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1045,25 +1045,22 @@ def _feature_encoder_runner(batch_sizes, fixed_seq_len=1500): def test_feature_encoder_capture_keys_are_all_reachable(self) -> None: # Every request contributes exactly fixed_seq_len positions, so the - # only reachable key per batch size is (bs, bs * fixed, fixed). The - # token path's cross product would also emit keys whose token count no - # batch can produce, and capture_keys drives mixed encoder/decoder + # only reachable key per batch size is (bs, bs * fixed, fixed), and + # every slot in that layout is a full fixed_seq_len sequence. The token + # path's cross product would also emit keys whose token count no batch + # can produce, and capture_keys drives mixed encoder/decoder # decoder-graph warmup. fixed = 1500 batch_sizes = [1, 2, 4, 8] runner = self._feature_encoder_runner(batch_sizes, fixed) self.assertEqual( - sorted(runner.capture_keys), - [(bs, bs * fixed, fixed) for bs in batch_sizes], + runner._capture_sequence_lengths, + {(bs, bs * fixed, fixed): [fixed] * bs + for bs in batch_sizes}, ) - - def test_feature_encoder_capture_layout_is_uniform(self) -> None: - fixed = 1500 - runner = self._feature_encoder_runner([2], fixed) - self.assertEqual( - runner._capture_sequence_lengths[(2, 2 * fixed, fixed)], - [fixed, fixed]) + self.assertEqual(runner.capture_keys, + frozenset(runner._capture_sequence_lengths)) @staticmethod def _encoder_spec_engine(encoder_cuda_graph_config, @@ -1087,29 +1084,29 @@ def encoder_graph_spec(self): engine.mapping = SimpleNamespace(tp_size=tp_size) return engine, spec - def test_encoder_graph_spec_returns_spec_for_feature_model(self) -> None: + def test_encoder_graph_spec_selection(self) -> None: # The model selects feature mode, not the config: an encoder either - # takes fixed-shape features or it does not. - engine, spec = self._encoder_spec_engine( - EncodeCudaGraphConfig(batch_sizes=[1, 2]), declares_spec=True) - self.assertEqual(engine._encoder_graph_spec(), spec) - - def test_encoder_graph_spec_declines_token_model(self) -> None: - engine, _ = self._encoder_spec_engine(EncodeCudaGraphConfig( - batch_sizes=[1], num_tokens=[1500], seq_lens=[1500]), - declares_spec=False) - self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) - - def test_encoder_graph_spec_declines_without_config(self) -> None: - engine, _ = self._encoder_spec_engine(None, declares_spec=True) - self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) - - def test_encoder_graph_spec_declines_tensor_parallel(self) -> None: - engine, _ = self._encoder_spec_engine( - EncodeCudaGraphConfig(batch_sizes=[1]), - declares_spec=True, - tp_size=2) - self.assertEqual(engine._encoder_graph_spec(), (None, None, None)) + # takes fixed-shape features or it does not. TP > 1 is gated off + # because allreduce inside encoder capture is unverified. + declined = (None, None, None) + cases = [ + ("feature model", EncodeCudaGraphConfig(batch_sizes=[1, 2]), True, + 1, None), + ("token model", + EncodeCudaGraphConfig(batch_sizes=[1], + num_tokens=[1500], + seq_lens=[1500]), False, 1, declined), + ("no config", None, True, 1, declined), + ("tensor parallel", EncodeCudaGraphConfig(batch_sizes=[1]), True, 2, + declined), + ] + + for name, config, declares_spec, tp_size, expected in cases: + with self.subTest(name): + engine, spec = self._encoder_spec_engine( + config, declares_spec=declares_spec, tp_size=tp_size) + self.assertEqual(engine._encoder_graph_spec(), + expected if expected is not None else spec) def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( self) -> None: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 94a4c61f0bdd..9cb89ee89802 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -319,40 +319,6 @@ def test_rejects_non_positive(self, llm_args_cls, field_name, llm_args_cls(model=llama_model_path, **{field_name: invalid_value}) -@pytest.mark.cpu_only -class TestEncoderCudaGraphConfigValidation: - """`encoder_cuda_graph_config` accepts batch sizes on their own. - - An encoder whose input is a fixed-shape per-request feature tensor - (Whisper) derives `num_tokens` / `seq_lens` from the model, and only the - model engine knows which kind of encoder a model has, so those buckets are - checked there rather than at config validation. - """ - - def test_accepts_batch_sizes_without_token_buckets(self): - llm_args = TorchLlmArgs(model=llama_model_path, - encoder_max_batch_size=8, - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 2], enable_padding=True)) - assert llm_args.encoder_cuda_graph_config.batch_sizes == [1, 2] - assert not llm_args.encoder_cuda_graph_config.num_tokens - assert not llm_args.encoder_cuda_graph_config.seq_lens - - def test_still_requires_encoder_max_batch_size(self): - with pytest.raises(ValueError): - TorchLlmArgs(model=llama_model_path, - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 2])) - - def test_still_rejects_encode_only(self): - with pytest.raises(ValueError): - TorchLlmArgs(model=llama_model_path, - encode_only=True, - encoder_max_batch_size=8, - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 2])) - - @pytest.mark.cpu_only def test_decoding_type_eagle3_parses_to_eagle3_decoding_config(): adapter = TypeAdapter(SpeculativeConfig) @@ -1890,28 +1856,37 @@ def test_encoder_decoder_cuda_graph_user_interface(self): assert not disabled_args.enable_encoder_decoder_mixed_cuda_graph - def test_encoder_cuda_graph_config_validation(self): - invalid_cases = [ - ( - { - "encoder_cuda_graph_config": - EncodeCudaGraphConfig( - batch_sizes=[1, 4], - num_tokens=[16, 64], - seq_lens=[8, 32], - enable_padding=True, - ), - }, - "encoder_cuda_graph_config requires encoder_max_batch_size", + # Batch sizes alone are valid. An encoder whose input is a fixed-shape + # per-request feature tensor (Whisper) derives num_tokens / seq_lens + # from the model, and only the model engine knows which kind of encoder + # a model has, so those buckets are checked there, not here. + feature_args = TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=4, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + enable_padding=True, ), - ] + ) + + assert feature_args.encoder_cuda_graph_config.batch_sizes == [1, 4] + assert not feature_args.encoder_cuda_graph_config.num_tokens + assert not feature_args.encoder_cuda_graph_config.seq_lens - for kwargs, error_match in invalid_cases: - with pytest.raises(ValidationError, match=error_match): - TorchLlmArgs( - model=llama_model_path, - **kwargs, - ) + def test_encoder_cuda_graph_config_validation(self): + with pytest.raises( + ValidationError, + match="encoder_cuda_graph_config requires encoder_max_batch_size" + ): + TorchLlmArgs( + model=llama_model_path, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ), + ) def test_cuda_graph_config_infers_encode_mode_from_raw_dict(self): args = TorchLlmArgs( From a6377377c54ebc454bc7cf3bf8321a9f901ee984 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:22:58 -0700 Subject: [PATCH 6/8] [TRTLLM-14778][fix] Resolve encoder graph batch sizes in one pass `_filter_cuda_graph_batch_sizes` now takes the per-request token cost, so a feature encoder's `fixed_seq_len` budget is applied by the same filter the token path uses instead of a second cap layered on afterwards. A cap below one request returns no batch sizes, making the "nothing fits, stay eager" fallback reachable. Also annotate the spec helpers and retire the in-flight H2D before the feature staging buffer is reallocated. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 138 +++++++++++------- 1 file changed, 84 insertions(+), 54 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2869b27c778f..e7d51e4078d3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -218,11 +218,21 @@ def _filter_piecewise_capture_num_tokens( def _filter_cuda_graph_batch_sizes(cuda_graph_batch_sizes: list[int], max_batch_size: int, max_num_tokens: int, - max_total_draft_tokens: int, + tokens_per_request: int, enable_padding: bool) -> list[int]: - # This is the largest possible batch size for a pure decoding batch. + """Drop the batch sizes that exceed the request or token budget. + + `tokens_per_request` is what a single request costs against + `max_num_tokens`: `1 + max_total_draft_tokens` for a pure decoding batch, + or the fixed encoder output length for an encoder whose input is a + fixed-shape per-request feature tensor. + """ max_cuda_graph_bs = min(max_batch_size, - int(max_num_tokens / (1 + max_total_draft_tokens))) + max_num_tokens // tokens_per_request) + if max_cuda_graph_bs < 1: + # Not even a single request fits the token budget, so there is no + # capturable batch size and padding has nothing to pad to. + return [] result = [] # This function assumes cuda_graph_batch_sizes is sorted @@ -589,9 +599,10 @@ def __init__( 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, - self.original_max_total_draft_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 @@ -599,9 +610,18 @@ def __init__( self._encoder_cuda_graph_padding_enabled = ( encoder_cuda_graph_padding_enabled) + + # A feature-driven encoder (Whisper) declares a fixed-shape per-request + # contract instead of packed tokens: every request costs exactly + # `fixed_seq_len` of the encoder token budget, and the num_tokens / + # seq_lens buckets are derived from the model rather than configured. + # The model selects the mode; `encoder_cuda_graph_config` only opts in. + (self._encoder_feature_shape, self._encoder_feature_dtype, + self._encoder_fixed_seq_len) = self._encoder_graph_spec() + self._encoder_cuda_graph_batch_sizes = (_filter_cuda_graph_batch_sizes( encoder_cuda_graph_batch_sizes, self.encoder_batch_size, - self.encoder_max_num_tokens, 0, + self.encoder_max_num_tokens, self._encoder_fixed_seq_len or 1, self._encoder_cuda_graph_padding_enabled) if encoder_cuda_graph_batch_sizes else []) @@ -622,12 +642,41 @@ def __init__( self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) - use_encoder_cuda_graph = ( - (self._is_encoder_decoder_model() or self._is_encode_only) - and self.encoder_cuda_graph_config is not None - and (self._model_encoder_graph_spec() is not None or - (bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens)))) + # Resolve which capture mode has usable shapes. In feature mode the + # batch sizes *are* the whole key space, so an empty list after budget + # filtering leaves nothing to capture. A model that declares a feature + # contract cannot consume the packed token inputs the token-shaped + # capture path synthesizes, so when feature mode is unavailable for it + # the encoder stays eager instead of falling through to token capture. + if self._encoder_feature_shape is not None: + encoder_graph_shapes_available = bool( + self._encoder_cuda_graph_batch_sizes) + if not encoder_graph_shapes_available: + logger.warning( + "Feature-mode encoder CUDA graphs: no configured batch " + "size fits within encoder max_num_tokens " + f"({self.encoder_max_num_tokens}) // encoder output " + f"length ({self._encoder_fixed_seq_len}); the encoder " + "step stays eager.") + self._encoder_feature_shape = None + self._encoder_feature_dtype = None + self._encoder_fixed_seq_len = None + elif self._model_encoder_graph_spec() is not None: + encoder_graph_shapes_available = False + if self.encoder_cuda_graph_config is not None: + logger.warning( + "This model's encoder consumes fixed-shape features and " + "feature-mode encoder CUDA graphs are unavailable; the " + "encoder step stays eager.") + else: + encoder_graph_shapes_available = (bool(self._cuda_graph_num_tokens) + and bool( + self._cuda_graph_seq_lens)) + + use_encoder_cuda_graph = ((self._is_encoder_decoder_model() + or self._is_encode_only) + and self.encoder_cuda_graph_config is not None + and encoder_graph_shapes_available) self.torch_compile_config = self.llm_args.torch_compile_config torch_compile_enabled = bool(self.torch_compile_config is not None) @@ -870,45 +919,16 @@ def __init__( encoder_graph_batch_sizes = self._encoder_cuda_graph_batch_sizes encoder_graph_max_batch_size = (encoder_graph_batch_sizes[-1] if encoder_graph_batch_sizes else 0) - encoder_graph_max_num_tokens = self._max_cuda_graph_num_tokens - - # A feature-driven encoder (Whisper) declares a fixed-shape per-request - # contract instead of packed tokens, so its graph shapes follow from the - # batch sizes alone and `num_tokens` / `seq_lens` are ignored if set. - # Enablement is the same `encoder_cuda_graph_config` opt-in either way. - feature_shape, feature_dtype, fixed_seq_len = self._encoder_graph_spec() - if feature_shape is not None: - bs_cap = max( - 1, - min(self.batch_size, - self.encoder_max_num_tokens // fixed_seq_len)) - encoder_graph_batch_sizes = sorted( - bs for bs in encoder_graph_batch_sizes if bs <= bs_cap) - if not encoder_graph_batch_sizes: - logger.warning( - "Feature-mode encoder CUDA graphs: no configured batch " - f"size fits within max_num_tokens // encoder_output_len = " - f"{bs_cap}; the encoder step stays eager.") - feature_shape = feature_dtype = fixed_seq_len = None - # Without this the runner would be built in *token* mode with - # empty bucket lists: no capture keys, but `enabled` true and - # token-mode static buffers allocated for nothing. - use_encoder_cuda_graph = False - else: - encoder_graph_max_batch_size = encoder_graph_batch_sizes[-1] - encoder_graph_max_num_tokens = (encoder_graph_max_batch_size * - fixed_seq_len) - elif (use_encoder_cuda_graph - and self._model_encoder_graph_spec() is not None): - # Feature mode was declined above (TP > 1). This model's encoder - # cannot consume the packed token inputs the token-shaped capture - # path synthesizes, so keep it eager rather than let warmup drive - # tokens into it. - logger.warning( - "This model's encoder consumes fixed-shape features and " - "feature-mode encoder CUDA graphs are unavailable; the encoder " - "step stays eager.") - use_encoder_cuda_graph = False + # Feature mode's graph shapes follow from the batch sizes alone, so its + # token budget is one fixed-length encoder output per request; the + # token path uses the configured num_tokens buckets. + feature_shape = self._encoder_feature_shape + feature_dtype = self._encoder_feature_dtype + fixed_seq_len = self._encoder_fixed_seq_len + encoder_graph_max_num_tokens = (encoder_graph_max_batch_size * + fixed_seq_len + if feature_shape is not None else + self._max_cuda_graph_num_tokens) encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( use_cuda_graph=use_encoder_cuda_graph, @@ -3705,7 +3725,8 @@ def _is_encoder_decoder_model(self) -> bool: getattr(getattr(self.model, "model_config", None), "is_encoder_decoder", False)) - def _model_encoder_graph_spec(self): + def _model_encoder_graph_spec( + self) -> Optional[Tuple[Tuple[int, ...], torch.dtype, int]]: """The model's fixed-shape encoder contract, or None. Queried once.""" if not hasattr(self, "_cached_model_encoder_graph_spec"): # torch.compile wraps the model; the spec is on the original. @@ -3715,7 +3736,9 @@ def _model_encoder_graph_spec(self): is not None else None) return self._cached_model_encoder_graph_spec - def _encoder_graph_spec(self): + def _encoder_graph_spec( + self + ) -> Tuple[Optional[Tuple[int, ...]], Optional[torch.dtype], Optional[int]]: """Fixed-shape encoder contract, or (None, None, None) if unavailable. Returns ``(feature_shape, feature_dtype, fixed_seq_len)`` when the model @@ -7662,6 +7685,10 @@ def _pack_encoder_features(self, if (staging is None or staging.dtype != first.dtype or staging.shape[1:] != first.shape[1:] or staging.shape[0] < rows): + # Retire the previous batch's H2D before dropping the last + # reference to the buffer it reads from. + if staging is not None: + self._encoder_feature_staging_event.synchronize() staging = torch.empty((rows, *first.shape[1:]), dtype=first.dtype, pin_memory=prefer_pinned()) @@ -7669,8 +7696,11 @@ def _pack_encoder_features(self, self._encoder_feature_staging_event = torch.cuda.Event() # Dedicated copy stream: enqueued on the encoder stream the H2D # would queue behind the previous encoder forward, and the next - # batch's staging reuse would host-block on that forward. - self._encoder_feature_copy_stream = torch.cuda.Stream() + # batch's staging reuse would host-block on that forward. One + # stream for the runner's lifetime, so a reallocation cannot + # strand work on a stream nothing waits on again. + if getattr(self, '_encoder_feature_copy_stream', None) is None: + self._encoder_feature_copy_stream = torch.cuda.Stream() else: # The previous batch's H2D from this buffer must be complete # before its rows are overwritten. It ran on the copy stream, From f367474ecb577697f0f2b1dfac76d08e635f76f5 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:23:20 -0700 Subject: [PATCH 7/8] [TRTLLM-14778][test] Assert encoder graph replay, not just capture A populated `graphs` only proves warmup captured; `pad_batch` and the shape checks in `_maybe_forward_encoder_graph` can route every request to the eager encoder without emptying it. Count feature replays on the runner and assert it. Also cover the `encode_only` rejection and document that the 12.5% padding bound is unreachable for power-of-two bucket lists. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_whisper.py | 2 +- .../_torch/pyexecutor/cuda_graph_runner.py | 17 +++++++++++++++++ .../llmapi/test_llm_api_pytorch_whisper.py | 7 +++++-- tests/unittest/llmapi/test_llm_args.py | 18 ++++++++++++++++++ 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_whisper.py b/tensorrt_llm/_torch/models/modeling_whisper.py index edc3ae28f832..38d68d84b656 100644 --- a/tensorrt_llm/_torch/models/modeling_whisper.py +++ b/tensorrt_llm/_torch/models/modeling_whisper.py @@ -891,7 +891,7 @@ def __pp_init__(self): def config(self): return self.model_config.pretrained_config - def encoder_graph_spec(self): + def encoder_graph_spec(self) -> Tuple[Tuple[int, ...], torch.dtype, int]: """Fixed-shape encoder contract for enc-dec encoder CUDA graphs. Every Whisper encoder request is an fp32 waveform zero-padded by diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 85212dc4951d..3d5e76d74740 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1162,6 +1162,13 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): # H2D copies must be issued before graph replay instead of captured. self._capture_h2d_copy = prefer_pinned() + # Replays served from a captured feature graph. A populated `graphs` + # only proves capture happened; both `pad_batch` and the shape checks + # in `_maybe_forward_encoder_graph` can route every request to the + # eager encoder without emptying it, so tests need this to tell a + # working graph path from a silent eager fallback. + self.num_feature_replays = 0 + def _get_capture_stream(self) -> torch.cuda.Stream: """Return this runner's dedicated capture stream, creating it lazily.""" if self._capture_stream is None: @@ -1509,6 +1516,15 @@ def pad_batch(self, inputs: Dict[str, Any], # the 1-token pads of the token path. Fall back to eager across # large bucket gaps so graph replay cannot add more than 12.5% # encoder work. + # + # 12.5% is tight enough that consecutive power-of-two buckets never + # clear it: the next bucket is always >= 1.33x the current batch + # size. With the default generated bucket list, and with the + # [1, 2] the Whisper integration test configures, `enable_padding` + # is therefore inert and only exact batch sizes replay - which is + # what `_waiting_encoder_requests` forms microbatches to hit. + # Padding only becomes reachable with a dense list such as + # [1, 2, 3, 4]. if (batch_size == 0 or padded_batch_size > batch_size * self.MAX_FEATURE_PADDING_RATIO): yield inputs @@ -2007,6 +2023,7 @@ def _replay_features( self._feature_h2d_events[slot].record() self.graphs[key].replay() + self.num_feature_replays += 1 return self.graph_outputs[key] def replay( diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py index 0572ffa83d5c..4f10445b09e7 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py @@ -288,12 +288,15 @@ def _assert_cuda_graph_state(llm: LLM, captured: bool, encoder_captured: bool = assert not encoder_runner.enabled assert not encoder_runner.graphs return - # Capture must actually have happened: a silent fallback to the eager - # encoder path would otherwise pass every output assertion above. assert encoder_runner.enabled assert encoder_runner.graphs assert encoder_runner.feature_mode assert encoder_runner.is_encoder_decoder + # Capture alone is not enough: `pad_batch` and the shape checks in + # `_maybe_forward_encoder_graph` can route every request to the eager + # encoder while `graphs` stays populated, and that silent fallback would + # pass every output assertion above. Only the replay counter rules it out. + assert encoder_runner.num_feature_replays > 0 # Feature-combination matrix mirroring the T5/BART enc-dec coverage. Cases: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index e51da9bcdca9..719ce326aad8 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1915,6 +1915,24 @@ def test_encoder_cuda_graph_config_validation(self): ), ) + # `encoder_cuda_graph_config` is the encoder-decoder knob; an + # encode-only model configures its single forward through + # `cuda_graph_config` instead. + with pytest.raises( + ValidationError, + match="encoder_cuda_graph_config is for encoder-decoder"): + TorchLlmArgs( + model=llama_model_path, + encode_only=True, + encoder_max_batch_size=4, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ), + ) + def test_cuda_graph_config_infers_encode_mode_from_raw_dict(self): args = TorchLlmArgs( model=llama_model_path, From 76323d14afc74a7d6e6004d9aa01c0898cf420eb Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:18:17 -0700 Subject: [PATCH 8/8] [TRTLLM-14778][fix] Address review comments on encoder CUDA graphs Assert encoder replay against the warmup baseline: the capture pass replays each key once, so `> 0` was satisfied before any request was served. Require a token encoder's bucket lists at engine init instead of warning and running eager. Reject multi-row feature tensors, which would overrun the replay mirror. Stop admitting feature batches at a size the runner never captured. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- docs/source/models/encoder-decoder.md | 37 +++++++++++- .../_torch/pyexecutor/cuda_graph_runner.py | 9 ++- .../_torch/pyexecutor/model_engine.py | 60 ++++++++++++------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 ++- tensorrt_llm/llmapi/llm_args.py | 19 ++++-- .../llmapi/test_llm_api_pytorch_whisper.py | 7 ++- .../_torch/executor/test_py_executor.py | 15 +++++ .../executor/test_pytorch_model_engine.py | 36 +++++++++++ 8 files changed, 159 insertions(+), 33 deletions(-) diff --git a/docs/source/models/encoder-decoder.md b/docs/source/models/encoder-decoder.md index 8d09f6816282..57cbae010b02 100644 --- a/docs/source/models/encoder-decoder.md +++ b/docs/source/models/encoder-decoder.md @@ -45,7 +45,7 @@ The following table describes the supported and recommended configurations. | Beam search | Yes with V1 | Configure `max_beam_width` when constructing `LLM`, then set `use_beam_search=True` in `SamplingParams`. | | Attention backend | `TRTLLM` | Use this backend for encoder-decoder models. It is required when `tensor_parallel_size > 1`. | | Decoder CUDA graphs | Yes, except in FP32 | `CudaGraphConfig` captures decoder work. V1 supports greedy and beam search; V2 supports its single-beam path. FP32 encoder-decoder models decline capture at engine init and log a warning instead of failing. | -| Encoder CUDA graphs | Yes | Set `encoder_cuda_graph_config=EncodeCudaGraphConfig(...)` and `encoder_max_batch_size`. Usually set `encoder_max_batch_size` lower than `max_batch_size`. The `TRTLLM` attention backend is required. | +| Encoder CUDA graphs | Yes | Set `encoder_cuda_graph_config=EncodeCudaGraphConfig(...)` and `encoder_max_batch_size`. Usually set `encoder_max_batch_size` lower than `max_batch_size`. The `TRTLLM` attention backend is required. Text encoders also require `num_tokens` and `seq_lens`; a fixed-shape feature encoder such as Whisper derives both from the model and needs only `batch_sizes`. | | Overlap scheduler | Yes | Enabled by default. V1 supports greedy decoding and beam search; V2 remains limited to `max_beam_width=1`. | | Tensor parallelism | Yes | Use `tensor_parallel_size > 1` with `attn_backend="TRTLLM"`. Attention head counts must be divisible by the TP size. | | Pipeline parallelism | No | Keep `pipeline_parallel_size=1`. | @@ -437,6 +437,34 @@ size, total packed tokens, and maximum sequence length. The limit. With beam search, decoder graph batch sizes must cover the active decoder sequences after beam expansion. +Which encoder buckets you must supply depends on the model. A text encoder, +such as BART or T5, packs a variable number of tokens per request, so +`num_tokens` and `seq_lens` are part of its key space and are required; leaving +either unset is an error at engine construction. An encoder whose input is a +fixed-shape per-request feature tensor, such as Whisper's audio waveform, +produces the same number of encoder positions for every request, so both lists +follow from the model and are derived rather than configured. For those models +`batch_sizes` alone enables capture, and any `num_tokens` or `seq_lens` you set +is ignored: + +```python +from tensorrt_llm.llmapi import EncodeCudaGraphConfig + + +llm = LLM( + model="openai/whisper-large-v3", + backend="pytorch", + attn_backend="TRTLLM", + max_batch_size=8, + encoder_max_batch_size=8, + encoder_cuda_graph_config=EncodeCudaGraphConfig(batch_sizes=[1, 2, 4, 8]), + # ... the remaining Whisper settings from "Transcribe audio with Whisper" +) +``` + +Batch sizes that do not fit `encoder_max_num_tokens` divided by the model's +encoder output length are dropped, and the encoder stays eager when none fit. + `max_batch_size` controls the total decoder concurrency, while `encoder_max_batch_size` controls encoder microbatch admission. For better performance, tune `encoder_max_batch_size`, `encoder_max_num_tokens`, and the @@ -656,6 +684,13 @@ that the encoder graph buckets cover the request shape, and that `attn_backend="TRTLLM"`. Unsupported shapes and attention backends fall back to eager encoder execution. +### `num_tokens` or `seq_lens` unset is rejected at engine construction + +A text encoder needs both bucket lists, so the engine raises rather than +silently running eager. Supply them, or drop `encoder_cuda_graph_config` if you +do not want encoder graphs. A fixed-shape feature encoder such as Whisper does +not hit this: it derives both from the model and needs only `batch_sizes`. + ### Output quality differs from the Hugging Face example Confirm that the source uses the task prefix and language settings expected by diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 3d5e76d74740..df17cd2cdf45 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1154,7 +1154,9 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): # split-K GEMM captured into one graph spins forever in # cutlass::Semaphore::wait() once the other graph's matmuls leave that # region non-zero, because the captured graph has no node that re-zeroes - # it. Capture on our own stream instead. + # it. Capture on our own stream instead. The coupling is a property of + # the shared capture stream, not of the capture mode, so this applies to + # the token encoder path (T5/BART) as much as to feature mode. self._capture_stream: Optional[torch.cuda.Stream] = None # CUDA graph H2D memcpy nodes require pinned host sources. In CC mode @@ -1523,8 +1525,9 @@ def pad_batch(self, inputs: Dict[str, Any], # [1, 2] the Whisper integration test configures, `enable_padding` # is therefore inert and only exact batch sizes replay - which is # what `_waiting_encoder_requests` forms microbatches to hit. - # Padding only becomes reachable with a dense list such as - # [1, 2, 3, 4]. + # Padding is reachable only when a non-bucket batch has the next + # bucket within 12.5%, which first happens at a batch of 8 padding + # to a configured bucket of 9. if (batch_size == 0 or padded_batch_size > batch_size * self.MAX_FEATURE_PADDING_RATIO): yield inputs diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e7d51e4078d3..81f4ff67086e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -577,25 +577,8 @@ def __init__( self.encoder_cuda_graph_config.enable_padding if self.encoder_cuda_graph_config is not None else False) - # A fixed-shape feature encoder derives both bucket lists from the - # model's encoder output length, so only a token-driven encoder needs - # the user to supply them. - if (self.encoder_cuda_graph_config is not None - and self._model_encoder_graph_spec() is None - and (not encoder_cuda_graph_num_tokens - or not encoder_cuda_graph_seq_lens)): - missing = [] - if not encoder_cuda_graph_num_tokens: - missing.append("num_tokens/max_num_token") - if not encoder_cuda_graph_seq_lens: - missing.append("seq_lens/max_seq_len") - logger.warning( - f"Encoder CUDA graph configuration has " - f"{' and '.join(missing)} unset. Encoder CUDA graphs require " - f"both dimensions and will be disabled. " - f"To enable them, specify e.g. " - f"EncodeCudaGraphConfig(max_batch_size=64, num_tokens=[128, 256, " - f"512], max_seq_len=128, enable_padding=True).") + self._check_encoder_graph_bucket_config(encoder_cuda_graph_num_tokens, + encoder_cuda_graph_seq_lens) self._cuda_graph_padding_enabled = cuda_graph_padding_enabled @@ -944,7 +927,9 @@ def __init__( # The encoder runner takes its own graph pool. Encoder replay runs # on `encoder_stream`, device-concurrent with decoder replay, and # torch's pool-sharing contract assumes replays from a shared pool - # are not concurrent. + # are not concurrent. That concurrency exists for every + # encoder-decoder model, so the token encoder path (T5/BART) stops + # sharing the decoder pool too, at the cost of its own allocation. cuda_graph_mem_pool=None, is_encoder_decoder=self._is_encoder_decoder_model(), use_fixed_sequence_slots=(self._is_encoder_decoder_model() @@ -3736,6 +3721,35 @@ def _model_encoder_graph_spec( is not None else None) return self._cached_model_encoder_graph_spec + def _check_encoder_graph_bucket_config( + self, encoder_cuda_graph_num_tokens: List[int], + encoder_cuda_graph_seq_lens: List[int]) -> None: + """Reject an encoder graph config the model cannot complete. + + A fixed-shape feature encoder derives both bucket lists from the + model's encoder output length, so only a token-driven encoder needs the + user to supply them — and for that encoder the buckets are the whole + key space, so a config missing them can only ever run eager. Raise + rather than degrade silently: the request to capture was explicit. + """ + if (self.encoder_cuda_graph_config is None + or self._model_encoder_graph_spec() is not None): + return + missing = [] + if not encoder_cuda_graph_num_tokens: + missing.append("num_tokens/max_num_token") + if not encoder_cuda_graph_seq_lens: + missing.append("seq_lens/max_seq_len") + if not missing: + return + raise ValueError( + f"Encoder CUDA graph configuration has {' and '.join(missing)} " + f"unset. This model's encoder consumes packed tokens, so it needs " + f"both dimensions: specify e.g. " + f"EncodeCudaGraphConfig(max_batch_size=64, num_tokens=[128, 256, " + f"512], max_seq_len=128, enable_padding=True), or drop " + f"encoder_cuda_graph_config to run the encoder eagerly.") + def _encoder_graph_spec( self ) -> Tuple[Optional[Tuple[int, ...]], Optional[torch.dtype], Optional[int]]: @@ -7983,8 +7997,12 @@ def _maybe_forward_encoder_graph( features: List[torch.Tensor] = [] for request in encoder_requests: f = request.py_encoder_input_features + # Exactly one row per request: `_replay_features` copies + # `f.shape[0]` rows per request into a mirror slice the bucket + # sizes at one row per request, so a multi-row feature would + # overrun it. if (f is None or int(request.encoder_output_len) != fixed - or f.shape[1:] != runner.config.feature_shape + or tuple(f.shape) != (1, *runner.config.feature_shape) or f.dtype != runner.config.feature_dtype): return None features.append(f) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 1896a856a6ba..13d6e437e8d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5525,7 +5525,14 @@ def _waiting_encoder_requests( batch_size for batch_size in configured_batch_sizes if batch_size <= encoder_batch_size_limit ] - if (encoder_cuda_graph_config.enable_padding + # Targeting a size the runner never captured only pays off on the + # token path, where the pads are single tokens and the batch still + # rounds up into a captured bucket. A feature pad slot is a whole + # fixed_seq_len request, so `pad_batch` refuses any gap wider than + # 12.5% and such a batch would run eager; target the largest + # captured size within the limit instead. + if (not is_feature_encoder + and encoder_cuda_graph_config.enable_padding and any(batch_size > encoder_batch_size_limit for batch_size in configured_batch_sizes) and (not supported_batch_sizes diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 593dcd640465..f02500882a22 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -308,27 +308,35 @@ class EncodeCudaGraphConfig(BaseCudaGraphConfig): min_length=1, description= "List of total token counts (sum of all per-request sequence lengths " - "in a batch) to create encoder CUDA graphs for.") + "in a batch) to create encoder CUDA graphs for. Required for an " + "encoder that packs a variable number of tokens per request; ignored " + "by an encoder whose input is a fixed-shape per-request feature " + "tensor, which derives this from the model.") max_num_token: NonNegativeInt = Field( default=0, description="Maximum total number of tokens for encoder CUDA graphs. If " "`num_tokens` is provided, must equal max(num_tokens); otherwise " - "`num_tokens` is generated from this value.") + "`num_tokens` is generated from this value. Ignored by a fixed-shape " + "feature encoder.") seq_lens: Optional[List[PositiveInt]] = Field( default=None, min_length=1, description= "List of max per-request sequence lengths to create encoder CUDA " - "graphs for.") + "graphs for. Required for an encoder that packs a variable number of " + "tokens per request; ignored by an encoder whose input is a " + "fixed-shape per-request feature tensor, which derives this from the " + "model.") max_seq_len: NonNegativeInt = Field( default=0, description= "Maximum per-request sequence length for encoder CUDA graphs. If " "`seq_lens` is provided, must equal max(seq_lens); otherwise " - "`seq_lens` is generated from this value.") + "`seq_lens` is generated from this value. Ignored by a fixed-shape " + "feature encoder.") @model_validator(mode='after') def validate_encoder_cuda_graph_config(self) -> 'EncodeCudaGraphConfig': @@ -5195,7 +5203,8 @@ def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': # `num_tokens` / `seq_lens` are checked by the model engine rather than # here: an encoder whose input is a fixed-shape per-request feature # tensor derives both from the model, and only the engine knows which - # kind of encoder the model has. + # kind of encoder the model has. It still raises for the token encoders + # that require them. return self attn_backend: str = Field( diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py index 4f10445b09e7..57a073c2aff3 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py @@ -295,8 +295,11 @@ def _assert_cuda_graph_state(llm: LLM, captured: bool, encoder_captured: bool = # Capture alone is not enough: `pad_batch` and the shape checks in # `_maybe_forward_encoder_graph` can route every request to the eager # encoder while `graphs` stays populated, and that silent fallback would - # pass every output assertion above. Only the replay counter rules it out. - assert encoder_runner.num_feature_replays > 0 + # pass every output assertion above. Only the replay counter rules it out, + # and only against the warmup baseline: the capture pass replays each key + # once immediately after capturing it, so anything at or below + # `len(graphs)` is still explainable by warmup alone. + assert encoder_runner.num_feature_replays > len(encoder_runner.graphs) # Feature-combination matrix mirroring the T5/BART enc-dec coverage. Cases: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 8fa9f54cd595..eed875beb862 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -327,6 +327,21 @@ def test_encoder_microbatch_admission_uses_resolved_feature_batch_sizes( assert executor.encoder_batch_wait_iters_count == 0 +def test_encoder_microbatch_admission_skips_uncaptured_padded_size(): + # An `encoder_max_batch_size` above `max_batch_size` leaves a captured + # bucket beyond the admission limit. Padding admission up to the limit + # itself is a token-path move: a feature batch of 8 would have to pad to + # the captured 16, which `pad_batch` refuses, so it must target 4 instead. + executor = _make_feature_encoder_batch_wait_executor([1, 2, 4, 16], encoder_max_batch_size=16) + executor.max_batch_size = 8 + encoder_requests = [object() for _ in range(12)] + + scheduled = executor._waiting_encoder_requests(encoder_requests, [], [object()] * 2) + + assert scheduled == encoder_requests[:4] + assert executor.encoder_batch_wait_iters_count == 0 + + def test_encoder_microbatch_admission_ignores_disabled_feature_runner(): # supported_batch_sizes stays populated from the config even when capture # was declined (TP > 1, or no bucket fits), so waiting on those shapes diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 8abfbb181574..2e5ef9c80479 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1164,6 +1164,42 @@ def test_encoder_graph_spec_selection(self) -> None: self.assertEqual(engine._encoder_graph_spec(), expected if expected is not None else spec) + def test_encoder_graph_bucket_config_is_required_for_token_encoders( + self) -> None: + # A token encoder's num_tokens/seq_lens buckets are the whole key + # space, so a config missing them can only run eager — a loud failure, + # not a silent perf regression. A feature encoder derives both from the + # model, so the same config is complete there. + engine, _ = self._encoder_spec_engine( + EncodeCudaGraphConfig(batch_sizes=[1, 2]), declares_spec=False) + with self.assertRaisesRegex( + ValueError, "num_tokens/max_num_token and " + "seq_lens/max_seq_len"): + engine._check_encoder_graph_bucket_config([], []) + + engine, _ = self._encoder_spec_engine(EncodeCudaGraphConfig( + batch_sizes=[1, 2], num_tokens=[1500]), + declares_spec=False) + with self.assertRaisesRegex(ValueError, "seq_lens/max_seq_len unset"): + engine._check_encoder_graph_bucket_config([1500], []) + + for name, config, declares_spec in [ + ("token model with both buckets", + EncodeCudaGraphConfig(batch_sizes=[1], + num_tokens=[1500], + seq_lens=[1500]), False), + ("feature model derives both", + EncodeCudaGraphConfig(batch_sizes=[1, 2]), True), + ("no config", None, False), + ]: + with self.subTest(name): + engine, _ = self._encoder_spec_engine( + config, declares_spec=declares_spec) + num_tokens = config.num_tokens if config else [] + seq_lens = config.seq_lens if config else [] + engine._check_encoder_graph_bucket_config( + num_tokens or [], seq_lens or []) + def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( self) -> None: runner = EncoderCUDAGraphRunner.__new__(EncoderCUDAGraphRunner)