基于 PTO 指令集实现混合精度 Flash AttentionTFA 算子工程搭建、流水线优化与性能调优实战Ascend A2/A3【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa本指南以开源仓库 cann/pto-isa 中的flash_atten示例为主线完整讲解如何使用 PTOParallel Tile Operation虚拟指令集在 Ascend A2910B/ A3910C上实现一个 FlashAttention 2.0 风格的混合精度注意力算子TFA从工程目录、构建与运行命令、case 生成机制到compute_qk → compute_p → compute_pv → compute_gu四阶段内核实现、Cube/Vector 跨核流水线编排、多核切分与负载均衡再到可复现的参考性能数据。读完本文你将掌握这套算子的完整落地路径并能直接复现运行、按需扩展 case 或继续沿着文末的优化方向深入调优。1. 示例定位与适用平台flash_atten位于仓库 kernels/automode/a2a3/flash_atten 目录注意示例 README 内部书写的目录布局为kernels/manual/common/flash_atten/与当前仓库实际目录kernels/automode/a2a3/flash_atten/不一致本文一律以仓库实际路径为准。它演示了如何用 PTO 实现一个**混合精度mixed precision**的 Flash Attention 算子覆盖三类完整链路工程组织、构建、运行与性能测量。适用 AI 处理器Ascend A2Ascend 910BAscend A3Ascend 910C在 run.sh 中SOC 版本的正则校验为^Ascend910B|^Ascend910_9599即本示例在脚本层面同时兼容 A2/A3/A5 系列而 CMakeLists.txt 中仅对Ascend910B1映射dav-c220与Ascend910_9599映射dav-c310给出 aicore 架构映射其余 SOC 会直接报Unsupported SOC_VERSION错误。A2/A3 上物理核数为 25内核注释中Max-no-of-physical-cores (which is 25 in A2/A3)。2. 工程目录结构实际目录下包含内核实现、宿主侧驱动、构建脚本、case/数据生成脚本与流水线分析工具kernels/automode/a2a3/flash_atten/ ├── scripts/ │ ├── gen_data.py # 生成 Q/K/V 输入与 golden 输出FP16 输入、FP32 精度计算 │ ├── generate_cases.py # 生成 case 配置产出 generated_cases.h / generated_cases.json │ ├── pipeline_log_analysis.py # 流水线日志分析 │ ├── pipeline_schedule_gen.py # 流水线调度图生成 │ └── run_timeline.sh # 运行时间线timeline辅助脚本 ├── CMakeLists.txt # 构建配置cce 编译选项、SIM/NPU 两种 RUN_MODE ├── fa_performance_kernel.cpp # TFA 内核实现compute_qk / compute_p / compute_pv / compute_gu ├── fa_performance_kernel.h # LaunchTFA 模板声明与默认参数常量 ├── main.cpp # Host 侧独立驱动无 gtest ├── multiBuffer.hpp # 多缓冲 / 多级流水辅助模板 ├── pto_macro_fa_gu.hpp # GU 阶段宏TGU_ND / TGU_LAST_ND 累加与归一化 ├── pto_macro_fa_softmax.hpp # TSOFTMAXFA 流式 softmax 宏数值稳定递推 ├── pto_macro_matmul.hpp # matmul 宏matTile → accTile含多种 AccMode ├── fa_flows.svg # FA 计算流程示意图 ├── fa_pipeline.svg # 跨 CV FIFO 与阶段内 ping/pong 示意 ├── fa_pipeline_preload{0,2,4}_generated.svg # 不同预取深度的理论流水线图 ├── fa_sim_run_pipeline_h128_s0_128_s1_1024.svg # 仿真运行流水线图 ├── main.cpp # Host 侧入口 └── run.sh # 一键构建 运行脚本3. 构建与运行3.1 环境准备首先配置 Ascend CANN 环境示例路径source ${ASCEND_INSTALL_PATH}/bin/setenv.bashCMakeLists.txt 要求环境变量ASCEND_HOME_PATH已设置否则报Cannot find ASCEND_HOME_PATH, please run set_env.sh.构建同时依赖/usr/local/Ascend/driver路径下的 kernel 头文件。内核编译需要 PTO 编译器选项--cce-pto-enable --cce-pto-auto-enableC 标准为 C17编译器为 bisheng。3.2 运行示例默认 casecd ${git_clone_path}/kernels/automode/a2a3/flash_atten # 运行默认 case与 generated_cases.* 内置集合一致 bash run.sh -r npu -v Ascend910B1 # 从内置集合中只运行一个 case bash run.sh -r npu -v Ascend910B1 -c case_float_H_128_S0_128_S1_1024成功时输出test success该输出来自 main.cpp 中最终 golden 对比o_ok为真时打印test success否则test failed。同时会在工作目录追加写入report.csv记录每个 case 的case,HEAD,S0,S1,CUBE_S0,CUBE_S1,TILE_S1,start_time,end_time,duration_us,avg_block_us,GOPS,TFLOPS,result一行方便批量对比。3.3 run.sh 参数详解run.sh 通过getopt支持以下参数短参数长参数含义默认值-r--run-mode运行模式npu板上或sim仿真必填-v--soc-versionSOC 版本如Ascend910B1、Ascend910_9599必填-n--npuNPU 设备 ID0-c--case只运行指定 casecase 名或数值元组空运行全部-a--cases自定义 case 集合分号分隔空用内置集合-p--qk-preloadcube 侧 QK 预取深度qkPreloadNum4-i--intermediate使能中间张量qk/p/pv/exp_max FIFO逐块校验关-d--debugDebug 构建-DDEBUG_MODEON开启cce::printf关-k--mask使能 causal mask关注意事项来自脚本实现SOC 版本必须以Ascend910B或Ascend910_9599开头否则报Unsupported SocVersionAscend910B4-1不支持sim模式需用Ascend910B4或Ascend910_9599sim模式会通过LD_LIBRARY_PATH注入${ASCEND_HOME_PATH}/tools/simulator/${SOC_VERSION}/libCMake 侧RUN_MODEsim时链接runtime_camodelnpu时链接runtime脚本在Ascend910_9599A3 仿真下会给可执行文件追加--sys_cnt_multiple1.0A2/A3 板上默认 20.0见 main.cpp 中g_sys_cnt_multiple 20.0的注释与--sys_cnt_multiple参数解析脚本显式传入-DCMAKE_LINK_DEPENDS_USE_LINKEROFF因为 bisheng/cce-ld 不接受 CMake 3.27 默认开启的--dependency-file链接器选项。3.4 自定义 casecase 由run.sh转发给 generate_cases.py格式为分号分隔的元组HEAD_SIZE,S0,S1,CUBE_S0[,TILE_S1]。其中CUBE_S1固定为128TILE_S1支持128等于CUBE_S1、256、512缺省为256校验规则S1必须能被CUBE_S1128整除、TILE_S1必须是CUBE_S1的倍数、S1必须能被TILE_S1整除、qk_preload 1CUBE_S0大于S0或不整除时自动回退为S0。# 提供自定义 case分号分隔HEAD_SIZE,S0,S1,CUBE_S0,TILE_S1 bash run.sh -r npu -v Ascend910B1 --cases 128,128,1024,128,128;128,2048,2048,128,512 # 提供自定义 case并只运行其中一个 bash run.sh -r npu -v Ascend910B1 --cases 128,128,1024,128,128;128,512,2048,128,128 \ -c case_float_H_128_S0_128_S1_1024main.cpp的 case 筛选接受两种形式规范 case 名case_float_H_HEAD_S0_S0_S1_S1或数值元组HEAD,S0,S1[,CUBE_S0[,TILE_S1]]后者会被自动归一化为规范名。若未提供任何筛选条件则运行内置默认集合(H128,S0 128/512) × (S1 1024/2048/8192)共 6 个 case全部以qk_preload4、causal_maskfalse生成见 generate_cases.py 中DEFAULT_CASES。generate_cases.py每次运行会把 case 列表渲染为build/generated_cases.h宏TFA_FOR_EACH_CASE(MACRO)kGeneratedTfaCases[]数组和build/generated_cases.json宿主与内核两侧通过同一份宏实例化LaunchTFA与run_tfa保证 case 集合一致。数据由 gen_data.py 生成随机种子固定np.random.seed(7)Q/K 输入以 FP16 落盘q.bin/kt.bin/v.bingolden 以 FP32 计算qk.bin、softmaxx_exp、每 tile 的 global_sum/exp_max、最终o.bin等is_causal时对 QK 加上-3.40282e38的上三角 mask与内核TTRI的实现一致。4. 参考性能数据A2/A34.1 指标定义S0query 序列长度Q/O 的行数S1key/value 序列长度K/V 的行数Total task time (us)每个 task 的端到端 kernel 时间微秒GOps该 task 计数的总运算量main.cpp 中按S0 * S1 * HEAD_SIZE * 4 / 1e6统计即 QK 与 PV 两次矩阵乘的乘加共 4 次浮点操作TFLOPSGOps / timeNormalized TFLOPSTFLOPS × (24 / cores_used)用于估算在 24 核 A3 上的满芯吞吐。4.2 小结要点更大的S1能提升利用率归一化吞吐从S11024到S14096增幅非常明显对S181921–2 核的最佳归一化吞吐非常接近≈172.9 vs 171.4说明此时 kernel 更接近带宽/同步上限而非随核数线性扩展4/8 核在较小S1时归一化吞吐反而更低通常由固定开销主导流水线 warm-up、同步、内存搬运仿真Simulation数值可能显著高于板上实测模拟器不建模全部硬件争用/时延特性性能决策应以板上onboard数据为准。4.3 数据表归一化 TFLOPS越高越好CoresS0S11024S12048S14096S18192112838.2762.62147.08172.86225648.5173.04148.03171.43451238.6058.10138.19149.278102425.2837.5199.94120.04Total task timeus越低越好CoresS0S11024S12048S14096S18192112842.0851.4443.8074.54225633.20144.10143.52175.162451241.72155.44146.62186.3228102463.7285.88264.461107.342GOpsCoresS0S11024S12048S14096S18192112867.11134.22268.44536.872256134.22268.44536.871073.744512268.44536.871073.742147.4881024536.871073.742147.484294.97实测 TFLOPSCoresS0S11024S12048S14096S1819211281.592.616.137.2022564.046.0912.3414.2945126.439.6823.0324.88810248.4312.5033.3140.01仿真与板上对比Seq 2KRunTotal task time (us)Normalized TFLOPS板上NPU147.9221.78仿真Simulation28.59112.66上表均为示例目录记录的参考数据用于说明不同核数/序列长度下的相对趋势实际数字随 CANN 版本、环境与编译选项变化请以本地复现结果为准。5. 算子原理FlashAttention 2.0 分块计算流程5.1 数学形式令 Q ∈ ℝ^{S0×H}、K ∈ ℝ^{H×S1}、V ∈ ℝ^{S1×H}其中 H 为HEAD_SIZE。单头 attention 的标准形式省略 softmax 常数项为$$\text{QK} Q K^\top \in \mathbb{R}^{S0\times S1}$$ $$P \operatorname{softmax}!\left(\frac{\text{QK}}{\tau}\right)\in \mathbb{R}^{S0\times S1}$$ $$O P,V \in \mathbb{R}^{S0\times H}$$为降低显存占用并提升访存效率QK 与 softmax 按 (S0, S1) 分块tile流式计算在遍历 S1-tiles 的过程中持续更新输出 O 的 running sum。5.2 数值稳定的 tiled softmax 递推对每一行 (i)处理当前 tile 时做如下递推与常见数值稳定 softmax 写法等价步骤 1局部行最大值— (m_i \max_j X_{ij})local_max当前 tile 每行的最大值步骤 2更新全局最大值— (M_i \max!\big(M_{\text{prev},i},; m_i\big))new_global_max合并历史全局 max 与当前 local max步骤 3历史项重标定系数— (\text{exp_max}i \exp!\Big(s \cdot (M{\text{prev},i} - M_i)\Big))l1_exp_max当全局 max 增大时用该指数因子重标定历史累加项步骤 4逐元素指数— (e_{ij} \exp!\Big(s \cdot (X_{ij} - M_i)\Big))p_tile_fp32/x_expFP32 缓冲中保存逐元素指数x_expcast 为 fp16 供后续 matmul步骤 5本 tile 局部和— (\ell_i \sum_j e_{ij})local_sum当前 tile 的逐行指数和步骤 6更新全局和— (S_i \text{exp_max}i \cdot S{\text{prev},i} \ell_i)l2_global_sum数值稳定的递推累加处理完所有 tiles 后最终 softmax 概率为 (p_{ij} e_{ij} / S_i)。关键备注缩放系数 (s 1/\sqrt{\text{HEAD_SIZE}})在 pto_macro_fa_softmax.hpp 中通过constexpr_inv_sqrt(HEAD_SIZE)在编译期计算全局 max 增大时重标定历史累加项保证数值稳定kernel 保存x_exp供compute_pv使用同时保留l1_exp_max与l2_global_sum供compute_gu做 running 累加与最终归一化。计算流程整体如下图所示5.3 各阶段张量形状输入QS0 × HEAD_SIZEfp16KS1 × HEAD_SIZEfp16VS1 × HEAD_SIZEfp16每个 S1 tile 的中间量tile tqk_tileS0 × Cube_S1fp32 累加例如64×128/128×128p_tilex_expS0 × Cube_S1fp16用于 matmulpv_tileS0 × HEAD_SIZEfp32每个 tile 的部分结果输出OS0 × HEAD_SIZEfp16/fp32在源码层面fa_performance_kernel.cpp 定义了TileMatQDataCube_S0×HEAD_SIZE、TileMatKDataHEAD_SIZE×Cube_S1、TileQKData TileAccfloat, Cube_S0, Cube_S1、TileMatPDataCube_S0×Cube_S1、TileMatVDataCube_S1×HEAD_SIZE、TilePVData TileAccfloat, Cube_S0, HEAD_SIZE累加器一律使用 FP32。6. 分阶段实现与调参内核由四个阶段构成compute_qkcube 侧 QK 矩阵乘、compute_pvector 侧流式 softmaxTSOFTMAXFA、compute_pvcube 侧 P·V 矩阵乘、compute_guvector 侧归约与最终归一化。6.1compute_qkCube matmul作用计算单个 S1 tile 的 Q·K_tcube pipeline。实现要点Q tile leftTile 驻留当tile_idx0时加载一次 Q后续 tiles 只加载 K减少从全局内存GM的重复读取qk 部分结果写入紧凑的 ping/pong 全局缓冲复用pto_macro_matmulmatTile → accTile由它决定 CubeK 方向的 tilingleft/right tile 定义并维护 left/right tile 的 ping/pong 状态。调参点assign_running_acc_tile让输出 accTile 在compute_qk与compute_pv之间双缓冲避免生产者/消费者重叠时的写后读WAR冒险qkPreloadNum默认 4同时决定qkp_tile_fifo_size 1 qkPreloadNum在 fa_performance_kernel.cpp 中实际由CV_FIFO_SIZE模板参数统一控制默认kFaCvFifoSize 8用于 cube 生产者与 vector softmax 消费者之间的 FIFO 深度开启UF_ENABLE时使用AccMode::InitFinalSum关闭时退化为AccMode::Init。6.2compute_pVector softmaxTSOFTMAXFA作用在 S1 维度按 tile 增量计算并保持数值稳定的 tiled softmax。实现要点Vector tilingVec_S0 S0 / VEC_CORES每个 vector subblock 处理Vec_S0 × Cube_S1VEC_CORES默认 2控制 S0 行在 vector subblock 间的切分每个 vector core 用get_subblockid()计算全局张量 load/store 的 tile 索引与 qk/p/pv/o buffer 偏移天然形成 SPMD 并行TSOFTMAXFA微内核pto_macro_fa_softmax.hpp负责 softmax 递推保存每 tile 的l1_exp_max与l2_global_sum供compute_gu做 running 累加并在最后一步计算最终 O。其指令序列对应上文递推TROWMAXlocal/global max→TROWEXPANDSUB减 max→TMULS乘 scale→TEXP逐元素指数→TROWSUM逐行求和→TCVTFP32→FP16 cast 得到x_exp首个 tile 走softmax_opt_fa_init_implinit 专用后续 tile 走softmax_opt_fa_not_init_impl带exp_max重标定causal mask 通过TTRI生成上三角并乘-3.40282e38加到 QK 上实现对角线偏移1 s0_index % TileDataS1::Cols。实现取舍优先使用固定 tile 尺寸的TROWMAX/TROWSUM128/256/512/1024 reduce 轴上的实现通常更高效对动态有效行/列可先做TFILLPADPAD_MIN/-INF把动态 mask 转成静态例如处理动态 S0TROWEXPANDSUB支持原地计算dstsrc可减少临时 buffer精心交错 1D reduce tile 计算与 2D 计算以复用 vector 单元内的 pipe barrier bubble。UB 分配allocate_vec_tile_buffers目的为compute_p/compute_gu的 per-vector tiles 预先规划 UB 偏移让 vector cores 复用一小组固定 UB 地址从而可以用静态偏移直接TASSIGN/TSTORE/TLOAD模板参数SrcBuffers、XexpBuffers、pvVecBuffers、ExpMaxBuffersExpMaxBuffers通常等于qkp_tile_fifo_size保证每个 preload buffer 有独立的l1_exp_maxtile典型分配顺序qkvec tiles →m1_local_max→m2_global_max→input_reduce_tmp中间指数用单个 float tile→l1_local_sum→l2_global_sum→l1_exp_max[]数量 ExpMaxBuffers→x_exp[]数量 XexpBuffers→ 尾部放置runningOTile便于 GUTLOAD。6.3compute_pvP·V matmul作用把每个 tile 的 Psoftmax 输出与对应 V tile 相乘得到 PV 的部分累加cube matmul 风格。实现要点加载 V tile 与 P tileP 通过TPOP从 p FIFO 取把pv_tile_fifo写入全局 float buffer 的 per-tile ping/pong 缓冲使用AccMode细分InitFinalSum/InitPartialSum/AccFinalSum/AccPartialSum以支持 Tile_S1 内多个 Cube_S1 子块的连续累加以及 causal mask 下的跳过逻辑。调参点pv_tile_fifo_size通常与 qk FIFO 深度一致控制 P 生产与 GU 消费之间的 FIFO 深度。6.4compute_gu归约与归一化作用消费pv_tile_fifo并累加到runningOTile最后一个 tile 触发对l2_global_sum的最终除法得到输出 O。实现要点vector core 驱动使用TGU_ND/TGU_LAST_ND宏pto_macro_fa_gu.hpp 中的pto_macro_fa_gu/pto_macro_fa_gu_last/pto_macro_fa_gu_single_and_last_tile做 per-tile 累加首个 tile 直接把 PV 作为runningOTile的初值tile 0中间 tile 用l1_exp_max_ififo重标定后再累加末 tile 除以l2_global_sum得到最终 OGU_Phase::Epilogue阶段将runningOTile经TSTORE写回全局 O。实现取舍保持runningOTile绑定assigned避免重复分配TROWEXPANDMUL/TROWEXPANDDIV支持原地计算dstsrc减少临时 buffer。7. 流水线编排Cube/Vector 并行与预取深度7.1 跨阶段 CV FIFO 与阶段内 ping/pong跨阶段通过 CV FIFO 阶段内 ping/pong 做软件流水化S1 tiles 循环中的典型流程cubecompute_qk预加载下一批 QK tile并通过 flag 通知 vectorvectorcompute_p等待 qk 就绪在该 chunk 上运行TSOFTMAXFA产出 p tile并通知 pv 消费者cubecompute_pv消费 p 与 v生成pv_tile_fifo写回全局并通知 GU 消费者vectorcompute_gu消费pv_tile_fifo并累加到runningOTile。源码层面fa_performance_kernel.cpp 用三条TMPipeGM_FIFO把四阶段串起来QKPipeflagBUF0_QK_READYcubecompute_qk生产qk_tile_fifofp32vectorcompute_p消费PPipeflagBUF1_SM_READYvectorcompute_p生产p_tile_fifofp16x_exp与exp_max_ififocubecompute_pv消费PVPipeflagUPDATE_READYcubecompute_pv生产pv_tile_fifofp32vectorcompute_gu消费。FIFO 深度统一由模板参数CV_FIFO_SIZE默认kFaCvFifoSize 8控制消费同步周期CV_FIFO_CONS_SYNC_PERIOD默认kFaCvFifoConsSyncPeriod 4决定每隔几个 tile 才做一次消耗确认配合should_wait_consumption/should_notify_consumption把轻量 device flag 的同步开销摊薄kernel 尾部通过pending_consumption_events计算尚未被 wait 的 notify 并统一 drain。cube/vector 之间的跨核同步复用 FFTS 机制set_ffts_base_addr、TSYNC_CVID、ffts_cross_core_sync等。阶段内关键机制matmul_macro_ptoassign_running_acc_tileleftTile/rightTile/AccTile 双缓冲使 cube core 能在 preload 序列里交错compute_qk与compute_pvcompute_p的 qk 输入与 p 输出也做双缓冲expTl1_exp_max_ififo[qkp_tile_fifo_size]提供多 preload buffer支持更晚的结果转发K/P/V 的 Mat tilekMatTNBuffers2、pMatTNBuffers2、vMatTNBuffers2与 vector 侧的srcVecTNBuffers2、xexpVecTNBuffers2提供 L1/UB 层 ping-pong注释建议保持 1–2 层除非实测到明显 stall 气泡再加深。7.2 阶段重排用预取打破数据依赖以 Head128、S0128、S11024 为例CUBE_S1128时共有 4 个 loop每个 loop 执行compute_qk - compute_p - compute_pv - compute_gu。compute_qk与compute_pv在 cube core 执行compute_p与compute_gu在 vector core 执行。若不预执行preload0四段会趋向串行执行增大qkPreloadNum让 cube pipeline “跑在前面”提前产出后续 S1 的 QK tile从而摊平 vector 侧资源仿真H128Seqlen1024行为与理论趋势一致从该 case 的流水线图可以看到瓶颈更偏向 cube 侧的TSTORECube 利用率约 30%文档明确后续优化会优先围绕这一点展开。7.3 调参入口tuning knobsqkPreloadNum允许 cube pipeline 预先产出更多 QK tiles更大 preload 提升 Cube/VEC 并发度但增大 FIFO 占用qkp_tile_fifo_sizeqk/p FIFO 深度用于生产者/消费者重叠pv_tile_fifo_sizePV FIFO 深度通常与 qk FIFO 深度匹配用于 PV 生产与 GU 消费重叠同步优先用轻量 device flags 减少 stall在 UB/L1 允许的前提下尽量用更深的 FIFO/更大的 preload 拉开重叠。值得注意README 中写qkp_tile_fifo_size 1 qkPreloadNum而当前源码将两者解耦——qkPreloadNum模板参数QK_PRELOAD默认 4只控制预热深度FIFO 深度由CV_FIFO_SIZE默认 8独立控制此外 fa_performance_kernel.h 还固定kFaCubeS1128、kFaTileS1256。调参时应以当前源码行为为准。7.4 流水线辅助工具仓库 scripts 目录提供了 pipeline_schedule_gen.py生成理论流水线调度图、pipeline_log_analysis.py分析运行日志与 run_timeline.sh用于把内核 profile 数据与理论调度做对比分析README 中若干*_generated.svg图即由这些脚本产出。8. 多核切分与负载均衡8.1 多核 tiling 与工作分配QKV 输入为 BNSD 布局Batch、Head 数、Seqlen、HEAD_SIZE计算过程中产生中间 QK(S0,S1)。由于 S1 是归约轴多核切分通常按 (B, N, S/Cube-S0) 展开内核启动block_rows S0 / CUBE_S0个 blockLaunchTFAblock_rows, nullptr, stream每个 block 独立负责一段Cube_S0行通过get_block_idx()计算行偏移在 Flash-decoding 场景中(B, N, S/Cube-S0) 较小未来可以考虑沿 S1 轴切分文档标注 TODO每核保留部分 O再由另一个 kernel 做最终 GU对很大的 S0超过 A2/A3 最大物理核数 25时中间 FIFO buffer 可能引入浪费与不必要的 L2 回写可按 core id 做进一步优化TODO。值得一提的工程细节main.cpp 在LaunchTFA前先启动warmup_kernel24, nullptr, stream预热所有核再通过PTO_PREFETCH把 Q/K/V 预取到 L2A2/A3 上使用 SDMA 路径以减少冷启动与首次访存开销。8.2 负载均衡指导引入 causal attention mask 时需要考虑计算稀疏性TODO多核 tiling 需关注沿 S0 轴的负载不均当前做法大 S0 时采用 block 数 物理核数的多 block 启动方式用更多小任务分摊负载文档还提出未来可探索更多策略如按 mask 形状自适应分配行块。9. 源码级验证与调试建议中间张量校验bash run.sh -r npu -v Ascend910B1 -i开启INTERMEDIATE_CHECK后内核会额外TSTORE暴露 qk/p/exp_max/pv 各 FIFO 内容按 block 与 tile 落盘block{b}_qk_fifo.bin、block{b}_p_fifo.bin、block{b}_p_max_fifo.bin、block{b}_pv_fifo.bin等main.cpp 依据 golden 重建期望 FIFO 并逐块逐 tile 对比qk/p/pv 容差 1e-3exp_max 容差 1e-2输出[CHECK] FIFO intermediate ok或逐 tile 的失败清单最终输出对比无论是否开启中间校验都会对比o.bingolden与o_out.bin设备输出容差 1e-3打印test success/test failed性能统计profile buffer 记录每个 block 的 cube 与两个 vector subblock 起止时间每 block 3KBmain.cpp据此计算 block 时长、端到端时长与 TFLOPS并追加写report.csv--sys_cnt_multiple用于把系统计数器 tick 换算为时间A2/A3 默认 20.0A3 仿真 1.0Debug 构建-d参数使能-DDEBUG_MODEON_DEBUG--cce-enable-print内核内cce::printf会打印每个核/block 的起止时间。10. 总结与后续优化方向本文完整梳理了 TFA 示例从工程搭建到源码级实现的全链路运行层面环境、run.sh 参数、case 生成、数据生成、性能测量、原理层面FA2.0 数学、tiled softmax 递推、张量形状、实现层面四阶段内核、CV FIFO、ping/pong、preload 阶段重排、UB 分配、以及多核切分与负载均衡。结合 fa_performance_kernel.cpp、main.cpp 与脚本源码可以确认归一化吞吐在小S1受固定开销主导、S18192时逼近带宽/同步上限且当前瓶颈偏 cube 侧TSTORECube 利用率约 30%。文档与源码中标注的 TODO 方向即后续优化主线缓解 cubeTSTORE瓶颈、Flash-decoding 场景沿 S1 轴切分多核部分 O 独立 GU kernel、大 S0 时按 core id 裁剪中间 FIFO 以降低 L2 回写、以及 causal mask 稀疏性感知的多核负载均衡。复现与扩展该示例时建议以板上实测数据为准并从调整qkPreloadNum、FIFO 深度kFaCvFifoSize与TILE_S1入手观察流水线行为变化。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
基于 PTO 指令集实现混合精度 Flash AttentionTFA 算子工程搭建、流水线优化与性能调优实战Ascend A2/A3【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa本指南以开源仓库 cann/pto-isa 中的flash_atten示例为主线完整讲解如何使用 PTOParallel Tile Operation虚拟指令集在 Ascend A2910B/ A3910C上实现一个 FlashAttention 2.0 风格的混合精度注意力算子TFA从工程目录、构建与运行命令、case 生成机制到compute_qk → compute_p → compute_pv → compute_gu四阶段内核实现、Cube/Vector 跨核流水线编排、多核切分与负载均衡再到可复现的参考性能数据。读完本文你将掌握这套算子的完整落地路径并能直接复现运行、按需扩展 case 或继续沿着文末的优化方向深入调优。1. 示例定位与适用平台flash_atten位于仓库 kernels/automode/a2a3/flash_atten 目录注意示例 README 内部书写的目录布局为kernels/manual/common/flash_atten/与当前仓库实际目录kernels/automode/a2a3/flash_atten/不一致本文一律以仓库实际路径为准。它演示了如何用 PTO 实现一个**混合精度mixed precision**的 Flash Attention 算子覆盖三类完整链路工程组织、构建、运行与性能测量。适用 AI 处理器Ascend A2Ascend 910BAscend A3Ascend 910C在 run.sh 中SOC 版本的正则校验为^Ascend910B|^Ascend910_9599即本示例在脚本层面同时兼容 A2/A3/A5 系列而 CMakeLists.txt 中仅对Ascend910B1映射dav-c220与Ascend910_9599映射dav-c310给出 aicore 架构映射其余 SOC 会直接报Unsupported SOC_VERSION错误。A2/A3 上物理核数为 25内核注释中Max-no-of-physical-cores (which is 25 in A2/A3)。2. 工程目录结构实际目录下包含内核实现、宿主侧驱动、构建脚本、case/数据生成脚本与流水线分析工具kernels/automode/a2a3/flash_atten/ ├── scripts/ │ ├── gen_data.py # 生成 Q/K/V 输入与 golden 输出FP16 输入、FP32 精度计算 │ ├── generate_cases.py # 生成 case 配置产出 generated_cases.h / generated_cases.json │ ├── pipeline_log_analysis.py # 流水线日志分析 │ ├── pipeline_schedule_gen.py # 流水线调度图生成 │ └── run_timeline.sh # 运行时间线timeline辅助脚本 ├── CMakeLists.txt # 构建配置cce 编译选项、SIM/NPU 两种 RUN_MODE ├── fa_performance_kernel.cpp # TFA 内核实现compute_qk / compute_p / compute_pv / compute_gu ├── fa_performance_kernel.h # LaunchTFA 模板声明与默认参数常量 ├── main.cpp # Host 侧独立驱动无 gtest ├── multiBuffer.hpp # 多缓冲 / 多级流水辅助模板 ├── pto_macro_fa_gu.hpp # GU 阶段宏TGU_ND / TGU_LAST_ND 累加与归一化 ├── pto_macro_fa_softmax.hpp # TSOFTMAXFA 流式 softmax 宏数值稳定递推 ├── pto_macro_matmul.hpp # matmul 宏matTile → accTile含多种 AccMode ├── fa_flows.svg # FA 计算流程示意图 ├── fa_pipeline.svg # 跨 CV FIFO 与阶段内 ping/pong 示意 ├── fa_pipeline_preload{0,2,4}_generated.svg # 不同预取深度的理论流水线图 ├── fa_sim_run_pipeline_h128_s0_128_s1_1024.svg # 仿真运行流水线图 ├── main.cpp # Host 侧入口 └── run.sh # 一键构建 运行脚本3. 构建与运行3.1 环境准备首先配置 Ascend CANN 环境示例路径source ${ASCEND_INSTALL_PATH}/bin/setenv.bashCMakeLists.txt 要求环境变量ASCEND_HOME_PATH已设置否则报Cannot find ASCEND_HOME_PATH, please run set_env.sh.构建同时依赖/usr/local/Ascend/driver路径下的 kernel 头文件。内核编译需要 PTO 编译器选项--cce-pto-enable --cce-pto-auto-enableC 标准为 C17编译器为 bisheng。3.2 运行示例默认 casecd ${git_clone_path}/kernels/automode/a2a3/flash_atten # 运行默认 case与 generated_cases.* 内置集合一致 bash run.sh -r npu -v Ascend910B1 # 从内置集合中只运行一个 case bash run.sh -r npu -v Ascend910B1 -c case_float_H_128_S0_128_S1_1024成功时输出test success该输出来自 main.cpp 中最终 golden 对比o_ok为真时打印test success否则test failed。同时会在工作目录追加写入report.csv记录每个 case 的case,HEAD,S0,S1,CUBE_S0,CUBE_S1,TILE_S1,start_time,end_time,duration_us,avg_block_us,GOPS,TFLOPS,result一行方便批量对比。3.3 run.sh 参数详解run.sh 通过getopt支持以下参数短参数长参数含义默认值-r--run-mode运行模式npu板上或sim仿真必填-v--soc-versionSOC 版本如Ascend910B1、Ascend910_9599必填-n--npuNPU 设备 ID0-c--case只运行指定 casecase 名或数值元组空运行全部-a--cases自定义 case 集合分号分隔空用内置集合-p--qk-preloadcube 侧 QK 预取深度qkPreloadNum4-i--intermediate使能中间张量qk/p/pv/exp_max FIFO逐块校验关-d--debugDebug 构建-DDEBUG_MODEON开启cce::printf关-k--mask使能 causal mask关注意事项来自脚本实现SOC 版本必须以Ascend910B或Ascend910_9599开头否则报Unsupported SocVersionAscend910B4-1不支持sim模式需用Ascend910B4或Ascend910_9599sim模式会通过LD_LIBRARY_PATH注入${ASCEND_HOME_PATH}/tools/simulator/${SOC_VERSION}/libCMake 侧RUN_MODEsim时链接runtime_camodelnpu时链接runtime脚本在Ascend910_9599A3 仿真下会给可执行文件追加--sys_cnt_multiple1.0A2/A3 板上默认 20.0见 main.cpp 中g_sys_cnt_multiple 20.0的注释与--sys_cnt_multiple参数解析脚本显式传入-DCMAKE_LINK_DEPENDS_USE_LINKEROFF因为 bisheng/cce-ld 不接受 CMake 3.27 默认开启的--dependency-file链接器选项。3.4 自定义 casecase 由run.sh转发给 generate_cases.py格式为分号分隔的元组HEAD_SIZE,S0,S1,CUBE_S0[,TILE_S1]。其中CUBE_S1固定为128TILE_S1支持128等于CUBE_S1、256、512缺省为256校验规则S1必须能被CUBE_S1128整除、TILE_S1必须是CUBE_S1的倍数、S1必须能被TILE_S1整除、qk_preload 1CUBE_S0大于S0或不整除时自动回退为S0。# 提供自定义 case分号分隔HEAD_SIZE,S0,S1,CUBE_S0,TILE_S1 bash run.sh -r npu -v Ascend910B1 --cases 128,128,1024,128,128;128,2048,2048,128,512 # 提供自定义 case并只运行其中一个 bash run.sh -r npu -v Ascend910B1 --cases 128,128,1024,128,128;128,512,2048,128,128 \ -c case_float_H_128_S0_128_S1_1024main.cpp的 case 筛选接受两种形式规范 case 名case_float_H_HEAD_S0_S0_S1_S1或数值元组HEAD,S0,S1[,CUBE_S0[,TILE_S1]]后者会被自动归一化为规范名。若未提供任何筛选条件则运行内置默认集合(H128,S0 128/512) × (S1 1024/2048/8192)共 6 个 case全部以qk_preload4、causal_maskfalse生成见 generate_cases.py 中DEFAULT_CASES。generate_cases.py每次运行会把 case 列表渲染为build/generated_cases.h宏TFA_FOR_EACH_CASE(MACRO)kGeneratedTfaCases[]数组和build/generated_cases.json宿主与内核两侧通过同一份宏实例化LaunchTFA与run_tfa保证 case 集合一致。数据由 gen_data.py 生成随机种子固定np.random.seed(7)Q/K 输入以 FP16 落盘q.bin/kt.bin/v.bingolden 以 FP32 计算qk.bin、softmaxx_exp、每 tile 的 global_sum/exp_max、最终o.bin等is_causal时对 QK 加上-3.40282e38的上三角 mask与内核TTRI的实现一致。4. 参考性能数据A2/A34.1 指标定义S0query 序列长度Q/O 的行数S1key/value 序列长度K/V 的行数Total task time (us)每个 task 的端到端 kernel 时间微秒GOps该 task 计数的总运算量main.cpp 中按S0 * S1 * HEAD_SIZE * 4 / 1e6统计即 QK 与 PV 两次矩阵乘的乘加共 4 次浮点操作TFLOPSGOps / timeNormalized TFLOPSTFLOPS × (24 / cores_used)用于估算在 24 核 A3 上的满芯吞吐。4.2 小结要点更大的S1能提升利用率归一化吞吐从S11024到S14096增幅非常明显对S181921–2 核的最佳归一化吞吐非常接近≈172.9 vs 171.4说明此时 kernel 更接近带宽/同步上限而非随核数线性扩展4/8 核在较小S1时归一化吞吐反而更低通常由固定开销主导流水线 warm-up、同步、内存搬运仿真Simulation数值可能显著高于板上实测模拟器不建模全部硬件争用/时延特性性能决策应以板上onboard数据为准。4.3 数据表归一化 TFLOPS越高越好CoresS0S11024S12048S14096S18192112838.2762.62147.08172.86225648.5173.04148.03171.43451238.6058.10138.19149.278102425.2837.5199.94120.04Total task timeus越低越好CoresS0S11024S12048S14096S18192112842.0851.4443.8074.54225633.20144.10143.52175.162451241.72155.44146.62186.3228102463.7285.88264.461107.342GOpsCoresS0S11024S12048S14096S18192112867.11134.22268.44536.872256134.22268.44536.871073.744512268.44536.871073.742147.4881024536.871073.742147.484294.97实测 TFLOPSCoresS0S11024S12048S14096S1819211281.592.616.137.2022564.046.0912.3414.2945126.439.6823.0324.88810248.4312.5033.3140.01仿真与板上对比Seq 2KRunTotal task time (us)Normalized TFLOPS板上NPU147.9221.78仿真Simulation28.59112.66上表均为示例目录记录的参考数据用于说明不同核数/序列长度下的相对趋势实际数字随 CANN 版本、环境与编译选项变化请以本地复现结果为准。5. 算子原理FlashAttention 2.0 分块计算流程5.1 数学形式令 Q ∈ ℝ^{S0×H}、K ∈ ℝ^{H×S1}、V ∈ ℝ^{S1×H}其中 H 为HEAD_SIZE。单头 attention 的标准形式省略 softmax 常数项为$$\text{QK} Q K^\top \in \mathbb{R}^{S0\times S1}$$ $$P \operatorname{softmax}!\left(\frac{\text{QK}}{\tau}\right)\in \mathbb{R}^{S0\times S1}$$ $$O P,V \in \mathbb{R}^{S0\times H}$$为降低显存占用并提升访存效率QK 与 softmax 按 (S0, S1) 分块tile流式计算在遍历 S1-tiles 的过程中持续更新输出 O 的 running sum。5.2 数值稳定的 tiled softmax 递推对每一行 (i)处理当前 tile 时做如下递推与常见数值稳定 softmax 写法等价步骤 1局部行最大值— (m_i \max_j X_{ij})local_max当前 tile 每行的最大值步骤 2更新全局最大值— (M_i \max!\big(M_{\text{prev},i},; m_i\big))new_global_max合并历史全局 max 与当前 local max步骤 3历史项重标定系数— (\text{exp_max}i \exp!\Big(s \cdot (M{\text{prev},i} - M_i)\Big))l1_exp_max当全局 max 增大时用该指数因子重标定历史累加项步骤 4逐元素指数— (e_{ij} \exp!\Big(s \cdot (X_{ij} - M_i)\Big))p_tile_fp32/x_expFP32 缓冲中保存逐元素指数x_expcast 为 fp16 供后续 matmul步骤 5本 tile 局部和— (\ell_i \sum_j e_{ij})local_sum当前 tile 的逐行指数和步骤 6更新全局和— (S_i \text{exp_max}i \cdot S{\text{prev},i} \ell_i)l2_global_sum数值稳定的递推累加处理完所有 tiles 后最终 softmax 概率为 (p_{ij} e_{ij} / S_i)。关键备注缩放系数 (s 1/\sqrt{\text{HEAD_SIZE}})在 pto_macro_fa_softmax.hpp 中通过constexpr_inv_sqrt(HEAD_SIZE)在编译期计算全局 max 增大时重标定历史累加项保证数值稳定kernel 保存x_exp供compute_pv使用同时保留l1_exp_max与l2_global_sum供compute_gu做 running 累加与最终归一化。计算流程整体如下图所示5.3 各阶段张量形状输入QS0 × HEAD_SIZEfp16KS1 × HEAD_SIZEfp16VS1 × HEAD_SIZEfp16每个 S1 tile 的中间量tile tqk_tileS0 × Cube_S1fp32 累加例如64×128/128×128p_tilex_expS0 × Cube_S1fp16用于 matmulpv_tileS0 × HEAD_SIZEfp32每个 tile 的部分结果输出OS0 × HEAD_SIZEfp16/fp32在源码层面fa_performance_kernel.cpp 定义了TileMatQDataCube_S0×HEAD_SIZE、TileMatKDataHEAD_SIZE×Cube_S1、TileQKData TileAccfloat, Cube_S0, Cube_S1、TileMatPDataCube_S0×Cube_S1、TileMatVDataCube_S1×HEAD_SIZE、TilePVData TileAccfloat, Cube_S0, HEAD_SIZE累加器一律使用 FP32。6. 分阶段实现与调参内核由四个阶段构成compute_qkcube 侧 QK 矩阵乘、compute_pvector 侧流式 softmaxTSOFTMAXFA、compute_pvcube 侧 P·V 矩阵乘、compute_guvector 侧归约与最终归一化。6.1compute_qkCube matmul作用计算单个 S1 tile 的 Q·K_tcube pipeline。实现要点Q tile leftTile 驻留当tile_idx0时加载一次 Q后续 tiles 只加载 K减少从全局内存GM的重复读取qk 部分结果写入紧凑的 ping/pong 全局缓冲复用pto_macro_matmulmatTile → accTile由它决定 CubeK 方向的 tilingleft/right tile 定义并维护 left/right tile 的 ping/pong 状态。调参点assign_running_acc_tile让输出 accTile 在compute_qk与compute_pv之间双缓冲避免生产者/消费者重叠时的写后读WAR冒险qkPreloadNum默认 4同时决定qkp_tile_fifo_size 1 qkPreloadNum在 fa_performance_kernel.cpp 中实际由CV_FIFO_SIZE模板参数统一控制默认kFaCvFifoSize 8用于 cube 生产者与 vector softmax 消费者之间的 FIFO 深度开启UF_ENABLE时使用AccMode::InitFinalSum关闭时退化为AccMode::Init。6.2compute_pVector softmaxTSOFTMAXFA作用在 S1 维度按 tile 增量计算并保持数值稳定的 tiled softmax。实现要点Vector tilingVec_S0 S0 / VEC_CORES每个 vector subblock 处理Vec_S0 × Cube_S1VEC_CORES默认 2控制 S0 行在 vector subblock 间的切分每个 vector core 用get_subblockid()计算全局张量 load/store 的 tile 索引与 qk/p/pv/o buffer 偏移天然形成 SPMD 并行TSOFTMAXFA微内核pto_macro_fa_softmax.hpp负责 softmax 递推保存每 tile 的l1_exp_max与l2_global_sum供compute_gu做 running 累加并在最后一步计算最终 O。其指令序列对应上文递推TROWMAXlocal/global max→TROWEXPANDSUB减 max→TMULS乘 scale→TEXP逐元素指数→TROWSUM逐行求和→TCVTFP32→FP16 cast 得到x_exp首个 tile 走softmax_opt_fa_init_implinit 专用后续 tile 走softmax_opt_fa_not_init_impl带exp_max重标定causal mask 通过TTRI生成上三角并乘-3.40282e38加到 QK 上实现对角线偏移1 s0_index % TileDataS1::Cols。实现取舍优先使用固定 tile 尺寸的TROWMAX/TROWSUM128/256/512/1024 reduce 轴上的实现通常更高效对动态有效行/列可先做TFILLPADPAD_MIN/-INF把动态 mask 转成静态例如处理动态 S0TROWEXPANDSUB支持原地计算dstsrc可减少临时 buffer精心交错 1D reduce tile 计算与 2D 计算以复用 vector 单元内的 pipe barrier bubble。UB 分配allocate_vec_tile_buffers目的为compute_p/compute_gu的 per-vector tiles 预先规划 UB 偏移让 vector cores 复用一小组固定 UB 地址从而可以用静态偏移直接TASSIGN/TSTORE/TLOAD模板参数SrcBuffers、XexpBuffers、pvVecBuffers、ExpMaxBuffersExpMaxBuffers通常等于qkp_tile_fifo_size保证每个 preload buffer 有独立的l1_exp_maxtile典型分配顺序qkvec tiles →m1_local_max→m2_global_max→input_reduce_tmp中间指数用单个 float tile→l1_local_sum→l2_global_sum→l1_exp_max[]数量 ExpMaxBuffers→x_exp[]数量 XexpBuffers→ 尾部放置runningOTile便于 GUTLOAD。6.3compute_pvP·V matmul作用把每个 tile 的 Psoftmax 输出与对应 V tile 相乘得到 PV 的部分累加cube matmul 风格。实现要点加载 V tile 与 P tileP 通过TPOP从 p FIFO 取把pv_tile_fifo写入全局 float buffer 的 per-tile ping/pong 缓冲使用AccMode细分InitFinalSum/InitPartialSum/AccFinalSum/AccPartialSum以支持 Tile_S1 内多个 Cube_S1 子块的连续累加以及 causal mask 下的跳过逻辑。调参点pv_tile_fifo_size通常与 qk FIFO 深度一致控制 P 生产与 GU 消费之间的 FIFO 深度。6.4compute_gu归约与归一化作用消费pv_tile_fifo并累加到runningOTile最后一个 tile 触发对l2_global_sum的最终除法得到输出 O。实现要点vector core 驱动使用TGU_ND/TGU_LAST_ND宏pto_macro_fa_gu.hpp 中的pto_macro_fa_gu/pto_macro_fa_gu_last/pto_macro_fa_gu_single_and_last_tile做 per-tile 累加首个 tile 直接把 PV 作为runningOTile的初值tile 0中间 tile 用l1_exp_max_ififo重标定后再累加末 tile 除以l2_global_sum得到最终 OGU_Phase::Epilogue阶段将runningOTile经TSTORE写回全局 O。实现取舍保持runningOTile绑定assigned避免重复分配TROWEXPANDMUL/TROWEXPANDDIV支持原地计算dstsrc减少临时 buffer。7. 流水线编排Cube/Vector 并行与预取深度7.1 跨阶段 CV FIFO 与阶段内 ping/pong跨阶段通过 CV FIFO 阶段内 ping/pong 做软件流水化S1 tiles 循环中的典型流程cubecompute_qk预加载下一批 QK tile并通过 flag 通知 vectorvectorcompute_p等待 qk 就绪在该 chunk 上运行TSOFTMAXFA产出 p tile并通知 pv 消费者cubecompute_pv消费 p 与 v生成pv_tile_fifo写回全局并通知 GU 消费者vectorcompute_gu消费pv_tile_fifo并累加到runningOTile。源码层面fa_performance_kernel.cpp 用三条TMPipeGM_FIFO把四阶段串起来QKPipeflagBUF0_QK_READYcubecompute_qk生产qk_tile_fifofp32vectorcompute_p消费PPipeflagBUF1_SM_READYvectorcompute_p生产p_tile_fifofp16x_exp与exp_max_ififocubecompute_pv消费PVPipeflagUPDATE_READYcubecompute_pv生产pv_tile_fifofp32vectorcompute_gu消费。FIFO 深度统一由模板参数CV_FIFO_SIZE默认kFaCvFifoSize 8控制消费同步周期CV_FIFO_CONS_SYNC_PERIOD默认kFaCvFifoConsSyncPeriod 4决定每隔几个 tile 才做一次消耗确认配合should_wait_consumption/should_notify_consumption把轻量 device flag 的同步开销摊薄kernel 尾部通过pending_consumption_events计算尚未被 wait 的 notify 并统一 drain。cube/vector 之间的跨核同步复用 FFTS 机制set_ffts_base_addr、TSYNC_CVID、ffts_cross_core_sync等。阶段内关键机制matmul_macro_ptoassign_running_acc_tileleftTile/rightTile/AccTile 双缓冲使 cube core 能在 preload 序列里交错compute_qk与compute_pvcompute_p的 qk 输入与 p 输出也做双缓冲expTl1_exp_max_ififo[qkp_tile_fifo_size]提供多 preload buffer支持更晚的结果转发K/P/V 的 Mat tilekMatTNBuffers2、pMatTNBuffers2、vMatTNBuffers2与 vector 侧的srcVecTNBuffers2、xexpVecTNBuffers2提供 L1/UB 层 ping-pong注释建议保持 1–2 层除非实测到明显 stall 气泡再加深。7.2 阶段重排用预取打破数据依赖以 Head128、S0128、S11024 为例CUBE_S1128时共有 4 个 loop每个 loop 执行compute_qk - compute_p - compute_pv - compute_gu。compute_qk与compute_pv在 cube core 执行compute_p与compute_gu在 vector core 执行。若不预执行preload0四段会趋向串行执行增大qkPreloadNum让 cube pipeline “跑在前面”提前产出后续 S1 的 QK tile从而摊平 vector 侧资源仿真H128Seqlen1024行为与理论趋势一致从该 case 的流水线图可以看到瓶颈更偏向 cube 侧的TSTORECube 利用率约 30%文档明确后续优化会优先围绕这一点展开。7.3 调参入口tuning knobsqkPreloadNum允许 cube pipeline 预先产出更多 QK tiles更大 preload 提升 Cube/VEC 并发度但增大 FIFO 占用qkp_tile_fifo_sizeqk/p FIFO 深度用于生产者/消费者重叠pv_tile_fifo_sizePV FIFO 深度通常与 qk FIFO 深度匹配用于 PV 生产与 GU 消费重叠同步优先用轻量 device flags 减少 stall在 UB/L1 允许的前提下尽量用更深的 FIFO/更大的 preload 拉开重叠。值得注意README 中写qkp_tile_fifo_size 1 qkPreloadNum而当前源码将两者解耦——qkPreloadNum模板参数QK_PRELOAD默认 4只控制预热深度FIFO 深度由CV_FIFO_SIZE默认 8独立控制此外 fa_performance_kernel.h 还固定kFaCubeS1128、kFaTileS1256。调参时应以当前源码行为为准。7.4 流水线辅助工具仓库 scripts 目录提供了 pipeline_schedule_gen.py生成理论流水线调度图、pipeline_log_analysis.py分析运行日志与 run_timeline.sh用于把内核 profile 数据与理论调度做对比分析README 中若干*_generated.svg图即由这些脚本产出。8. 多核切分与负载均衡8.1 多核 tiling 与工作分配QKV 输入为 BNSD 布局Batch、Head 数、Seqlen、HEAD_SIZE计算过程中产生中间 QK(S0,S1)。由于 S1 是归约轴多核切分通常按 (B, N, S/Cube-S0) 展开内核启动block_rows S0 / CUBE_S0个 blockLaunchTFAblock_rows, nullptr, stream每个 block 独立负责一段Cube_S0行通过get_block_idx()计算行偏移在 Flash-decoding 场景中(B, N, S/Cube-S0) 较小未来可以考虑沿 S1 轴切分文档标注 TODO每核保留部分 O再由另一个 kernel 做最终 GU对很大的 S0超过 A2/A3 最大物理核数 25时中间 FIFO buffer 可能引入浪费与不必要的 L2 回写可按 core id 做进一步优化TODO。值得一提的工程细节main.cpp 在LaunchTFA前先启动warmup_kernel24, nullptr, stream预热所有核再通过PTO_PREFETCH把 Q/K/V 预取到 L2A2/A3 上使用 SDMA 路径以减少冷启动与首次访存开销。8.2 负载均衡指导引入 causal attention mask 时需要考虑计算稀疏性TODO多核 tiling 需关注沿 S0 轴的负载不均当前做法大 S0 时采用 block 数 物理核数的多 block 启动方式用更多小任务分摊负载文档还提出未来可探索更多策略如按 mask 形状自适应分配行块。9. 源码级验证与调试建议中间张量校验bash run.sh -r npu -v Ascend910B1 -i开启INTERMEDIATE_CHECK后内核会额外TSTORE暴露 qk/p/exp_max/pv 各 FIFO 内容按 block 与 tile 落盘block{b}_qk_fifo.bin、block{b}_p_fifo.bin、block{b}_p_max_fifo.bin、block{b}_pv_fifo.bin等main.cpp 依据 golden 重建期望 FIFO 并逐块逐 tile 对比qk/p/pv 容差 1e-3exp_max 容差 1e-2输出[CHECK] FIFO intermediate ok或逐 tile 的失败清单最终输出对比无论是否开启中间校验都会对比o.bingolden与o_out.bin设备输出容差 1e-3打印test success/test failed性能统计profile buffer 记录每个 block 的 cube 与两个 vector subblock 起止时间每 block 3KBmain.cpp据此计算 block 时长、端到端时长与 TFLOPS并追加写report.csv--sys_cnt_multiple用于把系统计数器 tick 换算为时间A2/A3 默认 20.0A3 仿真 1.0Debug 构建-d参数使能-DDEBUG_MODEON_DEBUG--cce-enable-print内核内cce::printf会打印每个核/block 的起止时间。10. 总结与后续优化方向本文完整梳理了 TFA 示例从工程搭建到源码级实现的全链路运行层面环境、run.sh 参数、case 生成、数据生成、性能测量、原理层面FA2.0 数学、tiled softmax 递推、张量形状、实现层面四阶段内核、CV FIFO、ping/pong、preload 阶段重排、UB 分配、以及多核切分与负载均衡。结合 fa_performance_kernel.cpp、main.cpp 与脚本源码可以确认归一化吞吐在小S1受固定开销主导、S18192时逼近带宽/同步上限且当前瓶颈偏 cube 侧TSTORECube 利用率约 30%。文档与源码中标注的 TODO 方向即后续优化主线缓解 cubeTSTORE瓶颈、Flash-decoding 场景沿 S1 轴切分多核部分 O 独立 GU kernel、大 S0 时按 core id 裁剪中间 FIFO 以降低 L2 回写、以及 causal mask 稀疏性感知的多核负载均衡。复现与扩展该示例时建议以板上实测数据为准并从调整qkPreloadNum、FIFO 深度kFaCvFifoSize与TILE_S1入手观察流水线行为变化。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
