@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a51376b3f0
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 22:45 UTC

[Performance][DSR1]: Fused RoPE+KVCache+q_concat for MLA (#40392)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Co-authored-by: ElizaWszola <ewszola@redhat.com>

Rohan Potdar committed May 11, 2026 at 07:10 UTC a51376b3f05a2f74eac6ceeed7e52598b871a0fb
12 files changed +966 -109
csrc/cache_kernels_fused.cu
+75 -60
@@ -21,28 +21,33 @@ namespace vllm {
21
22 // NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's
23 // using u16 as the backing type.
24 -template <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
25 - typename cache_t, Fp8KVCacheDataType kv_dt>
24 +template <typename qk_t, typename cos_sin_t, bool IS_NEOX,
25 + typename raw_kv_scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
26 __global__ void concat_and_cache_mla_rope_fused_kernel(
27 const int64_t* __restrict__ positions, // [num_tokens]
28 qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim]
29 qk_t* __restrict__ k_pe, // [num_tokens, rot_dim]
30 const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
31 - const qk_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
32 - // rot_dim // 2]
31 + const cos_sin_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
32 + // rot_dim // 2]
33 const int rot_dim, const int64_t q_pe_stride_token,
34 const int64_t q_pe_stride_head, const int64_t k_pe_stride,
35 const int64_t kv_c_stride, const int num_q_heads,
36 cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank +
37 // rot_dim)]
38 - const int64_t* __restrict__ kv_cache_slot_mapping, // [num_tokens]
38 + const int64_t* __restrict__ slot_mapping, // [num_tokens]
39 const int block_stride, const int entry_stride, const int kv_lora_rank,
40 const int block_size, const float* kv_cache_quant_scale) {
41 // Each thread block is responsible for one token.
42 const int64_t token_idx = blockIdx.x;
43 + const int64_t slot_idx = slot_mapping[token_idx];
44 + // NOTE: slot_idx can be -1 if the token is padded
45 + if (slot_idx < 0) {
46 + return;
47 + }
48 const int64_t pos = positions[token_idx];
49
45 - const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
50 + const cos_sin_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
51
52 const int embed_dim = rot_dim / 2;
53
@@ -54,8 +59,8 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
59
60 // NOTE: Would be nice to have interleaved sin/cos so we could just load
61 // both at the same time.
57 - qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
58 - qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
62 + qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
63 + qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
64
65 qk_t* q_pe_head_ptr =
66 q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head;
@@ -81,21 +86,15 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
86 q_pe_head_ptr[pair_idx_y] = y_dst;
87 }
88
84 - const int64_t slot_idx = kv_cache_slot_mapping[token_idx];
89 const int64_t block_idx = slot_idx / block_size;
90 const int64_t entry_idx = slot_idx % block_size;
91
88 - // NOTE: slot_idx can be -1 if the token is padded
89 - if (slot_idx < 0) {
90 - return;
91 - }
92 -
92 // K with 1 HEAD
93 for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) {
94 int pair_idx = i;
95
97 - qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
98 - qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
96 + qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
97 + qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
98
99 qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;
100
@@ -165,36 +164,43 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
164
165 } // namespace vllm
166
168 -#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
169 - do { \
170 - VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
171 - using qk_t = scalar_t; \
172 - if (rope_is_neox) { \
173 - vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, true, RAW_KV_T, \
174 - CACHE_T, KV_DTYPE> \
175 - <<<grid, block, 0, stream>>>( \
176 - positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
177 - k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
178 - rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
179 - q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
180 - num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
181 - kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
182 - entry_stride, kv_lora_rank, block_size, \
183 - kv_cache_quant_scale.data_ptr<float>()); \
184 - } else { \
185 - vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T, \
186 - CACHE_T, KV_DTYPE> \
187 - <<<grid, block, 0, stream>>>( \
188 - positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
189 - k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
190 - rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
191 - q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
192 - num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
193 - kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
194 - entry_stride, kv_lora_rank, block_size, \
195 - kv_cache_quant_scale.data_ptr<float>()); \
196 - } \
197 - }); \
167 +#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
168 + do { \
169 + VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
170 + using qk_t = scalar_t; \
171 + VLLM_DISPATCH_FLOATING_TYPES( \
172 + rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
173 + [&] { \
174 + using cos_sin_t = scalar_t; \
175 + if (rope_is_neox) { \
176 + vllm::concat_and_cache_mla_rope_fused_kernel< \
177 + qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
178 + <<<grid, block, 0, stream>>>( \
179 + positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
180 + k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
181 + rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
182 + q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
183 + kv_c_stride, num_q_heads, \
184 + reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
185 + slot_mapping.data_ptr<int64_t>(), block_stride, \
186 + entry_stride, kv_lora_rank, block_size, \
187 + kv_cache_quant_scale.data_ptr<float>()); \
188 + } else { \
189 + vllm::concat_and_cache_mla_rope_fused_kernel< \
190 + qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
191 + <<<grid, block, 0, stream>>>( \
192 + positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
193 + k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
194 + rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
195 + q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
196 + kv_c_stride, num_q_heads, \
197 + reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
198 + slot_mapping.data_ptr<int64_t>(), block_stride, \
199 + entry_stride, kv_lora_rank, block_size, \
200 + kv_cache_quant_scale.data_ptr<float>()); \
201 + } \
202 + }); \
203 + }); \
204 } while (false)
205
206 // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -208,43 +214,52 @@ void concat_and_cache_mla_rope_fused(
214 torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
215 torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
216 bool rope_is_neox,
211 - torch::Tensor&
212 - kv_cache_slot_mapping, // [num_tokens] or [num_actual_tokens]
217 + torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
218 torch::Tensor&
219 kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
220 const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
216 - const int64_t num_tokens = q_pe.size(0);
221 + // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
222 + // slot_mapping.size(0) because of padding for CUDA graphs.
223 + // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
224 + // both include padding.
225 + // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
226 + // since key includes padding for CUDA graphs, while slot_mapping does not.
227 + // In this case, slot_mapping.size(0) represents the actual number of tokens
228 + // before padding.
229 + // For compatibility with both cases, we use slot_mapping.size(0) as the
230 + // number of tokens.
231 + int num_tokens = slot_mapping.size(0);
232 + int num_padded_tokens = q_pe.size(0);
233 + TORCH_CHECK_GE(num_padded_tokens, num_tokens);
234
235 const int num_q_heads = q_pe.size(1);
236 const int rot_dim = q_pe.size(2);
237 const int kv_lora_rank = kv_c.size(1);
238
222 - TORCH_CHECK(positions.size(0) >=
223 - num_tokens); // CUDA Graphs might pad this for us
239 + TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
240 TORCH_CHECK_EQ(positions.dim(), 1);
241 TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
242
227 - TORCH_CHECK_EQ(q_pe.size(0), num_tokens);
243 + TORCH_CHECK_EQ(q_pe.dim(), 3);
244 + TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
245 TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
246 TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
230 - TORCH_CHECK_EQ(q_pe.dim(), 3);
247
232 - TORCH_CHECK_EQ(k_pe.size(0), num_tokens);
233 - TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
248 TORCH_CHECK_EQ(k_pe.dim(), 2);
249 + TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
250 + TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
251 TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
252
237 - TORCH_CHECK_EQ(kv_c.size(0), num_tokens);
238 - TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
253 TORCH_CHECK_EQ(kv_c.dim(), 2);
254 + TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
255 + TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
256 TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
257 TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
258
259 TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
244 - TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type());
260
246 - TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens);
247 - TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long);
261 + TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
262 + TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
263
264 TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
265 TORCH_CHECK_EQ(kv_cache.dim(), 3);
tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py new
+413
@@ -0,0 +1,413 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +import vllm.config
8 +from tests.compile.backend import TestBackend
9 +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
10 +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
11 +from vllm.compilation.passes.fusion.mla_rope_kvcache_cat_fusion import (
12 + MLARoPEKVCacheCatFusionPass,
13 +)
14 +from vllm.compilation.passes.utility.fix_functionalization import (
15 + FixFunctionalizationPass,
16 +)
17 +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
18 +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
19 +from vllm.config import (
20 + CacheConfig,
21 + CompilationConfig,
22 + CompilationMode,
23 + ModelConfig,
24 + PassConfig,
25 + VllmConfig,
26 +)
27 +from vllm.forward_context import get_forward_context, set_forward_context
28 +from vllm.model_executor.layers.attention import MLAAttention
29 +from vllm.model_executor.layers.linear import ColumnParallelLinear
30 +from vllm.model_executor.layers.rotary_embedding import (
31 + DeepseekScalingRotaryEmbedding,
32 + RotaryEmbedding,
33 +)
34 +from vllm.platforms import current_platform
35 +from vllm.utils.torch_utils import _encode_layer_name
36 +from vllm.v1.attention.backend import (
37 + AttentionBackend,
38 + CommonAttentionMetadata,
39 +)
40 +from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
41 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
42 +
43 +INDEX_SELECT_OP = torch.ops.aten.index.Tensor
44 +VLLM_UNIFIED_MLA_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_mla_kv_cache_update
45 +FP8_DTYPE = current_platform.fp8_dtype()
46 +
47 +
48 +class MLARoPEKVCacheCatTestModel(torch.nn.Module):
49 + def __init__(
50 + self,
51 + vllm_config: VllmConfig,
52 + attn_backend: AttentionBackendEnum,
53 + use_deepseek_scaling_rope: bool,
54 + num_heads: int,
55 + qk_nope_head_dim: int,
56 + qk_rope_head_dim: int,
57 + v_head_dim: int,
58 + q_lora_rank: int,
59 + kv_lora_rank: int,
60 + is_neox: bool,
61 + dtype: torch.dtype,
62 + device: torch.device,
63 + prefix: str = "model.layers.0.self_attn.attn",
64 + ):
65 + super().__init__()
66 + self.num_heads = num_heads
67 + self.qk_nope_head_dim = qk_nope_head_dim
68 + self.qk_rope_head_dim = qk_rope_head_dim
69 + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
70 + self.v_head_dim = v_head_dim
71 + self.q_lora_rank = q_lora_rank
72 + self.kv_lora_rank = kv_lora_rank
73 + self.dtype = dtype
74 + self.device = device
75 + self.layer_name = prefix
76 +
77 + self.num_kv_heads = 1
78 + self.head_size = kv_lora_rank + qk_rope_head_dim
79 + self.block_size = vllm_config.cache_config.block_size
80 + self.scale = self.qk_head_dim**-0.5
81 +
82 + if use_deepseek_scaling_rope:
83 + self.rotary_emb = DeepseekScalingRotaryEmbedding(
84 + head_size=qk_rope_head_dim,
85 + rotary_dim=qk_rope_head_dim,
86 + max_position_embeddings=4096,
87 + base=10000,
88 + is_neox_style=is_neox,
89 + scaling_factor=1.0,
90 + dtype=dtype,
91 + )
92 + else:
93 + self.rotary_emb = RotaryEmbedding(
94 + head_size=qk_rope_head_dim,
95 + rotary_dim=qk_rope_head_dim,
96 + max_position_embeddings=4096,
97 + base=10000,
98 + is_neox_style=is_neox,
99 + dtype=dtype,
100 + )
101 +
102 + # Initialize intermediate mm layers for unit test
103 + self.q_b_proj = ColumnParallelLinear(
104 + self.q_lora_rank,
105 + self.num_heads * self.qk_head_dim,
106 + bias=False,
107 + prefix=f"{prefix}.q_b_proj",
108 + ).to(device)
109 + self.kv_b_proj = ColumnParallelLinear(
110 + self.kv_lora_rank,
111 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
112 + bias=False,
113 + prefix=f"{prefix}.kv_b_proj",
114 + ).to(device)
115 +
116 + # ColumnParallelLinear default init in bf16 with seed 0 produces
117 + # near-zero weights (7/4.7M nonzero), making the GEMM output almost
118 + # entirely zero and masking correctness bugs. Reinitialize to get
119 + # dense outputs.
120 + with torch.no_grad():
121 + torch.nn.init.normal_(self.q_b_proj.weight, std=0.02)
122 + torch.nn.init.normal_(self.kv_b_proj.weight, std=0.02)
123 +
124 + # Register layer metadata for the fusion pass via MLAAttention
125 + self.mla_attn = MLAAttention(
126 + num_heads=self.num_heads,
127 + scale=self.scale,
128 + qk_nope_head_dim=self.qk_nope_head_dim,
129 + qk_rope_head_dim=self.qk_rope_head_dim,
130 + v_head_dim=self.v_head_dim,
131 + q_lora_rank=self.q_lora_rank,
132 + kv_lora_rank=self.kv_lora_rank,
133 + kv_b_proj=self.kv_b_proj,
134 + cache_config=vllm_config.cache_config,
135 + quant_config=vllm_config.quant_config,
136 + prefix=prefix,
137 + attn_backend=attn_backend.get_class(),
138 + )
139 + self.attn_backend: type[AttentionBackend] = self.mla_attn.get_attn_backend()
140 + self.mla_attn._k_scale = self.mla_attn._k_scale.to(device)
141 + self.mla_attn._v_scale = self.mla_attn._v_scale.to(device)
142 +
143 + # Keep both the string dtype (for ops) and torch dtype (for tensors)
144 + self.kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
145 + self.kv_cache_dtype = (
146 + FP8_DTYPE if self.kv_cache_dtype_str.startswith("fp8") else self.dtype
147 + )
148 +
149 + # Initialize attn MetadataBuilder
150 + self.builder = self.attn_backend.get_builder_cls()(
151 + kv_cache_spec=self.mla_attn.get_kv_cache_spec(vllm_config),
152 + layer_names=[self.mla_attn.layer_name],
153 + vllm_config=vllm_config,
154 + device=device,
155 + )
156 +
157 + def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
158 + """Initialize attention metadata."""
159 + # Create common attn metadata
160 + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
161 + common_attn_metadata = create_common_attn_metadata(
162 + batch_spec, self.block_size, self.device, arange_block_indices=True
163 + )
164 +
165 + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
166 + num_blocks = batch_size * max_blocks
167 +
168 + # Fetch the attention backend and kv cache shape and stride order
169 + kv_cache_shape = self.attn_backend.get_kv_cache_shape(
170 + num_blocks, self.block_size, self.num_kv_heads, self.head_size
171 + )
172 + try:
173 + kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
174 + except (AttributeError, NotImplementedError):
175 + kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
176 +
177 + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
178 + inv_order = [
179 + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
180 + ]
181 +
182 + raw_tensor = torch.zeros(
183 + num_blocks * self.block_size * self.num_kv_heads * self.head_size,
184 + dtype=self.kv_cache_dtype,
185 + device=self.device,
186 + )
187 + raw_tensor = raw_tensor.view(kv_cache_shape)
188 + kv_cache = raw_tensor.permute(*inv_order)
189 +
190 + self.mla_attn.kv_cache = kv_cache
191 +
192 + # Build attn metadata
193 + attn_metadata = self.builder.build(
194 + common_prefix_len=0, common_attn_metadata=common_attn_metadata
195 + )
196 +
197 + return attn_metadata
198 +
199 + def forward(
200 + self, qkv_lora: torch.Tensor, positions: torch.Tensor
201 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
202 + qkv_lora = qkv_lora.clone()
203 + q_c, kv_lora = qkv_lora.split(
204 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
205 + dim=-1,
206 + )
207 + q = self.q_b_proj(q_c)[0]
208 + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
209 +
210 + q = q.view(-1, self.num_heads, self.qk_head_dim)
211 + k_pe = k_pe.unsqueeze(1)
212 +
213 + q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
214 + positions, q[..., self.qk_nope_head_dim :], k_pe
215 + )
216 +
217 + dummy = torch.ops.vllm.unified_mla_kv_cache_update(
218 + kv_c,
219 + k_pe,
220 + _encode_layer_name(self.layer_name),
221 + self.kv_cache_dtype_str,
222 + self.mla_attn._k_scale,
223 + )
224 + return q, kv_c, k_pe, dummy
225 +
226 + def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
227 + ops = [
228 + INDEX_SELECT_OP,
229 + torch.ops.vllm.unified_mla_kv_cache_update.default,
230 + ]
231 + return ops
232 +
233 + def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
234 + return [torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default]
235 +
236 +
237 +MLA_BACKENDS = [AttentionBackendEnum.TRITON_MLA]
238 +if flash_attn_supports_mla():
239 + MLA_BACKENDS += [AttentionBackendEnum.FLASH_ATTN_MLA]
240 +if is_aiter_found_and_supported():
241 + MLA_BACKENDS += [AttentionBackendEnum.ROCM_AITER_MLA]
242 +
243 +
244 +@pytest.mark.parametrize("attn_backend", MLA_BACKENDS)
245 +@pytest.mark.parametrize("use_deepseek_scaling_rope", [True])
246 +@pytest.mark.parametrize("num_heads", [16])
247 +@pytest.mark.parametrize("qk_nope_head_dim", [128])
248 +@pytest.mark.parametrize("qk_rope_head_dim", [64])
249 +@pytest.mark.parametrize("v_head_dim", [128])
250 +@pytest.mark.parametrize("q_lora_rank", [1536])
251 +@pytest.mark.parametrize("kv_lora_rank", [512])
252 +@pytest.mark.parametrize("block_size", [16])
253 +@pytest.mark.parametrize("is_neox", [True, False])
254 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
255 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
256 +@pytest.mark.skipif(
257 + not current_platform.is_cuda_alike(),
258 + reason="MLA RoPE+KVCache+Cat fusion is only supported on CUDA and ROCm.",
259 +)
260 +def test_mla_rope_kvcache_cat_fusion(
261 + attn_backend: AttentionBackendEnum,
262 + use_deepseek_scaling_rope: bool,
263 + num_heads: int,
264 + qk_nope_head_dim: int,
265 + qk_rope_head_dim: int,
266 + v_head_dim: int,
267 + q_lora_rank: int,
268 + kv_lora_rank: int,
269 + block_size: int,
270 + is_neox: bool,
271 + dtype: torch.dtype,
272 + kv_cache_dtype: str,
273 + monkeypatch: pytest.MonkeyPatch,
274 +):
275 + torch.set_default_device("cuda")
276 + torch.set_default_dtype(dtype)
277 + torch.manual_seed(0)
278 +
279 + vllm_config = VllmConfig(
280 + model_config=ModelConfig(
281 + model="deepseek-ai/DeepSeek-V2-Lite",
282 + dtype=dtype,
283 + ),
284 + cache_config=CacheConfig(
285 + block_size=block_size,
286 + cache_dtype=kv_cache_dtype,
287 + ),
288 + compilation_config=CompilationConfig(
289 + mode=CompilationMode.VLLM_COMPILE,
290 + pass_config=PassConfig(
291 + fuse_rope_kvcache_cat_mla=True,
292 + eliminate_noops=True,
293 + ),
294 + ),
295 + )
296 +
297 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
298 + if not torch.distributed.is_initialized():
299 + from vllm.distributed.parallel_state import (
300 + init_distributed_environment,
301 + initialize_model_parallel,
302 + )
303 + from vllm.utils.system_utils import update_environment_variables
304 +
305 + update_environment_variables(
306 + {
307 + "RANK": "0",
308 + "LOCAL_RANK": "0",
309 + "WORLD_SIZE": "1",
310 + "MASTER_ADDR": "localhost",
311 + "MASTER_PORT": "54321",
312 + }
313 + )
314 + init_distributed_environment()
315 + initialize_model_parallel()
316 +
317 + if attn_backend == AttentionBackendEnum.ROCM_AITER_MLA:
318 + m.setenv("VLLM_ROCM_USE_AITER", "1")
319 + rocm_aiter_ops.refresh_env_variables()
320 +
321 + model = MLARoPEKVCacheCatTestModel(
322 + vllm_config=vllm_config,
323 + attn_backend=attn_backend,
324 + use_deepseek_scaling_rope=use_deepseek_scaling_rope,
325 + num_heads=num_heads,
326 + qk_nope_head_dim=qk_nope_head_dim,
327 + qk_rope_head_dim=qk_rope_head_dim,
328 + v_head_dim=v_head_dim,
329 + q_lora_rank=q_lora_rank,
330 + kv_lora_rank=kv_lora_rank,
331 + is_neox=is_neox,
332 + dtype=dtype,
333 + device=torch.get_default_device(),
334 + )
335 +
336 + fusion_pass = MLARoPEKVCacheCatFusionPass(vllm_config)
337 + # note: FixFunctionalizationPass is required to correctly lower
338 + # the fused op to its inplace version with auto-functionalization v1.
339 + # Without it, decompose_auto_functionalized calls clone_preserve_strides
340 + # on the non-contiguous q_pe slice directly, and inductor's lowering
341 + # of the resulting as_strided chain incorrectly drops the storage offset.
342 + # auto-functionalization v2 avoids this: it clones the contiguous base
343 + # tensor (_all_bases) and reconstructs the slice as a view, so the
344 + # offset is never passed through as_strided lowering.
345 + passes = [
346 + NoOpEliminationPass(vllm_config),
347 + fusion_pass,
348 + PostCleanupPass(vllm_config),
349 + FixFunctionalizationPass(vllm_config),
350 + ]
351 + backend = TestBackend(*passes)
352 +
353 + T = 5
354 +
355 + qkv_lora = torch.randn(
356 + T,
357 + q_lora_rank + kv_lora_rank + qk_rope_head_dim,
358 + dtype=dtype,
359 + )
360 + pos = torch.arange(T, dtype=torch.long)
361 +
362 + qkv_unfused = qkv_lora.clone()
363 + pos_unfused = pos.clone()
364 +
365 + # Run unfused version
366 + with set_forward_context(None, vllm_config):
367 + forward_context = get_forward_context()
368 + attn_metadata = model.build_attn_metadata(T)
369 + forward_context.slot_mapping = {
370 + model.layer_name: attn_metadata.slot_mapping
371 + }
372 + q_unfused, kv_c_unfused, k_pe_unfused, dummy = model(
373 + qkv_unfused, pos_unfused
374 + )
375 + attn_layer = forward_context.no_compile_layers[model.layer_name]
376 + kv_cache_unfused = attn_layer.kv_cache.clone()
377 + del dummy
378 +
379 + # Run fused version (compiled)
380 + torch._dynamo.mark_dynamic(qkv_lora, 0)
381 + torch._dynamo.mark_dynamic(pos, 0)
382 + with set_forward_context(None, vllm_config):
383 + model_fused = torch.compile(model, backend=backend)
384 + forward_context = get_forward_context()
385 + attn_metadata = model.build_attn_metadata(T)
386 + forward_context.slot_mapping = {
387 + model.layer_name: attn_metadata.slot_mapping
388 + }
389 + q_fused, kv_c_fused, k_pe_fused, dummy = model_fused(qkv_lora, pos)
390 + attn_layer = forward_context.no_compile_layers[model.layer_name]
391 + kv_cache_fused = attn_layer.kv_cache
392 + del dummy
393 +
394 + assert fusion_pass.matched_count == 1
395 +
396 + backend.check_before_ops(model.ops_in_model_before())
397 + backend.check_after_ops(model.ops_in_model_after())
398 +
399 + if dtype == torch.float16:
400 + ATOL, RTOL = (2e-3, 2e-3)
401 + else:
402 + ATOL, RTOL = (1e-2, 1e-2)
403 +
404 + torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
405 + torch.testing.assert_close(kv_c_unfused, kv_c_fused, atol=ATOL, rtol=RTOL)
406 + torch.testing.assert_close(k_pe_unfused, k_pe_fused, atol=ATOL, rtol=RTOL)
407 + # Cannot compare fp8_* directly here, cast to model dtype instead
408 + torch.testing.assert_close(
409 + kv_cache_unfused.view(dtype),
410 + kv_cache_fused.view(dtype),
411 + atol=ATOL,
412 + rtol=RTOL,
413 + )
tests/compile/passes/test_rope_kvcache_fusion.py
+4 -11
@@ -34,7 +34,6 @@ from vllm.v1.attention.backend import (
34 CommonAttentionMetadata,
35 )
36 from vllm.v1.attention.backends.registry import AttentionBackendEnum
37 -from vllm.v1.kv_cache_interface import AttentionSpec
37
38 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
39 VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
@@ -102,13 +101,8 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
101 )
102
103 # Initialize attn MetadataBuilder
105 - self.builder = self.attn.attn_backend.get_builder_cls()(
106 - kv_cache_spec=AttentionSpec(
107 - block_size=self.block_size,
108 - num_kv_heads=self.num_kv_heads,
109 - head_size=head_size,
110 - dtype=self.kv_cache_dtype,
111 - ),
104 + self.builder = self.attn_backend.get_builder_cls()(
105 + kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config),
106 layer_names=[self.attn.layer_name],
107 vllm_config=vllm_config,
108 device=device,
@@ -126,12 +120,11 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
120 num_blocks = batch_size * max_blocks
121
122 # Fetch the attention backend and kv cache shape and stride order
129 - attn_backend = self.attn.attn_backend
130 - kv_cache_shape = attn_backend.get_kv_cache_shape(
123 + kv_cache_shape = self.attn_backend.get_kv_cache_shape(
124 num_blocks, self.block_size, self.num_kv_heads, self.head_size
125 )
126 try:
134 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
127 + kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
128 except (AttributeError, NotImplementedError):
129 kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
130
vllm/compilation/passes/fusion/matcher_utils.py
+84
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 kNvfp4Dynamic,
24 )
25 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
26 +from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
27 + DeepseekScalingRotaryEmbedding,
28 +)
29 from vllm.platforms import current_platform
30
31 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
@@ -158,6 +161,87 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
161 return result
162
163
164 +class MatcherDeepseekScalingRotaryEmbedding(MatcherCustomOp):
165 + def __init__(
166 + self,
167 + is_neox: bool,
168 + head_size: int,
169 + num_heads: int,
170 + num_kv_heads: int,
171 + use_flashinfer: bool = False,
172 + enabled: bool | None = None,
173 + ) -> None:
174 + if enabled is None:
175 + enabled = DeepseekScalingRotaryEmbedding.enabled()
176 +
177 + super().__init__(enabled)
178 + self.is_neox = is_neox
179 + self.head_size = head_size
180 + self.num_heads = num_heads
181 + self.num_kv_heads = num_kv_heads
182 + self.q_size = self.num_heads * self.head_size
183 + self.kv_size = self.num_kv_heads * self.head_size
184 + self.rotary_dim = head_size
185 + self.use_flashinfer = use_flashinfer
186 +
187 + def inputs(self) -> list[torch.Tensor]:
188 + positions = self.empty_int64(5)
189 + query = self.empty(5, self.num_heads, self.head_size)
190 + key = self.empty(5, self.num_kv_heads, self.head_size)
191 + cos_sin_cache = self.empty(4096, self.rotary_dim)
192 + return [positions, query, key, cos_sin_cache]
193 +
194 + def forward_custom(
195 + self,
196 + positions: torch.Tensor,
197 + query: torch.Tensor,
198 + key: torch.Tensor | None,
199 + cos_sin_cache: torch.Tensor,
200 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
201 + if self.use_flashinfer:
202 + torch.ops.vllm.flashinfer_rotary_embedding(
203 + positions,
204 + query,
205 + key,
206 + self.head_size,
207 + cos_sin_cache,
208 + self.is_neox,
209 + )
210 + return query, key
211 + result: tuple[torch.Tensor, torch.Tensor | None] = (
212 + DeepseekScalingRotaryEmbedding.forward_static(
213 + positions,
214 + query,
215 + key,
216 + self.head_size,
217 + self.rotary_dim,
218 + cos_sin_cache,
219 + self.is_neox,
220 + )
221 + )
222 + return result
223 +
224 + def forward_native(
225 + self,
226 + positions: torch.Tensor,
227 + query: torch.Tensor,
228 + key: torch.Tensor | None,
229 + cos_sin_cache: torch.Tensor,
230 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
231 + result: tuple[torch.Tensor, torch.Tensor | None] = (
232 + DeepseekScalingRotaryEmbedding.forward_static(
233 + positions,
234 + query,
235 + key,
236 + self.head_size,
237 + self.rotary_dim,
238 + cos_sin_cache,
239 + self.is_neox,
240 + )
241 + )
242 + return result
243 +
244 +
245 class MatcherQuantFP8(MatcherCustomOp):
246 def __init__(
247 self,
vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py new
+271
@@ -0,0 +1,271 @@
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, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
39 + if layer_slot_mapping is not None:
40 + ops.concat_and_cache_mla_rope_fused(
41 + positions,
42 + q_pe,
43 + k_pe,
44 + kv_c,
45 + cos_sin_cache,
46 + is_neox,
47 + layer_slot_mapping,
48 + kv_cache,
49 + kv_cache_dtype,
50 + kv_cache_scale,
51 + )
52 + return torch.empty(0, device=kv_c.device, dtype=kv_c.dtype)
53 +
54 +
55 +def fused_rope_unified_mla_kv_cache_update_fake(
56 + positions: torch.Tensor,
57 + q_pe: torch.Tensor,
58 + k_pe: torch.Tensor,
59 + kv_c: torch.Tensor,
60 + cos_sin_cache: torch.Tensor,
61 + is_neox: bool,
62 + kv_cache_dtype: str,
63 + kv_cache_scale: torch.Tensor,
64 + layer_name: LayerNameType,
65 +) -> torch.Tensor:
66 + return torch.empty(0, dtype=kv_c.dtype, device=kv_c.device)
67 +
68 +
69 +direct_register_custom_op(
70 + op_name="fused_rope_unified_mla_kv_cache_update",
71 + op_func=fused_rope_unified_mla_kv_cache_update_impl,
72 + fake_impl=fused_rope_unified_mla_kv_cache_update_fake,
73 + mutates_args=["q_pe", "k_pe"],
74 +)
75 +
76 +
77 +class MLARoPEKVCacheCatPattern(VllmPatternReplacement):
78 + FUSED_OP = torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default
79 +
80 + def __init__(
81 + self,
82 + layer: MLAAttention,
83 + is_neox: bool,
84 + use_flashinfer: bool = False,
85 + use_deepseek_scaling: bool = False,
86 + ) -> None:
87 + self.layer_name = layer.layer_name
88 + self.kv_cache_dtype = layer.kv_cache_dtype
89 + self.num_heads = layer.num_heads
90 + self.num_kv_heads = layer.num_kv_heads
91 + self.kv_lora_rank = layer.kv_lora_rank
92 + self.qk_rope_head_dim = layer.qk_rope_head_dim
93 + self.is_neox = is_neox
94 + self.use_flashinfer = use_flashinfer
95 + self._ln = _encode_layer_name(self.layer_name)
96 +
97 + if use_deepseek_scaling:
98 + self.rope_matcher = MatcherDeepseekScalingRotaryEmbedding(
99 + is_neox=self.is_neox,
100 + head_size=self.qk_rope_head_dim,
101 + num_heads=self.num_heads,
102 + num_kv_heads=self.num_kv_heads,
103 + use_flashinfer=self.use_flashinfer,
104 + )
105 + else:
106 + self.rope_matcher = MatcherRotaryEmbedding( # type: ignore
107 + is_neox=self.is_neox,
108 + head_size=self.qk_rope_head_dim,
109 + num_heads=self.num_heads,
110 + num_kv_heads=self.num_kv_heads,
111 + use_flashinfer=self.use_flashinfer,
112 + )
113 +
114 + def get_inputs(self) -> list[torch.Tensor]:
115 + T = 5
116 + L = 4096
117 + q_pe = self.empty_bf16(T, self.num_heads, self.qk_rope_head_dim)
118 + k_pe = self.empty_bf16(T, self.qk_rope_head_dim)
119 + kv_c_normed = self.empty_bf16(T, self.kv_lora_rank)
120 + cos_sin_cache = self.empty_bf16(L, self.qk_rope_head_dim)
121 + positions = self.empty(T, dtype=torch.int64)
122 + k_scale = self.empty(0, dtype=torch.float32)
123 + inputs = [
124 + q_pe,
125 + k_pe,
126 + kv_c_normed,
127 + positions,
128 + cos_sin_cache,
129 + k_scale,
130 + ]
131 + if _USE_LAYERNAME:
132 + inputs.append(self._ln)
133 + return inputs
134 +
135 + @property
136 + def pattern(self):
137 + _ln = self._ln
138 +
139 + if _USE_LAYERNAME:
140 +
141 + def _pattern_with_ln(
142 + q_pe: torch.Tensor,
143 + k_pe: torch.Tensor,
144 + kv_c_normed: torch.Tensor,
145 + positions: torch.Tensor,
146 + cos_sin_cache: torch.Tensor,
147 + k_scale: torch.Tensor,
148 + layer_name: LayerNameType,
149 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
150 + k_pe_unsqueezed = k_pe.unsqueeze(1)
151 + q_pe, k_pe = self.rope_matcher(
152 + positions, q_pe, k_pe_unsqueezed, cos_sin_cache
153 + )
154 + dummy = torch.ops.vllm.unified_mla_kv_cache_update(
155 + kv_c_normed, k_pe, layer_name, self.kv_cache_dtype, k_scale
156 + )
157 + return dummy, q_pe, k_pe
158 +
159 + return _pattern_with_ln
160 +
161 + def _pattern(
162 + q_pe: torch.Tensor,
163 + k_pe: torch.Tensor,
164 + kv_c_normed: torch.Tensor,
165 + positions: torch.Tensor,
166 + cos_sin_cache: torch.Tensor,
167 + k_scale: torch.Tensor,
168 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
169 + k_pe_unsqueezed = k_pe.unsqueeze(1)
170 + q_pe, k_pe = self.rope_matcher(
171 + positions, q_pe, k_pe_unsqueezed, cos_sin_cache
172 + )
173 + dummy = torch.ops.vllm.unified_mla_kv_cache_update(
174 + kv_c_normed, k_pe, _ln, self.kv_cache_dtype, k_scale
175 + )
176 + return dummy, q_pe, k_pe
177 +
178 + return _pattern
179 +
180 + @property
181 + def replacement(self):
182 + _ln = self._ln
183 +
184 + if _USE_LAYERNAME:
185 +
186 + def _replacement_with_ln(
187 + q_pe: torch.Tensor,
188 + k_pe: torch.Tensor,
189 + kv_c_normed: torch.Tensor,
190 + positions: torch.Tensor,
191 + cos_sin_cache: torch.Tensor,
192 + k_scale: torch.Tensor,
193 + layer_name: LayerNameType,
194 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
195 + at = auto_functionalized(
196 + self.FUSED_OP,
197 + positions=positions,
198 + q_pe=q_pe,
199 + k_pe=k_pe,
200 + kv_c=kv_c_normed,
201 + cos_sin_cache=cos_sin_cache,
202 + is_neox=self.is_neox,
203 + kv_cache_dtype=self.kv_cache_dtype,
204 + kv_cache_scale=k_scale,
205 + layer_name=layer_name,
206 + )
207 + dummy, q_pe, k_pe_squeezed = at
208 + k_pe = k_pe_squeezed.unsqueeze(1)
209 + return dummy, q_pe, k_pe
210 +
211 + return _replacement_with_ln
212 +
213 + def _replacement(
214 + q_pe: torch.Tensor,
215 + k_pe: torch.Tensor,
216 + kv_c_normed: torch.Tensor,
217 + positions: torch.Tensor,
218 + cos_sin_cache: torch.Tensor,
219 + k_scale: torch.Tensor,
220 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
221 + at = auto_functionalized(
222 + self.FUSED_OP,
223 + positions=positions,
224 + q_pe=q_pe,
225 + k_pe=k_pe,
226 + kv_c=kv_c_normed,
227 + cos_sin_cache=cos_sin_cache,
228 + is_neox=self.is_neox,
229 + kv_cache_dtype=self.kv_cache_dtype,
230 + kv_cache_scale=k_scale,
231 + layer_name=_ln,
232 + )
233 + dummy, q_pe, k_pe_squeezed = at
234 + k_pe = k_pe_squeezed.unsqueeze(1)
235 + return dummy, q_pe, k_pe
236 +
237 + return _replacement
238 +
239 +
240 +class MLARoPEKVCacheCatFusionPass(VllmFusionPatternMatcherPass):
241 + def __init__(self, config: VllmConfig) -> None:
242 + super().__init__(config, "mla_rope_kv_cache_fusion_pass")
243 +
244 + attn_layers = get_layers_from_vllm_config(config, MLAAttention)
245 +
246 + for _, layer in attn_layers.items():
247 + for is_neox in [False, True]:
248 + for use_deepseek_scaling in [False, True]:
249 + if RotaryEmbedding.enabled():
250 + for use_flashinfer in [False, True]:
251 + self.register(
252 + MLARoPEKVCacheCatPattern(
253 + layer,
254 + is_neox,
255 + use_flashinfer,
256 + use_deepseek_scaling,
257 + )
258 + )
259 + else:
260 + self.register(
261 + MLARoPEKVCacheCatPattern(
262 + layer,
263 + is_neox,
264 + use_deepseek_scaling=use_deepseek_scaling,
265 + )
266 + )
267 +
268 + if _USE_LAYERNAME:
269 + break
270 +
271 + self.dump_patterns(config, self.pm_pass)
vllm/compilation/passes/pass_manager.py
+4
@@ -33,6 +33,7 @@ if current_platform.is_cuda_alike():
33 from .fusion.act_quant_fusion import ActivationQuantFusionPass
34 from .fusion.attn_quant_fusion import AttnQuantFusionPass
35 from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass
36 + from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass
37 from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
38 from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
39 from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
@@ -174,6 +175,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
175 self.passes += [ScatterSplitReplacementPass(config)]
176 self.passes += [RopeKVCacheFusionPass(config)]
177
178 + if self.pass_config.fuse_rope_kvcache_cat_mla:
179 + self.passes += [MLARoPEKVCacheCatFusionPass(config)]
180 +
181 if self.pass_config.fuse_attn_quant:
182 self.passes += [AttnQuantFusionPass(config)]
183 self.passes += [MLAAttnQuantFusionPass(config)]
vllm/compilation/passes/utility/fix_functionalization.py
+39
@@ -181,6 +181,45 @@ class FixFunctionalizationPass(VllmInductorPass):
181 2: "key",
182 }
183 self.defunctionalize(graph, node, mutated_args=mutated_args)
184 + elif (
185 + hasattr(torch.ops.vllm, "fused_rope_unified_mla_kv_cache_update")
186 + and at_target
187 + == torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default
188 + ):
189 + # AOTAutograd functionalizes `q[..., nope_dim:] = rope_result` into
190 + # a sequence of aten ops on q: view+slice+copy+slice_scatter.
191 + # Since the fused MLA RoPE op mutates q_pe in-place, we can remove
192 + # the redundant copy and slice_scatter ops during defunctionalization.
193 + getitem_nodes = self.getitem_users(node)
194 + q_pe_out = getitem_nodes[1]
195 +
196 + for user in list(q_pe_out.users):
197 + if is_func(user, torch.ops.aten.copy.default):
198 + copy_temp = user
199 + slice_temp = copy_temp.args[0]
200 + for user in list(copy_temp.users):
201 + if is_func(user, torch.ops.aten.slice_scatter.default):
202 + slice_scatter_temp = user
203 + view_temp = slice_scatter_temp.args[0]
204 +
205 + view_orig = slice_temp.args[0]
206 + slice_scatter_temp.replace_all_uses_with(view_orig)
207 + self._remove(slice_scatter_temp)
208 + self._remove(copy_temp)
209 + self._remove(slice_temp)
210 + self._remove(view_temp)
211 + self._remove(q_pe_out)
212 +
213 + # defunctionalize k_pe manually; self.replace_users_with_mutated_args
214 + # does not support only replacing specific kwargs
215 + k_pe_in = node.kwargs["k_pe"]
216 + k_pe_out = getitem_nodes[2]
217 + k_pe_out.replace_all_uses_with(k_pe_in)
218 + self._remove(k_pe_out)
219 +
220 + self.insert_defunctionalized(graph, node)
221 + self._remove(node)
222 +
223 # only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn
224 elif (
225 hasattr(torch.ops.vllm, "function_with_mutated_args_and_return")
vllm/config/compilation.py
+10 -1
@@ -136,8 +136,10 @@ class PassConfig:
136 """Enable flashinfer allreduce fusion."""
137 fuse_minimax_qk_norm: bool = None # type: ignore[assignment]
138 """Enable fused allreduce+RMSNorm for MiniMax QK norm."""
139 - enable_qk_norm_rope_fusion: bool = False
139 + enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment]
140 """Enable fused Q/K RMSNorm + RoPE pass."""
141 + fuse_rope_kvcache_cat_mla: bool = None # type: ignore[assignment]
142 + """Enable fused MLA KV cache update with RoPE."""
143
144 # ROCm/AITER specific fusions
145 fuse_act_padding: bool = None # type: ignore[assignment]
@@ -228,6 +230,7 @@ class PassConfig:
230 "fuse_act_padding",
231 "fuse_mla_dual_rms_norm",
232 "fuse_rope_kvcache",
233 + "fuse_rope_kvcache_cat_mla",
234 mode="wrap",
235 )
236 @classmethod
@@ -285,6 +288,12 @@ class PassConfig:
288 "The fusion will be disabled."
289 )
290 self.fuse_rope_kvcache = False
291 + if self.fuse_rope_kvcache_cat_mla and not current_platform.is_cuda_alike():
292 + logger.warning_once(
293 + "MLA KV cache update with RoPE fusion enabled but the "
294 + "current platform is not CUDA or ROCm. The fusion will be disabled."
295 + )
296 + self.fuse_rope_kvcache_cat_mla = False
297
298 def log_enabled_passes(self) -> None:
299 """
vllm/config/vllm.py
+13
@@ -155,6 +155,15 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:
155 )
156
157
158 +def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool:
159 + """Enable if use_inductor_graph_partition is enabled."""
160 +
161 + return (
162 + cfg.compilation_config.use_inductor_graph_partition
163 + or not cfg.compilation_config.splitting_ops_contain_kv_cache_update()
164 + )
165 +
166 +
167 def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
168 """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss."""
169
@@ -184,6 +193,7 @@ OPTIMIZATION_LEVEL_00 = {
193 "fuse_act_padding": False,
194 "fuse_mla_dual_rms_norm": False,
195 "fuse_rope_kvcache": False,
196 + "fuse_rope_kvcache_cat_mla": False,
197 },
198 "cudagraph_mode": CUDAGraphMode.NONE,
199 "use_inductor_graph_partition": False,
@@ -204,6 +214,7 @@ OPTIMIZATION_LEVEL_01 = {
214 "fuse_act_padding": enable_norm_pad_fusion,
215 "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
216 "fuse_rope_kvcache": False,
217 + "fuse_rope_kvcache_cat_mla": False,
218 },
219 "cudagraph_mode": CUDAGraphMode.PIECEWISE,
220 "use_inductor_graph_partition": False,
@@ -226,6 +237,7 @@ OPTIMIZATION_LEVEL_02 = {
237 "fuse_act_padding": enable_norm_pad_fusion,
238 "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
239 "fuse_rope_kvcache": enable_rope_kvcache_fusion,
240 + "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion,
241 },
242 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
243 "use_inductor_graph_partition": False,
@@ -248,6 +260,7 @@ OPTIMIZATION_LEVEL_03 = {
260 "fuse_act_padding": enable_norm_pad_fusion,
261 "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
262 "fuse_rope_kvcache": enable_rope_kvcache_fusion,
263 + "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion,
264 },
265 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
266 "use_inductor_graph_partition": False,
vllm/model_executor/layers/attention/mla_attention.py
+19 -24
@@ -345,6 +345,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
345 cache_config: CacheConfig | None = None,
346 quant_config: QuantizationConfig | None = None,
347 prefix: str = "",
348 + attn_backend: type[AttentionBackend] | None = None,
349 use_sparse: bool = False,
350 indexer: object | None = None,
351 **extra_impl_args,
@@ -374,14 +375,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):
375 self.quant_config = quant_config
376
377 dtype = torch.get_default_dtype()
377 - self.attn_backend = get_attn_backend(
378 - self.head_size,
379 - dtype,
380 - kv_cache_dtype,
381 - use_mla=True,
382 - use_sparse=use_sparse,
383 - num_heads=self.num_heads,
384 - )
378 + if attn_backend is not None:
379 + assert attn_backend.is_mla(), (
380 + f"MLAAttention: attn_backend must be an MLA backend, "
381 + f"got {attn_backend.get_name()} instead"
382 + )
383 + self.attn_backend = attn_backend
384 + else:
385 + self.attn_backend = get_attn_backend(
386 + self.head_size,
387 + dtype,
388 + kv_cache_dtype,
389 + use_mla=True,
390 + use_sparse=use_sparse,
391 + num_heads=self.num_heads,
392 + )
393
394 # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format
395 # Automatically convert fp8 kv-cache format to "fp8_ds_mla"
@@ -1008,23 +1016,9 @@ def unified_mla_kv_cache_update(
1016 the data dependency between them to ensure torch.compile preserves ordering.
1017 """
1018 layer_name = _resolve_layer_name(layer_name)
1011 - forward_context = get_forward_context()
1012 - attn_layer = forward_context.no_compile_layers[layer_name]
1013 - kv_cache = attn_layer.kv_cache
1014 -
1015 - # This needs to run even when we don't have metadata yet, so that the op
1016 - # is correctly captured.
1017 - if kv_cache.numel() == 0:
1018 - # Can't update an empty KV cache.
1019 - return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
1020 -
1021 - slot_mapping = forward_context.slot_mapping
1022 - assert isinstance(slot_mapping, dict), (
1023 - f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
1024 - )
1025 - layer_slot_mapping = slot_mapping.get(layer_name)
1019 + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
1020 if layer_slot_mapping is not None:
1027 - attn_layer.impl.do_kv_cache_update(
1021 + attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined]
1022 kv_c_normed,
1023 k_pe,
1024 kv_cache,
@@ -1113,6 +1107,7 @@ direct_register_custom_op(
1107 mutates_args=["output", "output_block_scale"],
1108 fake_impl=unified_mla_attention_with_output_fake,
1109 dispatch_key=current_platform.dispatch_key,
1110 + tags=(torch.Tag.flexible_layout,),
1111 )
1112
1113
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+32 -9
@@ -127,29 +127,52 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
127 ) -> tuple[torch.Tensor, torch.Tensor | None]:
128 """PyTorch-native implementation equivalent to forward()."""
129 assert key is not None
130 - cos_sin_cache = self._match_cos_sin_cache_dtype(query)
131 - query_rot = query[..., : self.rotary_dim]
132 - key_rot = key[..., : self.rotary_dim]
133 - if self.rotary_dim < self.head_size:
134 - query_pass = query[..., self.rotary_dim :]
135 - key_pass = key[..., self.rotary_dim :]
130 + return self.forward_static(
131 + positions,
132 + query,
133 + key,
134 + self.head_size,
135 + self.rotary_dim,
136 + self.cos_sin_cache,
137 + self.is_neox_style,
138 + offsets,
139 + )
140 +
141 + @staticmethod
142 + def forward_static(
143 + positions: torch.Tensor,
144 + query: torch.Tensor,
145 + key: torch.Tensor | None,
146 + head_size: int,
147 + rotary_dim: int,
148 + cos_sin_cache: torch.Tensor,
149 + is_neox_style: bool,
150 + offsets: torch.Tensor | None = None,
151 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
152 + """A static implementation of forward()."""
153 + assert key is not None
154 + query_rot = query[..., :rotary_dim]
155 + key_rot = key[..., :rotary_dim]
156 + if rotary_dim < head_size:
157 + query_pass = query[..., rotary_dim:]
158 + key_pass = key[..., rotary_dim:]
159
160 cos_sin = cos_sin_cache[
161 torch.add(positions, offsets) if offsets is not None else positions
162 ]
163 cos, sin = cos_sin.chunk(2, dim=-1)
141 - if self.is_neox_style:
164 + if is_neox_style:
165 cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
166 sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
167 else:
168 cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
169 sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
170
148 - rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj
171 + rotate_fn = rotate_neox if is_neox_style else rotate_gptj
172 query_rot = query_rot * cos + rotate_fn(query_rot) * sin
173 key_rot = key_rot * cos + rotate_fn(key_rot) * sin
174
152 - if self.rotary_dim < self.head_size:
175 + if rotary_dim < head_size:
176 query = torch.cat((query_rot, query_pass), dim=-1)
177 key = torch.cat((key_rot, key_pass), dim=-1)
178 else:
vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py
+2 -4
@@ -195,10 +195,8 @@ class DualChunkRotaryEmbedding(CustomOp):
195 def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass):
196 cos, sin = cos_sin.chunk(2, dim=-1)
197 if self.is_neox_style:
198 - # NOTE(woosuk): Here we assume that the positions tensor has the
199 - # shape [batch_size, seq_len].
200 - cos = cos.repeat(1, 1, 2).unsqueeze(-2)
201 - sin = sin.repeat(1, 1, 2).unsqueeze(-2)
198 + cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
199 + sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
200 else:
201 cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
202 sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)