@@ -1,81 +1,365 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List
|
||||
from dataclasses import dataclass, field
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
|
||||
has_kv_transfer_group,
|
||||
is_v1_kv_transfer_group)
|
||||
import torch.nn.functional as F
|
||||
from vllm.config import VllmConfig, get_current_vllm_config
|
||||
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group, is_v1_kv_transfer_group
|
||||
from vllm.forward_context import ForwardContext, get_forward_context
|
||||
from vllm.utils.torch_utils import get_dtype_size
|
||||
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
|
||||
|
||||
from vllm_ascend.device.utils import FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE
|
||||
from vllm_ascend.utils import (
|
||||
AscendDeviceType,
|
||||
get_ascend_config,
|
||||
get_ascend_device_type,
|
||||
is_pd_decode_recompute_scheduler_enabled,
|
||||
)
|
||||
from vllm_ascend.worker.kvcomp_utils import KVCompMetaData
|
||||
|
||||
SFA_QSFA_TILE_SIZE = 128
|
||||
|
||||
|
||||
def get_sfa_qsfa_packed_head_dim(
|
||||
kv_lora_rank: int,
|
||||
qk_rope_head_dim: int,
|
||||
tile_size: int = SFA_QSFA_TILE_SIZE,
|
||||
) -> int:
|
||||
if kv_lora_rank % tile_size != 0:
|
||||
raise ValueError(
|
||||
f"kv_lora_rank must be divisible by tile_size for SFA QSFA packed cache, "
|
||||
f"got {kv_lora_rank=} and {tile_size=}."
|
||||
)
|
||||
scale_metadata_bytes = (kv_lora_rank // tile_size) * get_dtype_size(torch.float32)
|
||||
return kv_lora_rank + qk_rope_head_dim * get_dtype_size(torch.bfloat16) + scale_metadata_bytes
|
||||
|
||||
|
||||
def cache_graph_workspace(
|
||||
graph_params,
|
||||
num_tokens: int,
|
||||
candidate_workspace: torch.Tensor,
|
||||
*,
|
||||
use_max_workspace: bool,
|
||||
) -> torch.Tensor:
|
||||
# Most models keep the original first-workspace cache behavior. Models with
|
||||
# mixed attention layer shapes may need the largest workspace for a graph
|
||||
# size because layers can require different FIA workspace sizes.
|
||||
current_workspace = graph_params.workspaces.get(num_tokens)
|
||||
if use_max_workspace:
|
||||
if current_workspace is None or (
|
||||
candidate_workspace.numel() * candidate_workspace.element_size()
|
||||
> current_workspace.numel() * current_workspace.element_size()
|
||||
):
|
||||
graph_params.workspaces[num_tokens] = candidate_workspace
|
||||
elif current_workspace is None:
|
||||
graph_params.workspaces[num_tokens] = candidate_workspace
|
||||
return graph_params.workspaces[num_tokens]
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def needs_layer_aware_fia_graph_replay() -> bool:
|
||||
vllm_config = get_current_vllm_config()
|
||||
model_config = vllm_config.model_config
|
||||
hf_config = getattr(model_config, "hf_config", None)
|
||||
hf_text_config = getattr(model_config, "hf_text_config", None)
|
||||
text_config = getattr(hf_config, "text_config", None)
|
||||
model_types = (
|
||||
getattr(hf_config, "model_type", None),
|
||||
getattr(hf_text_config, "model_type", None),
|
||||
getattr(text_config, "model_type", None),
|
||||
)
|
||||
return any(model_type in {"gemma4", "gemma4_text"} for model_type in model_types)
|
||||
|
||||
|
||||
def ascend_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int:
|
||||
scheduler_config = vllm_config.scheduler_config
|
||||
cache_config = vllm_config.cache_config
|
||||
model_config = vllm_config.model_config
|
||||
|
||||
chunked_prefill_workspace_size = min(
|
||||
# Make sure there is enough for 8 full length request or at least
|
||||
# 4 pages of cache per request
|
||||
max(8 * model_config.max_model_len, 4 * scheduler_config.max_num_seqs * cache_config.block_size),
|
||||
# For long-context models try not to over-allocate limiting
|
||||
# kv-cache space, limiting it to 128k tokens,
|
||||
# which would result in the workspace being:
|
||||
# 2*(576)*(128*1024) = 288mb
|
||||
# (assuming 576 MLA head dim, and fp16)
|
||||
# which would result in up-projected context being
|
||||
# 2*(192*128)*(128*1024) = 6gb
|
||||
# (assuming 192 QK head dim, 128 heads, and fp16)
|
||||
128 * 1024,
|
||||
)
|
||||
|
||||
chunked_prefill_workspace_size = max(
|
||||
chunked_prefill_workspace_size,
|
||||
scheduler_config.max_num_seqs * cache_config.block_size,
|
||||
)
|
||||
|
||||
return chunked_prefill_workspace_size
|
||||
|
||||
|
||||
def using_paged_attention(runtime_shape: int, vllm_config: VllmConfig, head_size: int | None = None) -> bool:
|
||||
if vllm_config.speculative_config is not None:
|
||||
return False
|
||||
if get_ascend_device_type() == AscendDeviceType.A5:
|
||||
return False
|
||||
# TODO: Remove this fallback when A2/A3 FIA TND supports Gemma4's
|
||||
# 512-dim global attention heads. Decode can use PA directly; prefill is
|
||||
# handled by the device adaptor.
|
||||
if head_size == FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE:
|
||||
return True
|
||||
from vllm.config.compilation import CUDAGraphMode
|
||||
|
||||
cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
|
||||
if cudagraph_mode != CUDAGraphMode.FULL_DECODE_ONLY:
|
||||
return False
|
||||
|
||||
return runtime_shape in get_ascend_config().pa_shape_list
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def enable_cp():
|
||||
prefill_config = get_current_vllm_config().parallel_config
|
||||
return prefill_config.prefill_context_parallel_size > 1 or prefill_config.decode_context_parallel_size > 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class AscendCommonAttentionMetadata:
|
||||
class AscendPrefillContextParallelMetadata:
|
||||
"""
|
||||
Metadata for Prefill Context Parallelism (PCP) in CommonAttentionMetadata.
|
||||
|
||||
Contains index tensors and sequence lengths for PCP operations.
|
||||
"""
|
||||
|
||||
pcp_allgather_restore_idx: torch.Tensor = None
|
||||
|
||||
num_actual_tokens_pcp_padded: int = 0
|
||||
|
||||
num_computed_tokens_of_pcp_dcp: list[list[list[int]]] | None = None
|
||||
|
||||
q_head_idx_tensor: torch.Tensor = None
|
||||
|
||||
q_tail_idx_tensor: torch.Tensor = None
|
||||
|
||||
kv_with_q_head_nomask_idx_tensor: torch.Tensor = None
|
||||
|
||||
kv_with_q_head_mask_idx_tensor: torch.Tensor = None
|
||||
|
||||
kv_with_q_tail_nomask_idx_tensor: torch.Tensor = None
|
||||
|
||||
kv_with_q_tail_mask_idx_tensor: torch.Tensor = None
|
||||
|
||||
kv_tail_proj_idx_tensor: torch.Tensor = None
|
||||
|
||||
kv_with_q_head_attn_idx_in_tail_tensor: torch.Tensor = None
|
||||
|
||||
kv_with_q_tail_attn_idx_in_tail_tensor: torch.Tensor = None
|
||||
|
||||
attn_mask_seqlens: torch.Tensor = None
|
||||
|
||||
head_attn_nomask_seqlens: torch.Tensor = None
|
||||
|
||||
tail_attn_nomask_seqlens: torch.Tensor = None
|
||||
|
||||
head_actual_seq_lengths_kv: list[int] | None = None
|
||||
|
||||
tail_actual_seq_lengths_kv: list[int] | None = None
|
||||
|
||||
q_full_idx: torch.Tensor = None
|
||||
|
||||
# original query_lens before pcp split
|
||||
query_lens_pcp_full_cpu: torch.Tensor = None
|
||||
|
||||
# original max_query_len before pcp split
|
||||
max_query_len_pcp_full: int = 0
|
||||
|
||||
# the following attributes are specifically used in hybrid-attn models.
|
||||
pcp_use_hybrid_attn: bool = False
|
||||
|
||||
pcp_unpad_mask: torch.Tensor = None
|
||||
|
||||
# to get the right order of query in prefill per rank
|
||||
pcp_fa_query_idx: torch.Tensor = None
|
||||
|
||||
# restore the full sequence across all pcp ranks
|
||||
# when entering from linear-attention to attention
|
||||
pcp_enter_fa_restore_idx: torch.Tensor = None
|
||||
|
||||
# restore the original FA padded layout without boolean-mask scatter
|
||||
pcp_fa_padding_restore_idx: torch.Tensor = None
|
||||
|
||||
# scatter the full sequence across all pcp ranks
|
||||
# when exiting from attention to linear-attention
|
||||
pcp_exit_fa_scatter_idx: torch.Tensor = None
|
||||
|
||||
# the number of tokens padded in linear-attn per rank
|
||||
pcp_padded_tokens_fla: int = 0
|
||||
|
||||
# the max number of unpadded tokens in all ranks
|
||||
max_num_tokens_across_pcp: int = 0
|
||||
|
||||
# the number of scheduled tokens on the current rank before padding
|
||||
total_num_scheduled_tokens: int = 0
|
||||
|
||||
# Because the sequence shard in linear attention layers does not include padding,
|
||||
# the full attention layers cannot obtain the correct query_lens with pcp pad for
|
||||
# chunked prefill calculation. Therefore, this value needs to be passed to the backend.
|
||||
# TODO:To be refactored.
|
||||
attn_chunk_seqlens: torch.Tensor = None
|
||||
dcp_mtp_attn_mask: torch.Tensor = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AscendCommonAttentionMetadata(CommonAttentionMetadata):
|
||||
"""
|
||||
Per-batch attention metadata, shared across layers and backends.
|
||||
AttentionMetadataBuilder instances use it to construct per-layer metadata.
|
||||
|
||||
For many of the tensors we keep both GPU and CPU versions.
|
||||
|
||||
For many of the tensors we keep both NPU and CPU versions.
|
||||
"""
|
||||
|
||||
query_start_loc: torch.Tensor
|
||||
query_start_loc_cpu: torch.Tensor
|
||||
"""(batch_size + 1,), the start location of each request in query Tensor"""
|
||||
# CPU tensor of sequence lengths for host-side operations.
|
||||
# E.g., tensor([128, 256, 64]) for 3 requests with different seq lengths.
|
||||
seq_lens_cpu: torch.Tensor = None
|
||||
|
||||
seq_lens_cpu: torch.Tensor
|
||||
"""(batch_size,), the length of each request including both computed tokens
|
||||
and newly scheduled tokens"""
|
||||
# CPU tensor of already computed tokens count per request.
|
||||
# E.g., tensor([100, 200, 50]) means req0 has 100 tokens already computed.
|
||||
num_computed_tokens_cpu: torch.Tensor = None
|
||||
|
||||
seq_lens: torch.Tensor
|
||||
"""same to seq_lens_cpu, for compatibility with some new attn metadata
|
||||
(such as GDN)."""
|
||||
# Number of decode tokens per request, used for speculative decoding.
|
||||
# E.g., 1 for normal decoding, >1 for speculative decoding.
|
||||
decode_token_per_req: int = 1
|
||||
|
||||
num_computed_tokens_cpu: torch.Tensor
|
||||
"""(batch_size,), the number of computed tokens for each request"""
|
||||
|
||||
num_reqs: int
|
||||
"""Number of requests"""
|
||||
num_actual_tokens: int
|
||||
"""Total number of tokens in batch"""
|
||||
|
||||
max_query_len: int
|
||||
"""Max token number of request in batch"""
|
||||
|
||||
decode_token_per_req: int
|
||||
"""decode token number per request"""
|
||||
|
||||
block_table_tensor: torch.Tensor
|
||||
|
||||
slot_mapping: torch.Tensor
|
||||
|
||||
actual_seq_lengths_q: list[int]
|
||||
# Actual query sequence lengths for each token in the batch (CPU list).
|
||||
# E.g., [1, 1, 1, 128] for 3 decode tokens and 1 prefill with 128 tokens.
|
||||
actual_seq_lengths_q: list[int] = field(default_factory=list)
|
||||
|
||||
# NPU tensor of position indices for rotary embeddings computation.
|
||||
# E.g., tensor([0, 1, 2, ...]) indicating token positions in sequence.
|
||||
positions: torch.Tensor = None
|
||||
positions_cpu: torch.Tensor = None
|
||||
|
||||
attn_mask: torch.Tensor = None
|
||||
|
||||
spec_attn_mask: torch.Tensor = None
|
||||
|
||||
# Current attention state (e.g., ChunkedPrefill, DecodeOnly).
|
||||
attn_state: Any = None
|
||||
|
||||
enable_dbo_across_dp: bool = False
|
||||
|
||||
is_only_prefill: bool = False
|
||||
|
||||
# Padding size for graph capture, -1 means not in graph mode.
|
||||
graph_pad_size: int = -1
|
||||
|
||||
# Total number of tokens including padding, used for padding operations.
|
||||
num_input_tokens: int = 0
|
||||
|
||||
# Metadata for Prefill Context Parallelism (PCP) operations.
|
||||
prefill_context_parallel_metadata: AscendPrefillContextParallelMetadata | None = None
|
||||
kvcomp_metadata: KVCompMetaData | None = None
|
||||
|
||||
# TODO: Remove it when vLLM no longer uses this function.
|
||||
def unpadded(self, num_actual_tokens: int, num_actual_reqs: int) -> "AscendCommonAttentionMetadata":
|
||||
# This only use to eagle now. It will be use to enforce_eager in future.
|
||||
# Helper to slice optional per-request tensors to ``num_actual_reqs``.
|
||||
def _slice_reqs(x):
|
||||
return x[:num_actual_reqs] if x is not None else None
|
||||
|
||||
return AscendCommonAttentionMetadata(
|
||||
query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
|
||||
query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
|
||||
seq_lens=self.seq_lens[:num_actual_reqs],
|
||||
seq_lens_cpu=_slice_reqs(self.seq_lens_cpu),
|
||||
num_computed_tokens_cpu=_slice_reqs(self.num_computed_tokens_cpu),
|
||||
num_reqs=num_actual_reqs,
|
||||
num_actual_tokens=num_actual_tokens,
|
||||
max_query_len=self.max_query_len,
|
||||
decode_token_per_req=self.decode_token_per_req,
|
||||
# NOTE: keep all tokens for block_table_tensor and slot_mapping otherwise
|
||||
# there will be error about shape mismatch during reshape and cache.
|
||||
# This is really strange since vLLM slices them as well
|
||||
block_table_tensor=self.block_table_tensor,
|
||||
slot_mapping=self.slot_mapping,
|
||||
causal=self.causal,
|
||||
actual_seq_lengths_q=self.actual_seq_lengths_q[:num_actual_tokens],
|
||||
positions=self.positions,
|
||||
positions_cpu=self.positions_cpu,
|
||||
attn_state=self.attn_state,
|
||||
graph_pad_size=-1, # It should be -1 when not run in fullgraph mode.
|
||||
num_input_tokens=self.num_input_tokens,
|
||||
prefill_context_parallel_metadata=self.prefill_context_parallel_metadata,
|
||||
seq_lens_cpu_upper_bound=self.seq_lens_cpu_upper_bound[:num_actual_reqs]
|
||||
if self.seq_lens_cpu_upper_bound is not None
|
||||
else None,
|
||||
max_seq_len=self.max_seq_len,
|
||||
# Propagate parent-class fields so the unpadded view is a
|
||||
# faithful sub-batch of the original. Missing any of these
|
||||
# would silently break downstream consumers (e.g. NPU
|
||||
# backends preferring ``_seq_lens_cpu`` over ``seq_lens_cpu``,
|
||||
# DCP backends needing ``dcp_local_seq_lens(_cpu)``,
|
||||
# encoder-decoder layers needing ``encoder_seq_lens``, the
|
||||
# mamba ``is_prefilling`` flag, and FastPrefill's
|
||||
# ``logits_indices_padded`` / ``num_logits_indices``).
|
||||
_seq_lens_cpu=_slice_reqs(self._seq_lens_cpu),
|
||||
_num_computed_tokens_cpu=_slice_reqs(self._num_computed_tokens_cpu),
|
||||
dcp_local_seq_lens=_slice_reqs(self.dcp_local_seq_lens),
|
||||
dcp_local_seq_lens_cpu=_slice_reqs(self.dcp_local_seq_lens_cpu),
|
||||
is_prefilling=_slice_reqs(self.is_prefilling),
|
||||
encoder_seq_lens=_slice_reqs(self.encoder_seq_lens),
|
||||
encoder_seq_lens_cpu=_slice_reqs(self.encoder_seq_lens_cpu),
|
||||
logits_indices_padded=self.logits_indices_padded,
|
||||
num_logits_indices=self.num_logits_indices,
|
||||
)
|
||||
|
||||
|
||||
def filter_chunked_req_indices(
|
||||
seq_len: torch.Tensor,
|
||||
mask_for_non_zero_chunk: list[bool] | None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
filter the reqs which are doing real chunk_prefill.
|
||||
|
||||
Args:
|
||||
seq_len: contains multi-req length: [req0_len, req1_len, ...]
|
||||
mask_for_non_zero_chunk: [True, False, True, False, ...]
|
||||
Returns:
|
||||
filtered_indices: the real chunked req's indices
|
||||
"""
|
||||
assert mask_for_non_zero_chunk is not None and len(seq_len) == len(mask_for_non_zero_chunk)
|
||||
offsets = torch.cumsum(torch.cat([torch.tensor([0]), seq_len[:-1]]), dim=0)
|
||||
filtered_indices = torch.cat(
|
||||
[
|
||||
torch.arange(offsets[i], offsets[i] + seq_len[i])
|
||||
for i in range(len(mask_for_non_zero_chunk))
|
||||
if mask_for_non_zero_chunk[i]
|
||||
]
|
||||
)
|
||||
return filtered_indices
|
||||
|
||||
|
||||
def split_decodes_and_prefills(
|
||||
common_attn_metadata: AscendCommonAttentionMetadata,
|
||||
decode_threshold: int = 1,
|
||||
require_uniform: bool = False,
|
||||
treat_short_extends_as_decodes: bool = True,
|
||||
) -> tuple[int, int, int, int]:
|
||||
"""
|
||||
Assuming a reordered batch, finds the boundary between prefill and decode
|
||||
requests.
|
||||
While pcp > 1, query_lens is split across pcp ranks, so we pass in the
|
||||
original query_lens and max_query_len to distinguish prefills and decodes.
|
||||
|
||||
The batch is expected to be ordered as:
|
||||
decode -> short_extend -> long_extend -> prefill
|
||||
|
||||
Args:
|
||||
common_attn_metadata: AscendCommonAttentionMetadata object containing the
|
||||
batch metadata.
|
||||
decode_threshold: The maximum query length to be considered a decode.
|
||||
require_uniform: If True, requires that all decode requests have the
|
||||
same query length. When set, some queries may be considered
|
||||
prefills even if they are <= decode_threshold, in order to ensure
|
||||
uniformity.
|
||||
treat_short_extends_as_decodes: If True (default), short extends
|
||||
(query_len <= threshold but still prefilling) are counted as
|
||||
decodes. If False, they are counted as prefills.
|
||||
|
||||
Returns:
|
||||
num_decodes: The number of decode requests.
|
||||
@@ -83,22 +367,57 @@ def split_decodes_and_prefills(
|
||||
num_decode_tokens: The number of tokens in the decode requests.
|
||||
num_prefill_tokens: The number of tokens in the prefill requests.
|
||||
"""
|
||||
max_query_len = common_attn_metadata.max_query_len
|
||||
long_seq_metadata = common_attn_metadata.prefill_context_parallel_metadata
|
||||
query_lens_pcp_full = long_seq_metadata.query_lens_pcp_full_cpu if long_seq_metadata else None
|
||||
max_query_len_pcp_full = long_seq_metadata.max_query_len_pcp_full if long_seq_metadata else 0
|
||||
max_query_len = common_attn_metadata.max_query_len if max_query_len_pcp_full == 0 else max_query_len_pcp_full
|
||||
num_reqs = common_attn_metadata.num_reqs
|
||||
if num_reqs == 0:
|
||||
return 0, 0, 0, 0
|
||||
|
||||
num_tokens = common_attn_metadata.num_actual_tokens
|
||||
query_start_loc = common_attn_metadata.query_start_loc_cpu
|
||||
|
||||
if max_query_len <= decode_threshold:
|
||||
# PD D + RecomputeScheduler: num_computed may be N-1 after KV recv while
|
||||
# this step is MTP decode (max_query_len <= threshold).
|
||||
if is_pd_decode_recompute_scheduler_enabled():
|
||||
treat_short_extends_as_decodes = True
|
||||
|
||||
if (
|
||||
max_query_len <= decode_threshold
|
||||
and (not require_uniform or decode_threshold <= 1)
|
||||
and treat_short_extends_as_decodes
|
||||
):
|
||||
return num_reqs, 0, num_tokens, 0
|
||||
|
||||
query_lens = query_start_loc[1:] - query_start_loc[:-1]
|
||||
is_prefill = query_lens > decode_threshold
|
||||
query_lens_sharded = query_start_loc[1:] - query_start_loc[:-1]
|
||||
query_lens = query_lens_sharded if query_lens_pcp_full is None else query_lens_pcp_full
|
||||
if query_lens[0].item() > decode_threshold:
|
||||
return 0, num_reqs, 0, num_tokens
|
||||
|
||||
if require_uniform:
|
||||
if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
|
||||
return num_reqs, 0, num_tokens, 0
|
||||
is_prefill = query_lens != query_lens[0]
|
||||
else:
|
||||
is_prefill = query_lens > decode_threshold
|
||||
|
||||
if not treat_short_extends_as_decodes:
|
||||
assert common_attn_metadata.is_prefilling is not None
|
||||
raw_is_prefilling = common_attn_metadata.is_prefilling
|
||||
is_prefilling = raw_is_prefilling[: query_lens.shape[0]]
|
||||
if is_prefilling.shape[0] < query_lens.shape[0]:
|
||||
is_prefilling = F.pad(
|
||||
is_prefilling,
|
||||
(0, query_lens.shape[0] - is_prefilling.shape[0]),
|
||||
value=False,
|
||||
)
|
||||
is_prefill |= is_prefilling
|
||||
|
||||
if not torch.any(is_prefill):
|
||||
return num_reqs, 0, num_tokens, 0
|
||||
|
||||
first_prefill = is_prefill.int().argmax(dim=-1).item()
|
||||
assert torch.all(query_lens[first_prefill:] >= decode_threshold)
|
||||
assert torch.all(query_lens[:first_prefill] <= decode_threshold)
|
||||
num_decodes = first_prefill
|
||||
num_prefills = num_reqs - num_decodes
|
||||
num_decode_tokens = query_start_loc[first_prefill].item()
|
||||
@@ -122,7 +441,7 @@ def wait_for_kv_layer_from_connector(layer_name: str):
|
||||
|
||||
def maybe_save_kv_layer_to_connector(
|
||||
layer_name: str,
|
||||
kv_cache_layer: List[torch.Tensor],
|
||||
kv_cache_layer: list[torch.Tensor],
|
||||
):
|
||||
if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
|
||||
return
|
||||
@@ -135,3 +454,68 @@ def maybe_save_kv_layer_to_connector(
|
||||
return
|
||||
# TODO: assert ascendMetadata
|
||||
connector.save_kv_layer(layer_name, kv_cache_layer, attn_metadata)
|
||||
|
||||
|
||||
def notify_kv_cache_written(layer_name: str = ""):
|
||||
"""Notify the connector that the paged KV cache for ``layer_name`` has been
|
||||
written for the current step.
|
||||
|
||||
The attention layer calls this unconditionally; each connector decides whether
|
||||
it needs to record a synchronization primitive (e.g. a compute-stream event
|
||||
later waited on by the resharding stream to overlap the outgoing KV copy).
|
||||
Connectors that don't need it -- such as the AscendStore pool connector, which
|
||||
records its own sync event at save time -- simply do not implement
|
||||
``on_kv_cache_written`` and this becomes a no-op.
|
||||
"""
|
||||
if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
|
||||
return
|
||||
|
||||
connector = get_kv_transfer_group()
|
||||
on_kv_cache_written = getattr(connector, "on_kv_cache_written", None)
|
||||
if on_kv_cache_written is not None:
|
||||
on_kv_cache_written(layer_name)
|
||||
|
||||
|
||||
def round_up(val: int, align: int) -> int:
|
||||
if align == 0:
|
||||
return 0
|
||||
return -(val // -align) * align
|
||||
|
||||
|
||||
def trans_rope_weight(weight, rope_dim):
|
||||
if rope_dim == 0:
|
||||
return weight.contiguous()
|
||||
nope_part = weight[..., :-rope_dim, :]
|
||||
rope_part = weight[..., -rope_dim:, :]
|
||||
reordered_rope_part = torch.cat((rope_part[..., ::2, :], rope_part[..., 1::2, :]), dim=-2)
|
||||
return torch.cat((nope_part, reordered_rope_part), dim=-2).contiguous()
|
||||
|
||||
|
||||
def transdata(nd_mat, block_size: tuple = (16, 16)):
|
||||
r = round_up(nd_mat.shape[0], block_size[0])
|
||||
c = round_up(nd_mat.shape[1], block_size[1])
|
||||
r_pad = r - nd_mat.shape[0]
|
||||
c_pad = c - nd_mat.shape[1]
|
||||
nd_mat = F.pad(nd_mat, (0, r_pad, 0, c_pad))
|
||||
nz_mat = torch.permute(
|
||||
torch.reshape(
|
||||
nd_mat,
|
||||
(r // block_size[0], block_size[0], c // block_size[1], block_size[1]),
|
||||
),
|
||||
[2, 0, 1, 3],
|
||||
)
|
||||
nz_mat = torch.reshape(nz_mat, (nz_mat.shape[0], nz_mat.shape[1] * nz_mat.shape[2], nz_mat.shape[3]))
|
||||
return nz_mat
|
||||
|
||||
|
||||
def enabling_mlapo(vllm_config: VllmConfig) -> bool:
|
||||
config_val = get_ascend_config().enable_mlapo
|
||||
if get_ascend_device_type() == AscendDeviceType.A5:
|
||||
return bool(config_val)
|
||||
|
||||
is_decode_instance = (
|
||||
vllm_config.kv_transfer_config is not None
|
||||
and vllm_config.kv_transfer_config.is_kv_consumer
|
||||
and not vllm_config.kv_transfer_config.is_kv_producer
|
||||
)
|
||||
return bool(config_val and is_decode_instance)
|
||||
|
||||
Reference in New Issue
Block a user