add pagedattention to support fullgraph.

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

Author: wangxiaoxin-sherie

tests/ut/attention/test_attention_v1.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, call
 
 import torch
 
@@ -409,6 +409,224 @@ def test_forward_decode_only(self, mock_paged_attention,
 
         mock_paged_attention.assert_called_once()
         assert output.shape == (10, 8 * 64)
+    
+    @patch('torch_npu._npu_paged_attention')
+    @patch('torch.npu.graph_task_group_end')
+    @patch('torch.npu.graph_task_group_begin')
+    @patch('torch.npu.ExternalEvent')
+    @patch('torch_npu.npu.current_stream')
+    @patch('torch_npu._npu_paged_attention_get_workspace')
+    def test_paged_attention_with_existing_workspace(self, mock_get_workspace, mock_current_stream,
+                                                     mock_external_event_class, mock_graph_begin,
+                                                     mock_graph_end, mock_paged_attention):
+        """
+        测试工作区已存在的情况。
+        """
+        # --- 设置 Mock 对象 ---
+        # 模拟 graph_params 和 attn_metadata
+        graph_params = MagicMock()
+        attn_metadata = MagicMock()
+        num_tokens = 10
+
+        # 模拟已存在的 workspace
+        graph_params.workspaces = {num_tokens: 10}
+        graph_params.events = {num_tokens: []}
+        graph_params.attn_params = {num_tokens: []}
+        graph_params.handles = {num_tokens: []}
+
+        # 模拟其他参数
+        query = torch.randn(2, 5, 8)  # [batch_size, seq_len, hidden_size]
+        key_cache = MagicMock()
+        value_cache = MagicMock()
+        num_kv_heads = 4
+        num_heads = 8
+        scale = 0.1
+        block_tables = MagicMock()
+        seq_lens = MagicMock()
+        output = torch.randn(2, 5, 8)
+
+        # self 对象 (模拟类实例)
+        self_obj = MagicMock()
+        self_obj.key_cache = key_cache
+        self_obj.value_cache = value_cache
+        self_obj.num_kv_heads = num_kv_heads
+        self_obj.num_heads = num_heads
+        self_obj.scale = scale
+
+        # Mock 流和事件
+        mock_stream = MagicMock()
+        mock_current_stream.return_value = mock_stream
+        mock_event_instance = MagicMock()
+        mock_external_event_class.return_value = mock_event_instance
+
+        # Mock graph_task_group_end 返回句柄
+        mock_handle = MagicMock()
+        mock_graph_end.return_value = mock_handle
+
+        # --- 调用被测试的代码逻辑 (模拟) ---
+        # 注意：这里直接模拟了代码逻辑，因为在实际应用中，这段代码可能在某个函数或方法内部
+        # 为了测试，我们需要将其封装或直接在这里模拟执行流程
+
+        # 1. 获取 workspace
+        workspace = graph_params.workspaces.get(num_tokens)
+        # 断言：工作区已存在，不应调用 get_workspace
+        mock_get_workspace.assert_not_called()
+        self.assertEqual(workspace, 10)
+
+        # 2. Handle graph capturing mode
+        stream = mock_current_stream()
+        event = mock_external_event_class()
+        event.wait(stream)
+        event.reset(stream)
+        graph_params.events[num_tokens].append(event)
+        graph_params.attn_params[num_tokens].append((
+            query,
+            self_obj.key_cache,
+            self_obj.value_cache,
+            self_obj.num_kv_heads,
+            self_obj.num_heads,
+            self_obj.scale,
+            attn_metadata.block_tables,
+            attn_metadata.seq_lens,
+            output,
+        ))
+
+        # 断言事件调用
+        mock_event_instance.wait.assert_called_once_with(mock_stream)
+        mock_event_instance.reset.assert_called_once_with(mock_stream)
+        self.assertEqual(len(graph_params.events[num_tokens]), 1)
+        self.assertEqual(len(graph_params.attn_params[num_tokens]), 1)
+
+        # 3. Execute graph task
+        mock_graph_begin(stream)
+        mock_paged_attention(
+            query=query,
+            key_cache=self_obj.key_cache,
+            value_cache=self_obj.value_cache,
+            num_kv_heads=self_obj.num_kv_heads,
+            num_heads=self_obj.num_heads,
+            scale_value=self_obj.scale,
+            block_table=attn_metadata.block_tables,
+            context_lens=attn_metadata.seq_lens,
+            out=output,
+            workspace=workspace
+        )
+        handle = mock_graph_end(stream)
+
+        # 断言图任务调用
+        mock_graph_begin.assert_called_once_with(mock_stream)
+        mock_graph_end.assert_called_once_with(mock_stream)
+        self.assertEqual(handle, mock_handle)
+        self.assertEqual(len(graph_params.handles[num_tokens]), 1)
+
+
+    @patch('torch_npu._npu_paged_attention')
+    @patch('torch.npu.graph_task_group_end')
+    @patch('torch.npu.graph_task_group_begin')
+    @patch('torch.npu.ExternalEvent')
+    @patch('torch_npu.npu.current_stream')
+    @patch('torch_npu._npu_paged_attention_get_workspace')
+    def test_paged_attention_with_new_workspace(self, mock_get_workspace, mock_current_stream,
+                                                mock_external_event_class, mock_graph_begin,
+                                                mock_graph_end, mock_paged_attention):
+        """
+        测试工作区不存在，需要创建的情况。
+        """
+        # --- 设置 Mock 对象 ---
+        graph_params = MagicMock()
+        attn_metadata = MagicMock()
+        num_tokens = 15
+
+        # 模拟不存在的 workspace
+        graph_params.workspaces = {}
+        graph_params.events = {num_tokens: []}
+        graph_params.attn_params = {num_tokens: []}
+        graph_params.handles = {num_tokens: []}
+
+
+        query = torch.randn(1, 3, 16)
+        key_cache = MagicMock()
+        value_cache = MagicMock()
+        num_kv_heads = 2
+        num_heads = 4
+        scale = 0.2
+        block_tables = MagicMock()
+        seq_lens = MagicMock()
+        output = torch.randn(1, 3, 16)
+
+        self_obj = MagicMock()
+        self_obj.key_cache = key_cache
+        self_obj.value_cache = value_cache
+        self_obj.num_kv_heads = num_kv_heads
+        self_obj.num_heads = num_heads
+        self_obj.scale = scale
+
+        mock_stream = MagicMock()
+        mock_current_stream.return_value = mock_stream
+        mock_event_instance = MagicMock()
+        mock_external_event_class.return_value = mock_event_instance
+
+        # 模拟创建新的 workspace
+        new_workspace = MagicMock()
+        mock_get_workspace.return_value = new_workspace
+
+        mock_handle = MagicMock()
+        mock_graph_end.return_value = mock_handle
+
+        graph_params.workspaces[num_tokens] = 10
+
+        # Handle graph capturing mode (同上，简化)
+        stream = mock_current_stream()
+        event = mock_external_event_class()
+        event.wait(stream)
+        event.reset(stream)
+        graph_params.events[num_tokens].append(event)
+        graph_params.attn_params[num_tokens].append((
+            query,
+            self_obj.key_cache,
+            self_obj.value_cache,
+            self_obj.num_kv_heads,
+            self_obj.num_heads,
+            self_obj.scale,
+            attn_metadata.block_tables,
+            attn_metadata.seq_lens,
+            output,
+        ))
+
+        # Execute graph task (同上，简化)
+        mock_graph_begin(stream)
+        mock_paged_attention(
+            query=query,
+            key_cache=self_obj.key_cache,
+            value_cache=self_obj.value_cache,
+            num_kv_heads=self_obj.num_kv_heads,
+            num_heads=self_obj.num_heads,
+            scale_value=self_obj.scale,
+            block_table=attn_metadata.block_tables,
+            context_lens=attn_metadata.seq_lens,
+            out=output,
+            workspace=10 # 使用新创建的 workspace
+        )
+        handle = mock_graph_end(stream)
+
+        # 断言图任务调用 (同上，简化)
+        mock_graph_begin.assert_called_once_with(mock_stream)
+        expected_paged_attn_call = call(
+            query=query,
+            key_cache=key_cache,
+            value_cache=value_cache,
+            num_kv_heads=num_kv_heads,
+            num_heads=num_heads,
+            scale_value=scale,
+            block_table=block_tables,
+            context_lens=seq_lens,
+            out=output,
+            workspace=new_workspace # 验证使用了新 workspace
+        )
+        mock_paged_attention.assert_called_once_with(**expected_paged_attn_call.kwargs)
+        mock_graph_end.assert_called_once_with(mock_stream)
+        self.assertEqual(handle, mock_handle)
+        self.assertEqual(len(graph_params.handles[num_tokens]), 1)
 
     @patch('torch_npu._npu_reshape_and_cache')
     @patch('torch_npu.npu_fused_infer_attention_score')

vllm_ascend/attention/attention_v1.py

@@ -432,13 +432,31 @@ def _forward_decode_only(
             forward_context: ForwardContext = get_forward_context()
             num_tokens = query.shape[0]
             if forward_context.capturing:
+                # Prepare tensors for attention output
+                # TODO: Refactor this to step-level instead of layer-level
+
+                # Get workspace from cache or calculate it if not present.
+                workspace = graph_params.workspaces.get(num_tokens)
+                if workspace is None:
+                    workspace = torch_npu._npu_paged_attention_get_workspace(
+                        query=query,
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        num_kv_heads=self.num_kv_heads,
+                        num_heads=self.num_heads,
+                        scale_value=self.scale,
+                        block_table=attn_metadata.block_tables,
+                        context_lens=attn_metadata.seq_lens,
+                        out=output)
+                    graph_params.workspaces[num_tokens] = workspace
+
+                # Handle graph capturing mode
                 stream = torch_npu.npu.current_stream()
 
                 event = torch.npu.ExternalEvent()
                 event.wait(stream)
                 event.reset(stream)
                 graph_params.events[num_tokens].append(event)
-
                 graph_params.attn_params[num_tokens].append((
                     query,
                     self.key_cache,
@@ -461,7 +479,8 @@ def _forward_decode_only(
                     scale_value=self.scale,
                     block_table=attn_metadata.block_tables,
                     context_lens=attn_metadata.seq_lens,
-                    out=output)
+                    out=output,
+                    workspace=workspace)
                 handle = torch.npu.graph_task_group_end(stream)
                 graph_params.handles[num_tokens].append(handle)
             else:

vllm_ascend/compilation/acl_graph.py

@@ -215,15 +215,17 @@ def update_attn_params(update_stream, forward_context, runtime_shape):
 
         with torch.npu.stream(update_stream):
             torch.npu.graph_task_update_begin(update_stream, handle)
-            torch_npu._npu_paged_attention(query=query,
-                                           key_cache=key_cache,
-                                           value_cache=value_cache,
-                                           num_kv_heads=num_kv_heads,
-                                           num_heads=num_heads,
-                                           scale_value=scale,
-                                           block_table=block_table,
-                                           context_lens=seq_lens,
-                                           out=output)
+            torch_npu._npu_paged_attention(
+                query=query,
+                key_cache=key_cache,
+                value_cache=value_cache,
+                num_kv_heads=num_kv_heads,
+                num_heads=num_heads,
+                scale_value=scale,
+                block_table=block_table,
+                context_lens=seq_lens,
+                out=output,
+                workspace=graph_params.workspaces.get(runtime_shape))
             torch.npu.graph_task_update_end(update_stream)
 
             event.record(update_stream)