@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cba06764d7

[Platform] Fill in the missing backend parameter for torch.compile (#51781)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

wangxiyuan committed Aug 19, 2026 at 18:31 UTC cba06764d7a9da41e6f535d6355c13f725574f07

5 files changed +15 -9

vllm/model_executor/models/diffusion_gemma.py

+3 -3

@@ -126,7 +126,7 @@ class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo):

126 return super().get_mm_max_tokens_per_item(seq_len, mm_counts)

127

128

129 -@torch.compile(dynamic=True)

129 +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

130 def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:

131 # fp32 before tanh for numerical stability (matches HF DiffusionGemma).

132 # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over

@@ -454,7 +454,7 @@ class DiffusionGemmaForConditionalGeneration(

454 raise ValueError(f"Unsupported modality: {modality}")

455

456

457 -@torch.compile(dynamic=True)

457 +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

458 def _compute_num_rejected(

459 num_logits: torch.Tensor,

460 num_sampled: torch.Tensor,

@@ -466,7 +466,7 @@ def _compute_num_rejected(

466 return torch.where(is_denoise, query_lens, num_rejected)

467

468

469 -@torch.compile(dynamic=True)

469 +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

470 def _compiled_sample_step(

471 # Logits from the model [num_decode * CL, vocab]

472 logits: torch.Tensor,

vllm/model_executor/models/kimi_k25_vit.py

+5 -1

@@ -61,7 +61,11 @@ def get_rope_shape_decorate(func):

63 @get_rope_shape_decorate

64 -@torch.compile(dynamic=True, disable=current_platform.simple_compile_backend == "tpu")

64 +@torch.compile(

65 + dynamic=True,

66 + backend=current_platform.simple_compile_backend,

67 + disable=current_platform.simple_compile_backend == "tpu",

68 +)

69 def get_rope_shape(org, interpolation_mode, shape):

70 return (

71 F.interpolate(

vllm/model_executor/models/parakeet.py

+4 -3

@@ -19,6 +19,7 @@ from vllm.logger import init_logger

19 from vllm.model_executor.layers.activation import ReLUSquaredActivation

20 from vllm.model_executor.layers.layernorm import RMSNorm

21 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

22 +from vllm.platforms import current_platform

23 from vllm.transformers_utils.configs.parakeet import ExtractorConfig, ParakeetConfig

25 logger = init_logger(__name__)

@@ -186,7 +187,7 @@ class ParakeetExtractor:

187 )

188 return self._apply_mel_filters(stft, mel_filters)

189

189 - @torch.compile(dynamic=True)

190 + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

191 def _apply_mel_filters(

192 self, stft_output: torch.Tensor, mel_filters: torch.Tensor

193 ) -> torch.Tensor:

@@ -195,7 +196,7 @@ class ParakeetExtractor:

196 mel_spec = torch.log(mel_spec + LOG_ZERO_GUARD_VALUE)

197 return mel_spec.permute(0, 2, 1)

198

198 - @torch.compile(dynamic=True)

199 + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

200 def _apply_preemphasis(

201 self, input_features: torch.Tensor, audio_lengths: torch.Tensor

202 ) -> torch.Tensor:

@@ -213,7 +214,7 @@ class ParakeetExtractor:

214 input_features = input_features.masked_fill(~timemask, 0.0)

215 return input_features

216

216 - @torch.compile(dynamic=True)

217 + @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

218 def _normalize_mel_features(

219 self, mel_features: torch.Tensor, audio_lengths: torch.Tensor

220 ) -> tuple[torch.Tensor, torch.Tensor]:

vllm/transformers_utils/processors/nano_nemotron_vl.py

+2 -1

@@ -26,6 +26,7 @@ from vllm.model_executor.models.parakeet import ParakeetExtractor

26 from vllm.multimodal.inputs import AudioItem

27 from vllm.multimodal.processing.processor import PromptUpdateDetails

28 from vllm.multimodal.video_prune.evs import compute_retained_tokens_count

29 +from vllm.platforms import current_platform

30 from vllm.tokenizers.hf import HfTokenizer

32 from .internvl import calculate_internvl_targets, get_internvl_target_ratios

@@ -56,7 +57,7 @@ def calculate_timestamps(

57 return timestamps

59 -@torch.compile(dynamic=True)

60 +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

61 def _bicubic_resize_and_normalize(

62 tensor: torch.Tensor,

63 size: tuple[int, int] | None = None,

vllm/v1/sample/ops/topk_topp_sampler.py

+1 -1

@@ -338,7 +338,7 @@ class TopKTopPSampler(nn.Module):

338

339 # Note: this is a workaround for

340 # https://github.com/pytorch/pytorch/pull/151218

341 -@torch.compile(dynamic=True)

341 +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

342 def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor:

343 probs = logits.softmax(dim=-1, dtype=torch.float32)

344 q = torch.empty_like(probs)