@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 994acec0cc

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 11, 12:30 UTC

[Bugfix] Fix fusion for VL models (#30244)

Signed-off-by: ElizaWszola <ewszola@redhat.com>

ElizaWszola committed Dec 14, 2025 at 14:22 UTC 994acec0cc9d6348268b5f371c66239fe75f928d

4 files changed +143 -72

tests/compile/distributed/test_fusions_e2e.py

+78

@@ -27,6 +27,7 @@ is_blackwell = lambda: current_platform.is_device_capability_family(100)

27 class Matches(NamedTuple):

28 attention_fusion: int = 0

29 allreduce_fusion: int = 0

30 + rms_quant_norm_fusion: int = 0

31 sequence_parallel: int = 0

32 async_tp: int = 0

@@ -40,6 +41,7 @@ class ModelBackendTestCase(NamedTuple):

42 MODELS_FP8: list[ModelBackendTestCase] = []

43 MODELS_FP4: list[ModelBackendTestCase] = []

44 +MODELS_GROUP_FP8: list[ModelBackendTestCase] = []

45 MODELS: list[ModelBackendTestCase] = [] # tp-only

47 if current_platform.is_cuda():

@@ -498,3 +500,79 @@ def run_model(compile_config: int | CompilationConfig, model: str, **model_kwarg

500 compilation_config.compile_ranges_split_points = (

501 llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points

502 )

503 +

504 +

505 +if current_platform.is_cuda():

506 + MODELS_GROUP_FP8 = [

507 + ModelBackendTestCase(

508 + model_name="Qwen/Qwen3-30B-A3B-FP8",

509 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),

510 + backend=AttentionBackendEnum.TRITON_ATTN,

511 + matches=Matches(

512 + rms_quant_norm_fusion=48,

513 + ),

514 + ),

515 + ]

516 +

517 +CUSTOM_OPS_QUANT_RMS_NORM = ["+quant_fp8,+rms_norm"]

518 +

519 +

520 +@pytest.mark.parametrize(

521 + "model_name, model_kwargs, backend, matches, custom_ops",

522 + # Test rms norm+group quant_fp8 fusion

523 + list[tuple[Any, ...]](flat_product(MODELS_GROUP_FP8, CUSTOM_OPS_QUANT_RMS_NORM)),

524 +)

525 +@pytest.mark.parametrize("inductor_graph_partition", [True, False])

526 +def test_rms_group_quant(

527 + model_name: str,

528 + model_kwargs: dict[str, Any],

529 + backend: AttentionBackendEnum,

530 + matches: Matches,

531 + custom_ops: str,

532 + inductor_graph_partition: bool,

533 + caplog_mp_spawn,

534 + monkeypatch,

535 +):

536 + if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):

537 + pytest.skip("Inductor graph partition requires torch>=2.9")

538 +

539 + custom_ops_list = custom_ops.split(",") if custom_ops else []

540 +

541 + if inductor_graph_partition:

542 + mode = CUDAGraphMode.FULL_AND_PIECEWISE

543 + splitting_ops: list[str] | None = None

544 + else:

545 + mode = CUDAGraphMode.FULL_DECODE_ONLY

546 + splitting_ops = []

547 +

548 + # Disable, compile cache to make sure custom passes run.

549 + # Otherwise, we can't verify fusion happened through the logs.

550 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")

551 +

552 + # To capture subprocess logs, we need to know whether spawn or fork is used.

553 + # Force spawn as it is more general.

554 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

555 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)

556 +

557 + compilation_config = CompilationConfig(

558 + # Testing properties

559 + custom_ops=custom_ops_list,

560 + use_inductor_graph_partition=inductor_graph_partition,

561 + cudagraph_mode=mode,

562 + splitting_ops=splitting_ops,

563 + # Common

564 + mode=CompilationMode.VLLM_COMPILE,

565 + pass_config=PassConfig(eliminate_noops=True, enable_fusion=True),

566 + # Inductor caches custom passes by default as well via uuid

567 + inductor_compile_config={"force_disable_caches": True},

568 + )

