[BugFix]Add check for empty _ATTN_KEYS_BUFFER (vllm-project#11214)

### What this PR does / why we need it?
Adds a check for empty ``_ATTN_KEYS_BUFFER`` and corresponding tests.

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@a30addc
---------
Signed-off-by: zhiyu-wa <1959864813@qq.com>

Author: zhiyu-wa

tests/e2e/pull_request/four_card/test_graph_mode.py

@@ -366,6 +366,11 @@
     "capture_mem_tolerance": 1.5,
 }
 
+CASE_DS_ACLGRAPH_ENPU = {
+    **CASE_DS_ACLGRAPH,
+    "env_vars": {"ENPU_ENABLE": "true"},
+}
+
 # inherit from tests/e2e/pull_request/utils.py::compare_logprobs
 ATOL = 0.0689
 
@@ -493,6 +498,9 @@ def _run_worker_process(
         }
     )
 
+    for key, value in cur_case.get("env_vars", {}).items():
+        os.environ[key] = str(value)
+
     # Apply hooks and run inference
     with _install_spies(metrics):
         short_prompts = cur_case["prompts"]["short"]
@@ -595,7 +603,7 @@ def check_capture_mem(capture_mem, baseline_capture_mem=0.2, capture_mem_toleran
 
 
 @wait_until_npu_memory_free(0.7)
-@pytest.mark.parametrize("cur_case", [CASE_QWEN_ACLGRAPH, CASE_DS_ACLGRAPH])
+@pytest.mark.parametrize("cur_case", [CASE_QWEN_ACLGRAPH, CASE_DS_ACLGRAPH, CASE_DS_ACLGRAPH_ENPU])
 def test_aclgraph(cur_case: dict, monkeypatch: pytest.MonkeyPatch):
     # Counter doesn't work in default "spawn" mode
     metrics = None

tests/ut/attention/a2/test_attention_v1.py

@@ -3,6 +3,7 @@
 
 import torch
 
+import vllm_ascend.attention.attention_v1 as attn_module
 from tests.ut.base import TestBase
 from vllm_ascend.attention.attention_v1 import (
     AscendAttentionBackend,
@@ -605,3 +606,54 @@ def test_get_kvcomp_params_decode_hamming(self, mock_hamming, mock_reshape):
         # Verify the result is recorded in hamming_output_records
         self.assertTrue(torch.equal(kvcomp_metadata.hamming_output, torch.ones(2, 5)))
         self.assertIsNotNone(kvcomp_metadata.seq_lens_for_hamming)
+
+    @patch("torch.npu.stream")
+    @patch("torch.npu.graph_task_update_begin")
+    @patch("torch.npu.graph_task_update_end")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("vllm_ascend.attention.attention_v1.get_graph_params")
+    @patch("vllm_ascend.attention.attention_v1._EXTRA_CTX")
+    @patch("vllm_ascend.attention.attention_v1.using_paged_attention", return_value=False)
+    @patch("vllm_ascend.attention.attention_v1.needs_layer_aware_fia_graph_replay", return_value=False)
+    @patch("vllm_ascend.attention.attention_v1._ATTN_KEYS_BUFFER", new=[])
+    def test_update_graph_params(
+        self,
+        mock_needs_layer_aware_fia_graph_replay,
+        mock_using_paged_attention,
+        mock_EXTRA_CTX,
+        mock_get_graph_params,
+        mock_fia,
+        mock_graph_task_update_end,
+        mock_graph_task_update_begin,
+        mock_stream,
+    ):
+        """Test behavior when _ATTN_KEYS_BUFFER is [] after dummy_run."""
+
+        mock_EXTRA_CTX.sinks = False
+        mock_EXTRA_CTX.is_draft_model = False
+
+        param: list[MagicMock | None] = [MagicMock()] * 21
+        param[16] = None
+        param[20] = None
+
+        mock_get_graph_params.return_value.attn_params = {1: [tuple(param)] * 3}
+        mock_get_graph_params.return_value.handles = {1: [MagicMock()] * 3}
+        mock_get_graph_params.return_value.events = {1: [MagicMock()] * 3}
+
+        attn_metadata_keys = [
+            "model.layers.10.self_attn.attn",
+            "model.layers.2.self_attn.attn",
+            "model.layers.5.self_attn.attn",
+        ]
+        forward_context = MagicMock()
+        forward_context.attn_metadata = {key: MagicMock() for key in attn_metadata_keys}
+        # breakpoint()
+        self.impl.update_graph_params(self.mock_stream, forward_context, 1, self.mock_vllm_config)
+
+        expected = [
+            "model.layers.2.self_attn.attn",
+            "model.layers.5.self_attn.attn",
+            "model.layers.10.self_attn.attn",
+        ]
+        self.assertEqual(attn_module._ATTN_KEYS_BUFFER, expected)
+        self.assertEqual(mock_fia.out.call_count, 3)

vllm_ascend/attention/attention_v1.py

@@ -644,7 +644,7 @@ def update_graph_params(
                     global _ATTN_KEYS_BUFFER
                     if attn_keys_length == 0:
                         return
-                    if _ATTN_KEYS_BUFFER is None or len(_ATTN_KEYS_BUFFER) != attn_keys_length:
+                    if not _ATTN_KEYS_BUFFER or len(_ATTN_KEYS_BUFFER) != attn_keys_length:
                         import regex as re
 
                         def extract_layer_index(key: str) -> int: