企业官网建设流程全解析

FALite v01 多 Mix 核组 Task 并行：CANN 上从单核闭环迈向整卡并行的第一步

【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples

导读

FALite（Flash Attention Lite）是 CANN 仓库cann-samples中基于 Ascend 950 的因果（causal）FlashAttention 前向教学样例，核心是一个 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）。在统一长序列规格B=N=1,S=131072,D=128下，v00 的 Task Duration 中位数为 2574789.75 μs，按整卡 432 TFLOP/s 峰值折算的有效 Cube MFU 只有 0.3954%——根本原因是其余 31 个 AIC 完全没有参与计算，而不是单个核组内部的流水不够高效。

v01 的思路非常克制：不改变任何数学公式、不改变矩阵乘/Vector 计算、不改变 SRAM 布局，只把 Query tile 任务按固定步长分发到多个 Mix 核组，让整卡算力参与进来。这使"是否用上了整卡"与"单个核组内部流水是否高效"两个问题可以分开观察，为后续 v02～v12 的流水优化建立了可比较的基线。完整版本路线与性能汇总见总 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.h）：

tr = 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 只有j=0...i共i+1个 item，因此越靠后的 task 工作量越大，这是后续讨论负载均衡的前提。

Host 侧核数检查与 useAicNum

Host 入口FlashAttnLiteNPU（flash_attn_lite_host.cpp）先做两件事：

1. 通过PlatformAscendCManager::GetInstance()->GetCoreNumAic()取得设备 AIC 核数；若请求核数超过设备 AIC 数，直接拒绝启动（requestedAicCoreNum > deviceAicCoreNum时打印错误并返回）；
2. 请求值为0表示"使用设备全部 AIC 作为上限"；合法请求再取useAicNum = min(核数上限, numTasks)，即实际核组数不会超过 task 总数。

TilingData 中的useAicNum字段由 ComputeFlashAttnLiteTilingData 计算并写入，Kernel 入口固定发射<<<data.useAicNum, 0, stream>>>（flash_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）。

图中B=2,N=2,S=640：每个(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 行 0～63，AIV1 负责 64～127。它们分别维护 64 行的m/l/alpha和OAcc，无需交换或合并这些状态。四个阶段的分工如下表：

GM 与片上数据布局

GM 输入输出与中间 workspace

Q/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 在 FlashAttnLiteNPU 中通过aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_FIRST)分配dS/dP/dDeltaO三块 workspace，并在释放前调用aclrtSynchronizeStream(stream)同步（v03 起 workspace 归零后即可提交即返回）。

AIC 缓冲与生命周期

Q/K/V/P的 L1 数据为 BF16 NZ（Cube 读取所需的分块排列），以下均为单槽、不存在 slot 翻转：

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 共占 115456 B（112.75 KiB），两路 AIV 各有一套。S、DeltaO、P 和 alpha 按 item 复用；m、l、OAcc 按 task 复用。AIV 侧这些地址同样由 Host TilingData 下发（flash_attn_lite_host.cpp）。

AIC 核内流水与同步

C1 与 C2 的数据通路

C1（falite_kernel_aic.h）按K GM→L1→L0A和Q L1→L0B准备输入，再执行 Mmad 和 Fixpipe。Q 的 MTE2 写入发生在 task 开始；MTE1 在j==0取得 Q 槽所有权，到j+1==kvTileCount（causal 裁剪后的实际 item 数）时才归还。CopyGmToL1通过Nd2NzParams完成 GM 的 ND 到 L1 的 NZ 转换，并在validRows < tileRows时用Fill对 L1 末尾行补零。

C2（falite_kernel_aic.h）先把 P 从 GM 搬入 L1，再将 P 转置装入 L0A（LoadData2DParamsV2的ifTranspose=true）；随后发射 V 的 GM→L1 和 L1→L0B 搬运。两次 Mmad 都初始化 L0C（cmatrixInitVal=true），而不是在 L0C 内跨 item 累加；跨 item 的输出累计由 AIV 的 OAcc 完成。

Mutex 槽位交接

每个 Pipe 用Mutex::Lock/Unlock取得和归还对应资源（ID 定义见 falite_kernel_aic.h）：

要点：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 读入 UB（MUTEX_S_UB=1在 MTE2 写入与 Vector 读取之间交接），再由 Vector 处理。OnlineColwiseSoftmaxVF将 FP32 未归一化权重写回 S，随后Cast（MUTEX_P_UB=0的 Vector Lock/Unlock 包住）转成普通二维 BF16 P；MTE3 用同一 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 搬入 UB（MUTEX_DELTA_O_UB=2），用OnlineUpdateVF对每行广播 alpha 顺序更新 OAcc。m/l/OAcc 不设多槽，也不在两路 AIV 之间做归约。

task 结束时，FusedDivCastVF将OAcc/l转为连续 BF16 输出，仍写进 P UB；Vector→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.h：

表中的 V、FIX 对应PIPE_V、PIPE_FIX。真机 mode2（GROUP_CROSS_MODE=2）下，一次 AIC Set 同时通知同组两路 AIV；AIC 等待 P_READY 或 DONE 时，必须等两路 AIV 都发出对应信号。SIM_COMPATIBLE=ON的仿真封装改用 mode4 分别处理两路（PAIR_CROSS_MODE=4，AIV1 使用flagId + 16的偏移），逻辑 flag 和阶段不变——该分支在 falite_kernel_common.h 的SetAicToAiv/WaitAicToAiv/SetAivToAic/WaitAivToAic四组封装中统一处理。