569 +

570 + with caplog_mp_spawn(logging.DEBUG) as log_holder:

571 + run_model(compilation_config, model_name, **model_kwargs)

572 +

573 + log_matches = re.findall(

574 + r"\[fusion.py:\d+] Replaced (\d+) patterns",

575 + log_holder.text,

576 + )

577 + assert len(log_matches) == 1, log_holder.text

578 + assert int(log_matches[0]) == matches.rms_quant_norm_fusion

vllm/compilation/fusion.py

+52 -48

@@ -23,17 +23,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

23 kNvfp4Quant,

24 kStaticTensorScale,

25 )

26 -from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

27 - cutlass_block_fp8_supported,

28 -)

26 from vllm.platforms import current_platform

30 -from vllm.utils.deep_gemm import (

31 - is_deep_gemm_e8m0_used,

32 - should_use_deepgemm_for_fp8_linear_for_nk,

33 -)

28 from .inductor_pass import enable_fake_mode

36 -from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm

29 +from .matcher_utils import (

30 + MatcherFusedAddRMSNorm,

31 + MatcherQuantFP8,

32 + MatcherRMSNorm,

33 +)

34 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

36 logger = init_logger(__name__)

@@ -118,21 +115,18 @@ FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {

115

116

117 class RMSNormQuantPattern:

121 - def __init__(self, epsilon: float, key: FusedRMSQuantKey):

118 + def __init__(

119 + self,

120 + epsilon: float,

121 + key: FusedRMSQuantKey,

122 + has_col_major_scales: bool = False,

123 + is_e8m0: bool = False,

124 + ):

125 self.epsilon = epsilon

126 self.quant_dtype = key.quant.dtype

127 config = get_current_vllm_config()

128 self.model_dtype = config.model_config.dtype if config.model_config else None

129

127 - # groupwise FP8 linear uses col major scales if deepgemm and cutlass

128 - using_deepgemm = should_use_deepgemm_for_fp8_linear_for_nk(

129 - self.model_dtype,

130 - config.model_config.hf_config.intermediate_size,

131 - config.model_config.hf_config.hidden_size,

132 - )

133 - use_col_major_scales = using_deepgemm or cutlass_block_fp8_supported()

134 - use_e8m0 = is_deep_gemm_e8m0_used() if using_deepgemm else False

135 -

130 assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"

131 self.FUSED_OP = FUSED_OPS[key]

132

@@ -142,7 +136,7 @@ class RMSNormQuantPattern:

136 else MatcherFusedAddRMSNorm(epsilon)

137 )

138 self.quant_matcher = MatcherQuantFP8(

145 - key.quant, use_col_major_scales=use_col_major_scales, use_e8m0=use_e8m0

139 + key.quant, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0

140 )

141

142

@@ -260,6 +254,8 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

254 quant_dtype: torch.dtype,

255 group_shape: GroupShape,

256 symmetric=True,

257 + has_col_major_scales: bool = False,

258 + is_e8m0: bool = False,

259 ):

260 scale = ScaleDesc(torch.float32, False, group_shape)

261 key = FusedRMSQuantKey(

@@ -267,7 +263,11 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

263 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),

264 )

265 self.group_shape = group_shape

270 - super().__init__(epsilon, key)

266 + self.has_col_major_scales = has_col_major_scales

267 + self.is_e8m0 = is_e8m0

268 + super().__init__(

269 + epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0

270 + )

271

272 def register(self, pm_pass: PatternMatcherPass):

273 def pattern(input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor):

@@ -283,9 +283,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

283 input = input.to(dtype=self.model_dtype)

284

285 result = torch.empty_like(input, dtype=self.quant_dtype)

286 - scale = self.quant_matcher.make_scale(

287 - input, transposed=self.quant_matcher.use_col_major_scales

288 - )

