@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 5.87 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for the DoubleQuant fan-out variants registered by
4 ``RocmAiterRMSNormQuantFusionPass``.
5
6 Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds
7 two distinct ``rocm_aiter_group_fp8_quant`` consumers and rewrite it into
8 two independent fused ``rms_norm + group_fp8_quant`` ops:
9
10 * ``DoubleAiterRMSFp8GroupQuantPattern`` matches the un-viewed shape
11 (e.g. Kimi-K2.5 / DSR1).
12 * ``DoubleAiterRMSFp8GroupQuantViewPattern`` (this PR) is the view-tolerant
13 sibling that additionally matches the
14 ``rms_norm -> view -> group_fp8_quant`` shape that DSv3.2's MLA indexer
15 q_c norm exposes through ``Fp8BlockScaledMMLinearKernel.apply_weights``'s
16 2D-flatten boilerplate.
17 """
18
19 import pytest
20 import torch
21
22 import vllm.config
23 from tests.compile.backend import TestBackend
24 from vllm._aiter_ops import rocm_aiter_ops
25 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
26 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
27 from vllm.config import (
28 CompilationConfig,
29 CompilationMode,
30 ModelConfig,
31 PassConfig,
32 VllmConfig,
33 )
34
35 EPS = 1e-5
36 HIDDEN_SIZE = 256
37 GROUP_SIZE = 128
38
39
40 class _NoViewDoubleQuantModel(torch.nn.Module):
41 """``rms_norm -> 2x group_fp8_quant`` fan-out (Kimi-K2.5 / DSR1 shape)."""
42
43 def __init__(self) -> None:
44 super().__init__()
45 self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
46
47 def forward(
48 self, x: torch.Tensor
49 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
50 # avoid graph input being a direct arg to a matched pattern node
51 x = torch.relu(x)
52 rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
53 q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
54 q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
55 return q1, s1, q2, s2
56
57
58 class _ViewDoubleQuantModel(torch.nn.Module):
59 """``rms_norm -> view -> 2x group_fp8_quant`` fan-out (DSv3.2 shape).
60
61 Reproduces the FX-graph shape produced by ``Fp8BlockScaledMMLinearKernel``'s
62 2D-flatten before the FP8 group quant op.
63 """
64
65 def __init__(self) -> None:
66 super().__init__()
67 self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
68
69 def forward(
70 self, x: torch.Tensor
71 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
72 x = torch.relu(x)
73 rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
74 view = rms.view(-1, rms.shape[-1])
75 q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
76 q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
77 return q1, s1, q2, s2
78
79
80 @pytest.mark.parametrize(
81 "model_cls",
82 [_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
83 ids=["no_view", "with_view"],
84 )
85 @pytest.mark.skip(
86 reason="Skipping for now because pytorch compiler removes one the two quant ops"
87 )
88 def test_double_aiter_rms_fp8_group_quant_fusion(
89 model_cls: type[torch.nn.Module],
90 monkeypatch: pytest.MonkeyPatch,
91 ) -> None:
92 """Both fan-out shapes (with and without an intermediate view) must fuse
93 into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via
94 ``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the
95 new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling.
96
97 A failure on the ``with_view`` parametrization is a regression on the
98 DSv3.2 q_c norm path that this PR's view-tolerant pattern is intended
99 to cover.
100 """
101 torch._dynamo.reset()
102
103 vllm_config = VllmConfig(
104 model_config=ModelConfig(dtype=torch.bfloat16),
105 compilation_config=CompilationConfig(
106 mode=CompilationMode.VLLM_COMPILE,
107 custom_ops=["+rms_norm", "+quant_fp8"],
108 pass_config=PassConfig(
109 fuse_norm_quant=True,
110 eliminate_noops=True,
111 ),
112 ),
113 )
114
115 with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
116 from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
117 RocmAiterRMSNormQuantFusionPass,
118 )
119
120 torch.set_default_device("cuda")
121 torch.set_default_dtype(torch.bfloat16)
122 torch.manual_seed(0)
123
124 m.setenv("VLLM_ROCM_USE_AITER", "1")
125 rocm_aiter_ops.refresh_env_variables()
126
127 fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
128 passes = [
129 NoOpEliminationPass(vllm_config),
130 fusion_pass,
131 PostCleanupPass(vllm_config),
132 ]
133 backend = TestBackend(*passes)
134 model = model_cls()
135
136 x = torch.randn(8, HIDDEN_SIZE)
137 torch._dynamo.mark_dynamic(x, 0)
138
139 outputs_unfused = model(x)
140 model_fused = torch.compile(model, backend=backend)
141 outputs_fused = model_fused(x)
142
143 # Both consumers must be rewritten into the fused op (one
144 # ``register_replacement`` rewrite covers the whole 1-to-2 fan-out).
145 assert fusion_pass.matched_count == 1, (
146 f"Expected the {model_cls.__name__} fan-out to fuse via the "
147 f"DoubleQuant pattern (matched_count == 1), got "
148 f"{fusion_pass.matched_count}"
149 )
150
151 fused_op = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
152 backend.check_after_ops([fused_op])
153
154 # Numerical parity sanity-check: the fused pair must match the
155 # unfused pair on FP8 outputs (exact byte-equality is the goal,
156 # but allow a tiny tolerance for any residual numeric noise).
157 for fused_t, unfused_t in zip(outputs_fused, outputs_unfused):
158 torch.testing.assert_close(
159 fused_t.to(torch.float32),
160 unfused_t.to(torch.float32),
161 atol=1e-2,
162 rtol=1e-2,
163 )