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 3, 22:45 UTC
[ROCm][Perf] Fix RMSNorm+Quant fusion for gfx950 (non-fnuz) (#41825)
Signed-off-by: Frida Andersson <fanderss@amd.com> Signed-off-by: Chuan Li <chuali@amd.com> Co-authored-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Chuan Li <chuali@amd.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Frida Andersson <frida-andersson@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
frida-andersson committed
May 11, 2026 at 21:00 UTC
a721315488429e728e25adabb6d94c35a4770b94
3 files changed
+342
-7
tests/compile/passes/test_double_aiter_rms_quant_fusion.py
new
+166
@@ -0,0 +1,166 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for the DoubleQuant fan-out variants registered by
5
+``RocmAiterRMSNormQuantFusionPass``.
6
+
7
+Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds
8
+two distinct ``rocm_aiter_group_fp8_quant`` consumers and rewrite it into
9
+two independent fused ``rms_norm + group_fp8_quant`` ops:
10
+
11
+* ``DoubleAiterRMSFp8GroupQuantPattern`` matches the un-viewed shape
12
+ (e.g. Kimi-K2.5 / DSR1).
13
+* ``DoubleAiterRMSFp8GroupQuantViewPattern`` (this PR) is the view-tolerant
14
+ sibling that additionally matches the
15
+ ``rms_norm -> view -> group_fp8_quant`` shape that DSv3.2's MLA indexer
16
+ q_c norm exposes through ``Fp8BlockScaledMMLinearKernel.apply_weights``'s
17
+ 2D-flatten boilerplate.
18
+"""
19
+
20
+import pytest
21
+import torch
22
+
23
+import vllm.config
24
+from tests.compile.backend import TestBackend
25
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
26
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
27
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
28
+from vllm.config import (
29
+ CompilationConfig,
30
+ CompilationMode,
31
+ ModelConfig,
32
+ PassConfig,
33
+ VllmConfig,
34
+)
35
+
36
+EPS = 1e-5
37
+HIDDEN_SIZE = 256
38
+GROUP_SIZE = 128
39
+
40
+
41
+class _NoViewDoubleQuantModel(torch.nn.Module):
42
+ """``rms_norm -> 2x group_fp8_quant`` fan-out (Kimi-K2.5 / DSR1 shape)."""
43
+
44
+ def __init__(self) -> None:
45
+ super().__init__()
46
+ self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
47
+
48
+ def forward(
49
+ self, x: torch.Tensor
50
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
51
+ # avoid graph input being a direct arg to a matched pattern node
52
+ x = torch.relu(x)
53
+ rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
54
+ q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
55
+ q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
56
+ return q1, s1, q2, s2
57
+
58
+
59
+class _ViewDoubleQuantModel(torch.nn.Module):
60
+ """``rms_norm -> view -> 2x group_fp8_quant`` fan-out (DSv3.2 shape).
61
+
62
+ Reproduces the FX-graph shape produced by ``Fp8BlockScaledMMLinearKernel``'s
63
+ 2D-flatten before the FP8 group quant op.
64
+ """
65
+
66
+ def __init__(self) -> None:
67
+ super().__init__()
68
+ self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
69
+
70
+ def forward(
71
+ self, x: torch.Tensor
72
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
73
+ x = torch.relu(x)
74
+ rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
75
+ view = rms.view(-1, rms.shape[-1])
76
+ q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
77
+ q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
78
+ return q1, s1, q2, s2
79
+
80
+
81
+@pytest.mark.parametrize(
82
+ "model_cls",
83
+ [_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
84
+ ids=["no_view", "with_view"],
85
+)
86
+@pytest.mark.skipif(
87
+ not is_aiter_found_and_supported(),
88
+ reason="Only test on ROCm with AITER installed and supported",
89
+)
90
+def test_double_aiter_rms_fp8_group_quant_fusion(
91
+ model_cls: type[torch.nn.Module],
92
+ monkeypatch: pytest.MonkeyPatch,
93
+) -> None:
94
+ """
95
+ Both fan-out shapes (with and without an intermediate view) must fuse
96
+ into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via
97
+ ``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the
98
+ new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling.
99
+
100
+ A failure on the ``with_view`` parametrization is a regression on the
101
+ DSv3.2 q_c norm path that this PR's view-tolerant pattern is intended
102
+ to cover.
103
+ """
104
+ torch._dynamo.reset()
105
+
106
+ vllm_config = VllmConfig(
107
+ model_config=ModelConfig(dtype=torch.bfloat16),
108
+ compilation_config=CompilationConfig(
109
+ mode=CompilationMode.VLLM_COMPILE,
110
+ custom_ops=["+rms_norm", "+quant_fp8"],
111
+ pass_config=PassConfig(
112
+ fuse_norm_quant=True,
113
+ eliminate_noops=True,
114
+ ),
115
+ ),
116
+ )
117
+
118
+ with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
119
+ from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
120
+ RocmAiterRMSNormQuantFusionPass,
121
+ )
122
+
123
+ torch.set_default_device("cuda")
124
+ torch.set_default_dtype(torch.bfloat16)
125
+ torch.manual_seed(0)
126
+
127
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
128
+ rocm_aiter_ops.refresh_env_variables()
129
+
130
+ fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
131
+ passes = [
132
+ NoOpEliminationPass(vllm_config),
133
+ fusion_pass,
134
+ PostCleanupPass(vllm_config),
135
+ ]
136
+ backend = TestBackend(*passes)
137
+ model = model_cls()
138
+
139
+ x = torch.randn(8, HIDDEN_SIZE)
140
+ torch._dynamo.mark_dynamic(x, 0)
141
+
142
+ outputs_unfused = model(x)
143
+ model_fused = torch.compile(model, backend=backend)
144
+ outputs_fused = model_fused(x)
145
+
146
+ # Both consumers must be rewritten into the fused op (one
147
+ # ``register_replacement`` rewrite covers the whole 1-to-2 fan-out).
148
+ assert fusion_pass.matched_count == 1, (
149
+ f"Expected the {model_cls.__name__} fan-out to fuse via the "
150
+ f"DoubleQuant pattern (matched_count == 1), got "
151
+ f"{fusion_pass.matched_count}"
152
+ )
153
+
154
+ fused_op = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
155
+ backend.check_after_ops([fused_op])
156
+
157
+ # Numerical parity sanity-check: the fused pair must match the
158
+ # unfused pair on FP8 outputs (exact byte-equality is the goal,
159
+ # but allow a tiny tolerance for any residual numeric noise).
160
+ for fused_t, unfused_t in zip(outputs_fused, outputs_unfused):
161
+ torch.testing.assert_close(
162
+ fused_t.to(torch.float32),
163
+ unfused_t.to(torch.float32),
164
+ atol=1e-2,
165
+ rtol=1e-2,
166
+ )
vllm/compilation/passes/fusion/matcher_utils.py
+1
-6
@@ -273,12 +273,7 @@ class MatcherQuantFP8(MatcherCustomOp):
273
"ROCm aiter fusion pass currently supports "
274
"quantization operation with group_size 128"
275
)
276
- if current_platform.is_fp8_fnuz():
277
- self.QUANT_OP = rocm_aiter_ops.get_group_quant_op()
278
- else:
279
- self.QUANT_OP = (
280
- torch.ops.vllm.triton_per_token_group_quant_fp8.default
281
- )
276
+ self.QUANT_OP = rocm_aiter_ops.get_group_quant_op()
277
278
else:
279
assert quant_key in QUANT_OPS, (
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+175
-1
@@ -6,6 +6,7 @@ from typing import Any
6
import torch
7
import torch._inductor.pattern_matcher as pm
8
from torch import fx
9
+from torch._inductor.fx_passes.post_grad import view_to_reshape
10
from torch._inductor.pattern_matcher import PatternMatcherPass
11
12
import vllm.ir.ops
@@ -293,6 +294,161 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
294
pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
295
296
297
+class DoubleAiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
298
+ """
299
+ Pattern matching ``rms_norm`` whose output feeds *two* distinct
300
+ ``rocm_aiter_group_fp8_quant`` consumers, replacing it with two
301
+ independent fused ``rms_norm_group_fp8_quant`` ops.
302
+
303
+ Repeating the rms_norm in the replacement is preferable to leaving
304
+ the fused 16-bit rms output materialized for two unfused quant
305
+ consumers, and matches what the previous manual graph surgery
306
+ achieved by cloning the rms_norm node.
307
+ """
308
+
309
+ FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
310
+
311
+ def __init__(
312
+ self,
313
+ epsilon: float,
314
+ quant_dtype: torch.dtype,
315
+ group_shape: GroupShape,
316
+ match_aiter_quant: bool = True,
317
+ symmetric: bool = True,
318
+ ) -> None:
319
+ scale = ScaleDesc(torch.float32, False, group_shape)
320
+ key = FusedRMSQuantKey(
321
+ fused_add=False,
322
+ quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
323
+ )
324
+
325
+ super().__init__(epsilon, key, match_aiter_quant)
326
+
327
+ def register(self, pm_pass: PatternMatcherPass) -> None:
328
+ def pattern(
329
+ input: torch.Tensor,
330
+ weight: torch.Tensor,
331
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
332
+ result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
333
+ result1, scale1 = self.quant_matcher(result_rms)
334
+ result2, scale2 = self.quant_matcher(result_rms)
335
+ return result1, scale1, result2, scale2
336
+
337
+ def replacement(
338
+ input: torch.Tensor,
339
+ weight: torch.Tensor,
340
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
341
+ at1 = self.FUSED_OP(
342
+ x=input,
343
+ weight=weight,
344
+ variance_epsilon=self.epsilon,
345
+ group_size=128,
346
+ )
347
+ at2 = self.FUSED_OP(
348
+ x=input,
349
+ weight=weight,
350
+ variance_epsilon=self.epsilon,
351
+ group_size=128,
352
+ )
353
+
354
+ return at1[0], at1[1], at2[0], at2[1]
355
+
356
+ pm.register_replacement(
357
+ pattern,
358
+ replacement,
359
+ # input, weight
360
+ [self.empty(5, 16), self.empty(16)],
361
+ pm.fwd_only,
362
+ pm_pass,
363
+ )
364
+
365
+
366
+class DoubleAiterRMSFp8GroupQuantViewPattern(AiterRMSNormQuantPattern):
367
+ """
368
+ View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``.
369
+
370
+ Matches the same 1-to-2 fan-out, but with a ``view``/``reshape`` between
371
+ the ``rms_norm`` output and the two ``rocm_aiter_group_fp8_quant``
372
+ consumers::
373
+
374
+ rms_norm -> view -> rocm_aiter_group_fp8_quant
375
+ \\-> view -> rocm_aiter_group_fp8_quant
376
+
377
+ This shape arises in DeepSeek-V3.2's MLA indexer q_c norm, where the
378
+ FP8 linear path's 2D-flatten boilerplate
379
+ (``Fp8BlockScaledMMLinearKernel.apply_weights``) inserts a view between
380
+ the rms_norm output and each FP8 group quant op. The non-view sibling
381
+ pattern silently no-ops on this graph because the pattern matcher
382
+ requires the in-graph and in-pattern node shapes to align.
383
+
384
+ The trace_fn runs Inductor's ``view_to_reshape`` post-grad pass to
385
+ normalize ``view`` to ``reshape`` in both the pattern and the input
386
+ graph, widening the match without touching the no-view sibling.
387
+ """
388
+
389
+ FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
390
+
391
+ def __init__(
392
+ self,
393
+ epsilon: float,
394
+ quant_dtype: torch.dtype,
395
+ group_shape: GroupShape,
396
+ match_aiter_quant: bool = True,
397
+ symmetric: bool = True,
398
+ ) -> None:
399
+ scale = ScaleDesc(torch.float32, False, group_shape)
400
+ key = FusedRMSQuantKey(
401
+ fused_add=False,
402
+ quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
403
+ )
404
+
405
+ super().__init__(epsilon, key, match_aiter_quant)
406
+
407
+ def register(self, pm_pass: PatternMatcherPass) -> None:
408
+ def pattern(
409
+ input: torch.Tensor,
410
+ weight: torch.Tensor,
411
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
412
+ result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
413
+ view_rms = result_rms.view(-1, result_rms.shape[-1])
414
+ result1, scale1 = self.quant_matcher(view_rms)
415
+ result2, scale2 = self.quant_matcher(view_rms)
416
+ return result1, scale1, result2, scale2
417
+
418
+ def replacement(
419
+ input: torch.Tensor,
420
+ weight: torch.Tensor,
421
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
422
+ at1 = self.FUSED_OP(
423
+ x=input,
424
+ weight=weight,
425
+ variance_epsilon=self.epsilon,
426
+ group_size=128,
427
+ )
428
+ at2 = self.FUSED_OP(
429
+ x=input,
430
+ weight=weight,
431
+ variance_epsilon=self.epsilon,
432
+ group_size=128,
433
+ )
434
+
435
+ return at1[0], at1[1], at2[0], at2[1]
436
+
437
+ def trace_with_view_to_reshape(*args: Any, **kwargs: Any) -> fx.GraphModule:
438
+ gm = pm.fwd_only(*args, **kwargs)
439
+ view_to_reshape(gm)
440
+ return gm
441
+
442
+ pm.register_replacement(
443
+ pattern,
444
+ replacement,
445
+ # input, weight
446
+ [self.empty(5, 16), self.empty(16)],
447
+ trace_with_view_to_reshape,
448
+ pm_pass,
449
+ )
450
+
451
+
452
class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
453
"""
454
This pass fuses aiter rms_norm & vllm/aiter quant custom ops
@@ -309,8 +465,24 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
465
)
466
467
# Make sure fused add patterns are before simple rms norm,
312
- # as the latter is a subset of the former in torch ops
468
+ # as the latter is a subset of the former in torch ops.
469
+ # The DoubleQuant patterns handle 1 rms_norm -> 2 group_fp8_quant
470
+ # fan-out (e.g. DSv3.2) and must be registered before the single
471
+ # group-quant pattern so they match first. The view-tolerant variant
472
+ # additionally covers the rms_norm -> view -> 2x quant shape that
473
+ # appears when the FP8 linear path inserts a 2D-flatten boilerplate
474
+ # (DSv3.2 MLA indexer q_c norm).
475
for epsilon in [1e-5, 1e-6]:
476
+ # Fuse aiter rms_norm + 2x aiter group fp8 quant
477
+ DoubleAiterRMSFp8GroupQuantPattern(
478
+ epsilon, FP8_DTYPE, GroupShape(1, 128)
479
+ ).register(self.patterns)
480
+
481
+ # View-tolerant sibling for DSv3.2 q_c norm fan-out
482
+ DoubleAiterRMSFp8GroupQuantViewPattern(
483
+ epsilon, FP8_DTYPE, GroupShape(1, 128)
484
+ ).register(self.patterns)
485
+
486
# Fuse aiter rms_norm + aiter dynamic group fp8 quant
487
AiterRMSFp8GroupQuantPattern(
488
epsilon, FP8_DTYPE, GroupShape(1, 128)
@@ -360,6 +532,8 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
532
AiterFusedAddRMSNormDynamicQuantPattern,
533
AiterRMSFp8GroupQuantPattern,
534
AiterFusedAddRMSFp8GroupQuantPattern,
535
+ DoubleAiterRMSFp8GroupQuantPattern,
536
+ DoubleAiterRMSFp8GroupQuantViewPattern,
537
]
538
return self.hash_source(self, *fusion_patterns)
539