通用 GQA/PA性能优化PR #15579已合入 main作者 solin28825 · 2026-09-04 · 改动 1 文件 +42/-10
图模式下走 _npu_paged_attention（PA）路径时，每步 decode 前要更新图参数
（update_paged_attention_graph_param，attention/attention_v1.py）。优化前的循环里，
每一层都调用一次 torch_npu._npu_paged_attention_get_workspace：
同一 decode 步内，所有层的形状和 seq_lens 完全相同——28 次查询查的是同一个结果， 纯粹是重复劳动，还把 host 侧 ACL API 时间线撑爆。
改动后：workspace 查询提到循环外，每轮更新只查一次、跨层复用；
只有当尺寸相关输入变化（seq_lens 张量身份 (data_ptr, shape) 及各形状/dtype 组成
的 key 不匹配）才重新查询。这与 FIA 路径已有的 graph_params.workspaces[num_tokens] 复用形态对齐。
单文件改动（attention/attention_v1.py，+42/-10）。三点设计值得注意：
| 设计点 | 说明 |
|---|---|
| ws_key 覆盖全部尺寸输入 | (data_ptr, shape) + query/key_cache/value_cache 的 shape 与 dtype + block_table.shape——异构层配置只是触发重查，绝不会复用错误尺寸 |
| 只做轮内复用，刻意不做跨步 | workspace 大小还依赖 seq_lens 的数值（实测 len≤80 时 36MB、len≤400 时 54MB，随 decode 增长）——上一步的缓存可能偏小，复用有风险 |
| 保守回退 | 任一层 key 不匹配 → 立即走一次全新 get_workspace，正确性优先 |
with torch.npu.stream(update_stream): # 图参数更新（每步一次，内含逐层循环）
for key, param, handle, event in zip( # ↳逐层循环：28 层模型每步 28 轮（问题所在）
forward_context.attn_metadata,
graph_params.attn_params[num_tokens],
graph_params.handles[num_tokens],
graph_params.events[num_tokens],
):
(
query,
key_cache,
value_cache,
num_kv_heads,
num_heads,
scale,
block_table,
seq_lens,
output,
) = param
seq_lens = forward_context.attn_metadata[key].seq_lens # 同一步内所有层共享同一 seq_lens（跨层复用的前提）
workspace = torch_npu._npu_paged_attention_get_workspace( # ⚡每层都查一次：36-54MB 分配 + Setup 往返（问题所在）
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.graph_task_update_begin(update_stream, handle) # 图任务重放绑定 + PA 执行
torch_npu._npu_paged_attention(
query=query,step_ws_key = None # ⚡本轮 workspace 缓存键（跨层复用的核心）
workspace = None
for key, param, handle, event in zip(
forward_context.attn_metadata,
graph_params.attn_params[num_tokens],
) = param
seq_lens = forward_context.attn_metadata[key].seq_lens
# The key covers every size-relevant input, so models
# with heterogeneous layer configs simply trigger a
# re-query instead of reusing a wrong-sized workspace.
ws_key = ( # 键覆盖全部尺寸相关输入：任何差异都触发重查（保守设计）
seq_lens.data_ptr(),
tuple(seq_lens.shape),
query.shape,
query.dtype,
key_cache.shape,
key_cache.dtype,
value_cache.shape,
value_cache.dtype,
block_table.shape if block_table is not None else None,
block_table.dtype if block_table is not None else None,
output.shape,
output.dtype,
num_kv_heads,
num_heads,
scale,
)
if step_ws_key != ws_key: # 键不匹配才重新查询（异构层配置自动回退）
workspace = torch_npu._npu_paged_attention_get_workspace( # 每轮最多查询一次（原为每层一次）
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,
)
step_ws_key = ws_key # 记录本轮键：后续层直接复用
torch.npu.graph_task_update_begin(update_stream, handle)
torch_npu._npu_paged_attention(效果（910B 实测，vLLM v0.26.0 / TP=1 / FULL_DECODE_ONLY / c=32）：
| 指标 | 改动前 | 改动后 |
|---|---|---|
| get_workspace Setup/Execute 比 | 2.000（22232/11116） | 1.036（14355/13860） |
| update_graph_params 均值 | 2.4-2.95 ms | 1.65-1.67 ms（每步省 1.1-1.3ms） |
| Host ACL API 总时长（5.1s 窗口） | 1335 ms | 1040 ms（-22%） |
| QPS c=1 / c=32 / c=64 | 1.27 / 16.73 / 21.49 | 1.35 / 17.13 / 21.90（+5.9%/+2.4%/+1.9%） |
功能输出不变（全功能套件 md5 一致）。