@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5e4e0e51f4

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 9, 18:15 UTC

[torch.compile] Compile `CustomOp.forward_native` for `SiluAndMul` and `QuantFP8` to avoid raw torch ops inside opaque custom ops (#32806)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Luka Govedič committed Jan 22, 2026 at 22:52 UTC 5e4e0e51f4fbb6d6fbf7a94feef07bd3617bfa5f

7 files changed +52 -13

tests/compile/test_silu_mul_quant_fusion.py

+2 -1

@@ -222,7 +222,7 @@ def test_fusion_silu_and_mul_quant(

222 x = torch.rand(num_tokens, hidden_size * 2)

223

224 # Reshape pass is needed for the fusion pass to work

225 - custom_ops = []

225 + custom_ops = ["none"]

226 if enable_silu_mul_custom_op:

227 custom_ops.append("+silu_and_mul")

228 if enable_quant_fp8_custom_op:

@@ -231,6 +231,7 @@ def test_fusion_silu_and_mul_quant(

231 compilation_config=CompilationConfig(

232 mode=CompilationMode.VLLM_COMPILE,

233 custom_ops=custom_ops,

234 + backend="eager", # avoid compilation for SiluAndMul and QuantFP8

235 pass_config=PassConfig(fuse_act_quant=True, eliminate_noops=True),

236 ),

237 )

tests/kernels/core/test_activation.py

+1 -1

@@ -57,7 +57,7 @@ def test_act_and_mul(

57 torch.set_default_device(device)

58 x = torch.randn(num_tokens, 2 * d, dtype=dtype)

59 if activation == "silu_and_mul":

60 - layer = SiluAndMul()

60 + layer = SiluAndMul(compile_native=False)

61 fn = torch.ops._C.silu_and_mul

62 if activation == "mul_and_silu":

63 layer = MulAndSilu()

vllm/compilation/matcher_utils.py

@@ -336,6 +336,7 @@ class MatcherQuantFP8(MatcherCustomOp):

336 quant_key.scale.group_shape,

337 column_major_scales=has_col_major_scales,

338 use_ue8m0=is_e8m0,

339 + compile_native=False,

340 )

341

342 def forward_rocm_aiter(

vllm/model_executor/custom_op.py

+42 -7

@@ -1,16 +1,15 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -

3 +import torch

4 import torch.nn as nn

6 from vllm.config import get_cached_compilation_config

7 from vllm.logger import init_logger

8 +from vllm.model_executor.utils import maybe_disable_graph_partition

9 from vllm.platforms import current_platform

11 logger = init_logger(__name__)

13 -

13 # Dictionary of all custom ops (classes, indexed by registered name).

14 # To check if an op with a name is enabled, call .enabled() on the class.

15 # Examples:

@@ -118,10 +117,10 @@ class CustomOp(nn.Module):

117 )

118 return super().__new__(op_cls_to_instantiate)

119

121 - def __init__(self, enforce_enable: bool = False):

120 + def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False):

121 super().__init__()

122 self._enforce_enable = enforce_enable

124 - self._forward_method = self.dispatch_forward()

123 + self._forward_method = self.dispatch_forward(compile_native=compile_native)

124

125 def forward(self, *args, **kwargs):

126 return self._forward_method(*args, **kwargs)

@@ -162,7 +161,7 @@ class CustomOp(nn.Module):

161 # PyTorch-native implementation.

162 return self.forward_native(*args, **kwargs)

163

165 - def dispatch_forward(self):

164 + def dispatch_forward(self, compile_native: bool):

165 # NOTE(woosuk): Here we assume that vLLM was built for only one

166 # specific backend. Currently, we do not support dynamic dispatching.

167 compilation_config = get_cached_compilation_config()

@@ -180,7 +179,9 @@ class CustomOp(nn.Module):

179 compilation_config.disabled_custom_ops.update([self.__class__.name])

180

181 if not enabled:

183 - return self.forward_native

182 + # Compile forward_native to avoid eager torch ops if inside

183 + # opaque torch custom op (e.g. fused_moe, unified_attention, etc.)

184 + return self.maybe_compile(self.forward_native, enable=compile_native)

185

186 if current_platform.is_rocm():

187 return self.forward_hip

@@ -195,6 +196,40 @@ class CustomOp(nn.Module):

196 else:

197 return self.forward_cuda

198

199 + def maybe_compile(self, fn, *, enable: bool = True):

200 + """

201 + Compile fn if compilation enabled.

202 + Useful for CustomOp instances called from within a torch custom op,

203 + meaning the forward call is hidden from the model-level torch.compile.

204 +

205 + NOTE: this does not enable fusion across ops, so opaque custom ops

206 + should still be unwrapped wherever possible.

207 + """

208 + # Do not compile if compilation disabled

209 + from vllm.config.compilation import CompilationMode

210 +

211 + if not enable:

212 + return fn

213 +

214 + # Do not compile if global compilation disabled

215 + compilation_config = get_cached_compilation_config()

216 + if compilation_config.mode == CompilationMode.NONE:

217 + return fn

218 +

219 + # If eager backend is used, do not compile either

220 + if compilation_config.backend == "eager":

221 + return fn

222 +

223 + # dynamic=True to avoid recompilations

224 + return torch.compile(

225 + fn,

226 + dynamic=True,

227 + backend=current_platform.simple_compile_backend,

228 + options=maybe_disable_graph_partition(

229 + current_platform.simple_compile_backend

230 + ),

231 + )

232 +

233 @classmethod

234 def enabled(cls) -> bool:

235 # if no name, then it was not registered

vllm/model_executor/layers/activation.py

+2 -2

@@ -75,8 +75,8 @@ class SiluAndMul(CustomOp):

76 # --8<-- [end:silu_and_mul]

78 - def __init__(self):

79 - super().__init__()

78 + def __init__(self, *, compile_native: bool = True):

79 + super().__init__(compile_native=compile_native)

80 if current_platform.is_cuda_alike():

81 self.op = torch.ops._C.silu_and_mul

82 elif current_platform.is_xpu():

vllm/model_executor/layers/quantization/input_quant_fp8.py

+3 -1

@@ -38,6 +38,7 @@ class QuantFP8(CustomOp):

38 column_major_scales: bool = False,

39 tma_aligned_scales: bool = False,

40 use_ue8m0: bool | None = None, # for Torch compile

41 + compile_native: bool = True,

42 ):

43 """

44 :param static: static or dynamic quantization

@@ -49,8 +50,9 @@ class QuantFP8(CustomOp):

50 TMA-aligned layout

51 :param column_major_scales: For group quantization, output scales in

52 column major format

53 + :param compile_native: Manually compile forward_native if compile mode > None

54 """

53 - super().__init__()

55 + super().__init__(compile_native=compile_native)

56 self.static = static

57 self.group_shape = group_shape

58 self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN

vllm/model_executor/layers/rotary_embedding/common.py

+1 -1

@@ -129,7 +129,7 @@ class ApplyRotaryEmb(CustomOp):

129 is_neox_style: bool = True,

130 enable_fp32_compute: bool = False,

131 ) -> None:

132 - super().__init__(enforce_enable)

132 + super().__init__(enforce_enable=enforce_enable)

133 self.is_neox_style = is_neox_style

134 self.enable_fp32_compute = enable_fp32_compute

135