@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 74c583bc50

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 00:15 UTC

[Core] Whisper support `torch.compile` (#30385)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Jan 19, 2026 at 11:02 UTC 74c583bc508c2dafb9e95bab3b635884e4a021f3

5 files changed +27 -1

tests/entrypoints/openai/correctness/test_transcription_api_correctness.py

+3 -1

@@ -156,7 +156,9 @@ def test_wer_correctness(

156 model_name, dataset_repo, expected_wer, n_examples=-1, max_concurrent_request=None

157 ):

158 # TODO refactor to use `ASRDataset`

159 - with RemoteOpenAIServer(model_name, ["--enforce-eager"]) as remote_server:

159 + with RemoteOpenAIServer(

160 + model_name, ["--enforce-eager"], max_wait_seconds=480

161 + ) as remote_server:

162 dataset = load_hf_dataset(dataset_repo)

163

164 if not max_concurrent_request:

vllm/compilation/decorators.py

@@ -25,6 +25,7 @@ from vllm.config import (

25 set_current_vllm_config,

26 )

27 from vllm.config.compilation import DynamicShapesType

28 +from vllm.forward_context import get_forward_context, is_forward_context_available

29 from vllm.logger import init_logger

30 from vllm.sequence import IntermediateTensors

31 from vllm.utils.import_utils import resolve_obj_by_qualname

@@ -388,6 +389,12 @@ def _support_torch_compile(

389 if self.do_not_compile or torch.compiler.is_compiling():

390 return self.forward(*args, **kwargs)

391

392 + # If skip_compiled is set, bypass compiled model call. This is used e.g. for

393 + # enc-dec models where tensor shapes/types vary across invocations, preventing

394 + # the capture of a single computational graph.

395 + if is_forward_context_available() and get_forward_context().skip_compiled:

396 + return self.forward(*args, **kwargs)

397 +

398 # if aot_compiled_fn is set, call it with partition wrapper context.

399 # The partition wrapper must be active at runtime for CUDA graph

400 # capture to work correctly with inductor graph partitioning.

vllm/forward_context.py

@@ -207,6 +207,9 @@ class ForwardContext:

207

208 ubatch_slices: UBatchSlices | None = None

209

210 + # If True, bypass the compiled model call, e.g. by using .forward() directly

211 + skip_compiled: bool = False

212 +

213 additional_kwargs: dict[str, Any] = field(default_factory=dict)

214

215 def __post_init__(self):

@@ -240,6 +243,7 @@ def create_forward_context(

243 batch_descriptor: BatchDescriptor | None = None,

244 ubatch_slices: UBatchSlices | None = None,

245 additional_kwargs: dict[str, Any] | None = None,

246 + skip_compiled: bool = False,

247 ):

248 return ForwardContext(

249 no_compile_layers=vllm_config.compilation_config.static_forward_context,

@@ -249,6 +253,7 @@ def create_forward_context(

253 cudagraph_runtime_mode=cudagraph_runtime_mode,

254 batch_descriptor=batch_descriptor,

255 ubatch_slices=ubatch_slices,

256 + skip_compiled=skip_compiled,

257 additional_kwargs=additional_kwargs or {},

258 )

259

@@ -278,6 +283,7 @@ def set_forward_context(

283 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

284 batch_descriptor: BatchDescriptor | None = None,

285 ubatch_slices: UBatchSlices | None = None,

286 + skip_compiled: bool = False,

287 ):

288 """A context manager that stores the current forward context,

289 can be attention metadata, etc.

@@ -336,6 +342,7 @@ def set_forward_context(

342 batch_descriptor,

343 ubatch_slices,

344 additional_kwargs,

345 + skip_compiled,

346 )

347

348 try:

vllm/model_executor/models/whisper.py

@@ -19,6 +19,7 @@ from transformers import (

19 from transformers.models.whisper.modeling_whisper import sinusoids

21 from vllm.attention.layer import Attention

22 +from vllm.compilation.decorators import support_torch_compile

23 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig

24 from vllm.config.multimodal import BaseDummyOptions

25 from vllm.distributed import get_tensor_model_parallel_world_size

@@ -561,6 +562,7 @@ class WhisperEncoder(nn.Module):

562 return self.forward_layers(hidden_states)

563

564

565 +@support_torch_compile(dynamic_arg_dims={"input_ids": 0, "positions": -1})

566 class WhisperDecoder(nn.Module):

567 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

568 super().__init__()

vllm/v1/worker/gpu_model_runner.py

@@ -3268,6 +3268,13 @@ class GPUModelRunner(

3268 # Mark KV scales as calculated after the first forward pass

3269 self.calculate_kv_scales = False

3270

3271 + # Encoder-decoder models can only compile the pure decode steps where no

3272 + # encoder inputs are present. Use eager for the first pass.

3273 + num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs)

3274 + has_encoder_input = (

3275 + self.model_config.is_encoder_decoder and num_encoder_reqs > 0

3276 + )

3277 +

3278 # Run the model.

3279 # Use persistent buffers for CUDA graphs.

3280 with (

@@ -3279,6 +3286,7 @@ class GPUModelRunner(

3286 cudagraph_runtime_mode=cudagraph_mode,

3287 batch_descriptor=batch_desc,

3288 ubatch_slices=ubatch_slices_padded,

3289 + skip_compiled=has_encoder_input,

3290 ),

3291 record_function_or_nullcontext("gpu_model_runner: forward"),

3292 self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,