146 lines
5.0 KiB
Python
146 lines
5.0 KiB
Python
import torch
|
|
import vllm.envs as envs_vllm
|
|
from flash_attn_npu_v3 import flash_attn_with_kvcache as _fa3_fn # type: ignore[import-not-found]
|
|
from vllm.v1.attention.backend import AttentionBackend # type: ignore
|
|
|
|
from vllm_ascend.attention.attention_v1 import (
|
|
AscendAttentionBackendImpl,
|
|
AscendAttentionMetadataBuilder,
|
|
)
|
|
|
|
|
|
class AscendFABackend(AttentionBackend):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "CUSTOM" if not envs_vllm.VLLM_USE_V2_MODEL_RUNNER else "FLASH_ATTN"
|
|
|
|
@staticmethod
|
|
def get_impl_cls() -> type["AscendFAImpl"]:
|
|
return AscendFAImpl
|
|
|
|
@staticmethod
|
|
def get_builder_cls() -> type["AscendAttentionMetadataBuilder"]:
|
|
return AscendAttentionMetadataBuilder
|
|
|
|
@staticmethod
|
|
def get_kv_cache_shape(
|
|
num_blocks: int,
|
|
block_size: int,
|
|
num_kv_heads: int,
|
|
head_size: int,
|
|
cache_type: str = "",
|
|
) -> tuple[int, ...]:
|
|
return (2, num_blocks, block_size, num_kv_heads, head_size)
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes() -> list[int]:
|
|
return [128]
|
|
|
|
|
|
class AscendFAImpl(AscendAttentionBackendImpl):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
if self.sliding_window is not None:
|
|
raise ValueError(
|
|
"AscendFAImpl does not support sliding window attention. "
|
|
"Please disable sliding window or use the default FIA backend."
|
|
)
|
|
if not self.vllm_config.model_config.enforce_eager:
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
|
|
cudagraph_mode = self.vllm_config.compilation_config.cudagraph_mode
|
|
if cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY:
|
|
raise ValueError(
|
|
"AscendFAImpl does not support ACL graph capture with "
|
|
"FULL_DECODE_ONLY mode. Please set enforce_eager=True or "
|
|
"not set FULL_DECODE_ONLY or use the default FIA backend."
|
|
)
|
|
|
|
def _flash_attn_with_kvcache(
|
|
self,
|
|
query: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
actual_seq_lengths: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
is_causal: bool,
|
|
max_seq_len: int,
|
|
):
|
|
num_block, block_size, _, _ = self.key_cache.shape # type: ignore
|
|
key_fa_blk = self.key_cache.view( # type: ignore
|
|
num_block, block_size, self.num_kv_heads, self.head_size
|
|
)
|
|
value_fa_blk = self.value_cache.view( # type: ignore
|
|
num_block, block_size, self.num_kv_heads, self.head_size
|
|
)
|
|
|
|
attn_output = _fa3_fn(
|
|
query,
|
|
key_fa_blk,
|
|
value_fa_blk,
|
|
cache_seqlens=seq_lens, # kv sequence length for each individual request (NOT cumulative)
|
|
page_table=block_table, # must match the block table for the corresponding q
|
|
cu_seqlens_q=actual_seq_lengths, # cumulative sequence length for q
|
|
max_seqlen_q=max_seq_len,
|
|
causal=is_causal,
|
|
window_size=[-1, -1],
|
|
rotary_interleaved=False,
|
|
num_splits=1,
|
|
softcap=0.0,
|
|
attention_chunk=0,
|
|
sm_margin=0,
|
|
return_softmax_lse=False,
|
|
)
|
|
|
|
return attn_output
|
|
|
|
def forward_impl(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: tuple[torch.Tensor],
|
|
attn_metadata,
|
|
output: torch.Tensor,
|
|
):
|
|
num_tokens = attn_metadata.actual_seq_lengths_q[-1]
|
|
query = query[:num_tokens]
|
|
|
|
num_decodes = attn_metadata.num_decodes
|
|
num_decode_tokens = attn_metadata.num_decode_tokens
|
|
num_prefills = attn_metadata.num_prefills
|
|
outputs = []
|
|
|
|
if num_decodes > 0:
|
|
outputs.append(
|
|
self._flash_attn_with_kvcache(
|
|
query[:num_decode_tokens],
|
|
attn_metadata.block_tables[:num_decodes, :],
|
|
attn_metadata.query_start_loc[: num_decodes + 1],
|
|
attn_metadata.seq_lens[:num_decodes].npu(),
|
|
False,
|
|
max(attn_metadata.seq_lens[:num_decodes]),
|
|
)
|
|
)
|
|
|
|
if num_prefills > 0:
|
|
outputs.append(
|
|
self._flash_attn_with_kvcache(
|
|
query[num_decode_tokens:],
|
|
attn_metadata.block_tables[num_decode_tokens:, :],
|
|
attn_metadata.query_start_loc[num_decodes:],
|
|
attn_metadata.seq_lens[num_decodes:].npu(),
|
|
True, # enable causal for prefill
|
|
max(attn_metadata.seq_lens[num_decodes:]),
|
|
)
|
|
)
|
|
|
|
if not outputs:
|
|
raise ValueError("No attention output available")
|
|
|
|
attn_output_fa = outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0)
|
|
output[:num_tokens] = attn_output_fa[:num_tokens]
|
|
return output
|