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/models/modeling_whisper.py b/tensorrt_llm/_torch/models/modeling_whisper.py index 420b7ea6f9de..38d68d84b656 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) -> 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 + `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 9f626eba6b8a..df17cd2cdf45 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1053,6 +1053,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. @@ -1067,6 +1076,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 @@ -1075,14 +1088,41 @@ 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: + # 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) @@ -1107,11 +1147,36 @@ 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. 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 # 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() + # 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: + 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 = ( @@ -1119,6 +1184,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), @@ -1406,6 +1512,33 @@ 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. + # + # 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 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 + 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 @@ -1721,6 +1854,58 @@ def capture( inputs: Dict[str, Any], ) -> Any: """Warm up and/or capture the forward pass for a graph key.""" + capture_inputs, capture_h2d = (self._prepare_feature_capture( + key, inputs) if self.feature_mode else self._prepare_token_capture( + key, inputs)) + + self.graph_metadata[key] = { + "attn_metadata": capture_inputs["attn_metadata"] + } + + 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) + + 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(), + capture_error_mode="thread_local"): + if capture_h2d is not None: + capture_h2d() + output = forward_fn(capture_inputs) + + if self._contains_nested_tensor(output): + raise TypeError( + "Encoder CUDA graph does not support nested tensor outputs. " + "Disable encoder CUDA graphs for models with ragged 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 _prepare_token_capture( + self, + key: EncoderKeyType, + inputs: Dict[str, Any], + ) -> Tuple[Dict[str, Any], Optional[Callable[[], None]]]: + """Capture setup for the packed-token mode. + + 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_num_tokens = key[1] sliced_static_tensors = { @@ -1739,10 +1924,13 @@ def capture( capture_inputs = dict(inputs) 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 @@ -1751,52 +1939,47 @@ def capture( # 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) + copy_inputs() 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. - for _ in range(self.WARMUP_STEPS): - output = forward_fn(capture_inputs) + if not self._capture_h2d_copy: + return capture_inputs, None - if self.is_warmup_only: - return output + def capture_h2d() -> None: + copy_inputs() + attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, - pool=self.memory_pool, - 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) - output = forward_fn(capture_inputs) + return capture_inputs, capture_h2d - if self._contains_nested_tensor(output): - raise TypeError( - "Encoder CUDA graph does not support nested tensor outputs. " - "Disable encoder CUDA graphs for models with ragged 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 _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) + + return capture_inputs, None def retire_staging(self) -> None: """Wait until a prior replay no longer reads shared staging buffers.""" @@ -1804,12 +1987,61 @@ def retire_staging(self) -> None: self._staging_retirement_event.synchronize() self._staging_retirement_event = None + 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() + + # Per-request CPU tensors straight from the requests — one copy into + # the host mirror, no intermediate packing. + mirror = self._feature_mirrors[slot] + 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_() + + # 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() + self.num_feature_replays += 1 + 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 4906ad52f5e5..81f4ff67086e 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 @@ -567,27 +577,15 @@ def __init__( 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 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 + 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 @@ -595,9 +593,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 []) @@ -618,11 +625,41 @@ def __init__( self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) + # 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 bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens)) + 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) @@ -865,7 +902,17 @@ 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 + # 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, cuda_graph_padding_enabled=( @@ -877,15 +924,29 @@ 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. 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() 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. @@ -2102,23 +2163,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 @@ -2379,10 +2453,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 +3710,75 @@ def _is_encoder_decoder_model(self) -> bool: getattr(getattr(self.model, "model_config", None), "is_encoder_decoder", False)) + 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. + 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 + + 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]]: + """Fixed-shape encoder contract, or (None, None, None) if unavailable. + + 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 + or not self._is_encoder_decoder_model()): + return none + + spec = self._model_encoder_graph_spec() + if spec is None: + return none + + 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) @@ -7517,17 +7670,71 @@ 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): + # 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()) + 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. 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, + # 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, @@ -7760,12 +7967,118 @@ 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 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 + # 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 tuple(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) + 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. + 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_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) + capture_output = runner.capture( + key, self._enc_dec_encoder_graph_forward_fn, padded_inputs) + + 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: + 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 _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 00037596ad4d..13d6e437e8d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5501,19 +5501,38 @@ 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 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 bool(encoder_cuda_graph_config.num_tokens) - and bool(encoder_cuda_graph_config.seq_lens)): + 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) - configured_batch_sizes = (encoder_cuda_graph_config.batch_sizes - or []) + if is_feature_encoder: + # Feature batch sizes may have been derived rather than + # configured, so take the ones the runner actually resolved. + # 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.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 ] - 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 01696cdb6068..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': @@ -5192,14 +5200,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.") - 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. 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 bb72ccb7121a..57a073c2aff3 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, + EncodeCudaGraphConfig, + 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,33 @@ 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; 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": EncodeCudaGraphConfig( + 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 +162,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 +269,66 @@ 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 + 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, + # 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: # (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 +336,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 +347,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 +367,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 +386,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_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index ff242644533f..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: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 161ac02eb995..eed875beb862 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -43,6 +43,7 @@ ScheduledRequests, SerializableSchedulerOutput, ) +from tensorrt_llm.llmapi.llm_args import EncodeCudaGraphConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError pytestmark = pytest.mark.cpu_only @@ -214,6 +215,38 @@ 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 + + +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. + + 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=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, + feature_mode=True, + ) + ) executor.batch_wait_timeout_iters = 48 executor.encoder_batch_wait_iters_count = 0 return executor @@ -225,6 +258,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 @@ -263,6 +297,69 @@ def test_encoder_graph_warmup_uses_runtime_encoder_stream(): ) +@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, + [], + [object()] * 20, + ) + + assert scheduled == encoder_requests[:expected] + 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 + # 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 ed0f3bd4d00b..2e5ef9c80479 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -18,13 +18,15 @@ 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, SeqLenAwareSparseAttentionConfig, TorchLlmArgs) @@ -1075,6 +1077,129 @@ 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), 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( + runner._capture_sequence_lengths, + {(bs, bs * fixed, fixed): [fixed] * bs + for bs in batch_sizes}, + ) + self.assertEqual(runner.capture_keys, + frozenset(runner._capture_sequence_lengths)) + + @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_selection(self) -> None: + # The model selects feature mode, not the config: an encoder either + # 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_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) 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 ceef1c8de786..719ce326aad8 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1883,41 +1883,55 @@ 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", - ), - ( - { - "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"), + # 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, + ), + ) + + # `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(