@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 274 lines 9.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 from typing import ParamSpec
6
7 import torch
8 import torch._inductor.pattern_matcher as pm
9 from torch import fx
10 from torch._higher_order_ops.auto_functionalize import auto_functionalized
11 from torch._inductor.pattern_matcher import PatternMatcherPass
12
13 import vllm.ir.ops
14 from vllm.config import VllmConfig, get_layers_from_vllm_config
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.attention import Attention
17 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
18
19 from ..inductor_pass import enable_fake_mode
20 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21 from .matcher_utils import MatcherRotaryEmbedding
22 from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64
23
24 logger = init_logger(__name__)
25
26 FUSED_QK_ROPE_OP = torch.ops._C.fused_qk_norm_rope.default
27
28 # Head dimensions supported by csrc/fused_qknorm_rope_kernel.cu's
29 # launchFusedQKNormRope and launchFusedQKNormRopeNTokenHeads dispatchers.
30 # Keep in sync with the switch statements in that file.
31 SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS: tuple[int, ...] = (64, 128, 256)
32
33 P = ParamSpec("P")
34
35
36 class QkNormRopePattern:
37 """Match the unfused sequence in attention blocks and replace with the fused op.
38
39 Unfused (conceptually):
40 q, k, v = split(qkv, [qsz, kvsz, kvsz], -1)
41 qh = reshape(q, [-1, num_heads, head_dim])
42 kh = reshape(k, [-1, num_kv_heads, head_dim])
43 qn = rms_norm(qh, q_weight, eps)
44 kn = rms_norm(kh, k_weight, eps)
45 qf = reshape(qn, [-1, num_heads * head_dim])
46 kf = reshape(kn, [-1, num_kv_heads * head_dim])
47 qf, kf = rotary_embedding(positions, qf, kf, head_dim, cos_sin_cache, is_neox)
48 return qf, kf, v
49
50 Fused replacement:
51 fused_qk_norm_rope(qkv, num_heads, num_kv_heads, num_kv_heads, head_dim,
52 eps, q_weight, k_weight, cos_sin_cache, is_neox,
53 positions.view(-1))
54 return split(qkv, [qsz, kvsz, kvsz], -1)
55 """
56
57 def __init__(
58 self,
59 head_dim: int,
60 num_heads: int,
61 num_kv_heads: int,
62 eps: float,
63 is_neox: bool,
64 rope_flashinfer: bool = False,
65 ) -> None:
66 self.num_heads = num_heads
67 self.num_kv_heads = num_kv_heads
68 self.head_dim = head_dim
69 self.q_size = self.num_heads * self.head_dim
70 self.kv_size = self.num_kv_heads * self.head_dim
71 self.eps = eps
72 self.is_neox = is_neox
73 self.rope_flashinfer = rope_flashinfer
74 self.rope_matcher = MatcherRotaryEmbedding(
75 is_neox=is_neox,
76 head_size=self.head_dim,
77 num_heads=self.num_heads,
78 num_kv_heads=self.num_kv_heads,
79 use_flashinfer=self.rope_flashinfer,
80 )
81
82 def get_inputs(self) -> list[torch.Tensor]:
83 # Sample inputs to help pattern tracing
84 T = 5
85 qkv = empty_bf16(T, self.q_size + 2 * self.kv_size)
86 positions = empty_i64(T)
87 q_weight = empty_bf16(1, self.head_dim)
88 k_weight = empty_bf16(1, self.head_dim)
89 if self.rope_flashinfer:
90 cos_sin_cache = empty_fp32(4096, self.head_dim)
91 else:
92 cos_sin_cache = empty_bf16(4096, self.head_dim)
93 return [
94 qkv,
95 positions,
96 q_weight,
97 k_weight,
98 cos_sin_cache,
99 ]
100
101 @staticmethod
102 def wrap_trace_fn(
103 trace_fn: Callable[P, fx.GraphModule],
104 *process_fx_fns: Callable[[fx.GraphModule], None],
105 ) -> Callable[P, fx.GraphModule]:
106 def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule:
107 gm = trace_fn(*args, **kwargs)
108 for process_fx in process_fx_fns:
109 process_fx(gm)
110
111 return gm
112
113 return wrapped
114
115 @staticmethod
116 def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
117 from torch._inductor.fx_passes.post_grad import view_to_reshape
118
119 view_to_reshape(gm)
120
121 def register(self, pm_pass: PatternMatcherPass) -> None:
122 def pattern(
123 qkv: torch.Tensor,
124 positions: torch.Tensor,
125 q_weight: torch.Tensor,
126 k_weight: torch.Tensor,
127 cos_sin_cache: torch.Tensor,
128 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
129 # split qkv -> q,k,v
130 try:
131 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
132 except ValueError as e:
133 # register_replacement treats RuntimeError as a trace mismatch.
134 raise RuntimeError from e
135
136 # Q path: view -> RMS -> view back to q.shape
137 q_by_head = q.view(
138 *q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
139 )
140 q_normed_by_head = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)
141 q_flat = q_normed_by_head.view(q.shape)
142
143 # K path: view -> RMS -> view back to k.shape
144 k_by_head = k.view(
145 *k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
146 )
147 k_normed_by_head = vllm.ir.ops.rms_norm(k_by_head, k_weight, self.eps)
148 k_flat = k_normed_by_head.view(k.shape)
149
150 # RoPE: apply to flattened q/k
151 q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)
152 return q_rope, k_rope, v
153
154 def replacement(
155 qkv: torch.Tensor,
156 positions: torch.Tensor,
157 q_weight: torch.Tensor,
158 k_weight: torch.Tensor,
159 cos_sin_cache: torch.Tensor,
160 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
161 # Run fused qk_norm_rope op
162 result = auto_functionalized(
163 FUSED_QK_ROPE_OP,
164 qkv=qkv,
165 num_heads_q=self.num_heads,
166 num_heads_k=self.num_kv_heads,
167 num_heads_v=self.num_kv_heads,
168 head_dim=self.head_dim,
169 eps=self.eps,
170 q_weight=q_weight,
171 k_weight=k_weight,
172 cos_sin_cache=cos_sin_cache,
173 is_neox=self.is_neox,
174 position_ids=positions.view(-1),
175 forced_token_heads_per_warp=-1,
176 )
177 result_qkv = result[1]
178
179 # Split back to q,k,v and return
180 return result_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # type: ignore[no-any-return]
181
182 # NOTE: use fx_view_to_reshape to unify view/reshape to simplify
183 # pattern and increase matching opportunities
184 pm.register_replacement(
185 pattern,
186 replacement,
187 self.get_inputs(),
188 QkNormRopePattern.wrap_trace_fn(
189 pm.fwd_only,
190 QkNormRopePattern.fx_view_to_reshape,
191 ),
192 pm_pass,
193 )
194
195
196 class QKNormRoPEFusionPass(VllmPatternMatcherPass):
197 """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists.
198
199 Registers patterns for both standard vLLM ops and ROCm AITER ops
200 (when AITER is enabled), so the fusion fires regardless of which
201 RMSNorm/RoPE implementation the graph uses.
202 """
203
204 @enable_fake_mode
205 def __init__(self, config: VllmConfig) -> None:
206 super().__init__(config)
207 self.patterns: PatternMatcherPass = PatternMatcherPass(
208 pass_name="qk_norm_rope_fusion_pass"
209 )
210 self._attention_geometries: tuple[tuple[int, int, int], ...] = ()
211
212 dtype = config.model_config.dtype
213 if dtype not in (torch.bfloat16, torch.float16):
214 logger.warning_once(
215 "QK Norm+RoPE fusion not enabled: unsupported dtype %s", dtype
216 )
217 return
218
219 attn_layers: dict[str, Attention] = get_layers_from_vllm_config(
220 config, Attention
221 )
222 if len(attn_layers) == 0:
223 logger.warning_once(
224 "QK Norm+RoPE fusion enabled, but no Attention layers were discovered."
225 )
226 return
227
228 for layer in attn_layers.values():
229 if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS:
230 logger.warning_once(
231 "QK Norm+RoPE fusion not enabled: layer head_size=%d is not "
232 "supported by fused_qk_norm_rope kernel (supported: %s). "
233 "Falling back to unfused QK norm + RoPE path.",
234 layer.head_size,
235 SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS,
236 )
237 return
238
239 self._attention_geometries = tuple(
240 sorted(
241 {
242 (layer.head_size, layer.num_heads, layer.num_kv_heads)
243 for layer in attn_layers.values()
244 }
245 )
246 )
247
248 rope_flashinfer_options = (
249 [False, True] if RotaryEmbedding.enabled() else [False]
250 )
251 for head_dim, num_heads, num_kv_heads in self._attention_geometries:
252 for epsilon in [1e-5, 1e-6]:
253 for neox in [True, False]:
254 for rope_flashinfer in rope_flashinfer_options:
255 QkNormRopePattern(
256 head_dim=head_dim,
257 num_heads=num_heads,
258 num_kv_heads=num_kv_heads,
259 eps=epsilon,
260 is_neox=neox,
261 rope_flashinfer=rope_flashinfer,
262 ).register(self.patterns)
263
264 self.dump_patterns(config, self.patterns)
265
266 @VllmInductorPass.time_and_log
267 def __call__(self, graph: fx.Graph) -> None:
268 self.matched_count = self.patterns.apply(graph)
269 logger.debug("Fused QK Norm+RoPE on %s sites", self.matched_count)
270
271 def uuid(self) -> str:
272 return VllmInductorPass.hash_source(
273 self, QkNormRopePattern, repr(self._attention_geometries)
274 )