Ascend 910B3: GEMM 性能正常 (可达 baseline 114.7%)，FlashAttention 性能仅 3.3 TFLOPS (baseline 的 4~9%)

[Alex4210987]

问题概述

Triton → DLCompiler → 昇腾 910B3 路径：纯 Cube 计算的 GEMM 性能可用，但涉及 Cube-Vector 交织的 kernel（如 FlashAttention）性能断崖式下跌，基本不可用。

• GEMM（纯矩阵乘，全 Cube 计算）：autotune + deeplink 编译提示 + 对角线分核调度加持下，大 shape 可达 vendor 库的 113%~115%
• FlashAttention（Q×Kᵀ → online softmax → P×V，Cube 与 Vector 紧密交织）：仅约 3.3 TFLOPS（昇腾 FP16 峰值的 1.3%），仅为 vendor 库的 4%~9%；head_dim=128 全部编译/运行失败

两者性能差距约 80 倍。GEMM 证明 BiShengHIR 在纯 Cube 场景是称职的——问题出在编译器对 Cube-Vector 交织结构的调度建模上。

---

运行环境

• 硬件: 昇腾 Ascend 910B3（NPU，aarch64 架构），24 个 AI Core
• 软件: CANN 8.3.RC1，Python 3.11，torch 2.7.1，triton-ascend
• 被测对象: DLCompiler 仓库 test/ascend/passed_tests/test_fa_v1.py（FlashAttention）以及 test/ascend/passed_tests/test_matrix_multiplication_optimized.py（GEMM，Triton 版本）
• Baseline 对比对象:
  • GEMM: torch_npu.npu_mm（华为官方矩阵乘库）
  • FlashAttention: torch_npu.npu_fusion_attention（华为官方融合注意力库）
• 计时方法: do_bench(warmup=25, rep=100)。FlashAttention 额外使用 do_bench(warmup=500, rep=500) 与原始 test_fa_v1.py 完全一致的参数做验证，结果一致。
• 计算量公式:
  • GEMM = 2 × M × N × K
  • FlashAttention = 4 × batch × head × seq_len² × head_dim

---

测试一：GEMM（性能良好）

计算结构：纯矩阵乘法。加载 A 块 → 矩阵乘（Cube）→ 累加 → 存储 C 块（Vector 仅做类型转换），几乎全部由 Cube 单元完成。

测试方法：直接使用 DLCompiler 官方 test/ascend/passed_tests/test_matrix_multiplication_optimized.py，kernel 代码一字未改。使用 triton.language.extra.deeplink 库的 compile_hint 提示 + @triton.autotune（官方提供的 11 组配置）+ 对角线分核调度。

测试结果：

矩阵规模	Baseline（华为 mm）	DLCompiler-TR	相对 Baseline	状态
48×7168×16384	21.27 TFLOPS	24.87 TFLOPS	117.0%	正常
96×4096×4096	49.63 TFLOPS	14.38 TFLOPS	29.0%	正常
96×18432×7168	49.69 TFLOPS	46.53 TFLOPS	93.6%	正常
2048³	200.94 TFLOPS	81.42 TFLOPS	40.5%	正常
4096³	221.02 TFLOPS	250.62 TFLOPS	113.4%	正常
8192³	228.62 TFLOPS	262.20 TFLOPS	114.7%	正常

六个规模均能成功编译并运行。大 shape（4096³、8192³）性能超越 vendor 库；小 shape（2048³ 及以下）性能偏低（29%~40.5%），autotune 在大 shape 下优势明显。

---

测试二：FlashAttention（性能极差）

计算结构：Q × Kᵀ（Cube 矩阵乘）→ online softmax：求最大值 / 指数运算 / 按行缩放（Vector 逐元素运算）→ P × V（Cube 矩阵乘），然后进入下一轮迭代。每轮 softmax 的缩放因子还需反向修正上一轮的累积结果（跨轮数据依赖）。Cube 与 Vector 高度交织。

测试方法：直接使用 DLCompiler 仓库 test/ascend/passed_tests/test_fa_v1.py，kernel 代码一字未改。核心计算模式为内层循环 Q×Kᵀ → online softmax → P×V → 跨轮状态修正。调用时附带编译提示 debug=True, multibuffer=True, set_workspace_multibuffer=4, tile_mix_vector_loop=4, tile_mix_cube_loop=4。

测试结果（BM=block_M, BN=block_N，Baseline 为 npu_fusion_attention，TFLOPS 为 do_bench 实测值）：

