已关闭
[Requirement|需求建议]: 为 UndGenQkvRmsNormRopeCache 接入 E2E XPU 参考实现 #5238
qianzehong创建于  13 天前关闭于  13 天前
qianzehong成员
13 天前 创建

Backgroud(背景信息)

UndGenQkvRmsNormRopeCache 是昇腾专有算子。Torch E2E 测试通过 TTK --xpu-perf / cross_check 调用远端 xpu-server 时,GPU 端不存在对应的 torch.ops.cann_ops_transformer API,需要显式提供可执行的 XPU 参考类。

Benefit / Necessity (价值/作用)

让服务端通过 TestSpec 的 third_party 配置加载参考实现,在目标设备上执行现有 golden,支持该算子的 XPU 测试接入。

Design(设计方案)

仅修改 tests/assets/golden.py:新增并导出 UndGenQkvRmsNormRopeCacheXpuRef,在 Torch TestSpec 中注册 third_party。参考类接收目标设备上的张量,规范化 head 数、eps、mrope_section 后调用现有 golden 主入口,使用 inplace=False 保留输入 KV cache,返回 q 和更新后的 cache 副本。

验证

提交 812b8f88e 已在 NVIDIA A100-PCIE-40GB、PyTorch 2.13.0+cu130 上验证:8 个 CUDA 用例和 16 次服务端 DATA 请求执行通过,覆盖新增类的实例化、调用及 third_party 解析,GPU 输出与 CPU golden 的精度比较通过,输入 cache 保持不变。完整 TTK NPU/GPU 三方交叉校验及性能采集尚未执行。

关联 PR:https://gitcode.com/cann/ops-transformer/merge_requests/11828

likedislike
Qqianzehong成员
13 天前 添加了label:requirement
qianzehong成员
13 天前 评论:

/assign @qianzehong

likedislike
CANN-robotCANN-robot成员
13 天前 将 qianzehong 设为负责人
CANN-robotCANN-robot成员
13 天前 关闭了 issue
CANN-robotCANN-robot成员
13 天前 添加了label:resolved