From a51376b3f05a2f74eac6ceeed7e52598b871a0fb Mon Sep 17 00:00:00 2001 From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Date: Mon, 11 May 2026 07:10:50 -0700 Subject: [PATCH] [Performance][DSR1]: Fused RoPE+KVCache+q_concat for MLA (#40392) Signed-off-by: Rohan138 Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Co-authored-by: ElizaWszola --- csrc/cache_kernels_fused.cu | 135 +++--- .../test_mla_rope_kvcache_cat_fusion.py | 413 ++++++++++++++++++ .../passes/test_rope_kvcache_fusion.py | 15 +- .../passes/fusion/matcher_utils.py | 84 ++++ .../fusion/mla_rope_kvcache_cat_fusion.py | 271 ++++++++++++ vllm/compilation/passes/pass_manager.py | 4 + .../passes/utility/fix_functionalization.py | 39 ++ vllm/config/compilation.py | 11 +- vllm/config/vllm.py | 13 + .../layers/attention/mla_attention.py | 43 +- .../rotary_embedding/deepseek_scaling_rope.py | 41 +- .../rotary_embedding/dual_chunk_rope.py | 6 +- 12 files changed, 966 insertions(+), 109 deletions(-) create mode 100644 tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py create mode 100644 vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py diff --git a/csrc/cache_kernels_fused.cu b/csrc/cache_kernels_fused.cu index be037b2fde..8687ebe1f1 100644 --- a/csrc/cache_kernels_fused.cu +++ b/csrc/cache_kernels_fused.cu @@ -21,28 +21,33 @@ namespace vllm { // NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's // using u16 as the backing type. -template +template __global__ void concat_and_cache_mla_rope_fused_kernel( const int64_t* __restrict__ positions, // [num_tokens] qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim] qk_t* __restrict__ k_pe, // [num_tokens, rot_dim] const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] - const qk_t* __restrict__ rope_cos_sin_cache, // [max_position, 2, - // rot_dim // 2] + const cos_sin_t* __restrict__ rope_cos_sin_cache, // [max_position, 2, + // rot_dim // 2] const int rot_dim, const int64_t q_pe_stride_token, const int64_t q_pe_stride_head, const int64_t k_pe_stride, const int64_t kv_c_stride, const int num_q_heads, cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank + // rot_dim)] - const int64_t* __restrict__ kv_cache_slot_mapping, // [num_tokens] + const int64_t* __restrict__ slot_mapping, // [num_tokens] const int block_stride, const int entry_stride, const int kv_lora_rank, const int block_size, const float* kv_cache_quant_scale) { // Each thread block is responsible for one token. const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0) { + return; + } const int64_t pos = positions[token_idx]; - const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim; + const cos_sin_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim; const int embed_dim = rot_dim / 2; @@ -54,8 +59,8 @@ __global__ void concat_and_cache_mla_rope_fused_kernel( // NOTE: Would be nice to have interleaved sin/cos so we could just load // both at the same time. - qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx); - qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim); + qk_t cos = static_cast(VLLM_LDG(cos_sin_ptr + pair_idx)); + qk_t sin = static_cast(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim)); qk_t* q_pe_head_ptr = 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( q_pe_head_ptr[pair_idx_y] = y_dst; } - const int64_t slot_idx = kv_cache_slot_mapping[token_idx]; const int64_t block_idx = slot_idx / block_size; const int64_t entry_idx = slot_idx % block_size; - // NOTE: slot_idx can be -1 if the token is padded - if (slot_idx < 0) { - return; - } - // K with 1 HEAD for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) { int pair_idx = i; - qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx); - qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim); + qk_t cos = static_cast(VLLM_LDG(cos_sin_ptr + pair_idx)); + qk_t sin = static_cast(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim)); qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride; @@ -165,36 +164,43 @@ __global__ void concat_and_cache_mla_rope_fused_kernel( } // namespace vllm -#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ - do { \ - VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \ - using qk_t = scalar_t; \ - if (rope_is_neox) { \ - vllm::concat_and_cache_mla_rope_fused_kernel \ - <<>>( \ - positions.data_ptr(), q_pe.data_ptr(), \ - k_pe.data_ptr(), kv_c.data_ptr(), \ - rope_cos_sin_cache.data_ptr(), rot_dim, \ - q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \ - num_q_heads, reinterpret_cast(kv_cache.data_ptr()), \ - kv_cache_slot_mapping.data_ptr(), block_stride, \ - entry_stride, kv_lora_rank, block_size, \ - kv_cache_quant_scale.data_ptr()); \ - } else { \ - vllm::concat_and_cache_mla_rope_fused_kernel \ - <<>>( \ - positions.data_ptr(), q_pe.data_ptr(), \ - k_pe.data_ptr(), kv_c.data_ptr(), \ - rope_cos_sin_cache.data_ptr(), rot_dim, \ - q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \ - num_q_heads, reinterpret_cast(kv_cache.data_ptr()), \ - kv_cache_slot_mapping.data_ptr(), block_stride, \ - entry_stride, kv_lora_rank, block_size, \ - kv_cache_quant_scale.data_ptr()); \ - } \ - }); \ +#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ + do { \ + VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \ + using qk_t = scalar_t; \ + VLLM_DISPATCH_FLOATING_TYPES( \ + rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \ + [&] { \ + using cos_sin_t = scalar_t; \ + if (rope_is_neox) { \ + vllm::concat_and_cache_mla_rope_fused_kernel< \ + qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \ + <<>>( \ + positions.data_ptr(), q_pe.data_ptr(), \ + k_pe.data_ptr(), kv_c.data_ptr(), \ + rope_cos_sin_cache.data_ptr(), rot_dim, \ + q_pe_stride_token, q_pe_stride_head, k_pe_stride, \ + kv_c_stride, num_q_heads, \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.data_ptr(), block_stride, \ + entry_stride, kv_lora_rank, block_size, \ + kv_cache_quant_scale.data_ptr()); \ + } else { \ + vllm::concat_and_cache_mla_rope_fused_kernel< \ + qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \ + <<>>( \ + positions.data_ptr(), q_pe.data_ptr(), \ + k_pe.data_ptr(), kv_c.data_ptr(), \ + rope_cos_sin_cache.data_ptr(), rot_dim, \ + q_pe_stride_token, q_pe_stride_head, k_pe_stride, \ + kv_c_stride, num_q_heads, \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.data_ptr(), block_stride, \ + entry_stride, kv_lora_rank, block_size, \ + kv_cache_quant_scale.data_ptr()); \ + } \ + }); \ + }); \ } while (false) // 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( torch::Tensor& kv_c, // [num_tokens, kv_lora_rank] torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim] bool rope_is_neox, - torch::Tensor& - kv_cache_slot_mapping, // [num_tokens] or [num_actual_tokens] + torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)] const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) { - const int64_t num_tokens = q_pe.size(0); + // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from + // slot_mapping.size(0) because of padding for CUDA graphs. + // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because + // both include padding. + // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) + // since key includes padding for CUDA graphs, while slot_mapping does not. + // In this case, slot_mapping.size(0) represents the actual number of tokens + // before padding. + // For compatibility with both cases, we use slot_mapping.size(0) as the + // number of tokens. + int num_tokens = slot_mapping.size(0); + int num_padded_tokens = q_pe.size(0); + TORCH_CHECK_GE(num_padded_tokens, num_tokens); const int num_q_heads = q_pe.size(1); const int rot_dim = q_pe.size(2); const int kv_lora_rank = kv_c.size(1); - TORCH_CHECK(positions.size(0) >= - num_tokens); // CUDA Graphs might pad this for us + TORCH_CHECK_EQ(positions.size(0), num_padded_tokens); TORCH_CHECK_EQ(positions.dim(), 1); TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long); - TORCH_CHECK_EQ(q_pe.size(0), num_tokens); + TORCH_CHECK_EQ(q_pe.dim(), 3); + TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens); TORCH_CHECK_EQ(q_pe.size(1), num_q_heads); TORCH_CHECK_EQ(q_pe.size(2), rot_dim); - TORCH_CHECK_EQ(q_pe.dim(), 3); - TORCH_CHECK_EQ(k_pe.size(0), num_tokens); - TORCH_CHECK_EQ(k_pe.size(1), rot_dim); TORCH_CHECK_EQ(k_pe.dim(), 2); + TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens); + TORCH_CHECK_EQ(k_pe.size(1), rot_dim); TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type()); - TORCH_CHECK_EQ(kv_c.size(0), num_tokens); - TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank); TORCH_CHECK_EQ(kv_c.dim(), 2); + TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens); + TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank); TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type()); TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype()); TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim); - TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type()); - TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens); - TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long); + TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens); + TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long); TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim); TORCH_CHECK_EQ(kv_cache.dim(), 3); diff --git a/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py b/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py new file mode 100644 index 0000000000..cc3bfb7693 --- /dev/null +++ b/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.config +from tests.compile.backend import TestBackend +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm.compilation.passes.fusion.mla_rope_kvcache_cat_fusion import ( + MLARoPEKVCacheCatFusionPass, +) +from vllm.compilation.passes.utility.fix_functionalization import ( + FixFunctionalizationPass, +) +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass +from vllm.config import ( + CacheConfig, + CompilationConfig, + CompilationMode, + ModelConfig, + PassConfig, + VllmConfig, +) +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.rotary_embedding import ( + DeepseekScalingRotaryEmbedding, + RotaryEmbedding, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import _encode_layer_name +from vllm.v1.attention.backend import ( + AttentionBackend, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +INDEX_SELECT_OP = torch.ops.aten.index.Tensor +VLLM_UNIFIED_MLA_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_mla_kv_cache_update +FP8_DTYPE = current_platform.fp8_dtype() + + +class MLARoPEKVCacheCatTestModel(torch.nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + attn_backend: AttentionBackendEnum, + use_deepseek_scaling_rope: bool, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + is_neox: bool, + dtype: torch.dtype, + device: torch.device, + prefix: str = "model.layers.0.self_attn.attn", + ): + super().__init__() + self.num_heads = num_heads + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.dtype = dtype + self.device = device + self.layer_name = prefix + + self.num_kv_heads = 1 + self.head_size = kv_lora_rank + qk_rope_head_dim + self.block_size = vllm_config.cache_config.block_size + self.scale = self.qk_head_dim**-0.5 + + if use_deepseek_scaling_rope: + self.rotary_emb = DeepseekScalingRotaryEmbedding( + head_size=qk_rope_head_dim, + rotary_dim=qk_rope_head_dim, + max_position_embeddings=4096, + base=10000, + is_neox_style=is_neox, + scaling_factor=1.0, + dtype=dtype, + ) + else: + self.rotary_emb = RotaryEmbedding( + head_size=qk_rope_head_dim, + rotary_dim=qk_rope_head_dim, + max_position_embeddings=4096, + base=10000, + is_neox_style=is_neox, + dtype=dtype, + ) + + # Initialize intermediate mm layers for unit test + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + prefix=f"{prefix}.q_b_proj", + ).to(device) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + prefix=f"{prefix}.kv_b_proj", + ).to(device) + + # ColumnParallelLinear default init in bf16 with seed 0 produces + # near-zero weights (7/4.7M nonzero), making the GEMM output almost + # entirely zero and masking correctness bugs. Reinitialize to get + # dense outputs. + with torch.no_grad(): + torch.nn.init.normal_(self.q_b_proj.weight, std=0.02) + torch.nn.init.normal_(self.kv_b_proj.weight, std=0.02) + + # Register layer metadata for the fusion pass via MLAAttention + self.mla_attn = MLAAttention( + num_heads=self.num_heads, + scale=self.scale, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + kv_b_proj=self.kv_b_proj, + cache_config=vllm_config.cache_config, + quant_config=vllm_config.quant_config, + prefix=prefix, + attn_backend=attn_backend.get_class(), + ) + self.attn_backend: type[AttentionBackend] = self.mla_attn.get_attn_backend() + self.mla_attn._k_scale = self.mla_attn._k_scale.to(device) + self.mla_attn._v_scale = self.mla_attn._v_scale.to(device) + + # Keep both the string dtype (for ops) and torch dtype (for tensors) + self.kv_cache_dtype_str = vllm_config.cache_config.cache_dtype + self.kv_cache_dtype = ( + FP8_DTYPE if self.kv_cache_dtype_str.startswith("fp8") else self.dtype + ) + + # Initialize attn MetadataBuilder + self.builder = self.attn_backend.get_builder_cls()( + kv_cache_spec=self.mla_attn.get_kv_cache_spec(vllm_config), + layer_names=[self.mla_attn.layer_name], + vllm_config=vllm_config, + device=device, + ) + + def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata: + """Initialize attention metadata.""" + # Create common attn metadata + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) + common_attn_metadata = create_common_attn_metadata( + batch_spec, self.block_size, self.device, arange_block_indices=True + ) + + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size + num_blocks = batch_size * max_blocks + + # Fetch the attention backend and kv cache shape and stride order + kv_cache_shape = self.attn_backend.get_kv_cache_shape( + num_blocks, self.block_size, self.num_kv_heads, self.head_size + ) + try: + kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order() + except (AttributeError, NotImplementedError): + kv_cache_stride_order = tuple(range(len(kv_cache_shape))) + + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) + inv_order = [ + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) + ] + + raw_tensor = torch.zeros( + num_blocks * self.block_size * self.num_kv_heads * self.head_size, + dtype=self.kv_cache_dtype, + device=self.device, + ) + raw_tensor = raw_tensor.view(kv_cache_shape) + kv_cache = raw_tensor.permute(*inv_order) + + self.mla_attn.kv_cache = kv_cache + + # Build attn metadata + attn_metadata = self.builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + + return attn_metadata + + def forward( + self, qkv_lora: torch.Tensor, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + qkv_lora = qkv_lora.clone() + q_c, kv_lora = qkv_lora.split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + dim=-1, + ) + q = self.q_b_proj(q_c)[0] + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + + q = q.view(-1, self.num_heads, self.qk_head_dim) + k_pe = k_pe.unsqueeze(1) + + q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( + positions, q[..., self.qk_nope_head_dim :], k_pe + ) + + dummy = torch.ops.vllm.unified_mla_kv_cache_update( + kv_c, + k_pe, + _encode_layer_name(self.layer_name), + self.kv_cache_dtype_str, + self.mla_attn._k_scale, + ) + return q, kv_c, k_pe, dummy + + def ops_in_model_before(self) -> list[torch._ops.OpOverload]: + ops = [ + INDEX_SELECT_OP, + torch.ops.vllm.unified_mla_kv_cache_update.default, + ] + return ops + + def ops_in_model_after(self) -> list[torch._ops.OpOverload]: + return [torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default] + + +MLA_BACKENDS = [AttentionBackendEnum.TRITON_MLA] +if flash_attn_supports_mla(): + MLA_BACKENDS += [AttentionBackendEnum.FLASH_ATTN_MLA] +if is_aiter_found_and_supported(): + MLA_BACKENDS += [AttentionBackendEnum.ROCM_AITER_MLA] + + +@pytest.mark.parametrize("attn_backend", MLA_BACKENDS) +@pytest.mark.parametrize("use_deepseek_scaling_rope", [True]) +@pytest.mark.parametrize("num_heads", [16]) +@pytest.mark.parametrize("qk_nope_head_dim", [128]) +@pytest.mark.parametrize("qk_rope_head_dim", [64]) +@pytest.mark.parametrize("v_head_dim", [128]) +@pytest.mark.parametrize("q_lora_rank", [1536]) +@pytest.mark.parametrize("kv_lora_rank", [512]) +@pytest.mark.parametrize("block_size", [16]) +@pytest.mark.parametrize("is_neox", [True, False]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="MLA RoPE+KVCache+Cat fusion is only supported on CUDA and ROCm.", +) +def test_mla_rope_kvcache_cat_fusion( + attn_backend: AttentionBackendEnum, + use_deepseek_scaling_rope: bool, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + block_size: int, + is_neox: bool, + dtype: torch.dtype, + kv_cache_dtype: str, + monkeypatch: pytest.MonkeyPatch, +): + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + torch.manual_seed(0) + + vllm_config = VllmConfig( + model_config=ModelConfig( + model="deepseek-ai/DeepSeek-V2-Lite", + dtype=dtype, + ), + cache_config=CacheConfig( + block_size=block_size, + cache_dtype=kv_cache_dtype, + ), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + pass_config=PassConfig( + fuse_rope_kvcache_cat_mla=True, + eliminate_noops=True, + ), + ), + ) + + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: + if not torch.distributed.is_initialized(): + from vllm.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.utils.system_utils import update_environment_variables + + update_environment_variables( + { + "RANK": "0", + "LOCAL_RANK": "0", + "WORLD_SIZE": "1", + "MASTER_ADDR": "localhost", + "MASTER_PORT": "54321", + } + ) + init_distributed_environment() + initialize_model_parallel() + + if attn_backend == AttentionBackendEnum.ROCM_AITER_MLA: + m.setenv("VLLM_ROCM_USE_AITER", "1") + rocm_aiter_ops.refresh_env_variables() + + model = MLARoPEKVCacheCatTestModel( + vllm_config=vllm_config, + attn_backend=attn_backend, + use_deepseek_scaling_rope=use_deepseek_scaling_rope, + num_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + is_neox=is_neox, + dtype=dtype, + device=torch.get_default_device(), + ) + + fusion_pass = MLARoPEKVCacheCatFusionPass(vllm_config) + # note: FixFunctionalizationPass is required to correctly lower + # the fused op to its inplace version with auto-functionalization v1. + # Without it, decompose_auto_functionalized calls clone_preserve_strides + # on the non-contiguous q_pe slice directly, and inductor's lowering + # of the resulting as_strided chain incorrectly drops the storage offset. + # auto-functionalization v2 avoids this: it clones the contiguous base + # tensor (_all_bases) and reconstructs the slice as a view, so the + # offset is never passed through as_strided lowering. + passes = [ + NoOpEliminationPass(vllm_config), + fusion_pass, + PostCleanupPass(vllm_config), + FixFunctionalizationPass(vllm_config), + ] + backend = TestBackend(*passes) + + T = 5 + + qkv_lora = torch.randn( + T, + q_lora_rank + kv_lora_rank + qk_rope_head_dim, + dtype=dtype, + ) + pos = torch.arange(T, dtype=torch.long) + + qkv_unfused = qkv_lora.clone() + pos_unfused = pos.clone() + + # Run unfused version + with set_forward_context(None, vllm_config): + forward_context = get_forward_context() + attn_metadata = model.build_attn_metadata(T) + forward_context.slot_mapping = { + model.layer_name: attn_metadata.slot_mapping + } + q_unfused, kv_c_unfused, k_pe_unfused, dummy = model( + qkv_unfused, pos_unfused + ) + attn_layer = forward_context.no_compile_layers[model.layer_name] + kv_cache_unfused = attn_layer.kv_cache.clone() + del dummy + + # Run fused version (compiled) + torch._dynamo.mark_dynamic(qkv_lora, 0) + torch._dynamo.mark_dynamic(pos, 0) + with set_forward_context(None, vllm_config): + model_fused = torch.compile(model, backend=backend) + forward_context = get_forward_context() + attn_metadata = model.build_attn_metadata(T) + forward_context.slot_mapping = { + model.layer_name: attn_metadata.slot_mapping + } + q_fused, kv_c_fused, k_pe_fused, dummy = model_fused(qkv_lora, pos) + attn_layer = forward_context.no_compile_layers[model.layer_name] + kv_cache_fused = attn_layer.kv_cache + del dummy + + assert fusion_pass.matched_count == 1 + + backend.check_before_ops(model.ops_in_model_before()) + backend.check_after_ops(model.ops_in_model_after()) + + if dtype == torch.float16: + ATOL, RTOL = (2e-3, 2e-3) + else: + ATOL, RTOL = (1e-2, 1e-2) + + torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL) + torch.testing.assert_close(kv_c_unfused, kv_c_fused, atol=ATOL, rtol=RTOL) + torch.testing.assert_close(k_pe_unfused, k_pe_fused, atol=ATOL, rtol=RTOL) + # Cannot compare fp8_* directly here, cast to model dtype instead + torch.testing.assert_close( + kv_cache_unfused.view(dtype), + kv_cache_fused.view(dtype), + atol=ATOL, + rtol=RTOL, + ) diff --git a/tests/compile/passes/test_rope_kvcache_fusion.py b/tests/compile/passes/test_rope_kvcache_fusion.py index bab70c12a8..b27adfc46f 100644 --- a/tests/compile/passes/test_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rope_kvcache_fusion.py @@ -34,7 +34,6 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import AttentionSpec INDEX_SELECT_OP = torch.ops.aten.index.Tensor VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update @@ -102,13 +101,8 @@ class QKRoPEKVCacheTestModel(torch.nn.Module): ) # Initialize attn MetadataBuilder - self.builder = self.attn.attn_backend.get_builder_cls()( - kv_cache_spec=AttentionSpec( - block_size=self.block_size, - num_kv_heads=self.num_kv_heads, - head_size=head_size, - dtype=self.kv_cache_dtype, - ), + self.builder = self.attn_backend.get_builder_cls()( + kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config), layer_names=[self.attn.layer_name], vllm_config=vllm_config, device=device, @@ -126,12 +120,11 @@ class QKRoPEKVCacheTestModel(torch.nn.Module): num_blocks = batch_size * max_blocks # Fetch the attention backend and kv cache shape and stride order - attn_backend = self.attn.attn_backend - kv_cache_shape = attn_backend.get_kv_cache_shape( + kv_cache_shape = self.attn_backend.get_kv_cache_shape( num_blocks, self.block_size, self.num_kv_heads, self.head_size ) try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() + kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order() except (AttributeError, NotImplementedError): kv_cache_stride_order = tuple(range(len(kv_cache_shape))) diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index e5130c19c3..d861c83496 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Dynamic, ) from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import ( + DeepseekScalingRotaryEmbedding, +) from vllm.platforms import current_platform RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default @@ -158,6 +161,87 @@ class MatcherRotaryEmbedding(MatcherCustomOp): return result +class MatcherDeepseekScalingRotaryEmbedding(MatcherCustomOp): + def __init__( + self, + is_neox: bool, + head_size: int, + num_heads: int, + num_kv_heads: int, + use_flashinfer: bool = False, + enabled: bool | None = None, + ) -> None: + if enabled is None: + enabled = DeepseekScalingRotaryEmbedding.enabled() + + super().__init__(enabled) + self.is_neox = is_neox + self.head_size = head_size + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.q_size = self.num_heads * self.head_size + self.kv_size = self.num_kv_heads * self.head_size + self.rotary_dim = head_size + self.use_flashinfer = use_flashinfer + + def inputs(self) -> list[torch.Tensor]: + positions = self.empty_int64(5) + query = self.empty(5, self.num_heads, self.head_size) + key = self.empty(5, self.num_kv_heads, self.head_size) + cos_sin_cache = self.empty(4096, self.rotary_dim) + return [positions, query, key, cos_sin_cache] + + def forward_custom( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None, + cos_sin_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if self.use_flashinfer: + torch.ops.vllm.flashinfer_rotary_embedding( + positions, + query, + key, + self.head_size, + cos_sin_cache, + self.is_neox, + ) + return query, key + result: tuple[torch.Tensor, torch.Tensor | None] = ( + DeepseekScalingRotaryEmbedding.forward_static( + positions, + query, + key, + self.head_size, + self.rotary_dim, + cos_sin_cache, + self.is_neox, + ) + ) + return result + + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None, + cos_sin_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + result: tuple[torch.Tensor, torch.Tensor | None] = ( + DeepseekScalingRotaryEmbedding.forward_static( + positions, + query, + key, + self.head_size, + self.rotary_dim, + cos_sin_cache, + self.is_neox, + ) + ) + return result + + class MatcherQuantFP8(MatcherCustomOp): def __init__( self, diff --git a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py new file mode 100644 index 0000000000..5a493149a9 --- /dev/null +++ b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +from torch._higher_order_ops.auto_functionalize import auto_functionalized + +import vllm._custom_ops as ops +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.attention.attention import get_attention_context +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.utils.torch_utils import ( + _USE_LAYERNAME, + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) + +from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement +from .matcher_utils import MatcherDeepseekScalingRotaryEmbedding, MatcherRotaryEmbedding + +logger = init_logger(__name__) + + +def fused_rope_unified_mla_kv_cache_update_impl( + positions: torch.Tensor, + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache_dtype: str, + kv_cache_scale: torch.Tensor, + layer_name: LayerNameType, +) -> torch.Tensor: + layer_name = _resolve_layer_name(layer_name) + attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name) + if layer_slot_mapping is not None: + ops.concat_and_cache_mla_rope_fused( + positions, + q_pe, + k_pe, + kv_c, + cos_sin_cache, + is_neox, + layer_slot_mapping, + kv_cache, + kv_cache_dtype, + kv_cache_scale, + ) + return torch.empty(0, device=kv_c.device, dtype=kv_c.dtype) + + +def fused_rope_unified_mla_kv_cache_update_fake( + positions: torch.Tensor, + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache_dtype: str, + kv_cache_scale: torch.Tensor, + layer_name: LayerNameType, +) -> torch.Tensor: + return torch.empty(0, dtype=kv_c.dtype, device=kv_c.device) + + +direct_register_custom_op( + op_name="fused_rope_unified_mla_kv_cache_update", + op_func=fused_rope_unified_mla_kv_cache_update_impl, + fake_impl=fused_rope_unified_mla_kv_cache_update_fake, + mutates_args=["q_pe", "k_pe"], +) + + +class MLARoPEKVCacheCatPattern(VllmPatternReplacement): + FUSED_OP = torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default + + def __init__( + self, + layer: MLAAttention, + is_neox: bool, + use_flashinfer: bool = False, + use_deepseek_scaling: bool = False, + ) -> None: + self.layer_name = layer.layer_name + self.kv_cache_dtype = layer.kv_cache_dtype + self.num_heads = layer.num_heads + self.num_kv_heads = layer.num_kv_heads + self.kv_lora_rank = layer.kv_lora_rank + self.qk_rope_head_dim = layer.qk_rope_head_dim + self.is_neox = is_neox + self.use_flashinfer = use_flashinfer + self._ln = _encode_layer_name(self.layer_name) + + if use_deepseek_scaling: + self.rope_matcher = MatcherDeepseekScalingRotaryEmbedding( + is_neox=self.is_neox, + head_size=self.qk_rope_head_dim, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + use_flashinfer=self.use_flashinfer, + ) + else: + self.rope_matcher = MatcherRotaryEmbedding( # type: ignore + is_neox=self.is_neox, + head_size=self.qk_rope_head_dim, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + use_flashinfer=self.use_flashinfer, + ) + + def get_inputs(self) -> list[torch.Tensor]: + T = 5 + L = 4096 + q_pe = self.empty_bf16(T, self.num_heads, self.qk_rope_head_dim) + k_pe = self.empty_bf16(T, self.qk_rope_head_dim) + kv_c_normed = self.empty_bf16(T, self.kv_lora_rank) + cos_sin_cache = self.empty_bf16(L, self.qk_rope_head_dim) + positions = self.empty(T, dtype=torch.int64) + k_scale = self.empty(0, dtype=torch.float32) + inputs = [ + q_pe, + k_pe, + kv_c_normed, + positions, + cos_sin_cache, + k_scale, + ] + if _USE_LAYERNAME: + inputs.append(self._ln) + return inputs + + @property + def pattern(self): + _ln = self._ln + + if _USE_LAYERNAME: + + def _pattern_with_ln( + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + k_scale: torch.Tensor, + layer_name: LayerNameType, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + k_pe_unsqueezed = k_pe.unsqueeze(1) + q_pe, k_pe = self.rope_matcher( + positions, q_pe, k_pe_unsqueezed, cos_sin_cache + ) + dummy = torch.ops.vllm.unified_mla_kv_cache_update( + kv_c_normed, k_pe, layer_name, self.kv_cache_dtype, k_scale + ) + return dummy, q_pe, k_pe + + return _pattern_with_ln + + def _pattern( + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + k_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + k_pe_unsqueezed = k_pe.unsqueeze(1) + q_pe, k_pe = self.rope_matcher( + positions, q_pe, k_pe_unsqueezed, cos_sin_cache + ) + dummy = torch.ops.vllm.unified_mla_kv_cache_update( + kv_c_normed, k_pe, _ln, self.kv_cache_dtype, k_scale + ) + return dummy, q_pe, k_pe + + return _pattern + + @property + def replacement(self): + _ln = self._ln + + if _USE_LAYERNAME: + + def _replacement_with_ln( + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + k_scale: torch.Tensor, + layer_name: LayerNameType, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + at = auto_functionalized( + self.FUSED_OP, + positions=positions, + q_pe=q_pe, + k_pe=k_pe, + kv_c=kv_c_normed, + cos_sin_cache=cos_sin_cache, + is_neox=self.is_neox, + kv_cache_dtype=self.kv_cache_dtype, + kv_cache_scale=k_scale, + layer_name=layer_name, + ) + dummy, q_pe, k_pe_squeezed = at + k_pe = k_pe_squeezed.unsqueeze(1) + return dummy, q_pe, k_pe + + return _replacement_with_ln + + def _replacement( + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + k_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + at = auto_functionalized( + self.FUSED_OP, + positions=positions, + q_pe=q_pe, + k_pe=k_pe, + kv_c=kv_c_normed, + cos_sin_cache=cos_sin_cache, + is_neox=self.is_neox, + kv_cache_dtype=self.kv_cache_dtype, + kv_cache_scale=k_scale, + layer_name=_ln, + ) + dummy, q_pe, k_pe_squeezed = at + k_pe = k_pe_squeezed.unsqueeze(1) + return dummy, q_pe, k_pe + + return _replacement + + +class MLARoPEKVCacheCatFusionPass(VllmFusionPatternMatcherPass): + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "mla_rope_kv_cache_fusion_pass") + + attn_layers = get_layers_from_vllm_config(config, MLAAttention) + + for _, layer in attn_layers.items(): + for is_neox in [False, True]: + for use_deepseek_scaling in [False, True]: + if RotaryEmbedding.enabled(): + for use_flashinfer in [False, True]: + self.register( + MLARoPEKVCacheCatPattern( + layer, + is_neox, + use_flashinfer, + use_deepseek_scaling, + ) + ) + else: + self.register( + MLARoPEKVCacheCatPattern( + layer, + is_neox, + use_deepseek_scaling=use_deepseek_scaling, + ) + ) + + if _USE_LAYERNAME: + break + + self.dump_patterns(config, self.pm_pass) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 5d4355a5b2..4cc6bc9e5f 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -33,6 +33,7 @@ if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass + from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass from .fusion.rms_quant_fusion import RMSNormQuantFusionPass from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass @@ -174,6 +175,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc] self.passes += [ScatterSplitReplacementPass(config)] self.passes += [RopeKVCacheFusionPass(config)] + if self.pass_config.fuse_rope_kvcache_cat_mla: + self.passes += [MLARoPEKVCacheCatFusionPass(config)] + if self.pass_config.fuse_attn_quant: self.passes += [AttnQuantFusionPass(config)] self.passes += [MLAAttnQuantFusionPass(config)] diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py index 15eb23e6f9..2887c19ad4 100644 --- a/vllm/compilation/passes/utility/fix_functionalization.py +++ b/vllm/compilation/passes/utility/fix_functionalization.py @@ -181,6 +181,45 @@ class FixFunctionalizationPass(VllmInductorPass): 2: "key", } self.defunctionalize(graph, node, mutated_args=mutated_args) + elif ( + hasattr(torch.ops.vllm, "fused_rope_unified_mla_kv_cache_update") + and at_target + == torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default + ): + # AOTAutograd functionalizes `q[..., nope_dim:] = rope_result` into + # a sequence of aten ops on q: view+slice+copy+slice_scatter. + # Since the fused MLA RoPE op mutates q_pe in-place, we can remove + # the redundant copy and slice_scatter ops during defunctionalization. + getitem_nodes = self.getitem_users(node) + q_pe_out = getitem_nodes[1] + + for user in list(q_pe_out.users): + if is_func(user, torch.ops.aten.copy.default): + copy_temp = user + slice_temp = copy_temp.args[0] + for user in list(copy_temp.users): + if is_func(user, torch.ops.aten.slice_scatter.default): + slice_scatter_temp = user + view_temp = slice_scatter_temp.args[0] + + view_orig = slice_temp.args[0] + slice_scatter_temp.replace_all_uses_with(view_orig) + self._remove(slice_scatter_temp) + self._remove(copy_temp) + self._remove(slice_temp) + self._remove(view_temp) + self._remove(q_pe_out) + + # defunctionalize k_pe manually; self.replace_users_with_mutated_args + # does not support only replacing specific kwargs + k_pe_in = node.kwargs["k_pe"] + k_pe_out = getitem_nodes[2] + k_pe_out.replace_all_uses_with(k_pe_in) + self._remove(k_pe_out) + + self.insert_defunctionalized(graph, node) + self._remove(node) + # only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn elif ( hasattr(torch.ops.vllm, "function_with_mutated_args_and_return") diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 0f02a92681..350621c5d3 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -136,8 +136,10 @@ class PassConfig: """Enable flashinfer allreduce fusion.""" fuse_minimax_qk_norm: bool = None # type: ignore[assignment] """Enable fused allreduce+RMSNorm for MiniMax QK norm.""" - enable_qk_norm_rope_fusion: bool = False + enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment] """Enable fused Q/K RMSNorm + RoPE pass.""" + fuse_rope_kvcache_cat_mla: bool = None # type: ignore[assignment] + """Enable fused MLA KV cache update with RoPE.""" # ROCm/AITER specific fusions fuse_act_padding: bool = None # type: ignore[assignment] @@ -228,6 +230,7 @@ class PassConfig: "fuse_act_padding", "fuse_mla_dual_rms_norm", "fuse_rope_kvcache", + "fuse_rope_kvcache_cat_mla", mode="wrap", ) @classmethod @@ -285,6 +288,12 @@ class PassConfig: "The fusion will be disabled." ) self.fuse_rope_kvcache = False + if self.fuse_rope_kvcache_cat_mla and not current_platform.is_cuda_alike(): + logger.warning_once( + "MLA KV cache update with RoPE fusion enabled but the " + "current platform is not CUDA or ROCm. The fusion will be disabled." + ) + self.fuse_rope_kvcache_cat_mla = False def log_enabled_passes(self) -> None: """ diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53f91e1b65..d220aa6503 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -155,6 +155,15 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool: ) +def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool: + """Enable if use_inductor_graph_partition is enabled.""" + + return ( + cfg.compilation_config.use_inductor_graph_partition + or not cfg.compilation_config.splitting_ops_contain_kv_cache_update() + ) + + def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool: """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss.""" @@ -184,6 +193,7 @@ OPTIMIZATION_LEVEL_00 = { "fuse_act_padding": False, "fuse_mla_dual_rms_norm": False, "fuse_rope_kvcache": False, + "fuse_rope_kvcache_cat_mla": False, }, "cudagraph_mode": CUDAGraphMode.NONE, "use_inductor_graph_partition": False, @@ -204,6 +214,7 @@ OPTIMIZATION_LEVEL_01 = { "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": False, + "fuse_rope_kvcache_cat_mla": False, }, "cudagraph_mode": CUDAGraphMode.PIECEWISE, "use_inductor_graph_partition": False, @@ -226,6 +237,7 @@ OPTIMIZATION_LEVEL_02 = { "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, + "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion, }, "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE, "use_inductor_graph_partition": False, @@ -248,6 +260,7 @@ OPTIMIZATION_LEVEL_03 = { "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, + "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion, }, "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE, "use_inductor_graph_partition": False, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 9cf7ebf5cd..9b4395f41d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -345,6 +345,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", + attn_backend: type[AttentionBackend] | None = None, use_sparse: bool = False, indexer: object | None = None, **extra_impl_args, @@ -374,14 +375,21 @@ class MLAAttention(nn.Module, AttentionLayerBase): self.quant_config = quant_config dtype = torch.get_default_dtype() - self.attn_backend = get_attn_backend( - self.head_size, - dtype, - kv_cache_dtype, - use_mla=True, - use_sparse=use_sparse, - num_heads=self.num_heads, - ) + if attn_backend is not None: + assert attn_backend.is_mla(), ( + f"MLAAttention: attn_backend must be an MLA backend, " + f"got {attn_backend.get_name()} instead" + ) + self.attn_backend = attn_backend + else: + self.attn_backend = get_attn_backend( + self.head_size, + dtype, + kv_cache_dtype, + use_mla=True, + use_sparse=use_sparse, + num_heads=self.num_heads, + ) # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format # Automatically convert fp8 kv-cache format to "fp8_ds_mla" @@ -1008,23 +1016,9 @@ def unified_mla_kv_cache_update( the data dependency between them to ensure torch.compile preserves ordering. """ layer_name = _resolve_layer_name(layer_name) - forward_context = get_forward_context() - attn_layer = forward_context.no_compile_layers[layer_name] - kv_cache = attn_layer.kv_cache - - # This needs to run even when we don't have metadata yet, so that the op - # is correctly captured. - if kv_cache.numel() == 0: - # Can't update an empty KV cache. - return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype) - - slot_mapping = forward_context.slot_mapping - assert isinstance(slot_mapping, dict), ( - f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. " - ) - layer_slot_mapping = slot_mapping.get(layer_name) + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name) if layer_slot_mapping is not None: - attn_layer.impl.do_kv_cache_update( + attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined] kv_c_normed, k_pe, kv_cache, @@ -1113,6 +1107,7 @@ direct_register_custom_op( mutates_args=["output", "output_block_scale"], fake_impl=unified_mla_attention_with_output_fake, dispatch_key=current_platform.dispatch_key, + tags=(torch.Tag.flexible_layout,), ) diff --git a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py index 9a06eedd0f..7362abcc8f 100644 --- a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py @@ -127,29 +127,52 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase): ) -> tuple[torch.Tensor, torch.Tensor | None]: """PyTorch-native implementation equivalent to forward().""" assert key is not None - cos_sin_cache = self._match_cos_sin_cache_dtype(query) - query_rot = query[..., : self.rotary_dim] - key_rot = key[..., : self.rotary_dim] - if self.rotary_dim < self.head_size: - query_pass = query[..., self.rotary_dim :] - key_pass = key[..., self.rotary_dim :] + return self.forward_static( + positions, + query, + key, + self.head_size, + self.rotary_dim, + self.cos_sin_cache, + self.is_neox_style, + offsets, + ) + + @staticmethod + def forward_static( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None, + head_size: int, + rotary_dim: int, + cos_sin_cache: torch.Tensor, + is_neox_style: bool, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """A static implementation of forward().""" + assert key is not None + query_rot = query[..., :rotary_dim] + key_rot = key[..., :rotary_dim] + if rotary_dim < head_size: + query_pass = query[..., rotary_dim:] + key_pass = key[..., rotary_dim:] cos_sin = cos_sin_cache[ torch.add(positions, offsets) if offsets is not None else positions ] cos, sin = cos_sin.chunk(2, dim=-1) - if self.is_neox_style: + if is_neox_style: cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2) else: cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) - rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj + rotate_fn = rotate_neox if is_neox_style else rotate_gptj query_rot = query_rot * cos + rotate_fn(query_rot) * sin key_rot = key_rot * cos + rotate_fn(key_rot) * sin - if self.rotary_dim < self.head_size: + if rotary_dim < head_size: query = torch.cat((query_rot, query_pass), dim=-1) key = torch.cat((key_rot, key_pass), dim=-1) else: diff --git a/vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py b/vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py index ec03fc6533..59ad6359aa 100644 --- a/vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py @@ -195,10 +195,8 @@ class DualChunkRotaryEmbedding(CustomOp): def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass): cos, sin = cos_sin.chunk(2, dim=-1) if self.is_neox_style: - # NOTE(woosuk): Here we assume that the positions tensor has the - # shape [batch_size, seq_len]. - cos = cos.repeat(1, 1, 2).unsqueeze(-2) - sin = sin.repeat(1, 1, 2).unsqueeze(-2) + cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) + sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2) else: cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)