已合并
knowledge(triton) 新增 A2 SparseFlashAttentionGrad triton 算子优化知识 #174
knowledge(triton) 新增 A2 SparseFlashAttentionGrad triton 算子优化知识 #174
已合并
guomc创建于 8 天前
guomc
guomc
8 天前

Pull Request 说明

变更说明

  • 关联 Issue:无
  • 主要改动:

将原 sparse_attention_grad_optimization.md 拆分为两张单算子 Runbook,分别维护块掩码与 topK 索引两种稀疏注意力反向形态:

  1. 新增 block_sparse_attention_grad_optimization.md,记录 device 侧掩码转置与选择表构建、连续写回、归约方向和独立 tile 规则。
  2. 新增 sparse_flash_attention_grad_optimization.md,记录微基准定位、dense-k 大 GEMM 重构、中间布局、精度修复和调试方法。
  3. 两张 runbook 的类别约束统一链接 attention_writing_classification.md,不新增设计期约束卡,既有 block_sparse_attention_constraints.md 不改动。
  4. 删除合并卡片,同步 performance 目录索引及当天日志。
  • 补充说明:知识来自 Ascend 910B3 上的两个算子完整开发实测轨迹,环境为 triton_ascend 3.6.0 与 CANN 9.1.0。轨迹沉淀的 Runbook 按治理规约使用 sources: [],正文保留触发条件、探针方法、原始观测与验证信号。所有结论仅声明在实测平台 a2 有效,frontmatter 的 platforms 与实测范围一致。

性能证据

SparseFlashAttentionGrad 评测集 18 个 case,覆盖 B 从 1 到 8、S1 从 512 到 8192 含非 2 幂边界、S1 与 S2 不等长、N1 从 4 到 128、K 四档、两种 sparseMode 和 fp16/bf16。Triton 实现经 4 轮 profile 驱动优化,总核时从 6764.3ms 降到 199.5ms,提升 33.9 倍。对 AscendC 标杆算子 npu_sparse_flash_attention_grad 的核时系数,逐 case 独立进程加 profiler kernel_details 口径,全 18 个 case 大于 0.6,最小 0.7424,几何平均 2.6177,官方精度验证 18 个 case 全部通过。

BlockSparseAttentionGrad 的对照数据为:总核时 5934.7ms 降到 21.67ms,提升 273.8 倍,几何平均系数 0.9283,17 个 case 全部大于 0.6。

SFAG 逐 case 明细如下,系数为 AscendC 核时除以 Triton 核时,大于 1 表示 Triton 更快:

case shape(BSND,D=512,N2=1) K mode dtype 直生成(ms) 优化后(ms) 提升 AscendC(ms) 系数
1 B1 S1=512 S2=1024 N1=4 1024 3 fp16 13.48 0.455 30x 2.027 4.457
2 B1 S1=1024 S2=1024 N1=8 1024 3 fp16 52.65 0.928 57x 2.789 3.006
3 B2 S1=1024 S2=2048 N1=16 2048 0 fp16 396.29 5.787 68x 26.471 4.575
4 B1 S1=2048 S2=2048 N1=32 1024 3 fp16 398.58 10.752 37x 9.205 0.856
5 B2 S1=2048 S2=2048 N1=8 2048 3 fp16 412.53 8.232 50x 22.087 2.683
6 B1 S1=4096 S2=4096 N1=4 4096 0 fp16 423.14 16.660 25x 78.720 4.725
7 B1 S1=4096 S2=8192 N1=8 2048 3 fp16 403.71 24.168 17x 42.986 1.779
8 B4 S1=1024 S2=1024 N1=16 1024 0 fp16 395.10 5.559 71x 21.876 3.936
9 B1 S1=8192 S2=8192 N1=4 1024 3 fp16 212.26 26.437 8x 34.688 1.312
10 B1 S1=2048 S2=2048 N1=64 1024 0 fp16 779.41 14.379 54x 13.572 0.944
11 B1 S1=1024 S2=1024 N1=128 1024 3 fp16 803.74 6.734 119x 4.999 0.742
12 B2 S1=4096 S2=4096 N1=4 2048 0 bf16 421.91 20.952 20x 74.146 3.539
13 B8 S1=512 S2=1024 N1=4 1024 0 fp16 106.07 3.034 35x 21.212 6.991
14 B1 S1=2000 S2=2048 N1=8 1024 3 fp16 98.84 2.848 35x 8.530 2.996
15 B1 S1=2048 S2=8192 N1=4 8192 0 fp16 425.14 22.344 19x 79.018 3.536
16 B2 S1=1536 S2=1536 N1=16 1024 3 fp16 301.83 5.083 59x 11.530 2.268
17 B1 S1=4096 S2=4096 N1=8 4096 3 fp16 817.12 18.657 44x 41.860 2.244
18 B3 S1=1024 S2=2048 N1=8 2048 0 bf16 302.47 6.475 47x 37.763 5.832
合计/几何平均 — — — — 6764.3 199.5 33.9x 533.5 2.6177

变更类型

验证

本地已通过全库 Lint,检查 9058 页且 P0/P1/P2 均为 0;索引检查 712 项且无过期;两个目标算子均可召回新卡;36 项 trace 测试和字段检查通过。完整 check.sh 未通过:检索质量门禁的失败项位于本次未修改的 Ascend C 用例,Windows 下部分 evals 存在编码与路径兼容性失败,本机也无法从可用包源安装 ruff。

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 guomc 的贡献)
guomcguomc
8 天前 创建了 pull request,commit 5e5ec807
CANN-robotCANN-robot成员
8 天前 添加了label:stat/needs-squash
CANN-robotCANN-robot成员
8 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
8 天前 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.


PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
repo-cann/cannbot-knowledge ✅ zhanghz1, qq_41667743 (2/2) ✅ zhanghz1 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

guomc9, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
CANN-robotCANN-robot成员
8 天前 将qq_41667743,wangsj0012,jon_jin,zhanghz1,fafahan,zhengshencn_hwca,linghuis,OpSphere,xfwang_ecnu设为评审人
此处折叠了213条消息 查看更多
CANN-robotCANN-robot成员
7 天前 添加了label:ci-pipeline-passed
汪超成员
7 天前 评论:

/lgtm

likedislike
zhanghz1
zhanghz1成员
7 天前 评论:

/approve

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label:lgtmapproved
CANN-robotCANN-robot成员
7 天前 合入了pull request