diff --git a/docs/source/models/encoder-decoder.md b/docs/source/models/encoder-decoder.md index 8d09f6816282..1cccfa4bfd7e 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,47 @@ 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 rejected at `LLM(...)` for any architecture TensorRT-LLM +recognizes, and at model engine initialization otherwise, since only the loaded +model states with certainty which kind of encoder it has. An encoder whose input is a +fixed-shape per-request feature tensor, such as Whisper's fixed 30-second +zero-padded audio waveform (the mel transform runs inside the encoder, so the +per-request input is the waveform itself, not a spectrogram), 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. + +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. +Size the encoder token budget for the largest bucket before setting the +buckets: Whisper emits 1500 encoder positions per request, so `batch_sizes` up +to 8 needs `encoder_max_num_tokens` of at least 12000. `encoder_max_num_tokens` +falls back to `max_num_tokens` when unset, which is a decoder-sized number and +usually too small. + +```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, + # 8 buckets * 1500 encoder positions. Leave this at the default and the + # 4 and 8 buckets are silently dropped. + encoder_max_num_tokens=12000, + encoder_cuda_graph_config=EncodeCudaGraphConfig(batch_sizes=[1, 2, 4, 8]), + # ... the remaining Whisper settings from "Transcribe audio with Whisper", + # whose `max_batch_size=4` this example raises to 8 +) +``` + `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 +697,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..275ffe69b4e4 100644 --- a/tensorrt_llm/_torch/models/modeling_whisper.py +++ b/tensorrt_llm/_torch/models/modeling_whisper.py @@ -381,6 +381,12 @@ def __init__(self, config: WhisperConfig): extractor = _load_hf_feature_extractor(config) self.n_fft = int(extractor.n_fft) self.hop_length = int(extractor.hop_length) + # The extractor's own padded window length, kept verbatim because + # `WhisperInputProcessor` pads every request to exactly this many + # samples. `encoder_graph_spec` reports it as the fixed encoder input + # shape, so deriving it a second way would risk disagreeing with the + # tensors the processor actually produces. + self.n_samples = int(extractor.n_samples) # Pre-STFT Gaussian noise, applied where HF applies it; 0.0 (all # official checkpoints) disables it. self.dither = float(getattr(extractor, "dither", 0.0)) @@ -654,6 +660,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._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 +897,29 @@ 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 extractor's window, which yields exactly + ``max_source_positions`` encoder positions — so the encoder graph key + degenerates to the batch size. + + The window is read from the feature extractor rather than recomputed. + ``max_source_positions * 2 * hop_length`` inverts the conv stem + correctly but only up to the truncation in the processor's own + ``n_samples // hop_length // 2`` check, so a checkpoint whose + ``n_samples`` is not an exact multiple of ``2 * hop_length`` would + pass that check while disagreeing with this shape — capturing graphs + that every runtime batch then misses. Taking the extractor's value + keeps both sides on one number. + + Returns ``(per_request_feature_shape, dtype, fixed_seq_len)``. + """ + fixed_seq_len = int(self.config.max_source_positions) + n_samples = int(self.model.encoder.log_mel.n_samples) + 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 32ea096b621f..c7a9a0c6e240 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1111,6 +1111,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. @@ -1125,6 +1134,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 @@ -1133,14 +1146,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) @@ -1165,18 +1205,84 @@ 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() - def _create_shared_static_tensors(self): + # 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) -> None: """Allocates static tensors sized for the largest supported num_tokens.""" max_total_tokens = ( self.config.max_num_tokens if self.is_encoder_decoder else min( 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), @@ -1464,6 +1570,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 @@ -1523,6 +1656,57 @@ def prepare_encoder_decoder_inputs( source_sequence_lengths)] return prepared_inputs + def _resolve_graph_key(self, inputs: Dict[str, + Any]) -> Optional[EncoderKeyType]: + """The capture key `inputs` maps to, or None if no graph can serve it. + + Everything here is derived from `inputs` alone, so it is safe to ask + before building attention metadata. + """ + if not self.enabled or ExpertStatistic.should_record(): + return None + + if len(inputs['seq_lens']) not in self.supported_batch_sizes: + return 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: + return None + if (not self.padding_enabled and is_padding_performed) \ + or not is_padding_successful: + return None + return key + + def _captured_metadata_for_key(self, key: EncoderKeyType) -> Optional[Any]: + """Graph-resident metadata for an already-resolved `key`, or None.""" + if key not in self.graph_metadata: + return None + # Token-path graph keys all alias the same host staging buffers, so + # retire a prior graph's captured reads before the caller updates + # them. Feature mode stages elsewhere and this is a no-op there. + self.retire_staging() + return self.graph_metadata[key]["attn_metadata"] + + def captured_graph_metadata( + self, + inputs: Dict[str, Any], + ) -> Tuple[Optional[Any], Optional[EncoderKeyType]]: + """Graph-resident metadata for `inputs`, if its key is already captured. + + Resolvable from `inputs` alone, so the runtime path can call this + before building eager attention metadata that a graph hit would never + read. Returns (None, None) on a miss, leaving the caller to build that + metadata and take the full path. + """ + key = self._resolve_graph_key(inputs) + if key is None: + return None, None + graph_attn_metadata = self._captured_metadata_for_key(key) + if graph_attn_metadata is None: + return None, None + return graph_attn_metadata, key + def maybe_get_cuda_graph( self, inputs: Dict[str, Any], @@ -1551,28 +1735,13 @@ def maybe_get_cuda_graph( key="encoder_cuda_graph_backend_warning") return None, None - if ExpertStatistic.should_record(): - return None, None - - seq_lens = inputs['seq_lens'] - padded_batch_size = len(seq_lens) - if padded_batch_size not in self.supported_batch_sizes: - 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: - return None, None - padded_max_seq_len = key[2] - if (not self.padding_enabled and is_padding_performed) \ - or not is_padding_successful: + key = self._resolve_graph_key(inputs) + if key is None: return None, None - if key in self.graph_metadata: - # Every graph key aliases the same host staging buffers. Retire a - # prior graph's captured reads before the caller updates them. - self.retire_staging() - return self.graph_metadata[key]["attn_metadata"], key + graph_attn_metadata = self._captured_metadata_for_key(key) + if graph_attn_metadata is not None: + return graph_attn_metadata, key # New key not yet captured. Only create graph metadata during explicit # startup warmup; unseen runtime keys fall back to eager execution. @@ -1596,6 +1765,8 @@ def maybe_get_cuda_graph( # First sighting of this key: create graph-resident metadata and bind # it to stable pinned seq_lens storage for future replays. + padded_batch_size = len(inputs['seq_lens']) + padded_max_seq_len = key[2] graph_attn_metadata = attn_metadata.create_cuda_graph_metadata( padded_batch_size, False, @@ -1779,6 +1950,61 @@ 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() + # Do not keep the eager result live from this runner across graph + # setup/capture; release its reference before entering. + output = None + 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 = { @@ -1797,10 +2023,14 @@ 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: + """Refill the captured device token inputs from their host mirrors.""" + 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 @@ -1809,55 +2039,53 @@ 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: + """Stage this replay's token inputs and sequence lengths onto the device. - graph = torch.cuda.CUDAGraph() - # Do not keep the eager result live from this runner across graph - # setup/capture; release its reference before entering. - output = None - 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) + Returned to the caller so the copies are issued per replay, outside + the captured graph, rather than being baked into it. + """ + copy_inputs() + attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) - 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 + 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: mirrors are + pooled across buckets rather than owned per bucket + (`FEATURE_MIRROR_SLOTS` of them, rotated), and consecutive encoder + batches can be enqueued back-to-back, so a captured H2D would read a + mirror at replay-execution time, after the host had already rotated + back onto it. 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.""" @@ -1865,12 +2093,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 359c16f6322b..e37a0477bb11 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -42,6 +42,9 @@ PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) + +# isort: split +from tensorrt_llm.llmapi.llm_args import validate_token_encoder_bucket_config from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping @@ -260,11 +263,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 @@ -678,27 +691,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 @@ -706,9 +707,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 []) @@ -726,11 +736,41 @@ def __init__( self._encoder_cuda_graph_padding_enabled) if encoder_cuda_graph_seq_lens else []) + # 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 self.prefill_cuda_graph_backend = self.llm_args.prefill_cuda_graph_backend @@ -972,7 +1012,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=( @@ -984,15 +1034,35 @@ 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 must never be handed the decoder's 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. + # + # Literal None rather than `self._cuda_graph_mem_pool` states that + # as a requirement instead of leaving it to a coincidence. The + # engine attribute is None for the engine's whole life (nothing + # assigns it after its declaration), so each runner already + # created its own pool at its first capture and this allocates + # nothing new — but reading it here would silently start sharing + # the day someone gives that attribute a value. + 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. @@ -2240,23 +2310,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 @@ -2541,10 +2624,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: @@ -4029,6 +4122,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 feature encoder derives both bucket lists from the model, so only a + token encoder needs them supplied — and there they are the whole key + space, so a config missing them can only run eager. The request to + capture was explicit, so raise rather than degrade silently. + + Encode-only warns instead: its buckets arrive through + `cuda_graph_config` (see `__init__`), a slot that has always accepted a + batch-sizes-only config and run eager, so raising would break + deployments predating feature mode. + """ + if (self.encoder_cuda_graph_config is None + or self._model_encoder_graph_spec() is not None): + return + bucket_config_error = validate_token_encoder_bucket_config( + encoder_cuda_graph_num_tokens, + encoder_cuda_graph_seq_lens, + stays_eager=self._is_encode_only) + if bucket_config_error is None: + return + if self._is_encode_only: + logger.warning(bucket_config_error) + return + raise ValueError(bucket_config_error) + + 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) @@ -7787,17 +7949,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, @@ -8029,12 +8245,151 @@ 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): + # A shape the model's `encoder_graph_spec()` did not predict + # misses on every request, not just this one: capture already + # spent its time and memory and nothing will ever replay. Say + # so once — silence here reads as "graphs are working". + logger.warning_once( + "Encoder CUDA graph: request features do not match the " + "captured contract (expected shape " + f"{(1, *runner.config.feature_shape)} dtype " + f"{runner.config.feature_dtype} encoder_output_len " + f"{fixed}, got shape " + f"{None if f is None else tuple(f.shape)} dtype " + f"{None if f is None else f.dtype} encoder_output_len " + f"{int(request.encoder_output_len)}); the encoder step " + "stays eager.", + key="encoder_cuda_graph_feature_contract_warning") + 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)) + ] + # A captured bucket is served from graph-resident metadata, so + # only a miss pays for a fresh `TrtllmAttentionMetadata` + + # `prepare_encoder_only()`. + graph_attn_metadata, key = runner.captured_graph_metadata( + padded_inputs) + if key is None: + 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: + """Run the encoder step over a graph runner's capture inputs. + + Adapts the runner's flat capture dict to `_forward_step_encoder`'s + keyword names. Passed to `EncoderCUDAGraphRunner.capture` so the body + traced into the graph is the same code path a replay stands in for. + + Args: + capture_inputs: Padded encoder inputs owned by the graph runner. + + Returns: + Encoder hidden states, `[padded_batch, fixed_seq_len, hidden]`. + """ + 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 a27a2a417275..c4c930565e78 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6074,19 +6074,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.py b/tensorrt_llm/llmapi/llm.py index c1168a105f68..fbdcc2bc9c94 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -56,7 +56,8 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams -from .llm_args import TORCH_LLMARGS_EXPLICIT_DOCSTRING, TorchLlmArgs +from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING, TorchLlmArgs, + validate_token_encoder_bucket_config) from .llm_utils import (CachedModelLoader, KvCacheRetentionConfig, LlmBuildStats, ModelLoader) from .mpi_session import MpiPoolSession, external_mpi_comm_available @@ -1821,6 +1822,45 @@ def _collective_rpc( f"Executor type {type(self._executor)} does not support collective RPC." ) + def _reject_token_encoder_config_without_buckets(self) -> None: + """Reject a bucket-less token-encoder config before weights load. + + `PyTorchModelEngine.__init__` asks the instantiated model and keeps the + last word; this asks the class its architecture resolves to. It defers + wherever that class might not be the one the engine builds, so it can + only reject earlier, never differently. A model that installs + `encoder_graph_spec` at construction instead of declaring it on the + class would be misjudged here; no in-tree model does. + """ + config = self.args.encoder_cuda_graph_config + # `checkpoint_loader`, a non-HF `checkpoint_format` and `model_kwargs` + # all feed `checkpoint_loader.load_config()`, which is where the engine + # gets its class from, so the on-disk architecture may not be the one + # it builds. Decoder-only models are left to the engine too: its + # "consumes packed tokens" wording would only confuse there. + if (config is None or not self._is_encoder_decoder_model() + or self.args.model_kwargs is not None + or self.args.checkpoint_loader is not None + or self.args.checkpoint_format not in (None, "HF")): + return + bucket_config_error = validate_token_encoder_bucket_config( + config.num_tokens, config.seq_lens) + if bucket_config_error is None: + return + architectures = getattr(self._hf_model_config, "architectures", + None) if self._hf_model_config else None + if not architectures: + return + # Local: resolving an architecture imports its model module, which the + # client process otherwise never loads. + from tensorrt_llm._torch.models.modeling_utils import \ + get_registered_model_class + model_cls = get_registered_model_class(architectures[0]) + # Unresolved, or a feature encoder that derives both dimensions itself. + if model_cls is None or hasattr(model_cls, "encoder_graph_spec"): + return + raise ValueError(bucket_config_error) + def _build_model(self): super()._build_model() assert self._engine_dir is None @@ -1829,6 +1869,7 @@ def _build_model(self): # It should also be before bindings ExecutorConfig, which may depend on tokenizer info. self._tokenizer = self._try_load_tokenizer() self._hf_model_config = self._try_load_hf_model_config() + self._reject_token_encoder_config_without_buckets() self._generation_config = self._try_load_generation_config() self._generation_config_explicit_values = self._try_load_generation_config_explicit_values( ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 645d715e721a..080cc75586df 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': @@ -446,6 +454,36 @@ def _generate_cuda_graph_seq_lens(max_seq_len: int, return sizes +def validate_token_encoder_bucket_config( + num_tokens: Optional[List[int]], + seq_lens: Optional[List[int]], + *, + stays_eager: bool = False) -> Optional[str]: + """Why a packed-token encoder cannot capture with these buckets, or None. + + Answers for token encoders (T5, BART) only; a fixed-shape feature encoder + derives both dimensions from the model, so the caller establishes which + kind it has first. Returned rather than raised because callers differ: + encode-only warns and stays eager (`stays_eager`), the rest raise. + """ + missing = [] + if not num_tokens: + missing.append("num_tokens/max_num_token") + if not seq_lens: + missing.append("seq_lens/max_seq_len") + if not missing: + return None + head = (f"Encoder CUDA graph configuration has {' and '.join(missing)} " + "unset. This model's encoder consumes packed tokens, so it needs " + "both dimensions") + if stays_eager: + return f"{head}; the encode step stays eager." + return (f"{head}: specify e.g. EncodeCudaGraphConfig(max_batch_size=64, " + "num_tokens=[128, 256, 512], max_seq_len=128, " + "enable_padding=True), or drop encoder_cuda_graph_config to run " + "the encoder eagerly.") + + # For CudaGraphConfig's backward compatibility CudaGraphConfig = DecodeCudaGraphConfig @@ -5331,15 +5369,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 533cfb269ca6..90ea9d6128f3 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,14 +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; requesting graphs must - # work for every dtype. + """Build a Whisper LLM for the test matrix, optionally with encoder CUDA graphs.""" + # CudaGraphConfig captures the decode step; the enc-dec encoder step opts in + # separately through `encoder_cuda_graph_config`. Requesting decoder graphs + # must 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 @@ -133,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, ) @@ -239,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 -# covers fp32 enc-dec capturing decoder graphs. +# 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 covers fp32 enc-dec capturing decoder graphs. 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], True, 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], True, 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), ), @@ -278,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( @@ -288,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. @@ -307,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): @@ -325,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 acaf44ff5818..b25b3a99733b 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -48,9 +48,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 607f2b0871bc..48d6f7a01434 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -49,7 +49,7 @@ ScheduledRequests, SerializableSchedulerOutput, ) -from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig +from tensorrt_llm.llmapi.llm_args import EncodeCudaGraphConfig, MTPDecodingConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError pytestmark = pytest.mark.cpu_only @@ -202,6 +202,7 @@ def _make_async_encoder_executor(future): def _make_encoder_batch_wait_executor(batch_sizes=None, encoder_max_batch_size=8): + """Build a PyExecutor stub wired for token-path encoder batch-wait admission.""" executor = object.__new__(PyExecutor) executor.max_batch_size = 32 batch_sizes = batch_sizes or [1, 2, 4, 8] @@ -214,17 +215,51 @@ 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 def _make_encoder_fallback_batch_wait_executor(): + """Build a PyExecutor stub with no encoder graph config, for the fallback path.""" executor = object.__new__(PyExecutor) executor.llm_args = types.SimpleNamespace( 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 +298,72 @@ 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 +): + """Feature admission targets the runner's resolved sizes, not the configured ones.""" + 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(): + """Feature admission never targets a batch size the runner did not capture.""" + # 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(): + """A runner that declined capture must not make admission wait on unreplayable shapes.""" + # 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 193652e7c001..ad8e35934b78 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -5,6 +5,7 @@ from contextlib import nullcontext from dataclasses import dataclass from types import SimpleNamespace +from typing import List, Optional, Tuple from unittest.mock import Mock, patch import torch @@ -19,13 +20,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) + _filter_cuda_graph_batch_sizes, _make_single_token_context_graph_batch) from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + EncodeCudaGraphConfig, PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, TorchLlmArgs) @@ -1196,6 +1199,228 @@ def test_global_incompatibilities_bypass_candidate_selection(self) -> None: class PyTorchModelEngineTestCase(unittest.TestCase): + @staticmethod + def _feature_encoder_runner( + batch_sizes: List[int], + fixed_seq_len: int = 1500) -> EncoderCUDAGraphRunner: + """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 feature capture key matches a batch shape the runtime can actually produce.""" + # 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: Optional[EncodeCudaGraphConfig], + declares_spec: bool, + tp_size: int = 1, + is_encode_only: bool = False + ) -> Tuple[PyTorchModelEngine, Tuple[Tuple[int, ...], torch.dtype, int]]: + """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) -> Tuple[Tuple[int, ...], torch.dtype, int]: + """Stand-in fixed-shape encoder contract for the test model.""" + return spec + + engine = PyTorchModelEngine.__new__(PyTorchModelEngine) + engine.encoder_cuda_graph_config = encoder_cuda_graph_config + engine.is_draft_model = False + engine._is_encode_only = is_encode_only + engine.model = _Model() + engine.mapping = SimpleNamespace(tp_size=tp_size) + return engine, spec + + def test_encoder_graph_spec_selection(self) -> None: + """The model, not the config, selects feature mode; TP > 1 stays eager.""" + # 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 missing its bucket lists fails loudly instead of silently running eager.""" + # 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_graph_bucket_config_warns_for_encode_only(self) -> None: + """An encode-only model warns and stays eager rather than raising.""" + # An encode-only model receives its buckets through `cuda_graph_config`, + # a slot that has always accepted a batch-sizes-only + # EncodeCudaGraphConfig and run eager. Raising there would break + # deployments that predate feature mode, so warn and stay eager. + engine, _ = self._encoder_spec_engine( + EncodeCudaGraphConfig(batch_sizes=[1, 2]), + declares_spec=False, + is_encode_only=True) + with patch("tensorrt_llm._torch.pyexecutor.model_engine.logger.warning" + ) as warning: + engine._check_encoder_graph_bucket_config([], []) + warning.assert_called_once() + self.assertIn("stays eager", warning.call_args.args[0]) + + def test_feature_encoder_batch_sizes_drop_past_the_token_budget( + self) -> None: + """The encoder token budget caps feature bucket sizes; an empty list means stay eager.""" + # A feature request costs a whole fixed_seq_len against the encoder + # token budget, so the budget caps the bucket list far below + # encoder_max_batch_size. An empty result is the signal to stay eager; + # a floor of 1 here would capture a graph larger than the metadata + # budget the encoder step actually builds. + fixed = 1500 + for name, max_batch_size, max_num_tokens, expected in [ + ("budget allows every bucket", 8, 8 * fixed, [1, 2, 4, 8]), + ("budget truncates the tail", 8, 2 * fixed, [1, 2]), + ("budget below one request", 8, fixed - 1, []), + ("batch size caps below the budget", 2, 8 * fixed, [1, 2]), + ]: + with self.subTest(name): + self.assertEqual( + _filter_cuda_graph_batch_sizes([1, 2, 4, 8], + max_batch_size, + max_num_tokens, + fixed, + enable_padding=False), + expected) + + def test_feature_pad_batch_refuses_wide_bucket_gaps(self) -> None: + """Feature padding is refused once it would cost more than 12.5% extra work.""" + # A feature pad slot is a full fixed_seq_len encoder forward, unlike the + # 1-token pads of the token path, so padding is bounded at 12.5% extra + # work. Consecutive powers of two never clear that bound; a batch of 8 + # padding to a configured bucket of 9 is the first case that does. + fixed = 1500 + runner = self._feature_encoder_runner([1, 2, 4, 9], fixed) + runner.enabled = True + + for name, batch_size, expected_seq_lens in [ + ("exact bucket yields unchanged", 4, [fixed] * 4), + ("8 -> 9 is within 12.5%", 8, [fixed] * 9), + ("3 -> 4 exceeds 12.5%", 3, [fixed] * 3), + ("5 -> 9 exceeds 12.5%", 5, [fixed] * 5), + ]: + with self.subTest(name): + inputs = {'seq_lens': [fixed] * batch_size} + with runner.pad_batch(inputs, batch_size) as padded: + self.assertEqual(padded['seq_lens'], expected_seq_lens) + + def test_captured_graph_metadata_skips_the_eager_metadata_build( + self) -> None: + """A graph hit resolves captured metadata without an eager attn_metadata build.""" + # The runtime path asks for captured metadata before building any, so a + # graph hit must not need an attn_metadata argument at all: on a hit + # `maybe_get_cuda_graph` only reads it for a backend check. + fixed = 1500 + runner = self._feature_encoder_runner([1, 2], fixed) + runner.enabled = True + runner.retire_staging = Mock() + sentinel = object() + key = (2, 2 * fixed, fixed) + runner.graph_metadata[key] = {"attn_metadata": sentinel} + + metadata, hit_key = runner.captured_graph_metadata( + {'seq_lens': [fixed] * 2}) + self.assertIs(metadata, sentinel) + self.assertEqual(hit_key, key) + runner.retire_staging.assert_called_once() + + # An uncaptured bucket must miss, leaving the caller to build metadata + # and take the full path rather than silently reusing another key's. + runner.retire_staging.reset_mock() + self.assertEqual(runner.captured_graph_metadata({'seq_lens': [fixed]}), + (None, None)) + runner.retire_staging.assert_not_called() + 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 11f48edaf444..598cbb3223d9 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -179,6 +179,7 @@ class TestWarmupCleanup(unittest.TestCase): """Lock in warmup-cleanup behavior introduced by PR #14609 (Plan B).""" def test_encoder_decoder_encoder_warmup_is_deferred_and_uses_two_passes(self): + """Enc-dec encoder warmup is deferred and runs as two passes.""" model_engine = object.__new__(PyTorchModelEngine) model_engine.cuda_graph_runner = SimpleNamespace( enabled=True, @@ -196,6 +197,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 f1713d6c40df..164170c82e21 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -7,6 +7,7 @@ from dataclasses import is_dataclass from enum import Enum from pathlib import Path +from types import SimpleNamespace from typing import Annotated, Any, ClassVar, Literal, get_args, get_origin from unittest.mock import patch @@ -1907,6 +1908,7 @@ def test_cuda_graph_config_accepts_encoder_config(self): assert args.cuda_graph_config.max_seq_len == 32 def test_encoder_decoder_cuda_graph_user_interface(self): + """EncodeCudaGraphConfig round-trips through TorchLlmArgs as the user writes it.""" encoder_config = EncodeCudaGraphConfig( batch_sizes=[1, 4], num_tokens=[16, 64], @@ -1937,41 +1939,154 @@ 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, so which kind of encoder the model has decides + # whether they are required — a question the config cannot answer. + # `LLM._reject_token_encoder_config_without_buckets` asks the model + # class, and the model engine asks the loaded model. + feature_args = TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=4, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + enable_padding=True, ), - ( - { - "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"), + ) + + 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 + + def test_encoder_cuda_graph_config_validation(self): + """encoder_cuda_graph_config is rejected without encoder_max_batch_size.""" + 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, + ), + ) + + @staticmethod + def _bucketless_encoder_llm(architectures, is_encoder_decoder=True): + """A bare LLM carrying only what the bucket pre-check reads.""" + llm = TorchLLM.__new__(TorchLLM) + llm.args = TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=4, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + enable_padding=True, ), - ] + ) + llm._hf_model_config = (SimpleNamespace( + architectures=architectures, is_encoder_decoder=is_encoder_decoder) + if architectures is not None else None) + return llm - for kwargs, error_match in invalid_cases: - with pytest.raises(ValidationError, match=error_match): - TorchLlmArgs( - model=llama_model_path, - **kwargs, - ) + def test_token_encoder_without_buckets_is_rejected_before_weights_load( + self): + """A token encoder missing buckets is rejected at config time, before weights load.""" + + # The model class settles token-vs-feature without an instance, so + # the user learns now rather than after weights load. + class _TokenEncoder: + pass + + llm = self._bucketless_encoder_llm(["T5ForConditionalGeneration"]) + with patch( + "tensorrt_llm._torch.models.modeling_utils." + "get_registered_model_class", + return_value=_TokenEncoder): + with pytest.raises(ValueError, match="consumes packed tokens"): + llm._reject_token_encoder_config_without_buckets() + + def test_feature_encoder_without_buckets_is_accepted(self): + """A feature encoder may omit num_tokens/seq_lens; it derives both from the model.""" + + # Whisper derives both lists from the model: the same config is + # complete for it, and rejecting it would break every Whisper run. + class _FeatureEncoder: + + def encoder_graph_spec(self): + """Stand-in Whisper-shaped encoder contract.""" + return ((480000, ), torch.float32, 1500) + + llm = self._bucketless_encoder_llm(["WhisperForConditionalGeneration"]) + with patch( + "tensorrt_llm._torch.models.modeling_utils." + "get_registered_model_class", + return_value=_FeatureEncoder): + llm._reject_token_encoder_config_without_buckets() + + @pytest.mark.parametrize("architectures", [[], ["SomeUnregisteredArch"]]) + def test_an_unresolved_architecture_defers_to_the_model_engine( + self, architectures): + """An unregistered architecture defers the bucket check to the model engine.""" + # Out-of-tree models register in the worker, not here. Guessing would + # reject a feature encoder the engine goes on to accept. + llm = self._bucketless_encoder_llm(architectures) + with patch( + "tensorrt_llm._torch.models.modeling_utils." + "get_registered_model_class", + return_value=None): + llm._reject_token_encoder_config_without_buckets() + + @pytest.mark.parametrize("field,value", [ + ("model_kwargs", { + "architectures": ["WhisperForConditionalGeneration"] + }), + ("checkpoint_format", "MX"), + ("checkpoint_loader", object()), + ]) + def test_a_loader_that_picks_the_class_is_never_second_guessed( + self, field, value): + """When a checkpoint loader picks the class, the registry must not be consulted.""" + # These all reach `checkpoint_loader.load_config()`, where the engine + # gets its class, so the on-disk architecture may not be the one it + # builds. The registry must not even be consulted. + llm = self._bucketless_encoder_llm(["T5ForConditionalGeneration"]) + setattr(llm.args, field, value) + with patch("tensorrt_llm._torch.models.modeling_utils." + "get_registered_model_class") as resolve: + llm._reject_token_encoder_config_without_buckets() + resolve.assert_not_called() + + def test_a_decoder_only_model_is_left_to_the_model_engine(self): + """A decoder-only model skips the encoder bucket check entirely.""" + # No encoder at all: the engine's "consumes packed tokens" wording + # would only mislead, and skipping keeps the model import off this path. + llm = self._bucketless_encoder_llm(["LlamaForCausalLM"], + is_encoder_decoder=False) + with patch("tensorrt_llm._torch.models.modeling_utils." + "get_registered_model_class") as resolve: + llm._reject_token_encoder_config_without_buckets() + resolve.assert_not_called() def test_cuda_graph_config_infers_encode_mode_from_raw_dict(self): args = TorchLlmArgs(