import os from unittest.mock import MagicMock, patch import torch from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig from vllm.distributed.parallel_state import GroupCoordinator from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod 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.mla_v1 import ( AscendMLABackend, AscendMLADecodeMetadata, AscendMLAImpl, AscendMLAMetadata, AscendMLAMetadataBuilder, AscendMLAPrefillMetadata, ChunkedContextMetadata, DecodeMLAPreprocessResult, PrefillMLAPreprocessResult, ) from vllm_ascend.attention.utils import AscendCommonAttentionMetadata class TestAscendMLABackend(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(AscendMLABackend.get_name(), "ASCEND_MLA") def test_get_builder_cls(self): self.assertEqual(AscendMLABackend.get_builder_cls(), AscendMLAMetadataBuilder) def test_get_kv_cache_shape(self): result = AscendMLABackend.get_kv_cache_shape(2, 4, 8, 128) self.assertEqual(result, (2, 4, 8, 128)) def test_get_impl_cls(self): result = AscendMLABackend.get_impl_cls() self.assertEqual(result, AscendMLAImpl) def test_get_supported_kernel_block_sizes(self): result = AscendMLABackend.get_supported_kernel_block_sizes() self.assertEqual(result, [128]) @patch("vllm_ascend.attention.mla_v1.enable_cp") def test_get_builder_cls_with_cp(self, mock_enable_cp): mock_enable_cp.return_value = True builder_cls = AscendMLABackend.get_builder_cls() self.assertIsNotNone(builder_cls) @patch("vllm_ascend.attention.mla_v1.enable_cp") def test_get_impl_cls_with_cp(self, mock_enable_cp): mock_enable_cp.return_value = True impl_cls = AscendMLABackend.get_impl_cls() self.assertIsNotNone(impl_cls) class TestDecodeMLAPreprocessResult(TestBase): def test_decode_mla_preprocess_result_default(self): result = DecodeMLAPreprocessResult() self.assertIsNone(result.ql_nope) self.assertIsNone(result.q_pe) self.assertIsNone(result.k_nope) self.assertIsNone(result.k_pe) self.assertIsNone(result.decode_q_wo_k_up) self.assertIsNone(result.dequant_scale_q_nope) def test_decode_mla_preprocess_result_with_values(self): ql_nope = torch.randn(2, 4, 8) q_pe = torch.randn(2, 4, 8) k_nope = torch.randn(2, 4, 8) k_pe = torch.randn(2, 4, 8) decode_q_wo_k_up = torch.randn(2, 4, 8) dequant_scale_q_nope = torch.randn(2, 4, 8) result = DecodeMLAPreprocessResult( ql_nope=ql_nope, q_pe=q_pe, k_nope=k_nope, k_pe=k_pe, decode_q_wo_k_up=decode_q_wo_k_up, dequant_scale_q_nope=dequant_scale_q_nope, ) self.assertIs(result.ql_nope, ql_nope) self.assertIs(result.q_pe, q_pe) self.assertIs(result.k_nope, k_nope) self.assertIs(result.k_pe, k_pe) self.assertIs(result.decode_q_wo_k_up, decode_q_wo_k_up) self.assertIs(result.dequant_scale_q_nope, dequant_scale_q_nope) class TestPrefillMLAPreprocessResult(TestBase): def test_prefill_mla_preprocess_result_default(self): result = PrefillMLAPreprocessResult() self.assertIsNone(result.q_nope) self.assertIsNone(result.q_pe) self.assertIsNone(result.k_nope) self.assertIsNone(result.k_pe) self.assertIsNone(result.value) def test_prefill_mla_preprocess_result_with_values(self): q_nope = torch.randn(2, 4, 8) q_pe = torch.randn(2, 4, 8) k_nope = torch.randn(2, 4, 8) k_pe = torch.randn(2, 4, 8) value = torch.randn(2, 4, 8) result = PrefillMLAPreprocessResult(q_nope=q_nope, q_pe=q_pe, k_nope=k_nope, k_pe=k_pe, value=value) self.assertIs(result.q_nope, q_nope) self.assertIs(result.q_pe, q_pe) self.assertIs(result.k_nope, k_nope) self.assertIs(result.k_pe, k_pe) self.assertIs(result.value, value) class TestAscendMLAPrefillMetadata(TestBase): def test_ascend_mla_prefill_metadata_default(self): attn_mask = torch.tensor([[1, 0], [1, 1]], dtype=torch.bool) query_lens = [1, 2] seq_lens = [2, 2] context_lens = torch.tensor([1, 2]) input_positions = torch.tensor([0, 1, 0, 1]) query_start_loc = torch.tensor([0, 1, 3]) block_table = torch.tensor([[0, 1], [2, 3]]) max_query_len = 2 max_seq_lens = 2 metadata = AscendMLAPrefillMetadata( attn_mask=attn_mask, query_lens=query_lens, seq_lens=seq_lens, context_lens=context_lens, input_positions=input_positions, query_start_loc=query_start_loc, block_table=block_table, max_query_len=max_query_len, max_seq_lens=max_seq_lens, ) self.assertIs(metadata.attn_mask, attn_mask) self.assertEqual(metadata.query_lens, query_lens) self.assertEqual(metadata.seq_lens, seq_lens) self.assertIs(metadata.context_lens, context_lens) self.assertIs(metadata.input_positions, input_positions) self.assertIs(metadata.query_start_loc, query_start_loc) self.assertIs(metadata.block_table, block_table) self.assertEqual(metadata.max_query_len, max_query_len) self.assertEqual(metadata.max_seq_lens, max_seq_lens) self.assertIsNone(metadata.chunked_context) def test_ascend_mla_prefill_metadata_with_chunked_context(self): cu_seq_lens = torch.tensor([0, 2, 4]) starts = torch.tensor([0, 2]) seq_tot = [2, 2] max_seq_lens = [2, 2] workspace = torch.randn(2, 4) chunk_seq_lens = torch.tensor([2, 2]) chunked_context = ChunkedContextMetadata( cu_seq_lens=cu_seq_lens, starts=starts, seq_tot=seq_tot, max_seq_lens=max_seq_lens, workspace=workspace, chunk_seq_lens=chunk_seq_lens, chunk_seq_lens_npu=chunk_seq_lens, chunk_actual_seq_lengths_kv_list=[[2, 4]], ) metadata = AscendMLAPrefillMetadata( attn_mask=torch.tensor([[1, 0], [1, 1]], dtype=torch.bool), query_lens=[1, 2], seq_lens=[2, 2], context_lens=torch.tensor([1, 2]), input_positions=torch.tensor([0, 1, 0, 1]), query_start_loc=torch.tensor([0, 1, 3]), block_table=torch.tensor([[0, 1], [2, 3]]), max_query_len=2, max_seq_lens=2, chunked_context=chunked_context, ) self.assertIsNotNone(metadata.chunked_context) self.assertIs(metadata.chunked_context.cu_seq_lens, cu_seq_lens) self.assertIs(metadata.chunked_context.starts, starts) self.assertEqual(metadata.chunked_context.seq_tot, seq_tot) self.assertEqual(metadata.chunked_context.max_seq_lens, max_seq_lens) self.assertIs(metadata.chunked_context.workspace, workspace) self.assertIs(metadata.chunked_context.chunk_seq_lens, chunk_seq_lens) self.assertIs(metadata.chunked_context.chunk_seq_lens_npu, chunk_seq_lens) class TestAscendMLADecodeMetadata(TestBase): def test_ascend_mla_decode_metadata_default(self): input_positions = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]]) block_table = torch.tensor([[0, 3, 2, 1], [0, 2, 1, 3]]) seq_lens = torch.tensor([[2], [3]]) max_seq_lens = 4 seq_lens_list = [2, 3] attn_mask = None cp_seq_len = torch.tensor([2, 3]) metadata = AscendMLADecodeMetadata( input_positions=input_positions, block_table=block_table, seq_lens=seq_lens, max_seq_lens=max_seq_lens, seq_lens_list=seq_lens_list, attn_mask=attn_mask, cp_seq_len=cp_seq_len, ) self.assertIs(metadata.input_positions, input_positions) self.assertIs(metadata.block_table, block_table) self.assertIs(metadata.seq_lens, seq_lens) self.assertEqual(metadata.max_seq_lens, max_seq_lens) self.assertEqual(metadata.seq_lens_list, seq_lens_list) self.assertIsNone(attn_mask) self.assertIs(metadata.cp_seq_len, cp_seq_len) class TestAscendMLAMetadata(TestBase): def test_ascend_mla_metadata_default(self): num_actual_tokens_pcp_padded = 100 num_actual_tokens = 100 slot_mapping = torch.randn(100, 4, 1024) query_start_loc = torch.tensor([1, 2, 3, 4]) seq_lens = [30, 50] block_tables = torch.randint(0, 100, (100, 4)) num_decodes = 4 num_decode_tokens = 8 num_prefills = 8 num_input_tokens = 2 query_lens = None head_dim = None attn_mask = None attn_state = AscendAttentionState.ChunkedPrefill decode = None prefill = None metadata = AscendMLAMetadata( num_actual_tokens_pcp_padded, num_actual_tokens, slot_mapping, query_start_loc, seq_lens, seq_lens, block_tables, num_decodes, num_decode_tokens, num_prefills, num_input_tokens, query_lens, head_dim, attn_mask, attn_state, decode, prefill, ) self.assertEqual(metadata.num_actual_tokens, num_actual_tokens) self.assertIs(metadata.slot_mapping, slot_mapping) self.assertIs(metadata.query_start_loc, query_start_loc) self.assertEqual(metadata.seq_lens, seq_lens) self.assertIs(metadata.block_tables, block_tables) self.assertEqual(metadata.num_decodes, num_decodes) self.assertEqual(metadata.num_decode_tokens, num_decode_tokens) self.assertEqual(metadata.num_prefills, num_prefills) self.assertEqual(metadata.num_input_tokens, num_input_tokens) self.assertEqual(metadata.query_lens, query_lens) self.assertEqual(metadata.head_dim, head_dim) self.assertEqual(metadata.attn_mask, attn_mask) self.assertEqual(metadata.attn_state, attn_state) self.assertEqual(metadata.decode, decode) self.assertEqual(metadata.prefill, prefill) class TestAscendMLAMetadataBuilder(TestBase): def setUp(self): # 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() def tearDown(self): self.parent_init_patcher.stop() def test_ascend_mla_metadata_builder_default(self): mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_vllm_config.speculative_config = None ascend_config = MagicMock() with patch("vllm_ascend.attention.mla_v1.get_ascend_config", return_value=ascend_config): builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) self.assertEqual(builder.block_size, mock_vllm_config.cache_config.block_size) self.assertEqual(builder.chunked_prefill_enabled, mock_vllm_config.scheduler_config.enable_chunked_prefill) def test_ascend_mla_metadata_builder_spec_decode(self): mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_spec_config = MagicMock() mock_spec_config.num_speculative_tokens = 3 mock_vllm_config.speculative_config = mock_spec_config ascend_config = MagicMock() with patch("vllm_ascend.attention.mla_v1.get_ascend_config", return_value=ascend_config): builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) self.assertEqual(builder.block_size, mock_vllm_config.cache_config.block_size) self.assertEqual(builder.chunked_prefill_enabled, mock_vllm_config.scheduler_config.enable_chunked_prefill) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") @patch("vllm_ascend.attention.attention_mask.get_pcp_group") @patch("vllm.distributed.parallel_state.get_pcp_group") def test_ascend_mla_metadata_builder_build_full_graph( self, mock_get_pcp_group, mock_get_pcp_group_mask, mock_get_cos_and_sin_mla ): pcp_group = MagicMock() pcp_group.world_size = 1 mock_get_pcp_group.return_value = pcp_group mock_get_pcp_group_mask.return_value = pcp_group mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.chunked_prefill_enabled = False mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" torch.Tensor.pin_memory = lambda x: x # noqa mock_spec_config = MagicMock() mock_spec_config.num_speculative_tokens = 1 mock_spec_config.disable_padded_drafter_batch = True mock_vllm_config.speculative_config = mock_spec_config builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) common_metadata = MagicMock() common_metadata.graph_pad_size = 8 common_metadata.num_reqs = 4 common_metadata.num_actual_tokens = 5 common_metadata.max_query_len = 5 common_metadata.seq_lens_cpu = torch.Tensor([9, 10, 8, 8]).int() common_metadata.query_start_loc = torch.Tensor([0, 1, 2, 4, 5]).int() common_metadata.query_start_loc_cpu = torch.Tensor([0, 1, 2, 4, 5]).int() common_metadata.positions = torch.Tensor([1, 2, 3, 4, 5, 6]).int() block_table = torch.Tensor([[1, 0], [2, 0], [3, 0], [4, 0]]).int() common_metadata.block_table_tensor = block_table common_metadata.prefill_context_parallel_metadata = None mock_get_cos_and_sin_mla.return_value = (torch.tensor([6, 6]), torch.Tensor([6, 6])) metadata = builder.build(0, common_metadata) self.assertEqual(metadata.decode.actual_seq_lengths_q, [1, 2, 4, 5, 6, 6, 7, 8]) self.assertEqual(metadata.decode.block_table.shape[0], 8) def test_reorder_batch(self): ascend_config = MagicMock() mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_vllm_config.speculative_config = None with patch("vllm_ascend.attention.mla_v1.get_ascend_config", return_value=ascend_config): builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) builder.decode_threshold = 1 input_batch = MagicMock() input_batch.req_ids = [0, 1, 2, 3] scheduler_output = MagicMock() scheduler_output.num_scheduled_tokens = {0: 1, 1: 3, 2: 1, 3: 2} scheduler_output.scheduled_spec_decode_tokens = {0: [], 1: [1], 2: [], 3: []} input_batch.swap_states = MagicMock() modified = builder.reorder_batch(input_batch, scheduler_output) self.assertTrue(modified) input_batch.swap_states.assert_called_once_with(1, 2) def test_determine_chunked_prefill_workspace_size(self): mock_vllm_config = MagicMock() mock_vllm_config.scheduler_config.enable_chunked_prefill = True mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 128 mock_vllm_config.scheduler_config.max_num_batched_tokens = 4096 mock_vllm_config.model_config.max_model_len = 4096 result = AscendMLAMetadataBuilder.determine_chunked_prefill_workspace_size(mock_vllm_config) self.assertGreater(result, 0) def test_get_cudagraph_support(self): mock_vllm_config = MagicMock() mock_kv_cache_spec = MagicMock() result = AscendMLAMetadataBuilder.get_cudagraph_support(mock_vllm_config, mock_kv_cache_spec) from vllm.v1.attention.backend import AttentionCGSupport self.assertEqual(result, AttentionCGSupport.UNIFORM_BATCH) def test_set_num_actual_tokens(self): mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_vllm_config.speculative_config = None builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) common_attn_metadata = MagicMock() common_attn_metadata.num_actual_tokens = 100 builder.set_num_actual_tokens(common_attn_metadata) self.assertEqual(builder.num_actual_tokens, 100) def test_pad_actual_seq_lens_q_mtp_disable_pad(self): mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.chunked_prefill_enabled = False mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_vllm_config.speculative_config = None builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) input_seq_lens = [1, 2, 4, 5] expect_output = [1, 2, 4, 5, 6, 6, 7, 8] num_reqs = 4 num_reqs_pad_size = 4 output_seq_lens = builder.pad_actual_seq_len_q_mtp_disable_pad(num_reqs_pad_size, num_reqs, input_seq_lens) self.assertEqual(output_seq_lens, expect_output) def test_pad_actual_seq_lens_q_mtp_enable_pad(self): mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.chunked_prefill_enabled = False mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_vllm_config.speculative_config = None common_metadata = MagicMock() common_metadata.actual_seq_lengths_q = [2, 4, 6, 8] builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) input_seq_lens = [2, 4, 6] expect_output = [2, 4, 6, 8] num_reqs = 3 num_reqs_pad_size = 1 output_seq_lens = builder.pad_actual_seq_len_q_mtp_enable_pad( num_reqs_pad_size, num_reqs, input_seq_lens, common_metadata ) self.assertEqual(output_seq_lens, expect_output) def test_pad_actual_seq_lens_q_mtp_enable_pad_with_padding(self): mock_vllm_config = MagicMock() mock_vllm_config.model_config.max_model_len = 1024 mock_vllm_config.model_config.get_head_size.return_value = 64 mock_vllm_config.model_config.dtype = torch.float16 mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64 mock_vllm_config.cache_config.block_size = 16 mock_vllm_config.scheduler_config.max_num_seqs = 4 mock_vllm_config.scheduler_config.chunked_prefill_enabled = False mock_vllm_config.scheduler_config.enable_chunked_prefill = False mock_device = "cpu" mock_vllm_config.speculative_config = None common_metadata = MagicMock() common_metadata.actual_seq_lengths_q = [2, 4, 6, 100] builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config, mock_device) input_seq_lens = [2, 4, 6] num_reqs = 3 num_reqs_pad_size = 1 output_seq_lens = builder.pad_actual_seq_len_q_mtp_enable_pad( num_reqs_pad_size, num_reqs, input_seq_lens, common_metadata ) self.assertEqual(len(output_seq_lens), 4) self.assertEqual(output_seq_lens[:3], [2, 4, 6]) self.assertEqual(output_seq_lens[-1], 100) class TestAscendMLAMetadataBuilderBuild(TestBase): def setUp(self): # 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() self.mock_vllm_config = MagicMock(spec=VllmConfig) self.mock_vllm_config.cache_config = CacheConfig(block_size=32) mock_scheduler_config = MagicMock(spec=SchedulerConfig) mock_scheduler_config.max_num_seqs = 8 mock_scheduler_config.chunked_prefill_enabled = True mock_scheduler_config.enable_chunked_prefill = True self.mock_vllm_config.scheduler_config = mock_scheduler_config self.mock_vllm_config.speculative_config = None self.mock_device = torch.device("cpu") fake_weight_path = os.path.join(os.path.dirname(__file__), "..", "..", "_fake_weight") model_config = ModelConfig( model=fake_weight_path, skip_tokenizer_init=True, ) model_config.hf_text_config.head_dim = 128 model_config.hf_text_config.qk_rope_head_dim = 32 self.mock_vllm_config.model_config = model_config self.kv_cache_spec = MagicMock() self.kv_cache_spec.num_layers = 32 self.kv_cache_spec.head_size = 64 self.kv_cache_spec.num_heads = 32 def tearDown(self): self.parent_init_patcher.stop() @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") @patch("vllm_ascend.attention.attention_mask.get_pcp_group") @patch("vllm.distributed.parallel_state.get_pcp_group") @patch("vllm_ascend.attention.mla_v1.torch.zeros", wraps=torch.zeros) @patch("torch.Tensor.npu", new=lambda self: self) @patch("torch.npu.is_available") def test_build_prefix_no_cache_metadata( self, mock_npu_available, mock_zeros, mock_get_pcp_group, mock_get_pcp_group_mask, mock_get_cos_and_sin_mla ): mock_npu_available.return_value = False torch.Tensor.pin_memory = lambda x: x # noqa pcp_group = MagicMock() pcp_group.world_size = 1 mock_get_pcp_group.return_value = pcp_group mock_get_pcp_group_mask.return_value = pcp_group def zeros_override(*args, **kwargs): kwargs.pop("pin_memory", None) return mock_zeros._mock_wraps(*args, **kwargs) mock_zeros.side_effect = zeros_override common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=torch.tensor([0, 3, 7]), query_start_loc_cpu=torch.tensor([0, 3, 7]), seq_lens_cpu=torch.tensor([5, 6]), num_reqs=2, num_actual_tokens=10, max_query_len=5, decode_token_per_req=torch.tensor([1, 1]), block_table_tensor=torch.zeros((10, 10)), slot_mapping=torch.tensor(range(20)), actual_seq_lengths_q=torch.tensor([0, 1]), positions=torch.tensor([10, 10]), attn_state=AscendAttentionState.PrefillNoCache, num_computed_tokens_cpu=None, seq_lens=None, max_seq_len=6, ) base_inputs = { "num_actual_tokens": 10, "slot_mapping": torch.tensor(range(10)), "query_start_loc": torch.tensor([0, 3, 7]), "seq_lens": torch.tensor([5, 6]), "block_tables": torch.zeros((10, 10)), "num_prefills": 2, } builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) mock_get_cos_and_sin_mla.return_value = (torch.tensor(10), torch.Tensor(10)) metadata = builder.build(1, common_attn_metadata) self.assertIsInstance(metadata, AscendMLAMetadata) self.assertEqual(metadata.num_actual_tokens, base_inputs["num_actual_tokens"]) self.assertTrue(torch.all(metadata.slot_mapping == base_inputs["slot_mapping"])) self.assertEqual(metadata.head_dim, self.kv_cache_spec.head_size) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") @patch("vllm_ascend.attention.attention_mask.get_pcp_group") @patch("vllm.distributed.parallel_state.get_pcp_group") @patch("vllm_ascend.attention.mla_v1.torch.zeros", wraps=torch.zeros) @patch("torch.Tensor.npu", new=lambda self: self) @patch("torch.npu.is_available") def test_build_chunked_prefix_metadata( self, mock_npu_available, mock_zeros, mock_get_pcp_group, mock_get_pcp_group_mask, mock_get_cos_and_sin_mla ): mock_npu_available.return_value = False torch.Tensor.pin_memory = lambda x: x # noqa pcp_group = MagicMock() pcp_group.world_size = 1 mock_get_pcp_group.return_value = pcp_group mock_get_pcp_group_mask.return_value = pcp_group def zeros_override(*args, **kwargs): kwargs.pop("pin_memory", None) return mock_zeros._mock_wraps(*args, **kwargs) mock_zeros.side_effect = zeros_override 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, ) base_inputs = { "num_actual_tokens": 15, "slot_mapping": torch.tensor(range(15)), "query_start_loc": torch.tensor([0, 2, 5, 9]), "seq_lens": torch.tensor([4, 5, 6]), "block_tables": torch.zeros((10, 10)), "num_prefills": 3, } builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) mock_get_cos_and_sin_mla.return_value = (torch.tensor(10), torch.Tensor(10)) metadata = builder.build(1, common_attn_metadata) self.assertIsInstance(metadata, AscendMLAMetadata) self.assertEqual(metadata.num_actual_tokens, base_inputs["num_actual_tokens"]) self.assertTrue(torch.all(metadata.slot_mapping == base_inputs["slot_mapping"])) self.assertEqual(metadata.head_dim, self.kv_cache_spec.head_size) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") @patch("vllm_ascend.attention.attention_mask.get_pcp_group") @patch("vllm.distributed.parallel_state.get_pcp_group") def test_build_decode_only_metadata(self, mock_get_pcp_group, mock_get_pcp_group_mask, mock_get_cos_and_sin_mla): torch.Tensor.pin_memory = lambda x: x # noqa pcp_group = MagicMock() pcp_group.world_size = 1 mock_get_pcp_group.return_value = pcp_group mock_get_pcp_group_mask.return_value = pcp_group common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=torch.tensor([0, 1, 2, 3]), query_start_loc_cpu=torch.tensor([0, 1, 2, 3]), seq_lens_cpu=torch.tensor([4, 5, 6]), num_reqs=3, num_actual_tokens=3, max_query_len=1, block_table_tensor=torch.zeros((10, 10)), slot_mapping=torch.tensor(range(3)), actual_seq_lengths_q=torch.tensor([0, 1, 2]), decode_token_per_req=torch.tensor([1, 1, 1]), positions=torch.tensor([10, 10]), attn_state=AscendAttentionState.DecodeOnly, num_computed_tokens_cpu=None, seq_lens=None, max_seq_len=6, ) base_inputs = { "num_actual_tokens": 3, "slot_mapping": torch.tensor(range(3)), "query_start_loc": torch.tensor([0, 1, 2, 3]), "seq_lens": torch.tensor([4, 5, 6]), "num_decodes": 3, } builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) mock_get_cos_and_sin_mla.return_value = (torch.tensor([10, 10]), torch.Tensor([10, 10])) metadata = builder.build(1, common_attn_metadata) self.assertIsInstance(metadata, AscendMLAMetadata) self.assertEqual(metadata.num_actual_tokens, base_inputs["num_actual_tokens"]) self.assertTrue(torch.all(metadata.slot_mapping == base_inputs["slot_mapping"])) self.assertEqual(metadata.head_dim, self.kv_cache_spec.head_size) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") def test_build_decode_metadata_without_disable_padded_drafter_batch(self, mock_get_cos_and_sin_mla): common_attn_metadata = MagicMock() common_attn_metadata.num_reqs = 3 common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 1, 2, 3]) common_attn_metadata.positions = torch.tensor([10, 10, 10]) common_attn_metadata.decode_token_per_req = 1 builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) builder.num_actual_tokens = 3 builder.num_decode_tokens = 3 builder.num_decodes = 3 builder.graph_pad_size = 5 # > num_reqs builder.seq_lens = torch.tensor([4, 5, 6]) builder.slot_mapping = torch.tensor(range(3)) builder.block_table = torch.zeros((3, 10)) mock_speculative_config = MagicMock() mock_speculative_config.disable_padded_drafter_batch = False builder.speculative_config = mock_speculative_config builder.attn_mask_builder = MagicMock() builder.attn_mask_builder.get_splitfuse_attn_mask.return_value = torch.randn(1, 1, 5, 5) mock_get_cos_and_sin_mla.return_value = (torch.randn(5, 32), torch.randn(5, 32)) metadata = builder.build_decode_metadata(0, common_attn_metadata) self.assertIsInstance(metadata, AscendMLADecodeMetadata) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") @patch("vllm_ascend.attention.attention_mask.get_pcp_group") @patch("vllm.distributed.parallel_state.get_pcp_group") def test_build_for_graph_capture_decode_only( self, mock_get_pcp_group, mock_get_pcp_group_mask, mock_get_cos_and_sin_mla ): torch.Tensor.pin_memory = lambda x: x # noqa pcp_group = MagicMock() pcp_group.world_size = 1 mock_get_pcp_group.return_value = pcp_group mock_get_pcp_group_mask.return_value = pcp_group common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=torch.tensor([0, 1, 2, 3]), query_start_loc_cpu=torch.tensor([0, 1, 2, 3]), seq_lens_cpu=torch.tensor([4, 5, 6]), num_reqs=3, num_actual_tokens=3, max_query_len=1, block_table_tensor=torch.zeros((10, 10)), slot_mapping=torch.tensor(range(3)), actual_seq_lengths_q=torch.tensor([0, 1, 2]), decode_token_per_req=torch.tensor([1, 1, 1]), positions=torch.tensor([10, 10]), attn_state=AscendAttentionState.DecodeOnly, num_computed_tokens_cpu=None, seq_lens=None, max_seq_len=6, ) base_inputs = { "num_actual_tokens": 3, "slot_mapping": torch.tensor(range(3)), "query_start_loc": torch.tensor([0, 1, 2, 3]), "seq_lens": torch.tensor([4, 5, 6]), "num_decodes": 3, } builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) mock_get_cos_and_sin_mla.return_value = (torch.tensor([10, 10]), torch.Tensor([10, 10])) metadata = builder.build_for_graph_capture(common_attn_metadata, AscendAttentionState.DecodeOnly) self.assertIsInstance(metadata, AscendMLAMetadata) self.assertEqual(metadata.num_actual_tokens, base_inputs["num_actual_tokens"]) self.assertTrue(torch.all(metadata.slot_mapping == base_inputs["slot_mapping"])) self.assertEqual(metadata.head_dim, self.kv_cache_spec.head_size) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") def test_build_for_graph_capture_prefill(self, mock_get_cos_and_sin_mla): torch.Tensor.pin_memory = lambda x: x # noqa common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=torch.tensor([0, 3, 7]), query_start_loc_cpu=torch.tensor([0, 3, 7]), seq_lens_cpu=torch.tensor([5, 6]), num_reqs=2, num_actual_tokens=10, max_query_len=5, decode_token_per_req=torch.tensor([1, 1]), block_table_tensor=torch.zeros((10, 10)), slot_mapping=torch.tensor(range(20)), actual_seq_lengths_q=torch.tensor([0, 1]), positions=torch.tensor([10, 10]), attn_state=AscendAttentionState.PrefillNoCache, num_computed_tokens_cpu=None, seq_lens=None, max_seq_len=6, ) builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) mock_get_cos_and_sin_mla.return_value = (torch.tensor(10), torch.Tensor(10)) with self.assertRaises(NotImplementedError) as ctx: builder.build_for_graph_capture(common_attn_metadata, AscendAttentionState.PrefillNoCache) self.assertIn( "Currently we only support building dummy metadata for DecodeOnly and SpecDecoding state", str(ctx.exception), ) @patch("vllm_ascend.attention.mla_v1.get_cos_and_sin_mla") @patch("vllm_ascend.attention.attention_mask.get_pcp_group") @patch("vllm.distributed.parallel_state.get_pcp_group") def test_build_with_seq_lens_only(self, mock_get_pcp_group, mock_get_pcp_group_mask, mock_get_cos_and_sin_mla): torch.Tensor.pin_memory = lambda x: x # noqa pcp_group = MagicMock() pcp_group.world_size = 1 mock_get_pcp_group.return_value = pcp_group mock_get_pcp_group_mask.return_value = pcp_group common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=torch.tensor([0, 2, 5, 8]), query_start_loc_cpu=torch.tensor([0, 2, 5, 8]), seq_lens_cpu=None, num_reqs=3, num_actual_tokens=8, max_query_len=3, block_table_tensor=torch.zeros((10, 10)), slot_mapping=torch.tensor(range(8)), actual_seq_lengths_q=torch.tensor([2, 3, 3]), decode_token_per_req=torch.tensor([0, 0, 0]), positions=torch.tensor([0, 1, 0, 1, 2, 0, 1, 2]), attn_state=AscendAttentionState.PrefillNoCache, num_computed_tokens_cpu=None, seq_lens=torch.tensor([2, 3, 3]), max_seq_len=3, ) common_attn_metadata._seq_lens_cpu = None builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) mock_get_cos_and_sin_mla.return_value = (torch.randn(3, 32), torch.randn(3, 32)) metadata = builder.build(0, common_attn_metadata) self.assertIsInstance(metadata, AscendMLAMetadata) def test_build_chunked_metadata_without_chunked_prefill(self): common_attn_metadata = MagicMock() common_attn_metadata.num_reqs = 3 builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) builder.chunked_prefill_enabled = False result = builder.build_chunked_metadata(0, common_attn_metadata) self.assertIsNone(result) def test_build_chunked_metadata_with_no_context(self): common_attn_metadata = MagicMock() common_attn_metadata.num_reqs = 3 builder = AscendMLAMetadataBuilder( kv_cache_spec=self.kv_cache_spec, layer_names=["layer_0", "layer_1"], vllm_config=self.mock_vllm_config, device=self.mock_device, ) builder.chunked_prefill_enabled = True builder.seq_lens = torch.tensor([2, 2, 2]) builder.query_lens = torch.tensor([2, 2, 2]) builder.num_decodes = 0 result = builder.build_chunked_metadata(0, common_attn_metadata) self.assertIsNone(result) class TestAscendMLAImpl(TestBase): @patch("vllm.distributed.parallel_state._TP", new_callable=lambda: MagicMock(spec=GroupCoordinator)) @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") def setUp(self, get_current_vllm_config, 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 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 = AscendMLAImpl( 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, ) self.impl.fa_quant_layer = False 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) # 256 is power of 2, so padding should be 0 self.assertEqual(self.impl.num_heads_padded, 256) self.assertEqual(self.impl.head_padding, 0) @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") def test_init_head_padding_for_non_power_of_two(self, mock_get_current_vllm_config): """Test head padding computation for num_heads that are not power of 2 (e.g. GLM-4.7-Flash with 20 heads).""" mock_get_current_vllm_config.return_value = MagicMock() 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": MagicMock(), "rotary_emb": MagicMock(), } impl = AscendMLAImpl( num_heads=20, head_size=1024, scale=0.1, num_kv_heads=20, alibi_slopes=None, sliding_window=None, kv_cache_dtype="auto", blocksparse_params=None, logits_soft_cap=None, attn_type=None, kv_sharing_target_layer_name=None, **kwargs, ) self.assertEqual(impl.num_heads, 20) self.assertEqual(impl.num_heads_padded, 32) # next power of 2 self.assertEqual(impl.head_padding, 12) # 32 - 20 @patch("vllm_ascend.attention.mla_v1.get_ascend_config") @patch("vllm_ascend.attention.mla_v1.register_all_layers_to_shard_weight_series") @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") def test_init_with_layer_sharding(self, mock_get_current_vllm_config, mock_register, mock_get_ascend_config): mock_config = MagicMock() mock_config.layer_sharding = ["layer_0", "layer_1"] mock_get_ascend_config.return_value = mock_config mock_vllm_config = MagicMock() mock_speculative_config = MagicMock() mock_model_config = MagicMock() mock_parallel_config = MagicMock() mock_parallel_config.prefill_context_parallel_size = 1 mock_speculative_config.num_speculative_tokens = 4 mock_vllm_config.speculative_config = mock_speculative_config mock_model_config.dtype = torch.float16 mock_vllm_config.model_config = mock_model_config mock_vllm_config.additional_config = {"refresh": True} mock_vllm_config.parallel_config = mock_parallel_config mock_get_current_vllm_config.return_value = mock_vllm_config kwargs = { "layer_name": "layer_0", "layer_0": "sharding_config_0", "layer_2": "sharding_config_2", "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": MagicMock(), "rotary_emb": MagicMock(), } impl = AscendMLAImpl( num_heads=256, head_size=1024, scale=0.1, num_kv_heads=8, alibi_slopes=None, sliding_window=None, kv_cache_dtype="auto", blocksparse_params=None, logits_soft_cap=None, attn_type=None, kv_sharing_target_layer_name=None, **kwargs, ) self.assertEqual(len(impl.layer_sharding_kwargs), 1) # only include layer_0 self.assertEqual(impl.layer_sharding_kwargs[0], "sharding_config_0") mock_register.assert_called_once_with(impl.layer_sharding_kwargs) def test_q_proj_and_k_up_proj(self): batch_size = 4 x = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_head_dim) q_proj_output = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_head_dim) self.impl.q_proj.return_value = (q_proj_output,) if not hasattr(self.impl, "W_UK_T") or self.impl.W_UK_T is None: self.impl.W_UK_T = torch.randn(self.impl.num_heads, self.impl.qk_nope_head_dim, self.impl.kv_lora_rank) result = self.impl._q_proj_and_k_up_proj(x) ql_nope, q_pe = result self.assertEqual(ql_nope.shape[0], batch_size) self.assertEqual(ql_nope.shape[1], self.impl.num_heads) self.assertEqual(ql_nope.shape[2], self.impl.kv_lora_rank) self.assertEqual(q_pe.shape[0], batch_size) self.assertEqual(q_pe.shape[1], self.impl.num_heads) self.assertEqual(q_pe.shape[2], self.impl.qk_rope_head_dim) @patch("torch_npu.npu_interleave_rope") def test_rope_single(self, mock_npu_interleave_rope): batch_size = 2 seq_len = 10 dim = 32 x = torch.randn(batch_size, seq_len, dim) cos = torch.randn(seq_len, dim) sin = torch.randn(seq_len, dim) mock_npu_interleave_rope.return_value = torch.randn(batch_size, seq_len, 1, dim) result = self.impl.rope_single(x, cos, sin) self.assertEqual(result.shape, (batch_size, seq_len, dim)) mock_npu_interleave_rope.assert_called_once() def test_forward_mha_not_implemented(self): layer_name = "layer_0" hidden_states = torch.randn(2, 10, 768) kv_cache = [torch.randn(10, 1, 1, 768), torch.randn(10, 1, 1, 768)] attn_metadata = MagicMock() with self.assertRaises(NotImplementedError) as ctx: self.impl.forward_mha(layer_name, hidden_states, kv_cache, attn_metadata) self.assertIn( "forward_mha is not supported for MLA attention. Use forward() instead.", str(ctx.exception), ) def test_forward_mqa_not_implemented(self): layer_name = "layer_0" hidden_states = torch.randn(2, 10, 768) kv_cache = [torch.randn(10, 1, 1, 768), torch.randn(10, 1, 1, 768)] attn_metadata = MagicMock() with self.assertRaises(NotImplementedError) as ctx: self.impl.forward_mqa(layer_name, hidden_states, kv_cache, attn_metadata) self.assertIn( "forward_mqa is not supported for MLA attention. Use forward() instead.", str(ctx.exception), ) @patch("vllm_ascend.attention.mla_v1.torch_npu") def test_v_up_proj(self, mock_torch_npu): batch_size = 4 x = torch.randn(self.impl.num_heads, batch_size, self.impl.kv_lora_rank) if not hasattr(self.impl, "W_UV") or self.impl.W_UV is None: self.impl.W_UV = torch.randn(self.impl.num_heads, self.impl.kv_lora_rank, self.impl.v_head_dim) expected_shape = (batch_size, self.impl.num_heads * self.impl.v_head_dim) mock_torch_npu.npu_transpose_batchmatmul.return_value = torch.randn(*expected_shape) result = self.impl._v_up_proj(x) self.assertEqual(result.shape[0], batch_size) self.assertEqual(result.shape[1], self.impl.num_heads * self.impl.v_head_dim) @patch("vllm_ascend.attention.mla_v1.get_draft_graph_params") @patch("vllm_ascend.attention.mla_v1.get_graph_params") @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_v2.out") @patch("vllm_ascend.ascend_forward_context.get_forward_context") def test_update_graph_params( self, mock_get_forward_context, mock_fia, mock_update_end, mock_update_begin, mock_stream, mock_get_graph_params, mock_get_draft_graph_params, ): mock_update_stream = MagicMock() mock_forward_context = MagicMock() mock_attn_metadata = MagicMock() mock_forward_context.attn_metadata = {"layer_0": mock_attn_metadata} mock_attn_metadata.decode.seq_lens_list = [10, 20, 30] mock_attn_metadata.decode.actual_seq_lengths_q = [10, 20, 30] mock_attn_metadata.decode.block_table = torch.randint(0, 100, (3, 4)) mock_graph_params = MagicMock() mock_graph_params.attn_params = { 100: [ ( MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), ), ] } mock_graph_params.handles = {100: [MagicMock()]} mock_graph_params.events = {100: [MagicMock()]} mock_graph_params.workspaces = {100: MagicMock()} mock_get_graph_params.return_value = mock_graph_params mock_get_draft_graph_params.return_value = mock_graph_params # forward context mock_ctx = MagicMock() mock_get_forward_context.return_value = mock_ctx # speculative_config mock_speculative_config = MagicMock() mock_speculative_config.disable_padded_drafter_batch = False # Test non-draft model mock_ctx.is_draft_model = False AscendMLAImpl.update_graph_params( mock_update_stream, mock_forward_context, 100, speculative_config=mock_speculative_config, ) @patch("vllm_ascend.ascend_forward_context.get_forward_context") def test_update_graph_params_empty_layers(self, mock_get_forward_context): # if num_layers == 0 mock_update_stream = MagicMock() mock_forward_context = MagicMock() mock_forward_context.attn_metadata = {} mock_ctx = MagicMock() mock_ctx.is_draft_model = False mock_get_forward_context.return_value = mock_ctx AscendMLAImpl.update_graph_params(mock_update_stream, mock_forward_context, 100) @patch("vllm_ascend.attention.mla_v1.get_graph_params") @patch("torch_npu.npu_fused_infer_attention_score_v2") @patch("torch.npu.graph_task_update_end") @patch("torch.npu.graph_task_update_begin") @patch("torch.npu.stream") @patch("vllm_ascend.ascend_forward_context.get_forward_context") def test_update_graph_params_with_mtp( self, mock_get_forward_context, mock_npu_stream, mock_graph_task_update_begin, mock_graph_task_update_end, mock_npu_fused_infer, mock_get_graph_params, ): mock_update_stream = MagicMock() mock_forward_context = MagicMock() mock_attn_metadata = MagicMock() mock_attn_metadata.decode = MagicMock() mock_attn_metadata.decode.seq_lens_list = [10, 20, 30] mock_attn_metadata.decode.actual_seq_lengths_q = [10, 20, 30] mock_forward_context.attn_metadata = {"layer_0": mock_attn_metadata} # forward context mock_ctx = MagicMock() mock_ctx.is_draft_model = False mock_get_forward_context.return_value = mock_ctx mock_stream_context = MagicMock() mock_npu_stream.return_value = mock_stream_context mock_out = MagicMock() mock_npu_fused_infer.out = mock_out mock_graph_params = MagicMock() mock_graph_params.attn_params = { 100: [ ( MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock(), ), ] } mock_graph_params.handles = {100: [MagicMock()]} mock_graph_params.events = {100: [MagicMock()]} mock_get_graph_params.return_value = mock_graph_params mock_speculative_config = MagicMock() mock_speculative_config.method = "mtp" mock_speculative_config.num_speculative_tokens = 4 AscendMLAImpl.update_graph_params( mock_update_stream, mock_forward_context, 100, speculative_config=mock_speculative_config ) def test_get_context_seq_len_npu(self): mock_attn_metadata = MagicMock() mock_prefill_metadata = MagicMock() mock_chunked_context = MagicMock() mock_chunked_context.chunk_seq_lens_npu = torch.tensor([10, 20, 30]) mock_chunked_context.seq_tot = [10, 30, 60] mock_prefill_metadata.chunked_context = mock_chunked_context mock_attn_metadata.prefill = mock_prefill_metadata result = self.impl.get_context_seq_len_npu(1, mock_attn_metadata) self.assertEqual(result, 20) def test_reorg_kvcache(self): kv_c_normed = torch.randn(2, 4, 8) k_pe = torch.randn(2, 4, 8) mock_chunked_context = MagicMock() result_kv, result_k_pe = self.impl._reorg_kvcache(kv_c_normed, k_pe, mock_chunked_context, 0, 10) self.assertIs(result_kv, kv_c_normed) self.assertIs(result_k_pe, k_pe) @patch("vllm_ascend.attention.mla_v1.maybe_trans_nz") def test_process_weights_for_fused_fa_quant(self, mock_maybe_trans_nz): self.impl.fa_quant_layer = True self.impl.q_a_layernorm = MagicMock() self.impl.q_a_layernorm.weight.data = torch.randn(128) self.impl.kv_a_layernorm = MagicMock() self.impl.kv_a_layernorm.weight.data = torch.randn(128) self.impl.q_proj = MagicMock() self.impl.q_proj.weight.data = torch.randn(128, 128) self.impl.q_proj.weight_scale.data = torch.randn(128, 128) self.impl.fused_qkv_a_proj = MagicMock() self.impl.fused_qkv_a_proj.weight.data = torch.randn(128, 128, 64) self.impl.fused_qkv_a_proj.weight_scale = torch.randn(64) mock_layer = MagicMock() mock_layer.quant_kscale = torch.randn(128) mock_layer.fak_descale_float = torch.randn(1) self.impl.vllm_config = MagicMock() self.impl.vllm_config.compilation_config = MagicMock() self.impl.vllm_config.compilation_config.static_forward_context = {"layer_0": mock_layer} self.impl.layer_name = "layer_0" self.impl._process_weights_for_fused_fa_quant() self.assertTrue(hasattr(self.impl, "gamma1")) self.assertTrue(hasattr(self.impl, "gamma2")) self.assertTrue(hasattr(self.impl, "wu_q")) self.assertTrue(hasattr(self.impl, "wd_q")) self.assertTrue(hasattr(self.impl, "wd_kv")) @patch("vllm_ascend.attention.mla_v1.trans_rope_weight") @patch("vllm_ascend.attention.mla_v1.transdata") @patch("torch_npu.npu_format_cast") @patch("vllm_ascend.attention.mla_v1.torch_npu") def test_process_weights_for_fused_mlapo( self, mock_torch_npu, mock_format_cast, mock_transdata, mock_trans_rope_weight ): mock_format_cast.return_value = torch.randn(1, 128, 128) mock_transdata.return_value = torch.randn(128, 128) call_count = 0 def mock_trans_rope_weight_func(x, rope_dim): nonlocal call_count call_count += 1 if call_count == 2: # second return [64] tensor return torch.randn(64) else: # first return same shape tensor return torch.randn(x.shape[0], x.shape[1]) mock_trans_rope_weight.side_effect = mock_trans_rope_weight_func mock_torch_npu.npu_format_cast.return_value = torch.randn(1, 128, 128) self.impl.enable_mlapo = True self.impl.fused_qkv_a_proj = MagicMock() q_lora_rank = 32 kv_lora_rank_plus_rope = 32 + 32 # kv_lora_rank + qk_rope_head_dim total_rank = q_lora_rank + kv_lora_rank_plus_rope self.impl.fused_qkv_a_proj.weight.data = torch.randn(128, total_rank) # Fix the shape of deq_scale so that it matches the subsequent reshape operation # Ensure that the size of kv_a_proj_deq_scl is divisible by 64 self.impl.fused_qkv_a_proj.deq_scale = torch.randn(32 + 64 * 128) self.impl.fused_qkv_a_proj.quant_bias = torch.randn(total_rank) self.impl.fused_qkv_a_proj.input_scale.data = torch.randn(1) self.impl.fused_qkv_a_proj.input_offset.data = torch.randn(1) self.impl.q_proj = MagicMock() self.impl.q_proj.weight.data = torch.randn(128, 128) self.impl.q_proj.weight.device = torch.device("cpu") self.impl.q_proj.deq_scale.data = torch.randn(256 * 64) self.impl.q_proj.quant_bias.data = torch.randn(256 * 64) self.impl.q_proj.input_scale.data = torch.randn(1) self.impl.q_proj.input_offset.data = torch.randn(1) self.impl.q_a_layernorm = MagicMock() self.impl.q_a_layernorm.weight.data = torch.randn(128) self.impl.kv_a_layernorm = MagicMock() self.impl.kv_a_layernorm.weight.data = torch.randn(128) self.impl.q_lora_rank = q_lora_rank self.impl.kv_lora_rank = 32 self.impl.qk_rope_head_dim = 32 self.impl.hidden_size = 128 self.impl.num_heads = 256 self.impl.qk_nope_head_dim = 32 self.impl.vllm_config.scheduler_config.max_num_batched_tokens = 4096 self.impl.vllm_config.kv_transfer_config = MagicMock() self.impl.vllm_config.kv_transfer_config.is_kv_consumer = False self.impl._process_weights_for_fused_mlapo(torch.float16) self.assertTrue(hasattr(self.impl, "wd_qkv")) self.assertTrue(hasattr(self.impl, "deq_scale_qkv")) self.assertTrue(hasattr(self.impl, "quant_bias_qkv")) self.assertTrue(hasattr(self.impl, "wu_q")) @patch("torch_npu.npu_format_cast") def test_process_weights_for_fused_mlapo_a5(self, mock_format_cast): mock_format_cast.return_value = torch.randn(128, 128) self.impl.enable_mlapo = True self.impl.fused_qkv_a_proj = MagicMock() self.impl.fused_qkv_a_proj.weight.data = torch.randn(128, 128, 64) self.impl.fused_qkv_a_proj.weight_scale = torch.randn(64, 128, 128) self.impl.q_proj = MagicMock() self.impl.q_proj.weight.data = torch.randn(128, 128) self.impl.q_proj.weight_scale.data = torch.randn(128, 128, 128) self.impl.q_lora_rank = 32 self.impl._process_weights_for_fused_mlapo_a5(torch.float16) self.assertTrue(hasattr(self.impl, "weight_dq")) self.assertTrue(hasattr(self.impl, "weight_uq_qr")) self.assertTrue(hasattr(self.impl, "weight_dkv_kr")) self.assertTrue(hasattr(self.impl, "weight_dq_scale")) self.assertTrue(hasattr(self.impl, "weight_dkv_kr_scale")) @patch("vllm_ascend.attention.mla_v1.DeviceOperator") @patch("torch_npu.npu_fused_infer_attention_score") @patch("torch_npu.npu_attention_update") def test__forward_prefill(self, mock_npu_attention_update, mock_fia, mock_device_operator): batch_size = 2 # create input tensors q_nope = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_nope_head_dim) q_pe = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_rope_head_dim) k_nope = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_nope_head_dim) k_pe = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_rope_head_dim) value = torch.randn(batch_size, self.impl.num_heads, self.impl.v_head_dim) kv_c_and_k_pe_cache = [torch.randn(10, 1, 1, 192), torch.randn(10, 1, 1, 32)] attn_metadata = MagicMock() prefill_metadata = MagicMock() prefill_metadata.actual_seq_lengths_q = [10, 20] prefill_metadata.attn_mask = torch.randn(1, 1, 20, 20) prefill_metadata.chunked_context = MagicMock() prefill_metadata.chunked_context.seq_tot = [10, 10] prefill_metadata.chunked_context.starts = [0, 10] prefill_metadata.chunked_context.chunk_seq_lens_npu = [10, 10] prefill_metadata.chunked_context.chunk_actual_seq_lengths_kv_list = [[10], [10]] prefill_metadata.block_table = torch.randint(0, 100, (2, 4)) attn_metadata.prefill = prefill_metadata mock_device_operator.kv_cache_load = MagicMock() mock_fia.return_value = ( torch.randn(batch_size, self.impl.num_heads, self.impl.v_head_dim), torch.randn(self.impl.num_heads, batch_size), ) mock_npu_attention_update.return_value = ( torch.randn(batch_size, self.impl.num_heads, self.impl.v_head_dim), None, ) mock_kv_b_proj = MagicMock() # create [toks, num_heads, qk_nope_head_dim + v_head_dim] tensor toks = 10 kv_nope_shape = (toks, self.impl.num_heads, self.impl.qk_nope_head_dim + self.impl.v_head_dim) mock_kv_b_proj.return_value = (torch.randn(kv_nope_shape), None) self.impl.kv_b_proj = mock_kv_b_proj result = self.impl._forward_prefill(q_nope, q_pe, k_nope, k_pe, value, kv_c_and_k_pe_cache, attn_metadata) # verify result shape self.assertEqual(result.shape[0], batch_size) self.assertEqual(result.shape[1], self.impl.num_heads * self.impl.v_head_dim) @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") @patch("vllm_ascend.attention.mla_v1.DeviceOperator") @patch("torch_npu.npu_fused_infer_attention_score") def test_forward_prefill_non_power_of_two_heads(self, mock_fia, mock_device_operator, mock_get_current_vllm_config): """Test prefill with non-power-of-2 heads uses concat instead of query_rope/key_rope kwargs.""" mock_get_current_vllm_config.return_value = MagicMock() num_heads = 20 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": MagicMock(), "rotary_emb": MagicMock(), } impl = AscendMLAImpl( num_heads=num_heads, head_size=1024, scale=0.1, num_kv_heads=num_heads, alibi_slopes=None, sliding_window=None, kv_cache_dtype="auto", blocksparse_params=None, logits_soft_cap=None, attn_type=None, kv_sharing_target_layer_name=None, **kwargs, ) batch_size = 2 q_nope = torch.randn(batch_size, num_heads, impl.qk_nope_head_dim) q_pe = torch.randn(batch_size, num_heads, impl.qk_rope_head_dim) k_nope = torch.randn(batch_size, num_heads, impl.qk_nope_head_dim) k_pe = torch.randn(batch_size, num_heads, impl.qk_rope_head_dim) value = torch.randn(batch_size, num_heads, impl.v_head_dim) kv_c_and_k_pe_cache = [torch.randn(10, 1, 1, 192), torch.randn(10, 1, 1, 32)] attn_metadata = MagicMock() prefill_metadata = MagicMock() prefill_metadata.actual_seq_lengths_q = [10, 20] prefill_metadata.attn_mask = torch.randn(1, 1, 20, 20) prefill_metadata.chunked_context = None attn_metadata.prefill = prefill_metadata mock_device_operator.kv_cache_load = MagicMock() mock_fia.return_value = ( torch.randn(batch_size, num_heads, impl.v_head_dim), torch.randn(num_heads, batch_size), ) result = impl._forward_prefill(q_nope, q_pe, k_nope, k_pe, value, kv_c_and_k_pe_cache, attn_metadata) # FIA should be called without query_rope/key_rope when head_padding > 0 mock_fia.assert_called_once() call_kwargs = mock_fia.call_args.kwargs self.assertNotIn("query_rope", call_kwargs) self.assertNotIn("key_rope", call_kwargs) self.assertEqual(call_kwargs.get("num_heads"), num_heads) self.assertEqual(result.shape, (batch_size, num_heads * impl.v_head_dim)) @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading(self, mock_format_cast): layer = MagicMock(spec=LinearBase) layer.input_size_per_partition = 10 quant_method = MagicMock(spec=UnquantizedLinearMethod) layer.quant_method = quant_method shape_0 = self.impl.num_heads * (self.impl.qk_nope_head_dim + self.impl.v_head_dim) shape_1 = self.impl.kv_lora_rank layer.weight = torch.randn(shape_0, shape_1) self.impl.kv_b_proj = layer mock_format_cast.return_value = layer.weight self.impl.process_weights_after_loading(torch.bfloat16) self.assertEqual(self.impl.W_UK_T.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UK_T.shape[1], self.impl.qk_nope_head_dim) self.assertEqual(self.impl.W_UK_T.shape[2], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UV.shape[1], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[2], self.impl.v_head_dim) @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_mlapo(self, mock_format_cast): # test with enable_mlapo=True layer = MagicMock(spec=LinearBase) layer.input_size_per_partition = 10 quant_method = MagicMock(spec=UnquantizedLinearMethod) layer.quant_method = quant_method shape_0 = self.impl.num_heads * (self.impl.qk_nope_head_dim + self.impl.v_head_dim) shape_1 = self.impl.kv_lora_rank layer.weight = torch.randn(shape_0, shape_1) self.impl.kv_b_proj = layer mock_format_cast.return_value = layer.weight self.impl.enable_mlapo = True self.impl.process_weights_after_loading(torch.bfloat16) self.assertEqual(self.impl.W_UK_T.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UK_T.shape[1], self.impl.qk_nope_head_dim) self.assertEqual(self.impl.W_UK_T.shape[2], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UV.shape[1], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[2], self.impl.v_head_dim) @patch("vllm_ascend.attention.mla_v1.get_ascend_device_type") @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_mlapo_a5(self, mock_format_cast, mock_get_ascend_device_type): # test with enable_mlapo=True and device_type=A5 layer = MagicMock(spec=LinearBase) layer.input_size_per_partition = 10 quant_method = MagicMock(spec=UnquantizedLinearMethod) layer.quant_method = quant_method shape_0 = self.impl.num_heads * (self.impl.qk_nope_head_dim + self.impl.v_head_dim) shape_1 = self.impl.kv_lora_rank layer.weight = torch.randn(shape_0, shape_1) self.impl.kv_b_proj = layer mock_format_cast.return_value = layer.weight self.impl.enable_mlapo = True mock_fused_qkv_a_proj = MagicMock() mock_quant_method = MagicMock() from vllm_ascend.attention.mla_v1 import AscendW8A8LinearMethod mock_quant_method.quant_method = MagicMock(spec=AscendW8A8LinearMethod) mock_fused_qkv_a_proj.quant_method = mock_quant_method self.impl.fused_qkv_a_proj = mock_fused_qkv_a_proj # set device_type=A5 from vllm_ascend.attention.mla_v1 import AscendDeviceType mock_get_ascend_device_type.return_value = AscendDeviceType.A5 self.impl._process_weights_for_fused_mlapo_a5 = MagicMock() self.impl.process_weights_after_loading(torch.bfloat16) self.impl._process_weights_for_fused_mlapo_a5.assert_called_once_with(torch.bfloat16) self.assertEqual(self.impl.W_UK_T.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UK_T.shape[1], self.impl.qk_nope_head_dim) self.assertEqual(self.impl.W_UK_T.shape[2], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UV.shape[1], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[2], self.impl.v_head_dim) @patch("vllm_ascend.attention.mla_v1.get_ascend_device_type") @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_mlapo_non_a5(self, mock_format_cast, mock_get_ascend_device_type): # test with enable_mlapo=True and device_type!=A5 layer = MagicMock(spec=LinearBase) layer.input_size_per_partition = 10 quant_method = MagicMock(spec=UnquantizedLinearMethod) layer.quant_method = quant_method shape_0 = self.impl.num_heads * (self.impl.qk_nope_head_dim + self.impl.v_head_dim) shape_1 = self.impl.kv_lora_rank layer.weight = torch.randn(shape_0, shape_1) self.impl.kv_b_proj = layer mock_format_cast.return_value = layer.weight self.impl.enable_mlapo = True mock_fused_qkv_a_proj = MagicMock() mock_quant_method = MagicMock() from vllm_ascend.attention.mla_v1 import AscendW8A8LinearMethod mock_quant_method.quant_method = MagicMock(spec=AscendW8A8LinearMethod) mock_fused_qkv_a_proj.quant_method = mock_quant_method self.impl.fused_qkv_a_proj = mock_fused_qkv_a_proj from vllm_ascend.attention.mla_v1 import AscendDeviceType mock_get_ascend_device_type.return_value = AscendDeviceType.A2 self.impl._process_weights_for_fused_mlapo = MagicMock() self.impl.process_weights_after_loading(torch.bfloat16) self.impl._process_weights_for_fused_mlapo.assert_called_once_with(torch.bfloat16) self.assertEqual(self.impl.W_UK_T.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UK_T.shape[1], self.impl.qk_nope_head_dim) self.assertEqual(self.impl.W_UK_T.shape[2], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UV.shape[1], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[2], self.impl.v_head_dim) @patch("vllm_ascend.attention.mla_v1.maybe_trans_nz") @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_fa_quant(self, mock_format_cast, mock_maybe_trans_nz): # test with enable_mlapo=False and fa_quant_layer=True layer = MagicMock(spec=LinearBase) layer.input_size_per_partition = 10 quant_method = MagicMock(spec=UnquantizedLinearMethod) layer.quant_method = quant_method shape_0 = self.impl.num_heads * (self.impl.qk_nope_head_dim + self.impl.v_head_dim) shape_1 = self.impl.kv_lora_rank layer.weight = torch.randn(shape_0, shape_1) self.impl.kv_b_proj = layer mock_format_cast.return_value = layer.weight self.impl.enable_mlapo = False self.impl.fa_quant_layer = True self.impl._process_weights_for_fused_fa_quant = MagicMock() mock_maybe_trans_nz.return_value = torch.randn(1, 2, 3) self.impl.process_weights_after_loading(torch.bfloat16) self.impl._process_weights_for_fused_fa_quant.assert_called_once() self.assertEqual(self.impl.W_UK_T.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UK_T.shape[1], self.impl.qk_nope_head_dim) self.assertEqual(self.impl.W_UK_T.shape[2], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UV.shape[1], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[2], self.impl.v_head_dim) @patch("vllm_ascend.attention.mla_v1.post_process_after_loading_for_shard_weight_series") @patch("vllm_ascend.attention.mla_v1.is_hidden_layer") @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_layer_sharding( self, mock_format_cast, mock_is_hidden_layer, mock_post_process ): # test with layer_sharding_kwargs!=None layer = MagicMock(spec=LinearBase) layer.input_size_per_partition = 10 quant_method = MagicMock(spec=UnquantizedLinearMethod) layer.quant_method = quant_method shape_0 = self.impl.num_heads * (self.impl.qk_nope_head_dim + self.impl.v_head_dim) shape_1 = self.impl.kv_lora_rank layer.weight = torch.randn(shape_0, shape_1) self.impl.kv_b_proj = layer mock_format_cast.return_value = layer.weight self.impl.enable_mlapo = False self.impl.fa_quant_layer = False mock_layer1 = MagicMock() mock_layer2 = MagicMock() self.impl.layer_sharding_kwargs = [mock_layer1, mock_layer2] mock_is_hidden_layer.side_effect = [True, False] self.impl.process_weights_after_loading(torch.bfloat16) self.assertEqual(mock_is_hidden_layer.call_count, 2) mock_post_process.assert_called_once_with(mock_layer1) self.assertEqual(self.impl.W_UK_T.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UK_T.shape[1], self.impl.qk_nope_head_dim) self.assertEqual(self.impl.W_UK_T.shape[2], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[0], self.impl.num_heads) self.assertEqual(self.impl.W_UV.shape[1], self.impl.kv_lora_rank) self.assertEqual(self.impl.W_UV.shape[2], self.impl.v_head_dim) def test_compute_prefill_context_none(self): batch_size = 4 kv_cache = torch.randn(10, 1, 1, 192) query = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_head_dim) metadata = MagicMock() metadata.prefill = None prefix_out = torch.randn(2, 16, 128) prefix_lse = torch.randn(2, 16, 8) q_pe = query[..., self.impl.qk_nope_head_dim :] q_nope = query[..., : self.impl.qk_nope_head_dim] out, lse = self.impl._compute_prefill_context(q_nope, q_pe, kv_cache, 32, metadata, prefix_out, prefix_lse) self.assertTrue(torch.equal(prefix_out, out)) self.assertTrue(torch.equal(prefix_lse, lse)) def test_compute_prefill_context_empty_iters(self): # test compute_prefill_context with iters == 0 batch_size = 4 kv_cache = [torch.randn(10, 1, 1, 192), torch.randn(10, 1, 1, 32)] query = torch.randn(batch_size, self.impl.num_heads, self.impl.qk_head_dim) # Create a mock metadata where prefill_metadata.chunked_context.seq_tot is an empty list, so that iters == 0 metadata = MagicMock() prefill_metadata = MagicMock() chunked_context = MagicMock() chunked_context.seq_tot = [] # iters == 0 prefill_metadata.chunked_context = chunked_context metadata.prefill = prefill_metadata prefix_out = torch.randn(2, 16, 128) prefix_lse = torch.randn(2, 16, 8) q_pe = query[..., self.impl.qk_nope_head_dim :] q_nope = query[..., : self.impl.qk_nope_head_dim] out, lse = self.impl._compute_prefill_context(q_nope, q_pe, kv_cache, 32, metadata, prefix_out, prefix_lse) self.assertTrue(torch.equal(prefix_out, out)) self.assertTrue(torch.equal(prefix_lse, lse)) @patch("torch_npu.atb.npu_paged_cache_load") @patch("torch_npu.npu_attention_update") @patch("torch_npu.npu_fused_infer_attention_score") def test_compute_prefill_context(self, mock_fia, mock_update, mock_load): S, N, D, VD = 2, self.impl.num_heads, self.impl.qk_head_dim, self.impl.v_head_dim _, AND = self.impl.qk_rope_head_dim, self.impl.qk_nope_head_dim latent_kv_dim = self.impl.kv_lora_rank num_blocks, block_size = 100, 20 query = torch.randn(S, N, D) q_nope = query[..., : self.impl.qk_nope_head_dim] q_pe = query[..., self.impl.qk_nope_head_dim :] kv_cache_0 = torch.randn(num_blocks, block_size, N, latent_kv_dim) kv_cache_1 = torch.randn(num_blocks, block_size, N, D) kv_cache = [kv_cache_0, kv_cache_1] prefix_out = torch.randn(S, N, VD) prefix_lse = torch.randn(N, S) self.impl.kv_b_proj.return_value = (torch.randn(8, N, VD + AND),) # Mock FIA to return output and lse mock_fia.return_value = (torch.randn(S, N, VD), torch.randn(N, S)) # Mock attention_update to return merged output mock_update.return_value = (torch.randn(S * N, VD), None) chunk_ctx = MagicMock() chunk_ctx.seq_tot = [8] chunk_ctx.chunk_seq_lens = [torch.tensor([8])] chunk_ctx.chunk_seq_lens_npu = [torch.tensor([8])] chunk_ctx.starts = [torch.tensor([0])] prefill_meta = MagicMock() prefill_meta.chunked_context = chunk_ctx prefill_meta.query_lens = torch.tensor([S]) prefill_meta.block_table = torch.randint(0, 100, (S, 4)) meta = MagicMock() meta.prefill = prefill_meta self.impl.prefill_mask = torch.triu(torch.ones(512, 512, device=q_nope.device, dtype=q_nope.dtype), 1) out, lse = self.impl._compute_prefill_context(q_nope, q_pe, kv_cache, 32, meta, prefix_out, prefix_lse) mock_load.assert_called_once() mock_fia.assert_called_once() mock_update.assert_called_once() self.assertEqual(out.shape, prefix_out.shape) @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") @patch("torch_npu.atb.npu_paged_cache_load") @patch("torch_npu.npu_attention_update") @patch("torch_npu.npu_fused_infer_attention_score") def test_compute_prefill_context_non_power_of_two_heads( self, mock_fia, mock_update, mock_load, mock_get_current_vllm_config ): """Test prefill context with non-power-of-2 heads uses concat for query and key.""" mock_get_current_vllm_config.return_value = MagicMock() num_heads = 20 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": MagicMock(), "rotary_emb": MagicMock(), } impl = AscendMLAImpl( num_heads=num_heads, head_size=1024, scale=0.1, num_kv_heads=num_heads, alibi_slopes=None, sliding_window=None, kv_cache_dtype="auto", blocksparse_params=None, logits_soft_cap=None, attn_type=None, kv_sharing_target_layer_name=None, **kwargs, ) S, N, D, VD = 2, num_heads, impl.qk_head_dim, impl.v_head_dim latent_kv_dim = impl.kv_lora_rank num_blocks, block_size = 100, 20 query = torch.randn(S, N, D) q_nope = query[..., : impl.qk_nope_head_dim] q_pe = query[..., impl.qk_nope_head_dim :] kv_cache_0 = torch.randn(num_blocks, block_size, N, latent_kv_dim) kv_cache_1 = torch.randn(num_blocks, block_size, N, D) kv_cache = [kv_cache_0, kv_cache_1] prefix_out = torch.randn(S, N, VD) prefix_lse = torch.randn(N, S) impl.kv_b_proj.return_value = (torch.randn(8, N, VD + impl.qk_nope_head_dim),) mock_fia.return_value = (torch.randn(S, N, VD), torch.randn(N, S)) mock_update.return_value = (torch.randn(S * N, VD), None) chunk_ctx = MagicMock() chunk_ctx.seq_tot = [8] chunk_ctx.chunk_seq_lens = [torch.tensor([8])] chunk_ctx.chunk_seq_lens_npu = [torch.tensor([8])] chunk_ctx.starts = [torch.tensor([0])] chunk_ctx.chunk_actual_seq_lengths_kv_list = [[8]] prefill_meta = MagicMock() prefill_meta.chunked_context = chunk_ctx prefill_meta.query_lens = torch.tensor([S]) prefill_meta.block_table = torch.randint(0, 100, (S, 4)) meta = MagicMock() meta.prefill = prefill_meta out, lse = impl._compute_prefill_context(q_nope, q_pe, kv_cache, 32, meta, prefix_out, prefix_lse) mock_fia.assert_called_once() call_kwargs = mock_fia.call_args.kwargs self.assertNotIn("query_rope", call_kwargs) self.assertNotIn("key_rope", call_kwargs) self.assertEqual(out.shape, prefix_out.shape) @patch("vllm_ascend.ascend_forward_context.get_forward_context") @patch("vllm_ascend.attention.mla_v1.AscendMLAImpl._v_up_proj") @patch("torch_npu.npu_fused_infer_attention_score_v2") def test_forward_decode_without_graph( self, mock_npu_fused_infer_attention_score_v2, mock_up_proj, mock_get_forward_context ): num_tokens = 100 block_size = 4 q_nope = torch.randn(num_tokens, self.impl.num_heads, self.impl.qk_nope_head_dim) q_pe = torch.randn(num_tokens, self.impl.num_heads, self.impl.qk_rope_head_dim) k_nope = torch.randn(num_tokens, self.impl.num_heads, self.impl.qk_nope_head_dim) k_pe = torch.randn(num_tokens, self.impl.num_heads, self.impl.qk_rope_head_dim) metadata = MagicMock() metadata.decode = MagicMock() metadata.decode.block_table = MagicMock() metadata.decode.actual_seq_lengths = 10 mock_npu_fused_infer_attention_score_v2.return_value = [ torch.randn(num_tokens, self.impl.num_heads, self.impl.kv_lora_rank), None, ] mock_up_proj.return_value = torch.randn(num_tokens, self.impl.num_heads, self.impl.v_head_dim) mock_get_forward_context.return_value = MagicMock(capturing=False) result = self.impl._forward_decode(q_nope, q_pe, k_nope, k_pe, block_size, metadata) self.assertEqual(result.shape[0], num_tokens) self.assertEqual(result.shape[1], self.impl.num_heads) self.assertEqual(result.shape[2], self.impl.v_head_dim) mock_up_proj.assert_called_once() mock_npu_fused_infer_attention_score_v2.assert_called_once() @patch("torch.ops.vllm.maybe_all_gather_and_maybe_unpad") @patch("vllm_ascend.attention.mla_v1.get_weight_prefetch_method", return_value=MagicMock()) def test_mla_preprocess(self, mock_get_weight_prefetch_method, mock_maybe_all_gather_and_maybe_unpad): mock_maybe_all_gather_and_maybe_unpad.side_effect = lambda x, label: x batch_size = 4 seq_len = 8 hidden_size = 1024 hidden_states = torch.randn(batch_size * seq_len, hidden_size) kv_cache = MagicMock() attn_metadata = MagicMock() attn_metadata.num_decodes = 2 attn_metadata.num_prefills = 2 attn_metadata.num_decode_tokens = 2 attn_metadata.num_actual_tokens = 4 num_prefill_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) 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.num_heads, self.impl.qk_rope_head_dim ) self.impl.kv_a_proj_with_mqa = MagicMock() self.impl.kv_a_proj_with_mqa.return_value = [ torch.randn(num_prefill_tokens, 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( num_prefill_tokens, self.impl.num_heads, self.impl.qk_rope_head_dim + self.impl.kv_lora_rank + self.impl.q_lora_rank, ) ] self.impl.q_proj = MagicMock() self.impl.q_proj.return_value = [torch.randn(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(num_prefill_tokens, 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(num_prefill_tokens, self.impl.num_heads, self.impl.qk_rope_head_dim), torch.randn(num_prefill_tokens, self.impl.num_heads, self.impl.kv_lora_rank), ] self.impl._q_proj_and_k_up_proj = MagicMock() self.impl._q_proj_and_k_up_proj.return_value = [MagicMock(), MagicMock()] self.impl.num_kv_heads = self.impl.num_heads 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.assertIsNotNone(prefill_res) @patch("torch_npu.npu_kv_rmsnorm_rope_cache") def test_exec_kv_prefill(self, mock_kv_rmsnorm_rope_cache): B = 2 N = self.impl.num_kv_heads D = self.impl.kv_lora_rank + self.impl.qk_rope_head_dim kv_no_split = torch.randn(B, N, D) self.impl.enable_kv_nz = None self.impl.kv_a_layernorm.weight = MagicMock() self.impl.kv_a_layernorm.variance_epsilon = MagicMock() cos = MagicMock() sin = MagicMock() slots = MagicMock() kv_cache = [MagicMock(), MagicMock()] mock_kv_rmsnorm_rope_cache.return_value = [ None, None, torch.randn(B, N, 1, self.impl.qk_rope_head_dim), torch.randn(B, N, 1, self.impl.kv_lora_rank), ] k_pe, k_nope = self.impl.exec_kv_prefill(kv_no_split, cos, sin, kv_cache, slots) self.assertEqual(k_pe.shape[-1], self.impl.qk_rope_head_dim) self.assertEqual(k_nope.shape[-1], self.impl.kv_lora_rank) @patch("torch_npu.npu_kv_rmsnorm_rope_cache") def test_exec_kv_prefill_with_fa_quant(self, mock_kv_rmsnorm_rope_cache): # if fa_quant_layer is True B = 2 N = self.impl.num_kv_heads D = self.impl.kv_lora_rank + self.impl.qk_rope_head_dim kv_no_split = torch.randn(B, N, D) self.impl.enable_kv_nz = None self.impl.fa_quant_layer = True self.impl.kv_a_layernorm.weight = MagicMock() self.impl.kv_a_layernorm.variance_epsilon = MagicMock() cos = MagicMock() sin = MagicMock() slots = MagicMock() kv_cache = [MagicMock(), MagicMock()] block_size = 1 mock_kv_rmsnorm_rope_cache.return_value = [ None, None, torch.randn(B, N, block_size, self.impl.qk_rope_head_dim), torch.randn(B, N, block_size, self.impl.kv_lora_rank), ] k_pe, k_nope = self.impl.exec_kv_prefill(kv_no_split, cos, sin, kv_cache, slots) self.assertEqual(k_pe.shape[-1], self.impl.qk_rope_head_dim) self.assertEqual(k_nope.shape[-1], self.impl.kv_lora_rank) @patch("torch_npu.npu_kv_rmsnorm_rope_cache") def test_exec_kv_decode(self, mock_kv_rmsnorm_rope_cache): B = 2 N = self.impl.num_kv_heads D = self.impl.kv_lora_rank + self.impl.qk_rope_head_dim kv_no_split = torch.randn(B, N, D) self.impl.enable_kv_nz = None self.impl.kv_a_layernorm.weight = MagicMock() self.impl.kv_a_layernorm.variance_epsilon = MagicMock() cos = MagicMock() sin = MagicMock() slots = MagicMock() kv_cache = [MagicMock(), MagicMock()] mock_kv_rmsnorm_rope_cache.return_value = [ torch.randn(B, N, 1, self.impl.qk_rope_head_dim), torch.randn(B, N, 1, self.impl.kv_lora_rank), None, None, ] k_pe, k_nope = self.impl.exec_kv_decode(kv_no_split, cos, sin, kv_cache, slots) self.assertEqual(k_pe.shape[-1], self.impl.qk_rope_head_dim) self.assertEqual(k_nope.shape[-1], self.impl.kv_lora_rank) @patch("vllm_ascend.ascend_forward_context.get_forward_context") @patch("torch_npu.npu_fused_infer_attention_score_v2") def test_forward_decode(self, mock_npu_fused_infer_attention_score_v2, mock_get_forward_context): B = 2 N = self.impl.num_kv_heads BS = 100 HD = self.impl.v_head_dim self.impl.kv_lora_rank = 256 self.impl.spec_token_num = 1 self.impl._v_up_proj = MagicMock() self.impl._v_up_proj.return_value = torch.randn(B, N, HD) 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(BS, N, self.impl.kv_lora_rank) k_pe = torch.randn(BS, N, self.impl.qk_rope_head_dim) attn_metadata = MagicMock() attn_metadata.attn_state = AscendAttentionState.SpecDecoding attn_metadata.decode = MagicMock() attn_metadata.decode.actual_seq_qlen = MagicMock() attn_metadata.decode.actual_seq_kvlen = MagicMock() self.impl.enable_kv_nz = True mock_npu_fused_infer_attention_score_v2.return_value = [torch.randn(B, N, self.impl.kv_lora_rank), None] mock_get_forward_context.return_value = MagicMock(capturing=False) 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], N) self.assertEqual(result.shape[2], HD) @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") @patch("vllm_ascend.ascend_forward_context.get_forward_context") @patch("torch_npu.npu_fused_infer_attention_score_v2") def test_forward_decode_non_power_of_two_heads( self, mock_npu_fused_infer_attention_score_v2, mock_get_forward_context, mock_get_current_vllm_config ): """Test decode with non-power-of-2 heads pads to next power of 2 and slices output.""" mock_get_current_vllm_config.return_value = MagicMock() num_heads = 20 kwargs = { "kv_lora_rank": 256, "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": MagicMock(), "rotary_emb": MagicMock(), } impl = AscendMLAImpl( num_heads=num_heads, head_size=1024, scale=0.1, num_kv_heads=num_heads, alibi_slopes=None, sliding_window=None, kv_cache_dtype="auto", blocksparse_params=None, logits_soft_cap=None, attn_type=None, kv_sharing_target_layer_name=None, **kwargs, ) B = 2 BS = 100 HD = impl.v_head_dim impl.spec_token_num = 1 impl._v_up_proj = MagicMock() impl._v_up_proj.return_value = torch.randn(B, num_heads, HD) q_nope = torch.randn(B, num_heads, impl.qk_nope_head_dim) q_pe = torch.randn(B, num_heads, impl.qk_rope_head_dim) k_nope = torch.randn(BS, num_heads, impl.kv_lora_rank) k_pe = torch.randn(BS, num_heads, impl.qk_rope_head_dim) attn_metadata = MagicMock() attn_metadata.attn_state = AscendAttentionState.SpecDecoding attn_metadata.decode = MagicMock() attn_metadata.decode.actual_seq_qlen = MagicMock() attn_metadata.decode.actual_seq_kvlen = MagicMock() impl.enable_kv_nz = True impl.fa_quant_layer = False # Return padded output so slice logic works mock_npu_fused_infer_attention_score_v2.return_value = [ torch.randn(impl.num_heads_padded, B, impl.kv_lora_rank), None, ] mock_get_forward_context.return_value = MagicMock(capturing=False) result = 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], num_heads) self.assertEqual(result.shape[2], HD) # Verify num_query_heads passed to FIA is padded mock_npu_fused_infer_attention_score_v2.assert_called_once() call_kwargs = mock_npu_fused_infer_attention_score_v2.call_args.kwargs self.assertEqual(call_kwargs.get("num_query_heads"), impl.num_heads_padded) @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config") @patch("vllm_ascend.ascend_forward_context.get_forward_context") @patch("torch_npu.npu_fused_infer_attention_score_v2") def test_forward_decode_non_power_of_two_heads_normal( self, mock_npu_fused_infer_attention_score_v2, mock_get_forward_context, mock_get_current_vllm_config ): """Test normal decode (BNSD_NBSD) with non-power-of-2 heads pads q and slices output.""" mock_get_current_vllm_config.return_value = MagicMock() num_heads = 20 kwargs = { "kv_lora_rank": 256, "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": MagicMock(), "rotary_emb": MagicMock(), } impl = AscendMLAImpl( num_heads=num_heads, head_size=1024, scale=0.1, num_kv_heads=num_heads, alibi_slopes=None, sliding_window=None, kv_cache_dtype="auto", blocksparse_params=None, logits_soft_cap=None, attn_type=None, kv_sharing_target_layer_name=None, **kwargs, ) B = 2 BS = 100 HD = impl.v_head_dim impl.spec_token_num = 1 impl._v_up_proj = MagicMock() impl._v_up_proj.return_value = torch.randn(B, num_heads, HD) q_nope = torch.randn(B, num_heads, impl.qk_nope_head_dim) q_pe = torch.randn(B, num_heads, impl.qk_rope_head_dim) k_nope = torch.randn(BS, num_heads, impl.kv_lora_rank) k_pe = torch.randn(BS, num_heads, impl.qk_rope_head_dim) attn_metadata = MagicMock() attn_metadata.attn_state = AscendAttentionState.DecodeOnly attn_metadata.decode = MagicMock() attn_metadata.decode.actual_seq_qlen = MagicMock() attn_metadata.decode.actual_seq_kvlen = MagicMock() attn_metadata.decode.block_table = MagicMock() impl.enable_kv_nz = False impl.fa_quant_layer = False impl.speculative_config = None mock_npu_fused_infer_attention_score_v2.return_value = [ torch.randn(impl.num_heads_padded, B, 1, impl.kv_lora_rank), None, ] mock_get_forward_context.return_value = MagicMock(capturing=False) result = 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], num_heads) self.assertEqual(result.shape[2], HD) mock_npu_fused_infer_attention_score_v2.assert_called_once() call_kwargs = mock_npu_fused_infer_attention_score_v2.call_args.kwargs self.assertEqual(call_kwargs.get("num_query_heads"), impl.num_heads_padded) @patch("vllm_ascend.ascend_forward_context.get_forward_context") @patch("torch_npu.npu_fused_infer_attention_score_v2") def test_forward_decode_with_fa_quant(self, mock_npu_fused_infer_attention_score_v2, mock_get_forward_context): # test fa_quant_layer is True B = 2 N = self.impl.num_heads # use num_heads instead of num_kv_heads BS = 100 HD = self.impl.v_head_dim self.impl.kv_lora_rank = 256 self.impl.spec_token_num = 1 self.impl._v_up_proj = MagicMock() self.impl._v_up_proj.return_value = torch.randn(B, self.impl.num_kv_heads, HD) 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(BS, self.impl.num_kv_heads, self.impl.kv_lora_rank) k_pe = torch.randn(BS, self.impl.num_kv_heads, self.impl.qk_rope_head_dim) attn_metadata = MagicMock() attn_metadata.attn_state = AscendAttentionState.SpecDecoding attn_metadata.decode = MagicMock() attn_metadata.decode.actual_seq_qlen = MagicMock() attn_metadata.decode.actual_seq_kvlen = MagicMock() attn_metadata.decode.actual_seq_lengths_q = [10, 20] attn_metadata.decode.attn_mask = MagicMock() self.impl.fa_quant_layer = True self.impl.speculative_config = MagicMock() self.impl.fak_descale_float = torch.randn(1) # add fak_descale_float attribute mock_npu_fused_infer_attention_score_v2.return_value = [ torch.randn(B, self.impl.num_kv_heads, self.impl.kv_lora_rank), None, ] mock_get_forward_context.return_value = MagicMock(capturing=False) dequant_scale_q_nope = torch.randn(B, N) # shape is [B, num_heads] result = self.impl._forward_decode(q_nope, q_pe, k_nope, k_pe, BS, attn_metadata, dequant_scale_q_nope) self.assertEqual(result.shape[0], B) self.assertEqual(result.shape[1], self.impl.num_kv_heads) self.assertEqual(result.shape[2], HD)