FALite v01 多 Mix 核组 Task 并行CANN 上从单核闭环迈向整卡并行的第一步【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples导读FALiteFlash Attention Lite是 CANN 仓库cann-samples中基于 Ascend 950 的因果causalFlashAttention 前向教学样例核心是一个 Mix 核组内 1 个矩阵计算核AIC 2 个向量计算核AIV的协作流水。本文讲解其演进链中的 v01 版本在保持 v00 四阶段计算、GM 中间数据通路和单槽片上缓冲全部不变的前提下仅引入多个 Mix 核组按固定步长分担 Query tile 任务这一层 task 级并行。读完本文你将掌握 FALite 的 task/item 编号体系、Host 侧核数选择逻辑、AIC/AIV 两路分工、GM 与片上数据布局、核内 Mutex 与核间 CrossCore 双同步机制以及本版相对 v00 的性能收益与剩余瓶颈。v01 的定位先解决没用上整卡的问题v00 固定发射一个 Mix 核组由 1 个 AIC 和 2 个 AIV 顺序遍历全部 Query 分块任务task。在统一长序列规格BN1,S131072,D128下v00 的 Task Duration 中位数为 2574789.75 μs按整卡 432 TFLOP/s 峰值折算的有效 Cube MFU 只有 0.3954%——根本原因是其余 31 个 AIC 完全没有参与计算而不是单个核组内部的流水不够高效。v01 的思路非常克制不改变任何数学公式、不改变矩阵乘/Vector 计算、不改变 SRAM 布局只把 Query tile 任务按固定步长分发到多个 Mix 核组让整卡算力参与进来。这使是否用上了整卡与单个核组内部流水是否高效两个问题可以分开观察为后续 v02v12 的流水优化建立了可比较的基线。完整版本路线与性能汇总见总 README 的版本演进与性能优化。task 归属与两路 AIV 分工编号体系task、item 与 batch/head/tile 分解一个 task 完成某个(b,n)的一个 Query tile最多输出 128 行一个 item 是该 task 与一个 K/V tile 的组合即一组 C1/V1/C2/V2。Host 与 Kernel 使用同一套编号规则见 flash_attn_lite_host.cpp 与 falite_kernel_aic.htr ceil(S / 128) numTasks B * N * tr batchHeadIdx taskId / tr qTileIdx taskId % tr batchIdx batchHeadIdx / N headIdx batchHeadIdx % N其中tr对应数学符号 $T_r$每个(b,n)的 Query tile 数量numTasks是全部 task 总数。causal 模式下第i个 Query tile 只有j0...i共i1个 item因此越靠后的 task 工作量越大这是后续讨论负载均衡的前提。Host 侧核数检查与 useAicNumHost 入口FlashAttnLiteNPUflash_attn_lite_host.cpp先做两件事通过PlatformAscendCManager::GetInstance()-GetCoreNumAic()取得设备 AIC 核数若请求核数超过设备 AIC 数直接拒绝启动requestedAicCoreNum deviceAicCoreNum时打印错误并返回请求值为0表示使用设备全部 AIC 作为上限合法请求再取useAicNum min(核数上限, numTasks)即实际核组数不会超过 task 总数。TilingData 中的useAicNum字段由 ComputeFlashAttnLiteTilingData 计算并写入Kernel 入口固定发射data.useAicNum, 0, streamflash_attn_lite_kernel.asc。注意与 v00 不同v01 的--core-num参数会真正参与分核。固定步长分配与负载特征AIC 的核组编号是GetBlockIdx()AIV 的核组编号是GetBlockIdx()/GetSubBlockNum()同组 AIC 与两路 AIV 使用相同的起点和步长taskId aicIdx, aicIdx useAicNum, aicIdx 2 * useAicNum, ...Kernel 侧的循环对应 falite_kernel_aic.h 的for (taskId firstTaskId; taskId numTasks; taskId GetBlockNum())AIV 侧则按taskId data.useAicNum遍历falite_kernel_aiv.h。图中B2,N2,S640每个(b,n)有 5 个 Query tile共 20 个 task分给 4 个 Mix 核组Mix 0 处理0/4/8/12/16颜色仅表示归属。不同 task 只共享只读的 K/V各自维护自己的 Softmax 状态m/l/alpha、OAcc写入的输出行和 GM workspace 均不重叠因此不需要跨核组归约。固定步长分配把轻重任务轮转分散但不保证各组工作量完全相等——同一轮相邻taskId未跨过(b,n)边界时各核组会读取高度重合的 K/V 前缀可能命中 L2跨 Head/Batch 时则没有这类共享。两路 AIV 的 Query 行切片两路 AIV 用subAivIdx GetSubBlockIdx()区分AIV0 负责 Query 行 063AIV1 负责 64127。它们分别维护 64 行的m/l/alpha和OAcc无需交换或合并这些状态。四个阶段的分工如下表阶段核心本版职责C1AICK_j × Q_i^T → S^TBF16 输入、FP32 累加V1两路 AIV计算分块 Softmax更新m/l/alpha生成尚未除以完整分母的 BF16 PC2AICP_j × V_j → DeltaOBF16 输入、FP32 累加V2两路 AIVOAcc alpha × OAcc DeltaO保持 FP32task 末尾才除以 l、转 BF16 并写回GM 与片上数据布局GM 输入输出与中间 workspaceQ/K/V/O 在 GM 中紧凑排列单个 token 连续保存 128 个 BF16 通道某个 tile 的首元素偏移为(batchHeadIdx*S tileIdx*128)*128以元素为单位TilingData 中Addr字段以字节为单位Elems表示元素数见 flash_attn_lite_common.h。S 和 P 的 GM 物理形状均为[numTasks,128,128]每个 task 保存一个S^T/P^T[Key,Query]不是整个S×S矩阵元素偏移为taskId*128*128 keyLocal*128 queryLocal同 task 的各 item 重复使用这一块地址。两路 AIV 按行间隔读 S、写 P每个 Key 行只取属于自己的 64 个 Query 列。DeltaO 的物理形状为[numTasks,Query,128]每路取连续的 64 个 Query 行。Host 申请的 GM workspace每 task 大小使用范围FP32 S64 KiBC1 写入V1 读入BF16 P32 KiBV1 写入C2 读入FP32 DeltaO64 KiBC2 写入V2 读入合计160 KiB按 numTasks 分配不按 item 数累加Host 在 FlashAttnLiteNPU 中通过aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_FIRST)分配dS/dP/dDeltaO三块 workspace并在释放前调用aclrtSynchronizeStream(stream)同步v03 起 workspace 归零后即可提交即返回。AIC 缓冲与生命周期Q/K/V/P的 L1 数据为 BF16 NZCube 读取所需的分块排列以下均为单槽、不存在 slot 翻转物理空间起始地址大小生命周期P L10 KiB32 KiB每 item 生成一次C2 搬到 L0A 后不再读取Q L132 KiB32 KiB每 task 搬一次供全部 C1 复用K L164 KiB32 KiBC1 搬入MTE1 读完后可复用V L196 KiB32 KiBC2 搬入MTE1 读完后可复用L0A本空间 032 KiBC1 装 KC2 装 PMmad 读完后归还L0B本空间 032 KiBC1 装 QC2 装 VMmad 读完后归还L0C本空间 064 KiBC1/C2 交替写入Fixpipe 读完后归还L1 合计 128 KiB。L0A/L0B/L0C 属于不同的物理空间不能把三个起始地址 0 当成同一块内存这些地址和容量在 Host 侧 TilingData 中显式规划并经过 SRAM 容量校验flash_attn_lite_host.cpp对应 L1 512 KiB、L0A/L0B 64 KiB、L0C 256 KiB、UB 248 KiB 的片上容量常量。单路 AIV 缓冲与生命周期UB 区域起始字节地址大小用途S032 KiB[128,64]FP32 分数V1 原地改写为 FP32 未归一化权重DeltaO3276832 KiB[64,128]FP32当轮 V2 消费OAcc6553632 KiB[64,128]FP32整个 task 累计P/Output9830416 KiBV1 的 BF16 Ptask 末尾复用为 BF16 输出m / l / alpha114688 / 114944 / 115200各 256 B每个 Query 行一个 FP32 值每路 UB 共占 115456 B112.75 KiB两路 AIV 各有一套。S、DeltaO、P 和 alpha 按 item 复用m、l、OAcc 按 task 复用。AIV 侧这些地址同样由 Host TilingData 下发flash_attn_lite_host.cpp。AIC 核内流水与同步C1 与 C2 的数据通路C1falite_kernel_aic.h按K GM→L1→L0A和Q L1→L0B准备输入再执行 Mmad 和 Fixpipe。Q 的 MTE2 写入发生在 task 开始MTE1 在j0取得 Q 槽所有权到j1kvTileCountcausal 裁剪后的实际 item 数时才归还。CopyGmToL1通过Nd2NzParams完成 GM 的 ND 到 L1 的 NZ 转换并在validRows tileRows时用Fill对 L1 末尾行补零。C2falite_kernel_aic.h先把 P 从 GM 搬入 L1再将 P 转置装入 L0ALoadData2DParamsV2的ifTransposetrue随后发射 V 的 GM→L1 和 L1→L0B 搬运。两次 Mmad 都初始化 L0CcmatrixInitValtrue而不是在 L0C 内跨 item 累加跨 item 的输出累计由 AIV 的 OAcc 完成。Mutex 槽位交接每个 Pipe 用Mutex::Lock/Unlock取得和归还对应资源ID 定义见 falite_kernel_aic.hMutex ID资源Pipe 交接0Q L1MTE2 写 → MTE1 读 → 下一 task 的 MTE21K L1MTE2 写 → MTE1 读 → 下一 C1 的 MTE22V L1MTE2 写 → MTE1 读 → 下一 C2 的 MTE23P L1MTE2 写 → MTE1 读 → 下一 C2 的 MTE24L0A/L0B 共用所有权MTE1 写 → Mmad 读 → 下一阶段 MTE15L0CMmad 写 → Fixpipe 读 → 下一阶段 Mmad要点L0A/B 在 Mmad 之后即可归还L0C 则必须等 Fixpipe 读取不能把二者的释放合并为一个笼统的矩阵乘结束。本版没有手写 HardEvent 初始化或排空循环核内交接由上述 Mutex 表达Kernel 入口先执行InitSocState()flash_attn_lite_kernel.asc。Scalar 按函数顺序发射指令不等于所有 Pipe 都已完成——C2 中先发射 P、再发射 V不能单凭代码行次序推断两个 Pipe 的完成先后数据完成与槽位复用以 Mutex、CrossCore 和同一 Pipe 内的指令顺序为准。AIV 核内流水与同步以下只描述一路 AIV另一路执行相同控制流falite_kernel_aiv.h。V1 等S_READY后由 MTE2 把 GM 的半块 S 读入 UBMUTEX_S_UB1在 MTE2 写入与 Vector 读取之间交接再由 Vector 处理。OnlineColwiseSoftmaxVF将 FP32 未归一化权重写回 S随后CastMUTEX_P_UB0的 Vector Lock/Unlock 包住转成普通二维 BF16 PMTE3 用同一 Mutex 读取 P 并写入 GM。P 保留[Key,本路Query]排列L1 的 NZ 转换由 AIC MTE2 完成。Softmax 每次加载连续的 64 个 Query 列一个 Vector 寄存器的 64 个 FP32 lane 各对应一个 Query 行沿 Key 维循环即可分别维护 64 份最大值和分母——这正是转置计算 两路 AIV 各管 64 列布局带来的收益。V2 在O_READY后将半块 DeltaO 搬入 UBMUTEX_DELTA_O_UB2用OnlineUpdateVF对每行广播 alpha 顺序更新 OAcc。m/l/OAcc 不设多槽也不在两路 AIV 之间做归约。task 结束时FusedDivCastVF将OAcc/l转为连续 BF16 输出仍写进 P UBVector→MTE3 的 Mutex 交接保证结果生成后才能写回下一次 Vector 取得该槽又要等 MTE3 读完。OAcc 的下一 task 清零排在本核归一化之后无需等待 GM 写回才清零因为 MTE3 读取的是 BF16 工作区不是 OAcc。CV 核间同步单槽的主要数据消费依赖如下CV 指 AIC 与 AIV 的交接C1 --S_READY-- V1 --P_READY-- C2 --O_READY-- V2 ^ | ------------------ DONE -----------------------flag 定义见 falite_kernel_common.hflag 名称flag IDSet 端Wait 端数据含义S_READY0AIC FIXAIV MTE2S 已写入 GM可读入 UBO_READY1AIC FIXAIV MTE2DeltaO 已写入 GM可读入 UBDONE2AIV VAIC MTE1本 item 的 V2 已消费完成P_READY4AIV MTE3AIC MTE2两路 P 已写入 GM表中的 V、FIX 对应PIPE_V、PIPE_FIX。真机 mode2GROUP_CROSS_MODE2下一次 AIC Set 同时通知同组两路 AIVAIC 等待 P_READY 或 DONE 时必须等两路 AIV 都发出对应信号。SIM_COMPATIBLEON的仿真封装改用 mode4 分别处理两路PAIR_CROSS_MODE4AIV1 使用flagId 16的偏移逻辑 flag 和阶段不变——该分支在 falite_kernel_common.h 的SetAicToAiv/WaitAicToAiv/SetAivToAic/WaitAivToAic四组封装中统一处理。关于 DONE 有三点需要特别注意DONE 不是 task 输出完成信号。它在每次 V2 后发出最后一次 DONE 仍早于归一化与 O 写回其消费者是下一 item 的 AIC MTE1跨 task 时也照常消费。AIC 的首个 item 跳过 DONE 等待全部 task 结束后再消费最后一次 DONE没有预置的 DONE也不能在每个 task 开头都跳过它DONE 绑定 MTE1不是 Scalar 全局屏障。下一轮 K 或下一 task Q 的 MTE2 可以先进入自己的队列真正使用单槽结果的 MTE1 及其下游仍受 DONE 限制P 不单独设置已读完信号。C2 读 P 后才能产生 O_READY随后 V2 发 DONE、下一 C1 发 S_READYAIV 才生成下一份 P——这条依赖链保护了 P 的跨 item 复用核内 Mutex 另行保护各 Pipe 对本地工作区的读写。调度伪代码两侧保持独立的 task/item 循环with mutex(资源, Pipe)表示该 Pipe 的 Lock/UnlockC1/C2 的核内资源交接按前表执行不额外添加全 Pipe 等待# AIC first True for taskId in range(aicIdx, numTasks, useAicNum): i taskId % tr kvTileCount i 1 with mutex(Q_L1, MTE2): copy_q_to_l1(valid_rows(i)) # 尾行补零 for j in range(kvTileCount): if not first: WaitAivToAic(MTE1, DONE) first False # CubeStage1j0 取得 Q 的 MTE1 所有权 # j1kvTileCount 时归还内部按 Mutex 交接 L0。 CubeStage1(j) # Fixpipe - S GM SetAicToAiv(FIX, S_READY) WaitAivToAic(MTE2, P_READY) CubeStage2(j) # P GM - L0AFixpipe - DeltaO GM SetAicToAiv(FIX, O_READY) WaitAivToAic(MTE1, DONE) # 只消费末 item 的 DONE# AIV0 / AIV1各处理自己的 64 个 Query 行 for taskId in range(aicIdx, numTasks, useAicNum): i taskId % tr init(mFLOAT_LOWEST, l0, OAcc0) for j in range(i 1): WaitAicToAiv(MTE2, S_READY) with mutex(S_UB, MTE2): copy_s_half_from_gm() with mutex(S_UB, V): VectorStage1(j, i) # 内含 P_UB 的 V Lock/Cast/Unlock with mutex(P_UB, MTE3): copy_p_half_to_gm() SetAivToAic(MTE3, P_READY) WaitAicToAiv(MTE2, O_READY) with mutex(DeltaO_UB, MTE2): copy_delta_o_half_from_gm() with mutex(DeltaO_UB, V): VectorStage2() SetAivToAic(V, DONE) outputRows clamp(valid_rows(i) - subAivIdx * 64, 0, 64) if outputRows 0: with mutex(P_UB, V): FusedDivCastVF(outputRows) with mutex(P_UB, MTE3): copy_output_to_gm(outputRows * 128)注意 AIC 与 AIV 的 task 循环起点都取aicIdxAIV 侧通过aicIdx GetBlockIdx() / GetSubBlockNum()还原同组 AIC 的编号falite_kernel_aiv.h从而保证同组核心遍历完全相同的 task 序列。首轮、末轮与尾块每个 task 都重新初始化 m、l、OAcc。首个 V1 直接建立 m/l 并写alpha1OAcc 初值为 0所以首个 V2 仍调用普通更新函数OnlineUpdateVF天然兼容0 alpha*0 DeltaO。注意本核首个 item 不等 DONE与每个 task 的首个 V1不是同一个条件。causal 下只发射j0...i对角 item 的 V1 在求最大值与指数和的两次扫描中都屏蔽keyLocalqueryLocal逻辑 Query 行、Key 列矩阵的上三角。两路 AIV 分别使用 Query 局部起点 0、64queryColBegin subAivIdx * halfBr见 falite_kernel_aiv.h不能都从 0 计算 mask——这是两路切片正确性的关键细节。末尾不足 128 行时CopyGmToL1只读取紧凑 GM 的有效行并在 L1 补零Cube 仍执行固定 128 大小的矩阵乘V1 只扫描有效 KeyvalidBc将其余 Key 对应的 P 置零falite_kernel_aiv.h。GM 中间缓冲仍按完整 tile 保存最终输出才按有效 Query 行裁剪。即使 AIV1 的outputRows0也必须完成全部 V1/V2 和 P_READY/DONE只跳过最终归一化与 GM 输出否则同组 AIC 的聚合等待无法完成。本版逐 item 循环没有两个 item 一组的奇偶尾组分支。流水示意与性能证据示意图说明本版多核组 task 归属与组内 GM 交接的关系色块宽度和图中空隙不代表实际耗时跨 Pipe 的精确完成顺序应结合源码同步判断。截图来自B1,N1,S2048,D128的完整 PipeTimeline trace窗口为[232.419,272.419] μsAIV0/AIV1 的 MTE2 泳道对应中间块从 GM 读回 UB 的搬运单槽阶段之间仍有明显间隔。不能把某个空隙直接解释为 DONE 或 P_READY 的精确等待时长。统一长序列性能结果B1,N1,S131072,D12832 个 AICv01 的 Task Duration 中位数为88377.640625 μs有效 Cube MFU 为 11.52%。相较固定 1 个 AIC 的 v002574789.75 μs耗时缩短 96.57%、加速 29.13 倍。注意这一对比同时改变了实际核数只说明 task 级并行的效果不代表单个核组内部流水快了 29.13 倍。完整采集环境、核数差异和 MFU 口径统一见总 README 的统一性能结果。v01 还剩的问题多个核组已经用起来了但每个核组内部的三份中间结果S、P、DeltaO仍反复经过 GM。S和DeltaO都由 AIC 产生、由 AIV 消费Ascend 950 的 Fixpipe 本可以把它们直接写入 AIV UB——这正是 v02 的改法。因此 v01 的核心价值是隔离出 task 级并行这一单一变量为后续数据通路和流水优化建立清晰基线。代码阅读入口与运行阅读顺序文件与函数主要看什么1HostComputeFlashAttnLiteTilingData、FlashAttnLiteNPUtask 数、useAicNum、地址规划、workspace 申请与同步释放2TilingDataAddr/Elems单位Host 与 Kernel 共用布局3Kernel 入口InitSocState、__mix__(1,2)声明、causal 模板实例与useAicNum发射4AICKernelProcessForAIC、CubeStage1/2task/item 循环、Q 生命周期、GM→L1→L0、单槽复用5AIVKernelProcessForAIV、VectorStage1/2两路切片、S 原地改写和 Cast、输出工作区复用6核间封装flag ID、mode2/mode4 分支、GetKvTileCount与有效行数在 cann-samples 根目录执行以下命令即可构建并运行本版先按需安装依赖python3 -m pip install -r Samples/2_Performance/flash_attn_lite_story/requirements.txtcmake -S . -B build -DNPU_ARCHdav-3510 -DSIM_COMPATIBLEOFF cmake --build build --target falite_v01 -j ./build/Samples/2_Performance/flash_attn_lite_story/falite_v01 --core-num 2 --size 2 3 257该用例包含多个 Batch/Head 和非整块尾行--core-num以 AIC 数表示核组上限不传则使用设备全部 AIC--size支持 1/2/3 个参数形式。精度标准FP32 Golden 逐元素比对abs 0.004 0.004*abs(golden)、更多运行选项与性能采集命令msopprof --aic-metricsBasicInfo/PipeTimeline见总 README 的编译、运行与复现。与相邻版本的区别相对 v00v01 只增加三样东西Host 核数选择与useAicNum计算、按固定步长的 task 循环、各核心以aicIdx为起点的遍历逻辑矩阵乘、Vector 计算和 SRAM 布局完全不变。后续 v02 则移除 S/DeltaO 的 GM 往返仍不增加槽数——从此进入缩短数据通路的下一个优化维度。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
