@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 274 lines 9.23 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4 from torch._higher_order_ops.auto_functionalize import auto_functionalized
5
6 import vllm._custom_ops as ops
7 from vllm.config import VllmConfig, get_layers_from_vllm_config
8 from vllm.logger import init_logger
9 from vllm.model_executor.layers.attention import MLAAttention
10 from vllm.model_executor.layers.attention.attention import get_attention_context
11 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
12 from vllm.utils.torch_utils import (
13 _USE_LAYERNAME,
14 LayerNameType,
15 _encode_layer_name,
16 _resolve_layer_name,
17 direct_register_custom_op,
18 )
19
20 from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
21 from .matcher_utils import MatcherDeepseekScalingRotaryEmbedding, MatcherRotaryEmbedding
22
23 logger = init_logger(__name__)
24
25
26 def fused_rope_unified_mla_kv_cache_update_impl(
27 positions: torch.Tensor,
28 q_pe: torch.Tensor,
29 k_pe: torch.Tensor,
30 kv_c: torch.Tensor,
31 cos_sin_cache: torch.Tensor,
32 is_neox: bool,
33 kv_cache_dtype: str,
34 kv_cache_scale: torch.Tensor,
35 layer_name: LayerNameType,
36 ) -> torch.Tensor:
37 layer_name = _resolve_layer_name(layer_name)
38 attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
39 layer_name
40 )
41 attn_layer.prepare_kv_cache_update(attn_metadata)
42 if layer_slot_mapping is not None:
43 ops.concat_and_cache_mla_rope_fused(
44 positions,
45 q_pe,
46 k_pe,
47 kv_c,
48 cos_sin_cache,
49 is_neox,
50 layer_slot_mapping,
51 kv_cache.squeeze(1),
52 kv_cache_dtype,
53 kv_cache_scale,
54 )
55 return torch.empty(0, device=kv_c.device, dtype=kv_c.dtype)
56
57
58 def fused_rope_unified_mla_kv_cache_update_fake(
59 positions: torch.Tensor,
60 q_pe: torch.Tensor,
61 k_pe: torch.Tensor,
62 kv_c: torch.Tensor,
63 cos_sin_cache: torch.Tensor,
64 is_neox: bool,
65 kv_cache_dtype: str,
66 kv_cache_scale: torch.Tensor,
67 layer_name: LayerNameType,
68 ) -> torch.Tensor:
69 return torch.empty(0, dtype=kv_c.dtype, device=kv_c.device)
70
71
72 direct_register_custom_op(
73 op_name="fused_rope_unified_mla_kv_cache_update",
74 op_func=fused_rope_unified_mla_kv_cache_update_impl,
75 fake_impl=fused_rope_unified_mla_kv_cache_update_fake,
76 mutates_args=["q_pe", "k_pe"],
77 )
78
79
80 class MLARoPEKVCacheCatPattern(VllmPatternReplacement):
81 FUSED_OP = torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default
82
83 def __init__(
84 self,
85 layer: MLAAttention,
86 is_neox: bool,
87 use_flashinfer: bool = False,
88 use_deepseek_scaling: bool = False,
89 ) -> None:
90 self.layer_name = layer.layer_name
91 self.kv_cache_dtype = layer.kv_cache_dtype
92 self.num_heads = layer.num_heads
93 self.num_kv_heads = layer.num_kv_heads
94 self.kv_lora_rank = layer.kv_lora_rank
95 self.qk_rope_head_dim = layer.qk_rope_head_dim
96 self.is_neox = is_neox
97 self.use_flashinfer = use_flashinfer
98 self._ln = _encode_layer_name(self.layer_name)
99
100 if use_deepseek_scaling:
101 self.rope_matcher = MatcherDeepseekScalingRotaryEmbedding(
102 is_neox=self.is_neox,
103 head_size=self.qk_rope_head_dim,
104 num_heads=self.num_heads,
105 num_kv_heads=self.num_kv_heads,
106 use_flashinfer=self.use_flashinfer,
107 )
108 else:
109 self.rope_matcher = MatcherRotaryEmbedding( # type: ignore
110 is_neox=self.is_neox,
111 head_size=self.qk_rope_head_dim,
112 num_heads=self.num_heads,
113 num_kv_heads=self.num_kv_heads,
114 use_flashinfer=self.use_flashinfer,
115 )
116
117 def get_inputs(self) -> list[torch.Tensor]:
118 T = 5
119 L = 4096
120 q_pe = self.empty_bf16(T, self.num_heads, self.qk_rope_head_dim)
121 k_pe = self.empty_bf16(T, self.qk_rope_head_dim)
122 kv_c_normed = self.empty_bf16(T, self.kv_lora_rank)
123 cos_sin_cache = self.empty_bf16(L, self.qk_rope_head_dim)
124 positions = self.empty(T, dtype=torch.int64)
125 k_scale = self.empty(0, dtype=torch.float32)
126 inputs = [
127 q_pe,
128 k_pe,
129 kv_c_normed,
130 positions,
131 cos_sin_cache,
132 k_scale,
133 ]
134 if _USE_LAYERNAME:
135 inputs.append(self._ln)
136 return inputs
137
138 @property
139 def pattern(self):
140 _ln = self._ln
141
142 if _USE_LAYERNAME:
143
144 def _pattern_with_ln(
145 q_pe: torch.Tensor,
146 k_pe: torch.Tensor,
147 kv_c_normed: torch.Tensor,
148 positions: torch.Tensor,
149 cos_sin_cache: torch.Tensor,
150 k_scale: torch.Tensor,
151 layer_name: LayerNameType,
152 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
153 k_pe_unsqueezed = k_pe.unsqueeze(1)
154 q_pe, k_pe = self.rope_matcher(
155 positions, q_pe, k_pe_unsqueezed, cos_sin_cache
156 )
157 dummy = torch.ops.vllm.unified_mla_kv_cache_update(
158 kv_c_normed, k_pe, layer_name, self.kv_cache_dtype, k_scale
159 )
160 return dummy, q_pe, k_pe
161
162 return _pattern_with_ln
163
164 def _pattern(
165 q_pe: torch.Tensor,
166 k_pe: torch.Tensor,
167 kv_c_normed: torch.Tensor,
168 positions: torch.Tensor,
169 cos_sin_cache: torch.Tensor,
170 k_scale: torch.Tensor,
171 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
172 k_pe_unsqueezed = k_pe.unsqueeze(1)
173 q_pe, k_pe = self.rope_matcher(
174 positions, q_pe, k_pe_unsqueezed, cos_sin_cache
175 )
176 dummy = torch.ops.vllm.unified_mla_kv_cache_update(
177 kv_c_normed, k_pe, _ln, self.kv_cache_dtype, k_scale
178 )
179 return dummy, q_pe, k_pe
180
181 return _pattern
182
183 @property
184 def replacement(self):
185 _ln = self._ln
186
187 if _USE_LAYERNAME:
188
189 def _replacement_with_ln(
190 q_pe: torch.Tensor,
191 k_pe: torch.Tensor,
192 kv_c_normed: torch.Tensor,
193 positions: torch.Tensor,
194 cos_sin_cache: torch.Tensor,
195 k_scale: torch.Tensor,
196 layer_name: LayerNameType,
197 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
198 at = auto_functionalized(
199 self.FUSED_OP,
200 positions=positions,
201 q_pe=q_pe,
202 k_pe=k_pe,
203 kv_c=kv_c_normed,
204 cos_sin_cache=cos_sin_cache,
205 is_neox=self.is_neox,
206 kv_cache_dtype=self.kv_cache_dtype,
207 kv_cache_scale=k_scale,
208 layer_name=layer_name,
209 )
210 dummy, q_pe, k_pe_squeezed = at
211 k_pe = k_pe_squeezed.unsqueeze(1)
212 return dummy, q_pe, k_pe
213
214 return _replacement_with_ln
215
216 def _replacement(
217 q_pe: torch.Tensor,
218 k_pe: torch.Tensor,
219 kv_c_normed: torch.Tensor,
220 positions: torch.Tensor,
221 cos_sin_cache: torch.Tensor,
222 k_scale: torch.Tensor,
223 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
224 at = auto_functionalized(
225 self.FUSED_OP,
226 positions=positions,
227 q_pe=q_pe,
228 k_pe=k_pe,
229 kv_c=kv_c_normed,
230 cos_sin_cache=cos_sin_cache,
231 is_neox=self.is_neox,
232 kv_cache_dtype=self.kv_cache_dtype,
233 kv_cache_scale=k_scale,
234 layer_name=_ln,
235 )
236 dummy, q_pe, k_pe_squeezed = at
237 k_pe = k_pe_squeezed.unsqueeze(1)
238 return dummy, q_pe, k_pe
239
240 return _replacement
241
242
243 class MLARoPEKVCacheCatFusionPass(VllmFusionPatternMatcherPass):
244 def __init__(self, config: VllmConfig) -> None:
245 super().__init__(config, "mla_rope_kv_cache_fusion_pass")
246
247 attn_layers = get_layers_from_vllm_config(config, MLAAttention)
248
249 for _, layer in attn_layers.items():
250 for is_neox in [False, True]:
251 for use_deepseek_scaling in [False, True]:
252 if RotaryEmbedding.enabled():
253 for use_flashinfer in [False, True]:
254 self.register(
255 MLARoPEKVCacheCatPattern(
256 layer,
257 is_neox,
258 use_flashinfer,
259 use_deepseek_scaling,
260 )
261 )
262 else:
263 self.register(
264 MLARoPEKVCacheCatPattern(
265 layer,
266 is_neox,
267 use_deepseek_scaling=use_deepseek_scaling,
268 )
269 )
270
271 if _USE_LAYERNAME:
272 break
273
274 self.dump_patterns(config, self.pm_pass)