[Bugfix]:replace npu_incre_flash_attention with npu_fused_infer_attention_score

Signed-off-by: p00465316 <panchao13@huawei.com>

Author: p00465316

tests/e2e/multicard/test_torchair_graph_mode.py

@@ -21,8 +21,10 @@
 """
 import os
 from typing import Dict
+from unittest.mock import patch
 
 from tests.e2e.conftest import VllmRunner
+from vllm_ascend.ascend_forward_context import _get_fused_moe_state
 
 os.environ["PYTORCH_NPU_ALLOC_CONF"] = "max_split_size_mb:256"
 
@@ -167,53 +169,47 @@ def _qwen_torchair_test_fixture(
     tp,
     enable_expert_parallel,
 ):
-    # The current access control does not support 16 cards,
-    # so the MC2 operator in Qwen's graph mode cannot run.
-    # Once 16-card support is available,
-    # this e2e can be switched to graph mode.
-    example_prompts = [
-        "Hello, my name is",
-        "The president of the United States is",
-        "The capital of France is",
-        "The future of AI is",
-    ]
-
-    additional_config = {
-        "torchair_graph_config": {
-            "enabled": False,
-        },
-        "ascend_scheduler_config": {
-            "enabled": True,
-        },
-        "refresh": True,
-    }
 
-    with VllmRunner(
-            model,
-            dtype="half",
-            tensor_parallel_size=tp,
-            distributed_executor_backend="mp",
-            enforce_eager=True,
-            additional_config=additional_config,
-            enable_expert_parallel=enable_expert_parallel,
-    ) as vllm_model:
-        # use greedy sampler to make sure the generated results are fix
-        vllm_output = vllm_model.generate_greedy(example_prompts, 5)
-
-    # NOTE: vllm-ascend/pangu-pro-moe-pruing is only part of PanguProMoE
-    # with 2 hidden layers, thus the golden results seems inaccurate.
-    # This will only change if accuracy changes with the official weights
-    # of PanguProMoE.
-    golden_results = [
-        'Hello, my name is Remempondeprecatedmiot忱',
-        'The president of the United States is Remem下的一个 rever ceremoni Segnali',
-        'The capital of France is Rememvoud administrativ Remem投',
-        'The future of AI isotope Segnali Zoeken精细化 supus',
-    ]
+    def stubbed_get_state(ep_size, with_prefill, is_deepseek_v3_r1):
+        return _get_fused_moe_state(16, with_prefill, is_deepseek_v3_r1)
+
+    with patch("vllm_ascend.ascend_forward_context._get_fused_moe_state",
+               stubbed_get_state):
+        # The current access control does not support 16 cards,
+        # so the MC2 operator in Qwen's graph mode cannot run.
+        # Once 16-card support is available,
+        # this e2e can be switched to graph mode.
+        example_prompts = [
+            "Hello, my name is",
+            "The president of the United States is",
+            "The capital of France is",
+            "The future of AI is",
+        ]
+
+        additional_config = {
+            "torchair_graph_config": {
+                "enabled": True,
+            },
+            "ascend_scheduler_config": {
+                "enabled": True,
+            },
+            "refresh": True,
+        }
 
-    assert len(golden_results) == len(vllm_output)
-    for i in range(len(vllm_output)):
-        print(f"Generated text: {vllm_output[i][1]!r}")
+        with VllmRunner(
+                model,
+                dtype="half",
+                tensor_parallel_size=tp,
+                distributed_executor_backend="mp",
+                enforce_eager=False,
+                additional_config=additional_config,
+                enable_expert_parallel=enable_expert_parallel,
+        ) as vllm_model:
+            # use greedy sampler to make sure the generated results are fix
+            vllm_output = vllm_model.generate_greedy(example_prompts, 5)
+
+        for i in range(len(vllm_output)):
+            print(f"Generated text: {vllm_output[i][1]!r}")
 
 
 def test_e2e_qwen2_with_torchair():

tests/ut/torchair/test_torchair_attention.py

@@ -0,0 +1,94 @@
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.attention.backends.abstract import AttentionType
+from vllm.distributed.parallel_state import GroupCoordinator
+
+from tests.ut.base import TestBase
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+from vllm_ascend.torchair.torchair_attention import \
+    AscendAttentionTorchairBackendImpl
+
+
+class TestAscendAttentionTorchairBackendImpl(TestBase):
+
+    @patch("torch.zeros")
+    @patch('vllm.distributed.parallel_state._TP',
+           new_callable=lambda: MagicMock(spec=GroupCoordinator))  # TODO
+    @patch("vllm.distributed.get_tensor_model_parallel_world_size",
+           return_value=2)  # TODO
+    @patch("vllm.config.get_current_vllm_config")  # TODO
+    @patch("vllm_ascend.torchair.torchair_mla.get_ascend_config")  # TODO
+    def setUp(self, ascend_config, vllm_config, mock_get_tp_size, mock_tp,
+              mock_zeros):
+        mock_tp.world_size = 2  # TODO
+        ascend_config.torchair_graph_config.enabled = True  # TODO
+        ascend_config.torchair_graph_config.enable_kv_nz = False  # TODO
+        speculative_config = MagicMock()
+        speculative_config.num_speculative_tokens = 4
+        vllm_config.speculative_config = speculative_config
+
+        num_heads = 32
+        head_size = 128  # TODO
+        scale = 0.1  # TODO
+        num_kv_heads = 4
+        kv_cache_dtype = "auto"
+        attn_type = AttentionType.DECODER
+        mock_zeros.return_value = torch.ones((),
+                                             device='cpu',
+                                             dtype=torch.int32)
+
+        self.impl = AscendAttentionTorchairBackendImpl(
+            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=attn_type,
+            kv_sharing_target_layer_name=None)
+
+    @patch("torch_npu.npu_scatter_nd_update_")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    def test_forward_with_decode_only(self, mock_fused, _):
+        layer = MagicMock()
+        layer._k_scale_float = 1.0
+        layer._v_scale_float = 1.0
+
+        seq_len = 1
+        num_tokens = 100
+        num_blocks = 256
+        block_size = 4
+
+        query = torch.randn(num_tokens, seq_len,
+                            self.impl.num_heads * self.impl.head_size)
+        key = torch.randn(num_tokens, seq_len,
+                          self.impl.num_kv_heads * self.impl.head_size)
+        value = torch.randn(num_tokens, seq_len,
+                            self.impl.num_kv_heads * self.impl.head_size)
+        kv_cache = (torch.randn(num_blocks, block_size,
+                                self.impl.num_heads * self.impl.head_size),
+                    torch.randn(num_blocks, block_size,
+                                self.impl.num_heads * self.impl.head_size))
+        output = torch.randn(num_tokens, self.impl.num_heads,
+                             self.impl.head_size)
+
+        decode = MagicMock()  # TODO
+        decode.seq_lens_list = [2] * num_tokens
+        decode.block_table = torch.ones(num_tokens, 8, dtype=torch.int32)
+        decode.attn_mask = None
+
+        metadata = MagicMock()
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.decode = decode
+
+        mock_fused.return_value = (torch.ones(num_tokens, self.impl.num_heads,
+                                              self.impl.head_size),
+                                   torch.ones(1))
+
+        result = self.impl.forward(layer, query, key, value, kv_cache,
+                                   metadata, output, False)
+        self.assertEqual(result.shape[0], num_tokens)

vllm_ascend/torchair/torchair_attention.py

@@ -436,17 +436,24 @@ def forward(
             block_size = key_cache.shape[1]
             query = query.view(num_tokens, 1,
                                self.num_heads * self.head_size).contiguous()
-            output = torch_npu.npu_incre_flash_attention(
-                query,
-                key_cache,
-                value_cache,
-                num_key_value_heads=self.num_kv_heads,
+            output, _ = torch_npu.npu_fused_infer_attention_score(
+                query=query,
+                key=key_cache,
+                value=value_cache,
+                query_rope=None,
+                key_rope=None,
                 num_heads=self.num_heads,
-                actual_seq_lengths=seq_lens,
-                scale_value=self.scale,
-                block_table=block_table,
+                num_key_value_heads=self.num_kv_heads,
                 input_layout='BSH',
-                block_size=block_size)
+                atten_mask=decode_meta.attn_mask,
+                sparse_mode=0,
+                scale=self.scale,
+                antiquant_mode=0,
+                antiquant_scale=None,
+                block_table=block_table,
+                block_size=block_size,
+                actual_seq_lengths_kv=seq_lens,
+            )
         else:
             raise NotImplementedError(
                 "Torchair graph mode with non-MLA attention backend is still experimental."