@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 287 lines 9.11 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit test for the MLADualRMSNormFusionPass.
4
5 The pass fuses paired q/kv RMS norms in MLA attention into a single
6 fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel.
7 """
8
9 import pytest
10 import torch
11
12 import vllm.config
13 from tests.compile.backend import TestBackend
14 from vllm._aiter_ops import (
15 is_aiter_found_and_supported,
16 rocm_aiter_ops,
17 )
18 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
19 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
20 from vllm.config import (
21 CompilationConfig,
22 CompilationMode,
23 ModelConfig,
24 PassConfig,
25 VllmConfig,
26 )
27 from vllm.model_executor.layers.layernorm import RMSNorm
28 from vllm.platforms import current_platform
29
30 # MLA attention geometry for DeepSeek-V3 / Kimi-K2
31 Q_DIM = 1536
32 KV_C_DIM = 512
33 K_PE_DIM = 64
34 EPS = 1e-6
35
36 FP8_DTYPE = current_platform.fp8_dtype()
37
38
39 class MLADualRMSNormTestModel(torch.nn.Module):
40 """Minimal model reproducing the MLA dual RMS norm pattern:
41 linear -> split([q_dim, kv_dim])
42 +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear
43 +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
44 +-- kv_c (getitem 0) -> rms_norm(kv_w, eps)
45 +-- k_pe
46 """
47
48 def __init__(
49 self,
50 hidden_size: int,
51 q_dim: int = Q_DIM,
52 kv_c_dim: int = KV_C_DIM,
53 k_pe_dim: int = K_PE_DIM,
54 eps: float = EPS,
55 ):
56 super().__init__()
57 self.q_dim = q_dim
58 self.kv_dim = kv_c_dim + k_pe_dim
59 self.kv_c_dim = kv_c_dim
60 self.k_pe_dim = k_pe_dim
61
62 self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False)
63 self.q_norm = RMSNorm(q_dim, eps=eps)
64 self.kv_norm = RMSNorm(kv_c_dim, eps=eps)
65 self.q_b_proj = torch.nn.Linear(q_dim, hidden_size, bias=False)
66
67 def forward(
68 self, x: torch.Tensor
69 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
70 # Avoid graph input being a direct arg to a matched pattern node
71 x = torch.relu(x)
72
73 projected = self.proj(x)
74
75 q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1)
76 kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1)
77
78 q_normed = self.q_norm(q_c)
79 kv_normed = self.kv_norm(kv_c)
80
81 q_out = self.q_b_proj(q_normed)
82 return q_out, kv_normed, k_pe
83
84 def ops_in_model_before(self):
85 return [torch.ops.vllm_ir.rms_norm.default]
86
87 def ops_in_model_after(self):
88 return [torch.ops.vllm.fused_mla_dual_rms_norm.default]
89
90
91 @pytest.mark.parametrize("dtype", [torch.bfloat16])
92 @pytest.mark.parametrize("hidden_size", [7168])
93 @pytest.mark.skipif(
94 not is_aiter_found_and_supported(),
95 reason="Only test on ROCm with AITER installed and supported",
96 )
97 def test_fuse_mla_dual_rms_norm(
98 dtype: torch.dtype,
99 hidden_size: int,
100 monkeypatch: pytest.MonkeyPatch,
101 ):
102 torch._dynamo.reset()
103
104 vllm_config = VllmConfig(
105 model_config=ModelConfig(dtype=dtype),
106 compilation_config=CompilationConfig(
107 mode=CompilationMode.VLLM_COMPILE,
108 custom_ops=["+rms_norm"],
109 pass_config=PassConfig(
110 fuse_mla_dual_rms_norm=True,
111 eliminate_noops=True,
112 ),
113 ),
114 )
115
116 with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
117 from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
118 MLADualRMSNormFusionPass,
119 )
120
121 torch.set_default_device("cuda")
122 torch.set_default_dtype(dtype)
123 torch.manual_seed(42)
124
125 m.setenv("VLLM_ROCM_USE_AITER", "1")
126 rocm_aiter_ops.refresh_env_variables()
127
128 fusion_pass = MLADualRMSNormFusionPass(vllm_config)
129 passes = [
130 NoOpEliminationPass(vllm_config),
131 fusion_pass,
132 PostCleanupPass(vllm_config),
133 ]
134 backend = TestBackend(*passes)
135 model = MLADualRMSNormTestModel(hidden_size)
136
137 x = torch.randn(1, hidden_size)
138 torch._dynamo.mark_dynamic(x, 0)
139
140 outputs_unfused = model(x)
141
142 model_fused = torch.compile(model, backend=backend)
143 outputs_fused = model_fused(x)
144
145 torch.testing.assert_close(outputs_unfused, outputs_fused, atol=1e-2, rtol=1e-2)
146
147 assert fusion_pass.matched_count == 1, (
148 f"Expected 1 fused pair, got {fusion_pass.matched_count}"
149 )
150
151 backend.check_before_ops(model.ops_in_model_before())
152 backend.check_after_ops(model.ops_in_model_after())
153
154
155 class MLADualRMSNormFp8PerTokenTestModel(torch.nn.Module):
156 """Minimal model reproducing the FP8 MLA attention path with *per-token* quant:
157 linear -> split([q_dim, kv_dim])
158 +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant
159 +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
160 +-- kv_c (getitem 0) -> rms_norm (bf16)
161 +-- k_pe
162 """
163
164 def __init__(
165 self,
166 hidden_size: int,
167 q_dim: int = Q_DIM,
168 kv_c_dim: int = KV_C_DIM,
169 k_pe_dim: int = K_PE_DIM,
170 eps: float = EPS,
171 ):
172 super().__init__()
173 self.q_dim = q_dim
174 self.kv_dim = kv_c_dim + k_pe_dim
175 self.kv_c_dim = kv_c_dim
176 self.k_pe_dim = k_pe_dim
177 self.eps = eps
178
179 self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False)
180 self.q_weight = torch.nn.Parameter(torch.ones(q_dim))
181 self.kv_norm = RMSNorm(kv_c_dim, eps=eps)
182
183 def _dequant(self, x_fp8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
184 # Per-token: a single (M, 1) scale broadcast across the row.
185 return (x_fp8.to(torch.float32) * scale).to(torch.bfloat16)
186
187 def forward(self, x: torch.Tensor):
188 # Avoid graph input being a direct arg to a matched pattern node
189 x = torch.relu(x)
190
191 projected = self.proj(x)
192
193 q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1)
194 kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1)
195
196 q_fp8, q_scale = torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant(
197 q_c, self.q_weight, self.eps, FP8_DTYPE
198 )
199 kv_normed = self.kv_norm(kv_c)
200
201 return self._dequant(q_fp8, q_scale), kv_normed, k_pe
202
203 def ops_in_model_before(self):
204 return [
205 torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant.default,
206 torch.ops.vllm_ir.rms_norm.default,
207 ]
208
209 def ops_in_model_after(self):
210 return [torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default]
211
212
213 @pytest.mark.parametrize("dtype", [torch.bfloat16])
214 @pytest.mark.parametrize("hidden_size", [7168])
215 @pytest.mark.skipif(
216 not is_aiter_found_and_supported(),
217 reason="Only test on ROCm with AITER installed and supported",
218 )
219 def test_fuse_mla_dual_rms_norm_fp8_per_token(
220 dtype: torch.dtype,
221 hidden_size: int,
222 monkeypatch: pytest.MonkeyPatch,
223 ):
224 torch._dynamo.reset()
225
226 vllm_config = VllmConfig(
227 model_config=ModelConfig(dtype=dtype),
228 compilation_config=CompilationConfig(
229 mode=CompilationMode.VLLM_COMPILE,
230 custom_ops=["+rms_norm"],
231 pass_config=PassConfig(
232 fuse_mla_dual_rms_norm=True,
233 eliminate_noops=True,
234 ),
235 ),
236 )
237
238 with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
239 from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
240 MLADualRMSNormFusionPass,
241 )
242
243 torch.set_default_device("cuda")
244 torch.set_default_dtype(dtype)
245 torch.manual_seed(42)
246
247 m.setenv("VLLM_ROCM_USE_AITER", "1")
248 rocm_aiter_ops.refresh_env_variables()
249
250 fusion_pass = MLADualRMSNormFusionPass(vllm_config)
251 passes = [
252 NoOpEliminationPass(vllm_config),
253 fusion_pass,
254 PostCleanupPass(vllm_config),
255 ]
256 backend = TestBackend(*passes)
257 model = MLADualRMSNormFp8PerTokenTestModel(hidden_size)
258
259 x = torch.randn(4, hidden_size)
260 torch._dynamo.mark_dynamic(x, 0)
261
262 with torch.inference_mode():
263 outputs_unfused = model(x)
264
265 model_fused = torch.compile(model, backend=backend)
266 outputs_fused = model_fused(x)
267
268 q_deq_u, kv_normed_u, k_pe_u = outputs_unfused
269 q_deq_f, kv_normed_f, k_pe_f = outputs_fused
270
271 torch.testing.assert_close(k_pe_u, k_pe_f, atol=0, rtol=0)
272
273 torch.testing.assert_close(kv_normed_u, kv_normed_f, atol=1e-2, rtol=1e-2)
274
275 E4M3_STEP = 0.125
276 exact_frac = (q_deq_u == q_deq_f).float().mean().item()
277 assert exact_frac > 0.99, (
278 f"q: only {exact_frac:.4f} of elements bit-exact; scales likely differ"
279 )
280 torch.testing.assert_close(q_deq_u, q_deq_f, atol=1e-2, rtol=E4M3_STEP)
281
282 assert fusion_pass.matched_count == 1, (
283 f"Expected 1 fused pair, got {fusion_pass.matched_count}"
284 )
285
286 backend.check_before_ops(model.ops_in_model_before())
287 backend.check_after_ops(model.ops_in_model_after())