obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-21 05:00:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Rohan138 <[email protected]>
Signed-off-by: Rohan Potdar <[email protected]>
Co-authored-by: ElizaWszola <[email protected]>
This commit is contained in:
Rohan Potdar
2026-05-11 14:10:50 +00:00
committed by GitHub
co-authored by ElizaWszola
parent 8415bf2cdb
commit a51376b3f0
12 changed files with 966 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cache_kernels_fused.cu
+75 -60
View File
@@ -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 <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
typename cache_t, Fp8KVCacheDataType kv_dt>
template <typename qk_t, typename cos_sin_t, bool IS_NEOX,
typename raw_kv_scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
__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<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
qk_t sin = static_cast<qk_t>(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<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
qk_t sin = static_cast<qk_t>(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<qk_t, true, RAW_KV_T, \
CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} else { \
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T, \
CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} \
}); \
#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> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} else { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} \
}); \
}); \
} 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);
tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py
+413
View File
@@ -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,
)
tests/compile/passes/test_rope_kvcache_fusion.py
+4 -11
View File
@@ -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)))
vllm/compilation/passes/fusion/matcher_utils.py
+84
View File
@@ -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,
vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
+271
View File
@@ -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)
vllm/compilation/passes/pass_manager.py
+4
View File
@@ -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)]
vllm/compilation/passes/utility/fix_functionalization.py
+39
View File
@@ -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")
vllm/config/compilation.py
+10 -1
View File
@@ -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:
"""
vllm/config/vllm.py
+13
View File
@@ -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,
vllm/model_executor/layers/attention/mla_attention.py
+19 -24
View File
@@ -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,),
)
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+32 -9
View File
@@ -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:
vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py
+2 -4
View File
@@ -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)