1985 lines
86 KiB
Python
1985 lines
86 KiB
Python
#
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
|
|
import torch
|
|
import torch_npu
|
|
import vllm.envs as envs_vllm
|
|
from vllm.config import VllmConfig, get_current_vllm_config
|
|
from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
|
|
from vllm.utils.math_utils import cdiv
|
|
from vllm.v1.attention.backend import ( # type: ignore
|
|
AttentionBackend,
|
|
AttentionCGSupport,
|
|
AttentionImpl,
|
|
AttentionLayer,
|
|
AttentionMetadataBuilder,
|
|
AttentionType,
|
|
)
|
|
from vllm.v1.attention.backends.registry import ( # type: ignore
|
|
AttentionBackendEnum,
|
|
register_backend,
|
|
)
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.kv_cache_interface import AttentionSpec, CrossAttentionSpec
|
|
|
|
from vllm_ascend.ascend_forward_context import _EXTRA_CTX
|
|
from vllm_ascend.attention.attention_mask import AttentionMaskBuilder
|
|
from vllm_ascend.attention.context_parallel.common_cp import AscendMetadataForDecode, AscendMetadataForPrefill
|
|
from vllm_ascend.attention.kvcomp_attn.attention_utils import (
|
|
get_kvcomp_decode_params,
|
|
is_enable_hamming_sparse,
|
|
reshape_and_cache_kvcomp,
|
|
)
|
|
from vllm_ascend.attention.utils import (
|
|
AscendCommonAttentionMetadata,
|
|
cache_graph_workspace,
|
|
enable_cp,
|
|
needs_layer_aware_fia_graph_replay,
|
|
notify_kv_cache_written,
|
|
split_decodes_and_prefills,
|
|
using_paged_attention,
|
|
)
|
|
from vllm_ascend.compilation.acl_graph import (
|
|
get_draft_graph_params,
|
|
get_draft_graph_prefill_params,
|
|
get_graph_params,
|
|
update_draft_graph_params_workspaces,
|
|
update_graph_params_workspaces,
|
|
)
|
|
from vllm_ascend.device.device_op import DeviceOperator
|
|
from vllm_ascend.memcache_comm_fence import record_attention_compute_start
|
|
from vllm_ascend.ops.flashcomm2_oshard_manager import flashcomm2_oshard_manager
|
|
from vllm_ascend.utils import weak_ref_tensors
|
|
from vllm_ascend.worker.kvcomp_utils import KVCompMetaData
|
|
|
|
# default max value of sliding window size
|
|
SWA_INT_MAX = 2147483647
|
|
_ATTN_KEYS_BUFFER = None
|
|
|
|
|
|
@register_backend(AttentionBackendEnum.CUSTOM, "ASCEND")
|
|
class AscendAttentionBackend(AttentionBackend):
|
|
accept_output_buffer: bool = True
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
# HACK(Ronald1995): vllm `initialize_kv_cache` method in model runner v2 make
|
|
# attention name assertion, we just set name to FLASH_ATTN to avoid assertion error.
|
|
# rectify this when vllm disable the assertion.
|
|
return "CUSTOM" if not envs_vllm.VLLM_USE_V2_MODEL_RUNNER else "FLASH_ATTN"
|
|
|
|
@staticmethod
|
|
def get_impl_cls() -> type["AscendAttentionBackendImpl"]:
|
|
if enable_cp():
|
|
from vllm_ascend.attention.context_parallel.attention_cp import AscendAttentionCPImpl
|
|
|
|
return AscendAttentionCPImpl
|
|
return AscendAttentionBackendImpl
|
|
|
|
@staticmethod
|
|
def get_builder_cls() -> type["AscendAttentionMetadataBuilder"]:
|
|
if enable_cp():
|
|
from vllm_ascend.attention.context_parallel.attention_cp import AscendAttentionCPMetadataBuilder
|
|
|
|
return AscendAttentionCPMetadataBuilder
|
|
return AscendAttentionMetadataBuilder
|
|
|
|
@staticmethod
|
|
def get_kv_cache_shape(
|
|
num_blocks: int,
|
|
block_size: int,
|
|
num_kv_heads: int,
|
|
head_size: int,
|
|
cache_dtype_str: str = "",
|
|
) -> tuple[int, ...]:
|
|
return (2, num_blocks, block_size, num_kv_heads, head_size)
|
|
|
|
@staticmethod
|
|
def swap_blocks(
|
|
src_kv_cache: list[torch.Tensor],
|
|
dst_kv_cache: list[torch.Tensor],
|
|
src_to_dst: torch.Tensor,
|
|
) -> None:
|
|
src_key_cache, src_value_cache = src_kv_cache[0], src_kv_cache[1]
|
|
dst_key_cache, dst_value_cache = dst_kv_cache[0], dst_kv_cache[1]
|
|
src_indices = src_to_dst[:, 0]
|
|
dst_indices = src_to_dst[:, 1]
|
|
|
|
dst_key_cache[dst_indices] = src_key_cache[src_indices].to(dst_key_cache.device)
|
|
dst_value_cache[dst_indices] = src_value_cache[src_indices].to(dst_key_cache.device)
|
|
|
|
@staticmethod
|
|
def copy_blocks(
|
|
kv_caches: list[torch.Tensor],
|
|
src_to_dists: torch.Tensor,
|
|
) -> None:
|
|
src_indices = src_to_dists[:, 0]
|
|
dst_indices = src_to_dists[:, 1]
|
|
|
|
for kv_cache in kv_caches:
|
|
key_caches = kv_cache[0]
|
|
value_caches = kv_cache[1]
|
|
key_caches[dst_indices] = key_caches[src_indices]
|
|
value_caches[dst_indices] = value_caches[src_indices]
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes() -> list[int]:
|
|
return [128]
|
|
|
|
|
|
class AscendAttentionState(Enum):
|
|
PrefillNoCache = 0
|
|
PrefillCacheHit = 1
|
|
DecodeOnly = 2
|
|
ChunkedPrefill = 3
|
|
SpecDecoding = 4
|
|
|
|
|
|
@dataclass
|
|
class AscendMetadata:
|
|
"""
|
|
Per-layer attention metadata for Ascend FlashAttention backend.
|
|
|
|
Contains attention masks, token counts, sequence lengths and KV cache
|
|
related properties for attention computation.
|
|
"""
|
|
|
|
# **************************** Basic Properties ************************** #
|
|
attn_mask: torch.Tensor | None = None
|
|
# Current state of this attention run.
|
|
attn_state: AscendAttentionState = AscendAttentionState.ChunkedPrefill
|
|
|
|
# Number of tokens excluding padding.
|
|
num_actual_tokens_pcp_padded: int = 0
|
|
num_actual_tokens: int = 0
|
|
num_decode_tokens: int = 0
|
|
num_prefills: int = 0
|
|
num_decodes: int = 0
|
|
|
|
# The sequence length per sequence. Sequence length means the computed
|
|
# tokens + new tokens (is None if it is a decoding).
|
|
# (batch_size,)
|
|
# TODO(Angazenn): The following parameters are quite redundant and
|
|
# contains similar information (such as seq_lens seq_lens_list). We
|
|
# should simplified these parameters once attention schema in vLLM-Ascend
|
|
# is unified.
|
|
seq_lens: torch.Tensor = None
|
|
seq_lens_cpu: torch.Tensor = None
|
|
seq_lens_list: list[int] = None # type: ignore
|
|
actual_seq_lengths_q: list[int] = None # type: ignore
|
|
|
|
query_start_loc: torch.Tensor = None
|
|
# Maximum query length in the batch (None for decoding).
|
|
max_query_len: int | None = None
|
|
|
|
# ********************** KV Cache Related Properties ********************* #
|
|
# Block addresses per sequence (Seq id -> list of physical block).
|
|
# (batch_size, max_blocks_per_seq)
|
|
block_tables: torch.Tensor = None
|
|
|
|
# The indices of the token slots that input tokens will be stored into.
|
|
# E.g., if `slot_mapping` is [35, 2, 17] and the block size is 16, the
|
|
# three tokens are stored in the 3rd slot in block 2, 2nd slot in block 0,
|
|
# and 1st slot in block 1, respectively.
|
|
# (num_tokens,)
|
|
slot_mapping: torch.Tensor = None
|
|
# pcp
|
|
prefill: AscendMetadataForPrefill | None = None
|
|
# dcp
|
|
decode_meta: AscendMetadataForDecode | None = None
|
|
|
|
causal: bool = True
|
|
# runner_type in model_config.
|
|
model_runner_type: str = ""
|
|
# prefill reshape_and_cache event
|
|
reshape_cache_event: torch.npu.Event = None
|
|
|
|
kvcomp_metadata: KVCompMetaData | None = None
|
|
|
|
|
|
class AscendAttentionMetadataBuilder(AttentionMetadataBuilder[AscendMetadata]):
|
|
"""
|
|
Builder for constructing AscendMetadata from CommonAttentionMetadata.
|
|
|
|
Handles attention mask generation and metadata preparation for
|
|
Ascend FlashAttention backend.
|
|
"""
|
|
|
|
# Does this backend/builder reorder the batch?
|
|
# If not, set this to None. Otherwise set it to the query
|
|
# length that will be pulled into the front of the batch.
|
|
reorder_batch_threshold: int = 1
|
|
|
|
def __init__(
|
|
self,
|
|
kv_cache_spec: AttentionSpec,
|
|
layer_names: list[str],
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
):
|
|
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
|
self.vllm_config = vllm_config
|
|
self.model_config = vllm_config.model_config
|
|
self.compilation_config = vllm_config.compilation_config
|
|
self.device = device
|
|
self.max_num_blocks_per_req = cdiv(
|
|
self.model_config.max_model_len, AscendAttentionBackend.get_supported_kernel_block_sizes()[0]
|
|
)
|
|
|
|
self.speculative_config = vllm_config.speculative_config
|
|
self.decode_threshold = 1
|
|
if self.speculative_config:
|
|
spec_token_num = self.speculative_config.num_speculative_tokens
|
|
self.decode_threshold += spec_token_num
|
|
assert self.decode_threshold <= 16, (
|
|
f"decode_threshold exceeded \
|
|
npu_fused_infer_attention_score TND layout's limit of 16, \
|
|
got {self.decode_threshold}"
|
|
)
|
|
|
|
self.reorder_batch_threshold = self.decode_threshold
|
|
|
|
scheduler_config = vllm_config.scheduler_config
|
|
self.chunked_prefill_enabled = scheduler_config.enable_chunked_prefill
|
|
self.attn_mask_builder = AttentionMaskBuilder(self.device)
|
|
|
|
@classmethod
|
|
def get_cudagraph_support(
|
|
cls: type["AscendAttentionMetadataBuilder"],
|
|
vllm_config: VllmConfig,
|
|
kv_cache_spec: AttentionSpec,
|
|
) -> AttentionCGSupport:
|
|
# Explicit override in case the underlying builder specialized this getter.
|
|
# @override omitted only because of mypy limitation due to type variable.
|
|
return AttentionCGSupport.ALWAYS
|
|
|
|
def reorder_batch(self, input_batch, scheduler_output: "SchedulerOutput") -> bool:
|
|
return False
|
|
|
|
def build(
|
|
self,
|
|
common_prefix_len: int,
|
|
common_attn_metadata: AscendCommonAttentionMetadata,
|
|
fast_build: bool = False,
|
|
) -> AscendMetadata:
|
|
num_reqs = common_attn_metadata.num_reqs
|
|
num_actual_tokens = common_attn_metadata.num_actual_tokens
|
|
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
|
|
|
|
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills(
|
|
common_attn_metadata, decode_threshold=self.decode_threshold
|
|
)
|
|
|
|
block_table = common_attn_metadata.block_table_tensor
|
|
# Prefer _seq_lens_cpu (always available, updated during draft
|
|
# iterations) over seq_lens_cpu (None in async spec decode mode).
|
|
if common_attn_metadata._seq_lens_cpu is not None:
|
|
seq_lens = common_attn_metadata._seq_lens_cpu[:num_reqs]
|
|
elif common_attn_metadata.seq_lens_cpu is not None:
|
|
seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs]
|
|
else:
|
|
seq_lens = common_attn_metadata.seq_lens[:num_reqs].to("cpu")
|
|
|
|
slot_mapping = common_attn_metadata.slot_mapping[:num_actual_tokens]
|
|
# this slot_mapping override doesn't work since vllm will override it again. We should fix it vllm.
|
|
# see: https://github.com/vllm-project/vllm/blob/ce88756b967c2c5006746a424c15dd59a284ed8c/vllm/model_executor/layers/attention/cross_attention.py#L117
|
|
if isinstance(self.kv_cache_spec, CrossAttentionSpec):
|
|
seq_lens = common_attn_metadata.seq_lens
|
|
slot_mapping = common_attn_metadata.slot_mapping.to(torch.int32)
|
|
elif self.speculative_config and self.speculative_config.parallel_drafting:
|
|
seq_lens = common_attn_metadata.seq_lens
|
|
|
|
attn_state = common_attn_metadata.attn_state
|
|
|
|
# Get attn_mask from singleton AttentionMaskBuilder
|
|
attn_mask = self.attn_mask_builder.get_attention_mask(common_attn_metadata.causal, self.model_config)
|
|
|
|
# TODO: Yet another unnecessary H2D while we already have a query_start_loc on device
|
|
query_start_loc = query_start_loc_cpu.pin_memory().to(self.device, non_blocking=True)
|
|
|
|
actual_seq_lengths_q = query_start_loc_cpu[1:].tolist()
|
|
seq_lens_list = seq_lens.tolist()
|
|
# flashcomm1/SP (or cudagraph) padding makes the model runner insert a
|
|
# dummy padding request into query_start_loc to satisfy the FIA TND-layout
|
|
# constraint (sum of q lengths == hidden_states.shape[0]), bumping the
|
|
# q-derived batchSize by one. The query_start_loc buffer is sized
|
|
# `max_num_reqs + 2` to hold it, but the seq_lens and block_table buffers
|
|
# are only `max_num_reqs`, so when the batch is full the padded request
|
|
# overflows and `[:num_reqs_padded]` silently truncates them. FIA then
|
|
# fails (error 561002) checking, in order, the `actualSeqLengthsKv` length
|
|
# and then the block_table row count against batchSize. Pad them to match:
|
|
# the dummy request points at block 0, and its output is harmless because:
|
|
# (1) read side: the attention output for padding tokens is trimmed by
|
|
# `hidden_states = hidden_states[:-pad_size, :]` downstream;
|
|
# (2) write side: reshape_and_cache slices key/value/slot_mapping to
|
|
# `[:num_actual_tokens]` (unpadded count), so the dummy request
|
|
# never writes to KV cache.
|
|
# So any valid positive KV length / zero block row is fine. Pad both
|
|
# seq_lens_list and the seq_lens tensor: full_graph_fia_v2 passes the
|
|
# seq_lens tensor (not seq_lens_list) as actual_seq_kvlen during graph
|
|
# capture, and _get_fia_params derives the PrefillCacheHit batch size from
|
|
# seq_lens.shape[0], so the tensor has to carry the dummy request too.
|
|
num_reqs_fia = len(actual_seq_lengths_q)
|
|
if len(seq_lens_list) < num_reqs_fia:
|
|
padding_len = num_reqs_fia - len(seq_lens_list)
|
|
seq_lens_list = seq_lens_list + [1] * padding_len
|
|
seq_lens = torch.cat([seq_lens, seq_lens.new_ones(padding_len)])
|
|
if block_table is not None and block_table.shape[0] < num_reqs_fia:
|
|
block_table = torch.cat(
|
|
[
|
|
block_table,
|
|
block_table.new_zeros((num_reqs_fia - block_table.shape[0], block_table.shape[1])),
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
attn_metadata = AscendMetadata(
|
|
num_actual_tokens=num_actual_tokens,
|
|
num_decode_tokens=num_decode_tokens,
|
|
block_tables=block_table,
|
|
query_start_loc=query_start_loc,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens,
|
|
seq_lens_list=seq_lens_list,
|
|
max_query_len=common_attn_metadata.max_query_len,
|
|
actual_seq_lengths_q=actual_seq_lengths_q,
|
|
slot_mapping=slot_mapping,
|
|
attn_mask=attn_mask,
|
|
attn_state=attn_state,
|
|
num_prefills=num_prefills,
|
|
num_decodes=num_decodes,
|
|
causal=common_attn_metadata.causal,
|
|
model_runner_type=self.model_config.runner_type,
|
|
kvcomp_metadata=common_attn_metadata.kvcomp_metadata,
|
|
)
|
|
return attn_metadata
|
|
|
|
def build_for_graph_capture(
|
|
self,
|
|
common_attn_metadata: AscendCommonAttentionMetadata,
|
|
attn_state: AscendAttentionState = AscendAttentionState.DecodeOnly,
|
|
):
|
|
if attn_state in (
|
|
AscendAttentionState.DecodeOnly,
|
|
AscendAttentionState.ChunkedPrefill,
|
|
AscendAttentionState.SpecDecoding,
|
|
):
|
|
attn_metadata = self.build(
|
|
common_prefix_len=0,
|
|
common_attn_metadata=common_attn_metadata,
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
"Currently we only support building dummy metadata for DecodeOnly and ChunkedPrefill state"
|
|
)
|
|
|
|
attn_metadata.attn_state = attn_state
|
|
return attn_metadata
|
|
|
|
|
|
class AscendAttentionBackendImpl(AttentionImpl):
|
|
def __init__(
|
|
self,
|
|
num_heads: int,
|
|
head_size: int,
|
|
scale: float,
|
|
num_kv_heads: int,
|
|
alibi_slopes: list[float] | None,
|
|
sliding_window: int | None,
|
|
kv_cache_dtype: str,
|
|
logits_soft_cap: float | None,
|
|
attn_type: str,
|
|
kv_sharing_target_layer_name: str | None,
|
|
sinks: torch.Tensor = None,
|
|
**kwargs,
|
|
) -> None:
|
|
self.vllm_config = get_current_vllm_config()
|
|
self.num_heads = num_heads
|
|
self.head_size = head_size
|
|
self.scale = float(scale)
|
|
self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
|
|
self.hidden_size = self.num_heads * self.head_size
|
|
self.kv_cache_dtype = kv_cache_dtype
|
|
self.sliding_window = sliding_window
|
|
if alibi_slopes is not None:
|
|
alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32, device="npu")
|
|
self.alibi_slopes = alibi_slopes
|
|
self.attn_type = attn_type
|
|
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
|
|
|
|
assert self.num_heads % self.num_kv_heads == 0
|
|
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
|
|
self.key_cache = None
|
|
self.value_cache = None
|
|
self.is_kv_producer = (
|
|
self.vllm_config.kv_transfer_config is not None and self.vllm_config.kv_transfer_config.is_kv_producer
|
|
)
|
|
self.enable_c8_quant = self.vllm_config.quant_config is not None and getattr(
|
|
self.vllm_config.quant_config, "enable_c8_quant", False
|
|
)
|
|
self._use_layer_aware_fia_graph_replay = needs_layer_aware_fia_graph_replay()
|
|
self._use_max_workspace_for_fia_graph = self._use_layer_aware_fia_graph_replay
|
|
self.sinks = sinks
|
|
self.layerIndex = 0
|
|
self.enable_hamming_sparse = is_enable_hamming_sparse()
|
|
# Some mixed-attention models cannot rely on the iteration order of
|
|
# attn_metadata during graph replay. Record the captured layer name only
|
|
# for that path.
|
|
self._layer_name: str | None = None
|
|
|
|
def _graph_metadata_layer_name(self, layer: AttentionLayer | None = None) -> str | None:
|
|
layer_name = layer.layer_name if layer is not None else self._layer_name
|
|
# KV-sharing layers replay with the target layer's metadata instead of
|
|
# their own module name, matching vLLM's shared KV-cache ownership.
|
|
return self.kv_sharing_target_layer_name or layer_name
|
|
|
|
@staticmethod
|
|
def update_graph_params(
|
|
update_stream,
|
|
forward_context,
|
|
num_tokens,
|
|
vllm_config,
|
|
speculative_config=None,
|
|
num_dcp_pcp_tokens=None,
|
|
draft_attn_metadatas=None,
|
|
):
|
|
use_layer_aware_replay = needs_layer_aware_fia_graph_replay()
|
|
if using_paged_attention(num_tokens, vllm_config):
|
|
# Paged Attention update logic
|
|
if _EXTRA_CTX.is_draft_model:
|
|
if _EXTRA_CTX.is_draft_model_prefill:
|
|
graph_params = get_draft_graph_prefill_params()
|
|
else:
|
|
graph_params = get_draft_graph_params()
|
|
else:
|
|
graph_params = get_graph_params()
|
|
with torch.npu.stream(update_stream):
|
|
for key, param, handle, event in zip(
|
|
forward_context.attn_metadata,
|
|
graph_params.attn_params[num_tokens],
|
|
graph_params.handles[num_tokens],
|
|
graph_params.events[num_tokens],
|
|
):
|
|
(
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
num_kv_heads,
|
|
num_heads,
|
|
scale,
|
|
block_table,
|
|
seq_lens,
|
|
output,
|
|
) = param
|
|
seq_lens = forward_context.attn_metadata[key].seq_lens
|
|
|
|
workspace = torch_npu._npu_paged_attention_get_workspace(
|
|
query=query,
|
|
key_cache=key_cache,
|
|
value_cache=value_cache,
|
|
num_kv_heads=num_kv_heads,
|
|
num_heads=num_heads,
|
|
scale_value=scale,
|
|
block_table=block_table,
|
|
context_lens=seq_lens,
|
|
out=output,
|
|
)
|
|
torch.npu.graph_task_update_begin(update_stream, handle)
|
|
torch_npu._npu_paged_attention(
|
|
query=query,
|
|
key_cache=key_cache,
|
|
value_cache=value_cache,
|
|
num_kv_heads=num_kv_heads,
|
|
num_heads=num_heads,
|
|
scale_value=scale,
|
|
block_table=block_table,
|
|
context_lens=seq_lens,
|
|
out=output,
|
|
workspace=workspace,
|
|
)
|
|
torch.npu.graph_task_update_end(update_stream)
|
|
event.record(update_stream)
|
|
elif _EXTRA_CTX.sinks:
|
|
# FIA update logic
|
|
if _EXTRA_CTX.is_draft_model:
|
|
graph_params = get_draft_graph_params()
|
|
attn_metadata = draft_attn_metadatas
|
|
draft_attn_key_steps = [
|
|
(draft_step, key)
|
|
for draft_step, per_step_metadata in enumerate(attn_metadata)
|
|
for key in per_step_metadata
|
|
]
|
|
attn_keys = [key for _, key in draft_attn_key_steps]
|
|
else:
|
|
graph_params = get_graph_params()
|
|
attn_metadata = forward_context.attn_metadata
|
|
attn_keys = list(attn_metadata.keys())
|
|
# For Qwen3-next, since the kv_cache_config has already categorized
|
|
# linear_attn and self_attn, the attn_metadata is first arranged with
|
|
# self_attn followed by linear_attn. Therefore, using zip directly
|
|
# filters out the update operations for linear_attn.
|
|
# TODO: We use a new variable `attn_keys` to ensure the loop count is
|
|
# correct after get by `zip` because of the new structure of the attn_metadata
|
|
# when running with the merged full eagle-graph. Should check it with Qwen3-next.
|
|
num_layers = len(attn_keys)
|
|
if num_layers == 0:
|
|
return
|
|
captured_attn_params = graph_params.attn_params[num_tokens]
|
|
handles = graph_params.handles[num_tokens]
|
|
events = graph_params.events[num_tokens]
|
|
graph_param_count = len(captured_attn_params)
|
|
workspace = graph_params.workspaces.get(num_tokens)
|
|
if _EXTRA_CTX.is_draft_model:
|
|
if graph_param_count > len(draft_attn_key_steps):
|
|
repeat_count = cdiv(graph_param_count, len(draft_attn_key_steps))
|
|
draft_attn_key_steps = (draft_attn_key_steps * repeat_count)[:graph_param_count]
|
|
else:
|
|
draft_attn_key_steps = draft_attn_key_steps[:graph_param_count]
|
|
attn_keys = [key for _, key in draft_attn_key_steps]
|
|
elif use_layer_aware_replay:
|
|
# One graph size can contain captured FIA ops from all layers.
|
|
# Repeat attn keys to match the captured op count, then use the
|
|
# stored layer name in each op param to resolve the exact
|
|
# metadata entry during replay.
|
|
attn_keys = [attn_keys[index % num_layers] for index in range(graph_param_count)]
|
|
attn_count = 0
|
|
with torch.npu.stream(update_stream):
|
|
for key, param, handle, event in zip(
|
|
attn_keys,
|
|
captured_attn_params,
|
|
handles,
|
|
events,
|
|
):
|
|
(
|
|
query,
|
|
key_cache,
|
|
value,
|
|
block_tables,
|
|
attn_mask,
|
|
block_size,
|
|
seq_lens,
|
|
num_kv_heads,
|
|
num_heads,
|
|
scale,
|
|
sliding_window,
|
|
sinks,
|
|
attn_output,
|
|
softmax_lse,
|
|
layer_name,
|
|
) = param
|
|
|
|
if _EXTRA_CTX.is_draft_model:
|
|
draft_step, key = draft_attn_key_steps[attn_count]
|
|
seq_lens = attn_metadata[draft_step][key].seq_lens_list
|
|
actual_seq_lengths_q = attn_metadata[draft_step][key].actual_seq_lengths_q
|
|
attn_count = attn_count + 1
|
|
else:
|
|
metadata_key = layer_name if layer_name is not None and layer_name in attn_metadata else key
|
|
seq_lens = attn_metadata[metadata_key].seq_lens_list
|
|
actual_seq_lengths_q = attn_metadata[metadata_key].actual_seq_lengths_q
|
|
|
|
torch.npu.graph_task_update_begin(update_stream, handle)
|
|
torch_npu.npu_fused_infer_attention_score_v2.out(
|
|
query=query,
|
|
key=key_cache,
|
|
value=value,
|
|
block_table=block_tables,
|
|
atten_mask=attn_mask,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_qlen=actual_seq_lengths_q,
|
|
actual_seq_kvlen=seq_lens,
|
|
num_key_value_heads=num_kv_heads,
|
|
num_query_heads=num_heads,
|
|
sparse_mode=4 if sliding_window is not None else 3,
|
|
pre_tokens=sliding_window if sliding_window is not None else SWA_INT_MAX,
|
|
next_tokens=0,
|
|
softmax_scale=scale,
|
|
learnable_sink=sinks,
|
|
workspace=workspace,
|
|
out=[attn_output, softmax_lse],
|
|
)
|
|
torch.npu.graph_task_update_end(update_stream)
|
|
event.record(update_stream)
|
|
else:
|
|
# FIA update logic
|
|
if _EXTRA_CTX.is_draft_model:
|
|
if _EXTRA_CTX.is_draft_model_prefill:
|
|
graph_params = get_draft_graph_prefill_params()
|
|
else:
|
|
graph_params = get_draft_graph_params()
|
|
attn_metadata = draft_attn_metadatas
|
|
draft_attn_key_steps = [
|
|
(draft_step, key)
|
|
for draft_step, per_step_metadata in enumerate(attn_metadata)
|
|
for key in per_step_metadata
|
|
]
|
|
attn_keys = [key for _, key in draft_attn_key_steps]
|
|
else:
|
|
graph_params = get_graph_params()
|
|
attn_metadata = forward_context.attn_metadata
|
|
attn_keys = list(attn_metadata.keys())
|
|
if not use_layer_aware_replay:
|
|
# In some speculative methods (such as DFlash), the order of
|
|
# attn_keys in the Target model will be disrupted instead of
|
|
# increasing by layer index, so need regular expressions to
|
|
# reorder the attn_keys and store the results in
|
|
# _ATTN_KEYS_BUFFER.
|
|
attn_keys_length = len(graph_params.attn_params[num_tokens])
|
|
global _ATTN_KEYS_BUFFER
|
|
if attn_keys_length == 0:
|
|
return
|
|
if not _ATTN_KEYS_BUFFER or len(_ATTN_KEYS_BUFFER) != attn_keys_length:
|
|
import regex as re
|
|
|
|
def extract_layer_index(key: str) -> int:
|
|
match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", key)
|
|
return int(match.group(1)) if match else 0
|
|
|
|
def is_direct_target_attn_key(key: str) -> bool:
|
|
return (
|
|
re.search(
|
|
r"(?:^|\.)layers\.(\d+)\.self_attn\.attn$",
|
|
key,
|
|
)
|
|
is not None
|
|
)
|
|
|
|
attn_keys_to_order = attn_keys[:attn_keys_length]
|
|
if getattr(speculative_config, "method", None) == "mtp":
|
|
# Step3.5 MTP can expose draft KV-cache groups in the
|
|
# target runtime metadata. The target FULL graph only
|
|
# captures direct base-model self-attention handles, so
|
|
# select that target key domain instead of depending on
|
|
# the current draft module name.
|
|
direct_target_attn_keys = [key for key in attn_keys if is_direct_target_attn_key(key)]
|
|
if len(direct_target_attn_keys) >= attn_keys_length:
|
|
attn_keys_to_order = direct_target_attn_keys
|
|
|
|
attn_keys_tmp = attn_keys_to_order
|
|
attn_keys_tmp.sort(key=extract_layer_index)
|
|
_ATTN_KEYS_BUFFER = attn_keys_tmp[:attn_keys_length]
|
|
attn_keys[:attn_keys_length] = _ATTN_KEYS_BUFFER
|
|
# For Qwen3-next, since the kv_cache_config has already categorized
|
|
# linear_attn and self_attn, the attn_metadata is first arranged with
|
|
# self_attn followed by linear_attn. Therefore, using zip directly
|
|
# filters out the update operations for linear_attn.
|
|
# TODO: We use a new variable `attn_keys` to ensure the loop count is
|
|
# correct after get by `zip` because of the new structure of the attn_metadata
|
|
# when running with the merged full eagle-graph. Should check it with Qwen3-next.
|
|
num_layers = len(attn_keys)
|
|
if num_layers == 0:
|
|
return
|
|
captured_attn_params = graph_params.attn_params[num_tokens]
|
|
handles = graph_params.handles[num_tokens]
|
|
events = graph_params.events[num_tokens]
|
|
graph_param_count = len(captured_attn_params)
|
|
workspace = graph_params.workspaces.get(num_tokens)
|
|
if _EXTRA_CTX.is_draft_model:
|
|
if graph_param_count > len(draft_attn_key_steps):
|
|
repeat_count = cdiv(graph_param_count, len(draft_attn_key_steps))
|
|
draft_attn_key_steps = (draft_attn_key_steps * repeat_count)[:graph_param_count]
|
|
else:
|
|
draft_attn_key_steps = draft_attn_key_steps[:graph_param_count]
|
|
attn_keys = [key for _, key in draft_attn_key_steps]
|
|
elif use_layer_aware_replay:
|
|
# Keep the replay loop length aligned with captured FIA ops;
|
|
# layer-specific metadata lookup below prevents global/sliding
|
|
# window layers from accidentally sharing the same metadata.
|
|
attn_keys = [attn_keys[index % num_layers] for index in range(graph_param_count)]
|
|
attn_count = 0
|
|
with torch.npu.stream(update_stream):
|
|
for key, param, handle, event in zip(
|
|
attn_keys,
|
|
captured_attn_params,
|
|
handles,
|
|
events,
|
|
):
|
|
(
|
|
query,
|
|
key_cache,
|
|
value,
|
|
block_tables,
|
|
attn_mask,
|
|
block_size,
|
|
seq_lens,
|
|
query_start_loc,
|
|
num_kv_heads,
|
|
num_heads,
|
|
scale,
|
|
attn_output,
|
|
softmax_lse,
|
|
sparse_mode,
|
|
pre_tokens,
|
|
next_tokens,
|
|
sliding_window,
|
|
c8_k_aq_scale,
|
|
c8_k_aq_offset,
|
|
c8_v_aq_scale,
|
|
c8_v_aq_offset,
|
|
layer_name,
|
|
) = param
|
|
|
|
if _EXTRA_CTX.is_draft_model:
|
|
draft_step, key = draft_attn_key_steps[attn_count]
|
|
metadata = attn_metadata[draft_step][key]
|
|
seq_lens = metadata.seq_lens_list
|
|
actual_seq_lengths_q = metadata.actual_seq_lengths_q
|
|
block_tables = metadata.block_tables
|
|
attn_count = attn_count + 1
|
|
if not metadata.causal:
|
|
sparse_mode = 0
|
|
else:
|
|
metadata_key = layer_name if layer_name is not None and layer_name in attn_metadata else key
|
|
seq_lens = attn_metadata[metadata_key].seq_lens_list
|
|
actual_seq_lengths_q = attn_metadata[metadata_key].actual_seq_lengths_q
|
|
# NOTE:
|
|
# For models with sliding-window attention on the FIA full-graph replay path,
|
|
# rebinding `block_tables` to the latest metadata tensor causes corrupted /
|
|
# repeated outputs in our repro on Ascend NPU.
|
|
#
|
|
# Keep the captured block_tables tensor on this affected path.
|
|
# Non-SWA models preserve the original behavior and continue to refresh
|
|
# block_tables from attn_metadata.
|
|
if sliding_window is None:
|
|
block_tables = attn_metadata[metadata_key].block_tables
|
|
|
|
torch.npu.graph_task_update_begin(update_stream, handle)
|
|
input_layout = "TND"
|
|
extra_args = {}
|
|
if c8_k_aq_scale is not None:
|
|
extra_args = {
|
|
"key_antiquant_scale": c8_k_aq_scale,
|
|
"value_antiquant_scale": c8_v_aq_scale,
|
|
"key_antiquant_mode": 0,
|
|
"value_antiquant_mode": 0,
|
|
"inner_precise": 1,
|
|
}
|
|
input_layout = "BNSD"
|
|
sparse_mode = 0
|
|
torch_npu.npu_fused_infer_attention_score.out(
|
|
query=query,
|
|
key=key_cache,
|
|
value=value,
|
|
block_table=block_tables,
|
|
atten_mask=attn_mask,
|
|
input_layout=input_layout,
|
|
block_size=block_size,
|
|
actual_seq_lengths=actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=seq_lens,
|
|
num_key_value_heads=num_kv_heads,
|
|
num_heads=num_heads,
|
|
scale=scale,
|
|
sparse_mode=sparse_mode,
|
|
pre_tokens=pre_tokens,
|
|
next_tokens=next_tokens,
|
|
**extra_args,
|
|
workspace=workspace,
|
|
out=[attn_output, softmax_lse],
|
|
)
|
|
torch.npu.graph_task_update_end(update_stream)
|
|
|
|
event.record(update_stream)
|
|
|
|
def process_weights_after_loading(self, act_dtype: torch.dtype):
|
|
super().process_weights_after_loading(act_dtype)
|
|
if flashcomm2_oshard_manager.flashcomm2_oshard_enable():
|
|
flashcomm2_oshard_manager.post_process_after_loading()
|
|
|
|
def full_graph_fia(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
layer=None,
|
|
) -> torch.Tensor:
|
|
passed_key = key
|
|
key, value, block_size, block_table, actual_seq_lengths_kv = self._get_fia_params(key, value, attn_metadata)
|
|
if self.enable_hamming_sparse and attn_metadata.attn_state != AscendAttentionState.DecodeOnly:
|
|
reshape_and_cache_kvcomp(attn_metadata.kvcomp_metadata, self.layerIndex, passed_key)
|
|
elif self.enable_hamming_sparse:
|
|
block_table, actual_seq_lengths_kv = get_kvcomp_decode_params(
|
|
self.layerIndex, attn_metadata.kvcomp_metadata, query, passed_key, block_table, actual_seq_lengths_kv
|
|
)
|
|
|
|
num_tokens = attn_metadata.actual_seq_lengths_q[-1]
|
|
if _EXTRA_CTX.is_draft_model:
|
|
if _EXTRA_CTX.is_draft_model_prefill:
|
|
graph_params = get_draft_graph_prefill_params()
|
|
else:
|
|
graph_params = get_draft_graph_params()
|
|
else:
|
|
graph_params = get_graph_params()
|
|
actual_seq_lengths_q = attn_metadata.actual_seq_lengths_q
|
|
softmax_lse = torch.empty(1, dtype=query.dtype, device=query.device)
|
|
input_layout = "TND"
|
|
attn_mask = attn_metadata.attn_mask
|
|
sparse_mode = 4 if self.sliding_window else 3 if attn_metadata.causal else 0
|
|
pre_tokens = self.sliding_window or SWA_INT_MAX
|
|
next_tokens = 0 if self.sliding_window else SWA_INT_MAX
|
|
|
|
extra_args = {}
|
|
if self.enable_c8_quant and layer is not None:
|
|
extra_args = {
|
|
"key_antiquant_scale": layer._c8_k_aq_scale_nz_bnsd,
|
|
"value_antiquant_scale": layer._c8_v_aq_scale_nz_bnsd,
|
|
"key_antiquant_mode": 0,
|
|
"value_antiquant_mode": 0,
|
|
"inner_precise": 1,
|
|
}
|
|
|
|
# change key/value shape
|
|
_, block_size, _, _ = self.key_cache.shape # type: ignore
|
|
key = self._nz_5d_view(self.key_cache, block_size)
|
|
value = self._nz_5d_view(self.value_cache, block_size)
|
|
|
|
# TODO: change layerout from BNSD to TND.
|
|
input_layout = "BNSD"
|
|
query = query.unsqueeze(2)
|
|
output = output.unsqueeze(2)
|
|
attn_mask = None
|
|
sparse_mode = 0
|
|
use_max_workspace = self._use_max_workspace_for_fia_graph
|
|
workspace = graph_params.workspaces.get(num_tokens)
|
|
should_update_workspace_cache = False
|
|
if use_max_workspace:
|
|
# Some models mix attention layer shapes under the same graph size.
|
|
# During capture, keep the largest required workspace for that size.
|
|
candidate_workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_mask,
|
|
block_table=block_table,
|
|
input_layout=input_layout,
|
|
block_size=block_size,
|
|
actual_seq_lengths=actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
sparse_mode=sparse_mode,
|
|
pre_tokens=pre_tokens,
|
|
next_tokens=next_tokens,
|
|
scale=self.scale,
|
|
**extra_args,
|
|
)
|
|
workspace = cache_graph_workspace(
|
|
graph_params,
|
|
num_tokens,
|
|
candidate_workspace,
|
|
use_max_workspace=use_max_workspace,
|
|
)
|
|
should_update_workspace_cache = True
|
|
elif workspace is None:
|
|
workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_mask,
|
|
block_table=block_table,
|
|
input_layout=input_layout,
|
|
block_size=block_size,
|
|
actual_seq_lengths=actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
sparse_mode=sparse_mode,
|
|
pre_tokens=pre_tokens,
|
|
next_tokens=next_tokens,
|
|
scale=self.scale,
|
|
**extra_args,
|
|
)
|
|
should_update_workspace_cache = True
|
|
if should_update_workspace_cache:
|
|
if _EXTRA_CTX.is_draft_model:
|
|
update_draft_graph_params_workspaces(num_tokens, workspace)
|
|
else:
|
|
update_graph_params_workspaces(num_tokens, workspace)
|
|
|
|
# Handle graph capturing mode
|
|
stream = torch_npu.npu.current_stream()
|
|
|
|
event = torch.npu.ExternalEvent()
|
|
event.wait(stream)
|
|
event.reset(stream)
|
|
graph_params.events[num_tokens].append(event)
|
|
attn_params = (
|
|
weak_ref_tensors(query),
|
|
weak_ref_tensors(key),
|
|
weak_ref_tensors(value),
|
|
weak_ref_tensors(block_table),
|
|
weak_ref_tensors(attn_mask) if attn_mask is not None else None,
|
|
block_size,
|
|
actual_seq_lengths_kv,
|
|
actual_seq_lengths_q,
|
|
self.num_kv_heads,
|
|
self.num_heads,
|
|
self.scale,
|
|
weak_ref_tensors(output),
|
|
weak_ref_tensors(softmax_lse),
|
|
sparse_mode,
|
|
pre_tokens,
|
|
next_tokens,
|
|
self.sliding_window,
|
|
)
|
|
if self.enable_c8_quant and layer is not None:
|
|
attn_params = attn_params + (
|
|
weak_ref_tensors(layer._c8_k_aq_scale_nz_bnsd),
|
|
None,
|
|
weak_ref_tensors(layer._c8_v_aq_scale_nz_bnsd),
|
|
None,
|
|
) # type: ignore
|
|
else:
|
|
attn_params = attn_params + (None, None, None, None) # type: ignore
|
|
layer_name = self._graph_metadata_layer_name(layer) if self._use_layer_aware_fia_graph_replay else None
|
|
attn_params = attn_params + (layer_name,) # type: ignore
|
|
graph_params.attn_params[num_tokens].append(attn_params)
|
|
|
|
torch.npu.graph_task_group_begin(stream)
|
|
torch_npu.npu_fused_infer_attention_score.out(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_mask,
|
|
block_table=block_table,
|
|
input_layout=input_layout,
|
|
block_size=block_size,
|
|
actual_seq_lengths=actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale=self.scale,
|
|
sparse_mode=sparse_mode,
|
|
pre_tokens=pre_tokens,
|
|
next_tokens=next_tokens,
|
|
workspace=workspace,
|
|
out=[output, softmax_lse],
|
|
**extra_args,
|
|
)
|
|
|
|
output = output.view(num_tokens, self.num_heads, self.head_size)
|
|
|
|
handle = torch.npu.graph_task_group_end(stream)
|
|
graph_params.handles[num_tokens].append(handle)
|
|
return output, num_tokens
|
|
|
|
def full_graph_fia_v2(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
key, value, block_size, block_table, actual_seq_lengths_kv = self._get_fia_params(key, value, attn_metadata)
|
|
actual_seq_lengths_kv = attn_metadata.seq_lens
|
|
num_tokens = attn_metadata.actual_seq_lengths_q[-1]
|
|
if _EXTRA_CTX.is_draft_model:
|
|
graph_params = get_draft_graph_params()
|
|
else:
|
|
graph_params = get_graph_params()
|
|
|
|
actual_seq_lengths_q = attn_metadata.actual_seq_lengths_q
|
|
softmax_lse = torch.empty(1, dtype=query.dtype, device=query.device)
|
|
use_max_workspace = self._use_max_workspace_for_fia_graph
|
|
workspace = graph_params.workspaces.get(num_tokens)
|
|
should_update_workspace_cache = False
|
|
if use_max_workspace:
|
|
# See full_graph_fia: this path needs the max workspace across layer
|
|
# variants sharing the same graph size.
|
|
candidate_workspace = torch_npu._npu_fused_infer_attention_score_v2_get_max_workspace(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_qlen=actual_seq_lengths_q,
|
|
actual_seq_kvlen=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
softmax_scale=self.scale,
|
|
num_query_heads=self.num_heads,
|
|
sparse_mode=4 if self.sliding_window is not None else 3,
|
|
pre_tokens=self.sliding_window if self.sliding_window is not None else SWA_INT_MAX,
|
|
next_tokens=0,
|
|
learnable_sink=self.sinks,
|
|
)
|
|
workspace = cache_graph_workspace(
|
|
graph_params,
|
|
num_tokens,
|
|
candidate_workspace,
|
|
use_max_workspace=use_max_workspace,
|
|
)
|
|
should_update_workspace_cache = True
|
|
elif workspace is None:
|
|
workspace = torch_npu._npu_fused_infer_attention_score_v2_get_max_workspace(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_qlen=actual_seq_lengths_q,
|
|
actual_seq_kvlen=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
softmax_scale=self.scale,
|
|
num_query_heads=self.num_heads,
|
|
sparse_mode=4 if self.sliding_window is not None else 3,
|
|
pre_tokens=self.sliding_window if self.sliding_window is not None else SWA_INT_MAX,
|
|
next_tokens=0,
|
|
learnable_sink=self.sinks,
|
|
)
|
|
should_update_workspace_cache = True
|
|
if should_update_workspace_cache:
|
|
if _EXTRA_CTX.is_draft_model:
|
|
update_draft_graph_params_workspaces(num_tokens, workspace)
|
|
else:
|
|
update_graph_params_workspaces(num_tokens, workspace)
|
|
|
|
# Handle graph capturing mode
|
|
stream = torch_npu.npu.current_stream()
|
|
|
|
event = torch.npu.ExternalEvent()
|
|
event.wait(stream)
|
|
event.reset(stream)
|
|
graph_params.events[num_tokens].append(event)
|
|
graph_params.attn_params[num_tokens].append(
|
|
(
|
|
weak_ref_tensors(query),
|
|
weak_ref_tensors(key),
|
|
weak_ref_tensors(value),
|
|
weak_ref_tensors(block_table),
|
|
weak_ref_tensors(attn_metadata.attn_mask),
|
|
block_size,
|
|
actual_seq_lengths_kv,
|
|
self.num_kv_heads,
|
|
self.num_heads,
|
|
self.scale,
|
|
self.sliding_window,
|
|
self.sinks,
|
|
weak_ref_tensors(output),
|
|
weak_ref_tensors(softmax_lse),
|
|
self._graph_metadata_layer_name() if self._use_layer_aware_fia_graph_replay else None,
|
|
)
|
|
)
|
|
torch.npu.graph_task_group_begin(stream)
|
|
torch_npu.npu_fused_infer_attention_score_v2.out(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_qlen=actual_seq_lengths_q,
|
|
actual_seq_kvlen=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_query_heads=self.num_heads,
|
|
sparse_mode=4 if self.sliding_window is not None else 3,
|
|
pre_tokens=self.sliding_window if self.sliding_window is not None else SWA_INT_MAX,
|
|
next_tokens=0,
|
|
softmax_scale=self.scale,
|
|
learnable_sink=self.sinks,
|
|
workspace=workspace,
|
|
out=[output, softmax_lse],
|
|
)
|
|
handle = torch.npu.graph_task_group_end(stream)
|
|
graph_params.handles[num_tokens].append(handle)
|
|
return output, num_tokens
|
|
|
|
def full_graph_pa(
|
|
self,
|
|
query: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor | None = None,
|
|
):
|
|
graph_params = get_graph_params()
|
|
num_tokens = query.shape[0]
|
|
if _EXTRA_CTX.capturing:
|
|
# Get workspace from cache or calculate it if not present.
|
|
workspace = graph_params.workspaces.get(num_tokens)
|
|
if workspace is None:
|
|
workspace = torch_npu._npu_paged_attention_get_workspace(
|
|
query=query,
|
|
key_cache=self.key_cache,
|
|
value_cache=self.value_cache,
|
|
num_kv_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale_value=self.scale,
|
|
block_table=attn_metadata.block_tables,
|
|
context_lens=attn_metadata.seq_lens,
|
|
out=output,
|
|
)
|
|
update_graph_params_workspaces(num_tokens, workspace)
|
|
|
|
# Handle graph capturing mode
|
|
stream = torch_npu.npu.current_stream()
|
|
|
|
event = torch.npu.ExternalEvent()
|
|
event.wait(stream)
|
|
event.reset(stream)
|
|
graph_params.events[num_tokens].append(event)
|
|
graph_params.attn_params[num_tokens].append(
|
|
(
|
|
weak_ref_tensors(query),
|
|
weak_ref_tensors(self.key_cache),
|
|
weak_ref_tensors(self.value_cache),
|
|
self.num_kv_heads,
|
|
self.num_heads,
|
|
self.scale,
|
|
attn_metadata.block_tables,
|
|
attn_metadata.seq_lens,
|
|
weak_ref_tensors(output),
|
|
)
|
|
)
|
|
|
|
torch.npu.graph_task_group_begin(stream)
|
|
torch_npu._npu_paged_attention(
|
|
query=query,
|
|
key_cache=self.key_cache,
|
|
value_cache=self.value_cache,
|
|
num_kv_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale_value=self.scale,
|
|
block_table=attn_metadata.block_tables,
|
|
context_lens=attn_metadata.seq_lens,
|
|
out=output,
|
|
workspace=workspace,
|
|
)
|
|
handle = torch.npu.graph_task_group_end(stream)
|
|
graph_params.handles[num_tokens].append(handle)
|
|
return output
|
|
|
|
def _get_fia_params(self, key: torch.Tensor, value: torch.Tensor, attn_metadata: AscendMetadata, kv_cache=None):
|
|
# PrefillNoCache doesn't need key_cache, but other modes do
|
|
# Only initialize/require cache for modes that actually use it
|
|
if attn_metadata.attn_state != AscendAttentionState.PrefillNoCache:
|
|
# Initialize cache from kv_cache if not already set (for DecodeOnly mode)
|
|
if self.key_cache is None and kv_cache is not None:
|
|
if (
|
|
isinstance(kv_cache, torch.Tensor)
|
|
and kv_cache.dim() > 0
|
|
and kv_cache.shape[0] == 2
|
|
or isinstance(kv_cache, (list, tuple))
|
|
and len(kv_cache) >= 2
|
|
):
|
|
self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
|
|
|
|
if self.key_cache is None:
|
|
raise RuntimeError(
|
|
f"key_cache is None in _get_fia_params for mode {attn_metadata.attn_state}. kv_cache={kv_cache}"
|
|
)
|
|
|
|
if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
|
|
block_size = 128
|
|
block_table = None
|
|
actual_seq_lengths_kv = attn_metadata.actual_seq_lengths_q
|
|
if self.attn_type == AttentionType.ENCODER_DECODER:
|
|
actual_seq_lengths_kv = torch.cumsum(attn_metadata.seq_lens, dim=0).tolist()
|
|
elif attn_metadata.attn_state == AscendAttentionState.PrefillCacheHit:
|
|
batch_size = attn_metadata.seq_lens.shape[0]
|
|
block_table = attn_metadata.block_tables[:batch_size, :]
|
|
num_block, block_size, _, _ = self.key_cache.shape # type: ignore
|
|
key = self.key_cache.view( # type: ignore
|
|
num_block, block_size, -1
|
|
)
|
|
value = self.value_cache.view( # type: ignore
|
|
num_block, block_size, -1
|
|
)
|
|
actual_seq_lengths_kv = attn_metadata.seq_lens_list
|
|
elif attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
|
|
num_block, block_size, _, _ = self.key_cache.shape # type: ignore
|
|
key = self.key_cache.view( # type: ignore
|
|
num_block, block_size, -1
|
|
)
|
|
value = self.value_cache.view( # type: ignore
|
|
num_block, block_size, -1
|
|
)
|
|
block_table = attn_metadata.block_tables
|
|
actual_seq_lengths_kv = attn_metadata.seq_lens_list
|
|
# chunked prefill.
|
|
else:
|
|
num_block, block_size, _, _ = self.key_cache.shape # type: ignore
|
|
key = self.key_cache.view( # type: ignore
|
|
num_block, block_size, -1
|
|
)
|
|
value = self.value_cache.view( # type: ignore
|
|
num_block, block_size, -1
|
|
)
|
|
block_table = attn_metadata.block_tables
|
|
actual_seq_lengths_kv = attn_metadata.seq_lens_list
|
|
return key, value, block_size, block_table, actual_seq_lengths_kv
|
|
|
|
def forward_fused_infer_attention(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
kv_cache=None,
|
|
):
|
|
# we inherit ForwardContext in model runner v2, when enable model
|
|
# runner v2, there is not capturing attribute in forward_context,
|
|
# just use getattr to avoid attribute error.
|
|
if _EXTRA_CTX.capturing:
|
|
if self.sinks is not None:
|
|
attn_output, num_tokens = self.full_graph_fia_v2(query, key, value, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
else:
|
|
attn_output, num_tokens = self.full_graph_fia(query, key, value, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
passed_key = key
|
|
passed_value = value
|
|
key, value, block_size, block_table, actual_seq_lengths_kv = self._get_fia_params(
|
|
key, value, attn_metadata, kv_cache
|
|
)
|
|
if self.enable_hamming_sparse and attn_metadata.attn_state != AscendAttentionState.DecodeOnly:
|
|
reshape_and_cache_kvcomp(attn_metadata.kvcomp_metadata, self.layerIndex, passed_key)
|
|
elif self.enable_hamming_sparse:
|
|
block_table, actual_seq_lengths_kv = get_kvcomp_decode_params(
|
|
self.layerIndex, attn_metadata.kvcomp_metadata, query, passed_key, block_table, actual_seq_lengths_kv
|
|
)
|
|
num_tokens = attn_metadata.actual_seq_lengths_q[-1]
|
|
query = query[:num_tokens]
|
|
if (
|
|
attn_metadata.attn_state == AscendAttentionState.PrefillNoCache
|
|
and self.attn_type != AttentionType.ENCODER_DECODER
|
|
):
|
|
key = key[:num_tokens]
|
|
value = value[:num_tokens]
|
|
# Get workspace from cache or calculate it if not present.
|
|
if self.sinks is not None:
|
|
actual_seq_qlen = attn_metadata.actual_seq_lengths_q
|
|
if attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
|
|
actual_seq_qlen = torch.tensor([1] * len(attn_metadata.seq_lens_list), dtype=torch.int32).cumsum(dim=0)
|
|
if self.sliding_window is not None:
|
|
sparse_mode = 4
|
|
else:
|
|
sparse_mode = 3
|
|
attn_output, _ = torch_npu.npu_fused_infer_attention_score_v2(
|
|
query,
|
|
key.contiguous(),
|
|
value.contiguous(),
|
|
num_query_heads=self.num_heads,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
input_layout="TND",
|
|
pre_tokens=self.sliding_window if self.sliding_window is not None else SWA_INT_MAX,
|
|
next_tokens=0,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
sparse_mode=sparse_mode,
|
|
softmax_scale=self.scale,
|
|
block_table=block_table,
|
|
block_size=block_size,
|
|
actual_seq_qlen=actual_seq_qlen,
|
|
actual_seq_kvlen=actual_seq_lengths_kv,
|
|
learnable_sink=self.sinks,
|
|
)
|
|
else:
|
|
if not attn_metadata.causal:
|
|
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_lengths=attn_metadata.actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale=self.scale,
|
|
sparse_mode=0,
|
|
)
|
|
elif self.sliding_window is not None:
|
|
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_lengths=attn_metadata.actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale=self.scale,
|
|
pre_tokens=self.sliding_window,
|
|
next_tokens=0,
|
|
sparse_mode=4,
|
|
)
|
|
else:
|
|
attn_output, _ = DeviceOperator.npu_fused_infer_attention_score(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_lengths=attn_metadata.actual_seq_lengths_q,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
head_size=self.head_size,
|
|
scale=self.scale,
|
|
key_cache=self.key_cache,
|
|
value_cache=self.value_cache,
|
|
current_key=passed_key,
|
|
current_value=passed_value,
|
|
attn_metadata=attn_metadata,
|
|
is_prefill_no_cache=attn_metadata.attn_state == AscendAttentionState.PrefillNoCache,
|
|
sparse_mode=3,
|
|
)
|
|
|
|
attn_output = attn_output.view(num_tokens, self.num_heads, self.head_size)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
|
|
def forward_paged_attention(
|
|
self,
|
|
query: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if _EXTRA_CTX.capturing:
|
|
return self.full_graph_pa(query, attn_metadata, output)
|
|
torch_npu._npu_paged_attention(
|
|
query=query,
|
|
key_cache=self.key_cache,
|
|
value_cache=self.value_cache,
|
|
num_kv_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale_value=self.scale,
|
|
block_table=attn_metadata.block_tables,
|
|
context_lens=attn_metadata.seq_lens,
|
|
out=output,
|
|
)
|
|
return output
|
|
|
|
def _forward_encoder_attention(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
_: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
# use default sparse_mode 0 in normal scenario, which means no mask works on it
|
|
# Pad actual_seq_len with 0 when num_tokens > actual_seq_len in TND layout
|
|
actual_seq_qlen = attn_metadata.actual_seq_lengths_q
|
|
if query.shape[0] > actual_seq_qlen[-1]:
|
|
actual_seq_qlen = actual_seq_qlen + [0]
|
|
return torch_npu.npu_fusion_attention(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
head_num=self.num_heads,
|
|
input_layout="TND",
|
|
scale=self.scale,
|
|
actual_seq_qlen=actual_seq_qlen,
|
|
actual_seq_kvlen=actual_seq_qlen,
|
|
)[0]
|
|
|
|
def do_kv_cache_update(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: list[torch.Tensor],
|
|
slot_mapping: torch.Tensor,
|
|
) -> None:
|
|
if self.attn_type in (AttentionType.ENCODER_ONLY):
|
|
return
|
|
|
|
if self.key_cache is None:
|
|
self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
|
|
|
|
DeviceOperator.reshape_and_cache(
|
|
key=key,
|
|
value=value,
|
|
key_cache=self.key_cache,
|
|
value_cache=self.value_cache,
|
|
slot_mapping=slot_mapping,
|
|
)
|
|
|
|
def reshape_and_cache(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: tuple[torch.Tensor],
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
):
|
|
if len(kv_cache) > 1:
|
|
if self.key_cache is None:
|
|
self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
|
|
slots = attn_metadata.slot_mapping
|
|
encoder_decoder = self.attn_type == AttentionType.ENCODER_DECODER
|
|
DeviceOperator.reshape_and_cache(
|
|
key=key[: attn_metadata.num_actual_tokens] if not encoder_decoder else key,
|
|
value=value[: attn_metadata.num_actual_tokens] if not encoder_decoder else value,
|
|
key_cache=self.key_cache,
|
|
value_cache=self.value_cache,
|
|
# quick fix to make sure slots is int32 for cross attention case.
|
|
# see: https://github.com/vllm-project/vllm/blob/ce88756b967c2c5006746a424c15dd59a284ed8c/vllm/model_executor/layers/attention/cross_attention.py#L117
|
|
slot_mapping=slots[: attn_metadata.num_actual_tokens] if not encoder_decoder else slots.to(torch.int32),
|
|
)
|
|
notify_kv_cache_written()
|
|
return query, key, value, output
|
|
|
|
def forward_impl(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: tuple[torch.Tensor],
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
):
|
|
num_tokens = query.shape[0]
|
|
record_attention_compute_start()
|
|
|
|
if (
|
|
attn_metadata.attn_state == AscendAttentionState.DecodeOnly
|
|
and self.sliding_window is None
|
|
and using_paged_attention(num_tokens, self.vllm_config, self.head_size)
|
|
):
|
|
output = self.forward_paged_attention(query, attn_metadata, output)
|
|
else:
|
|
output = self.forward_fused_infer_attention(query, key, value, attn_metadata, output, kv_cache)
|
|
|
|
return output
|
|
|
|
def forward(
|
|
self,
|
|
layer: AttentionLayer,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: tuple[torch.Tensor],
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor | None = None,
|
|
output_scale: torch.Tensor | None = None,
|
|
output_block_scale: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward pass with Ascend attention.
|
|
Args:
|
|
query: shape = [num_tokens, num_heads, head_size]
|
|
key: shape = [num_tokens, num_kv_heads, head_size]
|
|
value: shape = [num_tokens, num_kv_heads, head_size]
|
|
kv_cache: shape =
|
|
[2, num_blocks, block_size, num_kv_heads, head_size]
|
|
attn_metadata: Metadata for attention.
|
|
Returns:
|
|
shape = [num_tokens, num_heads * head_size]
|
|
"""
|
|
assert output is not None, "Output tensor must be provided."
|
|
if self.enable_hamming_sparse:
|
|
self.layerIndex = int(layer.layer_name.split(".")[2])
|
|
if self._use_layer_aware_fia_graph_replay:
|
|
self._layer_name = layer.layer_name
|
|
|
|
if output_scale is not None or output_block_scale is not None:
|
|
raise NotImplementedError("fused output quantization is not yet supported for AscendAttentionBackendImpl")
|
|
|
|
assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
|
|
num_tokens = query.shape[0]
|
|
if attn_metadata is None:
|
|
return output.fill_(0)
|
|
|
|
# Initialize key_cache and value_cache from kv_cache if not already set.
|
|
# This is needed for DecodeOnly mode where key/value are None but we still
|
|
# need access to the cache for attention computation.
|
|
if self.key_cache is None and kv_cache is not None:
|
|
if (
|
|
isinstance(kv_cache, torch.Tensor)
|
|
and kv_cache.dim() > 0
|
|
and kv_cache.shape[0] == 2
|
|
or isinstance(kv_cache, (list, tuple))
|
|
and len(kv_cache) >= 2
|
|
):
|
|
self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
|
|
|
|
output_padded = None
|
|
if key is not None and value is not None:
|
|
output_padded = output
|
|
query, key, value, output_padded = self.reshape_and_cache(
|
|
query, key, value, kv_cache, attn_metadata, output
|
|
)
|
|
# pooling model branch
|
|
if attn_metadata.model_runner_type == "pooling" and not attn_metadata.causal:
|
|
attn_output = self._forward_encoder_attention(query, key, value, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
if output_padded is not None:
|
|
attn_output = self.forward_impl(query, key, value, kv_cache, attn_metadata, output_padded)
|
|
else:
|
|
attn_output = self.forward_impl(query, key, value, kv_cache, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
|
|
|
|
class AscendC8AttentionBackendImpl(AscendAttentionBackendImpl):
|
|
"""Attention backend implementation for INT8 KV cache (C8/QuaRot) models.
|
|
|
|
This subclass handles static per-channel INT8 KV cache quantization.
|
|
It is activated via class surgery in AscendC8KVCacheAttentionMethod.create_weights
|
|
(vllm_ascend/quantization/methods/kv_c8.py)
|
|
so that C8 attention layers automatically use this forward path.
|
|
"""
|
|
|
|
def forward(
|
|
self,
|
|
layer: AttentionLayer,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: tuple[torch.Tensor],
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor | None = None,
|
|
output_scale: torch.Tensor | None = None,
|
|
output_block_scale: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
assert output is not None, "Output tensor must be provided."
|
|
if self._use_layer_aware_fia_graph_replay:
|
|
self._layer_name = layer.layer_name
|
|
|
|
if output_scale is not None or output_block_scale is not None:
|
|
raise NotImplementedError("fused output quantization is not yet supported for AscendC8AttentionBackendImpl")
|
|
|
|
num_tokens = query.shape[0]
|
|
if attn_metadata is None:
|
|
return output.fill_(0)
|
|
|
|
self._prepare_c8_scales(layer, query.device)
|
|
float_key, float_value = None, None
|
|
if self.vllm_config.kv_transfer_config is None:
|
|
if key is not None and value is not None:
|
|
if attn_metadata.attn_state != AscendAttentionState.DecodeOnly:
|
|
float_key, float_value = key, value
|
|
key, value = self._quantize_kv_to_int8(key, value, layer, attn_metadata.num_actual_tokens)
|
|
query, key, value, _ = self._reshape_and_cache(query, key, value, kv_cache, attn_metadata, output)
|
|
# pooling model branch
|
|
if attn_metadata.model_runner_type == "pooling":
|
|
attn_output = self._forward_encoder_attention(query, key, value, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
if attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
|
|
if _EXTRA_CTX.capturing:
|
|
attn_output, num_tokens = self.full_graph_fia(query, key, value, attn_metadata, output, layer)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
return self._forward_c8_decode(query, attn_metadata, output, layer)
|
|
elif attn_metadata.attn_state == AscendAttentionState.ChunkedPrefill:
|
|
return self._forward_c8_chunked_prefill(query, float_key, float_value, attn_metadata, output, layer)
|
|
else:
|
|
return self._forward_c8_fused_infer_attention(
|
|
query,
|
|
float_key if float_key is not None else key,
|
|
float_value if float_value is not None else value,
|
|
attn_metadata,
|
|
output,
|
|
layer,
|
|
)
|
|
else:
|
|
if attn_metadata.attn_state != AscendAttentionState.DecodeOnly and self.is_kv_producer:
|
|
output_padded = None
|
|
if key is not None and value is not None:
|
|
output_padded = output
|
|
query, key, value, output_padded = self.reshape_and_cache(
|
|
query, key, value, kv_cache, attn_metadata, output
|
|
)
|
|
# pooling model branch
|
|
if attn_metadata.model_runner_type == "pooling":
|
|
attn_output = self._forward_encoder_attention(query, key, value, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
if output_padded is not None:
|
|
attn_output = self.forward_impl(query, key, value, kv_cache, attn_metadata, output_padded)
|
|
else:
|
|
attn_output = self.forward_impl(query, key, value, kv_cache, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
elif not self.is_kv_producer:
|
|
if key is not None and value is not None:
|
|
key, value = self._quantize_kv_to_int8(key, value, layer, attn_metadata.num_actual_tokens)
|
|
query, key, value, _ = self._reshape_and_cache(query, key, value, kv_cache, attn_metadata, output)
|
|
# pooling model branch
|
|
if attn_metadata.model_runner_type == "pooling":
|
|
attn_output = self._forward_encoder_attention(query, key, value, attn_metadata, output)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
if _EXTRA_CTX.capturing:
|
|
attn_output, num_tokens = self.full_graph_fia(query, key, value, attn_metadata, output, layer)
|
|
output[:num_tokens] = attn_output[:num_tokens]
|
|
return output
|
|
elif attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
|
|
return self._forward_c8_decode(query, attn_metadata, output, layer)
|
|
|
|
def _nz_5d_view(self, cache: torch.Tensor, block_size: int) -> torch.Tensor:
|
|
"""View a KV cache tensor in NZ 5D layout: (num_blocks, num_kv_heads, head_size//nz, block_size, nz)."""
|
|
NZ_FMT_LAST_DIM = 32
|
|
return cache.view(-1, self.num_kv_heads, self.head_size // NZ_FMT_LAST_DIM, block_size, NZ_FMT_LAST_DIM)
|
|
|
|
def _prepare_c8_scales(self, layer: AttentionLayer, device: torch.device) -> None:
|
|
"""Shard per-channel C8 scales/offsets to this TP rank and pre-compute
|
|
BF16 BNSD antiquant tensors for FIA V1 decode fast path.
|
|
"""
|
|
if hasattr(layer, "_c8_scales_prepared"):
|
|
return
|
|
|
|
def _shard_and_reshape(raw: torch.Tensor) -> torch.Tensor:
|
|
if raw.numel() == 1:
|
|
return raw.to(device=device)
|
|
expected = self.num_kv_heads * self.head_size
|
|
if raw.numel() != expected:
|
|
total_kv_heads = raw.numel() // self.head_size
|
|
tp_rank = get_tensor_model_parallel_rank()
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
kv_head_start = tp_rank * total_kv_heads // tp_size
|
|
raw = raw.view(total_kv_heads, self.head_size)[
|
|
kv_head_start : kv_head_start + self.num_kv_heads
|
|
].contiguous()
|
|
return raw.view(1, self.num_kv_heads, self.head_size).to(device=device)
|
|
|
|
layer._c8_k_scale = _shard_and_reshape(layer.k_cache_scale.data)
|
|
layer._c8_k_offset = _shard_and_reshape(layer.k_cache_offset.data)
|
|
layer._c8_v_scale = _shard_and_reshape(layer.v_cache_scale.data)
|
|
layer._c8_v_offset = _shard_and_reshape(layer.v_cache_offset.data)
|
|
|
|
layer._c8_k_inv_scale = 1.0 / layer._c8_k_scale
|
|
layer._c8_v_inv_scale = 1.0 / layer._c8_v_scale
|
|
|
|
nz_bnsd = (self.num_kv_heads, 1, self.head_size)
|
|
layer._c8_k_aq_scale_nz_bnsd = layer._c8_k_scale.view(nz_bnsd).contiguous()
|
|
layer._c8_v_aq_scale_nz_bnsd = layer._c8_v_scale.view(nz_bnsd).contiguous()
|
|
|
|
layer._c8_scales_prepared = True
|
|
|
|
def _dequant_paged_kv_to_dense(
|
|
self,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
seq_lens: list,
|
|
target_dtype: torch.dtype,
|
|
layer,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Gather paged INT8 KV blocks and dequantize."""
|
|
batch_size = block_table.shape[0]
|
|
max_blocks_per_seq = block_table.shape[1]
|
|
|
|
# NZ 5D view: (num_blocks, num_kv_heads, head_size//nz, block_size, nz)
|
|
block_size = self.key_cache.shape[1] # type: ignore[attr-defined]
|
|
max_tokens_padded = max_blocks_per_seq * block_size
|
|
|
|
flat_ids = block_table.reshape(-1)
|
|
key_nz = self._nz_5d_view(key, block_size)
|
|
value_nz = self._nz_5d_view(value, block_size)
|
|
|
|
# Gather: (batch*max_blocks, H, D//nz, S, nz)
|
|
gathered_k = key_nz[flat_ids]
|
|
gathered_v = value_nz[flat_ids]
|
|
# NZ→ND conversion: permute (S, H, D//nz, nz) → reshape (S, H, D)
|
|
gathered_k = (
|
|
gathered_k.permute(0, 3, 1, 2, 4)
|
|
.contiguous()
|
|
.view(batch_size, max_tokens_padded, self.num_kv_heads, self.head_size)
|
|
)
|
|
gathered_v = (
|
|
gathered_v.permute(0, 3, 1, 2, 4)
|
|
.contiguous()
|
|
.view(batch_size, max_tokens_padded, self.num_kv_heads, self.head_size)
|
|
)
|
|
|
|
seq_lens_t = torch.tensor(seq_lens, dtype=torch.long, device=key.device)
|
|
positions = torch.arange(max_tokens_padded, dtype=torch.long, device=key.device)
|
|
valid_mask = (positions.unsqueeze(0) < seq_lens_t.unsqueeze(1)).view(-1)
|
|
|
|
dense_k = gathered_k.view(-1, self.num_kv_heads, self.head_size)[valid_mask]
|
|
dense_v = gathered_v.view(-1, self.num_kv_heads, self.head_size)[valid_mask]
|
|
|
|
# Scale-only dequant for NZ (symmetric)
|
|
dense_k = dense_k.to(target_dtype) * layer._c8_k_scale
|
|
dense_v = dense_v.to(target_dtype) * layer._c8_v_scale
|
|
return dense_k, dense_v
|
|
|
|
def _quantize_kv_to_int8(
|
|
self,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
layer: AttentionLayer,
|
|
num_actual_tokens: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Quantize K/V from float to INT8 using static per-channel C8 scales."""
|
|
actual_key = key[:num_actual_tokens]
|
|
actual_value = value[:num_actual_tokens]
|
|
|
|
k_int8 = torch.clamp(
|
|
torch.round(actual_key * layer._c8_k_inv_scale + layer._c8_k_offset),
|
|
-128,
|
|
127,
|
|
).to(torch.int8)
|
|
v_int8 = torch.clamp(
|
|
torch.round(actual_value * layer._c8_v_inv_scale + layer._c8_v_offset),
|
|
-128,
|
|
127,
|
|
).to(torch.int8)
|
|
return k_int8, v_int8
|
|
|
|
def _forward_c8_decode(
|
|
self,
|
|
query: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
layer: AttentionLayer,
|
|
) -> torch.Tensor:
|
|
"""C8 decode via FIA V1 BNSD with native paged INT8 KV + perchannel antiquant."""
|
|
num_block, block_size, _, _ = self.key_cache.shape # type: ignore[attr-defined]
|
|
assert block_size % 32 == 0, f"C8 INT8 KV cache requires block_size to be a multiple of 32, got {block_size}"
|
|
batch_size = len(attn_metadata.seq_lens_list)
|
|
|
|
key = self._nz_5d_view(self.key_cache, block_size)
|
|
value = self._nz_5d_view(self.value_cache, block_size)
|
|
|
|
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
|
|
query[:batch_size].unsqueeze(2),
|
|
key,
|
|
value,
|
|
key_antiquant_scale=layer._c8_k_aq_scale_nz_bnsd,
|
|
value_antiquant_scale=layer._c8_v_aq_scale_nz_bnsd,
|
|
block_table=attn_metadata.block_tables,
|
|
actual_seq_lengths_kv=attn_metadata.seq_lens_list,
|
|
num_heads=self.num_heads,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
input_layout="BNSD",
|
|
scale=self.scale,
|
|
block_size=block_size,
|
|
antiquant_mode=0,
|
|
key_antiquant_mode=0,
|
|
value_antiquant_mode=0,
|
|
inner_precise=1,
|
|
sparse_mode=0,
|
|
)
|
|
attn_output = attn_output.squeeze(2)
|
|
output[:batch_size] = attn_output
|
|
return output
|
|
|
|
def _forward_c8_chunked_prefill(
|
|
self,
|
|
query: torch.Tensor,
|
|
float_key: torch.Tensor | None,
|
|
float_value: torch.Tensor | None,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
layer: AttentionLayer,
|
|
) -> torch.Tensor:
|
|
"""C8 ChunkedPrefill: decode via FIA V1 BNSD paged INT8 (zero gather),
|
|
prefill via FIA V1 TND with float KV (new) or gather+dequant (continuing).
|
|
"""
|
|
num_decode_tokens = attn_metadata.num_decode_tokens
|
|
num_decodes = attn_metadata.num_decodes
|
|
actual_seq_qlen = attn_metadata.actual_seq_lengths_q
|
|
num_tokens = int(actual_seq_qlen[-1]) # type: ignore[index]
|
|
|
|
if num_decode_tokens > 0:
|
|
num_block, block_size, _, _ = self.key_cache.shape # type: ignore[attr-defined]
|
|
assert block_size % 32 == 0, (
|
|
f"C8 INT8 KV cache requires block_size to be a multiple of 32, got {block_size}"
|
|
)
|
|
kv_k = self._nz_5d_view(self.key_cache, block_size)
|
|
kv_v = self._nz_5d_view(self.value_cache, block_size)
|
|
|
|
attn_out, _ = torch_npu.npu_fused_infer_attention_score(
|
|
query[:num_decode_tokens].unsqueeze(2),
|
|
kv_k,
|
|
kv_v,
|
|
key_antiquant_scale=layer._c8_k_aq_scale_nz_bnsd,
|
|
value_antiquant_scale=layer._c8_v_aq_scale_nz_bnsd,
|
|
block_table=attn_metadata.block_tables[:num_decodes],
|
|
actual_seq_lengths_kv=attn_metadata.seq_lens_list[:num_decodes],
|
|
num_heads=self.num_heads,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
input_layout="BNSD",
|
|
scale=self.scale,
|
|
block_size=block_size,
|
|
antiquant_mode=0,
|
|
key_antiquant_mode=0,
|
|
value_antiquant_mode=0,
|
|
inner_precise=1,
|
|
sparse_mode=0,
|
|
)
|
|
output[:num_decode_tokens] = attn_out.squeeze(2)
|
|
|
|
if attn_metadata.num_prefills > 0:
|
|
prefill_q = query[num_decode_tokens:num_tokens]
|
|
|
|
prefill_seq_qlen = [
|
|
actual_seq_qlen[i] - num_decode_tokens for i in range(num_decodes, len(actual_seq_qlen))
|
|
]
|
|
|
|
all_new_prefill = True
|
|
for i in range(num_decodes, len(attn_metadata.seq_lens_list)):
|
|
q_start = actual_seq_qlen[i - 1] if i > 0 else 0
|
|
qlen_i = actual_seq_qlen[i] - q_start
|
|
if attn_metadata.seq_lens_list[i] > qlen_i:
|
|
all_new_prefill = False
|
|
break
|
|
|
|
if all_new_prefill and float_key is not None and float_value is not None:
|
|
prefill_k = float_key[num_decode_tokens:num_tokens]
|
|
prefill_v = float_value[num_decode_tokens:num_tokens]
|
|
prefill_seq_kvlen = prefill_seq_qlen
|
|
else:
|
|
num_block, blk_size, _, _ = self.key_cache.shape # type: ignore[attr-defined]
|
|
paged_k = self._nz_5d_view(self.key_cache, blk_size)
|
|
paged_v = self._nz_5d_view(self.value_cache, blk_size)
|
|
|
|
prefill_bt = attn_metadata.block_tables[num_decodes:]
|
|
prefill_sl = attn_metadata.seq_lens_list[num_decodes:]
|
|
prefill_k, prefill_v = self._dequant_paged_kv_to_dense(
|
|
paged_k, paged_v, prefill_bt, prefill_sl, query.dtype, layer
|
|
)
|
|
prefill_seq_kvlen = torch.tensor(prefill_sl, dtype=torch.int32).cumsum(dim=0)
|
|
|
|
# block_table is None for prefill; FIA ignores block_size in this case.
|
|
# Use cache block_size for consistency rather than a magic number.
|
|
cache_block_size = self.key_cache.shape[1] # type: ignore[attr-defined]
|
|
attn_out, _ = torch_npu.npu_fused_infer_attention_score(
|
|
query=prefill_q,
|
|
key=prefill_k,
|
|
value=prefill_v,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=None,
|
|
input_layout="TND",
|
|
block_size=cache_block_size,
|
|
actual_seq_lengths=prefill_seq_qlen,
|
|
actual_seq_lengths_kv=prefill_seq_kvlen,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale=self.scale,
|
|
sparse_mode=3,
|
|
)
|
|
n_prefill = num_tokens - num_decode_tokens
|
|
attn_out = attn_out.view(n_prefill, self.num_heads, self.head_size)
|
|
output[num_decode_tokens:num_tokens] = attn_out[:n_prefill]
|
|
|
|
return output
|
|
|
|
def _forward_c8_fused_infer_attention(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
layer: AttentionLayer,
|
|
):
|
|
"""C8 FIA V1 TND for prefill states (PrefillNoCache uses float KV directly,
|
|
PrefillCacheHit gathers + dequants paged INT8 KV).
|
|
"""
|
|
key, value, block_size, block_table, actual_seq_lengths_kv = self._get_fia_params(key, value, attn_metadata)
|
|
|
|
actual_seq_qlen = attn_metadata.actual_seq_lengths_q
|
|
num_tokens = int(actual_seq_qlen[-1]) # type: ignore[index]
|
|
query = query[:num_tokens]
|
|
|
|
if (
|
|
attn_metadata.attn_state == AscendAttentionState.PrefillNoCache
|
|
and self.attn_type != AttentionType.ENCODER_DECODER
|
|
):
|
|
key = key[:num_tokens]
|
|
value = value[:num_tokens]
|
|
|
|
if key.dtype == torch.int8:
|
|
if block_table is not None:
|
|
seq_lens = (
|
|
actual_seq_lengths_kv if isinstance(actual_seq_lengths_kv, list) else actual_seq_lengths_kv.tolist()
|
|
)
|
|
key, value = self._dequant_paged_kv_to_dense(key, value, block_table, seq_lens, query.dtype, layer)
|
|
block_table = None
|
|
# block_table is None after dequant; FIA ignores block_size.
|
|
# Use cache block_size for consistency rather than a magic number.
|
|
block_size = self.key_cache.shape[1] # type: ignore[attr-defined]
|
|
actual_seq_lengths_kv = torch.tensor(seq_lens, dtype=torch.int32).cumsum(dim=0)
|
|
else:
|
|
key = (key.to(query.dtype) - layer._c8_k_offset) * layer._c8_k_scale
|
|
value = (value.to(query.dtype) - layer._c8_v_offset) * layer._c8_v_scale
|
|
|
|
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
|
|
query=query,
|
|
key=key,
|
|
value=value,
|
|
atten_mask=attn_metadata.attn_mask,
|
|
block_table=block_table,
|
|
input_layout="TND",
|
|
block_size=block_size,
|
|
actual_seq_lengths=actual_seq_qlen,
|
|
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
|
num_key_value_heads=self.num_kv_heads,
|
|
num_heads=self.num_heads,
|
|
scale=self.scale,
|
|
sparse_mode=3,
|
|
)
|
|
attn_output = attn_output.view(num_tokens, self.num_heads, self.head_size)
|
|
output[:num_tokens] = attn_output
|
|
return output
|
|
|
|
def _reshape_and_cache(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: tuple[torch.Tensor],
|
|
attn_metadata: AscendMetadata,
|
|
output: torch.Tensor,
|
|
):
|
|
if len(kv_cache) > 1:
|
|
if self.key_cache is None:
|
|
self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
|
|
slots = attn_metadata.slot_mapping
|
|
|
|
encoder_decoder = self.attn_type == AttentionType.ENCODER_DECODER
|
|
|
|
# NZ write path: 5D view + npu_scatter_pa_kv_cache
|
|
block_size = self.vllm_config.cache_config.block_size
|
|
k_cache_layer = self._nz_5d_view(self.key_cache, block_size)
|
|
v_cache_layer = self._nz_5d_view(self.value_cache, block_size)
|
|
|
|
torch_npu.npu_scatter_pa_kv_cache(
|
|
key=key[: attn_metadata.num_actual_tokens] if not encoder_decoder else key,
|
|
value=value[: attn_metadata.num_actual_tokens] if not encoder_decoder else value,
|
|
key_cache=k_cache_layer,
|
|
value_cache=v_cache_layer,
|
|
slot_mapping=slots[: attn_metadata.num_actual_tokens] if not encoder_decoder else slots,
|
|
)
|
|
|
|
notify_kv_cache_written()
|
|
return query, key, value, output
|