@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2196bac135

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 4, 17:30 UTC

[Compilation] Refactor SiluMul activation+quant Fusion Pass (#39684)

Signed-off-by: BadrBasowid <badr.basowid@gmail.com>

BadrBasowid committed Apr 23, 2026 at 21:10 UTC 2196bac1359a62a29f76936f37380235cc47f096

4 files changed +81 -103

tests/compile/fusions_e2e/conftest.py

+7 -1

@@ -189,7 +189,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):

189 # TODO: Remove log counting in unit tests

190 # once all matchers implement VllmFusionPatternMatcherPass

191 n_expected = tp_size * num_ranges_activated

192 - if match_name != "attn_quant_fusion":

192 + if match_name not in ("attn_quant_fusion", "act_quant_fusion"):

193 assert len(log_matches) == n_expected, (

194 f"Could not find {n_expected} {match_name} "

195 f"(found {len(log_matches)}) in:\n {log_holder.text}"

@@ -250,6 +250,12 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):

250 f"entries (SP took precedence), found: {log_matches}"

251 )

252

253 + elif match_name == "act_quant_fusion":

254 + actual_match = match_table.get("activation_quant_fusion_pass", 0)

255 + assert actual_match == expected_matches * n_expected, (

256 + f"Could not find {expected_matches * n_expected} "

257 + f"{match_name} (found {actual_match})."

258 + )

259 elif match_name == "attn_quant_fusion":

260 actual_match = match_table.get(

261 "attn_quant_fusion", 0

tests/compile/passes/test_silu_mul_quant_fusion.py

+1 -1

@@ -168,7 +168,7 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):

168

169 def forward(self, x):

170 y = self.silu_and_mul(x)

171 - x2 = self.w8a8_block_fp8_linear(y, self.w, self.wscale)

171 + x2 = self.w8a8_block_fp8_linear(y)

172 return x2

173

174 def ops_in_model_before(self):

vllm/compilation/passes/fusion/act_quant_fusion.py

+59 -72

@@ -1,16 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 -from abc import ABC, abstractmethod

4 +import itertools

5 from typing import Any

7 import torch

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

9 -from torch._inductor.pattern_matcher import (

10 - PatternMatcherPass,

11 - fwd_only,

12 - register_replacement,

13 -)

9 from torch._ops import OpOverload

11 from vllm.config import VllmConfig

@@ -24,8 +19,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

19 )

20 from vllm.platforms import current_platform

27 -from ..inductor_pass import enable_fake_mode

28 -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

22 +from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement

23 from .matcher_utils import MatcherQuantFP8, MatcherSiluAndMul

24 from .rms_quant_fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32

@@ -50,9 +44,9 @@ if current_platform.is_cuda_alike():

44 FUSED_OPS[kFp8Dynamic64Sym] = torch.ops._C.silu_and_mul_per_block_quant.default

53 -class ActivationQuantPattern(ABC):

47 +class ActivationQuantPattern(VllmPatternReplacement):

48 """

55 - The base class for Activation+Quant fusions.

49 + Base class for Activation+Quant fusions.

50 Should not be used directly.

51 """

@@ -79,10 +73,6 @@ class ActivationQuantPattern(ABC):

73 kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}

74 return torch.empty(*args, **kwargs)

82 - @abstractmethod

83 - def register(self, pm_pass: PatternMatcherPass) -> None:

84 - raise NotImplementedError

85 -

77 class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):

