293
vllm_ascend/ops/dsa.py
Normal file
293
vllm_ascend/ops/dsa.py
Normal file
@@ -0,0 +1,293 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user