CANN ops-transformer 中 BlitzSparseAttention 块稀疏注意力基准测试实战benchmark.py 调参、绘图管线与 16 种块粒度扫描【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文以 benchmark/README.md 为主体完整讲解 CANN ops-transformer 实验算子 BlitzSparseAttention基于 PromptFlashAttentionV3 扩展的块稀疏 prefill 注意力的基准测试与正确性验证流程从编译安装、torch_bsa接口构建到benchmark.py的全参数扫描、plot.py速度比绘图管线、16 种(BLOCK_SIZE_Q, BLOCK_SIZE_KV)块粒度sabi 粒度扫描结果解读以及test_lse.py的 softmax LSE 正确性测试设计。读完后你可以独立完成该算子的精度验证与性能基准并能依据数据自主选型稀疏块粒度。背景BlitzSparseAttention 算子与 benchmark 目录结构BlitzSparseAttention 是 ops-transformer 中基于 PromptFlashAttentionV3 的实验性块稀疏注意力内核新增sabi参数使 prefill 阶段可按块跳过未选中的 KV 块从而在生成式模型中稀疏度越高、加速越明显。sabi是uint16张量形状为[B, N, ceil(S/BLOCK_SIZE_Q), ceil(S/BLOCK_SIZE_KV)]每行列出该 Q 块保留的 KV 块列索引右端用0xFFFF哨兵填充。算子完整参数文档见 docs/aclnnBlitzSparseAttention.md算子总览见 README.md。benchmark目录包含基准测试与正确性测试的全部脚本文件作用benchmark.py基准测试驱动扫描全部(B, H, S, D, sparsity)×BLOCK_SHAPES组合输出延迟表并可与参考内核做正确性比对plot.py从 stdin 读取 benchmark 表格按(序列长度, block_shape)绘制 speedup-vs-sparsity 折线图test_attn.pyattention_out 正确性测试10k–30k 序列、1–4 个注意力头test_lse.pysoftmax_lse 第二路输出的正确性测试test_joint.py同时校验两路输出attention_out 与 softmax_lse需要特别注意的两个已知限制来自 README.md其一当前 128×128 sabi 粒度内部仍使用 128×512 的 matmul 瓦片未真正压缩未选中的 128×128 子块加速从稀疏度 ≥10% 才起步其二batch size 目前仅支持B1B1会产出错误结果——这也是 benchmark 文档中反复强调B_VALS [1]的原因。安装编译算子包并构建 torch 接口按文档要求完整安装分三步在项目根目录执行bash build.sh --make_clean --experimental -j96 --pkg --socascend910b --opsblitz_sparse_attention ./build/cann-ops-transformer-custom_linux-$(uname -i).run (cd experimental/attention/blitz_sparse_attention/torch_interface bash build.sh custom)第一步以--experimental标志构建实验算子包--socascend910b指定目标 SoC第二步安装 run 包第三步在 torch_interface/ 下构建torch_bsa扩展其 C 入口 torch_interface.cpp 通过aclnn_blitz_sparse_attention.h把 torch 张量转换为 aclTensor 并调用两段式 aclnn 接口block_shape参数的默认值为[128, 128]且强制校验两个分量必须取自{128, 256, 512, 1024}。如果只需要快速重编译 bf16 内核变体可追加--op-variant2bash build.sh --make_clean --experimental -j96 --pkg --socascend910b --opsblitz_sparse_attention --op-variant2 ./build/cann-ops-transformer-custom_linux-$(uname -i).runsmoketest 与基准测试的一条命令(cd experimental/attention/blitz_sparse_attention/benchmark python test_attn.py python benchmark.py)基准测试加绘图管线生成汇总所有(BLOCK_SIZE_Q, BLOCK_SIZE_KV, sparsity)组合的benchmark.png(cd experimental/attention/blitz_sparse_attention/benchmark python benchmark.py | tee bench.log /dev/tty | python plot.py)从源码看这条管道能成立是因为 benchmark.py 显式把日志写到 stdoutlogging.basicConfig(..., streamsys.stdout)而 plot.py 用正则ROW_RE从 stdin 逐行解析结果表无法匹配的行表头、分隔线自动忽略tee ... /dev/tty则让表格同时落盘和上屏。benchmark.py 参数详解扫描范围由文件顶部 Begin of Parameter Sweep Definitions 区间的ALL_CAPS常量控制。只要某个常量是列表扫描就会覆盖其全部笛卡尔积组合常量说明B_VALSbatch 大小。块稀疏模式目前仅支持B_VALS [1]H_VALS注意力头数S_VALS序列长度S_q S_kvD_VALShead 维度N_REPEATS用于估算耗时的正式运行次数N_WARMUP正式计时前的热身次数SPARSITY_VALS块稀疏度每行 Q 块随机保留相同数量的 KV 块块稀疏模式是 benchmark 唯一验证的内核模式BLOCK_SHAPES(BLOCK_SIZE_Q, BLOCK_SIZE_KV)对列表。内核在运行时通过block_shapeop-attr 接受两维各自取{128, 256, 512, 1024}的任意组合切换粒度无需重编译默认扫描完整的 4×416 组网格。输出表第一列标注每行生效的块形状BLOCK_SIZE_Q/BLOCK_SIZE_KV别名取自BLOCK_SHAPES[0]仅为向后兼容保留BLOCK_MASK_SEED随机采样的复现种子改值即改采样结果FRAMES_BY_BLOCK_SHAPE以(BLOCK_SIZE_Q, BLOCK_SIZE_KV)为键的二维字典返回FrameWidths(left_cols, right_cols, top_rows, bottom_rows)具名元组——在稀疏预算之外强制激活的边界块用于模拟典型 sink/streaming 模式各组合的取值按比例缩放使强制 token 占位在不同粒度下可比推导见下文示例 2PRINT_OUTPUTS打印输出张量以便人工核对PRINT_MASK打印 mask 以便目检是否符合预期PRINT_BLOCK_EQUALITY正确性失败时按块可视化输出与参考的逐块一致情况块粒度由PRINT_HEIGHT和PRINT_WIDTH指定RUN_REFERENCE除计时外同时校验正确性。仅在 mask 足够小时可行否则参考运行会 OOM因此超长上下文建议关闭、短序列才开启TORCH_REFERENCETrue类语义下为真时以 torch attention 实现为参考为假时块稀疏模式以 TorchNPU dense 对应 attention mask 为参考类似sparse_block模式。注意 torch_npu 不支持不同头使用不同 mask因此只能使用H1此外还有 benchmark.py 中的SABI_SORTED内核稀疏循环假设 sabi 行内索引升序排列FirstGreaterEqual线性扫描 chunkIdx 前向游走设为False会破坏正确性仅用于内核已适配的布局/性能实验。关于计时的实现细节_run_timed使用torch.npu.Event(enable_timingTrue)在 NPU 侧记录起止事件先执行N_WARMUP次不计时再计时N_REPEATS次取平均毫秒换算为微秒即表中Ref_Latency_[usec]/Our_Latency_[usec]的口径。正确性判定使用torch.allclose(out_our, out_ref, rtol0.01, atol0.001)对应表中的Outputs_equal列yes/no/N/A。sabi 的生成逻辑值得一看generate_sparse_blocks_by_row先把(s_q, s_kv)划分成n_block_rows × n_block_cols的块网格frame 边框块左侧left_cols列 右侧right_cols列对所有稀疏行强制保留顶部/底部top_rows/bottom_rows行整行稠密计入稀疏预算后剩余保留名额均摊到每个稀疏行若 frame 强制块占比超过1 − sparsity会抛出ValueError。is_block_sparse_pattern_feasible则在不分配任何张量的纯算术层面做同样的可行性预检被 pytest 用于收集阶段剔除不可满足的参数组合避免运行期报错与 LSE 哨兵值泄漏到下游算子。示例 1多值正确性与速度测试较短序列在 Parameter Sweep 区间调整如下完整继承自文档示例DTYPE torch.bfloat16 INPUT_LAYOUT BNSD # [batch_size, num_heads, seq_len, head_dim] B_VALS [1] H_VALS [16, 14] S_VALS [4096, 10_000] # S_q S_kv D_VALS [128] # head dimension N_REPEATS 20 N_WARMUP 2 SPARSITY_VALS [0.0, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] BLOCK_SHAPES [(128, 512)] SABI_SORTED True FRAMES_BY_BLOCK_SHAPE {(128, 512):None} # no frame forced to be selected PRINT_OUTPUTS False PRINT_MASK False PRINT_BLOCK_EQUALITY False PRINT_HEIGHT 128 PRINT_WIDTH 8 RUN_REFERENCE True TORCH_REFERENCE npu_fusion_attention输出节选 DTYPEtorch.bfloat16 INPUT_LAYOUTBNSD SABI_SORTEDTrue TORCH_REFERENCEnpu_fusion_attention block_shape H B s_q s_kv D frame(L,R,T,B) sparsity Outputs_equal Ref_Latency_[usec] Our_Latency_[usec] ------------------------------------------------------------------------------------------------------------------------ 128x512 16 1 4096 4096 128 - 0.00 yes 859.30 1619.89 128x512 16 1 4096 4096 128 - 0.05 yes 1000.11 1617.22 128x512 16 1 4096 4096 128 - 0.10 yes 997.50 1593.49 128x512 16 1 4096 4096 128 - 0.50 yes 997.88 1114.68 128x512 16 1 4096 4096 128 - 0.90 yes 997.56 752.91 128x512 16 1 10000 10000 128 - 0.00 yes 5696.26 6649.24 128x512 16 1 10000 10000 128 - 0.50 yes 7538.75 3560.27 128x512 16 1 10000 10000 128 - 0.90 yes 7533.11 1266.50 128x512 14 1 10000 10000 128 - 0.90 yes 6685.34 1186.83 要点解读与文档一致表中先看输入维度Outputs_equal列核对正确性随后是参考内核耗时与本内核耗时10000 序列、24 头时参考实现会 OOM文档未展示该数据短序列下加速不明显甚至为负如 4096 序列稀疏度 0 时本内核 1619.89μs 慢于参考的 859.30μs因为固定开销高于实际计算量加速收益只有在长序列上才体现frame(L,R,T,B)列在稀疏度为 0 时显示-因为 mask 稠密、frame 对保留块无影响。示例 2块粒度block_shape扫描S118806, H3, D128, BF16运行时block_shape属性让同一个内核二进制运行于(BLOCK_SIZE_Q, BLOCK_SIZE_KV)的 16 种组合各维取{128, 256, 512, 1024}。粒度更细的 sabi块更小提高每个 Q 行的稀疏模式分辨率代价是 sabi 张量的 HBM 流量增加和每 cube 瓦片更多的子块处理开销粒度更粗则元数据更便宜但分辨率下降。该扫描测量的正是 S118806 下的延迟/吞吐权衡。frame 宽度为何按 (29,15,8,4) / (29,15,8,4) 缩放benchmark 强制一个边框frame模拟生成式模型典型的 sink/streaming 模式注意力矩阵的前约 3600 行/列与末约 6 行/列永远被选中。由于 6 个 token 在任意维度上都不足一块底部/右侧始终向上取整为 1 个块行/块列。FrameWidths.left_cols/right_cols以BLOCK_SIZE_KV个 token 为单位top_rows/bottom_rows以BLOCK_SIZE_Q个 token 为单位因此各组合取值来自把固定的 3600 token 目标在两个维度上独立向上取整到整块sizeceil(3600 / size)1282925615512810244由此得到二维FRAMES_BY_BLOCK_SHAPE[(BLOCK_SIZE_Q, BLOCK_SIZE_KV)]表left_cols ∈ {29, 15, 8, 4}按BLOCK_SIZE_KV索引top_rows ∈ {29, 15, 8, 4}按BLOCK_SIZE_Q索引right_cols bottom_rows 1处处成立。例如(BLOCK_SIZE_Q1024, BLOCK_SIZE_KV1024)⇒FrameWidths(left_cols4, right_cols1, top_rows4, bottom_rows1)。测试文件如 test_attn.py 的small_frame用更小的目标约 256 token应用同样的缩放规则生成SPARSE_FRAME。这一逻辑在 benchmark.py 中由_LEFT_COLS_BY_KV {128: 29, 256: 15, 512: 8, 1024: 4}与_TOP_ROWS_BY_Q两个字典直接实现。输入配置B_VALS [1]; H_VALS [3]; S_VALS [118_806]; D_VALS [128] N_REPEATS 10; N_WARMUP 2 SPARSITY_VALS [0.0, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] BLOCK_SHAPES [(bsq, bskv) for bsq in (128, 256, 512, 1024) for bskv in (128, 256, 512, 1024)] # 4 × 4 16 pairs # Per-pair frame: left_cols/top_rows are scaled by 128 / size in each dim so # the forced-token footprint stays comparable across block shapes. FRAMES_BY_BLOCK_SHAPE { (bsq, bskv): FrameWidths( left_cols {128: 29, 256: 15, 512: 8, 1024: 4}[bskv], right_cols 1, top_rows {128: 29, 256: 15, 512: 8, 1024: 4}[bsq], bottom_rows 1, ) for bsq in (128, 256, 512, 1024) for bskv in (128, 256, 512, 1024) } TORCH_REFERENCE npu_fusion_attention输出与绘图执行python benchmark.py | tee bench.log /dev/tty | python plot.py会额外生成speedup-vs-sparsity图每个(sequence_length, block_shape)对一条曲线。plot.py的 y 值定义为Ref_Latency(sparsity0) / Our_Latency(sparsityp)图上一条水平线 y1 表示与 dense 参考持平由于参考只在稀疏度 0 测一次每个block_shape × seq_len对只有一格Ref_Latency同一曲线所有稀疏度共用同一基线。完整原始表16 种块形状 × 11 个稀疏度S118806, H3, D128, BF16摘录如下完整表见原文档折叠区块block_shape H B s_q s_kv D frame(L,R,T,B) sparsity Outputs_equal Ref_Latency_[usec] Our_Latency_[usec] ------------------------------------------------------------------------------------------------------------------------ 128x128 3 1 118806 118806 128 - 0.00 yes 157210.33 200781.80 128x128 3 1 118806 118806 128 (29,1,29,1) 0.50 N/A N/A 117734.07 128x128 3 1 118806 118806 128 (29,1,29,1) 0.90 N/A N/A 24609.61 128x512 3 1 118806 118806 128 - 0.00 yes 161129.31 198749.26 128x512 3 1 118806 118806 128 (8,1,29,1) 0.50 N/A N/A 85403.20 128x512 3 1 118806 118806 128 (8,1,29,1) 0.90 N/A N/A 17906.81 256x256 3 1 118806 118806 128 - 0.00 yes 163126.33 200421.86 256x256 3 1 118806 118806 128 (15,1,15,1) 0.50 N/A N/A 93020.79 512x512 3 1 118806 118806 128 - 0.00 yes 165493.05 202528.27 512x512 3 1 118806 118806 128 (8,1,8,1) 0.90 N/A N/A 17844.47 1024x1024 3 1 118806 118806 128 - 0.00 yes 166565.60 203121.20 1024x1024 3 1 118806 118806 128 (4,1,4,1) 0.90 N/A N/A 17044.48表中为简洁只保留了稀疏度0.0 / 0.5 / 0.9实际运行会对BLOCK_SHAPES中每个对输出SPARSITY_VALS中全部稀疏度。frame(L,R,T,B)列在稀疏度 0 时为-mask 稠密frame 对保留块无影响。相对 npu_fusion_attentiondense的速度比——完整 4×4 网格每格取值为Ref_Latency(sparsity0) / Our_Latency(sparsityp)。 1.0表示 BSA 慢于 PFA dense 1.0表示快于Q\KVsparsity12825651210241280.31.09×1.28×1.35×1.29×0.51.34×1.74×1.89×1.84×0.82.84×4.19×4.65×4.63×2560.31.12×1.30×1.33×1.28×0.51.38×1.75×1.84×1.80×0.82.94×4.30×4.73×4.63×5120.31.14×1.31×1.26×1.25×0.51.40×1.77×1.73×1.73×0.82.99×4.21×4.54×4.60×10240.31.14×1.30×1.22×1.22×0.51.42×1.79×1.73×1.77×0.83.05×4.36×4.60×4.57×文档给出的结论每个格子内加速随稀疏度单调增长——BSA 在稀疏度 0 时也不劣于约 0.78×稠密回退损耗约 17–22%到稀疏度 0.9 时全网格达到6.4×–9.8×BLOCK_SIZE_KV是主要调节杆。稀疏度 0.9 下把 KV 从 128 提到 1024在每个 Q 上都能再多赚 1.4×–1.5×稀疏度 0.8 同样成立KV128→1024 约 1.5×–1.6×。原因是 sabi 条目越宽每个保留块的 sabi 张量 HBM 流量越低。但 KV1024 并非稀疏度 0.9 下的绝对赢家Q128 与 Q512 时 KV512 以微弱优势胜出9.00× vs 8.94×9.27× vs 9.11×——最粗的 KV 粒度下强制 frame 的 token 占位开始挤占保留块预算BLOCK_SIZE_Q影响小得多。Q 瓦片重参数化不改变 cube 瓦片宽度D128 BF16 时basicSInnerSize 512也不改变每 Q 行的工作量唯一效果是把多个 128 token 的 Q 组合并到同一 sabi 行基本免费稀疏度 0.9 下同 KV 下 Q128 与 Q1024 的差距 ≤ 12%稠密模式回退稀疏度 0在 KV 维度上基本平坦——相对 PFA 的开销全网格稳定在约 17–22%0.78×–0.83×宽 KV packer 在稀疏度 0 下不再带来额外带宽收益128×128 始终是最贵的基线sabi 分辨率最细但每个保留块的元数据开销最大高稀疏度下最快格子是 Q∈{256, 1024} 且 KV1024——两者在稀疏度 0.9 时均达 9.77×相对 PFA denseKV512 也在约 5% 误差内9.00×–9.68×。因此粒度选择是模式保真度 vs 端到端延迟的应用层权衡内核并不强制某一个。PFA 参考npu_fusion_attentiondense是形状不变的——各块形状间Ref_Latency_[usec]列的微小差异只是同一负载重跑 16 次的计时噪声。示例 3softmax_lse 正确性测试test_lse.pytest_lse.py 校验内核的第二路输出softmax_lselog-sum-exp运行方式pytest test_lse.py -v每个测试用例都会在全部16 个默认块形状{128, 256, 512, 1024} × {128, 256, 512, 1024}下通过参数化block_shapefixture见 test_lse.py各跑一遍因此实际用例数约为下表每形状数量的 16 倍约 1600 个实例。少量小 S 或高稀疏度组合会被自动跳过——此时测试参数化本质上不可满足frame 强制密度超过1 − sparsity或某行的期望保留 KV 块数小于 1跳过规则即文件中的pytest.skip逻辑。各块形状下测试的内容仅 B1测试参考形状数据类型备注test_bsa_lse_vs_fiasFIAS dense14 个形状S 至 24000bfloat16, float16稠密模式无 sabitest_bsa_lse_vs_fias_sparseFIAS sparse_mode18 个形状S 至 24000bfloat16H1无边框 frametest_bsa_lse_vs_fias_sparse_framedFIAS sparse_mode12 个形状S ≥ 16384bfloat16H1FrameWidths(2,1,3,1)test_bsa_lse_sparse_multiheadPython float32 参考4 个形状S ≤ 1024bfloat16H2,4Python 循环 sabitest_bsa_lse_zero_input解析值log(S)5 个形状bfloat16, float16QKV0 ⟹ 均匀 softmaxtest_bsa_lse_shape_and_dtype—3 个形状bfloat16检查 shape[B,H,S]、dtypefloat32test_bsa_lse_disabled_returns_empty—3 个形状bfloat16flagFalse ⟹ numel0为什么以 FIAS 为 LSE 参考FIASnpu_fused_infer_attention_score已经通过softmaxLseFlagbool 属性暴露softmax_lse输出输出索引 1形状[B, N, S, 1]float32。BSA 计算的是同一物理量——对参与注意的 token 求log(Σ exp(q·kᵀ/√d))——因此配合匹配的 token 级注意力 mask 的 FIAS 就是 H1 形状的即插即用参考。H1 时不能直接用 FIAS其sparse_mode1要求[B,1,S,S]广播 mask传[B,H,S,S]会挂死 NPU所以改用循环遍历 sabi 索引的 Python float32 参考。容差ATOL 0.001RTOL 0.01。float16 被排除在稀疏测试之外FIAS 会按 flash-attention 分块处理带-infmask 的瓦片而 BSA 整体跳过这些块高稀疏度下两者 LSE 偏差超出任何合理容差。测试环境文档给出的验证环境结果均在该环境取得复现时请以仓库实际内容为准组件版本Host CPUaarch64DeviceAscend 910B2Device Driver25.3.rc1Docker imagedocker pull --platformarm64 swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-910b-ubuntu22.04-py3.10-opsOSUbuntu 22.04CANN8.5.0-beta.1Python3.11.10torch2.8.0cpuTorchNPU2.8.0小结与实操建议验证正确性短序列4k–10k用RUN_REFERENCE TrueTORCH_REFERENCE npu_fusion_attention快速比对系统回归用pytest test_attn.py/pytest test_lse.py/pytest test_joint.pytest_attn.py 在 24000/36000 序列、1–4 头、10 档稀疏度、16 种块形状下全量扫描且通过is_block_sparse_pattern_feasible在收集阶段剔除不可行组合基准性能修改benchmark.py顶部参数后直接python benchmark.py需要曲线图时接| python plot.py粒度选型追求模式分辨率选小BLOCK_SIZE_KV如 128追求低稀疏度下的延迟选大BLOCK_SIZE_KV512/1024BLOCK_SIZE_Q影响有限可按 sabi 张量大小约束自由选择注意限制仅B1正确TORCH_REFERENCE用 TorchNPU mask 参考时只能H1稀疏度 0 下相对 dense 参考约有 17–22% 的固定开销短序列约 ≤4k测得的负加速属预期现象长序列注意超时S 极大且 B·S·N·D 乘积较大时可能出现 aicore timeout建议做 S 切分详见 README.md 的约束说明。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考