78 """

@@ -100,8 +90,9 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):

90 scale,

91 ]

103 - def register(self, pm_pass: PatternMatcherPass) -> None:

104 - def pattern(

93 + @property

94 + def pattern(self):

95 + def _pattern(

96 input: torch.Tensor,

97 scale: torch.Tensor,

98 ) -> torch.Tensor:

@@ -109,7 +100,11 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):

100 result_quant = self.quant_matcher(result_silu_mul, scale)

101 return result_quant[0]

102

112 - def replacement(

103 + return _pattern

104 +

105 + @property

106 + def replacement(self):

107 + def _replacement(

108 input: torch.Tensor,

109 scale: torch.Tensor,

110 ) -> torch.Tensor:

@@ -123,10 +118,7 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):

118 )

119 return at[1]

120

126 - inps = self.get_inputs()

127 - pattern(*inps)

128 -

129 - register_replacement(pattern, replacement, inps, fwd_only, pm_pass)

121 + return _replacement

122

123

124 class SiluMulNvfp4QuantPattern(ActivationQuantPattern):

@@ -144,8 +136,9 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):

136 scale = empty_fp32(1, 1)

137 return [result, output_scale, input_, scale]

138

147 - def register(self, pm_pass: PatternMatcherPass) -> None:

148 - def pattern(

139 + @property

140 + def pattern(self):

141 + def _pattern(

142 result: torch.Tensor,

143 output_scale: torch.Tensor,

144 input: torch.Tensor,

@@ -162,7 +155,11 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):

155 )

156 return at[1], at[2]

157

165 - def replacement(

158 + return _pattern

159 +

160 + @property

161 + def replacement(self):

162 + def _replacement(

163 result: torch.Tensor,

164 output_scale: torch.Tensor,

165 input: torch.Tensor,

@@ -177,7 +174,7 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):

174 )

175 return at[1], at[2]

176

180 - register_replacement(pattern, replacement, self.get_inputs(), fwd_only, pm_pass)

177 + return _replacement

178

179

180 class SiluMulBlockQuantPattern(ActivationQuantPattern):

@@ -210,10 +207,9 @@ class SiluMulBlockQuantPattern(ActivationQuantPattern):

207 scale = self.quant_matcher.empty_f32(1, 1)

208 return self.silu_and_mul_matcher.inputs() + [scale]

209

213 - def register(self, pm_pass: PatternMatcherPass) -> None:

214 - is_scale_transposed = self.is_scale_transposed

215 -

216 - def pattern(

210 + @property

211 + def pattern(self):

212 + def _pattern(

213 input: torch.Tensor,

214 scale: torch.Tensor,

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

@@ -235,12 +231,16 @@ class SiluMulBlockQuantPattern(ActivationQuantPattern):

231 fp8_min=finfo.min,

232 fp8_max=finfo.max,

233 scale_ue8m0=self.is_e8m0,

238 - dummy_is_scale_transposed=is_scale_transposed,

234 + dummy_is_scale_transposed=self.is_scale_transposed,

235 dummy_is_tma_aligned=self.is_tma_aligned,

236 )

237 return result, scale

238

243 - def replacement(

239 + return _pattern

240 +

241 + @property

242 + def replacement(self):

243 + def _replacement(

244 input: torch.Tensor,

245 scale: torch.Tensor,

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

@@ -249,7 +249,7 @@ class SiluMulBlockQuantPattern(ActivationQuantPattern):

249 result = torch.empty(

250 output_shape, device=input.device, dtype=self.quant_dtype

251 )

252 - if is_scale_transposed:

252 + if self.is_scale_transposed:

253 scale = torch.empty(

254 (d // self.group_size, input.shape[0]),

255 device=input.device,

@@ -268,15 +268,14 @@ class SiluMulBlockQuantPattern(ActivationQuantPattern):

268 scales=scale,

269 group_size=self.group_size,

270 scale_ub=None,

271 - is_scale_transposed=is_scale_transposed,

271 + is_scale_transposed=self.is_scale_transposed,

272 )

273 return at[1], at[2]

274

275 - inps = self.get_inputs()

276 - register_replacement(pattern, replacement, inps, fwd_only, pm_pass)

275 + return _replacement

276

277

279 -class ActivationQuantFusionPass(VllmPatternMatcherPass):

278 +class ActivationQuantFusionPass(VllmFusionPatternMatcherPass):

279 """

280 This pass fuses a pre-defined set of custom ops into fused ops.

281 It uses the torch pattern matcher to find the patterns and replace them.

@@ -286,45 +285,33 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass):

285 https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980