#	batch	head	seq_len	head_dim	BM	BN	FLOPs	DLCompiler-TR	Baseline	相对 Baseline	状态
1	1	1	64	64	64	16	1 M	0.00	0.01	45.8%	正常
2	4	32	64	64	64	16	134 M	0.46	0.85	54.1%	正常
3	4	32	64	64	64	32	134 M	0.41	0.92	44.6%	正常
4	1	1	128	128	16	16	8 M	—	0.06	N/A	507015 aicore 异常
5	1	1	128	128	64	64	8 M	—	0.06	N/A	bishengir vcast 编译错误
6	1	2	1024	64	64	32	537 M	1.73	4.20	41.2%	正常
7	1	1	2048	64	64	32	1.1 G	3.29	7.50	43.9%	正常
8	2	32	512	64	64	32	4.3 G	3.29	39.21	8.4%	正常
9	1	1	4096	64	64	32	4.3 G	3.57	38.28	9.3%	正常
10	1	1	8192	64	64	32	17.2 G	3.41	60.35	5.7%	正常
11	2	32	1024	64	64	32	17.2 G	3.26	63.82	5.1%	正常
12	1	1	16384	64	64	32	68.7 G	3.34	67.17	5.0%	正常
13	2	32	2048	64	64	32	68.7 G	3.26	80.71	4.0%	正常
14	1	1	32768	64	64	32	274.9 G	3.28	78.62	4.2%	正常
15	1	32	4096	128	64	32	274.9 G	—	139.91	N/A	bishengir vcast 编译错误

关键数据点：

1. head_dim=64 时，DLCompiler-TR TFLOPS 锁定在约 3.3 TFLOPS——与问题规模完全无关（1M FLOPs → 275G FLOPs，性能几乎不变）。仅为昇腾 910B3 理论 FP16 峰值的 1.3%。

2. 相对 Baseline 从 54% 断崖式跌至 4~9%：小 shape（seq_len ≤ 2048）尚可维持 41%~54%；一旦进入计算密集区间（seq_len ≥ 512，batch≥2），立刻跌到 4%~9% 并永久锁死。Baseline 随规模线性增长（最高 140 TFLOPS），DLCompiler-TR 完全无法跟随。

3. head_dim=128 时，3/3 失败：2 个配置 507015 aicore 异常（BM=16/BN=16、BM=64/BN=64），1 个配置 bishengir vcast 编译错误（BM=64/BN=32）。

4. 与 GEMM 的对比直接定位问题：纯 Cube 的 GEMM 可达 262 TFLOPS（baseline 的 114.7%），Cube-Vector 交织的 FlashAttention 仅 3.3 TFLOPS，差距约 80 倍。

---

复现步骤

Triton kernel 通过 triton-ascend 的 triton.backends.dicp_triton 后端自动路由到 DLCompiler 的 bishengir 编译器链（bishengir-hivm-compile → bishengir-compile）。

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export LD_LIBRARY_PATH=${ASCEND_HOME_PATH}/aarch64-linux/lib64:${ASCEND_HOME_PATH}/hccl/lib64:${LD_LIBRARY_PATH}
export TORCH_DEVICE_BACKEND_AUTOLOAD=0

# GEMM（Triton → dicp_triton → bishengir）
python test/ascend/passed_tests/test_matrix_multiplication_optimized.py

# FlashAttention（同上）
python test/ascend/passed_tests/test_fa_v1.py

---

根因分析

昇腾 910B3 的 Cube 单元（矩阵乘）和 Vector 单元（逐元素运算）是两套独立的执行引擎，数据必须沿 HBM → L1 → L0A/L0B 的层次路径流动，调度必须在编译期静态确定。

• GEMM：纯 Cube 循环，BiShengHIR 能识别并生成基本可用的调度代码
• FlashAttention：每次迭代中 Cube（Q×Kᵀ、P×V）与 Vector（softmax 的 exp/max/rescale、跨轮状态修正）紧密交替。BiShengHIR 对这种 Cube-Vector 交织、跨轮数据依赖的结构缺乏有效的调度建模，编译出的代码几乎完全退化为 Vector 串行执行，Cube 单元基本空闲

这是编译器架构层面的限制。tile_mix_vector_loop=4, tile_mix_cube_loop=4 等提示参数只能微调现有策略，无法从根本上改变编译器对交织结构的建模方式。

---

问题与建议

1. DLCompiler（或底层 BiShengHIR）是否提供了针对 Cube-Vector 交织模式的调度标注机制（如 IR annotation 或 compile hint），可以强制生成流水线化的 Cube/Vector 交替执行？目前 tile_mix_vector_loop 和 tile_mix_cube_loop 提示参数的效果有限，无法从约 3.3 TFLOPS 提升到可用水平。

2. head_dim=128 时 UB 溢出：DLCompiler 是否有机制允许用户控制中间缓冲区的分配策略（如限制 buffer 大小、允许溢出到 L1/HBM），以避免超出硬件 UB 容量限制（1,572,864 比特）？

3. DLCompiler 是否有计划在 deeplink 库中提供更细粒度的调度控制接口，让用户可以显式标注 Cube 和 Vector 运算的流水线编排，而非依赖编译器自动推断？

[Truth-Ke]

感谢您的反馈，fa麻烦请拉最新的分支，并且fa的seq_len变大导致性能劣化厉害，是因为bm bn太小了，可以拉大！
可以使用我们的最新的FA的脚本，性能是ascend-c 0.8以上
TA.py