bugfix

Signed-off-by: Angazenn <supperccell@163.com>

Author: Angazenn

vllm_ascend/compilation/acl_graph.py

@@ -213,8 +213,7 @@ def update_attn_params(update_stream, forward_context, runtime_shape):
                 output,
             ) = param
             seq_lens = forward_context.attn_metadata[key].seq_lens
-            torch.npu.graph_task_update_begin(update_stream, handle)
-            torch_npu._npu_paged_attention(
+            workspace = torch_npu._npu_paged_attention_get_workspace(
                 query=query,
                 key_cache=key_cache,
                 value_cache=value_cache,
@@ -223,8 +222,19 @@ def update_attn_params(update_stream, forward_context, runtime_shape):
                 scale_value=scale,
                 block_table=block_table,
                 context_lens=seq_lens,
-                out=output,
-                workspace=graph_params.workspaces.get(runtime_shape))
+                out=output)
+
+            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,
+                                           workspace=workspace)
             torch.npu.graph_task_update_end(update_stream)
 
             event.record(update_stream)