@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 284 lines 9.84 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6 from torch._ops import OpOverload, OpOverloadPacket
7
8 from tests.compile.backend import TestBackend
9 from vllm.compilation.passes.fusion.matcher_utils import (
10 FLASHINFER_ROTARY_OP,
11 ROTARY_OP,
12 )
13 from vllm.compilation.passes.fusion.qk_norm_rope_fusion import (
14 FUSED_QK_ROPE_OP,
15 QKNormRoPEFusionPass,
16 )
17 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
18 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
19 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
20 from vllm.config import (
21 CompilationConfig,
22 CompilationMode,
23 ModelConfig,
24 PassConfig,
25 VllmConfig,
26 set_current_vllm_config,
27 )
28 from vllm.model_executor.layers.attention import Attention
29 from vllm.model_executor.layers.layernorm import RMSNorm
30 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
31 from vllm.platforms import current_platform
32 from vllm.v1.attention.backend import AttentionType
33
34 RSQRT_OP = torch.ops.aten.rsqrt.default
35 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
36
37
38 class QKNormRoPETestModel(torch.nn.Module):
39 def __init__(
40 self,
41 *,
42 num_heads: int,
43 num_kv_heads: int,
44 head_dim: int,
45 eps: float,
46 is_neox: bool,
47 vllm_config: VllmConfig,
48 dtype: torch.dtype,
49 test_scattered_split: bool = False,
50 prefix: str = "model.layers.0.self_attn.attn",
51 ) -> None:
52 super().__init__()
53 self.num_heads = num_heads
54 self.num_kv_heads = num_kv_heads
55 self.head_dim = head_dim
56 self.q_size = num_heads * head_dim
57 self.kv_size = num_kv_heads * head_dim
58 self.rotary_dim = head_dim
59 self.eps = eps
60 self.dtype = dtype
61
62 # Register layer metadata for the fusion pass via Attention.
63 self.attn = Attention(
64 num_heads=self.num_heads,
65 head_size=self.head_dim,
66 scale=1.0 / self.head_dim**0.5,
67 num_kv_heads=self.num_kv_heads,
68 cache_config=vllm_config.cache_config,
69 prefix=prefix,
70 attn_type=AttentionType.DECODER,
71 )
72
73 self.q_norm = RMSNorm(self.head_dim, eps=self.eps)
74 self.k_norm = RMSNorm(self.head_dim, eps=self.eps)
75 self.rotary_emb = RotaryEmbedding(
76 self.head_dim,
77 rotary_dim=self.rotary_dim,
78 max_position_embeddings=4096,
79 base=10000,
80 is_neox_style=is_neox,
81 dtype=self.dtype,
82 )
83 self.test_scattered_split = test_scattered_split
84 self.enable_rms_norm_custom_op = self.q_norm.enabled()
85 self.enable_rope_custom_op = self.rotary_emb.enabled()
86
87 def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
88 if self.test_scattered_split:
89 q, _, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
90 _, k, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
91 _, _, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
92 else:
93 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
94 q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
95 q_by_head = self.q_norm(q_by_head)
96 q = q_by_head.view(q.shape)
97 k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)
98 k_by_head = self.k_norm(k_by_head)
99 k = k_by_head.view(k.shape)
100 q, k = self.rotary_emb(positions, q, k)
101 return q, k, v
102
103 def ops_in_model_before(self) -> list[OpOverload | OpOverloadPacket]:
104 ops: list[OpOverload | OpOverloadPacket] = [torch.ops.vllm_ir.rms_norm]
105 if self.enable_rope_custom_op:
106 if self.rotary_emb.use_flashinfer:
107 ops.append(FLASHINFER_ROTARY_OP)
108 else:
109 ops.append(ROTARY_OP)
110 else:
111 ops.append(INDEX_SELECT_OP)
112 return ops
113
114 def ops_in_model_after(self) -> list[OpOverload | OpOverloadPacket]:
115 return [FUSED_QK_ROPE_OP]
116
117
118 @pytest.mark.parametrize("scattered_split", [True, False])
119 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
120 @pytest.mark.parametrize("is_neox", [True, False])
121 @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
122 @pytest.mark.parametrize("enable_rope_custom_op", [True])
123 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
124 @pytest.mark.skipif(
125 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
126 reason="Only test on cuda and rocm platform",
127 )
128 def test_qk_norm_rope_fusion(
129 eps,
130 is_neox,
131 enable_rms_norm_custom_op,
132 enable_rope_custom_op,
133 dtype,
134 scattered_split,
135 ):
136 if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
137 pytest.skip("fused_qk_norm_rope custom op not available")
138
139 torch.set_default_device(current_platform.device_type)
140 torch.set_default_dtype(dtype)
141 torch.manual_seed(0)
142
143 custom_ops: list[str] = []
144 if enable_rms_norm_custom_op:
145 custom_ops.append("+rms_norm")
146 if enable_rope_custom_op:
147 custom_ops.append("+rotary_embedding")
148
149 vllm_config = VllmConfig(
150 model_config=ModelConfig(dtype=dtype),
151 compilation_config=CompilationConfig(
152 mode=CompilationMode.VLLM_COMPILE,
153 custom_ops=custom_ops,
154 pass_config=PassConfig(
155 enable_qk_norm_rope_fusion=True,
156 eliminate_noops=True,
157 ),
158 ),
159 )
160
161 num_heads, num_kv_heads, head_dim = 16, 4, 128
162 T = 5
163
164 with (
165 set_current_vllm_config(vllm_config),
166 vllm_config.kernel_config.ir_op_priority.set_priority(),
167 ):
168 model = QKNormRoPETestModel(
169 num_heads=num_heads,
170 num_kv_heads=num_kv_heads,
171 head_dim=head_dim,
172 eps=eps,
173 is_neox=is_neox,
174 vllm_config=vllm_config,
175 dtype=dtype,
176 test_scattered_split=scattered_split,
177 )
178
179 noop_pass = NoOpEliminationPass(vllm_config)
180 coalesce_pass = SplitCoalescingPass(vllm_config)
181 fusion_pass = QKNormRoPEFusionPass(vllm_config)
182 cleanup_pass = PostCleanupPass(vllm_config)
183
184 backend = TestBackend(noop_pass, coalesce_pass, fusion_pass, cleanup_pass)
185 backend_baseline = TestBackend(noop_pass, cleanup_pass)
186
187 qkv = torch.randn(T, model.q_size + 2 * model.kv_size)
188 pos = torch.arange(T, dtype=torch.long, device=qkv.device)
189 qkv_unfused = qkv.clone()
190 pos_unfused = pos.clone()
191
192 torch._dynamo.mark_dynamic(qkv, 0)
193 torch._dynamo.mark_dynamic(pos, 0)
194 model_fused = torch.compile(model, backend=backend)
195 q_fused, k_fused, v_fused = model_fused(qkv, pos)
196
197 torch._dynamo.mark_dynamic(qkv_unfused, 0)
198 torch._dynamo.mark_dynamic(pos_unfused, 0)
199 model_unfused = torch.compile(model, backend=backend_baseline)
200 q_unfused, k_unfused, v_unfused = model_unfused(qkv_unfused, pos_unfused)
201
202 if dtype == torch.float16:
203 ATOL, RTOL = (2e-3, 2e-3)
204 else:
205 ATOL, RTOL = (1e-2, 1e-2)
206
207 torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
208 torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
209 torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
210
211 assert fusion_pass.matched_count == 1
212
213 backend.check_before_ops(model.ops_in_model_before())
214 backend.check_after_ops(model.ops_in_model_after())
215
216
217 @pytest.mark.skipif(
218 not current_platform.is_cuda_alike(),
219 reason="Only test on CUDA and ROCm platforms",
220 )
221 def test_qk_norm_rope_fusion_with_heterogeneous_draft_model():
222 """The fusion pass should support target and draft attention geometries."""
223 if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
224 pytest.skip("fused_qk_norm_rope custom op not available")
225
226 dtype = torch.bfloat16
227 torch.set_default_device(current_platform.device_type)
228 torch.set_default_dtype(dtype)
229 torch.manual_seed(0)
230
231 vllm_config = VllmConfig(
232 model_config=ModelConfig(dtype=dtype),
233 compilation_config=CompilationConfig(
234 mode=CompilationMode.VLLM_COMPILE,
235 custom_ops=["+rms_norm", "+rotary_embedding"],
236 pass_config=PassConfig(
237 enable_qk_norm_rope_fusion=True,
238 eliminate_noops=True,
239 ),
240 ),
241 )
242
243 with (
244 set_current_vllm_config(vllm_config),
245 vllm_config.kernel_config.ir_op_priority.set_priority(),
246 ):
247 target_model = QKNormRoPETestModel(
248 num_heads=32,
249 num_kv_heads=8,
250 head_dim=128,
251 eps=1e-6,
252 is_neox=True,
253 vllm_config=vllm_config,
254 dtype=dtype,
255 prefix="model.layers.0.self_attn.attn",
256 )
257 draft_model = QKNormRoPETestModel(
258 num_heads=16,
259 num_kv_heads=8,
260 head_dim=128,
261 eps=1e-6,
262 is_neox=True,
263 vllm_config=vllm_config,
264 dtype=dtype,
265 prefix="draft_model.layers.0.self_attn.attn",
266 )
267
268 fusion_pass = QKNormRoPEFusionPass(vllm_config)
269 backend = TestBackend(
270 NoOpEliminationPass(vllm_config),
271 SplitCoalescingPass(vllm_config),
272 fusion_pass,
273 PostCleanupPass(vllm_config),
274 )
275
276 for model in (target_model, draft_model):
277 num_tokens = 5
278 qkv = torch.randn(num_tokens, model.q_size + 2 * model.kv_size)
279 positions = torch.arange(num_tokens, dtype=torch.long, device=qkv.device)
280 torch._dynamo.mark_dynamic(qkv, 0)
281 torch._dynamo.mark_dynamic(positions, 0)
282 torch.compile(model, backend=backend)(qkv, positions)
283
284 assert fusion_pass.matched_count == 1