195 lines
7.0 KiB
Python
195 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Attention layer."""
|
|
|
|
from typing import cast
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import vllm.envs as envs
|
|
from vllm.config import CacheConfig, get_current_vllm_config
|
|
from vllm.config.vllm import VllmConfig
|
|
from vllm.model_executor.layers.attention.attention import _init_kv_cache_quant
|
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
|
|
|
# from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
|
|
from vllm.v1.attention.backend import AttentionBackend
|
|
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
|
|
from vllm.v1.kv_cache_interface import KVCacheSpec
|
|
|
|
from vllm_ascend.attention.abstract import DSAAttentionImpl
|
|
from vllm_ascend.attention.dsa_v1 import AscendDSABackend
|
|
from vllm_ascend.core.kv_cache_interface import AscendMLAAttentionSpec
|
|
from vllm_ascend.utils import (
|
|
AscendDeviceType,
|
|
get_ascend_device_type,
|
|
)
|
|
|
|
|
|
def get_dsv4_block_sizes():
|
|
# cache_config.block_size: [mla, swa, c4 state, c128 state], [page_size_padded_t1, page_size_padded_t2]
|
|
_DSV4_BLOCK_SIZES = {
|
|
128: [[128, 128, 8, 32], [16640, 131072]],
|
|
64: [[64, 64, 4, 16], [8320, 65536]],
|
|
32: [[32, 32, 2, 8], [4160, 32768]],
|
|
}
|
|
_DSV4_BLOCK_SIZES_A5 = {
|
|
128: [[128, 128, 8, 16], [16896, 81920]],
|
|
64: [[64, 64, 4, 8], [8448, 40960]],
|
|
32: [[32, 32, 2, 4], [4224, 20480]],
|
|
}
|
|
if get_ascend_device_type() in {AscendDeviceType.A5}:
|
|
return _DSV4_BLOCK_SIZES_A5
|
|
else:
|
|
return _DSV4_BLOCK_SIZES
|
|
|
|
|
|
DSV4_BLOCK_SIZES = get_dsv4_block_sizes()
|
|
|
|
|
|
class DSAAttention(nn.Module, AttentionLayerBase):
|
|
"""Multi-Head Latent Attention layer.
|
|
|
|
This class takes query, and compressed key/value tensors as input.
|
|
The class does the following:
|
|
|
|
1. Store the input key and value tensors in the KV cache.
|
|
2. Perform (multi-head/multi-query/grouped-query) attention.
|
|
3. Return the output tensor.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
n_heads: int,
|
|
scale: float,
|
|
n_local_heads: int,
|
|
q_lora_rank: int,
|
|
o_lora_rank: int,
|
|
head_dim: int,
|
|
rope_head_dim: int | None,
|
|
nope_head_dim: int,
|
|
n_groups: int,
|
|
n_local_groups: int,
|
|
window_size: int,
|
|
compress_ratio: int,
|
|
cache_config: CacheConfig | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
**extra_impl_args,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.n_heads = n_heads
|
|
self.scale = scale
|
|
self.n_local_heads = n_local_heads
|
|
self.q_lora_rank = q_lora_rank
|
|
self.o_lora_rank = o_lora_rank
|
|
self.head_dim = head_dim
|
|
self.rope_head_dim = rope_head_dim
|
|
self.nope_head_dim = nope_head_dim
|
|
self.n_groups = n_groups
|
|
self.n_local_groups = n_local_groups
|
|
self.window_size = window_size
|
|
self.compress_ratio = compress_ratio
|
|
self.layer_name = prefix
|
|
self.head_size = self.head_dim
|
|
self.swa_cache_layer: DeepseekV4SWACache = extra_impl_args.get("swa_cache_layer")
|
|
|
|
assert self.swa_cache_layer is not None
|
|
|
|
if cache_config is not None:
|
|
kv_cache_dtype = cache_config.cache_dtype
|
|
else:
|
|
kv_cache_dtype = "auto"
|
|
|
|
# Initialize KV cache quantization attributes
|
|
_init_kv_cache_quant(self, quant_config, prefix)
|
|
|
|
self.attn_backend = AscendDSABackend
|
|
|
|
# NOTE(zxr): vllm_is_batch_invariant is delete during updating to v0.20.1
|
|
if (
|
|
cache_config is not None
|
|
and cache_config.enable_prefix_caching
|
|
and (self.attn_backend.get_name() == "TRITON_MLA" or self.attn_backend.get_name() == "FLASHINFER")
|
|
):
|
|
cache_config.enable_prefix_caching = False
|
|
|
|
impl_cls = cast(type[DSAAttentionImpl], self.attn_backend.get_impl_cls())
|
|
self.impl = impl_cls(
|
|
dim=self.dim,
|
|
n_heads=self.n_heads,
|
|
scale=self.scale,
|
|
n_local_heads=self.n_local_heads,
|
|
q_lora_rank=self.q_lora_rank,
|
|
o_lora_rank=self.o_lora_rank,
|
|
head_dim=self.head_dim,
|
|
rope_head_dim=self.rope_head_dim,
|
|
nope_head_dim=self.nope_head_dim,
|
|
n_groups=self.n_groups,
|
|
n_local_groups=self.n_local_groups,
|
|
window_size=self.window_size,
|
|
compress_ratio=self.compress_ratio,
|
|
**extra_impl_args,
|
|
)
|
|
|
|
self.use_direct_call = not current_platform.opaque_attention_op()
|
|
|
|
compilation_config = get_current_vllm_config().compilation_config
|
|
if prefix in compilation_config.static_forward_context:
|
|
raise ValueError(f"Duplicate layer name: {prefix}")
|
|
compilation_config.static_forward_context[prefix] = self
|
|
|
|
self.kv_cache = [
|
|
torch.tensor([]) for _ in range(get_current_vllm_config().parallel_config.pipeline_parallel_size)
|
|
]
|
|
self.kv_cache_dtype = kv_cache_dtype
|
|
|
|
self.use_sparse = True
|
|
|
|
# Initialize q/k/v range constants.
|
|
self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
|
|
self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
|
|
self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
|
|
|
|
def forward(
|
|
self,
|
|
q: torch.Tensor,
|
|
kv_c_normed: torch.Tensor,
|
|
k_pe: torch.Tensor,
|
|
output_shape: torch.Size | None = None,
|
|
) -> torch.Tensor:
|
|
return q
|
|
|
|
def process_weights_after_loading(self, act_dtype: torch.dtype):
|
|
if hasattr(self.impl, "process_weights_after_loading"):
|
|
self.impl.process_weights_after_loading(act_dtype)
|
|
|
|
def get_attn_backend(self) -> type[AttentionBackend]:
|
|
return self.attn_backend
|
|
|
|
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
|
|
if self.compress_ratio <= 1: # SWA part. Allocated separately as DeepseekV4SWACache.
|
|
return None
|
|
kv_cache_dtype = kv_cache_dtype_str_to_dtype(self.kv_cache_dtype, vllm_config.model_config)
|
|
if get_ascend_device_type() in {AscendDeviceType.A5}:
|
|
kv_cache_dtype = torch.float8_e4m3fn
|
|
vllm_config.cache_config.cache_dtype = "float8_e4m3fn"
|
|
|
|
cached_head_size = (
|
|
(self.head_size + 128) if get_ascend_device_type() in {AscendDeviceType.A5} else self.head_size
|
|
)
|
|
return AscendMLAAttentionSpec(
|
|
block_size=DSV4_BLOCK_SIZES[vllm_config.cache_config.block_size][0][0],
|
|
num_kv_heads=1,
|
|
head_size=cached_head_size,
|
|
dtype=kv_cache_dtype,
|
|
model_version="deepseek_v4",
|
|
compress_ratio=self.compress_ratio,
|
|
cache_dtype_str=vllm_config.cache_config.cache_dtype,
|
|
)
|