提交提案之前，请先检索仓库内是否已有相同的提案，如已有请在同一提案中进行讨论。

性能优化具体描述

背景与优化特性

HSTU Paged Attention 从分页 KV cache 读取历史 K/V，并将其整理到 attention 计算使用的 workspace。当前 develop@64b6118d4f6fc534e1bc3f6792ad74c4974a05fa 的 v220 kernel 在 CopyFromKvCache() 中按 page 串行执行 K、V 的 GM→UB→GM 搬运：上一批数据写回 workspace 后才开始下一批读取，MTE2 输入与 MTE3 输出之间没有形成跨 chunk 的流水重叠。

在增量推理的小 Q、大 KV 场景中，分页 KV cache 搬运量远大于新增 Q 的数据量。建议复用现有 UB，通过 K/V 双缓冲和 MTE2/MTE3 event，将下一 chunk 的 GM→UB 输入与上一 chunk 的 UB→GM 输出重叠，降低搬运路径的等待时间。

使用场景与适用范围

本次性能结论只适用于上表中的 Ascend 910B3 小 Q、大 KV 矩阵。以下范围尚未形成性能结论：

• 其他 Ascend 型号；
• 目标矩阵以外的 Q、KV、head dimension、head 数和 page size；
• FP16/FP32 的性能；
• 多卡或端到端模型性能。

优化方案

• 只修改 CopyFromKvCache()，复用既有 queKv_ UB，不增加 UB 总预算。
• 将既有 UB 划分为 K/V 两类 tensor 和两个 ping-pong stage。
• 使用 MTE2→MTE3 与 MTE3→MTE2 event 管理输入就绪和输出完成状态。
• 在相邻 chunk 间重叠下一批 K/V 输入与上一批 workspace 输出。
• chunk 容量由 kvLtUbSize_、sizeof(qType) 和运行时 headDim_ 推导。
• 保持 host tiling、tiling key、算子 API、输入输出格式和数值语义不变。

性能劣化说明

测试方法

• Baseline：develop@64b6118d4f6fc534e1bc3f6792ad74c4974a05fa 的 RecSDK 算子包中的 paged hstu attention 算子；
• Optimized：基于上述优化后的 paged hstu attention 算子；
• 测试矩阵：BF16，H=16，D=256，page size=32，3 个 (B, KV) × 6 个 Q，共 18 个 shape；
• 正确性：每个 shape 使用 seed 2026/2029/2039，共 54 个用例；只有正确性全部通过才运行性能测试；
• 性能统计：固定随机 shape 顺序，每个 shape 预热 16 次，使用 NPU Event 仅统计 device kernel 执行时间，每种形状采样 51 次，以 median 作为该 shape 延迟；
• 聚合方法：对 18 个 shape 的 median 计算几何均值；
• 加速比定义：baseline latency / optimized latency，低于 1.0x 表示性能劣化。
• 硬件：Ascend 910B3 单卡

详细性能数据

汇总与劣化结论

• Baseline 的 18-shape median 几何均值：0.9582518845 ms；
• Optimized 的 18-shape median 几何均值：0.8952944909 ms；
• 总体加速比：1.070320318x，延迟下降 6.5700%；
• 加速 shape：18/18；劣化 shape：0/18；
• 单 shape 加速比范围：1.063392x～1.075944x。

其他相关讨论

执行了完整的正确性测试

目标矩阵的正确性验证情况：

• Baseline：54/54 通过 CPU FP32 reference，rtol=atol=8e-3；
• Optimized：54/54 通过，输出有限值；
• Optimized 对 Baseline：54/54 bitwise 一致，最大绝对误差为 0；
• Optimized 对 CPU FP32 reference：最大绝对误差 3.293827e-4；
• eager 重复执行：54/54 bitwise 一致；
• eager/Graph：54/54 bitwise 一致；
• 仓库既有代表性 pytest：5 passed，覆盖 BF16/FP32/FP16、page size 32/128/256、
  batch 32、GQA 和多种 mask。

接口、兼容性与回退

• 对外 API、参数、输入输出 shape/格式：无变化；
• host tiling、tiling key、环境变量：无变化；
• 第三方开源软件：无新增；
• 用户资料：调用方式和约束未变化，无需修改用户文档；
• 安全：未新增外部输入、网络访问、文件访问或权限处理逻辑；
• 回退：回退候选的单个 squash commit 即可，不涉及数据格式或配置迁移。

风险边界：改动位于 v220 通用 Paged 搬运路径，非目标平台和非目标 shape 的正确性依赖既有回归与 PR CI，性能尚未系统测量。当前实现没有运行时性能回退或 shape gate；这是为了避免把测试矩阵中的 Q、KV 或 D 数值硬编码成分派条件。若社区要求扩展到更激进的特化路径，应另建 Issue/PR，并基于 UB/L1/寄存器容量、页/chunk 数、MTE 搬运量、计算量及平台 API 返回的 core 数建立可解释的成本模型，再跨 dtype、D、page size 和 Q/K 范围确定分派边界。

待社区确认

1. 本次仅承诺 Ascend 910B 小 Q、大 KV 性能收益，但实现位于无 shape gate 的 v220 通用路径；请确认该最小通用流水方案及回归范围是否可接受。
2. 其他平台/shape 若需要性能分派，请确认是否接受后续基于资源/成本模型的独立 Issue/PR；本 issue 的方案不接受直接写入 Q、KV、D、SoC 型号或固定 core 数阈值。

环境信息

欢迎加入社区，感谢您对社区的贡献 🎉!