286 + scale = self.quant_matcher.make_scale(input, self.has_col_major_scales)

287 at = auto_functionalized(

288 self.FUSED_OP,

289 result=result,

@@ -296,7 +294,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

294 scale_ub=None,

295 residual=residual,

296 group_size=self.group_shape[1],

299 - is_scale_transposed=self.quant_matcher.use_col_major_scales,

297 + is_scale_transposed=self.has_col_major_scales,

298 )

299

300 # result, residual, scale

@@ -318,6 +316,8 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

316 quant_dtype: torch.dtype,

317 group_shape: GroupShape,

318 symmetric=True,

319 + has_col_major_scales: bool = False,

320 + is_e8m0: bool = False,

321 ):

322 scale = ScaleDesc(torch.float32, False, group_shape)

323 key = FusedRMSQuantKey(

@@ -325,7 +325,9 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

325 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),

326 )

327 self.group_shape = group_shape

328 - super().__init__(epsilon, key)

328 + super().__init__(

329 + epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0

330 + )

331

332 def register(self, pm_pass: PatternMatcherPass):

333 def pattern(input: torch.Tensor, weight: torch.Tensor):

@@ -340,7 +342,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

342

343 result = torch.empty_like(input, dtype=self.quant_dtype)

344 scale = self.quant_matcher.make_scale(

343 - input, transposed=self.quant_matcher.use_col_major_scales

345 + input, transposed=self.quant_matcher.has_col_major_scales

346 )

347 at = auto_functionalized(

348 self.FUSED_OP,

@@ -352,7 +354,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

354 scale_ub=None,

355 residual=None,

356 group_size=self.group_shape[1],

355 - is_scale_transposed=self.quant_matcher.use_col_major_scales,

357 + is_scale_transposed=self.quant_matcher.has_col_major_scales,

358 )

359

360 # result, scale

@@ -489,27 +491,6 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):

491 # Make sure fused add patterns are before simple rms norm,

492 # as the latter is a subset of the former in torch ops

493 for epsilon in [1e-5, 1e-6]:

492 - # Fuse fused_add_rms_norm + fp8 group quant

493 - # Only register group quant patterns on CUDA where the C++ op exists

494 - if current_platform.is_cuda():

495 - FusedAddRMSNormGroupQuantPattern(

496 - epsilon, FP8_DTYPE, group_shape=GroupShape(1, 128)

497 - ).register(self.patterns)

498 -

499 - # Fuse rms_norm + fp8 group quant

500 - RMSNormGroupQuantPattern(

501 - epsilon, FP8_DTYPE, group_shape=GroupShape(1, 128)

502 - ).register(self.patterns)

503 -

504 - FusedAddRMSNormGroupQuantPattern(

505 - epsilon, FP8_DTYPE, group_shape=GroupShape(1, 64)

506 - ).register(self.patterns)

507 -

508 - # Fuse rms_norm + fp8 group quant

509 - RMSNormGroupQuantPattern(

510 - epsilon, FP8_DTYPE, group_shape=GroupShape(1, 64)

511 - ).register(self.patterns)

512 -

494 # Fuse fused_add_rms_norm + static fp8 quant

495 FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(

496 self.patterns

@@ -526,6 +507,29 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):

507 # Fuse rms_norm + dynamic per-token fp8 quant

508 RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)

509

510 + # Only register group quant patterns on CUDA where the C++ op exists

511 + if current_platform.is_cuda():

512 + for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:

513 + for has_col_major_scales in [True, False]:

514 + for is_e8m0 in [True, False]:

515 + # Fuse fused_add_rms_norm + fp8 group quant

516 + FusedAddRMSNormGroupQuantPattern(

517 + epsilon,

518 + FP8_DTYPE,

519 + group_shape=group_shape,

520 + has_col_major_scales=has_col_major_scales,

521 + is_e8m0=is_e8m0,

522 + ).register(self.patterns)

523 +

