994 lines
35 KiB
Python
994 lines
35 KiB
Python
|
from functools import partial
|
||
|
|
from unittest.mock import MagicMock, patch
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
from vllm.config import set_current_vllm_config
|
||
|
|
from vllm.forward_context import set_forward_context
|
||
|
|
from vllm.utils.torch_utils import set_random_seed
|
||
|
|
|
||
|
|
from tests.ut.attention.utils import (
|
||
|
|
BatchSpec,
|
||
|
|
create_common_attn_metadata,
|
||
|
|
create_vllm_config,
|
||
|
|
patch_distributed_groups,
|
||
|
|
)
|
||
|
|
from vllm_ascend.attention.attention_mask import AttentionMaskBuilder
|
||
|
|
from vllm_ascend.attention.attention_v1 import AscendMetadata
|
||
|
|
from vllm_ascend.attention.context_parallel.attention_cp import (
|
||
|
|
AscendAttentionCPImpl,
|
||
|
|
)
|
||
|
|
from vllm_ascend.attention.context_parallel.common_cp import (
|
||
|
|
AscendMetadataForDecode,
|
||
|
|
AscendMetadataForPrefill,
|
||
|
|
AscendPCPMetadata,
|
||
|
|
)
|
||
|
|
|
||
|
|
BATCH_SPECS = {
|
||
|
|
"single_prefill": BatchSpec(seq_lens=[128], query_lens=[128]),
|
||
|
|
"small_prefill": BatchSpec(seq_lens=[32, 48], query_lens=[32, 48]),
|
||
|
|
"medium_prefill": BatchSpec(seq_lens=[256, 512], query_lens=[256, 512]),
|
||
|
|
"large_prefill": BatchSpec(seq_lens=[1024, 2048], query_lens=[1024, 2048]),
|
||
|
|
"single_decode": BatchSpec(seq_lens=[32], query_lens=[1]),
|
||
|
|
"small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
|
||
|
|
"medium_decode": BatchSpec(seq_lens=[128, 256, 512, 1024], query_lens=[1, 1, 1, 1]),
|
||
|
|
"mixed_small": BatchSpec(seq_lens=[32, 40, 5, 5], query_lens=[1, 1, 5, 5]),
|
||
|
|
"mixed_medium": BatchSpec(seq_lens=[256, 512, 7, 7], query_lens=[1, 1, 7, 7]),
|
||
|
|
"mixed_large": BatchSpec(seq_lens=[1024, 2048, 16, 16], query_lens=[1, 1, 16, 16]),
|
||
|
|
"mtp_1_plus_3_small": BatchSpec(seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4]),
|
||
|
|
"mtp_1_plus_3_medium": BatchSpec(seq_lens=[1024, 2048, 3072, 4096], query_lens=[4, 4, 4, 4]),
|
||
|
|
"mtp_1_plus_3_tiny": BatchSpec(seq_lens=[64, 128], query_lens=[4, 4]),
|
||
|
|
}
|
||
|
|
|
||
|
|
MODELS = [
|
||
|
|
"Qwen/Qwen3-8B",
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
class MockAttentionLayer:
|
||
|
|
def __init__(self, device: torch.device):
|
||
|
|
self._q_scale = torch.tensor(1.0, device=device)
|
||
|
|
self._k_scale = torch.tensor(1.0, device=device)
|
||
|
|
self._v_scale = torch.tensor(1.0, device=device)
|
||
|
|
self._q_scale_float = 1.0
|
||
|
|
self._k_scale_float = 1.0
|
||
|
|
self._v_scale_float = 1.0
|
||
|
|
self.layer_name = "model.layers.0"
|
||
|
|
|
||
|
|
|
||
|
|
def compute_sdpa_reference(
|
||
|
|
q: torch.Tensor,
|
||
|
|
k: torch.Tensor,
|
||
|
|
v: torch.Tensor,
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
scale: float,
|
||
|
|
num_q_heads: int,
|
||
|
|
num_kv_heads: int,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Compute reference attention output using PyTorch SDPA.
|
||
|
|
|
||
|
|
Iterates over each sequence in the batch and computes causal attention
|
||
|
|
matching the FIA sparse_mode=3 behavior with splitfuse layout.
|
||
|
|
Used for pure prefill tests where seq_lens == query_lens.
|
||
|
|
"""
|
||
|
|
enable_gqa = num_q_heads != num_kv_heads
|
||
|
|
all_sdpa_outputs = []
|
||
|
|
q_offset = 0
|
||
|
|
kv_offset = 0
|
||
|
|
|
||
|
|
for i in range(batch_spec.batch_size):
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
kv_len = batch_spec.seq_lens[i]
|
||
|
|
|
||
|
|
q_i = q[q_offset : q_offset + q_len]
|
||
|
|
k_i = k[kv_offset : kv_offset + kv_len]
|
||
|
|
v_i = v[kv_offset : kv_offset + kv_len]
|
||
|
|
|
||
|
|
q_sdpa = q_i.unsqueeze(0).transpose(1, 2)
|
||
|
|
k_sdpa = k_i.unsqueeze(0).transpose(1, 2)
|
||
|
|
v_sdpa = v_i.unsqueeze(0).transpose(1, 2)
|
||
|
|
|
||
|
|
context_len = kv_len - q_len
|
||
|
|
if context_len > 0:
|
||
|
|
attn_mask = torch.ones(q_len, kv_len, dtype=torch.bool, device=q.device)
|
||
|
|
causal_mask = torch.tril(torch.ones(q_len, q_len, device=q.device))
|
||
|
|
attn_mask[:, context_len:] = causal_mask
|
||
|
|
sdpa_out = torch.nn.functional.scaled_dot_product_attention(
|
||
|
|
q_sdpa,
|
||
|
|
k_sdpa,
|
||
|
|
v_sdpa,
|
||
|
|
attn_mask=attn_mask,
|
||
|
|
is_causal=False,
|
||
|
|
enable_gqa=enable_gqa,
|
||
|
|
scale=scale,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
sdpa_out = torch.nn.functional.scaled_dot_product_attention(
|
||
|
|
q_sdpa,
|
||
|
|
k_sdpa,
|
||
|
|
v_sdpa,
|
||
|
|
is_causal=True,
|
||
|
|
enable_gqa=enable_gqa,
|
||
|
|
scale=scale,
|
||
|
|
)
|
||
|
|
|
||
|
|
all_sdpa_outputs.append(sdpa_out.transpose(1, 2).squeeze(0))
|
||
|
|
q_offset += q_len
|
||
|
|
kv_offset += kv_len
|
||
|
|
|
||
|
|
return torch.cat(all_sdpa_outputs, dim=0)
|
||
|
|
|
||
|
|
|
||
|
|
def compute_mixed_sdpa_reference(
|
||
|
|
full_q: torch.Tensor,
|
||
|
|
full_k: torch.Tensor,
|
||
|
|
full_v: torch.Tensor,
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
scale: float,
|
||
|
|
num_q_heads: int,
|
||
|
|
num_kv_heads: int,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Compute per-sequence SDPA reference for mixed decode+prefill.
|
||
|
|
|
||
|
|
Each sequence gets its own Q/K/V with causal masking
|
||
|
|
(context tokens are visible to new tokens).
|
||
|
|
"""
|
||
|
|
enable_gqa = num_q_heads != num_kv_heads
|
||
|
|
all_outputs = []
|
||
|
|
q_offset = 0
|
||
|
|
kv_offset = 0
|
||
|
|
|
||
|
|
for i in range(batch_spec.batch_size):
|
||
|
|
s_len = batch_spec.seq_lens[i]
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
context_len = s_len - q_len
|
||
|
|
|
||
|
|
q_i = full_q[q_offset : q_offset + q_len]
|
||
|
|
k_i = full_k[kv_offset : kv_offset + s_len]
|
||
|
|
v_i = full_v[kv_offset : kv_offset + s_len]
|
||
|
|
|
||
|
|
q_sdpa = q_i.unsqueeze(0).transpose(1, 2)
|
||
|
|
k_sdpa = k_i.unsqueeze(0).transpose(1, 2)
|
||
|
|
v_sdpa = v_i.unsqueeze(0).transpose(1, 2)
|
||
|
|
|
||
|
|
if context_len > 0:
|
||
|
|
attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=full_q.device)
|
||
|
|
causal_mask = torch.tril(torch.ones(q_len, q_len, device=full_q.device))
|
||
|
|
attn_mask[:, context_len:] = causal_mask
|
||
|
|
sdpa_out = torch.nn.functional.scaled_dot_product_attention(
|
||
|
|
q_sdpa,
|
||
|
|
k_sdpa,
|
||
|
|
v_sdpa,
|
||
|
|
attn_mask=attn_mask,
|
||
|
|
is_causal=False,
|
||
|
|
enable_gqa=enable_gqa,
|
||
|
|
scale=scale,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
sdpa_out = torch.nn.functional.scaled_dot_product_attention(
|
||
|
|
q_sdpa,
|
||
|
|
k_sdpa,
|
||
|
|
v_sdpa,
|
||
|
|
is_causal=True,
|
||
|
|
enable_gqa=enable_gqa,
|
||
|
|
scale=scale,
|
||
|
|
)
|
||
|
|
|
||
|
|
all_outputs.append(sdpa_out.transpose(1, 2).squeeze(0))
|
||
|
|
q_offset += q_len
|
||
|
|
kv_offset += s_len
|
||
|
|
|
||
|
|
return torch.cat(all_outputs, dim=0)
|
||
|
|
|
||
|
|
|
||
|
|
def _make_kv_cache_for_decode(
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
num_kv_heads: int,
|
||
|
|
head_size: int,
|
||
|
|
block_size: int,
|
||
|
|
dtype: torch.dtype,
|
||
|
|
device: torch.device,
|
||
|
|
key: torch.Tensor,
|
||
|
|
value: torch.Tensor,
|
||
|
|
block_table: torch.Tensor,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
|
|
"""Pre-populate KV cache with ALL tokens for decode-only tests.
|
||
|
|
|
||
|
|
Places context + new tokens sequentially in cache blocks
|
||
|
|
so that FIA paged attention can read them via block_table.
|
||
|
|
"""
|
||
|
|
num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
|
||
|
|
num_blocks = max(num_blocks, 64)
|
||
|
|
k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
|
||
|
|
kv_offset = 0
|
||
|
|
block_start = 0
|
||
|
|
for i in range(batch_spec.batch_size):
|
||
|
|
s_len = batch_spec.seq_lens[i]
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
context_len = s_len - q_len
|
||
|
|
|
||
|
|
k_full = key[kv_offset : kv_offset + s_len]
|
||
|
|
v_full = value[kv_offset : kv_offset + s_len]
|
||
|
|
|
||
|
|
context_k = k_full[:context_len].contiguous()
|
||
|
|
context_v = v_full[:context_len].contiguous()
|
||
|
|
new_k = k_full[context_len:].contiguous()
|
||
|
|
new_v = v_full[context_len:].contiguous()
|
||
|
|
|
||
|
|
num_blocks_for_seq = (s_len + block_size - 1) // block_size
|
||
|
|
block_table[i, :num_blocks_for_seq] = torch.arange(
|
||
|
|
block_start,
|
||
|
|
block_start + num_blocks_for_seq,
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=device,
|
||
|
|
)
|
||
|
|
|
||
|
|
for t_idx in range(context_len):
|
||
|
|
blk = block_start + t_idx // block_size
|
||
|
|
pos = t_idx % block_size
|
||
|
|
k_cache[blk, pos] = context_k[t_idx]
|
||
|
|
v_cache[blk, pos] = context_v[t_idx]
|
||
|
|
|
||
|
|
for t_idx in range(q_len):
|
||
|
|
blk = block_start + (context_len + t_idx) // block_size
|
||
|
|
pos = (context_len + t_idx) % block_size
|
||
|
|
k_cache[blk, pos] = new_k[t_idx]
|
||
|
|
v_cache[blk, pos] = new_v[t_idx]
|
||
|
|
|
||
|
|
kv_offset += s_len
|
||
|
|
block_start += num_blocks_for_seq
|
||
|
|
|
||
|
|
return k_cache, v_cache
|
||
|
|
|
||
|
|
|
||
|
|
def _make_kv_cache_for_mixed(
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
num_kv_heads: int,
|
||
|
|
head_size: int,
|
||
|
|
block_size: int,
|
||
|
|
dtype: torch.dtype,
|
||
|
|
device: torch.device,
|
||
|
|
key: torch.Tensor,
|
||
|
|
value: torch.Tensor,
|
||
|
|
block_table: torch.Tensor,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
|
|
"""Pre-populate KV cache for mixed decode+prefill tests.
|
||
|
|
|
||
|
|
Only decode sequences have context tokens placed in the cache.
|
||
|
|
Prefill sequences (seq_lens == query_lens) have no context;
|
||
|
|
their KV goes through the direct FIA prefill path.
|
||
|
|
"""
|
||
|
|
num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
|
||
|
|
num_blocks = max(num_blocks, 64)
|
||
|
|
k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
|
||
|
|
kv_offset = 0
|
||
|
|
block_start = 0
|
||
|
|
for i in range(batch_spec.batch_size):
|
||
|
|
s_len = batch_spec.seq_lens[i]
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
context_len = s_len - q_len
|
||
|
|
is_decode = q_len == 1
|
||
|
|
|
||
|
|
k_full = key[kv_offset : kv_offset + s_len]
|
||
|
|
v_full = value[kv_offset : kv_offset + s_len]
|
||
|
|
|
||
|
|
if is_decode and context_len > 0:
|
||
|
|
num_blocks_for_seq = (s_len + block_size - 1) // block_size
|
||
|
|
block_table[i, :num_blocks_for_seq] = torch.arange(
|
||
|
|
block_start,
|
||
|
|
block_start + num_blocks_for_seq,
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=device,
|
||
|
|
)
|
||
|
|
|
||
|
|
for t_idx in range(context_len):
|
||
|
|
blk = block_start + t_idx // block_size
|
||
|
|
pos = t_idx % block_size
|
||
|
|
k_cache[blk, pos] = k_full[t_idx]
|
||
|
|
v_cache[blk, pos] = v_full[t_idx]
|
||
|
|
|
||
|
|
for t_idx in range(q_len):
|
||
|
|
blk = block_start + (context_len + t_idx) // block_size
|
||
|
|
pos = (context_len + t_idx) % block_size
|
||
|
|
k_cache[blk, pos] = k_full[context_len + t_idx]
|
||
|
|
v_cache[blk, pos] = v_full[context_len + t_idx]
|
||
|
|
|
||
|
|
block_start += num_blocks_for_seq
|
||
|
|
|
||
|
|
kv_offset += s_len
|
||
|
|
|
||
|
|
return k_cache, v_cache
|
||
|
|
|
||
|
|
|
||
|
|
def build_cp_attn_metadata(
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
vllm_config,
|
||
|
|
device: torch.device,
|
||
|
|
pcp_size: int = 1,
|
||
|
|
pcp_rank: int = 0,
|
||
|
|
kv_cache_prepopulated: bool = False,
|
||
|
|
decode_threshold: int = 1,
|
||
|
|
) -> AscendMetadata:
|
||
|
|
common_attn_metadata = create_common_attn_metadata(batch_spec, vllm_config.cache_config.block_size, device)
|
||
|
|
|
||
|
|
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 = sum(1 for ql in batch_spec.query_lens if ql <= decode_threshold)
|
||
|
|
num_prefills = batch_spec.batch_size - num_decodes
|
||
|
|
num_decode_tokens = sum(ql for ql in batch_spec.query_lens if ql <= decode_threshold)
|
||
|
|
num_prefill_tokens = num_actual_tokens - num_decode_tokens
|
||
|
|
|
||
|
|
block_table = common_attn_metadata.block_table_tensor
|
||
|
|
slot_mapping = common_attn_metadata.slot_mapping
|
||
|
|
query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
|
||
|
|
num_decodes_flatten = query_lens[:num_decodes].sum().item()
|
||
|
|
|
||
|
|
seq_lens_cpu = common_attn_metadata.seq_lens_cpu[:num_reqs]
|
||
|
|
|
||
|
|
num_actual_tokens_pcp_padded = num_actual_tokens * pcp_size
|
||
|
|
|
||
|
|
query_start_loc = query_start_loc_cpu.to(device, non_blocking=True)
|
||
|
|
|
||
|
|
attn_mask_builder = AttentionMaskBuilder(device)
|
||
|
|
attn_mask = attn_mask_builder.get_attention_mask(common_attn_metadata.causal, vllm_config.model_config)
|
||
|
|
|
||
|
|
prefill_metadata = None
|
||
|
|
if num_prefills > 0:
|
||
|
|
prefill_query_lens = query_lens[num_decodes:]
|
||
|
|
attn_mask_seqlens = torch.cumsum(prefill_query_lens, dim=0).tolist()
|
||
|
|
head_attn_nomask_seqlens = attn_mask_seqlens if pcp_rank > 0 else []
|
||
|
|
tail_attn_nomask_seqlens = attn_mask_seqlens
|
||
|
|
|
||
|
|
total_prefill_tokens = num_prefill_tokens
|
||
|
|
prefill_tokens_offset = num_decode_tokens
|
||
|
|
|
||
|
|
if pcp_size > 1:
|
||
|
|
chunk_size = total_prefill_tokens // pcp_size
|
||
|
|
rank_start = prefill_tokens_offset + pcp_rank * chunk_size
|
||
|
|
rank_end = rank_start + chunk_size
|
||
|
|
|
||
|
|
q_head_idx = torch.arange(rank_start, rank_end, device=device, dtype=torch.long)
|
||
|
|
q_tail_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
|
||
|
|
kv_total = total_prefill_tokens * pcp_size
|
||
|
|
kv_with_q_head_mask_idx = torch.arange(
|
||
|
|
prefill_tokens_offset,
|
||
|
|
prefill_tokens_offset + kv_total,
|
||
|
|
device=device,
|
||
|
|
dtype=torch.long,
|
||
|
|
)
|
||
|
|
kv_with_q_head_nomask_idx = (
|
||
|
|
torch.arange(prefill_tokens_offset, prefill_tokens_offset + kv_total, device=device, dtype=torch.long)
|
||
|
|
if pcp_rank > 0
|
||
|
|
else torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
)
|
||
|
|
q_full_idx = torch.arange(chunk_size, device=device, dtype=torch.long)
|
||
|
|
kv_with_q_tail_nomask_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
kv_with_q_tail_mask_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
pcp_allgather_restore_idx = list(range(total_prefill_tokens * pcp_size))
|
||
|
|
else:
|
||
|
|
q_head_idx = torch.arange(
|
||
|
|
prefill_tokens_offset,
|
||
|
|
prefill_tokens_offset + total_prefill_tokens,
|
||
|
|
device=device,
|
||
|
|
dtype=torch.long,
|
||
|
|
)
|
||
|
|
q_tail_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
kv_with_q_head_mask_idx = torch.arange(
|
||
|
|
prefill_tokens_offset,
|
||
|
|
prefill_tokens_offset + total_prefill_tokens,
|
||
|
|
device=device,
|
||
|
|
dtype=torch.long,
|
||
|
|
)
|
||
|
|
kv_with_q_head_nomask_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
kv_with_q_tail_nomask_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
kv_with_q_tail_mask_idx = torch.tensor([], device=device, dtype=torch.long)
|
||
|
|
q_full_idx = torch.arange(total_prefill_tokens, device=device, dtype=torch.long)
|
||
|
|
pcp_allgather_restore_idx = None
|
||
|
|
|
||
|
|
pcp_metadata = AscendPCPMetadata(
|
||
|
|
q_head_idx=q_head_idx,
|
||
|
|
q_tail_idx=q_tail_idx,
|
||
|
|
kv_with_q_head_nomask_idx=kv_with_q_head_nomask_idx,
|
||
|
|
kv_with_q_head_mask_idx=kv_with_q_head_mask_idx,
|
||
|
|
kv_with_q_tail_nomask_idx=kv_with_q_tail_nomask_idx,
|
||
|
|
kv_with_q_tail_mask_idx=kv_with_q_tail_mask_idx,
|
||
|
|
attn_mask_seqlens=attn_mask_seqlens,
|
||
|
|
head_attn_nomask_seqlens=head_attn_nomask_seqlens,
|
||
|
|
tail_attn_nomask_seqlens=tail_attn_nomask_seqlens,
|
||
|
|
q_full_idx=q_full_idx,
|
||
|
|
pcp_use_hybrid_attn=False,
|
||
|
|
pcp_allgather_restore_idx=pcp_allgather_restore_idx,
|
||
|
|
)
|
||
|
|
|
||
|
|
prefill_cumsum_q = torch.cumsum(query_lens[num_decodes:], dim=0).to(device)
|
||
|
|
|
||
|
|
prefill_metadata = AscendMetadataForPrefill(
|
||
|
|
pcp_metadata=pcp_metadata,
|
||
|
|
pcp_exit_fa_scatter_idx=None,
|
||
|
|
chunked_context=None,
|
||
|
|
block_tables=block_table[num_decodes_flatten:, ...],
|
||
|
|
actual_seq_lengths_q=prefill_cumsum_q,
|
||
|
|
)
|
||
|
|
|
||
|
|
decode_metadata = None
|
||
|
|
if num_decodes > 0:
|
||
|
|
decode_query_lens = query_lens[:num_decodes].tolist()
|
||
|
|
decode_seq_lens = seq_lens_cpu[:num_decodes].tolist()
|
||
|
|
|
||
|
|
if kv_cache_prepopulated:
|
||
|
|
num_computed_tokens_arr = np.zeros((num_decodes_flatten, pcp_size, 1), dtype=np.int32)
|
||
|
|
flat_idx = 0
|
||
|
|
for i in range(num_decodes):
|
||
|
|
s_len = int(decode_seq_lens[i])
|
||
|
|
q_len = decode_query_lens[i]
|
||
|
|
context_len = s_len - q_len
|
||
|
|
for t in range(q_len):
|
||
|
|
num_computed_tokens_arr[flat_idx, pcp_rank, 0] = context_len + t + 1
|
||
|
|
flat_idx += 1
|
||
|
|
else:
|
||
|
|
num_computed_tokens_arr = np.zeros((num_decodes_flatten, pcp_size, 1), dtype=np.int32)
|
||
|
|
|
||
|
|
# Tile block_table for MTP: each decode request may have multiple tokens
|
||
|
|
if num_decodes_flatten > num_decodes:
|
||
|
|
tiled_rows = []
|
||
|
|
for i in range(num_decodes):
|
||
|
|
q_len = decode_query_lens[i]
|
||
|
|
row = block_table[i : i + 1]
|
||
|
|
tiled_rows.append(row.repeat(q_len, 1))
|
||
|
|
decode_block_tables = torch.cat(tiled_rows, dim=0)
|
||
|
|
else:
|
||
|
|
decode_block_tables = block_table[:num_decodes_flatten]
|
||
|
|
|
||
|
|
decode_metadata = AscendMetadataForDecode(
|
||
|
|
num_computed_tokens_of_pcp_dcp=num_computed_tokens_arr,
|
||
|
|
block_tables=decode_block_tables,
|
||
|
|
)
|
||
|
|
|
||
|
|
actual_seq_lengths_q = (
|
||
|
|
torch.arange(num_decodes_flatten, device=device) + 1
|
||
|
|
if num_decodes_flatten > 0
|
||
|
|
else torch.tensor([], device=device)
|
||
|
|
).tolist() + torch.cumsum(query_lens[num_decodes:], dim=0).tolist()
|
||
|
|
|
||
|
|
attn_metadata = AscendMetadata(
|
||
|
|
num_actual_tokens=num_actual_tokens,
|
||
|
|
num_decode_tokens=num_decode_tokens,
|
||
|
|
num_actual_tokens_pcp_padded=num_actual_tokens_pcp_padded,
|
||
|
|
num_decodes_flatten=num_decodes_flatten,
|
||
|
|
block_tables=block_table,
|
||
|
|
query_start_loc=query_start_loc,
|
||
|
|
seq_lens=common_attn_metadata.seq_lens[:num_reqs],
|
||
|
|
seq_lens_cpu=seq_lens_cpu,
|
||
|
|
seq_lens_list=seq_lens_cpu.tolist(),
|
||
|
|
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=common_attn_metadata.attn_state,
|
||
|
|
num_prefills=num_prefills,
|
||
|
|
num_decodes=num_decodes,
|
||
|
|
prefill=prefill_metadata,
|
||
|
|
decode_meta=decode_metadata,
|
||
|
|
)
|
||
|
|
return attn_metadata
|
||
|
|
|
||
|
|
|
||
|
|
def run_cp_attention(
|
||
|
|
query: torch.Tensor,
|
||
|
|
key: torch.Tensor,
|
||
|
|
value: torch.Tensor,
|
||
|
|
kv_cache: tuple[torch.Tensor, torch.Tensor],
|
||
|
|
attn_metadata: AscendMetadata,
|
||
|
|
impl: AscendAttentionCPImpl,
|
||
|
|
device: torch.device,
|
||
|
|
vllm_config,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Run CP attention forward pass with proper setup.
|
||
|
|
|
||
|
|
Mocks reshape_and_cache (no-op) since KV cache is pre-populated
|
||
|
|
for decode tests and not needed for prefill tests.
|
||
|
|
"""
|
||
|
|
mock_layer_entry = MagicMock()
|
||
|
|
for layer_name in ["placeholder"]:
|
||
|
|
vllm_config.compilation_config.static_forward_context[layer_name] = mock_layer_entry
|
||
|
|
|
||
|
|
num_tokens = query.shape[0]
|
||
|
|
num_q_heads = query.shape[1]
|
||
|
|
head_size = query.shape[2]
|
||
|
|
output = torch.empty(num_tokens, num_q_heads, head_size, dtype=query.dtype, device=device)
|
||
|
|
|
||
|
|
with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
|
||
|
|
from vllm.forward_context import get_forward_context
|
||
|
|
|
||
|
|
forward_ctx = get_forward_context()
|
||
|
|
forward_ctx.num_tokens = num_tokens
|
||
|
|
forward_ctx.is_draft_model = False
|
||
|
|
forward_ctx.is_draft_model_prefill = False
|
||
|
|
forward_ctx.capturing = False
|
||
|
|
forward_ctx.flash_comm_v1_enabled = False
|
||
|
|
forward_ctx.flashcomm_v2_enabled = False
|
||
|
|
|
||
|
|
mock_layer = MockAttentionLayer(device)
|
||
|
|
|
||
|
|
import vllm_ascend.device.device_op as device_op_module
|
||
|
|
|
||
|
|
original_reshape_and_cache = device_op_module.DeviceOperator.reshape_and_cache
|
||
|
|
original_kv_cache_load = device_op_module.DeviceOperator.kv_cache_load
|
||
|
|
|
||
|
|
def _mock_reshape_and_cache(key, value, key_cache, value_cache, slot_mapping):
|
||
|
|
return
|
||
|
|
|
||
|
|
def _mock_kv_cache_load(key_cache, value_cache, block_tables, seq_lens_kv, starts, key, value):
|
||
|
|
return
|
||
|
|
|
||
|
|
device_op_module.DeviceOperator.reshape_and_cache = staticmethod(_mock_reshape_and_cache)
|
||
|
|
device_op_module.DeviceOperator.kv_cache_load = staticmethod(_mock_kv_cache_load)
|
||
|
|
|
||
|
|
try:
|
||
|
|
output = impl.forward(mock_layer, query, key, value, kv_cache, attn_metadata, output=output)
|
||
|
|
finally:
|
||
|
|
device_op_module.DeviceOperator.reshape_and_cache = original_reshape_and_cache
|
||
|
|
device_op_module.DeviceOperator.kv_cache_load = original_kv_cache_load
|
||
|
|
|
||
|
|
return output
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(autouse=True)
|
||
|
|
def default_mock_config():
|
||
|
|
mock_config = MagicMock()
|
||
|
|
mock_config.compilation_config = MagicMock()
|
||
|
|
mock_config.compilation_config.custom_ops = ["all"]
|
||
|
|
mock_config.compilation_config.static_forward_context = {}
|
||
|
|
mock_config.parallel_config = MagicMock()
|
||
|
|
mock_config.parallel_config.prefill_context_parallel_size = 1
|
||
|
|
mock_config.parallel_config.decode_context_parallel_size = 1
|
||
|
|
mock_config.parallel_config.tensor_parallel_size = 1
|
||
|
|
mock_config.model_config = MagicMock()
|
||
|
|
mock_config.model_config.dtype = torch.float16
|
||
|
|
mock_config.speculative_config = None
|
||
|
|
mock_config.cache_config = MagicMock()
|
||
|
|
mock_config.cache_config.block_size = 128
|
||
|
|
mock_config.kv_transfer_config = None
|
||
|
|
|
||
|
|
with set_current_vllm_config(mock_config):
|
||
|
|
yield mock_config
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(autouse=True)
|
||
|
|
def mock_graph_params():
|
||
|
|
with patch("vllm_ascend.compilation.acl_graph.get_graph_params") as mock_get_graph:
|
||
|
|
graph_params = MagicMock()
|
||
|
|
graph_params.workspaces = {}
|
||
|
|
graph_params.handles = {}
|
||
|
|
graph_params.attn_params = {}
|
||
|
|
graph_params.events = {}
|
||
|
|
mock_get_graph.return_value = graph_params
|
||
|
|
with patch("vllm_ascend.compilation.acl_graph.get_draft_graph_params", return_value=graph_params):
|
||
|
|
yield
|
||
|
|
|
||
|
|
|
||
|
|
def _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale):
|
||
|
|
return AscendAttentionCPImpl(
|
||
|
|
num_heads=num_q_heads,
|
||
|
|
head_size=head_size,
|
||
|
|
scale=scale,
|
||
|
|
num_kv_heads=num_kv_heads,
|
||
|
|
alibi_slopes=None,
|
||
|
|
sliding_window=None,
|
||
|
|
kv_cache_dtype="auto",
|
||
|
|
logits_soft_cap=None,
|
||
|
|
attn_type="decoder",
|
||
|
|
kv_sharing_target_layer_name=None,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _assert_close(output, reference, rtol, atol, backend_name):
|
||
|
|
assert output.shape == reference.shape, (
|
||
|
|
f"[{backend_name}] shape {output.shape} != reference shape {reference.shape}"
|
||
|
|
)
|
||
|
|
assert output.dtype == reference.dtype, (
|
||
|
|
f"[{backend_name}] dtype {output.dtype} != reference dtype {reference.dtype}"
|
||
|
|
)
|
||
|
|
assert torch.isfinite(output).all(), f"[{backend_name}] produced non-finite values"
|
||
|
|
|
||
|
|
def error_msg(msg: str, name: str):
|
||
|
|
return f"[{name}] output differs from SDPA baseline. {msg}"
|
||
|
|
|
||
|
|
torch.testing.assert_close(
|
||
|
|
output,
|
||
|
|
reference,
|
||
|
|
rtol=rtol,
|
||
|
|
atol=atol,
|
||
|
|
msg=partial(error_msg, name=backend_name),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Pure prefill (seq_lens == query_lens, no context)
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
def _test_cp_prefill_precision_no_cp(
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
model: str,
|
||
|
|
*,
|
||
|
|
block_size: int = 128,
|
||
|
|
atol: float = 1e-2,
|
||
|
|
rtol: float = 1e-2,
|
||
|
|
):
|
||
|
|
set_random_seed(42)
|
||
|
|
|
||
|
|
vllm_config = create_vllm_config(
|
||
|
|
model_name=model,
|
||
|
|
tensor_parallel_size=1,
|
||
|
|
max_model_len=max(batch_spec.seq_lens) + block_size,
|
||
|
|
block_size=block_size,
|
||
|
|
num_gpu_blocks=8192,
|
||
|
|
)
|
||
|
|
device = torch.device("npu")
|
||
|
|
|
||
|
|
num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
|
||
|
|
num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
|
||
|
|
head_size = vllm_config.model_config.get_head_size()
|
||
|
|
dtype = torch.bfloat16
|
||
|
|
scale = 1.0 / (head_size**0.5)
|
||
|
|
|
||
|
|
num_tokens = batch_spec.compute_num_tokens()
|
||
|
|
total_kv = sum(batch_spec.seq_lens)
|
||
|
|
|
||
|
|
query_vllm = torch.randn(num_tokens, num_q_heads, head_size, dtype=dtype, device=device)
|
||
|
|
key_vllm = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
value_vllm = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
|
||
|
|
sdpa_output = compute_sdpa_reference(
|
||
|
|
query_vllm,
|
||
|
|
key_vllm,
|
||
|
|
value_vllm,
|
||
|
|
batch_spec,
|
||
|
|
scale,
|
||
|
|
num_q_heads,
|
||
|
|
num_kv_heads,
|
||
|
|
)
|
||
|
|
|
||
|
|
attn_metadata = build_cp_attn_metadata(batch_spec, vllm_config, device, pcp_size=1, pcp_rank=0)
|
||
|
|
|
||
|
|
num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
|
||
|
|
num_blocks = max(num_blocks, 64)
|
||
|
|
kv_cache = (
|
||
|
|
torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device),
|
||
|
|
torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device),
|
||
|
|
)
|
||
|
|
|
||
|
|
impl = _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale)
|
||
|
|
output = run_cp_attention(query_vllm, key_vllm, value_vllm, kv_cache, attn_metadata, impl, device, vllm_config)
|
||
|
|
|
||
|
|
_assert_close(output, sdpa_output, rtol, atol, "CP_Prefill")
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Pure decode (query_lens == 1, context in KV cache)
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
def _test_cp_decode_precision_no_cp(
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
model: str,
|
||
|
|
*,
|
||
|
|
block_size: int = 128,
|
||
|
|
atol: float = 1e-2,
|
||
|
|
rtol: float = 1e-2,
|
||
|
|
decode_threshold: int = 1,
|
||
|
|
):
|
||
|
|
set_random_seed(42)
|
||
|
|
|
||
|
|
vllm_config = create_vllm_config(
|
||
|
|
model_name=model,
|
||
|
|
tensor_parallel_size=1,
|
||
|
|
max_model_len=max(batch_spec.seq_lens) + block_size,
|
||
|
|
block_size=block_size,
|
||
|
|
num_gpu_blocks=8192,
|
||
|
|
)
|
||
|
|
device = torch.device("npu")
|
||
|
|
|
||
|
|
num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
|
||
|
|
num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
|
||
|
|
head_size = vllm_config.model_config.get_head_size()
|
||
|
|
dtype = torch.bfloat16
|
||
|
|
scale = 1.0 / (head_size**0.5)
|
||
|
|
|
||
|
|
total_kv = sum(batch_spec.seq_lens)
|
||
|
|
|
||
|
|
query_full = torch.randn(total_kv, num_q_heads, head_size, dtype=dtype, device=device)
|
||
|
|
key_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
value_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
|
||
|
|
sdpa_output = compute_mixed_sdpa_reference(
|
||
|
|
query_full,
|
||
|
|
key_full,
|
||
|
|
value_full,
|
||
|
|
batch_spec,
|
||
|
|
scale,
|
||
|
|
num_q_heads,
|
||
|
|
num_kv_heads,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Backend inputs: only the new (decode) tokens
|
||
|
|
all_q, all_k, all_v = [], [], []
|
||
|
|
q_offset = 0
|
||
|
|
kv_offset = 0
|
||
|
|
for i in range(batch_spec.batch_size):
|
||
|
|
s_len = batch_spec.seq_lens[i]
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
context_len = s_len - q_len
|
||
|
|
|
||
|
|
q_i = query_full[q_offset : q_offset + q_len].contiguous()
|
||
|
|
k_full_i = key_full[kv_offset : kv_offset + s_len]
|
||
|
|
v_full_i = value_full[kv_offset : kv_offset + s_len]
|
||
|
|
|
||
|
|
all_q.append(q_i)
|
||
|
|
all_k.append(k_full_i[context_len:])
|
||
|
|
all_v.append(v_full_i[context_len:])
|
||
|
|
|
||
|
|
q_offset += q_len
|
||
|
|
kv_offset += s_len
|
||
|
|
|
||
|
|
query_vllm = torch.cat(all_q, dim=0)
|
||
|
|
key_vllm = torch.cat(all_k, dim=0)
|
||
|
|
value_vllm = torch.cat(all_v, dim=0)
|
||
|
|
|
||
|
|
attn_metadata = build_cp_attn_metadata(
|
||
|
|
batch_spec,
|
||
|
|
vllm_config,
|
||
|
|
device,
|
||
|
|
pcp_size=1,
|
||
|
|
pcp_rank=0,
|
||
|
|
kv_cache_prepopulated=True,
|
||
|
|
decode_threshold=decode_threshold,
|
||
|
|
)
|
||
|
|
|
||
|
|
k_cache, v_cache = _make_kv_cache_for_decode(
|
||
|
|
batch_spec,
|
||
|
|
num_kv_heads,
|
||
|
|
head_size,
|
||
|
|
block_size,
|
||
|
|
dtype,
|
||
|
|
device,
|
||
|
|
key_full,
|
||
|
|
value_full,
|
||
|
|
attn_metadata.block_tables,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Re-tile decode block tables after _make_kv_cache_for_decode updates
|
||
|
|
# block_table in-place. Without this, the tiled decode_block_tables
|
||
|
|
# (created during build_cp_attn_metadata) still contains the original
|
||
|
|
# sequential block indices, causing FIA to read from empty cache blocks.
|
||
|
|
if attn_metadata.num_decodes_flatten > attn_metadata.num_decodes:
|
||
|
|
tiled_rows = []
|
||
|
|
for i in range(attn_metadata.num_decodes):
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
row = attn_metadata.block_tables[i : i + 1]
|
||
|
|
tiled_rows.append(row.repeat(q_len, 1))
|
||
|
|
attn_metadata.decode_meta.block_tables = torch.cat(tiled_rows, dim=0)
|
||
|
|
|
||
|
|
kv_cache = (k_cache, v_cache)
|
||
|
|
|
||
|
|
impl = _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale)
|
||
|
|
output = run_cp_attention(query_vllm, key_vllm, value_vllm, kv_cache, attn_metadata, impl, device, vllm_config)
|
||
|
|
|
||
|
|
_assert_close(output, sdpa_output, rtol, atol, "CP_Decode")
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Mixed decode + prefill (prefill sequences have seq_lens == query_lens)
|
||
|
|
#
|
||
|
|
# The CP attention non-chunked prefill path calls FIA directly with only the
|
||
|
|
# new tokens as KV. Context tokens in the KV cache are NOT loaded for prefill
|
||
|
|
# without chunked context. Therefore prefill sequences in mixed mode must
|
||
|
|
# have seq_lens == query_lens (no context).
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
def _test_cp_mixed_precision_no_cp(
|
||
|
|
batch_spec: BatchSpec,
|
||
|
|
model: str,
|
||
|
|
*,
|
||
|
|
block_size: int = 128,
|
||
|
|
atol: float = 1e-2,
|
||
|
|
rtol: float = 1e-2,
|
||
|
|
):
|
||
|
|
set_random_seed(42)
|
||
|
|
|
||
|
|
vllm_config = create_vllm_config(
|
||
|
|
model_name=model,
|
||
|
|
tensor_parallel_size=1,
|
||
|
|
max_model_len=max(batch_spec.seq_lens) + block_size,
|
||
|
|
block_size=block_size,
|
||
|
|
num_gpu_blocks=8192,
|
||
|
|
)
|
||
|
|
device = torch.device("npu")
|
||
|
|
|
||
|
|
num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
|
||
|
|
num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
|
||
|
|
head_size = vllm_config.model_config.get_head_size()
|
||
|
|
dtype = torch.bfloat16
|
||
|
|
scale = 1.0 / (head_size**0.5)
|
||
|
|
|
||
|
|
total_kv = sum(batch_spec.seq_lens)
|
||
|
|
|
||
|
|
query_full = torch.randn(total_kv, num_q_heads, head_size, dtype=dtype, device=device)
|
||
|
|
key_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
value_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
|
||
|
|
|
||
|
|
sdpa_output = compute_mixed_sdpa_reference(
|
||
|
|
query_full,
|
||
|
|
key_full,
|
||
|
|
value_full,
|
||
|
|
batch_spec,
|
||
|
|
scale,
|
||
|
|
num_q_heads,
|
||
|
|
num_kv_heads,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Backend inputs: only the new tokens
|
||
|
|
all_q, all_k, all_v = [], [], []
|
||
|
|
q_offset = 0
|
||
|
|
kv_offset = 0
|
||
|
|
for i in range(batch_spec.batch_size):
|
||
|
|
s_len = batch_spec.seq_lens[i]
|
||
|
|
q_len = batch_spec.query_lens[i]
|
||
|
|
context_len = s_len - q_len
|
||
|
|
|
||
|
|
q_i = query_full[q_offset : q_offset + q_len].contiguous()
|
||
|
|
k_full_i = key_full[kv_offset : kv_offset + s_len]
|
||
|
|
v_full_i = value_full[kv_offset : kv_offset + s_len]
|
||
|
|
|
||
|
|
all_q.append(q_i)
|
||
|
|
all_k.append(k_full_i[context_len:])
|
||
|
|
all_v.append(v_full_i[context_len:])
|
||
|
|
|
||
|
|
q_offset += q_len
|
||
|
|
kv_offset += s_len
|
||
|
|
|
||
|
|
query_vllm = torch.cat(all_q, dim=0)
|
||
|
|
key_vllm = torch.cat(all_k, dim=0)
|
||
|
|
value_vllm = torch.cat(all_v, dim=0)
|
||
|
|
|
||
|
|
attn_metadata = build_cp_attn_metadata(
|
||
|
|
batch_spec,
|
||
|
|
vllm_config,
|
||
|
|
device,
|
||
|
|
pcp_size=1,
|
||
|
|
pcp_rank=0,
|
||
|
|
kv_cache_prepopulated=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
k_cache, v_cache = _make_kv_cache_for_mixed(
|
||
|
|
batch_spec,
|
||
|
|
num_kv_heads,
|
||
|
|
head_size,
|
||
|
|
block_size,
|
||
|
|
dtype,
|
||
|
|
device,
|
||
|
|
key_full,
|
||
|
|
value_full,
|
||
|
|
attn_metadata.block_tables,
|
||
|
|
)
|
||
|
|
kv_cache = (k_cache, v_cache)
|
||
|
|
|
||
|
|
impl = _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale)
|
||
|
|
output = run_cp_attention(query_vllm, key_vllm, value_vllm, kv_cache, attn_metadata, impl, device, vllm_config)
|
||
|
|
|
||
|
|
_assert_close(output, sdpa_output, rtol, atol, "CP_Mixed")
|
||
|
|
|
||
|
|
|
||
|
|
class TestCPAttentionPrecision:
|
||
|
|
"""Precision tests for AscendAttentionCPImpl.
|
||
|
|
|
||
|
|
Validates that CP attention produces results matching
|
||
|
|
PyTorch SDPA within 1e-2 tolerance.
|
||
|
|
|
||
|
|
Test scenarios:
|
||
|
|
- Pure prefill (seq_lens == query_lens), PCP=1, DCP=1
|
||
|
|
- Pure decode, PCP=1, DCP=1 (context in KV cache)
|
||
|
|
- Mixed decode+prefill (prefill has seq_lens == query_lens), PCP=1, DCP=1
|
||
|
|
- MTP (Multi-Token Prediction) decode, PCP=1, DCP=1
|
||
|
|
"""
|
||
|
|
|
||
|
|
@pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"batch_spec_name",
|
||
|
|
[
|
||
|
|
"single_prefill",
|
||
|
|
"small_prefill",
|
||
|
|
"medium_prefill",
|
||
|
|
"large_prefill",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
@pytest.mark.parametrize("model", MODELS)
|
||
|
|
@patch_distributed_groups(dcp_size=1, pcp_size=1)
|
||
|
|
def test_cp_prefill_precision(
|
||
|
|
self,
|
||
|
|
mock_all2all,
|
||
|
|
mock_dcp,
|
||
|
|
mock_pcp,
|
||
|
|
batch_spec_name,
|
||
|
|
model,
|
||
|
|
):
|
||
|
|
batch_spec = BATCH_SPECS[batch_spec_name]
|
||
|
|
_test_cp_prefill_precision_no_cp(batch_spec, model)
|
||
|
|
|
||
|
|
@pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"batch_spec_name",
|
||
|
|
[
|
||
|
|
"single_decode",
|
||
|
|
"small_decode",
|
||
|
|
"medium_decode",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
@pytest.mark.parametrize("model", MODELS)
|
||
|
|
@patch_distributed_groups(dcp_size=1, pcp_size=1)
|
||
|
|
def test_cp_decode_precision(
|
||
|
|
self,
|
||
|
|
mock_all2all,
|
||
|
|
mock_dcp,
|
||
|
|
mock_pcp,
|
||
|
|
batch_spec_name,
|
||
|
|
model,
|
||
|
|
):
|
||
|
|
batch_spec = BATCH_SPECS[batch_spec_name]
|
||
|
|
_test_cp_decode_precision_no_cp(batch_spec, model)
|
||
|
|
|
||
|
|
@pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"batch_spec_name",
|
||
|
|
[
|
||
|
|
"mixed_small",
|
||
|
|
"mixed_medium",
|
||
|
|
"mixed_large",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
@pytest.mark.parametrize("model", MODELS)
|
||
|
|
@patch_distributed_groups(dcp_size=1, pcp_size=1)
|
||
|
|
def test_cp_mixed_precision(
|
||
|
|
self,
|
||
|
|
mock_all2all,
|
||
|
|
mock_dcp,
|
||
|
|
mock_pcp,
|
||
|
|
batch_spec_name,
|
||
|
|
model,
|
||
|
|
):
|
||
|
|
batch_spec = BATCH_SPECS[batch_spec_name]
|
||
|
|
_test_cp_mixed_precision_no_cp(batch_spec, model)
|
||
|
|
|
||
|
|
@pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"batch_spec_name",
|
||
|
|
[
|
||
|
|
"mtp_1_plus_3_tiny",
|
||
|
|
"mtp_1_plus_3_small",
|
||
|
|
"mtp_1_plus_3_medium",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
@pytest.mark.parametrize("model", MODELS)
|
||
|
|
@patch_distributed_groups(dcp_size=1, pcp_size=1)
|
||
|
|
def test_cp_mtp_decode_precision(
|
||
|
|
self,
|
||
|
|
mock_all2all,
|
||
|
|
mock_dcp,
|
||
|
|
mock_pcp,
|
||
|
|
batch_spec_name,
|
||
|
|
model,
|
||
|
|
):
|
||
|
|
"""MTP decode: each request produces 1 target + 3 speculative tokens.
|
||
|
|
|
||
|
|
All tokens (context + new) are pre-populated in the KV cache.
|
||
|
|
FIA paged attention with actual_seq_lengths_kv enforces causal
|
||
|
|
masking per token within each request.
|
||
|
|
"""
|
||
|
|
batch_spec = BATCH_SPECS[batch_spec_name]
|
||
|
|
_test_cp_decode_precision_no_cp(batch_spec, model, decode_threshold=4)
|