0
tests/ut/attention/__init__.py
Normal file
0
tests/ut/attention/__init__.py
Normal file
0
tests/ut/attention/a2/__init__.py
Normal file
0
tests/ut/attention/a2/__init__.py
Normal file
623
tests/ut/attention/a2/test_attention_cp.py
Normal file
623
tests/ut/attention/a2/test_attention_cp.py
Normal file
@@ -0,0 +1,623 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.attention.utils import patch_distributed_groups
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.attention import utils as attention_utils
|
||||
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 (
|
||||
AscendMetadataForPrefill,
|
||||
AscendPCPMetadata,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendAttentionCPImpl(TestBase):
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def setUp(self):
|
||||
self.layer = MagicMock()
|
||||
self.layer.layer_name = "test_layer"
|
||||
self.layer._k_scale_float = 1.0
|
||||
self.layer._v_scale_float = 1.0
|
||||
|
||||
self.attention_type = MagicMock()
|
||||
self.attention_type.DECODER = "decoder"
|
||||
self.attention_type.ENCODER = "encoder"
|
||||
|
||||
self.attn_metadata = MagicMock()
|
||||
self.attn_metadata.return_value = "1"
|
||||
|
||||
self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
|
||||
self.layer_no_quant.layer_name = "test_layer"
|
||||
self.layer_no_quant._k_scale_float = 1.0
|
||||
self.layer_no_quant._v_scale_float = 1.0
|
||||
self.mock_vllm_config = MagicMock()
|
||||
self.mock_vllm_config.speculative_config = None
|
||||
self.config_patcher = patch(
|
||||
"vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
|
||||
)
|
||||
self.utils_config_patcher = patch(
|
||||
"vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_vllm_config
|
||||
)
|
||||
self.config_patcher.start()
|
||||
self.utils_config_patcher.start()
|
||||
attention_utils.needs_layer_aware_fia_graph_replay.cache_clear()
|
||||
self.addCleanup(attention_utils.needs_layer_aware_fia_graph_replay.cache_clear)
|
||||
self.addCleanup(self.utils_config_patcher.stop)
|
||||
self.addCleanup(self.config_patcher.stop)
|
||||
|
||||
self.impl = AscendAttentionCPImpl(
|
||||
num_heads=8,
|
||||
head_size=64,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
def test_init(self):
|
||||
self.assertEqual(self.impl.pcp_size, 2)
|
||||
self.assertEqual(self.impl.pcp_rank, 0)
|
||||
self.assertEqual(self.impl.dcp_size, 2)
|
||||
self.assertEqual(self.impl.dcp_rank, 0)
|
||||
|
||||
def test_forward_prefill_cp(self):
|
||||
query = torch.randn(2, 4, 128)
|
||||
key = torch.randn(4, 1, 128)
|
||||
value = torch.randn(4, 1, 128)
|
||||
|
||||
def mock_attention_with_nomask_and_mask(q, k_mask, **kwargs):
|
||||
mock_output = torch.randn_like(q)
|
||||
mock_lse = torch.randn_like(k_mask)
|
||||
return mock_output, mock_lse
|
||||
|
||||
self.impl._attention_with_nomask_and_mask = MagicMock()
|
||||
self.impl._attention_with_nomask_and_mask.side_effect = mock_attention_with_nomask_and_mask
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.prefill = MagicMock()
|
||||
attn_metadata.prefill.pcp_metadata.q_head_idx = torch.tensor([0])
|
||||
attn_metadata.prefill.pcp_metadata.q_tail_idx = torch.tensor([1])
|
||||
attn_metadata.prefill.pcp_metadata.q_full_idx = torch.tensor([0, 1])
|
||||
attn_metadata.prefill.pcp_metadata.kv_with_q_head_mask_idx = torch.tensor([0])
|
||||
attn_metadata.prefill.pcp_metadata.kv_with_q_tail_nomask_idx = torch.tensor([0])
|
||||
attn_metadata.prefill.pcp_metadata.kv_with_q_tail_mask_idx = torch.tensor([0])
|
||||
attn_metadata.prefill.pcp_metadata.pcp_fa_query_idx = torch.tensor([0, 1])
|
||||
attn_metadata.prefill.pcp_metadata.pcp_use_hybrid_attn = False
|
||||
|
||||
output, attn_lse = self.impl._forward_prefill_cp(query, key, value, attn_metadata)
|
||||
|
||||
self.assertEqual(output.shape[0], 2)
|
||||
self.assertEqual(output.shape[1], 4)
|
||||
self.assertEqual(output.shape[2], 128)
|
||||
|
||||
@patch("torch_npu.npu_attention_update")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2)
|
||||
def test_forward_decode_pcp_dcp(
|
||||
self,
|
||||
mock_all2all,
|
||||
mock_dcp,
|
||||
mock_pcp,
|
||||
mock_get_forward_context,
|
||||
mock_npu_fused_infer_attention_score,
|
||||
mock_npu_attention_update,
|
||||
):
|
||||
query = torch.randn(2, 4, 64)
|
||||
self.impl.key_cache = torch.randn(100, 64, 1, 64)
|
||||
self.impl.value_cache = torch.randn(100, 64, 1, 64)
|
||||
|
||||
# Mock output
|
||||
mock_npu_attention_update.return_value = (torch.randn(2 * 4, 64), None)
|
||||
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
def mock_npu_fused_infer_attention_score_func(query, k_nope, value, **common_kwargs):
|
||||
mock_output = torch.randn_like(query)
|
||||
mock_lse = torch.randn(query.shape[0], query.shape[1], 1)
|
||||
return mock_output, mock_lse
|
||||
|
||||
mock_npu_fused_infer_attention_score.side_effect = mock_npu_fused_infer_attention_score_func
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.decode_meta = MagicMock()
|
||||
attn_metadata.num_decodes_flatten = 5
|
||||
attn_metadata.decode_meta.batch_seq_mask = torch.tensor([1, 0], dtype=torch.bool)
|
||||
output = self.impl._forward_decode_pcp_dcp(query, attn_metadata)
|
||||
|
||||
self.assertEqual(output.shape[0], 2)
|
||||
self.assertEqual(output.shape[1], 4)
|
||||
self.assertEqual(output.shape[2], 64)
|
||||
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_prefill_query_all_gather(self):
|
||||
query = torch.randn(2, 4, 128)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.prefill = MagicMock()
|
||||
attn_metadata.prefill.chunked_context = MagicMock()
|
||||
attn_metadata.prefill.chunked_context.cp_kv_recover_idx_for_chunk = torch.tensor([1, 2, 3, 0])
|
||||
output = self.impl._prefill_query_all_gather(attn_metadata, query)
|
||||
|
||||
self.assertEqual(output.shape[0], 4)
|
||||
self.assertEqual(output.shape[1], 8)
|
||||
self.assertEqual(output.shape[2], 128)
|
||||
|
||||
@patch("torch.ops.npu.npu_fused_infer_attention_score")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_compute_prefill_context(self, mock_npu_attention):
|
||||
block_num = 100
|
||||
block_size = 128
|
||||
kv_num_heads = 1
|
||||
head_size = 128
|
||||
kv_cache = (
|
||||
torch.randn(block_num, block_size, kv_num_heads, head_size),
|
||||
torch.randn(block_num, block_size, kv_num_heads, head_size),
|
||||
)
|
||||
|
||||
batch_size = 1024
|
||||
self.impl.head_size = head_size
|
||||
self.impl.num_heads = 4
|
||||
num_heads = self.impl.num_heads * self.impl.dcp_size
|
||||
query = torch.randn(batch_size, num_heads, head_size)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.prefill = MagicMock()
|
||||
attn_metadata.prefill.chunked_context = MagicMock()
|
||||
local_context_lens_allranks = torch.tensor([[[256, 256], [256, 256]]])
|
||||
attn_metadata.prefill.chunked_context.local_context_lens_allranks = local_context_lens_allranks
|
||||
attn_metadata.prefill.chunked_context.local_total_toks = local_context_lens_allranks[:, 0, 0].sum()
|
||||
|
||||
def mock_load_kv_for_chunk(attn_metadata, kv_cache, local_chunked_kv_lens_rank, query, total_toks):
|
||||
return torch.randn(total_toks, kv_num_heads, head_size), torch.randn(total_toks, kv_num_heads, head_size)
|
||||
|
||||
self.impl._load_kv_for_chunk = MagicMock()
|
||||
self.impl._load_kv_for_chunk.side_effect = mock_load_kv_for_chunk
|
||||
|
||||
mock_npu_attention.return_value = (
|
||||
torch.randn(batch_size, num_heads, head_size),
|
||||
torch.randn(batch_size, num_heads, 1),
|
||||
)
|
||||
|
||||
context_output = self.impl._compute_prefill_context(query, kv_cache, attn_metadata)
|
||||
local_context_output = torch.cat(context_output, dim=-1).permute([1, 2, 0]).contiguous()
|
||||
global_context_output = self.impl._gather_global_context_output(local_context_output)
|
||||
global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
|
||||
result_output, result_lse = self.impl._update_global_context_output(global_context_output)
|
||||
|
||||
self.assertEqual(result_output.shape[0], batch_size)
|
||||
self.assertEqual(result_output.shape[1], self.impl.num_heads)
|
||||
self.assertEqual(result_output.shape[2], head_size)
|
||||
self.assertEqual(result_lse.shape[0], batch_size)
|
||||
self.assertEqual(result_lse.shape[1], self.impl.num_heads)
|
||||
self.assertEqual(result_lse.shape[2], 1)
|
||||
|
||||
@patch("torch_npu.atb.npu_paged_cache_load")
|
||||
def test_load_kv_for_chunk(self, mock_npu_paged_cache_load):
|
||||
block_num = 100
|
||||
block_size = 128
|
||||
num_heads = 1
|
||||
head_size = 128
|
||||
|
||||
kv_cache = (
|
||||
torch.randn(block_num, block_size, num_heads, head_size),
|
||||
torch.randn(block_num, block_size, num_heads, head_size),
|
||||
)
|
||||
query = torch.randn(4, 8, 128)
|
||||
total_toks = 256
|
||||
local_chunked_kv_lens_rank = torch.randn(total_toks)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
|
||||
key, value = self.impl._load_kv_for_chunk(
|
||||
attn_metadata, kv_cache, local_chunked_kv_lens_rank, query, total_toks
|
||||
)
|
||||
|
||||
self.assertEqual(key.shape[0], total_toks)
|
||||
self.assertEqual(key.shape[1], num_heads)
|
||||
self.assertEqual(key.shape[2], head_size)
|
||||
self.assertEqual(value.shape[0], total_toks)
|
||||
self.assertEqual(value.shape[1], num_heads)
|
||||
self.assertEqual(value.shape[2], head_size)
|
||||
|
||||
@patch("torch_npu.Event", create=True)
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_reshape_and_cache(self, mock_event_class, mock_npu_reshape_and_cache):
|
||||
num_tokens = 4
|
||||
block_num = 100
|
||||
block_size = 128
|
||||
num_heads = 1
|
||||
head_size = 128
|
||||
self.impl.head_size = head_size
|
||||
self.impl.is_kv_producer = False
|
||||
|
||||
kv_cache = (
|
||||
torch.randn(block_num, block_size, num_heads, head_size),
|
||||
torch.randn(block_num, block_size, num_heads, head_size),
|
||||
)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.num_decode_tokens = 1
|
||||
attn_metadata.num_decodes = 1
|
||||
attn_metadata.num_prefills = 1
|
||||
attn_metadata.slot_mapping = torch.randn(2)
|
||||
attn_metadata.num_actual_tokens_pcp_padded = num_tokens * self.impl.pcp_size
|
||||
attn_metadata.prefill = MagicMock()
|
||||
attn_metadata.prefill.pcp_metadata.pcp_allgather_restore_idx = torch.tensor([0, 3, 1, 2, 0, 0, 0, 0])
|
||||
attn_metadata.prefill.pcp_metadata.pcp_use_hybrid_attn = False
|
||||
attn_metadata.prefill.pcp_metadata.pcp_padded_tokens_fla = 0
|
||||
attn_metadata.prefill.pcp_metadata.pcp_enter_fa_restore_idx = torch.arange(num_tokens * 3 * self.impl.pcp_size)
|
||||
attn_metadata.prefill.pcp_metadata.pcp_unpad_mask = torch.tensor(
|
||||
[True, False, True, True, True, True, True, True]
|
||||
)
|
||||
|
||||
query = torch.rand(num_tokens, num_heads, head_size)
|
||||
key = torch.randn(num_tokens, num_heads, head_size)
|
||||
value = torch.randn(num_tokens, num_heads, head_size)
|
||||
output = torch.rand(num_tokens, num_heads * head_size)
|
||||
|
||||
query, key, value, output = self.impl.reshape_and_cache(query, key, value, kv_cache, attn_metadata, output)
|
||||
self.assertEqual(key.shape[0], num_tokens * self.impl.pcp_size)
|
||||
self.assertEqual(key.shape[1], num_heads)
|
||||
self.assertEqual(key.shape[2], head_size)
|
||||
self.assertEqual(value.shape[0], num_tokens * self.impl.pcp_size)
|
||||
self.assertEqual(value.shape[1], num_heads)
|
||||
self.assertEqual(value.shape[2], head_size)
|
||||
|
||||
|
||||
class TestUpdateNpuAttnOutLse(TestBase):
|
||||
@patch_distributed_groups(needs_mocks=False)
|
||||
def setUp(self):
|
||||
self.layer = MagicMock()
|
||||
self.layer.layer_name = "test_layer"
|
||||
self.layer._k_scale_float = 1.0
|
||||
self.layer._v_scale_float = 1.0
|
||||
|
||||
self.attention_type = MagicMock()
|
||||
self.attention_type.DECODER = "decoder"
|
||||
self.attention_type.ENCODER = "encoder"
|
||||
|
||||
self.attn_metadata = MagicMock()
|
||||
self.attn_metadata.return_value = "1"
|
||||
|
||||
self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
|
||||
self.layer_no_quant.layer_name = "test_layer"
|
||||
self.layer_no_quant._k_scale_float = 1.0
|
||||
self.layer_no_quant._v_scale_float = 1.0
|
||||
self.mock_vllm_config = MagicMock()
|
||||
self.mock_vllm_config.speculative_config = None
|
||||
self.config_patcher = patch(
|
||||
"vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
|
||||
)
|
||||
self.utils_config_patcher = patch(
|
||||
"vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_vllm_config
|
||||
)
|
||||
self.config_patcher.start()
|
||||
self.utils_config_patcher.start()
|
||||
attention_utils.needs_layer_aware_fia_graph_replay.cache_clear()
|
||||
self.addCleanup(attention_utils.needs_layer_aware_fia_graph_replay.cache_clear)
|
||||
self.addCleanup(self.utils_config_patcher.stop)
|
||||
self.addCleanup(self.config_patcher.stop)
|
||||
|
||||
self.impl = AscendAttentionCPImpl(
|
||||
num_heads=8,
|
||||
head_size=64,
|
||||
scale=0.125,
|
||||
num_kv_heads=2,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
self.impl.pcp_size = 1
|
||||
self.impl.dcp_size = 1
|
||||
|
||||
self.batch_size = 2
|
||||
|
||||
# sequence length per batch
|
||||
self.q_lens_per_batch = [32, 64]
|
||||
self.kv_lens_nomask_per_batch = [32, 64]
|
||||
self.kv_lens_mask_per_batch = [32, 64]
|
||||
|
||||
# TND layout requires cumulative sum computation.
|
||||
self.q_seqlens_cumsum = self._cumsum(self.q_lens_per_batch) # [32, 96]
|
||||
self.kv_seqlens_nomask_cumsum = self._cumsum(self.kv_lens_nomask_per_batch) # [32, 96]
|
||||
self.kv_seqlens_mask_cumsum = self._cumsum(self.kv_lens_mask_per_batch) # [32, 96]
|
||||
|
||||
# Compute T value in TND layout
|
||||
self.q_total_tokens = self.q_seqlens_cumsum[-1]
|
||||
self.kv_total_nomask = self.kv_seqlens_nomask_cumsum[-1] #
|
||||
self.kv_total_mask = self.kv_seqlens_mask_cumsum[-1]
|
||||
|
||||
def _cumsum(self, arr: list[int]) -> list[int]:
|
||||
result = []
|
||||
total = 0
|
||||
for val in arr:
|
||||
total += val
|
||||
result.append(total)
|
||||
return result
|
||||
|
||||
def _build_attn_metadata(self, with_chunked_context=False):
|
||||
attn_metadata = AscendMetadata()
|
||||
attn_metadata.num_prefills = self.batch_size
|
||||
attn_metadata.num_decodes = 0
|
||||
attn_metadata.num_actual_tokens = self.q_total_tokens
|
||||
|
||||
prefill_metadata = AscendMetadataForPrefill()
|
||||
pcp_metadata = AscendPCPMetadata()
|
||||
pcp_metadata.attn_mask_seqlens = self.kv_seqlens_mask_cumsum
|
||||
pcp_metadata.head_attn_nomask_seqlens = self.kv_seqlens_nomask_cumsum
|
||||
pcp_metadata.tail_attn_nomask_seqlens = self.kv_seqlens_nomask_cumsum
|
||||
prefill_metadata.pcp_metadata = pcp_metadata
|
||||
|
||||
prefill_metadata.actual_seq_lengths_q = torch.tensor(self.q_seqlens_cumsum)
|
||||
|
||||
if with_chunked_context:
|
||||
chunked_context = AscendMetadataForPrefill.ChunkedContextMetadata(
|
||||
actual_chunk_seq_lengths=self.kv_seqlens_mask_cumsum,
|
||||
actual_seq_lengths_kv=self.kv_seqlens_mask_cumsum,
|
||||
starts=None,
|
||||
chunk_seq_mask_filtered_indices=None,
|
||||
)
|
||||
prefill_metadata.chunked_context = chunked_context
|
||||
else:
|
||||
prefill_metadata.chunked_context = None
|
||||
|
||||
attn_metadata.prefill = prefill_metadata
|
||||
attn_metadata.decode_meta = None
|
||||
return attn_metadata
|
||||
|
||||
@patch("torch.ops.npu.npu_fused_infer_attention_score")
|
||||
def test_attention_with_nomask_none(self, mock_npu_attention):
|
||||
# Mock input data
|
||||
q = torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size)
|
||||
q_seqlens = self.q_seqlens_cumsum
|
||||
k_nomask = None
|
||||
v_nomask = None
|
||||
kv_seqlens_nomask = self.kv_seqlens_nomask_cumsum
|
||||
k_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
v_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
kv_seqlens_mask = self.kv_seqlens_mask_cumsum
|
||||
mask = torch.randn(self.q_total_tokens, self.kv_total_mask)
|
||||
attn_metadata = self._build_attn_metadata(with_chunked_context=False)
|
||||
# Mock output
|
||||
mock_npu_attention.return_value = torch.randn(96, 8, 64), torch.randn(96, 8, 1)
|
||||
|
||||
# Call the method under test
|
||||
output, attn_lse = self.impl._attention_with_nomask_and_mask(
|
||||
q, q_seqlens, k_nomask, v_nomask, kv_seqlens_nomask, k_mask, v_mask, kv_seqlens_mask, mask, attn_metadata
|
||||
)
|
||||
|
||||
# Verify only mask attention was invoked
|
||||
mock_npu_attention.assert_called_with(
|
||||
q,
|
||||
k_mask,
|
||||
v_mask,
|
||||
num_heads=self.impl.num_heads,
|
||||
num_key_value_heads=self.impl.num_kv_heads,
|
||||
input_layout="TND",
|
||||
atten_mask=mask,
|
||||
scale=self.impl.scale,
|
||||
sparse_mode=3,
|
||||
antiquant_mode=0,
|
||||
antiquant_scale=None,
|
||||
softmax_lse_flag=True,
|
||||
actual_seq_lengths_kv=kv_seqlens_mask,
|
||||
actual_seq_lengths=q_seqlens,
|
||||
)
|
||||
# Assert the method call
|
||||
self.assertEqual(mock_npu_attention.call_count, 1)
|
||||
self.assertIsInstance(output, torch.Tensor)
|
||||
self.assertIsInstance(attn_lse, torch.Tensor)
|
||||
self.assertEqual(output.shape, (96, 8, 64))
|
||||
self.assertEqual(attn_lse.shape, (96, 8, 1))
|
||||
|
||||
@patch("torch.ops.npu.npu_fused_infer_attention_score")
|
||||
@patch("vllm_ascend.attention.context_parallel.attention_cp._update_out_and_lse")
|
||||
def test_attention_with_nomask_and_mask_chunk(self, mock_update_out_and_lse, mock_npu_fused_infer_attention_score):
|
||||
# Mock input data
|
||||
q = torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size)
|
||||
k_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
v_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
k_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
v_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
|
||||
mask = torch.randn(self.q_total_tokens, self.kv_total_mask)
|
||||
attn_metadata = self._build_attn_metadata(with_chunked_context=True)
|
||||
|
||||
# Mock output
|
||||
mock_npu_fused_infer_attention_score.return_value = (
|
||||
torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size),
|
||||
torch.randn(self.q_total_tokens, self.impl.num_heads, 1),
|
||||
)
|
||||
mock_update_out_and_lse.return_value = (
|
||||
torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size),
|
||||
torch.randn(self.q_total_tokens, self.impl.num_heads, 1),
|
||||
)
|
||||
# Call the method under test
|
||||
output, attn_lse = self.impl._attention_with_nomask_and_mask(
|
||||
q=q,
|
||||
q_seqlens=self.q_seqlens_cumsum,
|
||||
k_nomask=k_nomask,
|
||||
v_nomask=v_nomask,
|
||||
kv_seqlens_nomask=self.kv_seqlens_nomask_cumsum,
|
||||
k_mask=k_mask,
|
||||
v_mask=v_mask,
|
||||
kv_seqlens_mask=self.kv_seqlens_mask_cumsum,
|
||||
mask=mask,
|
||||
attn_metadata=attn_metadata,
|
||||
)
|
||||
# Assert the method call
|
||||
self.assertEqual(mock_npu_fused_infer_attention_score.call_count, 2)
|
||||
self.assertIsNotNone(output)
|
||||
self.assertIsNotNone(attn_lse)
|
||||
|
||||
@patch("torch.ops.npu.npu_fused_infer_attention_score")
|
||||
@patch("vllm_ascend.attention.context_parallel.attention_cp._npu_attn_out_lse_update")
|
||||
def test_attention_with_nomask_and_mask_nochunk(
|
||||
self, mock_npu_attn_out_lse_update, mock_npu_fused_infer_attention_score
|
||||
):
|
||||
# Mock input data
|
||||
q = torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size)
|
||||
k_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
v_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
k_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
v_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
|
||||
mask = torch.randn(self.q_total_tokens, self.kv_total_mask)
|
||||
|
||||
attn_metadata = self._build_attn_metadata(with_chunked_context=True)
|
||||
attn_metadata.prefill.chunked_context = None
|
||||
|
||||
# Mock output
|
||||
mock_npu_fused_infer_attention_score.return_value = (
|
||||
torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size),
|
||||
torch.randn(self.q_total_tokens, self.impl.num_heads, 1),
|
||||
)
|
||||
mock_npu_attn_out_lse_update.return_value = torch.randn(
|
||||
self.q_total_tokens, self.impl.num_heads, self.impl.head_size
|
||||
)
|
||||
|
||||
# Call the method under test
|
||||
output, attn_lse = self.impl._attention_with_nomask_and_mask(
|
||||
q=q,
|
||||
q_seqlens=self.q_seqlens_cumsum,
|
||||
k_nomask=k_nomask,
|
||||
v_nomask=v_nomask,
|
||||
kv_seqlens_nomask=self.kv_seqlens_nomask_cumsum,
|
||||
k_mask=k_mask,
|
||||
v_mask=v_mask,
|
||||
kv_seqlens_mask=self.kv_seqlens_mask_cumsum,
|
||||
mask=mask,
|
||||
attn_metadata=attn_metadata,
|
||||
)
|
||||
# Assert the method call
|
||||
mock_npu_attn_out_lse_update.assert_called_once()
|
||||
self.assertEqual(mock_npu_fused_infer_attention_score.call_count, 2)
|
||||
self.assertIsNotNone(output)
|
||||
self.assertEqual(attn_lse, None)
|
||||
|
||||
@patch("vllm_ascend.attention.context_parallel.attention_cp._npu_attn_out_lse_update")
|
||||
def test_update_chunk_attn_out_lse_with_current_attn_out_lse(self, mock_npu_attn_out_lse_update):
|
||||
# Mock input data
|
||||
current_attn_output_prefill = torch.randn(32764, 8, 128)
|
||||
current_attn_lse_prefill = torch.randn(32764, 8, 1)
|
||||
attn_output_full_chunk = torch.randn(65528, 8, 128)
|
||||
attn_lse_full_chunk = torch.randn(65528, 8, 1)
|
||||
prefill_query = torch.randn(32764, 8, 128)
|
||||
# mock attn_metadata
|
||||
attn_metadata = self._build_attn_metadata(with_chunked_context=True)
|
||||
attn_metadata.prefill.chunked_context.chunk_seq_mask_filtered_indices = torch.arange(32764, dtype=torch.int32)
|
||||
attn_metadata.prefill.chunked_context.kv_inverse_idx_for_chunk = torch.arange(32764, dtype=torch.int32)
|
||||
# Mock output
|
||||
mock_npu_attn_out_lse_update.return_value = torch.randn(32764, 8, 128)
|
||||
# test pcp_size > 1
|
||||
self.impl.pcp_size = 2
|
||||
self.impl.pcp_rank = 0
|
||||
self.impl.dcp_group = None
|
||||
self.impl.pcp_group = None
|
||||
# Call the method under test
|
||||
self.impl._update_chunk_attn_out_lse_with_current_attn_out_lse(
|
||||
current_attn_output_prefill,
|
||||
current_attn_lse_prefill,
|
||||
attn_output_full_chunk,
|
||||
attn_lse_full_chunk,
|
||||
prefill_query,
|
||||
attn_metadata,
|
||||
)
|
||||
# Assert the method call
|
||||
mock_npu_attn_out_lse_update.assert_called_once()
|
||||
# test pcp_size = 1
|
||||
self.impl.pcp_size = 1
|
||||
self.impl._update_chunk_attn_out_lse_with_current_attn_out_lse(
|
||||
current_attn_output_prefill,
|
||||
current_attn_lse_prefill,
|
||||
attn_output_full_chunk,
|
||||
attn_lse_full_chunk,
|
||||
prefill_query,
|
||||
attn_metadata,
|
||||
)
|
||||
self.assertEqual(mock_npu_attn_out_lse_update.call_count, 2)
|
||||
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=3)
|
||||
def test_update_chunk_attn_out_lse_dcp2_pcp3(self, mock_all_to_all_single, mock_dcp, mock_pcp):
|
||||
# Mock input data
|
||||
prefix_chunk_output = torch.randn(2, 4, 8)
|
||||
prefix_chunk_lse = torch.randn(2, 4, 1)
|
||||
self.impl.dcp_size = 2
|
||||
self.impl.pcp_size = 3
|
||||
self.impl.head_size = 8
|
||||
|
||||
# Call the method under test
|
||||
chunk_data = torch.cat([prefix_chunk_output, prefix_chunk_lse], dim=-1).permute([1, 2, 0]).contiguous()
|
||||
global_context_output = self.impl._gather_global_context_output(chunk_data)
|
||||
global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
|
||||
output, lse = self.impl._update_global_context_output(global_context_output)
|
||||
|
||||
# Assert the method call
|
||||
self.assertIsInstance(output, torch.Tensor)
|
||||
self.assertIsInstance(lse, torch.Tensor)
|
||||
self.assertEqual(output.shape, (2, 2, 8))
|
||||
self.assertEqual(lse.shape, (2, 2, 1))
|
||||
|
||||
mock_all_to_all_single.assert_called_once()
|
||||
mock_pcp.all_gather.assert_called_once()
|
||||
|
||||
@patch_distributed_groups(dcp_size=2)
|
||||
def test_update_chunk_attn_out_lse_dcp2_pcp1(self, mock_all_to_all_single, mock_dcp, mock_pcp):
|
||||
# Mock input data
|
||||
prefix_chunk_output = torch.randn(2, 4, 8)
|
||||
prefix_chunk_lse = torch.randn(2, 4, 1)
|
||||
|
||||
self.impl.dcp_size = 2
|
||||
self.impl.pcp_size = 1
|
||||
self.impl.head_size = 8
|
||||
|
||||
# Call the method under test
|
||||
chunk_data = torch.cat([prefix_chunk_output, prefix_chunk_lse], dim=-1).permute([1, 2, 0]).contiguous()
|
||||
global_context_output = self.impl._gather_global_context_output(chunk_data)
|
||||
global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
|
||||
output, lse = self.impl._update_global_context_output(global_context_output)
|
||||
|
||||
# Assert the method call
|
||||
self.assertIsInstance(output, torch.Tensor)
|
||||
self.assertIsInstance(lse, torch.Tensor)
|
||||
self.assertEqual(output.shape, (2, 2, 8))
|
||||
self.assertEqual(lse.shape, (2, 2, 1))
|
||||
|
||||
mock_all_to_all_single.assert_called_once()
|
||||
mock_pcp.all_gather.assert_not_called()
|
||||
|
||||
@patch_distributed_groups(pcp_size=2)
|
||||
def test_update_chunk_attn_out_lse_dcp1_pcp2(self, mock_all_to_all_single, mock_dcp, mock_pcp):
|
||||
# Mock input data
|
||||
prefix_chunk_output = torch.randn(2, 4, 8)
|
||||
prefix_chunk_lse = torch.randn(2, 4, 1)
|
||||
|
||||
self.impl.dcp_size = 1
|
||||
self.impl.pcp_size = 2
|
||||
self.impl.head_size = 8
|
||||
|
||||
# Call the method under test
|
||||
chunk_data = torch.cat([prefix_chunk_output, prefix_chunk_lse], dim=-1).permute([1, 2, 0]).contiguous()
|
||||
global_context_output = self.impl._gather_global_context_output(chunk_data)
|
||||
global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
|
||||
output, lse = self.impl._update_global_context_output(global_context_output)
|
||||
|
||||
# Assert the method call
|
||||
self.assertIsInstance(output, torch.Tensor)
|
||||
self.assertIsInstance(lse, torch.Tensor)
|
||||
self.assertEqual(output.shape, (2, 4, 8))
|
||||
self.assertEqual(lse.shape, (2, 4, 1))
|
||||
|
||||
mock_all_to_all_single.assert_not_called()
|
||||
mock_pcp.all_gather.assert_called_once()
|
||||
993
tests/ut/attention/a2/test_attention_cp_precision.py
Normal file
993
tests/ut/attention/a2/test_attention_cp_precision.py
Normal file
@@ -0,0 +1,993 @@
|
||||
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)
|
||||
775
tests/ut/attention/a2/test_attention_v1.py
Normal file
775
tests/ut/attention/a2/test_attention_v1.py
Normal file
@@ -0,0 +1,775 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
import vllm_ascend.attention.attention_v1 as attn_module
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.attention.attention_v1 import (
|
||||
AscendAttentionBackend,
|
||||
AscendAttentionBackendImpl,
|
||||
AscendAttentionMetadataBuilder,
|
||||
AscendAttentionState,
|
||||
)
|
||||
from vllm_ascend.attention.kvcomp_attn.attention_utils import get_kvcomp_decode_params, reshape_and_cache_kvcomp
|
||||
from vllm_ascend.attention.utils import (
|
||||
AscendCommonAttentionMetadata,
|
||||
cache_graph_workspace,
|
||||
needs_layer_aware_fia_graph_replay,
|
||||
using_paged_attention,
|
||||
)
|
||||
from vllm_ascend.device.device_op import A5DeviceAdaptor
|
||||
from vllm_ascend.device.utils import FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE
|
||||
from vllm_ascend.utils import AscendDeviceType
|
||||
|
||||
LARGE_HEAD_PREFILL_PATH = "vllm_ascend.device.utils.npu_large_head_prefill_attention"
|
||||
|
||||
|
||||
class TestAttentionGraphHelpers(TestBase):
|
||||
def test_cache_graph_workspace_keeps_first_workspace_by_default(self):
|
||||
graph_params = SimpleNamespace(workspaces={1: torch.empty(4)})
|
||||
candidate_workspace = torch.empty(8)
|
||||
|
||||
result = cache_graph_workspace(graph_params, 1, candidate_workspace, use_max_workspace=False)
|
||||
|
||||
self.assertEqual(result.numel(), 4)
|
||||
self.assertEqual(graph_params.workspaces[1].numel(), 4)
|
||||
|
||||
def test_cache_graph_workspace_updates_to_larger_workspace(self):
|
||||
graph_params = SimpleNamespace(workspaces={1: torch.empty(4)})
|
||||
candidate_workspace = torch.empty(8)
|
||||
|
||||
result = cache_graph_workspace(graph_params, 1, candidate_workspace, use_max_workspace=True)
|
||||
|
||||
self.assertEqual(result.numel(), 8)
|
||||
self.assertEqual(graph_params.workspaces[1].numel(), 8)
|
||||
|
||||
def test_large_head_uses_paged_attention_on_a2(self):
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.speculative_config = None
|
||||
with patch("vllm_ascend.attention.utils.get_ascend_device_type", return_value=AscendDeviceType.A2):
|
||||
self.assertTrue(using_paged_attention(1, vllm_config, head_size=FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE))
|
||||
|
||||
|
||||
class TestAscendAttentionBackend(TestBase):
|
||||
def setUp(self):
|
||||
self.mock_config = MagicMock()
|
||||
|
||||
mock_parallel_config = MagicMock()
|
||||
mock_parallel_config.prefill_context_parallel_size = 1
|
||||
mock_parallel_config.decode_context_parallel_size = 1
|
||||
|
||||
self.mock_config.parallel_config = mock_parallel_config
|
||||
|
||||
self.utils_patcher = patch("vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_config)
|
||||
self.utils_patcher.start()
|
||||
|
||||
from vllm_ascend.attention.utils import enable_cp
|
||||
|
||||
enable_cp.cache_clear()
|
||||
|
||||
def test_get_name(self):
|
||||
self.assertEqual(AscendAttentionBackend.get_name(), "CUSTOM")
|
||||
|
||||
def test_get_impl_cls(self):
|
||||
self.assertEqual(AscendAttentionBackend.get_impl_cls(), AscendAttentionBackendImpl)
|
||||
|
||||
def test_get_builder_cls(self):
|
||||
self.assertEqual(AscendAttentionBackend.get_builder_cls(), AscendAttentionMetadataBuilder)
|
||||
|
||||
def test_get_kv_cache_shape_not(self):
|
||||
result = AscendAttentionBackend.get_kv_cache_shape(10, 20, 30, 40)
|
||||
self.assertEqual(result, (2, 10, 20, 30, 40))
|
||||
|
||||
def test_swap_blocks(self):
|
||||
src_kv_cache = [torch.zeros((10, 20)), torch.zeros((10, 20))]
|
||||
dst_kv_cache = [torch.zeros((10, 20)), torch.zeros((10, 20))]
|
||||
src_to_dst = torch.tensor([[0, 1], [2, 3]])
|
||||
AscendAttentionBackend.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
|
||||
self.assertTrue(torch.all(dst_kv_cache[0][1] == src_kv_cache[0][0]))
|
||||
self.assertTrue(torch.all(dst_kv_cache[1][3] == src_kv_cache[1][2]))
|
||||
|
||||
def test_copy_blocks(self):
|
||||
kv_caches = [torch.zeros((10, 20)), torch.zeros((10, 20))]
|
||||
src_to_dists = torch.tensor([[0, 1], [2, 3]])
|
||||
AscendAttentionBackend.copy_blocks(kv_caches, src_to_dists)
|
||||
self.assertTrue(torch.all(kv_caches[0][1] == kv_caches[0][0]))
|
||||
self.assertTrue(torch.all(kv_caches[1][3] == kv_caches[1][2]))
|
||||
|
||||
|
||||
class TestAscendAttentionMetadataBuilder(TestBase):
|
||||
def setUp(self):
|
||||
self.mock_vllm_config = MagicMock()
|
||||
self.mock_vllm_config.speculative_config = None
|
||||
self.mock_vllm_config.model_config.max_model_len = 640
|
||||
self.mock_vllm_config.model_config.hf_text_config.sliding_window = None
|
||||
self.mock_vllm_config.cache_config.block_size = 64
|
||||
self.mock_vllm_config.compilation_config.cudagraph_mode = None
|
||||
self.mock_vllm_config.scheduler_config.max_num_seqs = 10
|
||||
self.mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
|
||||
self.mock_device = "cpu:0"
|
||||
torch.Tensor.pin_memory = lambda x: x # noqa
|
||||
self.builder = AscendAttentionMetadataBuilder(None, None, self.mock_vllm_config, self.mock_device)
|
||||
|
||||
def test_reorder_batch(self):
|
||||
mock_input_batch = MagicMock()
|
||||
mock_scheduler_output = MagicMock()
|
||||
|
||||
result = self.builder.reorder_batch(mock_input_batch, mock_scheduler_output)
|
||||
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_unpadded_preserves_internal_seq_lens_cpu(self):
|
||||
internal_seq_lens_cpu = torch.tensor([4, 5, 6], dtype=torch.int32)
|
||||
common_attn_metadata = AscendCommonAttentionMetadata(
|
||||
query_start_loc=torch.tensor([0, 2, 5, 9]),
|
||||
query_start_loc_cpu=torch.tensor([0, 2, 5, 9]),
|
||||
seq_lens=torch.tensor([4, 5, 6], dtype=torch.int32),
|
||||
_seq_lens_cpu=internal_seq_lens_cpu,
|
||||
seq_lens_cpu=None,
|
||||
num_computed_tokens_cpu=None,
|
||||
num_reqs=3,
|
||||
num_actual_tokens=9,
|
||||
max_query_len=4,
|
||||
block_table_tensor=torch.zeros((3, 1), dtype=torch.int32),
|
||||
slot_mapping=torch.arange(9, dtype=torch.int32),
|
||||
causal=True,
|
||||
actual_seq_lengths_q=[2, 3, 4],
|
||||
positions=torch.arange(9),
|
||||
attn_state=AscendAttentionState.ChunkedPrefill,
|
||||
max_seq_len=6,
|
||||
)
|
||||
|
||||
unpadded_metadata = common_attn_metadata.unpadded(num_actual_tokens=5, num_actual_reqs=2)
|
||||
|
||||
self.assertTrue(torch.equal(unpadded_metadata._seq_lens_cpu, internal_seq_lens_cpu[:2]))
|
||||
self.assertIsNone(unpadded_metadata.seq_lens_cpu)
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.AscendMetadata")
|
||||
def test_build(self, mock_ascend_metadata):
|
||||
common_attn_metadata = AscendCommonAttentionMetadata(
|
||||
query_start_loc=torch.tensor([0, 2, 5, 9]),
|
||||
query_start_loc_cpu=torch.tensor([0, 2, 5, 9]),
|
||||
seq_lens_cpu=torch.tensor([4, 5, 6]),
|
||||
num_reqs=3,
|
||||
num_actual_tokens=15,
|
||||
max_query_len=6,
|
||||
decode_token_per_req=torch.tensor([1, 1, 1]),
|
||||
block_table_tensor=torch.zeros((10, 10)),
|
||||
slot_mapping=torch.tensor(range(20)),
|
||||
actual_seq_lengths_q=torch.tensor([0, 1, 2]),
|
||||
positions=torch.tensor([10, 10]),
|
||||
attn_state=AscendAttentionState.ChunkedPrefill,
|
||||
num_computed_tokens_cpu=None,
|
||||
seq_lens=None,
|
||||
max_seq_len=6,
|
||||
)
|
||||
mock_model = MagicMock()
|
||||
|
||||
self.builder.build(1, common_attn_metadata, mock_model)
|
||||
|
||||
|
||||
class TestAscendAttentionBackendImpl(TestBase):
|
||||
def setUp(self):
|
||||
self.mock_event = MagicMock()
|
||||
self.mock_event.record.return_value = None
|
||||
self.mock_event.wait.return_value = None
|
||||
|
||||
self.mock_stream = MagicMock()
|
||||
self.event_patcher = patch("torch_npu.npu.Event", return_value=self.mock_event)
|
||||
self.stream_patcher = patch("torch_npu.npu.current_stream", return_value=self.mock_stream)
|
||||
|
||||
self.event_patcher.start()
|
||||
self.stream_patcher.start()
|
||||
|
||||
self.layer = MagicMock()
|
||||
self.layer.layer_name = "test_layer"
|
||||
self.layer._k_scale_float = 1.0
|
||||
self.layer._v_scale_float = 1.0
|
||||
self.attention_type = MagicMock()
|
||||
self.attention_type.DECODER = "decoder"
|
||||
self.attention_type.ENCODER = "encoder"
|
||||
self.attn_metadata = MagicMock()
|
||||
self.attn_metadata.return_value = "1"
|
||||
self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
|
||||
self.layer_no_quant.layer_name = "test_layer"
|
||||
self.layer_no_quant._k_scale_float = 1.0
|
||||
self.layer_no_quant._v_scale_float = 1.0
|
||||
self.mock_vllm_config = MagicMock()
|
||||
self.config_patcher = patch(
|
||||
"vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
|
||||
)
|
||||
self.utils_config_patcher = patch(
|
||||
"vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_vllm_config
|
||||
)
|
||||
self.config_patcher.start()
|
||||
self.utils_config_patcher.start()
|
||||
needs_layer_aware_fia_graph_replay.cache_clear()
|
||||
self.addCleanup(needs_layer_aware_fia_graph_replay.cache_clear)
|
||||
self.addCleanup(self.utils_config_patcher.stop)
|
||||
self.addCleanup(self.config_patcher.stop)
|
||||
|
||||
self.impl = AscendAttentionBackendImpl(
|
||||
num_heads=8,
|
||||
head_size=64,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
self.impl_192 = AscendAttentionBackendImpl(
|
||||
num_heads=8,
|
||||
head_size=192,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
self.impl_error = AscendAttentionBackendImpl(
|
||||
num_heads=8,
|
||||
head_size=192,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=None,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
self.impl_swa = AscendAttentionBackendImpl(
|
||||
num_heads=8,
|
||||
head_size=64,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=1024,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
self.impl_swa_sink = AscendAttentionBackendImpl(
|
||||
num_heads=8,
|
||||
head_size=64,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=1024,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
sinks=torch.tensor([-3.4062], dtype=torch.bfloat16),
|
||||
)
|
||||
|
||||
self.impl_large_head = AscendAttentionBackendImpl(
|
||||
num_heads=8,
|
||||
head_size=FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE,
|
||||
scale=1.0,
|
||||
num_kv_heads=8,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="float16",
|
||||
logits_soft_cap=None,
|
||||
attn_type=self.attention_type.DECODER,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
def test_large_head_prefill_uses_device_operator_fallback(self, mock_get_forward_context):
|
||||
query = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
key = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
value = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
output = torch.empty_like(query)
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.PrefillNoCache
|
||||
metadata.actual_seq_lengths_q = [2]
|
||||
metadata.causal = True
|
||||
metadata.attn_mask = None
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
with patch(LARGE_HEAD_PREFILL_PATH, return_value=(torch.ones_like(query), None)) as mock_forward:
|
||||
result = self.impl_large_head.forward_impl(query, key, value, (), metadata, output)
|
||||
|
||||
mock_forward.assert_called_once()
|
||||
self.assertIs(result, output)
|
||||
self.assertTrue(torch.equal(result, torch.ones_like(query)))
|
||||
|
||||
def test_supported_head_prefill_uses_fia(self):
|
||||
query = torch.randn(2, 8, 64)
|
||||
key = torch.randn(2, 8, 64)
|
||||
value = torch.randn(2, 8, 64)
|
||||
output = torch.empty_like(query)
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.PrefillNoCache
|
||||
metadata.actual_seq_lengths_q = [2]
|
||||
|
||||
self.impl.forward_fused_infer_attention = MagicMock(return_value=output)
|
||||
with patch(LARGE_HEAD_PREFILL_PATH, return_value=(torch.empty_like(query), None)) as mock_forward:
|
||||
result = self.impl.forward_impl(query, key, value, (), metadata, output)
|
||||
|
||||
mock_forward.assert_not_called()
|
||||
self.impl.forward_fused_infer_attention.assert_called_once()
|
||||
self.assertIs(result, output)
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.using_paged_attention", return_value=True)
|
||||
def test_decode_uses_paged_attention(self, mock_using_pa):
|
||||
query = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
output = torch.empty_like(query)
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
|
||||
self.impl_large_head.forward_paged_attention = MagicMock(return_value=output)
|
||||
self.impl_large_head.forward_fused_infer_attention = MagicMock(return_value=output)
|
||||
|
||||
result = self.impl_large_head.forward_impl(query, None, None, (), metadata, output)
|
||||
|
||||
self.impl_large_head.forward_paged_attention.assert_called_once()
|
||||
self.impl_large_head.forward_fused_infer_attention.assert_not_called()
|
||||
self.assertIs(result, output)
|
||||
mock_using_pa.assert_called_once()
|
||||
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
def test_a5_device_operator_uses_fia_for_large_head(self, mock_fia):
|
||||
query = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
key = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
value = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.PrefillNoCache
|
||||
metadata.actual_seq_lengths_q = [2]
|
||||
|
||||
mock_fia.return_value = (torch.ones_like(query), None)
|
||||
with patch(LARGE_HEAD_PREFILL_PATH, return_value=(torch.empty_like(query), None)) as mock_forward:
|
||||
result = A5DeviceAdaptor.npu_fused_infer_attention_score(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
attn_metadata=metadata,
|
||||
key_cache=None,
|
||||
value_cache=None,
|
||||
current_key=key,
|
||||
current_value=value,
|
||||
num_heads=8,
|
||||
num_key_value_heads=8,
|
||||
head_size=FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE,
|
||||
scale=1.0,
|
||||
is_prefill_no_cache=True,
|
||||
block_table=None,
|
||||
input_layout="TND",
|
||||
block_size=128,
|
||||
actual_seq_lengths=[2],
|
||||
actual_seq_lengths_kv=[2],
|
||||
sparse_mode=3,
|
||||
)
|
||||
|
||||
mock_forward.assert_not_called()
|
||||
mock_fia.assert_called_once()
|
||||
self.assertEqual(result[0].shape, query.shape)
|
||||
|
||||
def test_forward_no_attn_metadata(self):
|
||||
"""Test forward pass when attn_metadata is None"""
|
||||
query = torch.randn(10, 8 * 64)
|
||||
key = torch.randn(10, 8 * 64)
|
||||
value = torch.randn(10, 8 * 64)
|
||||
kv_cache = torch.empty(2, 0, 0, 8, 64)
|
||||
layer = self.layer_no_quant
|
||||
output = torch.empty_like(query)
|
||||
|
||||
output = self.impl.forward(layer, query, key, value, kv_cache, None, output)
|
||||
|
||||
assert output.shape == (10, 8 * 64)
|
||||
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
def test_forward_fused_infer_attention(
|
||||
self, mock_get_forward_context, mock_npu_fused_infer_attention_score, mock_npu_reshape_and_cache
|
||||
):
|
||||
"""Test forward pass in PrefillCacheHit state"""
|
||||
query = torch.randn(10, 8, 64)
|
||||
key = torch.randn(10, 8, 64)
|
||||
value = torch.randn(10, 8, 64)
|
||||
kv_cache = torch.empty(2, 5, 128, 8, 64)
|
||||
output = torch.empty_like(query)
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.PrefillCacheHit
|
||||
metadata.attn_mask = torch.randn(1, 1, 10, 10)
|
||||
metadata.query_lens = torch.tensor([10])
|
||||
metadata.seq_lens = torch.tensor([10])
|
||||
metadata.actual_seq_lengths_q = [10]
|
||||
metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
|
||||
metadata.num_actual_tokens = 10
|
||||
metadata.num_decode_tokens = 0
|
||||
metadata.num_decodes = 0
|
||||
metadata.num_prefills = 10
|
||||
metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
|
||||
layer = self.layer_no_quant
|
||||
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
mock_npu_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), torch.ones(10, 8, 64))
|
||||
output = self.impl.forward(layer, query, key, value, kv_cache, metadata, output)
|
||||
|
||||
mock_npu_fused_infer_attention_score.assert_called_once()
|
||||
assert output.shape == (10, 8, 64)
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.using_paged_attention")
|
||||
@patch("torch_npu._npu_paged_attention")
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
def test_forward_paged_attention(
|
||||
self, mock_get_forward_context, mock_npu_reshape_and_cache, mock_paged_attention, mock_using_paged_attention
|
||||
):
|
||||
"""Test forward pass in DecodeOnly state"""
|
||||
query = torch.randn(4, 8 * 64)
|
||||
key = torch.randn(4, 8 * 64)
|
||||
value = torch.randn(4, 8 * 64)
|
||||
kv_cache = torch.empty(2, 5, 128, 8, 64)
|
||||
output = torch.empty_like(query)
|
||||
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
metadata.seq_lens = torch.tensor([4])
|
||||
metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
|
||||
metadata.num_actual_tokens = 4
|
||||
metadata.slot_mapping = torch.zeros(4, dtype=torch.long)
|
||||
metadata.num_decodes = 4
|
||||
metadata.num_prefills = 0
|
||||
layer = self.layer_no_quant
|
||||
mock_using_paged_attention.return_value = True
|
||||
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
output = self.impl.forward(layer, query, key, value, kv_cache, metadata, output)
|
||||
|
||||
mock_paged_attention.assert_called_once()
|
||||
assert output.shape == (4, 8 * 64)
|
||||
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
def test_forward_decode_only_swa(
|
||||
self, mock_npu_reshape_and_cache, mock_fused_infer_attention_score, mock_get_forward_context
|
||||
):
|
||||
"""Test forward pass in DecodeOnly state"""
|
||||
query = torch.randn(10, 8 * 64)
|
||||
key = torch.randn(10, 8 * 64)
|
||||
value = torch.randn(10, 8 * 64)
|
||||
kv_cache = torch.empty(2, 5, 128, 8, 64)
|
||||
output = torch.empty(10, 8, 64)
|
||||
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
metadata.seq_lens = torch.tensor([10] * 10)
|
||||
metadata.actual_seq_lengths_q = [10]
|
||||
metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
|
||||
metadata.num_actual_tokens = 100
|
||||
metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
|
||||
metadata.num_decodes = 10
|
||||
metadata.num_prefills = 0
|
||||
layer = self.layer_no_quant
|
||||
mock_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), 1)
|
||||
output = self.impl_swa.forward(layer, query, key, value, kv_cache, metadata, output)
|
||||
print(output.shape)
|
||||
mock_fused_infer_attention_score.assert_called_once()
|
||||
assert output.shape == (10, 8, 64)
|
||||
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score_v2")
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
def test_forward_decode_only_swa_sink(
|
||||
self, mock_npu_reshape_and_cache, mock_fused_infer_attention_score, mock_get_forward_context
|
||||
):
|
||||
"""Test forward pass in DecodeOnly state"""
|
||||
query = torch.randn(10, 8 * 64)
|
||||
key = torch.randn(10, 8 * 64)
|
||||
value = torch.randn(10, 8 * 64)
|
||||
kv_cache = torch.empty(2, 5, 128, 8, 64)
|
||||
output = torch.empty(10, 8, 64)
|
||||
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
metadata.seq_lens = torch.tensor([10] * 10)
|
||||
metadata.attn_mask = torch.randn(1, 1, 10, 10)
|
||||
metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
|
||||
metadata.num_actual_tokens = 100
|
||||
metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
|
||||
metadata.num_decodes = 10
|
||||
metadata.num_prefills = 0
|
||||
layer = self.layer_no_quant
|
||||
mock_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), 1)
|
||||
output = self.impl_swa_sink.forward(layer, query, key, value, kv_cache, metadata, output)
|
||||
print(output.shape)
|
||||
mock_fused_infer_attention_score.assert_called_once()
|
||||
assert output.shape == (10, 8, 64)
|
||||
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
@patch("torch_npu._npu_paged_attention")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
def test_forward_decode_only_swa_seq_len_mismatch(
|
||||
self,
|
||||
mock_npu_reshape_and_cache,
|
||||
mock_fused_infer_attention_score,
|
||||
mock_paged_attention,
|
||||
mock_get_forward_context,
|
||||
):
|
||||
"""Test forward pass in DecodeOnly state when seq)len_mismatch"""
|
||||
query = torch.randn(10, 8, 64)
|
||||
key = torch.randn(10, 8, 64)
|
||||
value = torch.randn(10, 8, 64)
|
||||
kv_cache = torch.empty(2, 5, 128, 8, 64)
|
||||
output = torch.empty_like(query)
|
||||
|
||||
metadata = self.attn_metadata
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
metadata.seq_lens = torch.tensor([10]) # len == 1 != query.size(0)==10
|
||||
metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
|
||||
metadata.num_actual_tokens = 10
|
||||
metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
|
||||
layer = self.layer_no_quant
|
||||
metadata.num_decodes = 10
|
||||
metadata.num_prefills = 0
|
||||
metadata.actual_seq_lengths_q = [10]
|
||||
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
mock_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), torch.ones(10, 8, 64))
|
||||
|
||||
output = self.impl_swa.forward(layer, query, key, value, kv_cache, metadata, output)
|
||||
|
||||
mock_paged_attention.assert_not_called()
|
||||
mock_fused_infer_attention_score.assert_called_once()
|
||||
|
||||
assert output.shape == (10, 8, 64)
|
||||
|
||||
def test_get_kvcomp_params_early_exit(self):
|
||||
"""
|
||||
Test that get_kvcomp_decode_params returns original values
|
||||
when kvcomp is disabled or hashk_cache is missing.
|
||||
"""
|
||||
query = torch.randn(10, 8, 64)
|
||||
key = torch.randn(10, 8, 64)
|
||||
block_table = torch.zeros(1, 5, dtype=torch.long)
|
||||
actual_seq_lengths_kv = [10]
|
||||
|
||||
metadata = MagicMock()
|
||||
# Mocking the case where hashk_caches is not properly initialized
|
||||
kvcomp_metadata = MagicMock()
|
||||
kvcomp_metadata.hashk_caches = [None]
|
||||
metadata.kvcomp_metadata = kvcomp_metadata
|
||||
|
||||
self.impl.enable_hamming_sparse = True
|
||||
self.impl.layerIndex = 0
|
||||
|
||||
res_bt, res_sl = get_kvcomp_decode_params(0, kvcomp_metadata, query, key, block_table, actual_seq_lengths_kv)
|
||||
|
||||
self.assertIs(res_bt, block_table)
|
||||
self.assertEqual(res_sl, actual_seq_lengths_kv)
|
||||
|
||||
def test_get_kvcomp_params_reuse(self):
|
||||
"""
|
||||
Test that in DecodeOnly state, if the current layer is a skip layer,
|
||||
it correctly reuses the Hamming results from a previous layer.
|
||||
"""
|
||||
query = torch.randn(10, 8, 64)
|
||||
key = torch.randn(10, 8, 64)
|
||||
block_table = torch.zeros(1, 5, dtype=torch.long)
|
||||
actual_seq_lengths_kv = [10]
|
||||
|
||||
self.impl.enable_hamming_sparse = True
|
||||
self.impl.layerIndex = 1
|
||||
|
||||
metadata = MagicMock()
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
expected_bt = torch.ones(1, 5)
|
||||
expected_sl = torch.tensor([5])
|
||||
|
||||
# Construct kvcomp_metadata
|
||||
kvcomp_metadata = MagicMock()
|
||||
kvcomp_metadata.hashk_caches = [MagicMock(), MagicMock()]
|
||||
kvcomp_metadata.hamming_output = expected_bt
|
||||
kvcomp_metadata.seq_lens_from_hamming = expected_sl
|
||||
|
||||
kvcomp_config = MagicMock()
|
||||
kvcomp_config.vllm_hash_attention_skip_layers = [False, True]
|
||||
kvcomp_config.top_k_index_reuse = [0, 0]
|
||||
kvcomp_metadata.kvcomp_config = kvcomp_config
|
||||
metadata.kvcomp_metadata = kvcomp_metadata
|
||||
|
||||
metadata.hamming_output_records = [{"new_block_table": expected_bt, "new_seq_lens_list": expected_sl}, None]
|
||||
|
||||
res_bt, res_sl = get_kvcomp_decode_params(1, kvcomp_metadata, query, key, block_table, actual_seq_lengths_kv)
|
||||
|
||||
self.assertTrue(torch.equal(res_bt, expected_bt))
|
||||
self.assertTrue(torch.equal(res_sl, expected_sl))
|
||||
|
||||
@patch("torch.ops._C_ascend.npu_reshape_and_cache_bnsd", create=True)
|
||||
def test_get_kvcomp_params_prefill(self, mock_reshape_and_cache):
|
||||
"""
|
||||
Test that in non-DecodeOnly state (e.g., Prefill), only Hash compute
|
||||
and Cache update are performed, and original params are returned.
|
||||
"""
|
||||
key = torch.randn(2, 8, 64)
|
||||
|
||||
self.impl.enable_hamming_sparse = True
|
||||
self.impl.layerIndex = 0
|
||||
|
||||
metadata = MagicMock()
|
||||
metadata.attn_state = AscendAttentionState.PrefillCacheHit
|
||||
metadata.slot_mapping = torch.zeros(2)
|
||||
metadata.actual_seq_lengths_q_device = torch.tensor([1, 1])
|
||||
metadata.num_actual_tokens = 2
|
||||
metadata.actual_query_lens = torch.tensor([1, 1], dtype=torch.int32)
|
||||
metadata.query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32)
|
||||
|
||||
kvcomp_metadata = MagicMock()
|
||||
kvcomp_metadata.hashk_caches = [MagicMock()]
|
||||
kvcomp_config = MagicMock()
|
||||
kvcomp_config.vllm_hash_attention_skip_layers = [False]
|
||||
kvcomp_metadata.kvcomp_config = kvcomp_config
|
||||
|
||||
# Mock HashEncoder
|
||||
hash_encoder = MagicMock()
|
||||
hash_encoder.compute_hash.return_value = torch.ones(2, 8, 8)
|
||||
kvcomp_metadata.hash_encoder = hash_encoder
|
||||
|
||||
metadata.kvcomp_metadata = kvcomp_metadata
|
||||
|
||||
reshape_and_cache_kvcomp(kvcomp_metadata, 0, key)
|
||||
|
||||
# Ensure cache update was called but Hamming was bypassed
|
||||
self.assertTrue(mock_reshape_and_cache.called)
|
||||
|
||||
@patch("torch.ops._C_ascend.npu_reshape_and_cache_bnsd", create=True)
|
||||
@patch("torch.ops._C_ascend.npu_hamming_dist_top_k", create=True)
|
||||
def test_get_kvcomp_params_decode_hamming(self, mock_hamming, mock_reshape):
|
||||
"""
|
||||
Test that in DecodeOnly state, the full flow including Hash computation
|
||||
and Hamming Distance Top-K operation is executed.
|
||||
"""
|
||||
query = torch.randn(2, 8, 64)
|
||||
key = torch.randn(2, 8, 64)
|
||||
block_table = torch.zeros(2, 5, dtype=torch.long)
|
||||
actual_seq_lengths_kv = [10, 10]
|
||||
|
||||
self.impl.enable_hamming_sparse = True
|
||||
self.impl.layerIndex = 0
|
||||
|
||||
metadata = MagicMock()
|
||||
metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
metadata.seq_lens = torch.tensor([10, 10])
|
||||
metadata.actual_seq_lengths_q_device = torch.tensor([1, 1])
|
||||
metadata.slot_mapping = torch.zeros(2)
|
||||
metadata.block_tables = block_table
|
||||
metadata.hamming_output_records = [None]
|
||||
metadata.num_actual_tokens = 2
|
||||
metadata.actual_query_lens = torch.tensor([1, 1], dtype=torch.int32)
|
||||
metadata.query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32)
|
||||
metadata.chunk_sizes_for_hamming = torch.tensor([64, 64], dtype=torch.int32)
|
||||
metadata.max_seq_len_for_hamming = 1024
|
||||
metadata.block_tables_for_hamming = torch.zeros(2, 10, dtype=torch.int32)
|
||||
metadata.new_seq_lens_list = torch.tensor([5, 5], dtype=torch.int32)
|
||||
|
||||
kvcomp_metadata = MagicMock()
|
||||
kvcomp_metadata.hashk_caches = [MagicMock()]
|
||||
|
||||
kvcomp_config = MagicMock()
|
||||
kvcomp_config.vllm_hash_attention_skip_layers = [False]
|
||||
kvcomp_config.chunk_size = 64
|
||||
kvcomp_metadata.kvcomp_config = kvcomp_config
|
||||
|
||||
# Mock necessary Hamming parameters
|
||||
kvcomp_metadata.chunk_sizes_for_hamming_full = torch.tensor([1, 1])
|
||||
kvcomp_metadata.topk_for_hamming_full = torch.tensor([1, 1])
|
||||
kvcomp_metadata.topk_for_hamming_full_cpu = torch.tensor([1, 1])
|
||||
kvcomp_metadata.hamming_output = torch.zeros(2, 1)
|
||||
|
||||
# Mock HashEncoder
|
||||
hash_encoder = MagicMock()
|
||||
hash_encoder.compute_hash.return_value = torch.ones(2, 8, 8)
|
||||
kvcomp_metadata.hash_encoder = hash_encoder
|
||||
|
||||
metadata.kvcomp_metadata = kvcomp_metadata
|
||||
|
||||
# Mock npu_hamming_dist_top_k output; note the squeeze(1) in implementation
|
||||
mock_hamming.return_value = torch.ones(2, 1, 5)
|
||||
|
||||
res_bt, res_sl = get_kvcomp_decode_params(0, kvcomp_metadata, query, key, block_table, actual_seq_lengths_kv)
|
||||
|
||||
self.assertTrue(mock_reshape.called)
|
||||
self.assertTrue(mock_hamming.called)
|
||||
# Verify shape after squeeze(1) becomes (2, 5)
|
||||
self.assertEqual(res_bt.shape, (2, 5))
|
||||
self.assertTrue(torch.equal(res_bt, torch.ones(2, 5)))
|
||||
# Verify the result is recorded in hamming_output_records
|
||||
self.assertTrue(torch.equal(kvcomp_metadata.hamming_output, torch.ones(2, 5)))
|
||||
self.assertIsNotNone(kvcomp_metadata.seq_lens_for_hamming)
|
||||
|
||||
@patch("torch.npu.stream")
|
||||
@patch("torch.npu.graph_task_update_begin")
|
||||
@patch("torch.npu.graph_task_update_end")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch("vllm_ascend.attention.attention_v1.get_graph_params")
|
||||
@patch("vllm_ascend.attention.attention_v1._EXTRA_CTX")
|
||||
@patch("vllm_ascend.attention.attention_v1.using_paged_attention", return_value=False)
|
||||
@patch("vllm_ascend.attention.attention_v1.needs_layer_aware_fia_graph_replay", return_value=False)
|
||||
@patch("vllm_ascend.attention.attention_v1._ATTN_KEYS_BUFFER", new=[])
|
||||
def test_update_graph_params(
|
||||
self,
|
||||
mock_needs_layer_aware_fia_graph_replay,
|
||||
mock_using_paged_attention,
|
||||
mock_EXTRA_CTX,
|
||||
mock_get_graph_params,
|
||||
mock_fia,
|
||||
mock_graph_task_update_end,
|
||||
mock_graph_task_update_begin,
|
||||
mock_stream,
|
||||
):
|
||||
"""Test behavior when _ATTN_KEYS_BUFFER is [] after dummy_run."""
|
||||
|
||||
mock_EXTRA_CTX.sinks = False
|
||||
mock_EXTRA_CTX.is_draft_model = False
|
||||
|
||||
param: list[MagicMock | None] = [MagicMock()] * 22
|
||||
param[16] = None # sliding_window
|
||||
param[17] = None # c8_k_aq_scale
|
||||
param[21] = None # layer_name
|
||||
|
||||
mock_get_graph_params.return_value.attn_params = {1: [tuple(param)] * 3}
|
||||
mock_get_graph_params.return_value.handles = {1: [MagicMock()] * 3}
|
||||
mock_get_graph_params.return_value.events = {1: [MagicMock()] * 3}
|
||||
|
||||
attn_metadata_keys = [
|
||||
"model.layers.10.self_attn.attn",
|
||||
"model.layers.2.self_attn.attn",
|
||||
"model.layers.5.self_attn.attn",
|
||||
]
|
||||
forward_context = MagicMock()
|
||||
forward_context.attn_metadata = {key: MagicMock() for key in attn_metadata_keys}
|
||||
# breakpoint()
|
||||
self.impl.update_graph_params(self.mock_stream, forward_context, 1, self.mock_vllm_config)
|
||||
|
||||
expected = [
|
||||
"model.layers.2.self_attn.attn",
|
||||
"model.layers.5.self_attn.attn",
|
||||
"model.layers.10.self_attn.attn",
|
||||
]
|
||||
self.assertEqual(attn_module._ATTN_KEYS_BUFFER, expected)
|
||||
self.assertEqual(mock_fia.out.call_count, 3)
|
||||
397
tests/ut/attention/a2/test_attention_v1_precision.py
Normal file
397
tests/ut/attention/a2/test_attention_v1_precision.py
Normal file
@@ -0,0 +1,397 @@
|
||||
from functools import partial
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
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 vllm.v1.attention.backend import AttentionType
|
||||
from vllm.v1.attention.selector import get_attn_backend
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
from tests.ut.attention.utils import (
|
||||
BatchSpec,
|
||||
create_and_prepopulate_kv_cache,
|
||||
create_common_attn_metadata,
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def default_vllm_config():
|
||||
mock_config = MagicMock()
|
||||
mock_config.compilation_config = MagicMock()
|
||||
mock_config.compilation_config.custom_ops = ["all"]
|
||||
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
|
||||
mock_config.additional_config = None
|
||||
mock_config.quant_config = None
|
||||
|
||||
with set_current_vllm_config(mock_config):
|
||||
yield mock_config
|
||||
|
||||
|
||||
BATCH_SPECS = {
|
||||
"small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
|
||||
"small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
|
||||
"mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
|
||||
"medium_decode": BatchSpec(
|
||||
seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
|
||||
query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
|
||||
),
|
||||
"medium_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]),
|
||||
"mixed_medium": BatchSpec(seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]),
|
||||
"large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
|
||||
"large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
|
||||
"mixed_large": BatchSpec(seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]),
|
||||
"single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
|
||||
"single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
|
||||
"small_encoder_prefill": BatchSpec(seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]),
|
||||
"medium_encoder_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]),
|
||||
"mtp_1_plus_3": BatchSpec(seq_lens=[256, 512, 1024, 1536], query_lens=[4, 4, 4, 4]),
|
||||
"mtp_1_plus_7": BatchSpec(seq_lens=[512, 1024, 2048, 3072], query_lens=[8, 8, 8, 8]),
|
||||
"mtp_small": BatchSpec(seq_lens=[64, 128, 256], query_lens=[4, 4, 4]),
|
||||
}
|
||||
|
||||
|
||||
class MockAttentionLayer:
|
||||
"""A mock attention layer for testing."""
|
||||
|
||||
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 run_attention_backend(
|
||||
kv_cache_spec: FullAttentionSpec,
|
||||
layer_names: list[str],
|
||||
vllm_config,
|
||||
device: torch.device,
|
||||
common_attn_metadata: AscendCommonAttentionMetadata,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
dtype: torch.dtype,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
sliding_window: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run attention computation using the specified backend's AttentionImpl."""
|
||||
|
||||
num_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()
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
|
||||
backend = get_attn_backend(head_size, dtype, None, use_mla=False, use_sparse=False, use_mm_prefix=False)
|
||||
impl_cls = backend.get_impl_cls()
|
||||
builder_cls = backend.get_builder_cls()
|
||||
|
||||
builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
|
||||
attn_metadata = builder.build(
|
||||
common_prefix_len=0,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
impl = impl_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
alibi_slopes=None,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type.value if hasattr(attn_type, "value") else attn_type,
|
||||
kv_cache_dtype="auto",
|
||||
logits_soft_cap=None,
|
||||
kv_sharing_target_layer_name=None,
|
||||
)
|
||||
|
||||
mock_layer = MockAttentionLayer(device)
|
||||
output = torch.empty_like(query)
|
||||
output = impl.forward(mock_layer, query, key, value, kv_cache, attn_metadata, output=output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def compute_sdpa_reference(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
k_contexts: list[torch.Tensor],
|
||||
v_contexts: list[torch.Tensor],
|
||||
batch_spec: BatchSpec,
|
||||
num_q_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
scale: float,
|
||||
device: torch.device,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
) -> torch.Tensor:
|
||||
"""Compute reference attention output using SDPA as golden baseline."""
|
||||
|
||||
all_sdpa_outputs = []
|
||||
query_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[query_offset : query_offset + q_len]
|
||||
k_new_i = key[kv_offset : kv_offset + q_len]
|
||||
v_new_i = value[kv_offset : kv_offset + q_len]
|
||||
|
||||
k_full_i = torch.cat([k_contexts[i], k_new_i], dim=0)
|
||||
v_full_i = torch.cat([v_contexts[i], v_new_i], dim=0)
|
||||
|
||||
q_sdpa_in = q_i.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_full_i.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_full_i.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
if attn_type == AttentionType.ENCODER_ONLY:
|
||||
attn_mask = None
|
||||
else:
|
||||
attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
|
||||
causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
|
||||
attn_mask[:, context_len:] = causal_mask
|
||||
|
||||
sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
|
||||
q_sdpa_in,
|
||||
k_sdpa_in,
|
||||
v_sdpa_in,
|
||||
attn_mask=attn_mask,
|
||||
is_causal=False,
|
||||
enable_gqa=(num_q_heads != num_kv_heads),
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))
|
||||
query_offset += q_len
|
||||
kv_offset += q_len
|
||||
|
||||
return torch.cat(all_sdpa_outputs, dim=0)
|
||||
|
||||
|
||||
def _test_npu_attention_correctness(
|
||||
batch_spec: BatchSpec,
|
||||
model: str,
|
||||
*,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
block_size: int = 128,
|
||||
atol: float = 1e-2,
|
||||
rtol: float = 1e-2,
|
||||
tensor_parallel_size: int = 1,
|
||||
):
|
||||
"""Test attention backend correctness with SDPA as reference."""
|
||||
|
||||
set_random_seed(42)
|
||||
|
||||
hf_config_override = None
|
||||
if tensor_parallel_size > 1:
|
||||
from vllm.config import ModelConfig
|
||||
|
||||
temp_config = ModelConfig(model=model, max_model_len=1)
|
||||
original_num_heads = temp_config.hf_text_config.num_attention_heads
|
||||
original_num_kv_heads = getattr(temp_config.hf_text_config, "num_key_value_heads", None)
|
||||
hf_config_override = {
|
||||
"num_attention_heads": original_num_heads // tensor_parallel_size,
|
||||
}
|
||||
if original_num_kv_heads is not None:
|
||||
hf_config_override["num_key_value_heads"] = max(1, original_num_kv_heads // tensor_parallel_size)
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
model_name=model,
|
||||
tensor_parallel_size=1,
|
||||
max_model_len=max(batch_spec.seq_lens),
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=8192,
|
||||
hf_config_override=hf_config_override,
|
||||
)
|
||||
device = torch.device("npu")
|
||||
|
||||
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
|
||||
|
||||
batch_size = batch_spec.batch_size
|
||||
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()
|
||||
sliding_window = vllm_config.model_config.get_sliding_window()
|
||||
dtype = torch.bfloat16
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
|
||||
k_contexts, v_contexts = [], []
|
||||
all_q, all_k, all_v = [], [], []
|
||||
|
||||
for i in range(batch_size):
|
||||
s_len = batch_spec.seq_lens[i]
|
||||
q_len = batch_spec.query_lens[i]
|
||||
context_len = s_len - q_len
|
||||
|
||||
q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
|
||||
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
|
||||
all_q.append(q)
|
||||
all_k.append(k_full[context_len:])
|
||||
all_v.append(v_full[context_len:])
|
||||
|
||||
k_contexts.append(k_full[:context_len])
|
||||
v_contexts.append(v_full[:context_len])
|
||||
|
||||
query = torch.cat(all_q, dim=0)
|
||||
key = torch.cat(all_k, dim=0)
|
||||
value = torch.cat(all_v, dim=0)
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(batch_spec, vllm_config.cache_config.block_size, device)
|
||||
if attn_type == AttentionType.ENCODER_ONLY:
|
||||
common_attn_metadata.causal = False
|
||||
|
||||
kv_cache = create_and_prepopulate_kv_cache(
|
||||
k_contexts=k_contexts,
|
||||
v_contexts=v_contexts,
|
||||
block_size=block_size,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_size=head_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
num_blocks=8192,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
randomize_blocks=False,
|
||||
)
|
||||
|
||||
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 = query.shape[0]
|
||||
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
|
||||
|
||||
backend_output = run_attention_backend(
|
||||
kv_cache_spec,
|
||||
["placeholder"],
|
||||
vllm_config,
|
||||
device,
|
||||
common_attn_metadata,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
kv_cache,
|
||||
dtype,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
)
|
||||
|
||||
sdpa_output = compute_sdpa_reference(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
k_contexts,
|
||||
v_contexts,
|
||||
batch_spec,
|
||||
num_q_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
scale,
|
||||
device,
|
||||
attn_type=attn_type,
|
||||
)
|
||||
|
||||
name = "GQA"
|
||||
assert backend_output.shape == sdpa_output.shape, (
|
||||
f"[{name}] shape {backend_output.shape} != SDPA shape {sdpa_output.shape}"
|
||||
)
|
||||
assert backend_output.dtype == sdpa_output.dtype, (
|
||||
f"[{name}] dtype {backend_output.dtype} != SDPA dtype {sdpa_output.dtype}"
|
||||
)
|
||||
|
||||
assert torch.isfinite(backend_output).all(), f"[{name}] produced non-finite values"
|
||||
|
||||
# Calculate and print differences for debugging
|
||||
diff = torch.abs(backend_output - sdpa_output)
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
print(f"\n[{name}] Max difference: {max_diff:.6f}, Mean difference: {mean_diff:.6f}")
|
||||
print(f"[{name}] Backend output range: [{backend_output.min().item():.6f}, {backend_output.max().item():.6f}]")
|
||||
print(f"[{name}] SDPA output range: [{sdpa_output.min().item():.6f}, {sdpa_output.max().item():.6f}]")
|
||||
|
||||
def error_msg(msg: str, backend_name: str):
|
||||
return f"[{backend_name}] output differs from SDPA baseline. {msg}"
|
||||
|
||||
torch.testing.assert_close(
|
||||
backend_output,
|
||||
sdpa_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=partial(error_msg, backend_name="GQA"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
[
|
||||
"small_decode",
|
||||
"small_prefill",
|
||||
"mixed_small",
|
||||
"medium_decode",
|
||||
"medium_prefill",
|
||||
"mixed_medium",
|
||||
"large_decode",
|
||||
"large_prefill",
|
||||
"single_decode",
|
||||
"single_prefill",
|
||||
"mtp_1_plus_3",
|
||||
"mtp_1_plus_7",
|
||||
"mtp_small",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["Qwen/Qwen3-8B"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
def test_causal_backend_correctness(default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int):
|
||||
"""Test backend's correctness with causal attention."""
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
|
||||
_test_npu_attention_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
[
|
||||
"small_encoder_prefill",
|
||||
"medium_encoder_prefill",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["Qwen/Qwen3-8B"])
|
||||
def test_encoder_only_backend_correctness(default_vllm_config, batch_spec_name: str, model: str):
|
||||
"""Test backend's correctness with encoder-only attention."""
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
|
||||
_test_npu_attention_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
attn_type=AttentionType.ENCODER_ONLY,
|
||||
)
|
||||
138
tests/ut/attention/a2/test_common_cp.py
Normal file
138
tests/ut/attention/a2/test_common_cp.py
Normal file
@@ -0,0 +1,138 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_ascend.attention.context_parallel.common_cp import (
|
||||
_npu_attention_update,
|
||||
_npu_attn_out_lse_update,
|
||||
_update_out_and_lse,
|
||||
)
|
||||
|
||||
|
||||
class TestCommonCP(unittest.TestCase):
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_decode_context_model_parallel_world_size")
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_dcp_group")
|
||||
@patch("torch.distributed.all_to_all_single")
|
||||
def test_process_attn_out_lse_complex(
|
||||
self, mock_all2all, mock_get_dcp_group, mock_get_dcp_size, mock_get_pcp_group
|
||||
):
|
||||
from vllm_ascend.attention.context_parallel.common_cp import _process_attn_out_lse
|
||||
|
||||
pcp_size = 2
|
||||
dcp_size = 2
|
||||
mock_get_pcp_group.return_value.world_size = pcp_size
|
||||
mock_get_dcp_size.return_value = dcp_size
|
||||
|
||||
mock_group = MagicMock()
|
||||
mock_get_dcp_group.return_value.device_group = mock_group
|
||||
|
||||
bs, num_heads, head_dim = 4, 8, 64
|
||||
attn_output = torch.randn(bs, num_heads, head_dim, dtype=torch.float16)
|
||||
softmax_lse = torch.randn(bs, num_heads, 1, dtype=torch.float16)
|
||||
|
||||
def fake_all_gather(tensor, dim=0):
|
||||
return torch.cat([tensor, tensor], dim=dim)
|
||||
|
||||
mock_get_pcp_group.return_value.all_gather = fake_all_gather
|
||||
|
||||
output = _process_attn_out_lse(attn_output, softmax_lse)
|
||||
|
||||
self.assertEqual(output.dtype, torch.float32)
|
||||
|
||||
# [4, 8, 64] + [4, 8, 1] -> [4, 8, 65] (Cat)
|
||||
# DCP All2All -> [4, 8, 65]
|
||||
# PCP AllGather (dim=0, size=2) -> [8, 8, 65]
|
||||
expected_shape = (bs * pcp_size, num_heads, head_dim + 1)
|
||||
self.assertEqual(output.shape, expected_shape)
|
||||
|
||||
mock_all2all.assert_called_once()
|
||||
called_args = mock_all2all.call_args
|
||||
self.assertEqual(called_args.kwargs["group"], mock_group)
|
||||
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_decode_context_model_parallel_world_size")
|
||||
def test_process_attn_out_lse_simple(self, mock_get_dcp_size, mock_get_pcp_group):
|
||||
from vllm_ascend.attention.context_parallel.common_cp import _process_attn_out_lse
|
||||
|
||||
mock_get_pcp_group.return_value.world_size = 1
|
||||
mock_get_dcp_size.return_value = 1
|
||||
|
||||
attn_output = torch.randn(2, 4, 16)
|
||||
softmax_lse = torch.randn(2, 4, 1)
|
||||
|
||||
output = _process_attn_out_lse(attn_output, softmax_lse)
|
||||
|
||||
# concat: [2, 4, 16+1]
|
||||
self.assertEqual(output.shape, (2, 4, 17))
|
||||
|
||||
@patch("torch_npu.npu_attention_update")
|
||||
def test_npu_attn_out_lse_update(self, mock_npu_attention_update):
|
||||
# Mock input data
|
||||
attn_lse_mask = torch.randn(8, 128, 1)
|
||||
attn_lse_nomask = torch.randn(8, 128, 1)
|
||||
attn_out_mask = torch.randn(8, 128, 128)
|
||||
attn_out_nomask = torch.randn(8, 128, 128)
|
||||
|
||||
mock_npu_attention_update.return_value = (torch.randn(8 * 128, 128), None)
|
||||
|
||||
# Call the method under test
|
||||
output = _npu_attn_out_lse_update(attn_lse_mask, attn_lse_nomask, attn_out_mask, attn_out_nomask)
|
||||
|
||||
# Assertions
|
||||
self.assertIsInstance(output, torch.Tensor)
|
||||
self.assertEqual(output.shape, (8, 128, 128))
|
||||
mock_npu_attention_update.assert_called_once()
|
||||
|
||||
def test_update_out_and_lse(self):
|
||||
# Mock input data: [N, batch, heads, head_size/1]
|
||||
out_list = torch.randn(3, 2, 4, 8)
|
||||
lse_list = torch.randn(3, 2, 4, 1)
|
||||
|
||||
# Call the method under test
|
||||
out_final, lse_final = _update_out_and_lse(out_list, lse_list)
|
||||
|
||||
# Assert shapes
|
||||
self.assertEqual(out_final.shape, (2, 4, 8))
|
||||
self.assertEqual(lse_final.shape, (2, 4, 1))
|
||||
self.assertIsInstance(out_final, torch.Tensor)
|
||||
self.assertIsInstance(lse_final, torch.Tensor)
|
||||
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_decode_context_model_parallel_world_size")
|
||||
@patch("torch_npu.npu_attention_update")
|
||||
def test_npu_attention_update(self, mock_npu_update, mock_get_dcp, mock_get_pcp):
|
||||
pcp_size = 2
|
||||
dcp_size = 2
|
||||
mock_get_pcp.return_value.world_size = pcp_size
|
||||
mock_get_dcp.return_value = dcp_size
|
||||
|
||||
head_size = 64
|
||||
S, H = 4, 8 # Sequence and Heads per segment
|
||||
|
||||
attn_out_lse = torch.randn(pcp_size * S, dcp_size * H, head_size + 1)
|
||||
|
||||
# mock NPU op return (N=PCP*DCP=4, S*H=32) -> [4*32, 64]
|
||||
mock_npu_update.return_value = (torch.randn(pcp_size * dcp_size * S * H, head_size), None)
|
||||
|
||||
# test
|
||||
result = _npu_attention_update(head_size, attn_out_lse)
|
||||
|
||||
self.assertEqual(result.shape, (pcp_size * dcp_size * S, H, head_size))
|
||||
mock_npu_update.assert_called_once()
|
||||
|
||||
def test_out_lse_reshape(self):
|
||||
# Mock input data
|
||||
out_list = torch.randn(3, 2, 4, 8) # [N, batch_size, num_heads, head_size]
|
||||
lse_list = torch.randn(3, 2, 4, 1) # [N, batch_size, num_heads, 1]
|
||||
|
||||
# Call the method under test
|
||||
out_final, lse_final = _update_out_and_lse(out_list, lse_list)
|
||||
|
||||
# Assert the method call
|
||||
self.assertEqual(out_final.shape, (2, 4, 8)) # [batch_size, num_heads, head_size]
|
||||
self.assertEqual(lse_final.shape, (2, 4, 1)) # [batch_size, num_heads, 1]
|
||||
|
||||
self.assertIsInstance(out_final, torch.Tensor)
|
||||
self.assertIsInstance(lse_final, torch.Tensor)
|
||||
899
tests/ut/attention/a2/test_mla_cp.py
Executable file
899
tests/ut/attention/a2/test_mla_cp.py
Executable file
@@ -0,0 +1,899 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
from vllm.distributed.parallel_state import GroupCoordinator
|
||||
|
||||
from tests.ut.attention.utils import patch_distributed_groups
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.ascend_config import init_ascend_config
|
||||
from vllm_ascend.attention.attention_v1 import AscendAttentionState
|
||||
from vllm_ascend.attention.context_parallel.common_cp import (
|
||||
CPChunkedContextMetadata,
|
||||
_npu_attention_update,
|
||||
_process_attn_out_lse,
|
||||
)
|
||||
from vllm_ascend.attention.context_parallel.mla_cp import AscendMlaCPImpl
|
||||
from vllm_ascend.attention.mla_v1 import ChunkedContextMetadata
|
||||
|
||||
|
||||
def get_pcp_split_info(pcp_rank, pcp_size, seq_lens):
|
||||
q_head_idx, q_tail_idx = [], []
|
||||
kv_with_q_head_nomask_idx, kv_with_q_head_mask_idx = [], []
|
||||
kv_with_q_tail_nomask_idx, kv_with_q_tail_mask_idx = [], []
|
||||
kv_tail_proj_idx: list[int] = []
|
||||
kv_with_q_head_attn_idx_in_tail = []
|
||||
kv_with_q_tail_attn_idx_in_tail = []
|
||||
chunk_seqlens = []
|
||||
kv_with_q_head_nomask_seqlens, kv_with_q_tail_nomask_seqlens = [], []
|
||||
head_actual_seq_lengths_kv = []
|
||||
tail_actual_seq_lengths_kv = []
|
||||
q_req_offset = 0
|
||||
kv_req_offset = 0
|
||||
q_head_chunk_id = pcp_rank
|
||||
q_tail_chunk_id = pcp_size * 2 - 1 - pcp_rank
|
||||
for i, seq_len in enumerate(seq_lens):
|
||||
chunk_len = seq_len // 2
|
||||
chunk_seqlens.append(chunk_len)
|
||||
q_head_idx.extend(list(range(q_req_offset, q_req_offset + chunk_len)))
|
||||
kv_with_q_head_nomask_idx.extend(list(range(kv_req_offset, kv_req_offset + chunk_len * q_head_chunk_id)))
|
||||
kv_with_q_head_mask_idx.extend(
|
||||
list(range(kv_req_offset + chunk_len * q_head_chunk_id, kv_req_offset + chunk_len * (q_head_chunk_id + 1)))
|
||||
)
|
||||
kv_with_q_head_nomask_seqlens.append(chunk_len * q_head_chunk_id)
|
||||
|
||||
q_tail_idx.extend(list(range(q_req_offset + chunk_len, q_req_offset + chunk_len * 2)))
|
||||
kv_with_q_tail_nomask_idx.extend(list(range(kv_req_offset, kv_req_offset + chunk_len * q_tail_chunk_id)))
|
||||
kv_with_q_tail_mask_idx.extend(
|
||||
list(range(kv_req_offset + chunk_len * q_tail_chunk_id, kv_req_offset + chunk_len * (q_tail_chunk_id + 1)))
|
||||
)
|
||||
kv_with_q_tail_nomask_seqlens.append(chunk_len * q_tail_chunk_id)
|
||||
|
||||
tail_proj_offset = len(kv_tail_proj_idx)
|
||||
tail_proj_len = chunk_len * (q_tail_chunk_id + 1)
|
||||
kv_tail_proj_idx.extend(list(range(kv_req_offset, kv_req_offset + tail_proj_len)))
|
||||
kv_with_q_head_attn_idx_in_tail.extend(
|
||||
list(range(tail_proj_offset, tail_proj_offset + chunk_len * (q_head_chunk_id + 1)))
|
||||
)
|
||||
kv_with_q_tail_attn_idx_in_tail.extend(list(range(tail_proj_offset, tail_proj_offset + tail_proj_len)))
|
||||
head_actual_seq_lengths_kv.append(len(kv_with_q_head_attn_idx_in_tail))
|
||||
tail_actual_seq_lengths_kv.append(len(kv_with_q_tail_attn_idx_in_tail))
|
||||
|
||||
q_req_offset += seq_len
|
||||
kv_req_offset += seq_len * pcp_size
|
||||
return (
|
||||
torch.tensor(q_head_idx),
|
||||
torch.tensor(q_tail_idx),
|
||||
torch.tensor(kv_with_q_head_nomask_idx),
|
||||
torch.tensor(kv_with_q_head_mask_idx),
|
||||
torch.tensor(kv_with_q_tail_nomask_idx),
|
||||
torch.tensor(kv_with_q_tail_mask_idx),
|
||||
chunk_seqlens,
|
||||
kv_with_q_head_nomask_seqlens,
|
||||
kv_with_q_tail_nomask_seqlens,
|
||||
torch.tensor(kv_tail_proj_idx),
|
||||
torch.tensor(kv_with_q_head_attn_idx_in_tail),
|
||||
torch.tensor(kv_with_q_tail_attn_idx_in_tail),
|
||||
head_actual_seq_lengths_kv,
|
||||
tail_actual_seq_lengths_kv,
|
||||
)
|
||||
|
||||
|
||||
def get_chunk_metadata(
|
||||
pcp_size,
|
||||
dcp_size,
|
||||
num_prefills,
|
||||
num_decodes,
|
||||
block_size,
|
||||
num_computed_tokens_cpu,
|
||||
num_reqs,
|
||||
chunked_prefill_workspace_size,
|
||||
num_computed_tokens_of_pcp_dcp,
|
||||
cp_local_block_size,
|
||||
):
|
||||
reqs_start = num_decodes
|
||||
context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
|
||||
max_context_len_cpu = context_lens_cpu.max().item()
|
||||
num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
|
||||
max_context_chunk = chunked_prefill_workspace_size // num_prefills_with_context_cpu
|
||||
max_context_chunk = max_context_chunk // block_size * block_size
|
||||
|
||||
assert max_context_chunk > 0
|
||||
num_chunks = (max_context_len_cpu + max_context_chunk - 1) // max_context_chunk
|
||||
chunk_starts = torch.arange(num_chunks, dtype=torch.int32).unsqueeze(1).expand(-1, num_prefills) * max_context_chunk
|
||||
chunk_ends = torch.min(context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk)
|
||||
chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
|
||||
cu_seq_lens_cpu = torch.zeros(num_chunks, num_prefills + 1, dtype=torch.int32)
|
||||
torch.cumsum(chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32)
|
||||
|
||||
def cdiv(a, b):
|
||||
return (a + b - 1) // b
|
||||
|
||||
if dcp_size * pcp_size > 1:
|
||||
if num_computed_tokens_of_pcp_dcp is not None:
|
||||
local_context_lens_allranks = torch.tensor(num_computed_tokens_of_pcp_dcp[reqs_start:num_reqs]).reshape(
|
||||
-1, dcp_size * pcp_size
|
||||
)
|
||||
# Note(qcs): The max local context lengths
|
||||
# padded to `cp_local_block_size`.
|
||||
padded_local_context_lens_cpu = (
|
||||
cdiv(
|
||||
context_lens_cpu,
|
||||
cp_local_block_size * pcp_size * dcp_size,
|
||||
)
|
||||
* cp_local_block_size
|
||||
)
|
||||
padded_local_max_context_chunk_across_ranks = (
|
||||
cdiv(
|
||||
max_context_chunk,
|
||||
cp_local_block_size * pcp_size * dcp_size,
|
||||
)
|
||||
* cp_local_block_size
|
||||
)
|
||||
local_chunk_starts = (
|
||||
torch.arange(num_chunks, dtype=torch.int32).unsqueeze(1).expand(-1, num_prefills)
|
||||
* padded_local_max_context_chunk_across_ranks
|
||||
)
|
||||
local_chunk_ends = torch.min(
|
||||
padded_local_context_lens_cpu.unsqueeze(0),
|
||||
local_chunk_starts + padded_local_max_context_chunk_across_ranks,
|
||||
)
|
||||
padded_local_chunk_seq_lens = (local_chunk_ends - local_chunk_starts).clamp(min=0)
|
||||
padded_local_cu_chunk_seq_lens_cpu = torch.zeros(num_chunks, num_prefills + 1, dtype=torch.int32)
|
||||
torch.cumsum(
|
||||
padded_local_chunk_seq_lens,
|
||||
dim=1,
|
||||
out=padded_local_cu_chunk_seq_lens_cpu[:, 1:],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
chunk_actual_seq_lengths_kv_list = [torch.cumsum(chunk_seq_lens[i], dim=0).tolist() for i in range(num_chunks)]
|
||||
chunked_context_metadata = CPChunkedContextMetadata(
|
||||
cu_seq_lens=cu_seq_lens_cpu.to(non_blocking=True),
|
||||
starts=local_chunk_starts.to(non_blocking=True),
|
||||
seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),
|
||||
max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
|
||||
chunk_seq_lens=chunk_seq_lens,
|
||||
chunk_seq_lens_npu=chunk_seq_lens,
|
||||
chunk_actual_seq_lengths_kv_list=chunk_actual_seq_lengths_kv_list,
|
||||
workspace=None,
|
||||
padded_chunk_seq_lens_npu=padded_local_chunk_seq_lens,
|
||||
padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
|
||||
local_context_lens_allranks=local_context_lens_allranks.tolist(),
|
||||
padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to(non_blocking=True),
|
||||
cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
|
||||
chunk_size=padded_local_max_context_chunk_across_ranks,
|
||||
)
|
||||
else:
|
||||
chunked_context_metadata = ChunkedContextMetadata(
|
||||
cu_seq_lens=cu_seq_lens_cpu.to(non_blocking=True),
|
||||
starts=chunk_starts.to(non_blocking=True),
|
||||
seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
|
||||
max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
|
||||
chunk_seq_lens=chunk_seq_lens,
|
||||
chunk_seq_lens_npu=chunk_seq_lens,
|
||||
workspace=None,
|
||||
)
|
||||
return chunked_context_metadata
|
||||
|
||||
|
||||
class TestAscendMLAImpl(TestBase):
|
||||
@patch("vllm.distributed.parallel_state._TP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
|
||||
@patch("vllm.distributed.get_tensor_model_parallel_world_size", return_value=2)
|
||||
@patch("vllm_ascend.attention.mla_v1.get_current_vllm_config")
|
||||
@patch("vllm_ascend.attention.mla_v1.get_ascend_config")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def setUp(self, ascend_config, get_current_vllm_config, mock_get_tp_size, mock_tp):
|
||||
mock_tp.world_size = 2
|
||||
mock_tp.rank_in_group = MagicMock()
|
||||
mock_tp.device_group = MagicMock()
|
||||
vllm_config = MagicMock()
|
||||
speculative_config = MagicMock()
|
||||
model_config = MagicMock()
|
||||
parallel_config = MagicMock()
|
||||
parallel_config.prefill_context_parallel_size = 1
|
||||
parallel_config.tensor_parallel_size = 2
|
||||
speculative_config.num_speculative_tokens = 4
|
||||
vllm_config.speculative_config = speculative_config
|
||||
model_config.dtype = torch.float16
|
||||
vllm_config.model_config = model_config
|
||||
get_current_vllm_config.return_value = vllm_config
|
||||
vllm_config.additional_config = {"refresh": True}
|
||||
vllm_config.parallel_config = parallel_config
|
||||
init_ascend_config(vllm_config)
|
||||
|
||||
num_heads = 256
|
||||
head_size = 1024
|
||||
scale = 0.1
|
||||
num_kv_heads = 8
|
||||
kv_cache_dtype = "auto"
|
||||
|
||||
kv_a_layernorm = MagicMock()
|
||||
kv_a_layernorm.weight = torch.randn(96)
|
||||
kv_a_layernorm.variance_epsilon = 1e-6
|
||||
kwargs = {
|
||||
"kv_lora_rank": 32,
|
||||
"qk_nope_head_dim": 64,
|
||||
"qk_rope_head_dim": 32,
|
||||
"qk_head_dim": 96,
|
||||
"v_head_dim": 128,
|
||||
"q_lora_rank": 64,
|
||||
"q_proj": MagicMock(),
|
||||
"q_b_proj": MagicMock(),
|
||||
"kv_b_proj": MagicMock(),
|
||||
"o_proj": MagicMock(),
|
||||
"kv_a_proj_with_mqa": MagicMock(),
|
||||
"fused_qkv_a_proj": MagicMock(),
|
||||
"kv_a_layernorm": kv_a_layernorm,
|
||||
"rotary_emb": MagicMock(),
|
||||
}
|
||||
|
||||
self.impl = AscendMlaCPImpl(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
blocksparse_params=None,
|
||||
logits_soft_cap=None,
|
||||
attn_type=None,
|
||||
kv_sharing_target_layer_name=None,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def test_init(self):
|
||||
self.assertEqual(self.impl.num_heads, 256)
|
||||
self.assertEqual(self.impl.head_size, 1024)
|
||||
self.assertEqual(self.impl.scale, 0.1)
|
||||
self.assertEqual(self.impl.num_kv_heads, 8)
|
||||
self.assertEqual(self.impl.kv_cache_dtype, "auto")
|
||||
self.assertEqual(self.impl.kv_lora_rank, 32)
|
||||
self.assertEqual(self.impl.qk_nope_head_dim, 64)
|
||||
self.assertEqual(self.impl.qk_rope_head_dim, 32)
|
||||
self.assertEqual(self.impl.qk_head_dim, 96)
|
||||
self.assertEqual(self.impl.v_head_dim, 128)
|
||||
self.assertIsNotNone(self.impl.q_proj)
|
||||
self.assertIsNotNone(self.impl.kv_b_proj)
|
||||
self.assertIsNotNone(self.impl.o_proj)
|
||||
self.assertIsNotNone(self.impl.kv_a_proj_with_mqa)
|
||||
self.assertIsNotNone(self.impl.kv_a_layernorm)
|
||||
self.assertEqual(self.impl.num_queries_per_kv, 32)
|
||||
self.assertEqual(self.impl.pcp_size, 2)
|
||||
self.assertEqual(self.impl.dcp_size, 2)
|
||||
|
||||
@patch("torch.ops.vllm.maybe_all_gather_and_maybe_unpad")
|
||||
@patch("vllm_ascend.attention.mla_v1.get_weight_prefetch_method", return_value=MagicMock())
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_mla_preprocess_dcp(self, mock_get_weight_prefetch_method, mock_maybe_all_gather_and_maybe_unpad):
|
||||
self.impl.num_kv_heads = 1
|
||||
self.impl.num_heads = 16
|
||||
self.impl.qk_rope_head_dim = 64
|
||||
self.impl.kv_lora_rank = 512
|
||||
self.impl.q_lora_rank = 1536
|
||||
self.impl.dcp_size = 2
|
||||
self.impl.pcp_size = 2
|
||||
block_num = 10
|
||||
block_size = 128
|
||||
batch_size = 2
|
||||
hidden_size = 1024
|
||||
hidden_states = torch.randn(batch_size, hidden_size)
|
||||
|
||||
kv_cache0 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.kv_lora_rank)
|
||||
kv_cache1 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.qk_rope_head_dim)
|
||||
kv_cache = (kv_cache0, kv_cache1)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.num_decodes = 2
|
||||
attn_metadata.num_prefills = 0
|
||||
attn_metadata.num_prefill_tokens = 0
|
||||
attn_metadata.num_decode_tokens = 2
|
||||
attn_metadata.num_actual_tokens = 2
|
||||
attn_metadata.slot_mapping = torch.arange(4)
|
||||
attn_metadata.decode.cos = torch.randn(2, 64)
|
||||
attn_metadata.decode.sin = torch.randn(2, 64)
|
||||
|
||||
self.impl.q_a_layernorm = MagicMock()
|
||||
self.impl.q_a_layernorm.return_value = torch.randn(attn_metadata.num_actual_tokens, self.impl.q_lora_rank)
|
||||
self.impl.kv_a_proj_with_mqa = MagicMock()
|
||||
self.impl.kv_a_proj_with_mqa.return_value = [
|
||||
torch.randn(batch_size, self.impl.num_heads, self.impl.qk_rope_head_dim + self.impl.kv_lora_rank)
|
||||
]
|
||||
self.impl.fused_qkv_a_proj = MagicMock()
|
||||
self.impl.fused_qkv_a_proj.return_value = [
|
||||
torch.randn(
|
||||
attn_metadata.num_actual_tokens,
|
||||
self.impl.qk_rope_head_dim + self.impl.kv_lora_rank + self.impl.q_lora_rank,
|
||||
)
|
||||
]
|
||||
|
||||
self.impl.rope_single = MagicMock(side_effect=lambda x, cos, sin: x)
|
||||
self.impl.exec_kv_decode = MagicMock()
|
||||
self.impl.exec_kv_decode.return_value = [MagicMock(), MagicMock()]
|
||||
|
||||
self.impl._q_proj_and_k_up_proj = MagicMock()
|
||||
self.impl._q_proj_and_k_up_proj.return_value = [
|
||||
torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.kv_lora_rank),
|
||||
torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.qk_rope_head_dim),
|
||||
]
|
||||
|
||||
mock_maybe_all_gather_and_maybe_unpad.side_effect = lambda x, label: x
|
||||
|
||||
decode_res, prefill_res = self.impl._mla_preprocess(
|
||||
"mock_layer", hidden_states, kv_cache, attn_metadata, need_gather_q_kv=False
|
||||
)
|
||||
|
||||
self.assertIsNotNone(decode_res)
|
||||
self.assertIsNone(prefill_res)
|
||||
|
||||
@patch("torch_npu._npu_reshape_and_cache")
|
||||
@patch("torch.ops.vllm.maybe_all_gather_and_maybe_unpad")
|
||||
@patch("vllm_ascend.attention.mla_v1.get_weight_prefetch_method", return_value=MagicMock())
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_mla_preprocess_pcp(
|
||||
self, mock_get_weight_prefetch_method, mock_maybe_all_gather_and_maybe_unpad, mock_npu_reshape_and_cache
|
||||
):
|
||||
self.impl.num_kv_heads = 1
|
||||
self.impl.num_heads = 16
|
||||
self.impl.qk_rope_head_dim = 64
|
||||
self.impl.kv_lora_rank = 512
|
||||
self.impl.q_lora_rank = 1536
|
||||
self.impl.dcp_size = 2
|
||||
self.impl.pcp_size = 2
|
||||
block_num = 10
|
||||
block_size = 128
|
||||
batch_size = 2
|
||||
hidden_size = 1024
|
||||
hidden_states = torch.randn(batch_size, hidden_size)
|
||||
|
||||
kv_cache0 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.kv_lora_rank)
|
||||
kv_cache1 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.qk_rope_head_dim)
|
||||
kv_cache = (kv_cache0, kv_cache1)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.num_decodes = 0
|
||||
attn_metadata.num_prefills = 2
|
||||
attn_metadata.num_prefill_tokens = 2
|
||||
attn_metadata.num_decode_tokens = 0
|
||||
attn_metadata.num_actual_tokens = 2
|
||||
attn_metadata.num_actual_tokens_pcp_padded = 4
|
||||
attn_metadata.prefill.pcp_metadata = MagicMock()
|
||||
attn_metadata.prefill.pcp_metadata.pcp_allgather_restore_idx = torch.arange(4)
|
||||
tail_projection_len = 3
|
||||
attn_metadata.prefill.pcp_metadata.kv_tail_proj_idx = torch.arange(tail_projection_len, dtype=torch.int64)
|
||||
attn_metadata.slot_mapping = torch.arange(4)
|
||||
attn_metadata.prefill.cos = torch.randn(2, 64)
|
||||
attn_metadata.prefill.sin = torch.randn(2, 64)
|
||||
|
||||
self.impl.q_a_layernorm = MagicMock()
|
||||
self.impl.q_a_layernorm.return_value = torch.randn(attn_metadata.num_actual_tokens, self.impl.q_lora_rank)
|
||||
self.impl.kv_a_proj_with_mqa = MagicMock()
|
||||
self.impl.kv_a_proj_with_mqa.return_value = [
|
||||
torch.randn(batch_size, self.impl.num_heads, self.impl.qk_rope_head_dim + self.impl.kv_lora_rank)
|
||||
]
|
||||
self.impl.fused_qkv_a_proj = MagicMock()
|
||||
self.impl.fused_qkv_a_proj.return_value = [
|
||||
torch.randn(
|
||||
attn_metadata.num_actual_tokens,
|
||||
self.impl.qk_rope_head_dim + self.impl.kv_lora_rank + self.impl.q_lora_rank,
|
||||
)
|
||||
]
|
||||
|
||||
self.impl.rope_single = MagicMock(side_effect=lambda x, cos, sin: x)
|
||||
self.impl.exec_kv_decode = MagicMock()
|
||||
self.impl.exec_kv_decode.return_value = [MagicMock(), MagicMock()]
|
||||
|
||||
self.impl._q_proj_and_k_up_proj = MagicMock()
|
||||
self.impl._q_proj_and_k_up_proj.return_value = [
|
||||
torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.kv_lora_rank),
|
||||
torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.qk_rope_head_dim),
|
||||
]
|
||||
|
||||
mock_maybe_all_gather_and_maybe_unpad.side_effect = lambda x, label: x
|
||||
|
||||
self.impl.kv_a_layernorm = MagicMock()
|
||||
self.impl.kv_a_layernorm.return_value = torch.randn(
|
||||
attn_metadata.num_prefill_tokens, self.impl.num_kv_heads, self.impl.kv_lora_rank
|
||||
)
|
||||
|
||||
self.impl.q_proj = MagicMock()
|
||||
self.impl.q_proj.return_value = [
|
||||
torch.randn(attn_metadata.num_prefill_tokens, self.impl.num_heads, self.impl.qk_head_dim)
|
||||
]
|
||||
self.impl.kv_b_proj = MagicMock()
|
||||
self.impl.kv_b_proj.return_value = [
|
||||
torch.randn(
|
||||
tail_projection_len,
|
||||
self.impl.num_heads,
|
||||
self.impl.v_head_dim + self.impl.qk_nope_head_dim,
|
||||
)
|
||||
]
|
||||
self.impl.rope_single = MagicMock(side_effect=lambda x, cos, sin: x)
|
||||
self.impl.exec_kv_decode = MagicMock()
|
||||
self.impl.exec_kv_decode.return_value = [MagicMock(), MagicMock()]
|
||||
self.impl.exec_kv_prefill = MagicMock()
|
||||
self.impl.exec_kv_prefill.return_value = [
|
||||
torch.randn(attn_metadata.num_prefill_tokens, self.impl.num_heads, self.impl.qk_rope_head_dim),
|
||||
torch.randn(attn_metadata.num_prefill_tokens, self.impl.num_heads, self.impl.kv_lora_rank),
|
||||
]
|
||||
|
||||
decode_res, prefill_res = self.impl._mla_preprocess(
|
||||
"mock_layer", hidden_states, kv_cache, attn_metadata, need_gather_q_kv=False
|
||||
)
|
||||
self.assertIsNone(decode_res)
|
||||
self.assertIsNotNone(prefill_res)
|
||||
self.impl.kv_b_proj.assert_called_once()
|
||||
self.assertEqual(self.impl.kv_b_proj.call_args.args[0].shape[0], tail_projection_len)
|
||||
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_process_attn_out_lse(self):
|
||||
self.impl.dcp_size = 2
|
||||
self.impl.pcp_size = 2
|
||||
|
||||
B = 2
|
||||
N = self.impl.num_heads
|
||||
self.impl.kv_lora_rank = 512
|
||||
|
||||
attn_output = torch.randn(B, N, self.impl.kv_lora_rank)
|
||||
softmax_lse = torch.randn(B, N, 1)
|
||||
|
||||
decode_metadata = MagicMock()
|
||||
decode_metadata.actual_seq_lengths_q = MagicMock()
|
||||
decode_metadata.seq_lens_list = MagicMock()
|
||||
result = _process_attn_out_lse(attn_output, softmax_lse)
|
||||
|
||||
self.assertEqual(result.shape[0], B * self.impl.pcp_size)
|
||||
self.assertEqual(result.shape[1], N)
|
||||
self.assertEqual(result.shape[2], self.impl.kv_lora_rank + 1)
|
||||
|
||||
@patch("vllm_ascend.ascend_forward_context.get_forward_context")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch("torch_npu.npu_attention_update")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_forward_decode_pcp_dcp(
|
||||
self, mock_npu_attention_update, mock_npu_fused_infer_attention_score, mock_get_forward_context
|
||||
):
|
||||
self.impl.dcp_size = 2
|
||||
self.impl.pcp_size = 2
|
||||
self.impl.num_kv_heads = 1
|
||||
self.impl.num_heads = 16
|
||||
self.impl.kv_lora_rank = 64
|
||||
self.impl.qk_nope_head_dim = 64
|
||||
self.impl.spec_token_num = 1
|
||||
B = 2
|
||||
N = self.impl.num_heads * self.impl.dcp_size
|
||||
BS = 128
|
||||
NB = 100
|
||||
|
||||
q_nope = torch.randn(B, N, self.impl.qk_nope_head_dim)
|
||||
q_pe = torch.randn(B, N, self.impl.qk_rope_head_dim)
|
||||
k_nope = torch.randn(NB, 1, BS, self.impl.kv_lora_rank)
|
||||
k_pe = torch.randn(NB, 1, BS, self.impl.qk_rope_head_dim)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.attn_state = AscendAttentionState.SpecDecoding
|
||||
attn_metadata.decode = MagicMock()
|
||||
attn_metadata.decode.actual_seq_lengths_q = MagicMock()
|
||||
attn_metadata.decode.seq_lens_list = MagicMock()
|
||||
|
||||
self.impl.enable_kv_nz = True
|
||||
self.impl.speculative_config = None
|
||||
mock_npu_attention_update.return_value = (torch.randn(B, self.impl.num_heads, self.impl.kv_lora_rank), None)
|
||||
mock_npu_fused_infer_attention_score.return_value = [
|
||||
torch.randn(B, N, 1, self.impl.kv_lora_rank),
|
||||
torch.randn(B, N, 1, 1),
|
||||
]
|
||||
mock_get_forward_context.return_value = MagicMock(capturing=False)
|
||||
|
||||
self.impl._v_up_proj = MagicMock()
|
||||
self.impl._v_up_proj.return_value = torch.randn(B, self.impl.v_head_dim)
|
||||
|
||||
result = self.impl._forward_decode(q_nope, q_pe, k_nope, k_pe, BS, attn_metadata)
|
||||
|
||||
self.assertEqual(result.shape[0], B)
|
||||
self.assertEqual(result.shape[1], self.impl.v_head_dim)
|
||||
|
||||
@patch("torch_npu.atb.npu_paged_cache_load")
|
||||
@patch("torch_npu.npu_attention_update")
|
||||
@patch("torch_npu.npu_fused_infer_attention_score")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2)
|
||||
def test_compute_prefill_context_with_dcp_pcp(
|
||||
self, mock_all2all, mock_dcp, mock_pcp, mock_fia, mock_update, mock_load
|
||||
):
|
||||
def mock_fia_attn(*args, **kwargs):
|
||||
q = args[0]
|
||||
v = args[2]
|
||||
return (
|
||||
torch.randn(q.shape[0], v.shape[1], v.shape[2], dtype=torch.float16),
|
||||
torch.randn(v.shape[1], q.shape[0], dtype=torch.float16),
|
||||
)
|
||||
|
||||
mock_fia.side_effect = mock_fia_attn
|
||||
|
||||
def mock_kv_b_proj(kv_c_normed):
|
||||
return (
|
||||
torch.randn(
|
||||
kv_c_normed.shape[0],
|
||||
self.impl.num_heads,
|
||||
self.impl.v_head_dim + self.impl.qk_nope_head_dim,
|
||||
dtype=torch.float16,
|
||||
),
|
||||
)
|
||||
|
||||
def mock_reorg_kvcache(
|
||||
allgatered_kv_c_normed: torch.Tensor,
|
||||
allgatered_k_pe: torch.Tensor,
|
||||
chunked_context: CPChunkedContextMetadata,
|
||||
chunk_idx: int,
|
||||
toks: int,
|
||||
):
|
||||
return torch.randn(
|
||||
chunked_context.cu_seq_lens_lst[chunk_idx][-1],
|
||||
allgatered_kv_c_normed.shape[1],
|
||||
allgatered_kv_c_normed.shape[2],
|
||||
), torch.randn(
|
||||
chunked_context.cu_seq_lens_lst[chunk_idx][-1], allgatered_k_pe.shape[1], allgatered_k_pe.shape[2]
|
||||
)
|
||||
|
||||
# mock proj
|
||||
self.impl.kv_b_proj.side_effect = mock_kv_b_proj
|
||||
|
||||
def mock_update_fn(lse_list, out_list, mode):
|
||||
total_len = out_list[0].shape[0]
|
||||
D = out_list[0].shape[1]
|
||||
return (torch.randn(total_len, D, dtype=torch.float32), None)
|
||||
|
||||
mock_update.side_effect = mock_update_fn
|
||||
NUM_BLOCKS, BLOCK_SIZE = 10, 32 # fixed
|
||||
USED_BLOCKS = 3
|
||||
# pcp_size, dcp_size, nums_tokens_per_rank, nums_all_rank_context, num_prefills, num_decodes,
|
||||
# num_seqs, cp_local_block_size, num_computed_tokens, num_computed_tokens_of_pcp_dcp
|
||||
test_cases = [
|
||||
(2, 2, [4], [128], 1, 0, 1, 1, [[[32, 32], [32, 32]]]),
|
||||
(1, 2, [4], [128], 1, 0, 1, 1, [[[64, 64]]]),
|
||||
(2, 1, [4], [128], 1, 0, 1, 1, [[[64], [64]]]),
|
||||
(2, 2, [4, 7], [128, 128], 2, 0, 2, 1, [[[32, 32], [32, 32]], [[32, 32], [32, 32]]]),
|
||||
]
|
||||
# kv cache tensor
|
||||
kv_cache_0 = torch.randn(
|
||||
NUM_BLOCKS, BLOCK_SIZE, self.impl.num_heads, self.impl.kv_lora_rank, dtype=torch.float16
|
||||
)
|
||||
kv_cache_1 = torch.randn(NUM_BLOCKS, BLOCK_SIZE, self.impl.num_heads, self.impl.v_head_dim, dtype=torch.float16)
|
||||
kv_cache = [kv_cache_0, kv_cache_1]
|
||||
max_model_len = 4096
|
||||
max_num_seqs = 25
|
||||
# create chunk context
|
||||
chunked_prefill_workspace_size = min(max(8 * max_model_len, 4 * max_num_seqs * BLOCK_SIZE), 128 * 1024)
|
||||
self.impl.prefill_mask = torch.triu(torch.ones(10, 10, dtype=torch.float16), 1)
|
||||
for test_case in test_cases:
|
||||
(
|
||||
pcp_size,
|
||||
dcp_size,
|
||||
nums_tokens_per_rank,
|
||||
nums_all_rank_context,
|
||||
num_prefills,
|
||||
num_decodes,
|
||||
num_seqs,
|
||||
cp_local_block_size,
|
||||
num_computed_tokens_of_pcp_dcp,
|
||||
) = test_case
|
||||
mock_dcp.world_size = dcp_size
|
||||
mock_pcp.world_size = pcp_size
|
||||
assert len(nums_tokens_per_rank) == len(nums_all_rank_context)
|
||||
nums_context_per_rank = []
|
||||
for num_all_rank_context in nums_all_rank_context:
|
||||
assert num_all_rank_context % (pcp_size * dcp_size) == 0
|
||||
nums_context_per_rank.append(num_all_rank_context // (pcp_size * dcp_size))
|
||||
self.impl.dcp_size = dcp_size
|
||||
self.impl.pcp_size = pcp_size
|
||||
# create input
|
||||
query = torch.randn(
|
||||
sum(nums_tokens_per_rank), self.impl.num_heads, self.impl.qk_head_dim, dtype=torch.float16
|
||||
)
|
||||
q_nope = query[..., : self.impl.qk_nope_head_dim]
|
||||
q_pe = query[..., self.impl.qk_nope_head_dim :]
|
||||
prefix_out = torch.randn(
|
||||
sum(nums_tokens_per_rank), self.impl.num_heads, self.impl.v_head_dim, dtype=torch.float16
|
||||
)
|
||||
prefix_lse = torch.randn(self.impl.num_heads, sum(nums_tokens_per_rank), dtype=torch.float16)
|
||||
chunk_ctx = get_chunk_metadata(
|
||||
pcp_size,
|
||||
dcp_size,
|
||||
num_prefills=num_prefills,
|
||||
num_decodes=num_decodes,
|
||||
block_size=BLOCK_SIZE,
|
||||
num_computed_tokens_cpu=torch.tensor(nums_all_rank_context),
|
||||
num_reqs=num_seqs,
|
||||
chunked_prefill_workspace_size=chunked_prefill_workspace_size,
|
||||
num_computed_tokens_of_pcp_dcp=num_computed_tokens_of_pcp_dcp,
|
||||
cp_local_block_size=cp_local_block_size,
|
||||
)
|
||||
meta = MagicMock()
|
||||
prefill_meta = MagicMock()
|
||||
prefill_meta.query_lens = torch.tensor(nums_tokens_per_rank)
|
||||
prefill_meta.block_table = torch.randint(0, USED_BLOCKS, (1, 64)) # (batch, max_blocks)
|
||||
prefill_meta.chunked_context = chunk_ctx
|
||||
meta.prefill = prefill_meta
|
||||
|
||||
with patch.object(self.impl, "_reorg_kvcache") as mock_reorg:
|
||||
mock_reorg.side_effect = mock_reorg_kvcache
|
||||
|
||||
out, lse = self.impl._compute_prefill_context(
|
||||
q_nope, q_pe, kv_cache, self.impl.qk_rope_head_dim, meta, prefix_out, prefix_lse
|
||||
)
|
||||
|
||||
iters = len(chunk_ctx.seq_tot)
|
||||
self.impl.dcp_size = 1
|
||||
self.impl.pcp_size = 1
|
||||
self.assertEqual(mock_reorg.call_count, iters * (1 if dcp_size * pcp_size > 1 else 0))
|
||||
self.assertEqual(mock_load.call_count, iters)
|
||||
self.assertEqual(mock_fia.call_count, iters)
|
||||
mock_reorg.reset_mock()
|
||||
mock_load.reset_mock()
|
||||
mock_fia.reset_mock()
|
||||
mock_update.reset_mock()
|
||||
mock_dcp.reset_mock()
|
||||
mock_pcp.reset_mock()
|
||||
self.assertEqual(out.shape, prefix_out.shape)
|
||||
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2)
|
||||
def test_reorg_kvcache_with_dcp_pcp(self, mock_all2all, mock_dcp, mock_pcp):
|
||||
BLOCK_SIZE = 128 # fixed
|
||||
max_model_len = 4096
|
||||
max_num_seqs = 25
|
||||
test_cases = [
|
||||
(2, 2, [4], [128], 1, 0, 1, 1, [[[32, 32], [32, 32]]]),
|
||||
(1, 2, [4], [128], 1, 0, 1, 1, [[[64, 64]]]),
|
||||
(2, 1, [4], [128], 1, 0, 1, 1, [[[64], [64]]]),
|
||||
(2, 2, [4, 7], [128, 128], 2, 0, 2, 1, [[[32, 32], [32, 32]], [[32, 32], [32, 32]]]),
|
||||
]
|
||||
for test_case in test_cases:
|
||||
(
|
||||
pcp_size,
|
||||
dcp_size,
|
||||
nums_tokens_per_rank,
|
||||
nums_all_rank_context,
|
||||
num_prefills,
|
||||
num_decodes,
|
||||
num_seqs,
|
||||
cp_local_block_size,
|
||||
num_computed_tokens_of_pcp_dcp,
|
||||
) = test_case
|
||||
if pcp_size * dcp_size == 1:
|
||||
continue
|
||||
self.impl.dcp_size = dcp_size
|
||||
mock_dcp.world_size = dcp_size
|
||||
mock_dcp.all_gather.reset_mock()
|
||||
self.impl.pcp_size = pcp_size
|
||||
mock_pcp.world_size = pcp_size
|
||||
mock_pcp.all_gather.reset_mock()
|
||||
|
||||
chunked_prefill_workspace_size = min(max(8 * max_model_len, 4 * max_num_seqs * BLOCK_SIZE), 128 * 1024)
|
||||
chunked_context = get_chunk_metadata(
|
||||
pcp_size,
|
||||
dcp_size,
|
||||
num_prefills=num_prefills,
|
||||
num_decodes=num_decodes,
|
||||
block_size=BLOCK_SIZE,
|
||||
num_computed_tokens_cpu=torch.tensor(nums_all_rank_context),
|
||||
num_reqs=num_seqs,
|
||||
chunked_prefill_workspace_size=chunked_prefill_workspace_size,
|
||||
num_computed_tokens_of_pcp_dcp=num_computed_tokens_of_pcp_dcp,
|
||||
cp_local_block_size=cp_local_block_size,
|
||||
)
|
||||
|
||||
for i in range(len(chunked_context.seq_tot)):
|
||||
allgatered_kv_c_normed = torch.randn(
|
||||
chunked_context.seq_tot[i], self.impl.num_heads, self.impl.kv_lora_rank
|
||||
)
|
||||
allgatered_k_pe = torch.randn(
|
||||
chunked_context.seq_tot[i], self.impl.num_heads, self.impl.qk_rope_head_dim
|
||||
)
|
||||
result_kv, result_k_pe = self.impl._reorg_kvcache(
|
||||
allgatered_kv_c_normed,
|
||||
allgatered_k_pe,
|
||||
chunked_context,
|
||||
chunk_idx=i,
|
||||
toks=chunked_context.seq_tot[i],
|
||||
)
|
||||
self.assertEqual(
|
||||
result_kv.shape,
|
||||
(chunked_context.cu_seq_lens_lst[i][-1], self.impl.num_heads, self.impl.kv_lora_rank),
|
||||
)
|
||||
self.assertEqual(
|
||||
result_k_pe.shape,
|
||||
(chunked_context.cu_seq_lens_lst[i][-1], self.impl.num_heads, self.impl.qk_rope_head_dim),
|
||||
)
|
||||
|
||||
self.assertEqual(result_kv.shape[0], chunked_context.cu_seq_lens_lst[i][-1])
|
||||
self.assertEqual(result_k_pe.shape[0], chunked_context.cu_seq_lens_lst[i][-1])
|
||||
|
||||
self.assertEqual(mock_dcp.all_gather.call_count, (1 if dcp_size > 1 else 0))
|
||||
self.assertEqual(mock_pcp.all_gather.call_count, (1 if pcp_size > 1 else 0))
|
||||
|
||||
def test_out_lse_reshape(self):
|
||||
test_cases = [10, 1, 128, 512]
|
||||
for test_case in test_cases:
|
||||
num_tokens = test_case
|
||||
num_heads, head_dim = self.impl.num_heads, self.impl.v_head_dim
|
||||
attn_out = torch.randn(num_tokens, num_heads, head_dim)
|
||||
attn_lse = torch.randn(num_tokens, num_heads, 1)
|
||||
|
||||
out, lse = self.impl._out_lse_reshape(attn_out, attn_lse)
|
||||
|
||||
assert out.shape == (num_tokens * num_heads, head_dim)
|
||||
assert out.is_contiguous()
|
||||
|
||||
assert lse.shape == (num_tokens * num_heads,)
|
||||
assert lse.is_contiguous()
|
||||
|
||||
expected_out = attn_out.contiguous().view(-1, head_dim)
|
||||
expected_lse = attn_lse.contiguous().view(-1)
|
||||
|
||||
assert torch.allclose(out, expected_out)
|
||||
assert torch.allclose(lse, expected_lse)
|
||||
|
||||
@patch("torch_npu.npu_attention_update")
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
|
||||
@patch("vllm.distributed.parallel_state._PCP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
|
||||
@patch("vllm_ascend.attention.context_parallel.common_cp.get_dcp_group")
|
||||
@patch("vllm.distributed.parallel_state._DCP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
|
||||
def test_npu_attention_update_with_dcp_pcp(
|
||||
self, mock_dcp, mock_get_dcp_group, mock_pcp, mock_get_pcp_group, mock_npu_attention_update
|
||||
):
|
||||
NUM_TOKENS = 10 # fixed
|
||||
test_cases = [(1, 1), (1, 2), (2, 1), (2, 2), (2, 3)]
|
||||
for test_case in test_cases:
|
||||
self.impl.dcp_size = test_case[0]
|
||||
self.impl.pcp_size = test_case[1]
|
||||
num_heads, head_dim = self.impl.num_heads, self.impl.kv_lora_rank + 1
|
||||
|
||||
def mock_out_lse_reshape(attn_out, attn_lse):
|
||||
attn_out = attn_out.contiguous().view(attn_out.shape[0] * attn_out.shape[1], attn_out.shape[2])
|
||||
attn_lse = attn_lse.contiguous().view(attn_lse.shape[0] * attn_lse.shape[1] * attn_lse.shape[2])
|
||||
return attn_out, attn_lse
|
||||
|
||||
self.impl._out_lse_reshape = MagicMock()
|
||||
self.impl._out_lse_reshape.side_effect = mock_out_lse_reshape
|
||||
|
||||
def mock_npu_attention_update_effect(attn_lse_split_cp, attn_out_split_cp, update_type):
|
||||
return torch.randn_like(attn_out_split_cp[0]), torch.randn_like(attn_lse_split_cp[0])
|
||||
|
||||
mock_npu_attention_update.side_effect = mock_npu_attention_update_effect
|
||||
|
||||
mock_pcp_group = MagicMock()
|
||||
mock_pcp_group.world_size = self.impl.pcp_size
|
||||
mock_get_pcp_group.return_value = mock_pcp_group
|
||||
|
||||
mock_dcp.world_size = self.impl.dcp_size
|
||||
mock_dcp_group = MagicMock()
|
||||
mock_get_dcp_group.return_value = mock_dcp_group
|
||||
attn_out_lse = torch.randn(self.impl.pcp_size * NUM_TOKENS, self.impl.dcp_size * num_heads, head_dim)
|
||||
out = _npu_attention_update(self.impl.kv_lora_rank, attn_out_lse)
|
||||
self.impl.dcp_size = 1
|
||||
self.impl.pcp_size = 1
|
||||
assert out.shape == (NUM_TOKENS, num_heads, self.impl.kv_lora_rank)
|
||||
|
||||
@patch("torch.ops.npu.npu_fused_infer_attention_score")
|
||||
def test_attention_with_optional_kv_select_with_dcp_pcp(self, mock_npu_fia):
|
||||
num_heads = self.impl.num_heads
|
||||
v_head_dim = self.impl.v_head_dim
|
||||
qk_nope_head_dim = self.impl.qk_nope_head_dim
|
||||
qk_rope_head_dim = self.impl.qk_rope_head_dim
|
||||
|
||||
def mock_npu_fia_effect(*args, **kwargs):
|
||||
q_nope = args[0]
|
||||
q_tokens = q_nope.shape[0]
|
||||
out = torch.randn(q_tokens, num_heads, v_head_dim, dtype=torch.float16)
|
||||
lse = torch.randn(q_tokens, num_heads, 1, dtype=torch.float32)
|
||||
return out, lse
|
||||
|
||||
mock_npu_fia.side_effect = mock_npu_fia_effect
|
||||
|
||||
test_cases = [([8], 2, 2), ([8, 12], 2, 2)]
|
||||
for test_case in test_cases:
|
||||
scheduled_tokens, pcp_size, dcp_size = test_case
|
||||
nums_tokens_per_rank = [num // pcp_size for num in scheduled_tokens]
|
||||
seq_len_q, seq_len_k = sum(nums_tokens_per_rank), sum(scheduled_tokens)
|
||||
|
||||
q_nope = torch.randn(seq_len_q, num_heads, qk_nope_head_dim, dtype=torch.float16)
|
||||
q_pe = torch.randn(seq_len_q, num_heads, qk_rope_head_dim, dtype=torch.float16)
|
||||
k_nope = torch.randn(seq_len_k, num_heads, qk_nope_head_dim, dtype=torch.float16)
|
||||
k_pe = torch.randn(seq_len_k, num_heads, qk_rope_head_dim, dtype=torch.float16)
|
||||
value = torch.randn(seq_len_k, num_heads, v_head_dim, dtype=torch.float16)
|
||||
mask = torch.triu(torch.ones(10, 10, dtype=torch.float16), 1)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.prefill = MagicMock()
|
||||
attn_metadata.prefill.chunked_context = None
|
||||
|
||||
for rank in range(pcp_size):
|
||||
info = get_pcp_split_info(rank, pcp_size, nums_tokens_per_rank)
|
||||
q_head_idx = info[0]
|
||||
chunk_seqlens = info[6]
|
||||
kv_tail_proj_idx = info[9]
|
||||
kv_with_q_head_attn_idx_in_tail = info[10]
|
||||
head_actual_seq_lengths_kv = info[12]
|
||||
attn_mask_seqlens = torch.cumsum(torch.tensor(chunk_seqlens, dtype=torch.int32), dim=0).tolist()
|
||||
|
||||
output_head, lse_head = self.impl._attention_with_optional_kv_select(
|
||||
q_nope=torch.index_select(q_nope, 0, q_head_idx),
|
||||
q_pe=torch.index_select(q_pe, 0, q_head_idx),
|
||||
k_nope=torch.index_select(k_nope, 0, kv_tail_proj_idx),
|
||||
k_pe=torch.index_select(k_pe, 0, kv_tail_proj_idx),
|
||||
value=torch.index_select(value, 0, kv_tail_proj_idx),
|
||||
kv_attn_idx=kv_with_q_head_attn_idx_in_tail,
|
||||
attn_mask_seqlens=attn_mask_seqlens,
|
||||
actual_seq_lengths_kv=head_actual_seq_lengths_kv,
|
||||
mask=mask,
|
||||
attn_metadata=attn_metadata,
|
||||
)
|
||||
|
||||
self.assertEqual(output_head.shape, (q_head_idx.shape[0], num_heads, v_head_dim))
|
||||
if lse_head is not None:
|
||||
self.assertEqual(lse_head.shape, (q_head_idx.shape[0], num_heads, 1))
|
||||
else:
|
||||
self.assertIsNone(lse_head)
|
||||
self.assertEqual(mock_npu_fia.call_count, 1)
|
||||
self.assertEqual(mock_npu_fia.call_args.kwargs["sparse_mode"], 3)
|
||||
|
||||
mock_npu_fia.reset_mock()
|
||||
|
||||
@patch("torch.ops.npu.npu_fused_infer_attention_score")
|
||||
def test_attention_with_optional_kv_select_trigger_chunked(self, mock_npu_fia):
|
||||
num_heads = self.impl.num_heads
|
||||
v_head_dim = self.impl.v_head_dim
|
||||
qk_nope_head_dim = self.impl.qk_nope_head_dim
|
||||
qk_rope_head_dim = self.impl.qk_rope_head_dim
|
||||
|
||||
def mock_npu_fia_effect(*args, **kwargs):
|
||||
q_tokens = args[0].shape[0]
|
||||
out = torch.randn(q_tokens, num_heads, v_head_dim, dtype=torch.float16)
|
||||
lse = torch.randn(q_tokens, num_heads, 1, dtype=torch.float32)
|
||||
return out, lse
|
||||
|
||||
mock_npu_fia.side_effect = mock_npu_fia_effect
|
||||
|
||||
test_cases = [([8], 2, 2)]
|
||||
for test_case in test_cases:
|
||||
scheduled_tokens, pcp_size, dcp_size = test_case
|
||||
nums_tokens_per_rank = [num // pcp_size for num in scheduled_tokens]
|
||||
|
||||
q_nope = torch.randn(sum(nums_tokens_per_rank), num_heads, qk_nope_head_dim, dtype=torch.float16)
|
||||
q_pe = torch.randn(sum(nums_tokens_per_rank), num_heads, qk_rope_head_dim, dtype=torch.float16)
|
||||
k_nope = torch.randn(sum(scheduled_tokens), num_heads, qk_nope_head_dim, dtype=torch.float16)
|
||||
k_pe = torch.randn(sum(scheduled_tokens), num_heads, qk_rope_head_dim, dtype=torch.float16)
|
||||
value = torch.randn(sum(scheduled_tokens), num_heads, v_head_dim, dtype=torch.float16)
|
||||
mask = torch.ones(10, 10, dtype=torch.float16)
|
||||
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.prefill = MagicMock()
|
||||
attn_metadata.prefill.chunked_context = "TRIGGER_CHUNKED"
|
||||
|
||||
for rank in range(pcp_size):
|
||||
info = get_pcp_split_info(rank, pcp_size, nums_tokens_per_rank)
|
||||
q_head_idx = info[0]
|
||||
chunk_seqlens = info[6]
|
||||
kv_tail_proj_idx = info[9]
|
||||
kv_with_q_head_attn_idx_in_tail = info[10]
|
||||
head_actual_seq_lengths_kv = info[12]
|
||||
attn_mask_seqlens = torch.cumsum(torch.tensor(chunk_seqlens, dtype=torch.int32), dim=0).tolist()
|
||||
|
||||
output_head, lse_head = self.impl._attention_with_optional_kv_select(
|
||||
q_nope=torch.index_select(q_nope, 0, q_head_idx),
|
||||
q_pe=torch.index_select(q_pe, 0, q_head_idx),
|
||||
k_nope=torch.index_select(k_nope, 0, kv_tail_proj_idx),
|
||||
k_pe=torch.index_select(k_pe, 0, kv_tail_proj_idx),
|
||||
value=torch.index_select(value, 0, kv_tail_proj_idx),
|
||||
kv_attn_idx=kv_with_q_head_attn_idx_in_tail,
|
||||
attn_mask_seqlens=attn_mask_seqlens,
|
||||
actual_seq_lengths_kv=head_actual_seq_lengths_kv,
|
||||
mask=mask,
|
||||
attn_metadata=attn_metadata,
|
||||
)
|
||||
|
||||
self.assertEqual(mock_npu_fia.call_count, 1)
|
||||
self.assertEqual(mock_npu_fia.call_args.kwargs["sparse_mode"], 3)
|
||||
self.assertIsNotNone(lse_head)
|
||||
self.assertEqual(output_head.shape, (q_head_idx.shape[0], num_heads, v_head_dim))
|
||||
|
||||
mock_npu_fia.reset_mock()
|
||||
1023
tests/ut/attention/a2/test_mla_cp_precision.py
Normal file
1023
tests/ut/attention/a2/test_mla_cp_precision.py
Normal file
File diff suppressed because it is too large
Load Diff
681
tests/ut/attention/a2/test_mla_precision.py
Normal file
681
tests/ut/attention/a2/test_mla_precision.py
Normal file
@@ -0,0 +1,681 @@
|
||||
import math
|
||||
from functools import partial
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backend import AttentionType
|
||||
from vllm.v1.attention.selector import get_attn_backend
|
||||
from vllm.v1.kv_cache_interface import MLAAttentionSpec
|
||||
|
||||
from tests.ut.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def default_vllm_config():
|
||||
mock_config = MagicMock()
|
||||
mock_config.compilation_config = MagicMock()
|
||||
mock_config.compilation_config.custom_ops = ["all"]
|
||||
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
|
||||
|
||||
|
||||
BATCH_SPECS = {
|
||||
"small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
|
||||
"small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
|
||||
"mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
|
||||
"medium_decode": BatchSpec(
|
||||
seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
|
||||
query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
|
||||
),
|
||||
"medium_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]),
|
||||
"mixed_medium": BatchSpec(seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]),
|
||||
"large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
|
||||
"large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
|
||||
"mixed_large": BatchSpec(seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]),
|
||||
"single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
|
||||
"single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
|
||||
# encoder-only
|
||||
"small_encoder_prefill": BatchSpec(seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]),
|
||||
"medium_encoder_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]),
|
||||
"mtp_1_plus_3": BatchSpec(seq_lens=[256, 512, 1024, 1536], query_lens=[4, 4, 4, 4]),
|
||||
}
|
||||
|
||||
|
||||
class MockLinear:
|
||||
def __init__(self, out_features=128, in_features=128, device=None, dtype=torch.bfloat16):
|
||||
self.weight = torch.randn(out_features, in_features, dtype=dtype, device=device) / math.sqrt(in_features)
|
||||
self.quant_method = UnquantizedLinearMethod()
|
||||
|
||||
def __call__(self, x, **kwargs):
|
||||
if x.size(-1) != self.weight.size(-1):
|
||||
self.weight = torch.randn(
|
||||
self.weight.size(0), x.size(-1), dtype=self.weight.dtype, device=x.device
|
||||
) / math.sqrt(x.size(-1))
|
||||
|
||||
return (x @ self.weight.T, None)
|
||||
|
||||
|
||||
class MockLayerNorm:
|
||||
def __init__(self, normalized_shape, device=None, dtype=torch.bfloat16):
|
||||
self.weight = torch.ones(normalized_shape, dtype=dtype, device=device)
|
||||
self.variance_epsilon = 1e-6
|
||||
|
||||
def __call__(self, x):
|
||||
return x
|
||||
|
||||
|
||||
class MockRotary:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
def create_mla_kv_cache(
|
||||
k_nope_contexts: list[torch.Tensor],
|
||||
k_pe_contexts: list[torch.Tensor],
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
kv_lora_rank: int,
|
||||
qk_rope_head_dim: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
num_blocks: int,
|
||||
common_attn_metadata: AscendCommonAttentionMetadata,
|
||||
):
|
||||
"""Create MLA KV cache with two separate tensors.
|
||||
|
||||
MLA KV cache layout:
|
||||
- k_cache (nope): (num_blocks, block_size, num_kv_heads, kv_lora_rank)
|
||||
- v_cache (rope): (num_blocks, block_size, num_kv_heads, qk_rope_head_dim)
|
||||
"""
|
||||
k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, kv_lora_rank, dtype=dtype, device=device)
|
||||
v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
seq_lens = common_attn_metadata.seq_lens.cpu()
|
||||
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
|
||||
query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
|
||||
batch_size = len(k_nope_contexts)
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
|
||||
start_block_idx = 0
|
||||
for i in range(batch_size):
|
||||
k_nope_ctx = k_nope_contexts[i]
|
||||
k_pe_ctx = k_pe_contexts[i]
|
||||
context_len = int(seq_lens[i]) - int(query_lens[i])
|
||||
|
||||
num_blocks_for_seq = (int(seq_lens[i]) + block_size - 1) // block_size
|
||||
block_table[i, :num_blocks_for_seq] = torch.arange(
|
||||
start_block_idx, start_block_idx + num_blocks_for_seq, dtype=torch.int32
|
||||
)
|
||||
|
||||
k_cache_flat = k_cache[start_block_idx : start_block_idx + num_blocks_for_seq].view(
|
||||
-1, num_kv_heads, kv_lora_rank
|
||||
)
|
||||
v_cache_flat = v_cache[start_block_idx : start_block_idx + num_blocks_for_seq].view(
|
||||
-1, num_kv_heads, qk_rope_head_dim
|
||||
)
|
||||
k_cache_flat[:context_len] = k_nope_ctx[:context_len]
|
||||
v_cache_flat[:context_len] = k_pe_ctx[:context_len]
|
||||
|
||||
start_block_idx += num_blocks_for_seq
|
||||
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
for i in range(batch_size):
|
||||
context_len_i = int(seq_lens[i]) - int(query_lens[i])
|
||||
token_offsets = torch.arange(int(query_lens[i])) + context_len_i
|
||||
block_indices = token_offsets // block_size
|
||||
token_inter_block_offsets = token_offsets % block_size
|
||||
start = int(query_start_loc_cpu[i])
|
||||
end = int(query_start_loc_cpu[i + 1])
|
||||
slot_mapping[start:end] = block_table[i, block_indices] * block_size + token_inter_block_offsets.to(device).to(
|
||||
torch.int32
|
||||
)
|
||||
|
||||
return (k_cache, v_cache)
|
||||
|
||||
|
||||
def run_mla_attention_backend(
|
||||
kv_cache_spec: MLAAttentionSpec,
|
||||
vllm_config,
|
||||
device: torch.device,
|
||||
common_attn_metadata: AscendCommonAttentionMetadata,
|
||||
hidden_states: torch.Tensor,
|
||||
kv_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
dtype: torch.bfloat16,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
):
|
||||
from vllm_ascend.ascend_config import init_ascend_config
|
||||
|
||||
init_ascend_config(vllm_config)
|
||||
|
||||
from vllm_ascend.ops import rotary_embedding
|
||||
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
|
||||
|
||||
rotary_embedding._cos_cache = torch.ones(8192, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
rotary_embedding._sin_cache = torch.zeros(8192, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
rotary_embedding._cos_mla = torch.ones(8192, 1, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
rotary_embedding._sin_mla = torch.zeros(8192, 1, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
from vllm.distributed.parallel_state import GroupCoordinator
|
||||
|
||||
mock_tp_group = MagicMock(spec=GroupCoordinator)
|
||||
mock_tp_group.world_size = 1
|
||||
mock_tp_group.rank = 0
|
||||
|
||||
mock_weight_prefetch = MagicMock()
|
||||
mock_weight_prefetch.maybe_prefetch_mla_or_sla_weight_in_current_stream = MagicMock()
|
||||
|
||||
import vllm_ascend.utils as utils_module
|
||||
|
||||
original_weight_prefetch = utils_module._WEIGHT_PREFETCH_METHOD
|
||||
utils_module._WEIGHT_PREFETCH_METHOD = mock_weight_prefetch
|
||||
try:
|
||||
with patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group):
|
||||
num_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()
|
||||
|
||||
kv_lora_rank = getattr(hf_config, "kv_lora_rank", 512)
|
||||
q_lora_rank = getattr(hf_config, "q_lora_rank", 1536)
|
||||
qk_nope_head_dim = getattr(hf_config, "qk_nope_head_dim", 128)
|
||||
qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
|
||||
qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||
v_head_dim = getattr(hf_config, "v_head_dim", 128)
|
||||
|
||||
backend = get_attn_backend(head_size, dtype, None, use_mla=True, use_sparse=False, use_mm_prefix=False)
|
||||
impl_cls = backend.get_impl_cls()
|
||||
builder_cls = backend.get_builder_cls()
|
||||
|
||||
mock_layer_entry = MagicMock()
|
||||
for layer_name in ["placeholder"]:
|
||||
vllm_config.compilation_config.static_forward_context[layer_name] = mock_layer_entry
|
||||
|
||||
builder = builder_cls(
|
||||
kv_cache_spec,
|
||||
["placeholder"],
|
||||
vllm_config,
|
||||
device,
|
||||
)
|
||||
attn_metadata = builder.build(
|
||||
common_prefix_len=0,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
hidden_size = num_heads * head_size
|
||||
q_proj_out_dim = num_heads * qk_head_dim
|
||||
q_b_proj_out_dim = num_heads * qk_head_dim
|
||||
kv_b_proj_out_dim = num_heads * (qk_nope_head_dim + v_head_dim)
|
||||
kv_b_proj_in_dim = kv_lora_rank
|
||||
o_proj_out_dim = head_size * num_heads
|
||||
o_proj_in_dim = num_heads * v_head_dim
|
||||
|
||||
impl = impl_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=1.0 / (head_size**0.5),
|
||||
num_kv_heads=num_kv_heads,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
attn_type=attn_type.value if hasattr(attn_type, "value") else attn_type,
|
||||
kv_cache_dtype="auto",
|
||||
logits_soft_cap=None,
|
||||
kv_sharing_target_layer_name=None,
|
||||
q_lora_rank=q_lora_rank,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
qk_head_dim=qk_head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
q_proj=MockLinear(out_features=q_proj_out_dim, in_features=hidden_size, device=device, dtype=dtype),
|
||||
q_b_proj=MockLinear(out_features=q_b_proj_out_dim, in_features=q_lora_rank, device=device, dtype=dtype),
|
||||
kv_b_proj=MockLinear(
|
||||
out_features=kv_b_proj_out_dim, in_features=kv_b_proj_in_dim, device=device, dtype=dtype
|
||||
),
|
||||
o_proj=MockLinear(out_features=o_proj_out_dim, in_features=o_proj_in_dim, device=device, dtype=dtype),
|
||||
kv_a_layernorm=MockLayerNorm(normalized_shape=kv_lora_rank, device=device, dtype=dtype),
|
||||
q_a_layernorm=MockLayerNorm(normalized_shape=q_lora_rank, device=device, dtype=dtype),
|
||||
rotary_emb=MockRotary(),
|
||||
fused_qkv_a_proj=None,
|
||||
kv_a_proj_with_mqa=MockLinear(
|
||||
out_features=num_kv_heads * (kv_lora_rank + qk_rope_head_dim),
|
||||
in_features=hidden_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
),
|
||||
)
|
||||
|
||||
impl.fa_quant_layer = False
|
||||
impl.enable_mlapo = False
|
||||
impl.process_weights_after_loading(dtype)
|
||||
|
||||
output = torch.empty_like(hidden_states)
|
||||
output = impl.forward("layer_0", hidden_states, kv_cache, attn_metadata, output=output)
|
||||
finally:
|
||||
utils_module._WEIGHT_PREFETCH_METHOD = original_weight_prefetch
|
||||
|
||||
return output, impl
|
||||
|
||||
|
||||
def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
|
||||
variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
|
||||
x_normed = x * torch.rsqrt(variance + eps)
|
||||
return (x_normed * weight).to(x.dtype)
|
||||
|
||||
|
||||
def npu_interleave_rope_simple(
|
||||
x: torch.Tensor, cos: torch.Tensor | None = None, sin: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
"""Simulate Ascend npu_interleave_rope with default cos=1, sin=0.
|
||||
|
||||
Interleave the last dimension: [x0, x1, x2, x3, ...] -> [x0, x2, ..., x1, x3, ...]
|
||||
With cos=1, sin=0 (rope disabled), the function just returns the interleaved result.
|
||||
"""
|
||||
even = x[..., 0::2]
|
||||
odd = x[..., 1::2]
|
||||
return torch.cat([even, odd], dim=-1).contiguous()
|
||||
|
||||
|
||||
def prefill_sdpa(
|
||||
q_nope: torch.Tensor,
|
||||
q_pe_raw: torch.Tensor,
|
||||
k_pe_c: torch.Tensor,
|
||||
k_nope_c: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
impl,
|
||||
scale: float,
|
||||
is_causal: bool = True,
|
||||
context_len: int = 0,
|
||||
) -> torch.Tensor:
|
||||
q_pe = npu_interleave_rope_simple(q_pe_raw)
|
||||
q_full = torch.cat([q_nope, q_pe], dim=-1)
|
||||
k_full = torch.cat([k_nope_c, k_pe_c], dim=-1)
|
||||
|
||||
q_sdpa = q_full.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa = v.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
if context_len > 0 and is_causal:
|
||||
q_len = q_full.shape[0]
|
||||
kv_len = k_full.shape[0]
|
||||
mask = torch.tril(
|
||||
torch.ones(q_len, kv_len, device=q_full.device, dtype=torch.bool),
|
||||
diagonal=context_len,
|
||||
)
|
||||
attn_out = torch.nn.functional.scaled_dot_product_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
attn_mask=mask,
|
||||
enable_gqa=False,
|
||||
scale=scale,
|
||||
)
|
||||
else:
|
||||
attn_out = torch.nn.functional.scaled_dot_product_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
is_causal=is_causal,
|
||||
enable_gqa=False,
|
||||
scale=scale,
|
||||
)
|
||||
return attn_out.transpose(1, 2).squeeze(0)
|
||||
|
||||
|
||||
def decode_sdpa(
|
||||
ql_nope: torch.Tensor,
|
||||
q_pe: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
k_nope: torch.Tensor,
|
||||
impl,
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Compute SDPA for decode path.
|
||||
|
||||
Decode path: q is W_UK projected (ql_nope), k/v are raw latent.
|
||||
"""
|
||||
q_full = torch.cat([ql_nope, q_pe], dim=-1)
|
||||
k_full = torch.cat([k_nope, k_pe], dim=-1)
|
||||
|
||||
q_sdpa = q_full.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa = k_nope.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
attn_out = torch.nn.functional.scaled_dot_product_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
is_causal=False,
|
||||
enable_gqa=(impl.num_heads != impl.num_kv_heads),
|
||||
scale=scale,
|
||||
)
|
||||
return attn_out.transpose(1, 2).squeeze(0)
|
||||
|
||||
|
||||
def compute_mla_sdpa_reference(
|
||||
hidden_states: torch.Tensor,
|
||||
k_nope_contexts: list[torch.Tensor],
|
||||
k_pe_contexts: list[torch.Tensor],
|
||||
impl,
|
||||
batch_spec: BatchSpec,
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Compute MLA reference using SDPA as golden baseline.
|
||||
|
||||
Handles three modes:
|
||||
- Decode only: q via W_UK, k/v raw latent, _v_up_proj
|
||||
- Prefill only: q raw, k/v via kv_b_proj, no _v_up_proj
|
||||
- Mixed: decode first, then prefill, combined before o_proj
|
||||
"""
|
||||
rms_eps = impl.kv_a_layernorm.variance_epsilon
|
||||
rms_w = impl.kv_a_layernorm.weight
|
||||
|
||||
# --- Identify decode vs prefill tokens ---
|
||||
num_decode_tokens = 0
|
||||
for ql in batch_spec.query_lens:
|
||||
if ql == 1:
|
||||
num_decode_tokens += 1
|
||||
|
||||
# --- KV projection and normalization for all tokens ---
|
||||
kv_no_split = impl.kv_a_proj_with_mqa(hidden_states)[0]
|
||||
k_nope_all = kv_no_split[:, : impl.kv_lora_rank]
|
||||
k_pe_all = kv_no_split[:, impl.kv_lora_rank :]
|
||||
k_nope_normed = rms_norm(k_nope_all, rms_w, rms_eps)
|
||||
|
||||
# --- Q projection for all tokens ---
|
||||
q_b = impl.q_proj(hidden_states)[0]
|
||||
q_nope_all, q_pe_all = q_b.view(-1, impl.num_heads, impl.qk_head_dim).split(
|
||||
[impl.qk_nope_head_dim, impl.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
|
||||
# Decode: q_nope is W_UK projected; Prefill: q_nope stays raw
|
||||
q_nope_decode = q_nope_all[:num_decode_tokens]
|
||||
q_nope_decode_t = q_nope_decode.transpose(0, 1).float()
|
||||
ql_nope_decode = torch.bmm(q_nope_decode_t, impl.W_UK_T.float()).to(q_nope_all.dtype)
|
||||
ql_nope_decode = ql_nope_decode.transpose(0, 1)
|
||||
|
||||
q_pe_decode = npu_interleave_rope_simple(q_pe_all[:num_decode_tokens])
|
||||
|
||||
q_nope_prefill = q_nope_all[num_decode_tokens:]
|
||||
q_pe_prefill_raw = q_pe_all[num_decode_tokens:]
|
||||
|
||||
# --- Process each sequence ---
|
||||
decode_outputs = []
|
||||
prefill_outputs = []
|
||||
|
||||
token_offset = 0
|
||||
for i in range(len(batch_spec.seq_lens)):
|
||||
s_len_i = batch_spec.seq_lens[i]
|
||||
q_len_i = batch_spec.query_lens[i]
|
||||
context_len_i = s_len_i - q_len_i
|
||||
is_decode = q_len_i == 1
|
||||
|
||||
if is_decode:
|
||||
ql_nope_i = ql_nope_decode[token_offset : token_offset + q_len_i]
|
||||
q_pe_i = q_pe_decode[token_offset : token_offset + q_len_i]
|
||||
|
||||
k_nope_dec = k_nope_normed[token_offset : token_offset + q_len_i].view(
|
||||
q_len_i, impl.num_kv_heads, impl.kv_lora_rank
|
||||
)
|
||||
k_pe_dec = k_pe_all[token_offset : token_offset + q_len_i].view(
|
||||
q_len_i, impl.num_kv_heads, impl.qk_rope_head_dim
|
||||
)
|
||||
k_pe_dec = npu_interleave_rope_simple(k_pe_dec)
|
||||
|
||||
k_nope_full_i = torch.cat([k_nope_contexts[i], k_nope_dec], dim=0)
|
||||
k_pe_full_i = torch.cat([k_pe_contexts[i], k_pe_dec], dim=0)
|
||||
|
||||
sdpa_out = decode_sdpa(
|
||||
ql_nope_i,
|
||||
q_pe_i,
|
||||
k_pe_full_i,
|
||||
k_nope_full_i,
|
||||
impl,
|
||||
scale,
|
||||
)
|
||||
# _v_up_proj
|
||||
sdpa_out = sdpa_out.transpose(0, 1).contiguous()
|
||||
sdpa_out = torch.bmm(sdpa_out.float(), impl.W_UV.float()).to(sdpa_out.dtype)
|
||||
sdpa_out = sdpa_out.permute(1, 0, 2)
|
||||
sdpa_out = sdpa_out.reshape(-1, impl.num_heads * impl.v_head_dim)
|
||||
decode_outputs.append(sdpa_out)
|
||||
|
||||
else:
|
||||
q_nope_i = q_nope_prefill[token_offset - num_decode_tokens : token_offset - num_decode_tokens + q_len_i]
|
||||
q_pe_raw_i = q_pe_prefill_raw[token_offset - num_decode_tokens : token_offset - num_decode_tokens + q_len_i]
|
||||
|
||||
k_nope_new = k_nope_normed[token_offset : token_offset + q_len_i].view(
|
||||
q_len_i, impl.num_kv_heads, impl.kv_lora_rank
|
||||
)
|
||||
k_pe_raw_new = k_pe_all[token_offset : token_offset + q_len_i].view(
|
||||
q_len_i, impl.num_kv_heads, impl.qk_rope_head_dim
|
||||
)
|
||||
|
||||
# kv_b_proj on RMS-normed new tokens: k_nope_proj + v_proj
|
||||
k_nope_new_flat = k_nope_new.view(-1, impl.kv_lora_rank)
|
||||
kv_b_new = impl.kv_b_proj(k_nope_new_flat)[0].view(
|
||||
q_len_i, impl.num_heads, impl.qk_nope_head_dim + impl.v_head_dim
|
||||
)
|
||||
k_nope_proj_new, v_proj_new = kv_b_new.split([impl.qk_nope_head_dim, impl.v_head_dim], dim=-1)
|
||||
|
||||
# kv_b_proj on raw context (cache stores raw data)
|
||||
k_nope_ctx_raw = (
|
||||
k_nope_contexts[i]
|
||||
.view(context_len_i, impl.num_kv_heads, impl.kv_lora_rank)
|
||||
.reshape(-1, impl.kv_lora_rank)
|
||||
)
|
||||
kv_b_ctx = impl.kv_b_proj(k_nope_ctx_raw)[0].view(
|
||||
context_len_i, impl.num_heads, impl.qk_nope_head_dim + impl.v_head_dim
|
||||
)
|
||||
k_nope_proj_ctx, v_proj_ctx = kv_b_ctx.split([impl.qk_nope_head_dim, impl.v_head_dim], dim=-1)
|
||||
|
||||
k_nope_proj = torch.cat([k_nope_proj_ctx, k_nope_proj_new], dim=0)
|
||||
v_proj = torch.cat([v_proj_ctx, v_proj_new], dim=0)
|
||||
|
||||
# k_pe: context raw (as stored in cache), new interleaved
|
||||
k_pe_raw_ctx = k_pe_contexts[i].view(context_len_i, impl.num_kv_heads, impl.qk_rope_head_dim)
|
||||
k_pe_new_interleaved = npu_interleave_rope_simple(k_pe_raw_new)
|
||||
k_pe_cat = torch.cat([k_pe_raw_ctx, k_pe_new_interleaved], dim=0)
|
||||
k_pe_expanded = k_pe_cat.expand(*k_nope_proj.shape[:-1], -1)
|
||||
|
||||
sdpa_out = prefill_sdpa(
|
||||
q_nope_i,
|
||||
q_pe_raw_i,
|
||||
k_pe_expanded,
|
||||
k_nope_proj,
|
||||
v_proj,
|
||||
impl,
|
||||
scale,
|
||||
is_causal=True,
|
||||
context_len=context_len_i,
|
||||
)
|
||||
sdpa_out = sdpa_out.reshape(q_len_i, impl.num_heads * impl.v_head_dim)
|
||||
prefill_outputs.append(sdpa_out)
|
||||
|
||||
token_offset += q_len_i
|
||||
|
||||
# --- Combine outputs in token order ---
|
||||
all_outputs = decode_outputs + prefill_outputs
|
||||
attn_output = torch.cat(all_outputs, dim=0)
|
||||
|
||||
# --- Output projection ---
|
||||
sdpa_output = impl.o_proj(attn_output)[0]
|
||||
return sdpa_output
|
||||
|
||||
|
||||
def _test_mla_attention_correctness(
|
||||
batch_spec: BatchSpec,
|
||||
model: str,
|
||||
*,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
block_size: int = 128,
|
||||
atol: float = 1e-2,
|
||||
rtol: float = 1e-2,
|
||||
tensor_parallel_size: 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,
|
||||
num_gpu_blocks=8192,
|
||||
)
|
||||
device = torch.device("npu")
|
||||
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
kv_lora_rank = getattr(hf_config, "kv_lora_rank", 512)
|
||||
qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
|
||||
|
||||
batch_size = batch_spec.batch_size
|
||||
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)
|
||||
|
||||
kv_cache_spec = MLAAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_size=kv_lora_rank,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
k_nope_contexts, k_pe_contexts = [], []
|
||||
for i in range(batch_size):
|
||||
s_len = batch_spec.seq_lens[i]
|
||||
q_len = batch_spec.query_lens[i]
|
||||
context_len = s_len - q_len
|
||||
|
||||
k_nope_full = torch.randn(s_len, num_kv_heads, kv_lora_rank, dtype=dtype, device=device)
|
||||
k_pe_full = torch.randn(s_len, num_kv_heads, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
k_nope_contexts.append(k_nope_full[:context_len])
|
||||
k_pe_contexts.append(k_pe_full[:context_len])
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(batch_spec, vllm_config.cache_config.block_size, device)
|
||||
if attn_type == AttentionType.ENCODER_ONLY:
|
||||
common_attn_metadata.causal = False
|
||||
|
||||
num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
|
||||
num_blocks = max(16, num_blocks)
|
||||
|
||||
kv_cache = create_mla_kv_cache(
|
||||
k_nope_contexts=k_nope_contexts,
|
||||
k_pe_contexts=k_pe_contexts,
|
||||
block_size=block_size,
|
||||
num_kv_heads=num_kv_heads,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
num_blocks=num_blocks,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
num_tokens = common_attn_metadata.num_actual_tokens
|
||||
hidden_size = num_q_heads * head_size
|
||||
hidden_states = torch.randn(num_tokens, hidden_size, dtype=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
|
||||
|
||||
backend_output, impl = run_mla_attention_backend(
|
||||
kv_cache_spec,
|
||||
vllm_config,
|
||||
device,
|
||||
common_attn_metadata,
|
||||
hidden_states,
|
||||
kv_cache,
|
||||
dtype,
|
||||
attn_type=attn_type,
|
||||
)
|
||||
|
||||
sdpa_output = compute_mla_sdpa_reference(
|
||||
hidden_states,
|
||||
k_nope_contexts,
|
||||
k_pe_contexts,
|
||||
impl,
|
||||
batch_spec,
|
||||
scale,
|
||||
)
|
||||
|
||||
# Compare (same pattern as test_gqa.py)
|
||||
name = "MLA"
|
||||
assert backend_output.shape == sdpa_output.shape, (
|
||||
f"[{name}] shape {backend_output.shape} != SDPA shape {sdpa_output.shape}"
|
||||
)
|
||||
assert backend_output.dtype == sdpa_output.dtype, (
|
||||
f"[{name}] dtype {backend_output.dtype} != SDPA dtype {sdpa_output.dtype}"
|
||||
)
|
||||
|
||||
assert torch.isfinite(backend_output).all(), f"[{name}] produced non-finite values"
|
||||
|
||||
def error_msg(msg: str, backend_name: str):
|
||||
return f"[{backend_name}] output differs from SDPA baseline. {msg}"
|
||||
|
||||
torch.testing.assert_close(
|
||||
backend_output,
|
||||
sdpa_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=partial(error_msg, backend_name="MLA"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
[
|
||||
"small_decode",
|
||||
"small_prefill",
|
||||
"mixed_small",
|
||||
"medium_decode",
|
||||
"medium_prefill",
|
||||
"mixed_medium",
|
||||
"large_decode",
|
||||
"large_prefill",
|
||||
"single_decode",
|
||||
"single_prefill",
|
||||
"mtp_1_plus_3",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V2"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1])
|
||||
def test_mla_backend_correctness(default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int):
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
|
||||
_test_mla_attention_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
2415
tests/ut/attention/a2/test_mla_v1.py
Executable file
2415
tests/ut/attention/a2/test_mla_v1.py
Executable file
File diff suppressed because it is too large
Load Diff
844
tests/ut/attention/a2/test_sfa_cp_precision.py
Normal file
844
tests/ut/attention/a2/test_sfa_cp_precision.py
Normal file
@@ -0,0 +1,844 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
import math
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
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 vllm.forward_context import set_forward_context # noqa: E402
|
||||
|
||||
from tests.ut.attention.utils import ( # noqa: E402
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm_ascend.attention.context_parallel.sfa_cp import AscendSFACPImpl # noqa: E402
|
||||
from vllm_ascend.utils import enable_custom_op
|
||||
|
||||
enable_custom_op()
|
||||
|
||||
SPARSE_COUNT = 2048
|
||||
_BLOCK_SIZE = 128
|
||||
_TEST_NUM_HEADS = 8
|
||||
DEFAULT_RTOL = 1e-2
|
||||
DEFAULT_ATOL = 1e-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 ``seq_len <= SPARSE_COUNT`` so sparse matches dense reference."""
|
||||
for s, q in zip(spec.seq_lens, spec.query_lens):
|
||||
assert q <= s, f"query_len ({q}) must not exceed seq_len ({s})"
|
||||
assert s <= SPARSE_COUNT, (
|
||||
f"seq_len ({s}) must be <= SPARSE_COUNT ({SPARSE_COUNT}) so the "
|
||||
"sparse attention degenerates into dense attention for the "
|
||||
"reference comparison."
|
||||
)
|
||||
|
||||
|
||||
_VLLM_CONFIG_CACHE: dict = {}
|
||||
|
||||
|
||||
def _get_vllm_config(
|
||||
model: str,
|
||||
dtype: torch.dtype,
|
||||
*,
|
||||
max_model_len: int = 4096,
|
||||
tensor_parallel_size: int = 1,
|
||||
) -> VllmConfig:
|
||||
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=64,
|
||||
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,
|
||||
},
|
||||
)
|
||||
_VLLM_CONFIG_CACHE[key] = cfg
|
||||
return cfg
|
||||
|
||||
|
||||
# Spec name prefixes drive the SFA-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=[512, 1024, 1536, 2048], query_lens=[1, 1, 1, 1], name="decode_small_batch"
|
||||
),
|
||||
"decode_large_batch": BatchSpec(seq_lens=[2048] * 8, query_lens=[1] * 8, name="decode_large_batch"),
|
||||
"mtp_1_plus_1": BatchSpec(seq_lens=[512, 1024, 1536], query_lens=[2, 2, 2], name="mtp_1_plus_1"),
|
||||
"mtp_1_plus_3": BatchSpec(seq_lens=[1024, 1536, 2048, 2048], query_lens=[4, 4, 4, 4], name="mtp_1_plus_3"),
|
||||
"mtp_1_plus_7": BatchSpec(seq_lens=[1024, 1536, 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=[256, 512, 384], query_lens=[256, 512, 384], name="prefill_small_batch"),
|
||||
"prefill_with_context": BatchSpec(seq_lens=[512, 1024], query_lens=[128, 256], name="prefill_with_context"),
|
||||
"mixed_small": BatchSpec(seq_lens=[512, 1024, 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 _build_topk_indices(
|
||||
seq_lens: list[int],
|
||||
query_lens: list[int],
|
||||
sparse_count: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Build causal topk indices with INVALID_IDX (-1) padding."""
|
||||
num_tokens = sum(query_lens)
|
||||
topk = torch.full((num_tokens, 1, sparse_count), -1, dtype=torch.int32, device=device)
|
||||
|
||||
cum_q = 0
|
||||
for b, s_len in enumerate(seq_lens):
|
||||
q_len = query_lens[b]
|
||||
ctx_len = s_len - q_len
|
||||
for j in range(q_len):
|
||||
valid_end = ctx_len + j + 1
|
||||
topk[cum_q + j, 0, :valid_end] = torch.arange(valid_end, dtype=torch.int32, device=device)
|
||||
cum_q += q_len
|
||||
|
||||
return topk
|
||||
|
||||
|
||||
def _reference_sparse_attention(
|
||||
ql_nope: torch.Tensor,
|
||||
q_pe: torch.Tensor,
|
||||
k_nope_full_per_req: list[torch.Tensor],
|
||||
k_rope_full_per_req: list[torch.Tensor],
|
||||
seq_lens: list[int],
|
||||
query_lens: list[int],
|
||||
scale: float,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Pure-PyTorch dense MQA baseline in fp32."""
|
||||
outputs: list[torch.Tensor] = []
|
||||
|
||||
cum_q = 0
|
||||
for b, s_len in enumerate(seq_lens):
|
||||
q_len = query_lens[b]
|
||||
ctx_len = s_len - q_len
|
||||
|
||||
K_nope = k_nope_full_per_req[b][:s_len].float()
|
||||
K_rope = k_rope_full_per_req[b][:s_len].float()
|
||||
K = torch.cat([K_nope, K_rope], dim=-1)
|
||||
V = K_nope
|
||||
|
||||
for j in range(q_len):
|
||||
t = cum_q + j
|
||||
valid_end = ctx_len + j + 1
|
||||
|
||||
q_n = ql_nope[t].float()
|
||||
q_p = q_pe[t].float()
|
||||
Q = torch.cat([q_n, q_p], dim=-1)
|
||||
|
||||
K_b = K[:valid_end]
|
||||
V_b = V[:valid_end]
|
||||
|
||||
scores = (Q @ K_b.transpose(0, 1)) * scale
|
||||
attn = torch.softmax(scores, dim=-1)
|
||||
out = attn @ V_b
|
||||
outputs.append(out.to(out_dtype))
|
||||
|
||||
cum_q += q_len
|
||||
|
||||
return torch.stack(outputs, dim=0)
|
||||
|
||||
|
||||
def _build_cp_paged_kv_cache(
|
||||
seq_lens: list[int],
|
||||
k_nope_contexts: list[torch.Tensor],
|
||||
k_rope_contexts: list[torch.Tensor],
|
||||
block_size: int,
|
||||
kv_lora_rank: int,
|
||||
qk_rope_head_dim: int,
|
||||
cp_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
||||
"""Allocate a paged KV cache that matches the *decode* CP gather/block_table layout.
|
||||
|
||||
Each request is padded so its block count is exactly ``cp_size * L``
|
||||
(i.e. divisible by the total CP world size). The local rank is treated
|
||||
as "rank 0" and owns the first ``L`` of each request's logical blocks
|
||||
(the actual values in ``local_block_table`` only matter via shape because
|
||||
``gather_kv_cross_cp`` is mocked).
|
||||
|
||||
The ``gathered_*`` tensors mirror what the impl would see after both
|
||||
DCP and PCP all-gathers; layout matches the indexing produced by
|
||||
``gather_block_table``:
|
||||
``gathered[i*N + r*L + b]`` holds K data for logical block
|
||||
``b*cp_size + i`` of request ``r`` (where ``N == batch_size * L``).
|
||||
"""
|
||||
batch_size = len(seq_lens)
|
||||
|
||||
raw_blocks_per_req = [(s + block_size - 1) // block_size for s in seq_lens]
|
||||
max_raw = max(raw_blocks_per_req)
|
||||
total_blocks_per_req = ((max_raw + cp_size - 1) // cp_size) * cp_size
|
||||
L = total_blocks_per_req // cp_size
|
||||
|
||||
total_blocks = batch_size * total_blocks_per_req + 1 # +1 reserves block 0
|
||||
|
||||
full_k_nope_cache = torch.zeros(total_blocks, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
|
||||
full_k_rope_cache = torch.zeros(total_blocks, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
full_block_table = torch.zeros(batch_size, total_blocks_per_req, dtype=torch.int32, device=device)
|
||||
|
||||
next_block = 1
|
||||
for r in range(batch_size):
|
||||
s_len = seq_lens[r]
|
||||
for p in range(total_blocks_per_req):
|
||||
full_block_table[r, p] = next_block
|
||||
tok_start = p * block_size
|
||||
tok_end = min(tok_start + block_size, s_len)
|
||||
length = max(tok_end - tok_start, 0)
|
||||
if length > 0:
|
||||
full_k_nope_cache[next_block, :length, 0, :] = k_nope_contexts[r][tok_start:tok_end]
|
||||
full_k_rope_cache[next_block, :length, 0, :] = k_rope_contexts[r][tok_start:tok_end]
|
||||
next_block += 1
|
||||
|
||||
local_block_table = full_block_table[:, :L].contiguous()
|
||||
|
||||
N = batch_size * L
|
||||
gathered_k_nope = torch.zeros(cp_size * N, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
|
||||
gathered_k_rope = torch.zeros(cp_size * N, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
for i in range(cp_size):
|
||||
for r in range(batch_size):
|
||||
for b in range(L):
|
||||
p_logical = b * cp_size + i
|
||||
phys_block = int(full_block_table[r, p_logical].item())
|
||||
dst = i * N + r * L + b
|
||||
gathered_k_nope[dst] = full_k_nope_cache[phys_block]
|
||||
gathered_k_rope[dst] = full_k_rope_cache[phys_block]
|
||||
|
||||
return (
|
||||
full_k_nope_cache,
|
||||
full_k_rope_cache,
|
||||
full_block_table,
|
||||
local_block_table,
|
||||
gathered_k_nope,
|
||||
gathered_k_rope,
|
||||
L,
|
||||
)
|
||||
|
||||
|
||||
def _build_cp_prefill_compact_metadata(
|
||||
prefill_full_block_table: torch.Tensor,
|
||||
full_k_nope_cache: torch.Tensor,
|
||||
full_k_rope_cache: torch.Tensor,
|
||||
cp_size: int,
|
||||
block_size: int,
|
||||
seq_lens: list[int],
|
||||
kv_lora_rank: int,
|
||||
qk_rope_head_dim: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build the prefill compact metadata + matching gathered KV view.
|
||||
|
||||
Reproduces the production
|
||||
``AscendSFACPMetadataBuilder.build_prefill_compact_block_metadata``
|
||||
formula: ``block_table_cp[r, b*cp_size + i] = new_block_table[r, b] + i*M``
|
||||
where ``M = num_unique_local_blocks``. Then materialises the
|
||||
``gathered_compact`` tensor that ``gather_kv_cross_cp_compact`` would
|
||||
produce so that the kernel, when walking ``block_table_cp[r, :]``,
|
||||
sees the contiguous K context of request ``r`` block-by-block.
|
||||
|
||||
Inputs
|
||||
------
|
||||
prefill_full_block_table:
|
||||
Shape ``(num_prefill_reqs, L * cp_size)``. The "full" (logical)
|
||||
block table for prefill requests as built by
|
||||
``_build_cp_paged_kv_cache``. Entry ``[r, p]`` holds the physical
|
||||
block id that owns logical position ``p`` of request ``r``.
|
||||
|
||||
The "local" (rank-0) prefill block table -- which the production code
|
||||
actually receives in ``attn_metadata.block_table[num_decodes:]`` -- is
|
||||
derived as ``prefill_full_block_table[:, :L]``. This matches the
|
||||
convention used by ``_build_cp_paged_kv_cache`` for the decode path.
|
||||
"""
|
||||
num_prefill_reqs, total_blocks_per_req = prefill_full_block_table.shape
|
||||
assert total_blocks_per_req % cp_size == 0, (
|
||||
f"prefill_full_block_table has {total_blocks_per_req} columns, which is not divisible by cp_size={cp_size}"
|
||||
)
|
||||
L = total_blocks_per_req // cp_size
|
||||
|
||||
prefill_local_block_table = prefill_full_block_table[:, :L].contiguous()
|
||||
|
||||
block_arange = torch.arange(cp_size, dtype=prefill_local_block_table.dtype, device=device)
|
||||
valid_block_ids, new_block_table_flat = prefill_local_block_table.flatten().unique(return_inverse=True)
|
||||
num_blocks = valid_block_ids.shape[0]
|
||||
block_table_cp = (
|
||||
new_block_table_flat.unsqueeze(-1).to(prefill_local_block_table)
|
||||
+ (block_arange * num_blocks).view(1, 1, -1).to(prefill_local_block_table)
|
||||
).reshape(prefill_local_block_table.shape[0], -1)
|
||||
|
||||
new_block_table_2d = new_block_table_flat.view(prefill_local_block_table.shape)
|
||||
|
||||
M = int(num_blocks)
|
||||
gathered_compact_nope = torch.zeros(cp_size * M, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
|
||||
gathered_compact_rope = torch.zeros(cp_size * M, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
for r in range(num_prefill_reqs):
|
||||
s_len = seq_lens[r]
|
||||
for b_local in range(L):
|
||||
for i in range(cp_size):
|
||||
p_logical = b_local * cp_size + i
|
||||
tok_start = p_logical * block_size
|
||||
tok_end = min(tok_start + block_size, s_len)
|
||||
length = max(tok_end - tok_start, 0)
|
||||
if length <= 0:
|
||||
continue
|
||||
src_block = int(prefill_full_block_table[r, p_logical].item())
|
||||
dst = int(new_block_table_2d[r, b_local].item()) + i * M
|
||||
gathered_compact_nope[dst, :length, 0, :] = full_k_nope_cache[src_block, :length, 0, :]
|
||||
gathered_compact_rope[dst, :length, 0, :] = full_k_rope_cache[src_block, :length, 0, :]
|
||||
|
||||
return valid_block_ids, block_table_cp, gathered_compact_nope, gathered_compact_rope
|
||||
|
||||
|
||||
def _make_fake_self(
|
||||
*,
|
||||
scale: float,
|
||||
pcp_size: int,
|
||||
dcp_size: int,
|
||||
gather_kv_cross_cp_fn: Callable | None = None,
|
||||
gather_kv_cross_cp_compact_fn: Callable | None = None,
|
||||
) -> MagicMock:
|
||||
"""Construct a ``MagicMock`` matching the slice of ``AscendSFACPImpl``
|
||||
that ``_execute_sparse_flash_attention_process`` reads.
|
||||
|
||||
Bound methods that don't depend on init-time state
|
||||
(``gather_block_table``, ``_execute_sparse_flash_attention``) are
|
||||
delegated to the real (unbound) implementations so the production
|
||||
kernel call stays under test. The collective gathers
|
||||
(``gather_kv_cross_cp`` / ``gather_kv_cross_cp_compact``) are mocked
|
||||
via ``side_effect`` because they require an initialised PCP / DCP
|
||||
``ProcessGroupHCCL``, which can't be created on a single rank.
|
||||
"""
|
||||
fake_self = MagicMock()
|
||||
fake_self.scale = scale
|
||||
fake_self.pcp_size = pcp_size
|
||||
fake_self.dcp_size = dcp_size
|
||||
|
||||
if gather_kv_cross_cp_fn is not None:
|
||||
fake_self.gather_kv_cross_cp = MagicMock(side_effect=gather_kv_cross_cp_fn)
|
||||
if gather_kv_cross_cp_compact_fn is not None:
|
||||
fake_self.gather_kv_cross_cp_compact = MagicMock(side_effect=gather_kv_cross_cp_compact_fn)
|
||||
|
||||
fake_self.gather_block_table = lambda block_num, block_tables, block_arange: AscendSFACPImpl.gather_block_table(
|
||||
fake_self, block_num, block_tables, block_arange
|
||||
)
|
||||
fake_self._execute_sparse_flash_attention = lambda *args, **kwargs: AscendSFACPImpl._execute_sparse_flash_attention(
|
||||
fake_self, *args, **kwargs
|
||||
)
|
||||
fake_self._align_to_graph_bucket_tokens = lambda x, m: x
|
||||
return fake_self
|
||||
|
||||
|
||||
def _run_sfa_cp_kernel(
|
||||
*,
|
||||
ql_nope: torch.Tensor,
|
||||
q_pe: torch.Tensor,
|
||||
full_k_nope_cache: torch.Tensor,
|
||||
full_k_rope_cache: torch.Tensor,
|
||||
local_block_table: torch.Tensor,
|
||||
topk_indices: torch.Tensor,
|
||||
cum_query_lens: torch.Tensor,
|
||||
seq_lens_tensor: torch.Tensor,
|
||||
scale: float,
|
||||
pcp_size: int,
|
||||
dcp_size: int,
|
||||
num_decodes: int,
|
||||
num_decode_tokens: int,
|
||||
num_prefills: int,
|
||||
device: torch.device,
|
||||
gathered_k_nope: torch.Tensor | None = None,
|
||||
gathered_k_rope: torch.Tensor | None = None,
|
||||
valid_block_ids: torch.Tensor | None = None,
|
||||
block_table_cp: torch.Tensor | None = None,
|
||||
gathered_compact_nope: torch.Tensor | None = None,
|
||||
gathered_compact_rope: torch.Tensor | None = None,
|
||||
prefill_q_cum_seqlens: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run the SFA-CP kernel with mocked CP collectives."""
|
||||
cp_size = pcp_size * dcp_size
|
||||
|
||||
gather_kv_cross_cp_fn = None
|
||||
gather_kv_cross_cp_compact_fn = None
|
||||
|
||||
if num_decode_tokens > 0:
|
||||
assert gathered_k_nope is not None and gathered_k_rope is not None, (
|
||||
"decode branch requires gathered_k_nope/gathered_k_rope"
|
||||
)
|
||||
gathered_lookup = {
|
||||
id(full_k_nope_cache): gathered_k_nope,
|
||||
id(full_k_rope_cache): gathered_k_rope,
|
||||
}
|
||||
|
||||
def gather_kv_cross_cp_fn(kv: torch.Tensor, block_tables: torch.Tensor):
|
||||
gathered = gathered_lookup.get(id(kv))
|
||||
assert gathered is not None, "gather_kv_cross_cp called with an unexpected kv tensor"
|
||||
|
||||
return gathered, block_tables.numel()
|
||||
|
||||
if num_prefills > 0:
|
||||
assert (
|
||||
valid_block_ids is not None
|
||||
and block_table_cp is not None
|
||||
and gathered_compact_nope is not None
|
||||
and gathered_compact_rope is not None
|
||||
and prefill_q_cum_seqlens is not None
|
||||
), (
|
||||
"prefill branch requires valid_block_ids / block_table_cp / "
|
||||
"gathered_compact_{nope,rope} / prefill_q_cum_seqlens"
|
||||
)
|
||||
gathered_compact_lookup = {
|
||||
id(full_k_nope_cache): gathered_compact_nope,
|
||||
id(full_k_rope_cache): gathered_compact_rope,
|
||||
}
|
||||
|
||||
def gather_kv_cross_cp_compact_fn(kv: torch.Tensor, vbid: torch.Tensor):
|
||||
gathered = gathered_compact_lookup.get(id(kv))
|
||||
assert gathered is not None, "gather_kv_cross_cp_compact called with an unexpected kv tensor"
|
||||
|
||||
return gathered
|
||||
|
||||
fake_self = _make_fake_self(
|
||||
scale=scale,
|
||||
pcp_size=pcp_size,
|
||||
dcp_size=dcp_size,
|
||||
gather_kv_cross_cp_fn=gather_kv_cross_cp_fn,
|
||||
gather_kv_cross_cp_compact_fn=gather_kv_cross_cp_compact_fn,
|
||||
)
|
||||
|
||||
fake_attn_metadata = MagicMock()
|
||||
fake_attn_metadata.block_table = local_block_table
|
||||
fake_attn_metadata.num_decodes = num_decodes
|
||||
fake_attn_metadata.num_decode_tokens = num_decode_tokens
|
||||
fake_attn_metadata.num_prefills = num_prefills
|
||||
|
||||
fake_sfa_cp_metadata = MagicMock()
|
||||
fake_sfa_cp_metadata.block_arange = torch.arange(cp_size, dtype=torch.int32, device=device)
|
||||
if num_prefills > 0:
|
||||
fake_sfa_cp_metadata.valid_block_ids = valid_block_ids
|
||||
fake_sfa_cp_metadata.block_table_cp = block_table_cp
|
||||
fake_sfa_cp_metadata.prefill_q_cum_seqlens = prefill_q_cum_seqlens
|
||||
fake_attn_metadata.sfa_cp_metadata = fake_sfa_cp_metadata
|
||||
|
||||
cum_lens_arg = cum_query_lens if num_decode_tokens > 0 else prefill_q_cum_seqlens
|
||||
|
||||
return AscendSFACPImpl._execute_sparse_flash_attention_process(
|
||||
fake_self,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
(full_k_nope_cache, full_k_rope_cache),
|
||||
topk_indices,
|
||||
fake_attn_metadata,
|
||||
cum_lens_arg,
|
||||
seq_lens_tensor,
|
||||
)
|
||||
|
||||
|
||||
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."""
|
||||
assert backend_output.shape == reference_output.shape, (
|
||||
f"[{tag}] backend shape {tuple(backend_output.shape)} != reference shape {tuple(reference_output.shape)}"
|
||||
)
|
||||
assert backend_output.dtype == reference_output.dtype, (
|
||||
f"[{tag}] backend dtype {backend_output.dtype} != reference dtype {reference_output.dtype}"
|
||||
)
|
||||
assert torch.isfinite(backend_output).all(), f"[{tag}] sparse flash attention produced non-finite values"
|
||||
|
||||
torch.testing.assert_close(
|
||||
backend_output,
|
||||
reference_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=lambda m: f"[SFA-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
|
||||
|
||||
assert max_sig_rel_err < _MAX_SIG_REL_ERR, (
|
||||
f"[SFA-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"[SFA-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"[SFA-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 _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 generated independently so values across
|
||||
requests are uncorrelated."""
|
||||
k_nope = [torch.randn(s, kv_lora_rank, dtype=dtype, device=device) * 0.1 for s in seq_lens]
|
||||
k_rope = [torch.randn(s, qk_rope_head_dim, dtype=dtype, device=device) * 0.1 for s in seq_lens]
|
||||
return k_nope, k_rope
|
||||
|
||||
|
||||
def _test_sfa_cp_correctness(
|
||||
batch_spec: BatchSpec,
|
||||
model: str,
|
||||
*,
|
||||
pcp_size: int,
|
||||
dcp_size: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
atol: float = 1e-2,
|
||||
rtol: float = 1e-2,
|
||||
tensor_parallel_size: int = 1,
|
||||
) -> None:
|
||||
"""Test ``AscendSFACPImpl`` against a fp32 dense MQA 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)
|
||||
batch_size = batch_spec.batch_size
|
||||
num_tokens = batch_spec.compute_num_tokens()
|
||||
cp_size = pcp_size * dcp_size
|
||||
|
||||
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_rope_head_dim = hf_text.qk_rope_head_dim
|
||||
num_heads = hf_text.num_attention_heads
|
||||
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim
|
||||
scale = 1.0 / math.sqrt(head_dim)
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec,
|
||||
block_size=block_size,
|
||||
device=device,
|
||||
)
|
||||
|
||||
k_nope_contexts, k_rope_contexts = _make_synthetic_kv_contexts(
|
||||
seq_lens,
|
||||
kv_lora_rank,
|
||||
qk_rope_head_dim,
|
||||
dtype,
|
||||
device,
|
||||
)
|
||||
(
|
||||
full_k_nope_cache,
|
||||
full_k_rope_cache,
|
||||
full_block_table,
|
||||
local_block_table,
|
||||
gathered_k_nope,
|
||||
gathered_k_rope,
|
||||
_L,
|
||||
) = _build_cp_paged_kv_cache(
|
||||
seq_lens=seq_lens,
|
||||
k_nope_contexts=k_nope_contexts,
|
||||
k_rope_contexts=k_rope_contexts,
|
||||
block_size=block_size,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
cp_size=cp_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
ql_nope = (
|
||||
torch.randn(
|
||||
num_tokens,
|
||||
num_heads,
|
||||
kv_lora_rank,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
q_pe = (
|
||||
torch.randn(
|
||||
num_tokens,
|
||||
num_heads,
|
||||
qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
topk_indices = _build_topk_indices(seq_lens, query_lens, SPARSE_COUNT, device)
|
||||
|
||||
cum_query_lens = common_attn_metadata.query_start_loc[1:].to(torch.int32)
|
||||
seq_lens_tensor = common_attn_metadata.seq_lens.to(torch.int32)
|
||||
|
||||
decode_gathered_k_nope: torch.Tensor | None = None
|
||||
decode_gathered_k_rope: torch.Tensor | None = None
|
||||
prefill_full_block_table: torch.Tensor | None = None
|
||||
|
||||
if mode in ("decode", "mtp"):
|
||||
num_decodes = batch_size
|
||||
num_decode_tokens = num_tokens
|
||||
num_prefills = 0
|
||||
decode_gathered_k_nope = gathered_k_nope
|
||||
decode_gathered_k_rope = gathered_k_rope
|
||||
n_decode_reqs = batch_size
|
||||
elif mode == "prefill":
|
||||
num_decodes = 0
|
||||
num_decode_tokens = 0
|
||||
num_prefills = batch_size
|
||||
prefill_full_block_table = full_block_table
|
||||
n_decode_reqs = 0
|
||||
else: # mixed (chunked-prefill)
|
||||
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"
|
||||
)
|
||||
num_decodes = n_decode_reqs
|
||||
num_decode_tokens = n_decode_reqs
|
||||
num_prefills = batch_size - n_decode_reqs
|
||||
prefill_full_block_table = full_block_table[n_decode_reqs:]
|
||||
|
||||
L_decode = local_block_table[:n_decode_reqs].shape[1]
|
||||
N_decode = n_decode_reqs * L_decode
|
||||
decode_gathered_k_nope = torch.zeros(
|
||||
cp_size * N_decode,
|
||||
block_size,
|
||||
1,
|
||||
kv_lora_rank,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
decode_gathered_k_rope = torch.zeros(
|
||||
cp_size * N_decode,
|
||||
block_size,
|
||||
1,
|
||||
qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
for i in range(cp_size):
|
||||
for r in range(n_decode_reqs):
|
||||
for b in range(L_decode):
|
||||
p_logical = b * cp_size + i
|
||||
phys_block = int(full_block_table[r, p_logical].item())
|
||||
dst = i * N_decode + r * L_decode + b
|
||||
decode_gathered_k_nope[dst] = full_k_nope_cache[phys_block]
|
||||
decode_gathered_k_rope[dst] = full_k_rope_cache[phys_block]
|
||||
|
||||
valid_block_ids = None
|
||||
block_table_cp = None
|
||||
gathered_compact_nope = None
|
||||
gathered_compact_rope = None
|
||||
prefill_q_cum_seqlens = None
|
||||
if num_prefills > 0:
|
||||
valid_block_ids, block_table_cp, gathered_compact_nope, gathered_compact_rope = (
|
||||
_build_cp_prefill_compact_metadata(
|
||||
prefill_full_block_table=prefill_full_block_table,
|
||||
full_k_nope_cache=full_k_nope_cache,
|
||||
full_k_rope_cache=full_k_rope_cache,
|
||||
cp_size=cp_size,
|
||||
block_size=block_size,
|
||||
seq_lens=seq_lens[n_decode_reqs:],
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
if n_decode_reqs > 0:
|
||||
prefill_q_cum_seqlens = cum_query_lens[n_decode_reqs:] - cum_query_lens[n_decode_reqs - 1]
|
||||
else:
|
||||
prefill_q_cum_seqlens = cum_query_lens
|
||||
|
||||
with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
|
||||
backend_output = _run_sfa_cp_kernel(
|
||||
ql_nope=ql_nope,
|
||||
q_pe=q_pe,
|
||||
full_k_nope_cache=full_k_nope_cache,
|
||||
full_k_rope_cache=full_k_rope_cache,
|
||||
local_block_table=local_block_table,
|
||||
topk_indices=topk_indices,
|
||||
cum_query_lens=cum_query_lens,
|
||||
seq_lens_tensor=seq_lens_tensor,
|
||||
scale=scale,
|
||||
pcp_size=pcp_size,
|
||||
dcp_size=dcp_size,
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
num_prefills=num_prefills,
|
||||
device=device,
|
||||
gathered_k_nope=decode_gathered_k_nope,
|
||||
gathered_k_rope=decode_gathered_k_rope,
|
||||
valid_block_ids=valid_block_ids,
|
||||
block_table_cp=block_table_cp,
|
||||
gathered_compact_nope=gathered_compact_nope,
|
||||
gathered_compact_rope=gathered_compact_rope,
|
||||
prefill_q_cum_seqlens=prefill_q_cum_seqlens,
|
||||
)
|
||||
reference_output = _reference_sparse_attention(
|
||||
ql_nope=ql_nope,
|
||||
q_pe=q_pe,
|
||||
k_nope_full_per_req=k_nope_contexts,
|
||||
k_rope_full_per_req=k_rope_contexts,
|
||||
seq_lens=seq_lens,
|
||||
query_lens=query_lens,
|
||||
scale=scale,
|
||||
out_dtype=dtype,
|
||||
)
|
||||
|
||||
dt = "bf16" if dtype == torch.bfloat16 else "fp16"
|
||||
tag = f"{mode}|{batch_spec.name}|pcp={pcp_size}|dcp={dcp_size}|tp={tensor_parallel_size}|{dt}"
|
||||
|
||||
_record_and_assert(
|
||||
backend_output,
|
||||
reference_output,
|
||||
tag,
|
||||
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_sfa_cp.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.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_sfa_cp_correctness(
|
||||
batch_spec_name: str,
|
||||
model: str,
|
||||
pcp_size: int,
|
||||
dcp_size: int,
|
||||
dtype: torch.dtype,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
"""Test SFA-CP correctness across workload, topology, dtype, and TP size."""
|
||||
_test_sfa_cp_correctness(
|
||||
BATCH_SPECS[batch_spec_name],
|
||||
model,
|
||||
pcp_size=pcp_size,
|
||||
dcp_size=dcp_size,
|
||||
dtype=dtype,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
742
tests/ut/attention/a2/test_sfa_v1.py
Normal file
742
tests/ut/attention/a2/test_sfa_v1.py
Normal file
@@ -0,0 +1,742 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.distributed.parallel_state import GroupCoordinator
|
||||
|
||||
from tests.ut.attention.utils import patch_distributed_groups
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.ascend_config import init_ascend_config
|
||||
from vllm_ascend.attention.attention_v1 import AscendAttentionState
|
||||
|
||||
if "torch_npu._inductor" not in sys.modules:
|
||||
sys.modules["torch_npu._inductor"] = MagicMock()
|
||||
|
||||
from vllm_ascend.attention.sfa_v1 import (
|
||||
AscendSFABackend,
|
||||
AscendSFAImpl,
|
||||
AscendSFAMetadata,
|
||||
AscendSFAMetadataBuilder,
|
||||
custom_kv_rmsnorm_rope,
|
||||
)
|
||||
from vllm_ascend.attention.utils import get_sfa_qsfa_packed_head_dim
|
||||
from vllm_ascend.device.device_op import DeviceOperator
|
||||
from vllm_ascend.utils import enable_dsa_cp
|
||||
|
||||
|
||||
class TestAscendSFABackend(TestBase):
|
||||
def setUp(self):
|
||||
self.mock_config = MagicMock()
|
||||
mock_parallel_config = MagicMock()
|
||||
mock_parallel_config.prefill_context_parallel_size = 1
|
||||
mock_parallel_config.decode_context_parallel_size = 1
|
||||
self.mock_config.parallel_config = mock_parallel_config
|
||||
self.mock_config.model_config = MagicMock(spec=[])
|
||||
self.config_context = set_current_vllm_config(self.mock_config)
|
||||
self.config_context.__enter__()
|
||||
|
||||
self.utils_patcher = patch("vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_config)
|
||||
self.utils_patcher.start()
|
||||
|
||||
from vllm_ascend.attention.utils import enable_cp
|
||||
|
||||
enable_cp.cache_clear()
|
||||
|
||||
def tearDown(self):
|
||||
self.utils_patcher.stop()
|
||||
self.config_context.__exit__(None, None, None)
|
||||
|
||||
def test_get_name(self):
|
||||
self.assertEqual(AscendSFABackend.get_name(), "ASCEND_SFA")
|
||||
|
||||
def test_get_builder_cls(self):
|
||||
self.assertEqual(AscendSFABackend.get_builder_cls(), AscendSFAMetadataBuilder)
|
||||
|
||||
def test_get_kv_cache_shape(self):
|
||||
result = AscendSFABackend.get_kv_cache_shape(2, 4, 8, 128)
|
||||
self.assertEqual(result, (2, 4, 8, 128))
|
||||
|
||||
def test_get_impl_cls(self):
|
||||
result = AscendSFABackend.get_impl_cls()
|
||||
self.assertEqual(result, AscendSFAImpl)
|
||||
|
||||
@patch("vllm_ascend.attention.sfa_v1.enable_cp")
|
||||
def test_get_builder_cls_with_cp(self, mock_enable_cp):
|
||||
mock_enable_cp.return_value = True
|
||||
builder_cls = AscendSFABackend.get_builder_cls()
|
||||
self.assertIsNotNone(builder_cls)
|
||||
|
||||
@patch("vllm_ascend.attention.sfa_v1.enable_cp")
|
||||
def test_get_impl_cls_with_cp(self, mock_enable_cp):
|
||||
mock_enable_cp.return_value = True
|
||||
impl_cls = AscendSFABackend.get_impl_cls()
|
||||
self.assertIsNotNone(impl_cls)
|
||||
|
||||
|
||||
class TestAscendSFABytePackedGather(TestBase):
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_tp_group")
|
||||
def test_byte_packed_gather_preserves_mixed_dtype_tensors(self, mock_get_tp_group):
|
||||
mock_get_tp_group.return_value = SimpleNamespace(world_size=1)
|
||||
sfa_kv = torch.arange(12, dtype=torch.float16).view(2, 6)
|
||||
k_li = torch.arange(16, dtype=torch.int8).view(2, 1, 8)
|
||||
k_li_scale = torch.arange(2, dtype=torch.float32).view(2, 1)
|
||||
|
||||
gathered, handle, metadata = AscendSFAImpl._all_gather_byte_packed_async(
|
||||
[
|
||||
("sfa_kv", sfa_kv),
|
||||
("k_li", k_li),
|
||||
("k_li_scale", k_li_scale),
|
||||
],
|
||||
async_op=True,
|
||||
)
|
||||
|
||||
self.assertIsNone(handle)
|
||||
self.assertEqual(gathered.dtype, torch.int8)
|
||||
restored = AscendSFAImpl._restore_byte_gathered_tensors(gathered, metadata)
|
||||
for name, expected in (("sfa_kv", sfa_kv), ("k_li", k_li), ("k_li_scale", k_li_scale)):
|
||||
self.assertEqual(restored[name].shape, expected.shape)
|
||||
self.assertEqual(restored[name].dtype, expected.dtype)
|
||||
self.assertTrue(torch.equal(restored[name], expected))
|
||||
|
||||
def test_byte_packed_gather_rejects_mismatched_token_counts(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "different token counts"):
|
||||
AscendSFAImpl._all_gather_byte_packed_async(
|
||||
[
|
||||
("sfa_kv", torch.zeros(2, 6, dtype=torch.float16)),
|
||||
("k_li", torch.zeros(3, 8, dtype=torch.int8)),
|
||||
],
|
||||
async_op=True,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendSFADeviceOperator(TestBase):
|
||||
def _make_common_inputs(self):
|
||||
ql_nope = torch.randn(3, 4, 8)
|
||||
q_pe = torch.randn(3, 4, 2)
|
||||
topk_indices = torch.zeros(3, 1, dtype=torch.int32)
|
||||
attn_metadata = MagicMock()
|
||||
attn_metadata.block_table = torch.zeros(1, 4, dtype=torch.int32)
|
||||
actual_seq_lengths_query = torch.tensor([3], dtype=torch.int32)
|
||||
actual_seq_lengths_key = torch.tensor([3], dtype=torch.int32)
|
||||
impl = MagicMock()
|
||||
impl.scale = 0.125
|
||||
impl.qk_rope_head_dim = 2
|
||||
impl.sfa_qsfa_tile_size = 128
|
||||
return (
|
||||
impl,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
topk_indices,
|
||||
attn_metadata,
|
||||
actual_seq_lengths_query,
|
||||
actual_seq_lengths_key,
|
||||
)
|
||||
|
||||
def test_execute_sparse_flash_attention_returns_lse(self):
|
||||
(
|
||||
impl,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
topk_indices,
|
||||
attn_metadata,
|
||||
actual_seq_lengths_query,
|
||||
actual_seq_lengths_key,
|
||||
) = self._make_common_inputs()
|
||||
kv_cache = (
|
||||
torch.randn(4, 1, 1, 8),
|
||||
torch.randn(4, 1, 1, 2),
|
||||
)
|
||||
attn_output = torch.randn(3, 4, 8)
|
||||
softmax_max = torch.zeros(1, 3, 4)
|
||||
softmax_sum = torch.full((1, 3, 4), 2.0)
|
||||
|
||||
with patch.object(
|
||||
torch.ops._C_ascend,
|
||||
"npu_sparse_flash_attention",
|
||||
create=True,
|
||||
return_value=(attn_output, softmax_max, softmax_sum),
|
||||
) as mock_sfa:
|
||||
output, softmax_lse = DeviceOperator.execute_sparse_flash_attention_process(
|
||||
impl,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
kv_cache,
|
||||
topk_indices,
|
||||
attn_metadata,
|
||||
actual_seq_lengths_query,
|
||||
actual_seq_lengths_key,
|
||||
return_lse=True,
|
||||
)
|
||||
|
||||
self.assertIs(output, attn_output)
|
||||
self.assertEqual(softmax_lse.shape, (3, 4, 1))
|
||||
expected_lse = torch.full((3, 4, 1), torch.log(torch.tensor(2.0)).item())
|
||||
self.assertTrue(torch.allclose(softmax_lse, expected_lse))
|
||||
self.assertTrue(mock_sfa.call_args.kwargs["return_softmax_lse"])
|
||||
|
||||
def test_execute_sparse_flash_attention_c8_returns_lse(self):
|
||||
(
|
||||
impl,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
topk_indices,
|
||||
attn_metadata,
|
||||
actual_seq_lengths_query,
|
||||
actual_seq_lengths_key,
|
||||
) = self._make_common_inputs()
|
||||
packed_kv_cache = (torch.empty(4, 1, 1, 12, dtype=torch.int8),)
|
||||
attn_output = torch.randn(3, 4, 8)
|
||||
softmax_max = torch.ones(1, 3, 4)
|
||||
softmax_sum = torch.full((1, 3, 4), 3.0)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
torch.ops._C_ascend,
|
||||
"npu_kv_quant_sparse_flash_attention",
|
||||
create=True,
|
||||
return_value=(attn_output, softmax_max, softmax_sum),
|
||||
) as mock_qsfa,
|
||||
patch(
|
||||
"vllm_ascend.device.device_op.torch_npu.npu_kv_quant_sparse_flash_attention",
|
||||
create=True,
|
||||
side_effect=AssertionError("C8 SFA with LSE must use the custom op"),
|
||||
),
|
||||
):
|
||||
output, softmax_lse = DeviceOperator.execute_sparse_flash_attention_process(
|
||||
impl,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
packed_kv_cache,
|
||||
topk_indices,
|
||||
attn_metadata,
|
||||
actual_seq_lengths_query,
|
||||
actual_seq_lengths_key,
|
||||
sparse_mode=0,
|
||||
return_lse=True,
|
||||
)
|
||||
|
||||
self.assertIs(output, attn_output)
|
||||
expected_lse = torch.full((3, 4, 1), 1.0 + torch.log(torch.tensor(3.0)).item())
|
||||
self.assertTrue(torch.allclose(softmax_lse, expected_lse))
|
||||
call_kwargs = mock_qsfa.call_args.kwargs
|
||||
self.assertIs(call_kwargs["key"], packed_kv_cache[0])
|
||||
self.assertIs(call_kwargs["value"], packed_kv_cache[0])
|
||||
self.assertEqual(call_kwargs["query"].shape, (3, 4, 10))
|
||||
self.assertEqual(call_kwargs["sparse_mode"], 0)
|
||||
self.assertTrue(call_kwargs["return_softmax_lse"])
|
||||
|
||||
|
||||
class TestAscendSFAKVQuantSparseAttention(TestBase):
|
||||
@patch("vllm_ascend.attention.sfa_v1.torch_npu.npu_dynamic_block_quant")
|
||||
@patch("vllm_ascend.attention.sfa_v1.torch_npu.npu_interleave_rope")
|
||||
@patch("vllm_ascend.attention.sfa_v1.torch_npu.npu_rms_norm")
|
||||
def test_pack_prefill_kv_cache(self, mock_rms_norm, mock_rope, mock_block_quant):
|
||||
k_nope = torch.randn(2, 1, 1, 256, dtype=torch.bfloat16)
|
||||
k_pe = torch.randn(2, 1, 1, 16, dtype=torch.bfloat16)
|
||||
quantized = torch.randint(-128, 127, (2, 1, 256), dtype=torch.int8)
|
||||
scales = torch.arange(1, 5, dtype=torch.float32).view(2, 1, 2)
|
||||
mock_rms_norm.return_value = k_nope, None
|
||||
mock_rope.return_value = k_pe
|
||||
mock_block_quant.return_value = quantized, scales
|
||||
|
||||
actual_k_pe, actual_k_nope, actual_scales = custom_kv_rmsnorm_rope(
|
||||
torch.randn(2, 1, 1, 272, dtype=torch.bfloat16),
|
||||
torch.ones(256, dtype=torch.bfloat16),
|
||||
torch.randn(2, 1, 1, 16),
|
||||
torch.randn(2, 1, 1, 16),
|
||||
256,
|
||||
16,
|
||||
dst_type=1,
|
||||
tile_size=128,
|
||||
)
|
||||
packed_kv = torch.cat([actual_k_nope, actual_k_pe, actual_scales], dim=-1)
|
||||
|
||||
self.assertEqual(mock_block_quant.call_args.kwargs["dst_type"], 1)
|
||||
self.assertEqual(mock_block_quant.call_args.kwargs["row_block_size"], 1)
|
||||
self.assertEqual(mock_block_quant.call_args.kwargs["col_block_size"], 128)
|
||||
self.assertEqual(packed_kv.shape, (2, 1, 1, 296))
|
||||
self.assertTrue(torch.equal(packed_kv[..., :256], quantized.view_as(k_nope)))
|
||||
self.assertTrue(torch.equal(packed_kv[..., 256:288], k_pe.contiguous().view(torch.int8)))
|
||||
self.assertTrue(torch.equal(packed_kv[..., 288:], scales.view(2, 1, 1, 2).view(torch.int8)))
|
||||
|
||||
def test_execute_kv_quant_sparse_flash_attention(self):
|
||||
impl = AscendSFAImpl.__new__(AscendSFAImpl)
|
||||
impl.enable_sparse_sfa_c8 = True
|
||||
impl.scale = 0.125
|
||||
impl.sfa_qsfa_tile_size = 128
|
||||
impl.qk_rope_head_dim = 16
|
||||
ql_nope = torch.randn(3, 2, 32)
|
||||
q_pe = torch.randn(3, 2, 16)
|
||||
kv_cache = (torch.empty(4, 16, 1, 80, dtype=torch.int8),)
|
||||
topk_indices = torch.zeros(3, 1, dtype=torch.int32)
|
||||
attn_metadata = SimpleNamespace(block_table=torch.zeros(1, 4, dtype=torch.int32))
|
||||
actual_seq_lengths = torch.tensor([3], dtype=torch.int32)
|
||||
expected = torch.randn(3, 2, 32)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
torch.ops._C_ascend,
|
||||
"npu_kv_quant_sparse_flash_attention",
|
||||
create=True,
|
||||
return_value=(expected, torch.empty(0), torch.empty(0)),
|
||||
) as mock_qsfa,
|
||||
patch(
|
||||
"vllm_ascend.device.device_op.torch_npu.npu_kv_quant_sparse_flash_attention",
|
||||
create=True,
|
||||
side_effect=AssertionError("Base must use _C_ascend custom op"),
|
||||
),
|
||||
):
|
||||
result = impl._execute_sparse_flash_attention_process(
|
||||
ql_nope,
|
||||
q_pe,
|
||||
kv_cache,
|
||||
topk_indices,
|
||||
attn_metadata,
|
||||
actual_seq_lengths,
|
||||
actual_seq_lengths,
|
||||
)
|
||||
|
||||
self.assertIs(result, expected)
|
||||
call_kwargs = mock_qsfa.call_args.kwargs
|
||||
self.assertIs(call_kwargs["key"], kv_cache[0])
|
||||
self.assertEqual(call_kwargs["query"].shape, (3, 2, 48))
|
||||
self.assertEqual(call_kwargs["key_quant_mode"], 2)
|
||||
self.assertEqual(call_kwargs["tile_size"], 128)
|
||||
self.assertEqual(call_kwargs["return_softmax_lse"], False)
|
||||
|
||||
def test_prolog_v3_enables_packed_int8_kv_cache(self):
|
||||
impl = AscendSFAImpl.__new__(AscendSFAImpl)
|
||||
impl.enable_sparse_sfa_c8 = True
|
||||
impl.has_indexer = True
|
||||
impl.sfa_qsfa_tile_size = 128
|
||||
impl.sfa_qsfa_k_nope_clip_alpha = torch.ones(1)
|
||||
impl.sfa_qsfa_kr_cache_dummy = torch.empty(0, dtype=torch.bfloat16)
|
||||
impl.local_num_heads = 2
|
||||
impl.kv_lora_rank = 128
|
||||
impl.qk_rope_head_dim = 16
|
||||
impl.q_lora_rank = 8
|
||||
impl.q_a_layernorm = SimpleNamespace(weight=SimpleNamespace(data=torch.ones(8)), variance_epsilon=1e-5)
|
||||
impl.kv_a_layernorm = SimpleNamespace(weight=SimpleNamespace(data=torch.ones(128)), variance_epsilon=1e-5)
|
||||
impl.weight_dq = torch.empty(1)
|
||||
impl.weight_uq_qr = torch.empty(1)
|
||||
impl.W_UK_T = torch.empty(1)
|
||||
impl.weight_dkv_kr = torch.empty(1)
|
||||
impl.dequant_scale_w_dq = torch.empty(1)
|
||||
impl.dequant_scale_w_uq_qr = torch.empty(1)
|
||||
impl.dequant_scale_w_dkv_kr = torch.empty(1)
|
||||
k_cache = torch.empty(4, 16, 1, get_sfa_qsfa_packed_head_dim(128, 16), dtype=torch.int8)
|
||||
dsa_k_cache = torch.empty(4, 16, 1, 128, dtype=torch.bfloat16)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm_ascend.device.device_op.torch_npu.npu_dynamic_quant",
|
||||
return_value=(torch.empty(2, 8, dtype=torch.int8), torch.ones(2, 1)),
|
||||
),
|
||||
patch(
|
||||
"vllm_ascend.device.device_op.torch_npu.npu_mla_prolog_v3",
|
||||
create=True,
|
||||
return_value=(torch.randn(2, 2, 128), torch.randn(2, 2, 16), None, torch.randn(2, 8), None),
|
||||
) as mock_prolog,
|
||||
):
|
||||
impl._sfa_preprocess_with_prolog_v3(
|
||||
hidden_states=torch.randn(2, 8),
|
||||
kv_cache=(k_cache, dsa_k_cache),
|
||||
cos=torch.randn(2, 1, 1, 16),
|
||||
sin=torch.randn(2, 1, 1, 16),
|
||||
slot_mapping=torch.arange(2),
|
||||
cache_mode="PA_BSND",
|
||||
)
|
||||
|
||||
call_kwargs = mock_prolog.call_args.kwargs
|
||||
self.assertIs(call_kwargs["kv_cache"], k_cache)
|
||||
self.assertIs(call_kwargs["kr_cache"], impl.sfa_qsfa_kr_cache_dummy)
|
||||
self.assertEqual(call_kwargs["kv_cache_quant_mode"], 3)
|
||||
self.assertEqual(call_kwargs["ckvkr_repo_mode"], 1)
|
||||
self.assertEqual(call_kwargs["quant_scale_repo_mode"], 1)
|
||||
|
||||
|
||||
class TestAscendSFAMetadata(TestBase):
|
||||
def test_ascend_sfa_metadata_default(self):
|
||||
num_actual_tokens = 100
|
||||
slot_mapping = torch.randn(100, 4, 1024)
|
||||
seq_lens = torch.tensor([30, 50])
|
||||
cum_query_lens = torch.tensor([0, 30, 80])
|
||||
block_table = torch.randint(0, 100, (100, 4))
|
||||
|
||||
rope_dim = 32
|
||||
max_seq_len = int(seq_lens.max().item())
|
||||
sin = torch.randn(max_seq_len, rope_dim)
|
||||
cos = torch.randn(max_seq_len, rope_dim)
|
||||
|
||||
num_input_tokens = 2
|
||||
head_dim = None
|
||||
attn_mask = None
|
||||
attn_state = AscendAttentionState.ChunkedPrefill
|
||||
|
||||
metadata = AscendSFAMetadata(
|
||||
num_actual_tokens=num_actual_tokens,
|
||||
slot_mapping=slot_mapping,
|
||||
seq_lens=seq_lens,
|
||||
seq_lens_cpu=seq_lens,
|
||||
cum_query_lens=cum_query_lens,
|
||||
block_table=block_table,
|
||||
sin=sin,
|
||||
cos=cos,
|
||||
num_input_tokens=num_input_tokens,
|
||||
head_dim=head_dim,
|
||||
attn_mask=attn_mask,
|
||||
attn_state=attn_state,
|
||||
)
|
||||
|
||||
self.assertEqual(metadata.num_actual_tokens, num_actual_tokens)
|
||||
self.assertIs(metadata.slot_mapping, slot_mapping)
|
||||
self.assertTrue(torch.equal(metadata.seq_lens, seq_lens))
|
||||
self.assertTrue(torch.equal(metadata.cum_query_lens, cum_query_lens))
|
||||
self.assertIs(metadata.block_table, block_table)
|
||||
self.assertIs(metadata.sin, sin)
|
||||
self.assertIs(metadata.cos, cos)
|
||||
self.assertEqual(metadata.num_input_tokens, num_input_tokens)
|
||||
self.assertIs(metadata.head_dim, head_dim)
|
||||
self.assertIs(metadata.attn_mask, attn_mask)
|
||||
self.assertEqual(metadata.attn_state, attn_state)
|
||||
|
||||
|
||||
class TestAscendSFAMetadataBuilder(TestBase):
|
||||
@patch("vllm.distributed.parallel_state._TP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
|
||||
def setUp(self, mock_tp):
|
||||
mock_tp.world_size = 2
|
||||
mock_tp.rank_in_group = MagicMock()
|
||||
mock_tp.device_group = MagicMock()
|
||||
|
||||
self.mock_cfg = MagicMock()
|
||||
|
||||
self.mock_cfg.parallel_config = MagicMock()
|
||||
self.mock_cfg.parallel_config.tensor_parallel_size = 1
|
||||
self.mock_cfg.parallel_config.prefill_context_parallel_size = 1
|
||||
self.mock_cfg.parallel_config.decode_context_parallel_size = 1
|
||||
|
||||
self.mock_cfg.compilation_config = MagicMock()
|
||||
self.mock_cfg.compilation_config.pass_config = MagicMock()
|
||||
self.mock_cfg.compilation_config.pass_config.enable_sp = False
|
||||
|
||||
self.mock_cfg.speculative_config.num_speculative_tokens = 0
|
||||
|
||||
self.mock_cfg.additional_config = {"refresh": True}
|
||||
init_ascend_config(self.mock_cfg)
|
||||
|
||||
self.patcher = patch("vllm.config.get_current_vllm_config", return_value=self.mock_cfg)
|
||||
self.patcher.start()
|
||||
|
||||
mock_ascend_config = MagicMock()
|
||||
mock_ascend_config.c8_enable_reshape_optim = False
|
||||
mock_ascend_config.enable_mlapo = True
|
||||
mock_ascend_config.enable_shared_expert_dp = False
|
||||
mock_ascend_config.layer_sharding = None
|
||||
self.ascend_config_patcher = patch(
|
||||
"vllm_ascend.attention.sfa_v1.get_ascend_config",
|
||||
return_value=mock_ascend_config,
|
||||
)
|
||||
self.ascend_config_patcher.start()
|
||||
|
||||
# Mock parent class __init__ to avoid complex initialization,
|
||||
# but still set the essential attributes that child class needs
|
||||
def mock_parent_init(
|
||||
self, kv_cache_spec, layer_names, vllm_config, device, metadata_cls, supports_dcp_with_varlen
|
||||
):
|
||||
self.metadata_cls = metadata_cls
|
||||
self.kv_cache_spec = kv_cache_spec
|
||||
self.model_config = vllm_config.model_config
|
||||
self.vllm_config = vllm_config
|
||||
self.device = device
|
||||
self.chunked_prefill_workspace_size = 128 * 1024
|
||||
self.chunked_prefill_workspace = torch.empty(
|
||||
(self.chunked_prefill_workspace_size, vllm_config.model_config.get_head_size()),
|
||||
dtype=vllm_config.model_config.dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.parent_init_patcher = patch(
|
||||
"vllm.model_executor.layers.attention.mla_attention.MLACommonMetadataBuilder.__init__", mock_parent_init
|
||||
)
|
||||
self.parent_init_patcher.start()
|
||||
|
||||
if hasattr(enable_dsa_cp, "cache_clear"):
|
||||
enable_dsa_cp.cache_clear()
|
||||
|
||||
def tearDown(self):
|
||||
self.patcher.stop()
|
||||
self.ascend_config_patcher.stop()
|
||||
self.parent_init_patcher.stop()
|
||||
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_ascend_sfa_metadata_builder_default(self):
|
||||
kv_cache_spec = MagicMock()
|
||||
kv_cache_spec.block_size = 128
|
||||
layer_names = ["layer1", "layer2"]
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.cache_config.block_size = 16
|
||||
vllm_config.scheduler_config.max_num_seqs = 16
|
||||
vllm_config.model_config.max_model_len = 1024
|
||||
vllm_config.model_config.get_head_size.return_value = 64
|
||||
vllm_config.model_config.dtype = torch.float16
|
||||
vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
|
||||
speculative_config = MagicMock()
|
||||
speculative_config.num_speculative_tokens = 4
|
||||
vllm_config.speculative_config = speculative_config
|
||||
device = torch.device("cpu")
|
||||
|
||||
builder = AscendSFAMetadataBuilder(
|
||||
kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
|
||||
)
|
||||
|
||||
assert builder.device == device
|
||||
assert builder.vllm_config == vllm_config
|
||||
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config")
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
|
||||
@patch("vllm_ascend.attention.sfa_v1.enable_dsa_cp")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_ascend_sfa_metadata_builder_build(
|
||||
self,
|
||||
mock_enable_dsa_cp,
|
||||
mock_get_cos_and_sin_mla,
|
||||
mock_get_current_vllm_config,
|
||||
):
|
||||
mock_enable_dsa_cp.return_value = False
|
||||
|
||||
cfg = MagicMock()
|
||||
cfg.model_config = MagicMock()
|
||||
cfg.model_config.hf_text_config = MagicMock()
|
||||
|
||||
mock_get_current_vllm_config.return_value = cfg
|
||||
kv_cache_spec = MagicMock()
|
||||
kv_cache_spec.block_size = 128
|
||||
layer_names = ["layer1", "layer2"]
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.cache_config.block_size = 16
|
||||
vllm_config.scheduler_config.max_num_seqs = 16
|
||||
vllm_config.model_config.max_model_len = 1024
|
||||
vllm_config.model_config.get_head_size.return_value = 64
|
||||
vllm_config.model_config.dtype = torch.float16
|
||||
vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
|
||||
speculative_config = MagicMock()
|
||||
speculative_config.num_speculative_tokens = 4
|
||||
vllm_config.speculative_config = speculative_config
|
||||
device = torch.device("cpu")
|
||||
|
||||
builder = AscendSFAMetadataBuilder(
|
||||
kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
|
||||
)
|
||||
|
||||
common_attn_metadata = MagicMock()
|
||||
common_attn_metadata.num_reqs = 10
|
||||
common_attn_metadata.num_actual_tokens = 100
|
||||
common_attn_metadata.query_start_loc = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
|
||||
common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
|
||||
common_attn_metadata.slot_mapping = torch.randn(100, 4, 1024)
|
||||
common_attn_metadata.seq_lens_cpu = torch.tensor([2] * 10)
|
||||
common_attn_metadata.positions = torch.randn(100)
|
||||
common_attn_metadata.attn_mask = None
|
||||
common_attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
|
||||
common_attn_metadata.block_table_tensor = torch.randn(100, 4)
|
||||
common_attn_metadata.cos = None
|
||||
common_attn_metadata.sin = None
|
||||
common_attn_metadata.num_input_tokens = 100
|
||||
|
||||
mock_get_cos_and_sin_mla.return_value = (torch.randn(100), torch.randn(100))
|
||||
|
||||
metadata = builder.build(
|
||||
common_prefix_len=10,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
assert isinstance(metadata, AscendSFAMetadata)
|
||||
assert metadata.num_actual_tokens == common_attn_metadata.num_actual_tokens
|
||||
assert metadata.slot_mapping.shape == (100, 4, 1024)
|
||||
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_tp_group")
|
||||
def test_dsa_cp_metadata_builder_masks_graph_padding(
|
||||
self,
|
||||
mock_get_tp_group,
|
||||
mock_get_cos_and_sin_mla,
|
||||
):
|
||||
# TP8, graph size 80 and MTP3 produce 20 four-token request slots. With
|
||||
# nine real requests, rank 6 splits a padded slot at its local boundary.
|
||||
tp_group = MagicMock()
|
||||
tp_group.world_size = 8
|
||||
tp_group.rank_in_group = 6
|
||||
mock_get_tp_group.return_value = tp_group
|
||||
mock_get_cos_and_sin_mla.return_value = (
|
||||
torch.zeros(80, 1, 1, 64),
|
||||
torch.zeros(80, 1, 1, 64),
|
||||
)
|
||||
|
||||
builder = AscendSFAMetadataBuilder.__new__(AscendSFAMetadataBuilder)
|
||||
builder.kernel_block_size = 128
|
||||
builder.model_config = MagicMock()
|
||||
builder.model_config.get_head_size.return_value = 64
|
||||
builder.attn_mask_builder = MagicMock()
|
||||
builder.enable_dsa_cp = True
|
||||
builder.actual_seq_lengths_query = torch.zeros(21, dtype=torch.int32)
|
||||
builder.actual_seq_lengths_key = torch.zeros(21, dtype=torch.int32)
|
||||
builder.spec_actual_seq_lengths_query = None
|
||||
builder.spec_actual_seq_lengths_key = None
|
||||
builder.metadata_cls = AscendSFAMetadata
|
||||
|
||||
common_attn_metadata = MagicMock()
|
||||
common_attn_metadata.num_reqs = 20
|
||||
common_attn_metadata.num_actual_tokens = 36
|
||||
common_attn_metadata.num_input_tokens = 80
|
||||
common_attn_metadata.query_start_loc = torch.arange(0, 81, 4, dtype=torch.int32)
|
||||
common_attn_metadata.seq_lens = torch.zeros(20, dtype=torch.int32)
|
||||
common_attn_metadata.seq_lens[:9] = torch.arange(128, 137, dtype=torch.int32)
|
||||
common_attn_metadata._seq_lens_cpu = common_attn_metadata.seq_lens.clone()
|
||||
common_attn_metadata.seq_lens_cpu = common_attn_metadata.seq_lens.clone()
|
||||
common_attn_metadata.block_table_tensor = torch.zeros(20, 1, dtype=torch.int32)
|
||||
common_attn_metadata.slot_mapping = torch.arange(80, dtype=torch.int64)
|
||||
common_attn_metadata.positions = torch.arange(80, dtype=torch.int64)
|
||||
common_attn_metadata.attn_state = AscendAttentionState.DecodeOnly
|
||||
common_attn_metadata.causal = True
|
||||
|
||||
metadata = builder._build(common_attn_metadata)
|
||||
|
||||
local_seq_lens = metadata.dsa_cp_context.actual_seq_lengths_key
|
||||
assert local_seq_lens[17].item() == 0
|
||||
assert torch.all(local_seq_lens >= 0)
|
||||
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config")
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
|
||||
@patch("vllm_ascend.attention.sfa_v1.enable_dsa_cp", return_value=False)
|
||||
@patch("vllm.distributed.parallel_state.get_tp_group")
|
||||
@patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
|
||||
def test_ascend_sfa_metadata_builder_build_for_graph_capture(
|
||||
self, mock_get_tp_group, mock_enable_dsa_cp, mock_get_cos_and_sin_mla, mock_get_current_vllm_config
|
||||
):
|
||||
cfg = MagicMock()
|
||||
cfg.model_config = MagicMock()
|
||||
cfg.model_config.hf_text_config = MagicMock()
|
||||
|
||||
mock_get_current_vllm_config.return_value = cfg
|
||||
|
||||
kv_cache_spec = MagicMock()
|
||||
kv_cache_spec.block_size = 128
|
||||
layer_names = ["layer1", "layer2"]
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.cache_config.block_size = 16
|
||||
vllm_config.scheduler_config.max_num_seqs = 16
|
||||
vllm_config.model_config.max_model_len = 1024
|
||||
vllm_config.model_config.get_head_size.return_value = 64
|
||||
vllm_config.model_config.dtype = torch.float16
|
||||
vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
|
||||
speculative_config = MagicMock()
|
||||
speculative_config.num_speculative_tokens = 4
|
||||
vllm_config.speculative_config = speculative_config
|
||||
device = torch.device("cpu")
|
||||
|
||||
builder = AscendSFAMetadataBuilder(
|
||||
kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
|
||||
)
|
||||
|
||||
common_attn_metadata = MagicMock()
|
||||
common_attn_metadata.num_reqs = 10
|
||||
common_attn_metadata.num_actual_tokens = 100
|
||||
common_attn_metadata.query_start_loc = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
|
||||
common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
|
||||
common_attn_metadata.slot_mapping = torch.randn(100, 4, 1024)
|
||||
common_attn_metadata.seq_lens_cpu = torch.tensor([2] * 10)
|
||||
common_attn_metadata.positions = torch.randn(100)
|
||||
common_attn_metadata.attn_mask = None
|
||||
common_attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
|
||||
common_attn_metadata.block_table_tensor = torch.randn(100, 4)
|
||||
common_attn_metadata.cos = None
|
||||
common_attn_metadata.sin = None
|
||||
common_attn_metadata.num_input_tokens = 100
|
||||
|
||||
mock_get_cos_and_sin_mla.return_value = (torch.randn(100), torch.randn(100))
|
||||
|
||||
attn_metadata = builder.build_for_graph_capture(
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
attn_state=AscendAttentionState.DecodeOnly,
|
||||
)
|
||||
|
||||
assert isinstance(attn_metadata, AscendSFAMetadata)
|
||||
assert attn_metadata.attn_state == AscendAttentionState.DecodeOnly
|
||||
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config")
|
||||
@patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
|
||||
@patch("vllm_ascend.attention.sfa_v1.enable_dsa_cp", return_value=False)
|
||||
@patch("torch.ops._C_ascend.store_kv_block_metadata", create=True)
|
||||
def test_ascend_sfa_metadata_builder_build_with_c8_reshape_optim(
|
||||
self,
|
||||
store_kv_block_metadata,
|
||||
mock_enable_dsa_cp,
|
||||
mock_get_cos_and_sin_mla,
|
||||
mock_get_current_vllm_config,
|
||||
):
|
||||
cfg = MagicMock()
|
||||
cfg.model_config = MagicMock()
|
||||
cfg.model_config.hf_text_config = MagicMock()
|
||||
|
||||
mock_get_current_vllm_config.return_value = cfg
|
||||
kv_cache_spec = MagicMock()
|
||||
kv_cache_spec.block_size = 128
|
||||
layer_names = ["layer1", "layer2"]
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.cache_config.block_size = 16
|
||||
vllm_config.scheduler_config.max_num_seqs = 16
|
||||
vllm_config.model_config.max_model_len = 1024
|
||||
vllm_config.model_config.get_head_size.return_value = 64
|
||||
vllm_config.model_config.dtype = torch.float16
|
||||
vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
|
||||
speculative_config = MagicMock()
|
||||
speculative_config.num_speculative_tokens = 4
|
||||
vllm_config.speculative_config = speculative_config
|
||||
device = torch.device("cpu")
|
||||
|
||||
builder = AscendSFAMetadataBuilder(
|
||||
kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
|
||||
)
|
||||
|
||||
common_attn_metadata = MagicMock()
|
||||
common_attn_metadata.num_reqs = 10
|
||||
common_attn_metadata.num_actual_tokens = 100
|
||||
common_attn_metadata.query_start_loc = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
|
||||
common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
|
||||
common_attn_metadata.slot_mapping = torch.randn(100, 4, 1024)
|
||||
common_attn_metadata.seq_lens_cpu = torch.tensor([2] * 10)
|
||||
common_attn_metadata.positions = torch.randn(100)
|
||||
common_attn_metadata.attn_mask = None
|
||||
common_attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
|
||||
common_attn_metadata.block_table_tensor = torch.randn(100, 4)
|
||||
common_attn_metadata.cos = None
|
||||
common_attn_metadata.sin = None
|
||||
common_attn_metadata.num_input_tokens = 100
|
||||
|
||||
mock_get_cos_and_sin_mla.return_value = (torch.randn(100), torch.randn(100))
|
||||
|
||||
with patch("vllm_ascend.attention.sfa_v1.get_ascend_config") as mock_get_ascend_config:
|
||||
mock_ascend_config = MagicMock()
|
||||
mock_ascend_config.c8_enable_reshape_optim = True
|
||||
mock_get_ascend_config.return_value = mock_ascend_config
|
||||
|
||||
metadata = builder.build(
|
||||
common_prefix_len=10,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
assert isinstance(metadata, AscendSFAMetadata)
|
||||
assert metadata.num_actual_tokens == common_attn_metadata.num_actual_tokens
|
||||
assert metadata.slot_mapping.shape == (100, 4, 1024)
|
||||
|
||||
store_kv_block_metadata.assert_called_once()
|
||||
actual_args, _ = store_kv_block_metadata.call_args
|
||||
assert torch.equal(actual_args[0], common_attn_metadata.slot_mapping)
|
||||
assert actual_args[4] == 128
|
||||
|
||||
assert metadata.block_size == 128
|
||||
assert metadata.group_len is actual_args[1]
|
||||
assert metadata.group_key_idx is actual_args[2]
|
||||
assert metadata.group_key_cache_idx is actual_args[3]
|
||||
487
tests/ut/attention/a2/test_sfa_v1_precision.py
Normal file
487
tests/ut/attention/a2/test_sfa_v1_precision.py
Normal file
@@ -0,0 +1,487 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_ascend.utils import enable_custom_op
|
||||
|
||||
enable_custom_op()
|
||||
|
||||
# Metrics log: survives vLLM stdout/stderr redirection; path via env.
|
||||
_METRICS_LOG_PATH = os.environ.get(
|
||||
"SFA_V1_PRECISION_METRICS_LOG",
|
||||
"/tmp/sfa_v1_precision_metrics.log",
|
||||
)
|
||||
with open(_METRICS_LOG_PATH, "w", encoding="utf-8"):
|
||||
pass
|
||||
|
||||
|
||||
def _emit_metric(line: str) -> None:
|
||||
"""Append ``line`` to the metrics log and echo it to stderr."""
|
||||
with open(_METRICS_LOG_PATH, "a", encoding="utf-8") as f:
|
||||
f.write(line + "\n")
|
||||
try:
|
||||
sys.stderr.write(line + "\n")
|
||||
sys.stderr.flush()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
if "torch_npu._inductor" not in sys.modules:
|
||||
sys.modules["torch_npu._inductor"] = MagicMock()
|
||||
|
||||
from vllm.forward_context import set_forward_context # noqa: E402
|
||||
|
||||
from tests.ut.attention.utils import ( # noqa: E402
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm_ascend.attention.sfa_v1 import AscendSFAImpl # noqa: E402
|
||||
|
||||
SPARSE_COUNT = 2048 # indexer_select_post_process (sfa_v1)
|
||||
|
||||
DEFAULT_RTOL = 1e-2
|
||||
DEFAULT_ATOL = 1e-2
|
||||
|
||||
# Signal-relative checks (per-element |err|/|ref| near zero ref is unstable).
|
||||
_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|
|
||||
|
||||
_BLOCK_SIZE = 128
|
||||
_TEST_NUM_HEADS = 8
|
||||
|
||||
BATCH_SPECS: dict[str, BatchSpec] = {
|
||||
"pure_decode_single": BatchSpec(
|
||||
seq_lens=[1024],
|
||||
query_lens=[1],
|
||||
name="pure_decode_single",
|
||||
),
|
||||
"pure_decode_small_batch": BatchSpec(
|
||||
seq_lens=[64, 128, 256, 512],
|
||||
query_lens=[1, 1, 1, 1],
|
||||
name="pure_decode_small_batch",
|
||||
),
|
||||
"pure_decode_large_batch": BatchSpec(
|
||||
seq_lens=[2048] * 16,
|
||||
query_lens=[1] * 16,
|
||||
name="pure_decode_large_batch",
|
||||
),
|
||||
"pure_prefill_single": BatchSpec(
|
||||
seq_lens=[256],
|
||||
query_lens=[256],
|
||||
name="pure_prefill_single",
|
||||
),
|
||||
"pure_prefill_small_batch": BatchSpec(
|
||||
seq_lens=[128, 256, 384],
|
||||
query_lens=[128, 256, 384],
|
||||
name="pure_prefill_small_batch",
|
||||
),
|
||||
"pure_prefill_with_context": BatchSpec(
|
||||
seq_lens=[512, 1024],
|
||||
query_lens=[128, 256],
|
||||
name="pure_prefill_with_context",
|
||||
),
|
||||
"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",
|
||||
),
|
||||
"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",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _validate_spec(spec: BatchSpec) -> None:
|
||||
"""Require seq_len <= SPARSE_COUNT so sparse matches dense reference."""
|
||||
for s, q in zip(spec.seq_lens, spec.query_lens):
|
||||
assert q <= s, f"query_len ({q}) must not exceed seq_len ({s})"
|
||||
assert s <= SPARSE_COUNT, (
|
||||
f"seq_len ({s}) must be <= SPARSE_COUNT ({SPARSE_COUNT}) so the "
|
||||
"sparse attention degenerates into dense attention for the "
|
||||
"reference comparison."
|
||||
)
|
||||
|
||||
|
||||
_VLLM_CONFIG_CACHE: dict = {}
|
||||
|
||||
|
||||
def _get_vllm_config(
|
||||
model: str,
|
||||
dtype: torch.dtype,
|
||||
*,
|
||||
max_model_len: int = 4096,
|
||||
tensor_parallel_size: int = 1,
|
||||
):
|
||||
"""Cached ``VllmConfig`` for DSA/SFA (fp8 quant stripped; heads capped for UT)."""
|
||||
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"
|
||||
cfg = create_vllm_config(
|
||||
model_name=model,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
max_model_len=max_model_len,
|
||||
dtype=dtype_str,
|
||||
block_size=_BLOCK_SIZE,
|
||||
num_gpu_blocks=4096,
|
||||
max_num_seqs=64,
|
||||
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": _TEST_NUM_HEADS,
|
||||
"num_key_value_heads": 1,
|
||||
},
|
||||
)
|
||||
_VLLM_CONFIG_CACHE[key] = cfg
|
||||
return cfg
|
||||
|
||||
|
||||
def _build_paged_kv_cache_from_metadata(
|
||||
common_attn_metadata,
|
||||
seq_lens: list[int],
|
||||
k_nope_contexts: list[torch.Tensor],
|
||||
k_rope_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_rope caches from metadata block layout."""
|
||||
blocks_per_seq = [(s + block_size - 1) // block_size for s in seq_lens]
|
||||
total_blocks = sum(blocks_per_seq) + 1
|
||||
|
||||
k_nope_cache = torch.zeros(total_blocks, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
|
||||
k_rope_cache = torch.zeros(total_blocks, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
block_table.zero_()
|
||||
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_rope_cache[block_id, :length, 0, :] = k_rope_contexts[b][tok_start:tok_end]
|
||||
next_block_id += 1
|
||||
|
||||
return k_nope_cache, k_rope_cache, block_table
|
||||
|
||||
|
||||
def _build_topk_indices(
|
||||
seq_lens: list[int],
|
||||
query_lens: list[int],
|
||||
sparse_count: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Causal top-k indices; shape ``(T, 1, sparse_count)`` int32, -1 pad."""
|
||||
num_tokens = sum(query_lens)
|
||||
topk = torch.full((num_tokens, 1, sparse_count), -1, dtype=torch.int32, device=device)
|
||||
|
||||
cum_q = 0
|
||||
for b, s_len in enumerate(seq_lens):
|
||||
q_len = query_lens[b]
|
||||
ctx_len = s_len - q_len
|
||||
for j in range(q_len):
|
||||
valid_end = ctx_len + j + 1
|
||||
topk[cum_q + j, 0, :valid_end] = torch.arange(valid_end, dtype=torch.int32, device=device)
|
||||
cum_q += q_len
|
||||
|
||||
return topk
|
||||
|
||||
|
||||
def _reference_sparse_attention(
|
||||
ql_nope: torch.Tensor,
|
||||
q_pe: torch.Tensor,
|
||||
k_nope_cache: torch.Tensor,
|
||||
k_rope_cache: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
seq_lens: list[int],
|
||||
query_lens: list[int],
|
||||
scale: float,
|
||||
block_size: int,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Fp32 dense MQA softmax baseline over causal prefix."""
|
||||
batch_size = len(seq_lens)
|
||||
outputs: list[torch.Tensor] = []
|
||||
|
||||
cum_q = 0
|
||||
for b in range(batch_size):
|
||||
s_len = seq_lens[b]
|
||||
q_len = query_lens[b]
|
||||
ctx_len = s_len - q_len
|
||||
|
||||
n_blocks = (s_len + block_size - 1) // block_size
|
||||
block_ids = block_table[b, :n_blocks].long()
|
||||
k_blocks = k_nope_cache[block_ids]
|
||||
k_rope_blocks = k_rope_cache[block_ids]
|
||||
|
||||
k_full = k_blocks.reshape(n_blocks * block_size, -1)[:s_len]
|
||||
k_rope_full = k_rope_blocks.reshape(n_blocks * block_size, -1)[:s_len]
|
||||
|
||||
K = torch.cat([k_full, k_rope_full], dim=-1).float()
|
||||
V = k_full.float()
|
||||
|
||||
for j in range(q_len):
|
||||
t = cum_q + j
|
||||
valid_end = ctx_len + j + 1
|
||||
|
||||
q_n = ql_nope[t].float()
|
||||
q_p = q_pe[t].float()
|
||||
Q = torch.cat([q_n, q_p], dim=-1)
|
||||
|
||||
K_b = K[:valid_end]
|
||||
V_b = V[:valid_end]
|
||||
|
||||
scores = (Q @ K_b.transpose(0, 1)) * scale
|
||||
attn = torch.softmax(scores, dim=-1)
|
||||
out = attn @ V_b
|
||||
outputs.append(out.to(out_dtype))
|
||||
|
||||
cum_q += q_len
|
||||
|
||||
return torch.stack(outputs, dim=0)
|
||||
|
||||
|
||||
def _run_sfa_kernel(
|
||||
ql_nope: torch.Tensor,
|
||||
q_pe: torch.Tensor,
|
||||
k_nope_cache: torch.Tensor,
|
||||
k_rope_cache: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
topk_indices: torch.Tensor,
|
||||
cum_query_lens: torch.Tensor,
|
||||
seq_lens_tensor: torch.Tensor,
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Call kernel via MagicMock self (only ``scale`` needed)."""
|
||||
fake_self = MagicMock()
|
||||
fake_self.scale = scale
|
||||
|
||||
fake_attn_metadata = MagicMock()
|
||||
fake_attn_metadata.block_table = block_table
|
||||
|
||||
return AscendSFAImpl._execute_sparse_flash_attention_process(
|
||||
fake_self,
|
||||
ql_nope,
|
||||
q_pe,
|
||||
(k_nope_cache, k_rope_cache),
|
||||
topk_indices,
|
||||
fake_attn_metadata,
|
||||
cum_query_lens,
|
||||
seq_lens_tensor,
|
||||
)
|
||||
|
||||
|
||||
def _run_precision_check(
|
||||
spec: BatchSpec,
|
||||
dtype: torch.dtype,
|
||||
vllm_config,
|
||||
*,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
torch.manual_seed(2026)
|
||||
_validate_spec(spec)
|
||||
|
||||
device = torch.device("npu")
|
||||
seq_lens = list(spec.seq_lens)
|
||||
query_lens = list(spec.query_lens)
|
||||
batch_size = spec.batch_size
|
||||
num_tokens = 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
|
||||
qk_rope_head_dim = hf_text.qk_rope_head_dim
|
||||
kv_lora_rank = hf_text.kv_lora_rank
|
||||
num_heads = hf_text.num_attention_heads
|
||||
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim
|
||||
scale = 1.0 / math.sqrt(head_dim)
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(spec, block_size=block_size, device=device)
|
||||
|
||||
k_nope_contexts = [torch.randn(s, kv_lora_rank, dtype=dtype, device=device) for s in seq_lens]
|
||||
k_rope_contexts = [torch.randn(s, qk_rope_head_dim, dtype=dtype, device=device) for s in seq_lens]
|
||||
|
||||
k_nope_cache, k_rope_cache, block_table = _build_paged_kv_cache_from_metadata(
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
seq_lens=seq_lens,
|
||||
k_nope_contexts=k_nope_contexts,
|
||||
k_rope_contexts=k_rope_contexts,
|
||||
block_size=block_size,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
ql_nope = torch.randn(num_tokens, num_heads, kv_lora_rank, dtype=dtype, device=device)
|
||||
q_pe = torch.randn(num_tokens, num_heads, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
topk_indices = _build_topk_indices(seq_lens, query_lens, SPARSE_COUNT, device)
|
||||
|
||||
cum_query_lens = torch.tensor(
|
||||
[sum(query_lens[: i + 1]) for i in range(batch_size)],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device)
|
||||
|
||||
with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
|
||||
backend_output = _run_sfa_kernel(
|
||||
ql_nope=ql_nope,
|
||||
q_pe=q_pe,
|
||||
k_nope_cache=k_nope_cache,
|
||||
k_rope_cache=k_rope_cache,
|
||||
block_table=block_table,
|
||||
topk_indices=topk_indices,
|
||||
cum_query_lens=cum_query_lens,
|
||||
seq_lens_tensor=seq_lens_tensor,
|
||||
scale=scale,
|
||||
)
|
||||
reference_output = _reference_sparse_attention(
|
||||
ql_nope=ql_nope,
|
||||
q_pe=q_pe,
|
||||
k_nope_cache=k_nope_cache,
|
||||
k_rope_cache=k_rope_cache,
|
||||
block_table=block_table,
|
||||
seq_lens=seq_lens,
|
||||
query_lens=query_lens,
|
||||
scale=scale,
|
||||
block_size=block_size,
|
||||
out_dtype=dtype,
|
||||
)
|
||||
|
||||
tag = f"{spec.name},tp={tensor_parallel_size}"
|
||||
assert backend_output.shape == reference_output.shape, (
|
||||
f"[{tag}] backend shape {tuple(backend_output.shape)} != reference shape {tuple(reference_output.shape)}"
|
||||
)
|
||||
assert backend_output.dtype == reference_output.dtype, (
|
||||
f"[{tag}] backend dtype {backend_output.dtype} != reference dtype {reference_output.dtype}"
|
||||
)
|
||||
assert torch.isfinite(backend_output).all(), f"[{tag}] sparse flash attention produced non-finite values"
|
||||
|
||||
torch.testing.assert_close(
|
||||
backend_output,
|
||||
reference_output,
|
||||
rtol=DEFAULT_RTOL,
|
||||
atol=DEFAULT_ATOL,
|
||||
msg=lambda m: f"[SFA:{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
|
||||
|
||||
_emit_metric(
|
||||
f"[SFA:{spec.name}] tp={tensor_parallel_size} dtype={dtype} "
|
||||
f"peak={peak:.4e} "
|
||||
f"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"[SFA:{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"[SFA:{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"[SFA:{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})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V3.2-Exp"])
|
||||
@pytest.mark.parametrize("batch_spec_name", list(BATCH_SPECS.keys()))
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
def test_sfa_sparse_flash_attention_precision(
|
||||
batch_spec_name: str,
|
||||
model: str,
|
||||
dtype: torch.dtype,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
"""SFA kernel vs fp32 dense MQA reference (decode, prefill, mixed, MTP)."""
|
||||
vllm_config = _get_vllm_config(model, dtype, tensor_parallel_size=tensor_parallel_size)
|
||||
_run_precision_check(
|
||||
BATCH_SPECS[batch_spec_name],
|
||||
dtype,
|
||||
vllm_config,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
205
tests/ut/attention/test_attention_fa3.py
Normal file
205
tests/ut/attention/test_attention_fa3.py
Normal file
@@ -0,0 +1,205 @@
|
||||
from importlib import import_module, util
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import torch_npu
|
||||
|
||||
|
||||
def _fa3_available() -> bool:
|
||||
try:
|
||||
if util.find_spec("flash_attn_npu_v3") is None:
|
||||
return False
|
||||
mod = import_module("flash_attn_npu_v3")
|
||||
return hasattr(mod, "flash_attn_with_kvcache")
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def ref_fused_infer_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
block_table,
|
||||
block_size,
|
||||
actual_seq_lengths_q,
|
||||
actual_seq_lengths_kv,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
scale,
|
||||
attn_mask,
|
||||
causal,
|
||||
):
|
||||
if not causal:
|
||||
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
block_table=block_table,
|
||||
input_layout="TND",
|
||||
block_size=block_size,
|
||||
actual_seq_lengths=actual_seq_lengths_q,
|
||||
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
||||
num_key_value_heads=num_kv_heads,
|
||||
num_heads=num_heads,
|
||||
scale=scale,
|
||||
sparse_mode=0,
|
||||
)
|
||||
else:
|
||||
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
atten_mask=attn_mask,
|
||||
block_table=block_table,
|
||||
input_layout="TND",
|
||||
block_size=block_size,
|
||||
actual_seq_lengths=actual_seq_lengths_q,
|
||||
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
||||
num_key_value_heads=num_kv_heads,
|
||||
num_heads=num_heads,
|
||||
scale=scale,
|
||||
sparse_mode=3,
|
||||
)
|
||||
|
||||
attn_output = attn_output.view(-1, num_heads, head_size)
|
||||
return attn_output
|
||||
|
||||
|
||||
test_cases = [
|
||||
# (data_type, batch_size, num_heads, kv_heads, q_seqlen, kv_seqlen, head_size, block_size, is_causal)
|
||||
(torch.bfloat16, 1, 1, 1, 1024, 1024, 128, 128, False),
|
||||
(torch.bfloat16, 5, 4, 1, 1024, 1024, 128, 128, True),
|
||||
(torch.float16, 7, 16, 8, 512, 512, 128, 128, False),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _fa3_available(), reason="flash_attn_npu_v3 is not installed")
|
||||
@pytest.mark.parametrize(
|
||||
"data_type, batch_size, num_heads, kv_heads, q_seqlen, kv_seqlen, head_size, block_size, is_causal", test_cases
|
||||
)
|
||||
def test_fa_custom_ops_tnd(
|
||||
data_type, batch_size, num_heads, kv_heads, q_seqlen, kv_seqlen, head_size, block_size, is_causal
|
||||
):
|
||||
q_min_range = -1.0
|
||||
q_max_range = 1.0
|
||||
kv_min_range = -1.0
|
||||
kv_max_range = 1.0
|
||||
block_size = 128
|
||||
num_blocks = 64
|
||||
|
||||
q_sequences = sorted(
|
||||
torch.randint(low=1, high=q_seqlen + 1, size=(batch_size,)).tolist(), reverse=False
|
||||
) # actual_seq_lengths in fia need in ascending order
|
||||
kv_sequences = [torch.randint(low=q, high=kv_seqlen + 1, size=(1,)).item() for q in q_sequences]
|
||||
|
||||
t_q_sum = sum(q_sequences)
|
||||
|
||||
query = (q_min_range + (q_max_range - q_min_range) * torch.rand(t_q_sum, num_heads, head_size)).to(data_type).npu()
|
||||
|
||||
key_cache = None
|
||||
value_cache = None
|
||||
block_tables = []
|
||||
|
||||
key_cache = (
|
||||
(kv_min_range + (kv_max_range - kv_min_range) * torch.rand(num_blocks, block_size, kv_heads, head_size))
|
||||
.to(data_type)
|
||||
.npu()
|
||||
)
|
||||
value_cache = (
|
||||
(kv_min_range + (kv_max_range - kv_min_range) * torch.rand(num_blocks, block_size, kv_heads, head_size))
|
||||
.to(data_type)
|
||||
.npu()
|
||||
)
|
||||
max_num_blocks_per_seq = (kv_seqlen + block_size - 1) // block_size
|
||||
for i in range(batch_size):
|
||||
block_table = [max_num_blocks_per_seq * i + j for j in range(max_num_blocks_per_seq)]
|
||||
block_tables.append(block_table)
|
||||
block_tables = torch.tensor(block_tables, dtype=torch.int32).npu()
|
||||
|
||||
q_seqlen_list = q_sequences
|
||||
kv_seqlen_list = kv_sequences
|
||||
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
window_size_left = -1
|
||||
window_size_right = -1
|
||||
is_rotary_interleaved = False
|
||||
num_splits = 0
|
||||
kv_seqlen_list = torch.tensor(kv_seqlen_list, dtype=torch.int32).npu()
|
||||
rotary_cos = None
|
||||
rotary_sin = None
|
||||
cache_batch_idx = None
|
||||
leftpad_k = None
|
||||
new_q_seqlen_list = None
|
||||
|
||||
new_q_seqlen_list = [0]
|
||||
pre_seq_sum = 0
|
||||
for i in range(batch_size):
|
||||
pre_seq_sum += q_seqlen_list[i]
|
||||
new_q_seqlen_list.append(pre_seq_sum)
|
||||
new_q_seqlen_list = torch.tensor(new_q_seqlen_list, dtype=torch.int32).npu()
|
||||
|
||||
from flash_attn_npu_v3 import flash_attn_with_kvcache # type: ignore[import-not-found]
|
||||
|
||||
out_out = flash_attn_with_kvcache(
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
rotary_cos=rotary_cos,
|
||||
rotary_sin=rotary_sin,
|
||||
cache_seqlens=kv_seqlen_list,
|
||||
cache_batch_idx=cache_batch_idx,
|
||||
cache_leftpad=leftpad_k,
|
||||
page_table=block_tables,
|
||||
cu_seqlens_q=new_q_seqlen_list,
|
||||
cu_seqlens_k_new=None,
|
||||
max_seqlen_q=q_seqlen,
|
||||
rotary_seqlens=None,
|
||||
q_descale=None,
|
||||
k_descale=None,
|
||||
v_descale=None,
|
||||
softmax_scale=None,
|
||||
causal=is_causal,
|
||||
window_size=[window_size_left, window_size_right],
|
||||
attention_chunk=0,
|
||||
softcap=0.0,
|
||||
rotary_interleaved=is_rotary_interleaved,
|
||||
scheduler_metadata=None,
|
||||
num_splits=num_splits,
|
||||
pack_gqa=None,
|
||||
sm_margin=0,
|
||||
return_softmax_lse=False,
|
||||
)
|
||||
|
||||
ref_out = torch.empty((t_q_sum, num_heads, head_size), dtype=data_type)
|
||||
|
||||
if is_causal:
|
||||
attn_mask = torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8).to(query.device)
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
q_cumsum = torch.tensor(np.cumsum(q_sequences), dtype=torch.int32, device=query.device)
|
||||
|
||||
ref_out = ref_fused_infer_attention(
|
||||
query,
|
||||
key_cache.permute(0, 2, 1, 3).contiguous(), # [num_blocks, num_kv_heads, block_size, head_size]
|
||||
value_cache.permute(0, 2, 1, 3).contiguous(),
|
||||
block_tables,
|
||||
block_size,
|
||||
q_cumsum,
|
||||
kv_seqlen_list,
|
||||
num_heads,
|
||||
kv_heads,
|
||||
head_size,
|
||||
scale,
|
||||
attn_mask,
|
||||
is_causal,
|
||||
)
|
||||
|
||||
rtol = 1e-2
|
||||
atol = 1e-2
|
||||
torch.testing.assert_close(out_out.cpu(), ref_out.cpu(), rtol=rtol, atol=atol)
|
||||
@@ -1,5 +1,5 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
@@ -20,114 +20,29 @@ from vllm_ascend.attention.attention_mask import AttentionMaskBuilder
|
||||
|
||||
|
||||
class TestAttentionMaskBuilder(TestBase):
|
||||
|
||||
def test_init_attention_mask_builder(self):
|
||||
# generate attention_mask_builder with float16
|
||||
attention_mask_builder = AttentionMaskBuilder(max_seq_len=1024,
|
||||
dtype=torch.float16)
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 1024)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.dtype,
|
||||
torch.float16)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
|
||||
(1024, 1024))
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
|
||||
torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
|
||||
# generate attention_mask_builder with bfloat16
|
||||
attention_mask_builder = AttentionMaskBuilder(max_seq_len=2048,
|
||||
dtype=torch.bfloat16)
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 2048)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.dtype,
|
||||
torch.bfloat16)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
|
||||
(2048, 2048))
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
|
||||
torch.tensor(1, dtype=torch.bfloat16))
|
||||
|
||||
def test_get_mask_scale_factor(self):
|
||||
# supported data types
|
||||
self.assertEqual(
|
||||
AttentionMaskBuilder.get_mask_scale_factor(torch.float16), 1)
|
||||
self.assertEqual(
|
||||
AttentionMaskBuilder.get_mask_scale_factor(torch.bfloat16), -10000)
|
||||
# mask_scale_factor now only supports data types: torch.float16 and torch.bfloat16
|
||||
# Otherwise raise ValueError
|
||||
with self.assertRaises(ValueError):
|
||||
AttentionMaskBuilder.get_mask_scale_factor(torch.int8)
|
||||
|
||||
def test_get_attn_mask(self):
|
||||
# if the len is less than max_seq_len, the attn_mask_cache will not be updated
|
||||
attention_mask_builder = AttentionMaskBuilder(max_seq_len=1024,
|
||||
dtype=torch.float16)
|
||||
attn_mask = attention_mask_builder.get_attn_mask(
|
||||
max_seq_len=512, dtype=torch.float16, device=torch.device("cpu"))
|
||||
attention_mask_builder = AttentionMaskBuilder(torch.device("cpu"))
|
||||
attn_mask = attention_mask_builder.get_attn_mask(max_seq_len=512, dtype=torch.float16)
|
||||
self.assertEqual(attn_mask.shape, (512, 512))
|
||||
self.assertEqual(attn_mask[0][-1],
|
||||
torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 1024)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
|
||||
(1024, 1024))
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
|
||||
torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
self.assertEqual(attn_mask[0][-1], torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 512)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.shape, (512, 512))
|
||||
self.assertEqual(
|
||||
attention_mask_builder.attn_mask_cache[0][-1], torch.tensor(float("-inf"), dtype=torch.float16)
|
||||
)
|
||||
|
||||
# if the len is greater than max_seq_len, the attn_mask_cache will be updated
|
||||
attn_mask = attention_mask_builder.get_attn_mask(
|
||||
max_seq_len=2048, dtype=torch.float16, device=torch.device("cpu"))
|
||||
attn_mask = attention_mask_builder.get_attn_mask(max_seq_len=2048, dtype=torch.float16)
|
||||
self.assertEqual(attn_mask.shape, (2048, 2048))
|
||||
self.assertEqual(attn_mask[0][-1],
|
||||
torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
self.assertEqual(attn_mask[0][-1], torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 2048)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
|
||||
(2048, 2048))
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
|
||||
torch.tensor(float("-inf"), dtype=torch.float16))
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache.shape, (2048, 2048))
|
||||
self.assertEqual(
|
||||
attention_mask_builder.attn_mask_cache[0][-1], torch.tensor(float("-inf"), dtype=torch.float16)
|
||||
)
|
||||
|
||||
def test_get_splitfuse_attn_mask(self):
|
||||
attention_mask_builder = AttentionMaskBuilder(max_seq_len=1024,
|
||||
dtype=torch.float16)
|
||||
attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
|
||||
seq_lens=torch.tensor([10, 20, 100]),
|
||||
position=torch.tensor([7, 8, 9, 18, 19, 99]),
|
||||
dtype=torch.float16,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
self.assertEqual(attn_mask.shape, (6, 100))
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 1024)
|
||||
|
||||
attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
|
||||
seq_lens=torch.tensor([10, 3000, 2000]),
|
||||
position=torch.tensor([7, 8, 9, 2999, 1999]),
|
||||
dtype=torch.float16,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
self.assertEqual(attn_mask.shape, (5, 3000))
|
||||
self.assertEqual(attention_mask_builder._seq_len_cached, 3000)
|
||||
|
||||
# splitfuse_attn_mask now only supports data types: torch.float16 and torch.bfloat16
|
||||
# otherwise raise ValueError
|
||||
with self.assertRaises(ValueError):
|
||||
attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
|
||||
seq_lens=torch.tensor([10, 20, 100]),
|
||||
position=torch.tensor([7, 8, 9, 18, 19, 99]),
|
||||
dtype=torch.int8,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
def test_mask_value_cleanliness(self):
|
||||
attention_mask_builder = AttentionMaskBuilder(max_seq_len=6,
|
||||
dtype=torch.bfloat16)
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache[-2][-1],
|
||||
torch.tensor(1, dtype=torch.bfloat16))
|
||||
|
||||
attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
|
||||
seq_lens=torch.tensor([6]),
|
||||
position=torch.tensor([3, 4, 5]),
|
||||
dtype=torch.bfloat16,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
self.assertEqual(
|
||||
attn_mask[-2][-1],
|
||||
torch.tensor(-10000, dtype=torch.bfloat16,
|
||||
device=attn_mask.device))
|
||||
self.assertEqual(attention_mask_builder.attn_mask_cache[-2][-1],
|
||||
torch.tensor(1, dtype=torch.bfloat16))
|
||||
attention_mask_builder = AttentionMaskBuilder(torch.device("cpu"))
|
||||
attn_mask = attention_mask_builder.get_splitfuse_attn_mask()
|
||||
self.assertEqual(attn_mask.shape, (2048, 2048))
|
||||
|
||||
1724
tests/ut/attention/test_sfa_cp.py
Normal file
1724
tests/ut/attention/test_sfa_cp.py
Normal file
File diff suppressed because it is too large
Load Diff
103
tests/ut/attention/test_sfa_o_proj_tp.py
Normal file
103
tests/ut/attention/test_sfa_o_proj_tp.py
Normal file
@@ -0,0 +1,103 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
|
||||
if "torch_npu._inductor" not in sys.modules:
|
||||
sys.modules["torch_npu._inductor"] = MagicMock()
|
||||
|
||||
from vllm_ascend.attention.sfa_v1 import AscendSFAImpl
|
||||
|
||||
|
||||
class TestAscendSFAOProjTPParams(TestBase):
|
||||
class _OProj(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(torch.randn(4, 3), requires_grad=False)
|
||||
self.aclnn_input_scale = torch.nn.Parameter(torch.randn(3), requires_grad=False)
|
||||
self.weight_scale_second = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
|
||||
self.weight_scale_second.input_dim = 1
|
||||
self.weight_offset_second = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
|
||||
self.weight_offset_second.input_dim = 1
|
||||
self.extra_input_scale = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
|
||||
self.extra_input_scale.input_dim = 1
|
||||
self.weight_scale = torch.nn.Parameter(torch.randn(4), requires_grad=False)
|
||||
|
||||
def setUp(self):
|
||||
AscendSFAImpl.o_proj_full_pools.clear()
|
||||
|
||||
def _make_impl(self):
|
||||
impl = AscendSFAImpl.__new__(AscendSFAImpl)
|
||||
impl.tp_size = 2
|
||||
impl.o_proj = self._OProj()
|
||||
impl._is_o_proj_unquantized = lambda: False
|
||||
return impl
|
||||
|
||||
def test_o_proj_tp_params_alias_original_storage(self):
|
||||
impl = self._make_impl()
|
||||
o_proj = impl.o_proj
|
||||
|
||||
impl._init_o_proj_tp_full_params()
|
||||
|
||||
self.assertEqual(impl.o_proj_tp_weight.data_ptr(), o_proj.weight.data_ptr())
|
||||
self.assertEqual(
|
||||
impl.o_proj_tp_aclnn_input_params["aclnn_input_scale"].data_ptr(),
|
||||
o_proj.aclnn_input_scale.data_ptr(),
|
||||
)
|
||||
self.assertEqual(
|
||||
impl.o_proj_tp_input_sharded_quant_params["weight_scale_second"].data_ptr(),
|
||||
o_proj.weight_scale_second.data_ptr(),
|
||||
)
|
||||
self.assertEqual(
|
||||
impl.o_proj_tp_input_sharded_quant_params["weight_offset_second"].data_ptr(),
|
||||
o_proj.weight_offset_second.data_ptr(),
|
||||
)
|
||||
self.assertEqual(
|
||||
impl.o_proj_tp_input_sharded_quant_params["extra_input_scale"].data_ptr(),
|
||||
o_proj.extra_input_scale.data_ptr(),
|
||||
)
|
||||
self.assertNotIn("weight_scale", impl.o_proj_tp_input_sharded_quant_params)
|
||||
|
||||
def test_o_proj_full_weight_forward_restores_tp_storage(self):
|
||||
impl = self._make_impl()
|
||||
impl._init_o_proj_tp_full_params()
|
||||
original_weight_ptr = impl.o_proj.weight.data_ptr()
|
||||
original_scale_ptr = impl.o_proj.weight_scale_second.data_ptr()
|
||||
full_weight_ptr = impl.o_proj_full_pool.data_ptr()
|
||||
full_scale_ptr = impl.o_proj_full_input_sharded_quant_params["weight_scale_second"].data_ptr()
|
||||
|
||||
def _apply_with_full_weight(_attn_output):
|
||||
self.assertEqual(impl.o_proj.weight.data_ptr(), full_weight_ptr)
|
||||
self.assertEqual(impl.o_proj.weight_scale_second.data_ptr(), full_scale_ptr)
|
||||
return torch.ones(2, 4)
|
||||
|
||||
impl._apply_o_proj_full_weight = MagicMock(side_effect=_apply_with_full_weight)
|
||||
|
||||
output, require_o_proj_forward = impl._handle_o_proj_weight_switch_and_forward(
|
||||
attn_output=torch.randn(2, 3),
|
||||
output=torch.empty(2, 4),
|
||||
o_proj_full_handle=None,
|
||||
o_proj_full_param_handles=[],
|
||||
should_shard_weight=True,
|
||||
)
|
||||
|
||||
self.assertEqual(impl.o_proj.weight.data_ptr(), original_weight_ptr)
|
||||
self.assertEqual(impl.o_proj.weight_scale_second.data_ptr(), original_scale_ptr)
|
||||
self.assertFalse(require_o_proj_forward)
|
||||
self.assertTrue(torch.equal(output, torch.ones(2, 4)))
|
||||
341
tests/ut/attention/utils.py
Normal file
341
tests/ut/attention/utils.py
Normal file
@@ -0,0 +1,341 @@
|
||||
from dataclasses import dataclass
|
||||
from functools import wraps
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
CompilationConfig,
|
||||
DeviceConfig,
|
||||
LoadConfig,
|
||||
ModelConfig,
|
||||
ParallelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.config.model import ModelDType
|
||||
from vllm.distributed.parallel_state import all_gather_fake
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
|
||||
|
||||
|
||||
def patch_distributed_groups(dcp_size=1, dcp_rank=0, pcp_size=1, pcp_rank=0, needs_mocks=True):
|
||||
"""
|
||||
Decorator to patch common distributed group mocks with configuration
|
||||
|
||||
Args:
|
||||
dcp_size: DCP world size (default: 1)
|
||||
dcp_rank: DCP rank (default: 0)
|
||||
pcp_size: PCP world size (default: 1)
|
||||
pcp_rank: PCP rank (default: 0)
|
||||
needs_mocks: Whether to pass mock objects as the first arguments
|
||||
after 'self' to the decorated function.
|
||||
If True, the decorated function receives:
|
||||
func(self, mock_all_to_all_single, mock_dcp, mock_pcp, *args, **kwargs)
|
||||
If False, mocks are not passed and function receives:
|
||||
func(self, *args, **kwargs)
|
||||
(default: True)
|
||||
"""
|
||||
|
||||
def decorator(func):
|
||||
@wraps(func)
|
||||
@patch("torch.distributed.all_to_all_single")
|
||||
@patch("vllm.distributed.parallel_state._PCP")
|
||||
@patch("vllm.distributed.parallel_state._DCP")
|
||||
def wrapper(self, mock_dcp, mock_pcp, mock_all_to_all_single, *args, **kwargs):
|
||||
mock_dcp.rank_in_group = dcp_rank
|
||||
mock_dcp.world_size = dcp_size
|
||||
mock_dcp.device_group = MagicMock()
|
||||
|
||||
mock_dcp.all_gather = MagicMock()
|
||||
mock_dcp.all_gather.side_effect = lambda input_, dim: all_gather_fake(
|
||||
input_, dim, mock_dcp.world_size, "mock_dcp_group"
|
||||
)
|
||||
|
||||
mock_pcp.rank_in_group = pcp_rank
|
||||
mock_pcp.world_size = pcp_size
|
||||
mock_pcp.device_group = MagicMock()
|
||||
|
||||
mock_pcp.all_gather = MagicMock()
|
||||
mock_pcp.all_gather.side_effect = lambda input_, dim: all_gather_fake(
|
||||
input_, dim, mock_pcp.world_size, "mock_pcp_group"
|
||||
)
|
||||
|
||||
mock_all_to_all_single.side_effect = lambda output, input, *a, **kw: output.copy_(input)
|
||||
|
||||
if needs_mocks:
|
||||
return func(self, mock_all_to_all_single, mock_dcp, mock_pcp, *args, **kwargs)
|
||||
else:
|
||||
return func(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchSpec:
|
||||
"""Specification for a batch configuration (workload shape only)."""
|
||||
|
||||
seq_lens: list[int]
|
||||
query_lens: list[int]
|
||||
|
||||
name: str = "unnamed"
|
||||
|
||||
@property
|
||||
def batch_size(self):
|
||||
return len(self.seq_lens)
|
||||
|
||||
def __post_init__(self):
|
||||
assert len(self.seq_lens) == len(self.query_lens)
|
||||
|
||||
def compute_num_tokens(self):
|
||||
return sum(self.query_lens)
|
||||
|
||||
|
||||
def create_common_attn_metadata(
|
||||
batch_spec: BatchSpec,
|
||||
block_size: int,
|
||||
device: torch.device,
|
||||
max_block_idx: int = 1000,
|
||||
arange_block_indices: bool = True,
|
||||
) -> AscendCommonAttentionMetadata:
|
||||
"""Create CommonAttentionMetadata from a BatchSpec and ModelParams."""
|
||||
# Create query start locations
|
||||
query_start_loc = torch.zeros(batch_spec.batch_size + 1, dtype=torch.int32, device=device)
|
||||
query_start_loc[1:] = torch.tensor(batch_spec.query_lens, dtype=torch.int32, device=device).cumsum(0)
|
||||
query_start_loc_cpu = query_start_loc.cpu()
|
||||
num_tokens = batch_spec.compute_num_tokens()
|
||||
# Create sequence lengths
|
||||
seq_lens = torch.tensor(batch_spec.seq_lens, dtype=torch.int32, device=device)
|
||||
seq_lens_cpu = seq_lens.cpu()
|
||||
max_seq_len = int(seq_lens_cpu.max())
|
||||
|
||||
# Create computed tokens (context length for each sequence)
|
||||
context_lens = [batch_spec.seq_lens[i] - batch_spec.query_lens[i] for i in range(batch_spec.batch_size)]
|
||||
num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
|
||||
|
||||
# Create block table and slot mapping
|
||||
max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
|
||||
num_blocks = batch_spec.batch_size * max_blocks
|
||||
block_table_tensor = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
|
||||
batch_spec.batch_size, max_blocks
|
||||
)
|
||||
slot_mapping = torch.arange(num_tokens, dtype=torch.int32, device=device).view(num_tokens)
|
||||
|
||||
# Calculate max query length
|
||||
max_query_len = max(batch_spec.query_lens)
|
||||
|
||||
# Create positions tensor
|
||||
positions = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
return AscendCommonAttentionMetadata(
|
||||
query_start_loc=query_start_loc,
|
||||
query_start_loc_cpu=query_start_loc_cpu,
|
||||
seq_lens=seq_lens,
|
||||
seq_lens_cpu=seq_lens_cpu,
|
||||
_num_computed_tokens_cpu=num_computed_tokens_cpu,
|
||||
num_reqs=batch_spec.batch_size,
|
||||
num_actual_tokens=num_tokens,
|
||||
max_query_len=max_query_len,
|
||||
max_seq_len=max_seq_len,
|
||||
block_table_tensor=block_table_tensor,
|
||||
slot_mapping=slot_mapping,
|
||||
causal=True,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
|
||||
def create_vllm_config(
|
||||
model_name: str = "meta-llama/Meta-Llama-3-8B",
|
||||
tensor_parallel_size: int = 1,
|
||||
max_model_len: int = 1024,
|
||||
dtype: ModelDType | torch.dtype = "auto",
|
||||
num_gpu_blocks: int = 1000,
|
||||
block_size: int = 16,
|
||||
max_num_seqs: int = 256,
|
||||
max_num_batched_tokens: int = 8192,
|
||||
enable_chunked_prefill: bool = True,
|
||||
add_mock_model_methods: bool = True,
|
||||
hf_config_override: dict | None = None,
|
||||
hf_overrides: dict | None = None,
|
||||
) -> VllmConfig:
|
||||
"""Create a VllmConfig for testing with reasonable defaults.
|
||||
|
||||
``hf_overrides`` is forwarded to ``ModelConfig.__init__`` (vLLM-native
|
||||
mechanism); this is required for fp8-quantized DSA models such as
|
||||
``deepseek-ai/DeepSeek-V3.2-Exp`` whose ``quantization_config`` would
|
||||
otherwise be rejected by Ascend's ``ModelConfig`` validator.
|
||||
"""
|
||||
|
||||
model_config = ModelConfig(
|
||||
model=model_name,
|
||||
tokenizer=model_name,
|
||||
trust_remote_code=False,
|
||||
dtype=dtype,
|
||||
seed=0,
|
||||
max_model_len=max_model_len,
|
||||
hf_overrides=hf_overrides or {},
|
||||
)
|
||||
|
||||
cache_config = CacheConfig(
|
||||
block_size=block_size,
|
||||
cache_dtype="auto",
|
||||
)
|
||||
# Set cache blocks for testing
|
||||
# (these may be set during initialization normally)
|
||||
cache_config.num_gpu_blocks = num_gpu_blocks
|
||||
cache_config.num_cpu_blocks = 0
|
||||
|
||||
parallel_config = ParallelConfig(
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=max_num_seqs,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
enable_chunked_prefill=enable_chunked_prefill,
|
||||
max_model_len=model_config.max_model_len,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
)
|
||||
|
||||
device_config = DeviceConfig()
|
||||
load_config = LoadConfig()
|
||||
compilation_config = CompilationConfig()
|
||||
|
||||
if add_mock_model_methods:
|
||||
# Add mock methods to satisfy backends that need them
|
||||
# This is a workaround because tests don't build full, real models,
|
||||
# but some backends expect to query the model for layer-specific
|
||||
# parameters
|
||||
import types
|
||||
|
||||
model_config.get_num_layers = types.MethodType(lambda self: 1, model_config)
|
||||
model_config.get_sliding_window_for_layer = types.MethodType(lambda self, i: None, model_config)
|
||||
model_config.get_logits_soft_cap_for_layer = types.MethodType(lambda self, i: 0.0, model_config)
|
||||
model_config.get_sm_scale_for_layer = types.MethodType(
|
||||
lambda self, i: 1.0 / model_config.get_head_size() ** 0.5, model_config
|
||||
)
|
||||
|
||||
if hf_config_override:
|
||||
# Apply the override to BOTH ``hf_config`` and ``hf_text_config`` so
|
||||
# the attribute is visible to backends regardless of which view they
|
||||
# reach for. For text-only models these usually point to the same
|
||||
# object; for multimodal models ``hf_text_config`` may be different.
|
||||
for k, v in hf_config_override.items():
|
||||
setattr(model_config.hf_config, k, v)
|
||||
if model_config.hf_text_config is not model_config.hf_config:
|
||||
setattr(model_config.hf_text_config, k, v)
|
||||
|
||||
return VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
parallel_config=parallel_config,
|
||||
scheduler_config=scheduler_config,
|
||||
device_config=device_config,
|
||||
load_config=load_config,
|
||||
compilation_config=compilation_config,
|
||||
)
|
||||
|
||||
|
||||
def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
|
||||
"""Create a FullAttentionSpec from ModelParams only."""
|
||||
return FullAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
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=vllm_config.model_config.dtype,
|
||||
sliding_window=vllm_config.model_config.get_sliding_window(),
|
||||
)
|
||||
|
||||
|
||||
def create_and_prepopulate_kv_cache(
|
||||
k_contexts: list[torch.Tensor],
|
||||
v_contexts: list[torch.Tensor],
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
num_blocks: int,
|
||||
common_attn_metadata: AscendCommonAttentionMetadata,
|
||||
randomize_blocks: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Create and prepopulate a KV cache with context data.
|
||||
|
||||
Args:
|
||||
k_contexts: List of key context tensors for each sequence
|
||||
v_contexts: List of value context tensors for each sequence
|
||||
seq_lens: List of sequence lengths
|
||||
block_size: Size of each block
|
||||
num_kv_heads: Number of KV heads
|
||||
head_size: Size of each head
|
||||
dtype: Data type for the cache
|
||||
device: Device to create the cache on
|
||||
num_blocks: Total number of blocks in the cache
|
||||
block_table: Block table tensor to populate
|
||||
randomize_blocks: Whether to randomly permute blocks
|
||||
or use sequential order
|
||||
|
||||
Returns:
|
||||
Tuple of (kv_cache, updated_block_table)
|
||||
"""
|
||||
batch_size = len(k_contexts)
|
||||
seq_lens = common_attn_metadata.seq_lens.cpu()
|
||||
query_lens = common_attn_metadata.query_start_loc_cpu[1:] - common_attn_metadata.query_start_loc_cpu[:-1]
|
||||
context_lens = seq_lens - query_lens
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
# Create KV cache
|
||||
kv_cache = torch.zeros(2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)
|
||||
|
||||
# Populate the cache with the context tokens
|
||||
# Start from block_id=0
|
||||
start_block_idx = 0
|
||||
for i in range(batch_size):
|
||||
k_context, v_context = k_contexts[i], v_contexts[i]
|
||||
start = start_block_idx * block_size
|
||||
end = start + k_context.shape[0]
|
||||
kv_cache_flat[0, start:end, ...] = k_context
|
||||
kv_cache_flat[1, start:end, ...] = v_context
|
||||
|
||||
# Stay block aligned and allocate enough blocks for the new tokens
|
||||
start_block_idx += cdiv(int(seq_lens[i]), block_size)
|
||||
|
||||
blocks_end = start_block_idx
|
||||
|
||||
# Permute the context blocks
|
||||
if randomize_blocks:
|
||||
# Random permutation starting from block 0
|
||||
perm = torch.randperm(blocks_end)
|
||||
else:
|
||||
# Sequential order starting from block 0
|
||||
perm = torch.arange(blocks_end)
|
||||
inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
|
||||
inv_perm = torch.argsort(perm)
|
||||
kv_cache[:, :blocks_end, ...] = kv_cache[:, perm, ...]
|
||||
# Construct the right block table
|
||||
# Start from block_id=0
|
||||
start_block_idx = 0
|
||||
for i in range(batch_size):
|
||||
num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
|
||||
start = start_block_idx
|
||||
end = start + num_blocks_for_seq
|
||||
block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
|
||||
start_block_idx += num_blocks_for_seq
|
||||
|
||||
# Create a realistic slot mapping that corresponds to the block table
|
||||
for i in range(batch_size):
|
||||
token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
|
||||
block_indices = token_offsets // block_size
|
||||
token_inter_block_offsets = token_offsets % block_size
|
||||
start = common_attn_metadata.query_start_loc_cpu[i]
|
||||
end = common_attn_metadata.query_start_loc_cpu[i + 1]
|
||||
slot_mapping[start:end] = block_table[i, block_indices] * block_size + token_inter_block_offsets.to(device).to(
|
||||
torch.int32
|
||||
)
|
||||
|
||||
return kv_cache
|
||||
Reference in New Issue
Block a user