524 + # Fuse rms_norm + fp8 group quant

525 + RMSNormGroupQuantPattern(

526 + epsilon,

527 + FP8_DTYPE,

528 + group_shape=group_shape,

529 + has_col_major_scales=has_col_major_scales,

530 + is_e8m0=is_e8m0,

531 + ).register(self.patterns)

532 +

533 self.dump_patterns(config, self.patterns)

534

535 @VllmInductorPass.time_and_log

vllm/compilation/matcher_utils.py

+13 -7

@@ -234,24 +234,30 @@ class MatcherQuantFP8(MatcherCustomOp):

234 self,

235 quant_key: QuantKey,

236 enabled: bool | None = None,

237 - use_col_major_scales: bool = False,

238 - use_e8m0: bool = False,

237 + has_col_major_scales: bool = False,

238 + is_e8m0: bool = False,

239 ):

240 if enabled is None:

241 enabled = QuantFP8.enabled()

242

243 super().__init__(enabled)

244 self.quant_key = quant_key

245 - self.use_col_major_scales = use_col_major_scales

246 - self.use_e8m0 = use_e8m0

245 assert quant_key in QUANT_OPS, f"unsupported quantization scheme {quant_key}"

246 self.QUANT_OP = QUANT_OPS[quant_key]

247

248 + self.has_col_major_scales = has_col_major_scales

249 + self.is_e8m0 = is_e8m0

250 +

251 assert quant_key.dtype == current_platform.fp8_dtype(), (

252 "Only QuantFP8 supported by"

253 )

254 assert quant_key.scale2 is None

254 - self.quant_fp8 = QuantFP8(quant_key.scale.static, quant_key.scale.group_shape)

255 + self.quant_fp8 = QuantFP8(

256 + quant_key.scale.static,

257 + quant_key.scale.group_shape,

258 + column_major_scales=has_col_major_scales,

259 + use_ue8m0=is_e8m0,

260 + )

261

262 def forward_custom(

263 self,

@@ -264,7 +270,7 @@ class MatcherQuantFP8(MatcherCustomOp):

270

271 if self.quant_key.scale.group_shape.is_per_group():

272 assert scale is None

267 - scale = self.make_scale(input, transposed=self.use_col_major_scales)

273 + scale = self.make_scale(input, transposed=self.has_col_major_scales)

274

275 finfo = torch.finfo(self.quant_key.dtype)

276 fp8_min = finfo.min

@@ -279,7 +285,7 @@ class MatcherQuantFP8(MatcherCustomOp):

285 eps=1e-10,

286 fp8_min=fp8_min,

287 fp8_max=fp8_max,

282 - scale_ue8m0=self.use_e8m0,

288 + scale_ue8m0=self.is_e8m0,

289 )

290 return result, scale

291

vllm/utils/deep_gemm.py

-17

@@ -381,22 +381,6 @@ def should_use_deepgemm_for_fp8_linear(

381 )

382

383

384 -def should_use_deepgemm_for_fp8_linear_for_nk(

385 - output_dtype: torch.dtype,

386 - shape0: int,

387 - shape1: int,

388 - supports_deep_gemm: bool | None = None,

389 -):

390 - if supports_deep_gemm is None:

391 - supports_deep_gemm = is_deep_gemm_supported()

392 - return (

393 - supports_deep_gemm

394 - and output_dtype == torch.bfloat16

395 - and shape0 % 128 == 0

396 - and shape1 % 128 == 0

397 - )

398 -

399 -

384 __all__ = [

385 "calc_diff",

386 "DeepGemmQuantScaleFMT",

@@ -411,7 +395,6 @@ __all__ = [

395 "is_deep_gemm_supported",

396 "get_num_sms",

397 "should_use_deepgemm_for_fp8_linear",

414 - "should_use_deepgemm_for_fp8_linear_for_nk",

398 "get_col_major_tma_aligned_tensor",

399 "get_mk_alignment_for_contiguous_layout",

400 ]