294 lines
10 KiB
Python
294 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2023 The vLLM team.
|
|
# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
|
# and OPT implementations in this library. It has been modified from its
|
|
# original forms to accommodate minor architectural differences compared
|
|
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torch import nn
|
|
from vllm.config import CacheConfig, get_current_vllm_config
|
|
from vllm.forward_context import ForwardContext, get_forward_context
|
|
from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
from vllm.utils.torch_utils import direct_register_custom_op
|
|
from vllm.v1.attention.backend import AttentionMetadata
|
|
|
|
from vllm_ascend.models.layer.attention.layer import DSAAttention
|
|
from vllm_ascend.utils import (
|
|
AscendDeviceType,
|
|
get_ascend_device_type,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class DSAModules:
|
|
"""Modules used in SFA V2."""
|
|
|
|
wq_a: torch.nn.Module
|
|
q_norm: torch.nn.Module
|
|
q_norm_without_weight: torch.nn.Module
|
|
wq_b: torch.nn.Module
|
|
wkv: torch.nn.Module
|
|
kv_norm: torch.nn.Module
|
|
wo_a: torch.nn.Module
|
|
wo_b: torch.nn.Module
|
|
attn_sink: torch.nn.Module
|
|
indexer: torch.nn.Module | None
|
|
compressor: torch.nn.Module | None
|
|
swa_cache_layer: torch.nn.Module
|
|
topk_indices_buffer: torch.Tensor | None
|
|
indexer_rotary_emb: torch.nn.Module | None = None
|
|
skip_topk: bool = False
|
|
|
|
|
|
class AscendDeepseekSparseAttention(MultiHeadLatentAttentionWrapper):
|
|
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,
|
|
eps: float,
|
|
n_groups: int,
|
|
n_local_groups: int,
|
|
window_size: int,
|
|
compress_ratio: int,
|
|
dsa_modules: DSAModules,
|
|
cache_config: CacheConfig | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
nn.Module.__init__(self)
|
|
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.eps = eps
|
|
self.n_groups = n_groups
|
|
self.n_local_groups = n_local_groups
|
|
self.window_size = window_size
|
|
self.compress_ratio = compress_ratio
|
|
|
|
self.wq_a = dsa_modules.wq_a
|
|
self.q_norm = dsa_modules.q_norm
|
|
self.q_norm_without_weight = dsa_modules.q_norm_without_weight
|
|
self.wq_b = dsa_modules.wq_b
|
|
self.wkv = dsa_modules.wkv
|
|
self.kv_norm = dsa_modules.kv_norm
|
|
self.wo_a = dsa_modules.wo_a
|
|
self.wo_b = dsa_modules.wo_b
|
|
self.attn_sink = dsa_modules.attn_sink
|
|
self.indexer = dsa_modules.indexer
|
|
self.compressor = dsa_modules.compressor
|
|
self.topk_indices_buffer = dsa_modules.topk_indices_buffer
|
|
self.indexer_rotary_emb = dsa_modules.indexer_rotary_emb
|
|
self.skip_topk = dsa_modules.skip_topk
|
|
self.prefix = prefix
|
|
|
|
self.swa_cache_layer = dsa_modules.swa_cache_layer
|
|
|
|
self.dsa_attn = DSAAttention(
|
|
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,
|
|
cache_config=cache_config,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.attn",
|
|
# extra
|
|
wq_a=self.wq_a,
|
|
wq_b=self.wq_b,
|
|
wkv=self.wkv,
|
|
q_norm=self.q_norm,
|
|
q_norm_without_weight=self.q_norm_without_weight,
|
|
kv_norm=self.kv_norm,
|
|
indexer=self.indexer,
|
|
compressor=self.compressor,
|
|
wo_a=self.wo_a,
|
|
wo_b=self.wo_b,
|
|
attn_sink=self.attn_sink,
|
|
eps=self.eps,
|
|
swa_cache_layer=self.swa_cache_layer,
|
|
skip_topk=self.skip_topk,
|
|
topk_indices_buffer=self.topk_indices_buffer,
|
|
)
|
|
|
|
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
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
kv_cache: torch.Tensor | None = None,
|
|
attn_metadata: AttentionMetadata | None = None,
|
|
) -> torch.Tensor:
|
|
need_gather_q_kv = get_forward_context().flash_comm_v1_enabled
|
|
output_shape = hidden_states.shape
|
|
|
|
output = torch.empty(output_shape, dtype=hidden_states.dtype, device=hidden_states.device)
|
|
|
|
# All DSA forward paths (attention + o_proj, including OTP HCCL
|
|
# collectives) run inside the dsa_forward custom op, which is required
|
|
# for ACL graph capture (registered with dispatch_key="PrivateUse1").
|
|
torch.ops.vllm.dsa_forward(hidden_states, need_gather_q_kv, output, self.prefix)
|
|
|
|
output = output.view(-1, output_shape[-1])
|
|
return output
|
|
|
|
|
|
def dsa_forward(
|
|
hidden_states: torch.Tensor,
|
|
need_gather_q_kv: bool,
|
|
output: torch.Tensor,
|
|
layer_name: str,
|
|
) -> None:
|
|
forward_context: ForwardContext = get_forward_context()
|
|
self = forward_context.no_compile_layers[layer_name]
|
|
if forward_context.attn_metadata:
|
|
attn_metadata = filter_metadata(forward_context.attn_metadata, self.prefix)
|
|
else:
|
|
attn_metadata = forward_context.attn_metadata
|
|
|
|
if attn_metadata is None:
|
|
# Profiling run: forward() handles OTP by running _forward_o_proj on a
|
|
# zero input so HCCL collectives are captured by the ACL graph.
|
|
self.dsa_attn.impl.forward(self.dsa_attn.layer_name, hidden_states, None, None, need_gather_q_kv, output)
|
|
return
|
|
|
|
kv_cache = _build_kv_cache(self, forward_context)
|
|
|
|
self.dsa_attn.impl.forward(
|
|
self.dsa_attn.layer_name, hidden_states, kv_cache, attn_metadata, need_gather_q_kv, output
|
|
)
|
|
return
|
|
|
|
|
|
def dsa_forward_fake(
|
|
hidden_states: torch.Tensor,
|
|
need_gather_q_kv: bool,
|
|
output: torch.Tensor,
|
|
layer_name: str,
|
|
) -> None:
|
|
return
|
|
|
|
|
|
direct_register_custom_op(
|
|
op_name="dsa_forward",
|
|
op_func=dsa_forward,
|
|
mutates_args=["output"],
|
|
fake_impl=dsa_forward_fake,
|
|
dispatch_key="PrivateUse1",
|
|
)
|
|
|
|
|
|
def filter_metadata(metadata, prefix):
|
|
# filter using prefix, sort by key for deterministic order
|
|
return [v for k, v in sorted(metadata.items()) if k.startswith(prefix)]
|
|
|
|
|
|
def _build_kv_cache(self, forward_context):
|
|
"""Construct the 6-tuple KV cache used by impl.forward()."""
|
|
compress_kv_cache = None
|
|
swa_kv_cache = self.swa_cache_layer.kv_cache
|
|
state_cache = None
|
|
indexer_state_cache = None
|
|
indexer_k_cache = None
|
|
indexer_scale_cache = None
|
|
indexer_full_cache = None
|
|
|
|
if self.compress_ratio > 1:
|
|
state_cache = self.compressor.state_cache.kv_cache
|
|
compress_kv_cache = self.dsa_attn.kv_cache
|
|
virtual_engine = getattr(forward_context, "virtual_engine", None)
|
|
if virtual_engine is not None and isinstance(compress_kv_cache, (list, tuple)):
|
|
compress_kv_cache = compress_kv_cache[virtual_engine]
|
|
if self.compress_ratio == 4:
|
|
indexer_state_cache = self.indexer.compressor.state_cache.kv_cache
|
|
if get_ascend_device_type() in {AscendDeviceType.A5}:
|
|
indexer_k_cache, indexer_scale_cache, indexer_full_cache = (
|
|
self.indexer.k_cache.kv_cache[0][0],
|
|
self.indexer.k_cache.kv_cache[0][1],
|
|
self.indexer.k_cache.kv_cache[0][2],
|
|
)
|
|
else:
|
|
indexer_k_cache, indexer_scale_cache = (
|
|
self.indexer.k_cache.kv_cache[0][0],
|
|
self.indexer.k_cache.kv_cache[0][1],
|
|
)
|
|
|
|
if get_ascend_device_type() in {AscendDeviceType.A5}:
|
|
kv_cache = tuple(
|
|
[
|
|
unfold_kvcache(cache)
|
|
for cache in (
|
|
compress_kv_cache,
|
|
swa_kv_cache,
|
|
state_cache,
|
|
indexer_state_cache,
|
|
indexer_k_cache,
|
|
indexer_scale_cache,
|
|
indexer_full_cache,
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
kv_cache = tuple(
|
|
[
|
|
unfold_kvcache(cache)
|
|
for cache in (
|
|
compress_kv_cache,
|
|
swa_kv_cache,
|
|
state_cache,
|
|
indexer_state_cache,
|
|
indexer_k_cache,
|
|
indexer_scale_cache,
|
|
)
|
|
]
|
|
)
|
|
return kv_cache
|
|
|
|
|
|
def unfold_kvcache(kvcache):
|
|
while isinstance(kvcache, list) and len(kvcache) == 1:
|
|
kvcache = kvcache[0]
|
|
return kvcache
|