| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch._inductor.pattern_matcher as pm |
| 6 | from torch import fx |
| 7 | from torch._higher_order_ops import auto_functionalized |
| 8 | from torch._inductor.fx_passes.post_grad import view_to_reshape |
| 9 | from torch._inductor.pattern_matcher import PatternMatcherPass |
| 10 | |
| 11 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 12 | from vllm.config.utils import Range |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.attention.attention import ( |
| 15 | Attention, |
| 16 | get_attention_context, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.utils.torch_utils import ( |
| 20 | _USE_LAYERNAME, |
| 21 | LayerNameType, |
| 22 | _encode_layer_name, |
| 23 | _resolve_layer_name, |
| 24 | direct_register_custom_op, |
| 25 | ) |
| 26 | |
| 27 | from ..inductor_pass import enable_fake_mode |
| 28 | from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass |
| 29 | from .matcher_utils import ( |
| 30 | MatcherRotaryEmbedding, |
| 31 | ) |
| 32 | from .rms_quant_fusion import ( |
| 33 | empty_bf16, |
| 34 | empty_i64, |
| 35 | ) |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | FP8_DTYPE = current_platform.fp8_dtype() |
| 39 | STATIC_FP8_QUANT_OP = getattr(torch.ops._C, "static_scaled_fp8_quant", None) |
| 40 | |
| 41 | |
| 42 | def _supports_static_q_fp8_quant_fusion() -> bool: |
| 43 | return STATIC_FP8_QUANT_OP is not None and FP8_DTYPE is not None |
| 44 | |
| 45 | |
| 46 | def fused_rope_and_unified_kv_cache_update_impl( |
| 47 | query: torch.Tensor, |
| 48 | key: torch.Tensor, |
| 49 | value: torch.Tensor, |
| 50 | positions: torch.Tensor, |
| 51 | cos_sin_cache: torch.Tensor, |
| 52 | is_neox: bool, |
| 53 | layer_name: LayerNameType, |
| 54 | ) -> torch.Tensor: |
| 55 | """This impl fetches the KV cache and slot mapping from the forward context, |
| 56 | then calls the layer impl's `AttentionImpl.do_rope_and_kv_cache_update` method. |
| 57 | It also returns a dummy tensor, similar to `Attention.unified_kv_cache_update`, |
| 58 | that is passed to unified_attention to signal a side effect and |
| 59 | the data dependency between them to ensure torch.compile preserves ordering. |
| 60 | """ |
| 61 | layer_name = _resolve_layer_name(layer_name) |
| 62 | _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name) |
| 63 | if layer_slot_mapping is not None: |
| 64 | attn_layer.impl.do_rope_and_kv_cache_update( |
| 65 | attn_layer, |
| 66 | query, |
| 67 | key, |
| 68 | value, |
| 69 | positions, |
| 70 | cos_sin_cache, |
| 71 | is_neox, |
| 72 | kv_cache, |
| 73 | layer_slot_mapping, |
| 74 | ) |
| 75 | |
| 76 | return query.new_empty(0) |
| 77 | |
| 78 | |
| 79 | def fused_rope_and_unified_kv_cache_update_fake( |
| 80 | query: torch.Tensor, |
| 81 | key: torch.Tensor, |
| 82 | value: torch.Tensor, |
| 83 | positions: torch.Tensor, |
| 84 | cos_sin_cache: torch.Tensor, |
| 85 | is_neox: bool, |
| 86 | layer_name: LayerNameType, |
| 87 | ) -> torch.Tensor: |
| 88 | return torch.empty(0, device=query.device, dtype=query.dtype) |
| 89 | |
| 90 | |
| 91 | direct_register_custom_op( |
| 92 | op_name="fused_rope_and_unified_kv_cache_update", |
| 93 | op_func=fused_rope_and_unified_kv_cache_update_impl, |
| 94 | mutates_args=["query", "key"], |
| 95 | fake_impl=fused_rope_and_unified_kv_cache_update_fake, |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | class RopeStaticQQuantKVCachePattern: |
| 100 | """Fuse rope + static Q fp8 quant while preserving explicit KV-cache update |
| 101 | dependency ordering. |
| 102 | """ |
| 103 | |
| 104 | FUSED_ROPE_KV_OP = torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default |
| 105 | |
| 106 | def __init__( |
| 107 | self, |
| 108 | layer: Attention, |
| 109 | is_neox: bool, |
| 110 | ) -> None: |
| 111 | self.layer_name = layer.layer_name |
| 112 | self.num_heads = layer.num_heads |
| 113 | self.num_kv_heads = layer.num_kv_heads |
| 114 | self.head_size = layer.head_size |
| 115 | self.head_size_v = layer.head_size_v |
| 116 | self.is_neox = is_neox |
| 117 | |
| 118 | self.q_size = self.num_heads * self.head_size |
| 119 | self.k_size = self.num_kv_heads * self.head_size |
| 120 | self.v_size = self.num_kv_heads * self.head_size_v |
| 121 | |
| 122 | self.rope_matcher = MatcherRotaryEmbedding( |
| 123 | is_neox=self.is_neox, |
| 124 | head_size=self.head_size, |
| 125 | num_heads=self.num_heads, |
| 126 | num_kv_heads=self.num_kv_heads, |
| 127 | ) |
| 128 | |
| 129 | def get_inputs(self) -> list: |
| 130 | T = 5 |
| 131 | L = 4096 |
| 132 | qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size) |
| 133 | positions = empty_i64(T) |
| 134 | cos_sin_cache = empty_bf16(L, self.head_size) |
| 135 | q_scale = torch.empty((), dtype=torch.float32, device=qkv.device) |
| 136 | inputs: list = [qkv, positions, cos_sin_cache, q_scale] |
| 137 | if _USE_LAYERNAME: |
| 138 | inputs.append(_encode_layer_name(self.layer_name)) |
| 139 | return inputs |
| 140 | |
| 141 | def _mk_pattern_with_layer_name_input(self, _ln): |
| 142 | def pattern(qkv, positions, cos_sin_cache, q_scale, layer_name): |
| 143 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 144 | q, k = self.rope_matcher(positions, q, k, cos_sin_cache) |
| 145 | q_fp8 = torch.empty(q.shape, device=q.device, dtype=FP8_DTYPE) |
| 146 | _, q_fp8 = auto_functionalized( |
| 147 | torch.ops._C.static_scaled_fp8_quant.default, |
| 148 | result=q_fp8, |
| 149 | input=q, |
| 150 | scale=q_scale, |
| 151 | group_shape=(-1, -1), |
| 152 | ) |
| 153 | q_view = q_fp8.view(-1, self.num_heads, self.head_size) |
| 154 | k_view = k.view(-1, self.num_kv_heads, self.head_size) |
| 155 | v_view = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 156 | kv_cache_dummy = torch.ops.vllm.unified_kv_cache_update( |
| 157 | k_view, v_view, layer_name |
| 158 | ) |
| 159 | return kv_cache_dummy, q_view, k_view, v_view |
| 160 | |
| 161 | def replacement(qkv, positions, cos_sin_cache, q_scale, layer_name): |
| 162 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 163 | q_view = q.view(-1, self.num_heads, self.head_size) |
| 164 | k_view = k.view(-1, self.num_kv_heads, self.head_size) |
| 165 | v_view = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 166 | rope_kv_results = auto_functionalized( |
| 167 | self.FUSED_ROPE_KV_OP, |
| 168 | query=q_view, |
| 169 | key=k_view, |
| 170 | value=v_view, |
| 171 | positions=positions, |
| 172 | cos_sin_cache=cos_sin_cache, |
| 173 | is_neox=self.is_neox, |
| 174 | layer_name=layer_name, |
| 175 | ) |
| 176 | kv_cache_dummy = rope_kv_results[0] |
| 177 | q_after = rope_kv_results[1] |
| 178 | k_after = rope_kv_results[2] |
| 179 | q_after_flat = q_after.view(-1, self.q_size) |
| 180 | q_fp8 = torch.empty( |
| 181 | q_after_flat.shape, device=q_after_flat.device, dtype=FP8_DTYPE |
| 182 | ) |
| 183 | _, q_fp8 = auto_functionalized( |
| 184 | torch.ops._C.static_scaled_fp8_quant.default, |
| 185 | result=q_fp8, |
| 186 | input=q_after_flat, |
| 187 | scale=q_scale, |
| 188 | group_shape=(-1, -1), |
| 189 | ) |
| 190 | return ( |
| 191 | kv_cache_dummy, |
| 192 | q_fp8.view(-1, self.num_heads, self.head_size), |
| 193 | k_after, |
| 194 | v_view, |
| 195 | ) |
| 196 | |
| 197 | return pattern, replacement |
| 198 | |
| 199 | def _mk_pattern_with_layer_name_closure(self, _ln): |
| 200 | def pattern(qkv, positions, cos_sin_cache, q_scale): |
| 201 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 202 | q, k = self.rope_matcher(positions, q, k, cos_sin_cache) |
| 203 | q_fp8 = torch.empty(q.shape, device=q.device, dtype=FP8_DTYPE) |
| 204 | _, q_fp8 = auto_functionalized( |
| 205 | torch.ops._C.static_scaled_fp8_quant.default, |
| 206 | result=q_fp8, |
| 207 | input=q, |
| 208 | scale=q_scale, |
| 209 | group_shape=(-1, -1), |
| 210 | ) |
| 211 | q_view = q_fp8.view(-1, self.num_heads, self.head_size) |
| 212 | k_view = k.view(-1, self.num_kv_heads, self.head_size) |
| 213 | v_view = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 214 | kv_cache_dummy = torch.ops.vllm.unified_kv_cache_update(k_view, v_view, _ln) |
| 215 | return kv_cache_dummy, q_view, k_view, v_view |
| 216 | |
| 217 | def replacement(qkv, positions, cos_sin_cache, q_scale): |
| 218 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 219 | q_view = q.view(-1, self.num_heads, self.head_size) |
| 220 | k_view = k.view(-1, self.num_kv_heads, self.head_size) |
| 221 | v_view = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 222 | rope_kv_results = auto_functionalized( |
| 223 | self.FUSED_ROPE_KV_OP, |
| 224 | query=q_view, |
| 225 | key=k_view, |
| 226 | value=v_view, |
| 227 | positions=positions, |
| 228 | cos_sin_cache=cos_sin_cache, |
| 229 | is_neox=self.is_neox, |
| 230 | layer_name=_ln, |
| 231 | ) |
| 232 | kv_cache_dummy = rope_kv_results[0] |
| 233 | q_after = rope_kv_results[1] |
| 234 | k_after = rope_kv_results[2] |
| 235 | q_after_flat = q_after.view(-1, self.q_size) |
| 236 | q_fp8 = torch.empty( |
| 237 | q_after_flat.shape, device=q_after_flat.device, dtype=FP8_DTYPE |
| 238 | ) |
| 239 | _, q_fp8 = auto_functionalized( |
| 240 | torch.ops._C.static_scaled_fp8_quant.default, |
| 241 | result=q_fp8, |
| 242 | input=q_after_flat, |
| 243 | scale=q_scale, |
| 244 | group_shape=(-1, -1), |
| 245 | ) |
| 246 | return ( |
| 247 | kv_cache_dummy, |
| 248 | q_fp8.view(-1, self.num_heads, self.head_size), |
| 249 | k_after, |
| 250 | v_view, |
| 251 | ) |
| 252 | |
| 253 | return pattern, replacement |
| 254 | |
| 255 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 256 | _ln = _encode_layer_name(self.layer_name) |
| 257 | |
| 258 | if _USE_LAYERNAME: |
| 259 | pattern, replacement = self._mk_pattern_with_layer_name_input(_ln) |
| 260 | else: |
| 261 | pattern, replacement = self._mk_pattern_with_layer_name_closure(_ln) |
| 262 | |
| 263 | def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule: |
| 264 | gm = pm.fwd_only(*args, **kwargs) |
| 265 | view_to_reshape(gm) |
| 266 | return gm |
| 267 | |
| 268 | pm.register_replacement( |
| 269 | pattern, |
| 270 | replacement, |
| 271 | self.get_inputs(), |
| 272 | fwd_and_view_to_reshape, |
| 273 | pm_pass, |
| 274 | ) |
| 275 | |
| 276 | |
| 277 | class RopeReshapeKVCachePattern: |
| 278 | """This pattern matches the following unfused inplace ops: |
| 279 | q, k = rotary_embedding(positions, q, k, head_size, cos_sin_cache, is_neox) |
| 280 | kv_cache_dummy = unified_kv_cache_update(k, v, layer_name) |
| 281 | |
| 282 | and replaces it with the fused inplace op: |
| 283 | kv_cache_dummy = fused_rope_and_unified_kv_cache_update( |
| 284 | q, k, v, positions, cos_sin_cache, is_neox, layer_name |
| 285 | ) |
| 286 | """ |
| 287 | |
| 288 | FUSED_OP = torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default |
| 289 | |
| 290 | def __init__( |
| 291 | self, |
| 292 | layer: Attention, |
| 293 | is_neox: bool, |
| 294 | ) -> None: |
| 295 | self.layer_name = layer.layer_name |
| 296 | self.num_heads = layer.num_heads |
| 297 | self.num_kv_heads = layer.num_kv_heads |
| 298 | self.head_size = layer.head_size |
| 299 | self.head_size_v = layer.head_size_v |
| 300 | self.is_neox = is_neox |
| 301 | |
| 302 | self.q_size = self.num_heads * self.head_size |
| 303 | self.k_size = self.num_kv_heads * self.head_size |
| 304 | self.v_size = self.num_kv_heads * self.head_size_v |
| 305 | |
| 306 | self.rope_matcher = MatcherRotaryEmbedding( |
| 307 | is_neox=self.is_neox, |
| 308 | head_size=self.head_size, |
| 309 | num_heads=self.num_heads, |
| 310 | num_kv_heads=self.num_kv_heads, |
| 311 | ) |
| 312 | |
| 313 | def get_inputs(self) -> list: |
| 314 | # Sample inputs to help pattern tracing |
| 315 | T = 5 |
| 316 | L = 4096 |
| 317 | qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size) |
| 318 | positions = empty_i64(T) |
| 319 | cos_sin_cache = empty_bf16(L, self.head_size) |
| 320 | inputs: list = [qkv, positions, cos_sin_cache] |
| 321 | if _USE_LAYERNAME: |
| 322 | inputs.append(_encode_layer_name(self.layer_name)) |
| 323 | return inputs |
| 324 | |
| 325 | def _mk_pattern_with_layer_name_input(self, _ln): |
| 326 | """Pattern/replacement with layer_name as an explicit input.""" |
| 327 | |
| 328 | def pattern(qkv, positions, cos_sin_cache, layer_name): |
| 329 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 330 | q, k = self.rope_matcher(positions, q, k, cos_sin_cache) |
| 331 | q = q.view(-1, self.num_heads, self.head_size) |
| 332 | k = k.view(-1, self.num_kv_heads, self.head_size) |
| 333 | v = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 334 | return torch.ops.vllm.unified_kv_cache_update(k, v, layer_name), q, k, v |
| 335 | |
| 336 | def replacement(qkv, positions, cos_sin_cache, layer_name): |
| 337 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 338 | q = q.view(-1, self.num_heads, self.head_size) |
| 339 | k = k.view(-1, self.num_kv_heads, self.head_size) |
| 340 | v = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 341 | results = auto_functionalized( |
| 342 | self.FUSED_OP, |
| 343 | query=q, |
| 344 | key=k, |
| 345 | value=v, |
| 346 | positions=positions, |
| 347 | cos_sin_cache=cos_sin_cache, |
| 348 | is_neox=self.is_neox, |
| 349 | layer_name=layer_name, |
| 350 | ) |
| 351 | return results[0], results[1], results[2], v |
| 352 | |
| 353 | return pattern, replacement |
| 354 | |
| 355 | def _mk_pattern_with_layer_name_closure(self, _ln): |
| 356 | """Pattern/replacement with layer_name as a closure constant.""" |
| 357 | |
| 358 | def pattern(qkv, positions, cos_sin_cache): |
| 359 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 360 | q, k = self.rope_matcher(positions, q, k, cos_sin_cache) |
| 361 | q = q.view(-1, self.num_heads, self.head_size) |
| 362 | k = k.view(-1, self.num_kv_heads, self.head_size) |
| 363 | v = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 364 | return torch.ops.vllm.unified_kv_cache_update(k, v, _ln), q, k, v |
| 365 | |
| 366 | def replacement(qkv, positions, cos_sin_cache): |
| 367 | q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) |
| 368 | q = q.view(-1, self.num_heads, self.head_size) |
| 369 | k = k.view(-1, self.num_kv_heads, self.head_size) |
| 370 | v = v.view(-1, self.num_kv_heads, self.head_size_v) |
| 371 | results = auto_functionalized( |
| 372 | self.FUSED_OP, |
| 373 | query=q, |
| 374 | key=k, |
| 375 | value=v, |
| 376 | positions=positions, |
| 377 | cos_sin_cache=cos_sin_cache, |
| 378 | is_neox=self.is_neox, |
| 379 | layer_name=_ln, |
| 380 | ) |
| 381 | return results[0], results[1], results[2], v |
| 382 | |
| 383 | return pattern, replacement |
| 384 | |
| 385 | def register(self, pm_pass: PatternMatcherPass) -> None: |
| 386 | _ln = _encode_layer_name(self.layer_name) |
| 387 | |
| 388 | if _USE_LAYERNAME: |
| 389 | pattern, replacement = self._mk_pattern_with_layer_name_input(_ln) |
| 390 | else: |
| 391 | pattern, replacement = self._mk_pattern_with_layer_name_closure(_ln) |
| 392 | |
| 393 | # NOTE: use view_to_reshape to unify view/reshape to simplify |
| 394 | # pattern and increase matching opportunities |
| 395 | def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule: |
| 396 | gm = pm.fwd_only(*args, **kwargs) |
| 397 | view_to_reshape(gm) |
| 398 | return gm |
| 399 | |
| 400 | pm.register_replacement( |
| 401 | pattern, |
| 402 | replacement, |
| 403 | self.get_inputs(), |
| 404 | fwd_and_view_to_reshape, |
| 405 | pm_pass, |
| 406 | ) |
| 407 | |
| 408 | |
| 409 | class RopeKVCacheFusionPass(VllmPatternMatcherPass): |
| 410 | """This pass fuses the rotary embedding and KV cache update operations |
| 411 | into a single fused kernel if available. |
| 412 | |
| 413 | It uses the pattern matcher and matches each layer manually, as strings |
| 414 | cannot be wildcarded. This also lets us check support on attention layers |
| 415 | upon registration instead of during pattern matching. |
| 416 | |
| 417 | This fusion eliminates the need for separate kernel launches and |
| 418 | intermediate memory operations between the RoPE and cache update steps. |
| 419 | """ |
| 420 | |
| 421 | @enable_fake_mode |
| 422 | def __init__(self, config: VllmConfig) -> None: |
| 423 | super().__init__(config) |
| 424 | |
| 425 | self.patterns: PatternMatcherPass = PatternMatcherPass( |
| 426 | pass_name="rope_kv_cache_fusion_pass" |
| 427 | ) |
| 428 | |
| 429 | cc = config.compilation_config |
| 430 | self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num |
| 431 | |
| 432 | attn_layers = get_layers_from_vllm_config(config, Attention) |
| 433 | # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all |
| 434 | # layers produce the same pattern — register once then break. |
| 435 | for _, layer in attn_layers.items(): |
| 436 | if layer.impl.fused_rope_kvcache_supported(): |
| 437 | for is_neox in [True, False]: |
| 438 | if _supports_static_q_fp8_quant_fusion(): |
| 439 | RopeStaticQQuantKVCachePattern( |
| 440 | layer=layer, |
| 441 | is_neox=is_neox, |
| 442 | ).register(self.patterns) |
| 443 | RopeReshapeKVCachePattern( |
| 444 | layer=layer, |
| 445 | is_neox=is_neox, |
| 446 | ).register(self.patterns) |
| 447 | if _USE_LAYERNAME: |
| 448 | break |
| 449 | |
| 450 | self.dump_patterns(config, self.patterns) |
| 451 | |
| 452 | @VllmInductorPass.time_and_log |
| 453 | def __call__(self, graph: fx.Graph) -> None: |
| 454 | self.matched_count = self.patterns.apply(graph) |
| 455 | logger.debug("Replaced %s patterns", self.matched_count) |
| 456 | |
| 457 | def is_applicable_for_range(self, compile_range: Range) -> bool: |
| 458 | # This pass works best for the small-batch decode setting. |
| 459 | # For large-batch e.g. prefill, it is better to use two separate kernels |
| 460 | # since they are compute bound and the fused kernels require further tuning. |
| 461 | return compile_range.end <= self.max_token_num |
| 462 | |
| 463 | def uuid(self) -> str: |
| 464 | return VllmInductorPass.hash_source( |
| 465 | self, RopeStaticQQuantKVCachePattern, RopeReshapeKVCachePattern |
| 466 | ) |