286 """

287

289 - @enable_fake_mode

288 def __init__(self, config: VllmConfig) -> None:

291 - super().__init__(config)

292 -

293 - self.patterns: PatternMatcherPass = PatternMatcherPass(

294 - pass_name="activation_quant_fusion_pass"

295 - )

289 + super().__init__(config, "activation_quant_fusion_pass")

290

297 - pattern_silu_mul_fp8 = SiluMulFp8StaticQuantPattern()

298 - pattern_silu_mul_fp8.register(self.patterns)

291 + self.register(SiluMulFp8StaticQuantPattern())

292

293 if silu_and_mul_nvfp4_quant_supported:

301 - pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern()

302 - pattern_silu_mul_nvfp4.register(self.patterns)

294 + self.register(SiluMulNvfp4QuantPattern())

295

296 if current_platform.is_cuda():

305 - for quant_key in [kFp8Dynamic128Sym, kFp8Dynamic64Sym]:

306 - for is_scale_transposed in [False, True]:

307 - for is_e8m0 in [True, False]:

308 - for is_tma_aligned in [False, True]:

309 - SiluMulBlockQuantPattern(

310 - quant_key,

311 - is_scale_transposed=is_scale_transposed,

312 - is_e8m0=is_e8m0,

313 - is_tma_aligned=is_tma_aligned,

314 - ).register(self.patterns)

315 -

316 - self.dump_patterns(config, self.patterns)

317 -

318 - @VllmInductorPass.time_and_log

319 - def __call__(self, graph: torch.fx.Graph) -> None:

320 - self.matched_count = self.patterns.apply(graph)

321 - logger.debug("Replaced %s patterns", self.matched_count)

322 -

323 - def uuid(self) -> str:

324 - return VllmInductorPass.hash_source(

325 - self,

326 - ActivationQuantPattern,

327 - SiluMulFp8StaticQuantPattern,

328 - SiluMulNvfp4QuantPattern,

329 - SiluMulBlockQuantPattern,

330 - )

297 + for (

298 + quant_key,

299 + is_scale_transposed,

300 + is_e8m0,

301 + is_tma_aligned,

302 + ) in itertools.product(

303 + [kFp8Dynamic128Sym, kFp8Dynamic64Sym],

304 + [False, True],

305 + [True, False],

306 + [False, True],

307 + ):

308 + self.register(

309 + SiluMulBlockQuantPattern(

310 + quant_key,

311 + is_scale_transposed=is_scale_transposed,

312 + is_e8m0=is_e8m0,

313 + is_tma_aligned=is_tma_aligned,

314 + )

315 + )

316 +

317 + self.dump_patterns(config, self.pm_pass)

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

+14 -29

@@ -28,7 +28,6 @@ from ..vllm_inductor_pass import (

28 VllmPatternMatcherPass,

29 VllmPatternReplacement,

30 )

31 -from .act_quant_fusion import ActivationQuantPattern

31 from .matcher_utils import (

32 MatcherFusedAddRMSNorm,

33 MatcherQuantFP8,

@@ -345,7 +344,7 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):

344 return self.hash_source(self, *fusion_patterns)

345

346

348 -class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):

347 +class AiterSiluMulFp8GroupQuantPattern(VllmPatternReplacement):

348 """

349 This pattern fuses aiter silu_and_mul & group fp8 quant custom

350 ops into an aiter silu_and_mul_group_fp8_quant op.

@@ -364,26 +363,29 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):

363 self.silu_and_mul_matcher.inputs()[0],

364 ]

365

367 - def register(self, pm_pass: PatternMatcherPass) -> None:

368 - def pattern(

366 + @property

367 + def pattern(self):

368 + def _pattern(

369 input: torch.Tensor,

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

371 at1 = self.silu_and_mul_matcher(input)

372 at2 = self.quant_matcher(at1)

373 return at2[0], at2[1]

374

375 - def replacement(

375 + return _pattern

376 +

377 + @property

378 + def replacement(self):

379 + def _replacement(

380 input: torch.Tensor,

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

382 at = self.FUSED_SILU_MUL_QUANT_OP(x=input, group_size=128)

383 return at[0], at[1]

384

381 - pm.register_replacement(

382 - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass

383 - )

385 + return _replacement

386

387

386 -class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):

388 +class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmFusionPatternMatcherPass):

389 """

390 This pass fuses a pre-defined set of custom ops into fused ops.

391 It uses the torch pattern matcher to find the patterns and replace them.

@@ -393,29 +395,12 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):

395 https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980

396 """

397

396 - @enable_fake_mode

398 def __init__(self, config: VllmConfig) -> None:

398 - super().__init__(config)

399 -

400 - self.patterns: PatternMatcherPass = PatternMatcherPass(

401 - pass_name="rocm_aiter_silu_mul_fp8_group_quant_fusion_pass"

402 - )

403 -

404 - AiterSiluMulFp8GroupQuantPattern().register(self.patterns)

405 -

406 - self.dump_patterns(config, self.patterns)

399 + super().__init__(config, "rocm_aiter_silu_mul_fp8_group_quant_fusion_pass")

400

408 - @VllmInductorPass.time_and_log

409 - def __call__(self, graph: torch.fx.Graph) -> None:

410 - self.matched_count = self.patterns.apply(graph)

411 - logger.debug("Replaced %s patterns", self.matched_count)

401 + self.register(AiterSiluMulFp8GroupQuantPattern())

402

413 - def uuid(self) -> str:

414 - fusion_patterns = [

415 - ActivationQuantPattern,

416 - AiterSiluMulFp8GroupQuantPattern,

417 - ]

418 - return VllmInductorPass.hash_source(self, *fusion_patterns)

403 + self.dump_patterns(config, self.pm_pass)

404

405

406 class AddAiterRMSNormPadPattern: