已合并
[Feature][op][eqbsa-02] add eqbsa operator, attention infra, and enable build #514
[Feature][op][eqbsa-02] add eqbsa operator, attention infra, and enable build #514
已合并
lanwangli创建于 13 天前
lanwangli成员
13 天前

Which issue(s) this PR fixes or accomplishes

Fix part of #197

关联 Issue:https://gitcode.com/Ascend/MindIE-SD/issues/197

Purpose

上库 eagle_quant_block_sparse_attention 算子和 plugin 适配,关联 RFC: #197。

该算子用于叠加量化和块稀疏 Attention 能力:Q/K 采用 INT8 量化,P/V 相关路径采用 FP8 低精度表示,并通过 block sparse mask/index 降低长序列 Attention 的计算和访存开销。

本 PR 落地 eqbsa 可编译算子

  • kernel 依赖的 attention infra / TLA scaffolding(op_kernel/attn_infraop_kernel/tla
  • 算子本体:op_host(tiling / infershape / OWNERS)、op_api、arch35 kernel、kernel 入口
  • eagle_quant_block_sparse_attention 写入 build_ops.sh 默认算子列表

infra 是被 kernel 复用的搬运 / GEMM / epilogue / layout 基础设施,可按 vendored scaffolding 评审(目录边界与许可证)。tiling、aclnn 封装、arch35 kernel 与 build_ops.sh 注册是需要按算子逻辑评审的部分。

Scope

  • build/build_ops.sh
  • csrc/ops/eagle_quant_block_sparse_attention/op_api/**
  • csrc/ops/eagle_quant_block_sparse_attention/op_host/**
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/**
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/**
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/arch35/**
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention.cpp
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention_kernel_interface.cpp
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention_kernel_regular_arch32.h
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention_tilingkey.h
  • csrc/ops/eagle_quant_block_sparse_attention/op_kernel/kernel_common.hpp

Overall Sequence

  1. 构建兼容性: !512
  2. 算子源码 + attention infra + 构建启用: !514
  3. Plugin API + golden 测试: !515

Merge Safety

本 PR 已包含 kernel 编译所需的 attn_infra / tla 与算子本体,并启用 build_ops.sh。相对当前 dev 可独立编译 eqbsa。CANN 9.1 环境还需先合入 !512 的 msprof include。本 PR 不暴露 PyTorch plugin / Python API。

Test Plan

  1. 算子编译
    • build/build_ops.sh
  2. 评审重点
    • tiling、kernel 入口、arch35 kernel、OWNERS、build_ops.sh 注册名
    • attn_infra / tla 按 vendored scaffolding 看目录边界与许可证

Test Report

测试环境: Ascend 950,CANN 9.x

测试项 命令/步骤 结果
源码范围 83 文件:infra/tla + 算子本体 + build_ops.sh
eqbsa 上机 Ascend 950 上完成 eqbsa 精度验证 ✅ 已验证

Ascend 950 golden 精度测试通过

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 lanwangli 的贡献)
Llanwangli成员
13 天前 创建了 pull request,commit fa5effc0
Llanwangli成员
13 天前 关联了issue:[RFC]: 上库 eagle_quant_block_sparse_attention 算子和 plugin 适配
atomgit-bot
atomgit-bot
13 天前 评论:

🤖 正在生成合并请求摘要,请稍候…

likedislike
atomgit-bot
atomgit-bot
13 天前 评论:

🤖 AI 代码检视正在进行中,请稍候…

likedislike
Xxiangjie10成员
13 天前 添加了label:pr-audit-failed
此处折叠了64条消息 查看更多
lanwangli成员
5 天前 评论:

1.大型算子代码,总共2w行,当前算子代码结构 kernel和host不好拆分,已经把plugin和test拆分了,放在下一个PR
2.test在下一个PR合入

likedislike
blian
blian成员
5 天前 评论:

/merge

likedislike
ascend-robotascend-robot成员
5 天前 添加了label:keeper_approved
ascend-robotascend-robot成员
5 天前 合入了pull request
Llanwangli成员
5 天前 删除了label:stat/needs-squashkeeper_approved