1024 lines
35 KiB
Python
1024 lines
35 KiB
Python
#
|
|
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2024 The vLLM team.
|
|
#
|
|
# 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.
|
|
#
|
|
"""NPU precision tests: ``AscendMlaCPImpl`` vs fp32 reference (``mla_cp.py``)."""
|
|
|
|
import math
|
|
import sys
|
|
from contextlib import ExitStack
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
if "torch_npu._inductor" not in sys.modules:
|
|
sys.modules["torch_npu._inductor"] = MagicMock()
|
|
|
|
from vllm.config import VllmConfig # noqa: E402
|
|
|
|
from tests.ut.attention.utils import BatchSpec, create_vllm_config # noqa: E402
|
|
from vllm_ascend.attention.attention_v1 import AscendAttentionState # noqa: E402
|
|
from vllm_ascend.attention.context_parallel.mla_cp import AscendMlaCPImpl # noqa: E402
|
|
from vllm_ascend.attention.mla_v1 import AscendMLAImpl # noqa: E402
|
|
|
|
_BLOCK_SIZE = 128
|
|
_TEST_NUM_HEADS = 16
|
|
_KV_LORA_RANK = 512
|
|
_QK_NOPE_HEAD_DIM = 128
|
|
_QK_ROPE_HEAD_DIM = 64
|
|
_V_HEAD_DIM = 128
|
|
|
|
DEFAULT_RTOL = 1e-2
|
|
DEFAULT_ATOL = 1e-2
|
|
# fp16 carries an extra ULP of cumulative rounding through MLA's W_UV
|
|
# down-projection that SFA does not have, so single edge-cases at
|
|
# near-zero reference values can land 1 ULP outside the bf16 envelope.
|
|
# Widening here to 2e-2 still represents <0.5% peak-relative error
|
|
# (peaks are ~4 in mixed_prefill) -- well inside the "<1%" requirement
|
|
# enforced by the signal-relative assertion in ``_record_and_assert``.
|
|
FP16_RTOL = 2e-2
|
|
FP16_ATOL = 2e-2
|
|
|
|
_MAX_SIG_REL_ERR = 1e-2 # max |out-ref| / peak |ref|
|
|
_MAX_MEAN_SIG_ERR = 5e-3 # mean |out-ref| / mean |ref|
|
|
_MAX_REL_ERR = 1e-2 # max per-element rel err where |ref| >= floor
|
|
_SIG_FLOOR_FRAC = 5e-1 # floor = this fraction of peak |ref|
|
|
|
|
|
|
def _validate_spec(spec: BatchSpec) -> None:
|
|
"""Require ``query_len <= seq_len`` per request."""
|
|
for s, q in zip(spec.seq_lens, spec.query_lens):
|
|
assert q <= s, f"query_len ({q}) must not exceed seq_len ({s})"
|
|
|
|
|
|
_VLLM_CONFIG_CACHE: dict = {}
|
|
|
|
|
|
def _get_vllm_config(
|
|
model: str,
|
|
dtype: torch.dtype,
|
|
*,
|
|
max_model_len: int = 8192,
|
|
tensor_parallel_size: int = 1,
|
|
) -> VllmConfig:
|
|
"""Cached ``VllmConfig`` for MLA-CP (fp8 quant stripped; heads capped).
|
|
|
|
Mirrors the ``test_sfa_cp_precision.py`` strategy: ``tensor_parallel_size > 1``
|
|
is **simulated** by keeping real TP=1 on the underlying ``VllmConfig``
|
|
(no distributed init required) and instead dividing the test-friendly
|
|
head count by TP so the kernel runs with the post-TP-shard head budget.
|
|
"""
|
|
key = (model, dtype, tensor_parallel_size)
|
|
cfg = _VLLM_CONFIG_CACHE.get(key)
|
|
if cfg is not None:
|
|
return cfg
|
|
dtype_str = "bfloat16" if dtype == torch.bfloat16 else "float16"
|
|
sim_num_heads = max(1, _TEST_NUM_HEADS // tensor_parallel_size)
|
|
cfg = create_vllm_config(
|
|
model_name=model,
|
|
tensor_parallel_size=1, # always TP=1; head split is simulated
|
|
max_model_len=max_model_len,
|
|
dtype=dtype_str,
|
|
block_size=_BLOCK_SIZE,
|
|
num_gpu_blocks=4096,
|
|
max_num_seqs=256,
|
|
max_num_batched_tokens=max(8192, max_model_len * 2),
|
|
enable_chunked_prefill=True,
|
|
hf_overrides={"quantization_config": None},
|
|
hf_config_override={
|
|
"num_attention_heads": sim_num_heads,
|
|
"num_key_value_heads": 1,
|
|
"kv_lora_rank": _KV_LORA_RANK,
|
|
"qk_nope_head_dim": _QK_NOPE_HEAD_DIM,
|
|
"qk_rope_head_dim": _QK_ROPE_HEAD_DIM,
|
|
"v_head_dim": _V_HEAD_DIM,
|
|
},
|
|
)
|
|
_VLLM_CONFIG_CACHE[key] = cfg
|
|
return cfg
|
|
|
|
|
|
# Spec name prefixes drive the MLA-CP branch under test.
|
|
BATCH_SPECS: dict[str, BatchSpec] = {
|
|
"decode_single": BatchSpec(seq_lens=[1024], query_lens=[1], name="decode_single"),
|
|
"decode_small_batch": BatchSpec(seq_lens=[64, 128, 256, 512], query_lens=[1, 1, 1, 1], name="decode_small_batch"),
|
|
"decode_large_batch": BatchSpec(seq_lens=[2048] * 16, query_lens=[1] * 16, name="decode_large_batch"),
|
|
"mtp_1_plus_1": BatchSpec(seq_lens=[256, 512, 1024], query_lens=[2, 2, 2], name="mtp_1_plus_1"),
|
|
"mtp_1_plus_3": BatchSpec(seq_lens=[256, 512, 1024, 1536], query_lens=[4, 4, 4, 4], name="mtp_1_plus_3"),
|
|
"mtp_1_plus_7": BatchSpec(seq_lens=[512, 1024, 2048], query_lens=[8, 8, 8], name="mtp_1_plus_7"),
|
|
"prefill_single": BatchSpec(seq_lens=[256], query_lens=[256], name="prefill_single"),
|
|
"prefill_small_batch": BatchSpec(seq_lens=[128, 256, 384], query_lens=[128, 256, 384], name="prefill_small_batch"),
|
|
"prefill_medium_batch": BatchSpec(seq_lens=[512, 1024], query_lens=[512, 1024], name="prefill_medium_batch"),
|
|
"mixed_small": BatchSpec(seq_lens=[64, 128, 256, 512], query_lens=[1, 1, 64, 128], name="mixed_small"),
|
|
"mixed_medium": BatchSpec(
|
|
seq_lens=[1024, 1536, 2048, 256, 512], query_lens=[1, 1, 1, 64, 128], name="mixed_medium"
|
|
),
|
|
}
|
|
|
|
|
|
def _infer_mode(spec: BatchSpec) -> str:
|
|
"""Return one of: ``decode``, ``mtp``, ``prefill``, ``mixed``."""
|
|
name = spec.name
|
|
for prefix in ("decode_", "mtp_", "prefill_", "mixed_"):
|
|
if name.startswith(prefix):
|
|
return prefix.rstrip("_")
|
|
raise ValueError(
|
|
f"BatchSpec name {name!r} does not start with a known mode prefix ('decode_', 'mtp_', 'prefill_', 'mixed_')"
|
|
)
|
|
|
|
|
|
def _make_synthetic_kv_contexts(
|
|
seq_lens: list[int],
|
|
kv_lora_rank: int,
|
|
qk_rope_head_dim: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
|
|
"""Per-request K context (latent space) generated independently so values
|
|
across requests are uncorrelated."""
|
|
k_nope = [torch.randn(s, kv_lora_rank, dtype=dtype, device=device) for s in seq_lens]
|
|
k_pe = [torch.randn(s, qk_rope_head_dim, dtype=dtype, device=device) for s in seq_lens]
|
|
return k_nope, k_pe
|
|
|
|
|
|
def _build_paged_kv_cache(
|
|
seq_lens: list[int],
|
|
k_nope_contexts: list[torch.Tensor],
|
|
k_pe_contexts: list[torch.Tensor],
|
|
block_size: int,
|
|
kv_lora_rank: int,
|
|
qk_rope_head_dim: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Paged ``k_nope`` / ``k_pe`` caches for the decode / MTP path."""
|
|
batch_size = len(seq_lens)
|
|
blocks_per_seq = [(s + block_size - 1) // block_size for s in seq_lens]
|
|
total_blocks = sum(blocks_per_seq) + 1
|
|
max_blocks_per_seq = max(blocks_per_seq)
|
|
|
|
k_nope_cache = torch.zeros(
|
|
total_blocks,
|
|
block_size,
|
|
1,
|
|
kv_lora_rank,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
k_pe_cache = torch.zeros(
|
|
total_blocks,
|
|
block_size,
|
|
1,
|
|
qk_rope_head_dim,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
block_table = torch.zeros(
|
|
batch_size,
|
|
max_blocks_per_seq,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
|
|
next_block_id = 1
|
|
for b, s_len in enumerate(seq_lens):
|
|
n_blocks = blocks_per_seq[b]
|
|
for i in range(n_blocks):
|
|
block_id = next_block_id
|
|
block_table[b, i] = block_id
|
|
tok_start = i * block_size
|
|
tok_end = min(tok_start + block_size, s_len)
|
|
length = tok_end - tok_start
|
|
k_nope_cache[block_id, :length, 0, :] = k_nope_contexts[b][tok_start:tok_end]
|
|
k_pe_cache[block_id, :length, 0, :] = k_pe_contexts[b][tok_start:tok_end]
|
|
next_block_id += 1
|
|
|
|
return k_nope_cache, k_pe_cache, block_table
|
|
|
|
|
|
def _make_w_uv(
|
|
num_heads: int,
|
|
kv_lora_rank: int,
|
|
v_head_dim: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
return torch.randn(
|
|
num_heads,
|
|
kv_lora_rank,
|
|
v_head_dim,
|
|
dtype=dtype,
|
|
device=device,
|
|
) * (1.0 / math.sqrt(kv_lora_rank))
|
|
|
|
|
|
def _decode_reference(
|
|
q_nope_latent: torch.Tensor,
|
|
q_pe_latent: torch.Tensor,
|
|
k_nope_contexts: list[torch.Tensor],
|
|
k_pe_contexts: list[torch.Tensor],
|
|
seq_lens: list[int],
|
|
query_lens: list[int],
|
|
W_UV: torch.Tensor,
|
|
scale: float,
|
|
causal: bool,
|
|
num_heads: int,
|
|
v_head_dim: int,
|
|
out_dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
"""Pure-PyTorch dense MLA decode baseline in fp32 (latent-space attention
|
|
followed by ``W_UV`` projection back into ``v_head_dim``)."""
|
|
outputs: list[torch.Tensor] = []
|
|
cum_q = 0
|
|
for b, (s_len, q_len) in enumerate(zip(seq_lens, query_lens)):
|
|
ctx_len = s_len - q_len
|
|
|
|
K_nope = k_nope_contexts[b].float()
|
|
K_pe = k_pe_contexts[b].float()
|
|
V_lat = K_nope
|
|
K_full = torch.cat([K_nope, K_pe], dim=-1)
|
|
|
|
for j in range(q_len):
|
|
t = cum_q + j
|
|
valid_end = ctx_len + j + 1 if causal else s_len
|
|
|
|
q_n = q_nope_latent[t].float()
|
|
q_p = q_pe_latent[t].float()
|
|
Q = torch.cat([q_n, q_p], dim=-1)
|
|
|
|
K_b = K_full[:valid_end]
|
|
V_b = V_lat[:valid_end]
|
|
|
|
scores = (Q @ K_b.transpose(0, 1)) * scale
|
|
attn = torch.softmax(scores, dim=-1)
|
|
outputs.append(attn @ V_b)
|
|
|
|
cum_q += q_len
|
|
|
|
O_lat = torch.stack(outputs, dim=0)
|
|
O_proj = torch.bmm(O_lat.transpose(0, 1), W_UV.float())
|
|
O_final = O_proj.transpose(0, 1).contiguous()
|
|
return O_final.reshape(O_final.shape[0], num_heads * v_head_dim).to(out_dtype)
|
|
|
|
|
|
def _prefill_reference(
|
|
q_nope_full: list[torch.Tensor],
|
|
q_pe_full: list[torch.Tensor],
|
|
k_nope_full_per_req: list[torch.Tensor],
|
|
k_pe_full_per_req: list[torch.Tensor],
|
|
v_full_per_req: list[torch.Tensor],
|
|
scale: float,
|
|
num_heads: int,
|
|
v_head_dim: int,
|
|
out_dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
"""Pure-PyTorch dense MLA prefill baseline in fp32 (causal)."""
|
|
outputs: list[torch.Tensor] = []
|
|
for q_nope, q_pe, k_nope, k_pe, v in zip(
|
|
q_nope_full,
|
|
q_pe_full,
|
|
k_nope_full_per_req,
|
|
k_pe_full_per_req,
|
|
v_full_per_req,
|
|
):
|
|
q_len = q_nope.shape[0]
|
|
Q = torch.cat([q_nope, q_pe], dim=-1).float()
|
|
K = torch.cat([k_nope, k_pe], dim=-1).float()
|
|
V = v.float()
|
|
|
|
scores = (
|
|
torch.matmul(
|
|
Q.transpose(0, 1),
|
|
K.transpose(0, 1).transpose(-1, -2),
|
|
)
|
|
* scale
|
|
)
|
|
causal_mask = torch.triu(
|
|
torch.ones(q_len, q_len, dtype=torch.bool, device=Q.device),
|
|
diagonal=1,
|
|
)
|
|
scores = scores.masked_fill(causal_mask, float("-inf"))
|
|
attn = torch.softmax(scores, dim=-1)
|
|
o = torch.matmul(attn, V.transpose(0, 1)).transpose(0, 1).contiguous()
|
|
outputs.append(o.reshape(q_len, num_heads * v_head_dim).to(out_dtype))
|
|
|
|
return torch.cat(outputs, dim=0)
|
|
|
|
|
|
def _make_fake_cp_group(world_size: int) -> MagicMock:
|
|
"""Single-rank fake (PCP or DCP) collective group.
|
|
|
|
``all_gather(t, dim)`` is mocked as ``cat([t] * world_size, dim=dim)``
|
|
so downstream view / permute / npu_attention_update shapes line up.
|
|
For ``world_size == 1`` it is the identity.
|
|
"""
|
|
grp = MagicMock()
|
|
grp.world_size = world_size
|
|
grp.rank_in_group = 0
|
|
grp.device_group = MagicMock() if world_size > 1 else None
|
|
if world_size > 1:
|
|
grp.all_gather = MagicMock(
|
|
side_effect=lambda t, dim: torch.cat([t] * world_size, dim=dim),
|
|
)
|
|
else:
|
|
grp.all_gather = MagicMock(side_effect=lambda t, dim: t)
|
|
return grp
|
|
|
|
|
|
def _patch_distributed_groups_cp(pcp_size: int, dcp_size: int) -> list:
|
|
"""Patch CP distributed groups + collectives for single-rank simulation.
|
|
|
|
Mirrors the ``test_sfa_cp_precision.py`` strategy: ``pcp_size`` /
|
|
``dcp_size`` can be > 1 even though only one NPU is available, by
|
|
providing fake groups + side-effect mocks for the collectives that
|
|
``_process_attn_out_lse`` / ``_npu_attention_update`` exercise on the
|
|
decode path. ``torch.distributed.all_to_all_single`` is mocked as
|
|
``output.copy_(input)``: in a real multi-rank run, all-to-all permutes
|
|
head-group ownership across DCP ranks; with identical per-rank data
|
|
the no-op preserves the data the kernel-update pipeline expects.
|
|
"""
|
|
fake_pcp = _make_fake_cp_group(pcp_size)
|
|
fake_dcp = _make_fake_cp_group(dcp_size)
|
|
|
|
def _fake_all_to_all_single(output, input_, *args, **kwargs):
|
|
output.copy_(input_)
|
|
return None
|
|
|
|
common_cp = "vllm_ascend.attention.context_parallel.common_cp"
|
|
return [
|
|
patch(f"{common_cp}.get_pcp_group", return_value=fake_pcp),
|
|
patch(f"{common_cp}.get_dcp_group", return_value=fake_dcp),
|
|
patch(f"{common_cp}.get_decode_context_model_parallel_world_size", return_value=dcp_size),
|
|
patch("torch.distributed.all_to_all_single", side_effect=_fake_all_to_all_single),
|
|
]
|
|
|
|
|
|
def _patch_extra_ctx(module_path: str):
|
|
"""Patch ``_EXTRA_CTX`` so the kernel sees a benign forward context."""
|
|
fake_ctx = MagicMock()
|
|
fake_ctx.is_draft_model = False
|
|
fake_ctx.is_draft_model_prefill = False
|
|
fake_ctx.capturing = False
|
|
return patch(f"{module_path}._EXTRA_CTX", fake_ctx)
|
|
|
|
|
|
def _populate_impl_attrs(
|
|
impl: Any,
|
|
*,
|
|
scale: float,
|
|
num_heads: int,
|
|
kv_lora_rank: int,
|
|
qk_nope_head_dim: int,
|
|
qk_rope_head_dim: int,
|
|
v_head_dim: int,
|
|
W_UV: torch.Tensor,
|
|
vllm_config: VllmConfig,
|
|
speculative_config,
|
|
pcp_size: int,
|
|
dcp_size: int,
|
|
) -> None:
|
|
"""Populate ``AscendMlaCPImpl`` attributes that the kernel reads."""
|
|
impl.vllm_config = vllm_config
|
|
impl.scale = scale
|
|
impl.num_heads = num_heads
|
|
impl.num_heads_padded = 1 << (num_heads - 1).bit_length()
|
|
impl.head_padding = impl.num_heads_padded - num_heads
|
|
impl.num_kv_heads = 1
|
|
impl.kv_lora_rank = kv_lora_rank
|
|
impl.qk_nope_head_dim = qk_nope_head_dim
|
|
impl.qk_rope_head_dim = qk_rope_head_dim
|
|
impl.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
|
impl.v_head_dim = v_head_dim
|
|
impl.fa_quant_layer = False
|
|
impl.enable_kv_nz = False
|
|
impl.W_UV = W_UV
|
|
impl.layer_name = "test_layer"
|
|
impl.speculative_config = speculative_config
|
|
impl.pcp_size = pcp_size
|
|
impl.dcp_size = dcp_size
|
|
impl.pcp_rank = 0
|
|
impl.dcp_rank = 0
|
|
impl.pcp_group = None if pcp_size == 1 else MagicMock()
|
|
impl.dcp_group = None if dcp_size == 1 else MagicMock()
|
|
|
|
|
|
def _make_fake_self(*, dtype: torch.dtype, **kwargs) -> MagicMock:
|
|
"""``MagicMock`` self for the decode path (delegates ``_v_up_proj`` and
|
|
``_compute_prefill_context`` to the real unbound implementations)."""
|
|
fake_self = MagicMock()
|
|
_populate_impl_attrs(fake_self, **kwargs)
|
|
fake_self.dtype = dtype
|
|
fake_self._v_up_proj = lambda x: AscendMlaCPImpl._v_up_proj(fake_self, x)
|
|
fake_self._compute_prefill_context = lambda *a, **kw: AscendMLAImpl._compute_prefill_context(fake_self, *a, **kw)
|
|
return fake_self
|
|
|
|
|
|
def _make_real_impl(**kwargs) -> AscendMlaCPImpl:
|
|
"""Real ``AscendMlaCPImpl`` instance (used by the prefill path which
|
|
relies on inherited ``AscendMLAImpl._compute_prefill_context``)."""
|
|
impl = object.__new__(AscendMlaCPImpl)
|
|
_populate_impl_attrs(impl, **kwargs)
|
|
return impl
|
|
|
|
|
|
def _build_prefill_attn_mask(device: torch.device) -> torch.Tensor:
|
|
return torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8).to(device)
|
|
|
|
|
|
def _make_decode_metadata(
|
|
*,
|
|
seq_lens: list[int],
|
|
query_lens: list[int],
|
|
block_table: torch.Tensor,
|
|
attn_state,
|
|
attn_mask: torch.Tensor | None,
|
|
cp_seq_len: list[int],
|
|
) -> MagicMock:
|
|
num_decode_tokens = sum(query_lens)
|
|
decode_meta = MagicMock()
|
|
decode_meta.block_table = block_table
|
|
decode_meta.seq_lens_list = list(seq_lens)
|
|
decode_meta.actual_seq_lengths_q = list(range(1, num_decode_tokens + 1))
|
|
decode_meta.attn_mask = attn_mask
|
|
decode_meta.cp_seq_len = cp_seq_len
|
|
decode_meta.dcp_mtp_attn_mask = None
|
|
|
|
attn_metadata = MagicMock()
|
|
attn_metadata.attn_state = attn_state
|
|
attn_metadata.decode = decode_meta
|
|
attn_metadata.num_decodes = len(seq_lens)
|
|
attn_metadata.query_lens = query_lens
|
|
return attn_metadata
|
|
|
|
|
|
def _make_prefill_metadata(
|
|
*,
|
|
query_lens: list[int],
|
|
attn_mask: torch.Tensor,
|
|
) -> MagicMock:
|
|
actual_seq_lengths_q = [sum(query_lens[: i + 1]) for i in range(len(query_lens))]
|
|
prefill_meta = MagicMock()
|
|
prefill_meta.actual_seq_lengths_q = actual_seq_lengths_q
|
|
prefill_meta.attn_mask = attn_mask
|
|
prefill_meta.chunked_context = None
|
|
prefill_meta.pcp_metadata = None
|
|
|
|
attn_metadata = MagicMock()
|
|
attn_metadata.prefill = prefill_meta
|
|
return attn_metadata
|
|
|
|
|
|
def _run_mla_cp_decode_kernel(
|
|
*,
|
|
q_nope_latent: torch.Tensor,
|
|
q_pe_latent: torch.Tensor,
|
|
k_nope_cache: torch.Tensor,
|
|
k_pe_cache: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
seq_lens: list[int],
|
|
query_lens: list[int],
|
|
attn_state,
|
|
causal: bool,
|
|
speculative_config,
|
|
block_size: int,
|
|
pcp_size: int,
|
|
dcp_size: int,
|
|
dtype: torch.dtype,
|
|
impl_kwargs: dict,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
"""Drive ``AscendMlaCPImpl._forward_decode`` with mocked CP collectives.
|
|
|
|
Single-rank simulation of PCP / DCP > 1:
|
|
- We provide the FULL local KV (rank 0 "owns everything"),
|
|
``cp_seq_len`` = full per-request seq_lens, and the full block
|
|
table. The kernel produces the *full* attention for every head it
|
|
sees (no real per-rank chunking).
|
|
- When ``dcp_size > 1`` the production preprocess (``reorg_decode_q``)
|
|
all-gathers ``q_nope`` / ``q_pe`` along the head dim. We emulate
|
|
that by ``repeat(1, dcp_size, 1)`` so the kernel runs with
|
|
``num_heads * dcp_size`` heads (matching ``_forward_decode``'s own
|
|
head-count branch).
|
|
- The downstream ``_process_attn_out_lse`` + ``_npu_attention_update``
|
|
reduction over ``PCP * DCP`` contributions is therefore over
|
|
identical full-attention outputs (one per virtual rank), which
|
|
reduces back to the same full attention.
|
|
"""
|
|
attn_mask = _build_prefill_attn_mask(device) if causal else None
|
|
|
|
if attn_state == AscendAttentionState.SpecDecoding:
|
|
cp_seq_len: list[int] = []
|
|
for s_len, q_len in zip(seq_lens, query_lens):
|
|
for j in range(q_len):
|
|
cp_seq_len.append(s_len - q_len + j + 1)
|
|
per_token_rows = []
|
|
for b, q_len in enumerate(query_lens):
|
|
for _ in range(q_len):
|
|
per_token_rows.append(block_table[b])
|
|
block_table = torch.stack(per_token_rows, dim=0).contiguous()
|
|
else:
|
|
cp_seq_len = list(seq_lens)
|
|
|
|
attn_metadata = _make_decode_metadata(
|
|
seq_lens=seq_lens,
|
|
query_lens=query_lens,
|
|
block_table=block_table,
|
|
attn_state=attn_state,
|
|
attn_mask=attn_mask,
|
|
cp_seq_len=cp_seq_len,
|
|
)
|
|
|
|
q_nope = q_nope_latent
|
|
q_pe = q_pe_latent
|
|
if dcp_size > 1:
|
|
q_nope = q_nope.repeat(1, dcp_size, 1)
|
|
q_pe = q_pe.repeat(1, dcp_size, 1)
|
|
|
|
fake_self = _make_fake_self(
|
|
dtype=dtype,
|
|
speculative_config=speculative_config,
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
**impl_kwargs,
|
|
)
|
|
|
|
with ExitStack() as stack:
|
|
for p in _patch_distributed_groups_cp(pcp_size, dcp_size):
|
|
stack.enter_context(p)
|
|
stack.enter_context(_patch_extra_ctx("vllm_ascend.attention.context_parallel.mla_cp"))
|
|
return AscendMlaCPImpl._forward_decode(
|
|
fake_self,
|
|
q_nope,
|
|
q_pe,
|
|
k_nope_cache,
|
|
k_pe_cache,
|
|
block_size,
|
|
attn_metadata,
|
|
)
|
|
|
|
|
|
def _run_mla_cp_prefill_kernel(
|
|
*,
|
|
q_nope: torch.Tensor,
|
|
q_pe: torch.Tensor,
|
|
k_nope: torch.Tensor,
|
|
k_pe: torch.Tensor,
|
|
value: torch.Tensor,
|
|
k_nope_cache: torch.Tensor,
|
|
k_pe_cache: torch.Tensor,
|
|
query_lens: list[int],
|
|
pcp_size: int,
|
|
dcp_size: int,
|
|
impl_kwargs: dict,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
"""Drive ``AscendMlaCPImpl._forward_prefill`` with mocked CP collectives.
|
|
|
|
Single-rank simulation only covers ``pcp_size == 1`` (DCP > 1 OK);
|
|
PCP > 1 prefill requires per-rank PCP scheduler metadata that only a
|
|
real multi-rank job can produce faithfully (consistent with the SFA-CP
|
|
precision matrix, see the topology whitelist below).
|
|
"""
|
|
assert pcp_size == 1, "PCP > 1 prefill is out of scope"
|
|
|
|
attn_mask = _build_prefill_attn_mask(device)
|
|
attn_metadata = _make_prefill_metadata(
|
|
query_lens=query_lens,
|
|
attn_mask=attn_mask,
|
|
)
|
|
|
|
impl = _make_real_impl(
|
|
speculative_config=None,
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
**impl_kwargs,
|
|
)
|
|
|
|
with ExitStack() as stack:
|
|
for p in _patch_distributed_groups_cp(pcp_size, dcp_size):
|
|
stack.enter_context(p)
|
|
stack.enter_context(_patch_extra_ctx("vllm_ascend.attention.mla_v1"))
|
|
stack.enter_context(_patch_extra_ctx("vllm_ascend.attention.context_parallel.mla_cp"))
|
|
return AscendMlaCPImpl._forward_prefill(
|
|
impl,
|
|
q_nope,
|
|
q_pe,
|
|
k_nope,
|
|
k_pe,
|
|
value,
|
|
(k_nope_cache, k_pe_cache),
|
|
attn_metadata,
|
|
)
|
|
|
|
|
|
def _record_and_assert(
|
|
backend_output: torch.Tensor,
|
|
reference_output: torch.Tensor,
|
|
tag: str,
|
|
*,
|
|
dtype: torch.dtype,
|
|
atol: float = DEFAULT_ATOL,
|
|
rtol: float = DEFAULT_RTOL,
|
|
) -> tuple[float, float]:
|
|
"""Assert numerical closeness and record signal-relative metrics.
|
|
|
|
The MLA-CP decode kernel returns fp32 (online-softmax LSE reduction is
|
|
accumulated in fp32 and never down-cast), so we cast the backend output
|
|
to the reference dtype before comparing -- this differs from SFA-CP,
|
|
which preserves input dtype end-to-end.
|
|
"""
|
|
assert backend_output.shape == reference_output.shape, (
|
|
f"[{tag}] backend shape {tuple(backend_output.shape)} != reference shape {tuple(reference_output.shape)}"
|
|
)
|
|
assert torch.isfinite(backend_output).all(), f"[{tag}] MLA-CP attention produced non-finite values"
|
|
if backend_output.dtype != reference_output.dtype:
|
|
backend_output = backend_output.to(reference_output.dtype)
|
|
|
|
torch.testing.assert_close(
|
|
backend_output,
|
|
reference_output,
|
|
rtol=rtol,
|
|
atol=atol,
|
|
msg=lambda m: f"[MLA-CP:{tag}] kernel output diverges from baseline. {m}",
|
|
)
|
|
|
|
ref_f32 = reference_output.float()
|
|
out_f32 = backend_output.float()
|
|
diff = (out_f32 - ref_f32).abs()
|
|
ref_abs = ref_f32.abs()
|
|
peak = float(ref_abs.max())
|
|
mean_ref_abs = float(ref_abs.mean())
|
|
sig_floor = peak * _SIG_FLOOR_FRAC
|
|
|
|
max_abs_err = float(diff.max())
|
|
mean_abs_err = float(diff.mean())
|
|
max_sig_rel_err = max_abs_err / peak if peak > 0 else 0.0
|
|
mean_sig_rel_err = mean_abs_err / mean_ref_abs if mean_ref_abs > 0 else 0.0
|
|
|
|
significant_mask = ref_abs >= sig_floor
|
|
if significant_mask.any():
|
|
per_elem_rel = diff[significant_mask] / ref_abs[significant_mask]
|
|
max_rel_err_sig = float(per_elem_rel.max())
|
|
else:
|
|
max_rel_err_sig = 0.0
|
|
|
|
print(
|
|
f"[MLA-CP-precision] {tag} "
|
|
f"peak={peak:.4e} max_abs_err={max_abs_err:.4e} "
|
|
f"max_sig_rel_err={max_sig_rel_err * 100:.4f}% "
|
|
f"mean_sig_rel_err={mean_sig_rel_err * 100:.4f}% "
|
|
f"max_rel_err_sig(>={int(_SIG_FLOOR_FRAC * 100)}%peak)="
|
|
f"{max_rel_err_sig * 100:.4f}%"
|
|
)
|
|
|
|
assert max_sig_rel_err < _MAX_SIG_REL_ERR, (
|
|
f"[MLA-CP:{tag}] dtype={dtype} signal-relative max error "
|
|
f"{max_sig_rel_err * 100:.4f}% exceeds 1% budget "
|
|
f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})"
|
|
)
|
|
assert mean_sig_rel_err < _MAX_MEAN_SIG_ERR, (
|
|
f"[MLA-CP:{tag}] dtype={dtype} signal-relative mean error "
|
|
f"{mean_sig_rel_err * 100:.4f}% exceeds 0.5% drift budget "
|
|
f"(mean_ref_abs={mean_ref_abs:.4e}, mean_abs_err={mean_abs_err:.4e})"
|
|
)
|
|
assert max_rel_err_sig < _MAX_REL_ERR, (
|
|
f"[MLA-CP:{tag}] dtype={dtype} per-element relative error on "
|
|
f">={int(_SIG_FLOOR_FRAC * 100)}%-of-peak elements "
|
|
f"{max_rel_err_sig * 100:.4f}% exceeds 1% budget "
|
|
f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})"
|
|
)
|
|
|
|
return max_abs_err, max_rel_err_sig
|
|
|
|
|
|
def _test_mla_cp_correctness(
|
|
batch_spec: BatchSpec,
|
|
model: str,
|
|
*,
|
|
pcp_size: int,
|
|
dcp_size: int,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
atol: float = DEFAULT_ATOL,
|
|
rtol: float = DEFAULT_RTOL,
|
|
tensor_parallel_size: int = 1,
|
|
) -> None:
|
|
"""Test ``AscendMlaCPImpl`` against a fp32 dense MLA reference."""
|
|
mode = _infer_mode(batch_spec)
|
|
assert not (pcp_size > 1 and mode in ("prefill", "mixed")), (
|
|
f"PCP>1 {mode} is out of scope; the parametrize whitelist should not have generated this combination"
|
|
)
|
|
|
|
torch.manual_seed(2026)
|
|
_validate_spec(batch_spec)
|
|
|
|
vllm_config = _get_vllm_config(
|
|
model,
|
|
dtype,
|
|
tensor_parallel_size=tensor_parallel_size,
|
|
)
|
|
device = torch.device("npu")
|
|
|
|
seq_lens = list(batch_spec.seq_lens)
|
|
query_lens = list(batch_spec.query_lens)
|
|
num_tokens = batch_spec.compute_num_tokens()
|
|
|
|
cache_config = vllm_config.cache_config
|
|
hf_text = vllm_config.model_config.hf_text_config
|
|
block_size = cache_config.block_size
|
|
kv_lora_rank = hf_text.kv_lora_rank
|
|
qk_nope_head_dim = hf_text.qk_nope_head_dim
|
|
qk_rope_head_dim = hf_text.qk_rope_head_dim
|
|
v_head_dim = hf_text.v_head_dim
|
|
num_heads = hf_text.num_attention_heads
|
|
|
|
qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
|
scale = 1.0 / math.sqrt(qk_head_dim)
|
|
|
|
k_nope_contexts, k_pe_contexts = _make_synthetic_kv_contexts(
|
|
seq_lens,
|
|
kv_lora_rank,
|
|
qk_rope_head_dim,
|
|
dtype,
|
|
device,
|
|
)
|
|
k_nope_cache, k_pe_cache, block_table = _build_paged_kv_cache(
|
|
seq_lens,
|
|
k_nope_contexts,
|
|
k_pe_contexts,
|
|
block_size,
|
|
kv_lora_rank,
|
|
qk_rope_head_dim,
|
|
dtype,
|
|
device,
|
|
)
|
|
q_nope_latent = torch.randn(
|
|
num_tokens,
|
|
num_heads,
|
|
kv_lora_rank,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
q_pe_latent = torch.randn(
|
|
num_tokens,
|
|
num_heads,
|
|
qk_rope_head_dim,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
W_UV = _make_w_uv(num_heads, kv_lora_rank, v_head_dim, dtype, device)
|
|
|
|
q_nope_full = [torch.randn(q, num_heads, qk_nope_head_dim, dtype=dtype, device=device) for q in query_lens]
|
|
q_pe_full = [torch.randn(q, num_heads, qk_rope_head_dim, dtype=dtype, device=device) for q in query_lens]
|
|
k_nope_full = [torch.randn(q, num_heads, qk_nope_head_dim, dtype=dtype, device=device) for q in query_lens]
|
|
k_pe_full = [torch.randn(q, num_heads, qk_rope_head_dim, dtype=dtype, device=device) for q in query_lens]
|
|
v_full = [torch.randn(q, num_heads, v_head_dim, dtype=dtype, device=device) for q in query_lens]
|
|
|
|
impl_kwargs = dict(
|
|
scale=scale,
|
|
num_heads=num_heads,
|
|
kv_lora_rank=kv_lora_rank,
|
|
qk_nope_head_dim=qk_nope_head_dim,
|
|
qk_rope_head_dim=qk_rope_head_dim,
|
|
v_head_dim=v_head_dim,
|
|
W_UV=W_UV,
|
|
vllm_config=vllm_config,
|
|
)
|
|
|
|
dt = "bf16" if dtype == torch.bfloat16 else "fp16"
|
|
tag_base = f"{mode}|{batch_spec.name}|pcp={pcp_size}|dcp={dcp_size}|tp={tensor_parallel_size}|{dt}"
|
|
|
|
if mode in ("decode", "mtp"):
|
|
causal = mode == "mtp"
|
|
if mode == "mtp":
|
|
spec_window = query_lens[0]
|
|
speculative_config = MagicMock()
|
|
speculative_config.num_speculative_tokens = spec_window - 1
|
|
attn_state = AscendAttentionState.SpecDecoding
|
|
else:
|
|
speculative_config = None
|
|
attn_state = AscendAttentionState.DecodeOnly
|
|
|
|
backend_output = _run_mla_cp_decode_kernel(
|
|
q_nope_latent=q_nope_latent,
|
|
q_pe_latent=q_pe_latent,
|
|
k_nope_cache=k_nope_cache,
|
|
k_pe_cache=k_pe_cache,
|
|
block_table=block_table,
|
|
seq_lens=seq_lens,
|
|
query_lens=query_lens,
|
|
attn_state=attn_state,
|
|
causal=causal,
|
|
speculative_config=speculative_config,
|
|
block_size=block_size,
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
dtype=dtype,
|
|
impl_kwargs=impl_kwargs,
|
|
device=device,
|
|
)
|
|
reference_output = _decode_reference(
|
|
q_nope_latent=q_nope_latent,
|
|
q_pe_latent=q_pe_latent,
|
|
k_nope_contexts=k_nope_contexts,
|
|
k_pe_contexts=k_pe_contexts,
|
|
seq_lens=seq_lens,
|
|
query_lens=query_lens,
|
|
W_UV=W_UV,
|
|
scale=scale,
|
|
causal=causal,
|
|
num_heads=num_heads,
|
|
v_head_dim=v_head_dim,
|
|
out_dtype=dtype,
|
|
)
|
|
_record_and_assert(
|
|
backend_output,
|
|
reference_output,
|
|
tag_base,
|
|
dtype=dtype,
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
return
|
|
|
|
if mode == "prefill":
|
|
backend_output = _run_mla_cp_prefill_kernel(
|
|
q_nope=torch.cat(q_nope_full, dim=0),
|
|
q_pe=torch.cat(q_pe_full, dim=0),
|
|
k_nope=torch.cat(k_nope_full, dim=0),
|
|
k_pe=torch.cat(k_pe_full, dim=0),
|
|
value=torch.cat(v_full, dim=0),
|
|
k_nope_cache=k_nope_cache,
|
|
k_pe_cache=k_pe_cache,
|
|
query_lens=query_lens,
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
impl_kwargs=impl_kwargs,
|
|
device=device,
|
|
)
|
|
reference_output = _prefill_reference(
|
|
q_nope_full=q_nope_full,
|
|
q_pe_full=q_pe_full,
|
|
k_nope_full_per_req=k_nope_full,
|
|
k_pe_full_per_req=k_pe_full,
|
|
v_full_per_req=v_full,
|
|
scale=scale,
|
|
num_heads=num_heads,
|
|
v_head_dim=v_head_dim,
|
|
out_dtype=dtype,
|
|
)
|
|
_record_and_assert(
|
|
backend_output,
|
|
reference_output,
|
|
tag_base,
|
|
dtype=dtype,
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
return
|
|
|
|
n_decode_reqs = sum(1 for q in query_lens if q == 1)
|
|
assert all(q == 1 for q in query_lens[:n_decode_reqs]), (
|
|
f"mixed spec {batch_spec.name}: decode requests must come first"
|
|
)
|
|
assert all(q > 1 for q in query_lens[n_decode_reqs:]), (
|
|
f"mixed spec {batch_spec.name}: prefill requests must come after decode"
|
|
)
|
|
|
|
decode_backend = _run_mla_cp_decode_kernel(
|
|
q_nope_latent=q_nope_latent[:n_decode_reqs],
|
|
q_pe_latent=q_pe_latent[:n_decode_reqs],
|
|
k_nope_cache=k_nope_cache,
|
|
k_pe_cache=k_pe_cache,
|
|
block_table=block_table[:n_decode_reqs],
|
|
seq_lens=seq_lens[:n_decode_reqs],
|
|
query_lens=query_lens[:n_decode_reqs],
|
|
attn_state=AscendAttentionState.ChunkedPrefill,
|
|
causal=False,
|
|
speculative_config=None,
|
|
block_size=block_size,
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
dtype=dtype,
|
|
impl_kwargs=impl_kwargs,
|
|
device=device,
|
|
)
|
|
decode_reference = _decode_reference(
|
|
q_nope_latent=q_nope_latent[:n_decode_reqs],
|
|
q_pe_latent=q_pe_latent[:n_decode_reqs],
|
|
k_nope_contexts=k_nope_contexts[:n_decode_reqs],
|
|
k_pe_contexts=k_pe_contexts[:n_decode_reqs],
|
|
seq_lens=seq_lens[:n_decode_reqs],
|
|
query_lens=query_lens[:n_decode_reqs],
|
|
W_UV=W_UV,
|
|
scale=scale,
|
|
causal=False,
|
|
num_heads=num_heads,
|
|
v_head_dim=v_head_dim,
|
|
out_dtype=dtype,
|
|
)
|
|
_record_and_assert(
|
|
decode_backend,
|
|
decode_reference,
|
|
f"{tag_base}|mixed_decode",
|
|
dtype=dtype,
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
prefill_backend = _run_mla_cp_prefill_kernel(
|
|
q_nope=torch.cat(q_nope_full[n_decode_reqs:], dim=0),
|
|
q_pe=torch.cat(q_pe_full[n_decode_reqs:], dim=0),
|
|
k_nope=torch.cat(k_nope_full[n_decode_reqs:], dim=0),
|
|
k_pe=torch.cat(k_pe_full[n_decode_reqs:], dim=0),
|
|
value=torch.cat(v_full[n_decode_reqs:], dim=0),
|
|
k_nope_cache=k_nope_cache,
|
|
k_pe_cache=k_pe_cache,
|
|
query_lens=query_lens[n_decode_reqs:],
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
impl_kwargs=impl_kwargs,
|
|
device=device,
|
|
)
|
|
prefill_reference = _prefill_reference(
|
|
q_nope_full=q_nope_full[n_decode_reqs:],
|
|
q_pe_full=q_pe_full[n_decode_reqs:],
|
|
k_nope_full_per_req=k_nope_full[n_decode_reqs:],
|
|
k_pe_full_per_req=k_pe_full[n_decode_reqs:],
|
|
v_full_per_req=v_full[n_decode_reqs:],
|
|
scale=scale,
|
|
num_heads=num_heads,
|
|
v_head_dim=v_head_dim,
|
|
out_dtype=dtype,
|
|
)
|
|
_record_and_assert(
|
|
prefill_backend,
|
|
prefill_reference,
|
|
f"{tag_base}|mixed_prefill",
|
|
dtype=dtype,
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
# Whitelist of ``(batch_spec_name, pcp_size, dcp_size)`` combinations that are
|
|
# in scope for the single-rank precision suite. ``PCP > 1`` prefill / mixed
|
|
# requires per-rank PCP scheduler metadata (``q_head_idx``, ``q_tail_idx``,
|
|
# ``q_full_idx``, ...) that only a multi-rank job can produce faithfully -- so
|
|
# those combos are intentionally absent from the matrix (covered structurally
|
|
# by ``tests/ut/attention/test_mla_cp.py``). The non-CP ``(1, 1)`` baseline is
|
|
# covered by ``tests/ut/attention/test_mla_v1_precision.py``.
|
|
_TOPOLOGIES_ALL = [(2, 2), (2, 1), (1, 2), (1, 4)]
|
|
_TOPOLOGIES_PCP1 = [(1, 2), (1, 4)]
|
|
_TEST_CASES: list[tuple[str, int, int]] = [
|
|
(name, pcp, dcp)
|
|
for name in BATCH_SPECS
|
|
for (pcp, dcp) in (_TOPOLOGIES_ALL if _infer_mode(BATCH_SPECS[name]) in ("decode", "mtp") else _TOPOLOGIES_PCP1)
|
|
]
|
|
|
|
|
|
@pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
|
|
@pytest.mark.parametrize(
|
|
"batch_spec_name,pcp_size,dcp_size",
|
|
_TEST_CASES,
|
|
ids=[f"{n}-pcp{p}-dcp{d}" for (n, p, d) in _TEST_CASES],
|
|
)
|
|
@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V3.2-Exp"])
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
|
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
|
def test_mla_cp_correctness(
|
|
batch_spec_name: str,
|
|
model: str,
|
|
pcp_size: int,
|
|
dcp_size: int,
|
|
dtype: torch.dtype,
|
|
tensor_parallel_size: int,
|
|
) -> None:
|
|
"""Test MLA-CP correctness across workload, topology, dtype, and TP size."""
|
|
atol = FP16_ATOL if dtype == torch.float16 else DEFAULT_ATOL
|
|
rtol = FP16_RTOL if dtype == torch.float16 else DEFAULT_RTOL
|
|
_test_mla_cp_correctness(
|
|
BATCH_SPECS[batch_spec_name],
|
|
model,
|
|
pcp_size=pcp_size,
|
|
dcp_size=dcp_size,
|
|
dtype=dtype,
|
|
atol=atol,
|
|
rtol=rtol,
|
|
tensor_parallel_size=tensor_parallel_size,
|
|
)
|