关于 DONE 有三点需要特别注意：

1. DONE 不是 task 输出完成信号。它在每次 V2 后发出，最后一次 DONE 仍早于归一化与 O 写回；其消费者是下一 item 的 AIC MTE1，跨 task 时也照常消费。AIC 的首个 item 跳过 DONE 等待，全部 task 结束后再消费最后一次 DONE；没有预置的 DONE，也不能在每个 task 开头都跳过它；
2. DONE 绑定 MTE1，不是 Scalar 全局屏障。下一轮 K 或下一 task Q 的 MTE2 可以先进入自己的队列，真正使用单槽结果的 MTE1 及其下游仍受 DONE 限制；
3. P 不单独设置"已读完"信号。C2 读 P 后才能产生 O_READY，随后 V2 发 DONE、下一 C1 发 S_READY，AIV 才生成下一份 P——这条依赖链保护了 P 的跨 item 复用；核内 Mutex 另行保护各 Pipe 对本地工作区的读写。

调度伪代码

两侧保持独立的 task/item 循环；with mutex(资源, Pipe)表示该 Pipe 的 Lock/Unlock，C1/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 # CubeStage1：j==0 取得 Q 的 MTE1 所有权， # j+1==kvTileCount 时归还；内部按 Mutex 交接 L0。 CubeStage1(j) # Fixpipe -> S GM SetAicToAiv(FIX, S_READY) WaitAivToAic(MTE2, P_READY) CubeStage2(j) # P GM -> L0A；Fixpipe -> 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(m=FLOAT_LOWEST, l=0, OAcc=0) 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 循环起点都取aicIdx：AIV 侧通过aicIdx = GetBlockIdx() / GetSubBlockNum()还原同组 AIC 的编号（falite_kernel_aiv.h），从而保证同组核心遍历完全相同的 task 序列。

首轮、末轮与尾块

每个 task 都重新初始化 m、l、OAcc。首个 V1 直接建立 m/l 并写alpha=1；OAcc 初值为 0，所以首个 V2 仍调用普通更新函数（OnlineUpdateVF天然兼容0 = alpha*0 + DeltaO）。注意："本核首个 item 不等 DONE"与"每个 task 的首个 V1"不是同一个条件。

causal 下只发射j=0...i；对角 item 的 V1 在求最大值与指数和的两次扫描中都屏蔽keyLocal>queryLocal（逻辑 Query 行、Key 列矩阵的上三角）。两路 AIV 分别使用 Query 局部起点 0、64（queryColBegin = subAivIdx * halfBr，见 falite_kernel_aiv.h），不能都从 0 计算 mask——这是两路切片正确性的关键细节。

末尾不足 128 行时，CopyGmToL1只读取紧凑 GM 的有效行并在 L1 补零，Cube 仍执行固定 128 大小的矩阵乘；V1 只扫描有效 Key（validBc），将其余 Key 对应的 P 置零（falite_kernel_aiv.h）。GM 中间缓冲仍按完整 tile 保存，最终输出才按有效 Query 行裁剪。即使 AIV1 的outputRows=0，也必须完成全部 V1/V2 和 P_READY/DONE，只跳过最终归一化与 GM 输出，否则同组 AIC 的聚合等待无法完成。本版逐 item 循环，没有两个 item 一组的奇偶尾组分支。

流水示意与性能证据

示意图说明本版多核组 task 归属与组内 GM 交接的关系；色块宽度和图中空隙不代表实际耗时，跨 Pipe 的精确完成顺序应结合源码同步判断。

截图来自B=1,N=1,S=2048,D=128的完整 PipeTimeline trace，窗口为[232.419,272.419] μs；AIV0/AIV1 的 MTE2 泳道对应中间块从 GM 读回 UB 的搬运，单槽阶段之间仍有明显间隔。不能把某个空隙直接解释为 DONE 或 P_READY 的精确等待时长。

统一长序列性能结果（B=1,N=1,S=131072,D=128，32 个 AIC）：v01 的 Task Duration 中位数为88377.640625 μs，有效 Cube MFU 为 11.52%。相较固定 1 个 AIC 的 v00（2574789.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 级并行这一单一变量，为后续数据通路和流水优化建立清晰基线。

代码阅读入口与运行

在 cann-samples 根目录执行以下命令即可构建并运行本版（先按需安装依赖python3 -m pip install -r Samples/2_Performance/flash_attn_lite_story/requirements.txt）：

cmake -S . -B build -DNPU_ARCH=dav-3510 -DSIM_COMPATIBLE=OFF 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-metrics=BasicInfo/PipeTimeline）见总 README 的编译、运行与复现。

与相邻版本的区别

相对 v00，v01 只增加三样东西：Host 核数选择与useAicNum计算、按固定步长的 task 循环、各核心以aicIdx为起点的遍历逻辑；矩阵乘、Vector 计算和 SRAM 布局完全不变。后续 v02 则移除 S/DeltaO 的 GM 往返，仍不增加槽数——从此进入"缩短数据通路"的下一个优化维度。

【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples