CANN 算子性能调优实战KvRmsNormRopeCache 从 MemBase 到 RegBase 的平滑迁移与寄存器级 VF 计算【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples本文基于 cann-samples 仓库中的KvRmsNormRopeCache实战样例完整讲解一个融合了 RMSNorm、RoPE 与 KV Cache 写回的 AscendC 算子如何从 MemBase以LocalTensor为中心的 UB 计算平滑迁移到 RegBase以RegTensor为中心的寄存器级 VF 计算。读者将掌握 MemBase 与 RegBase 的差异边界、VFVector Function设计方法、Reg::LoadAlign/Reg::StoreAlign的显式寄存器搬运范式以及一套外层数据流不动、热计算链下沉寄存器的可复用调优路线。1. 样例定位与硬件编程模型该样例位于Samples/2_Performance/kv_rms_norm_rope_cache_story/目录结构如下MemBase 版本membase/full_load.ascRegBase 版本regbase/full_load.asc公共 host 侧、tiling、数据生成与精度校验include/sample_common.h、scripts/gen_data.py构建入口CMakeLists.txt1.1 目标硬件与数据规格样例面向 Ascend 950PR/950DT编译时通过NPU_ARCHdav-3510指定架构固定走 AIV Vector kernel数据类型为 BF16。典型 shape 如下表所示张量Shape说明kv[B, N, S, (dv dk)]最后一维前dv段给 RMSNorm后dk段给 RoPEgamma[dv]RMSNorm 缩放参数cos/sin[B, N, S, dk]RoPE 旋转角三角函数k_cache[B, N, S, dk]K 缓存按 index 写回v_cache[B, N, S, dv]V 缓存按 index 写回样例固定参数定义在 sample_common.hSAMPLE_BATCH8、SAMPLE_NUM_HEAD1、SAMPLE_SEQ128、SAMPLE_DV512、SAMPLE_DK128、SAMPLE_UB_FACTOR8、SAMPLE_EPSILON1e-5f、BF16 比对阈值BF16_COMPARE_TOL6e-2f。1.2 三层存储/计算模型从性能角度看算子运行在三层存储/计算模型上GM 层输入、输出和 cache 位于全局内存通过DataCopy/DataCopyPad与 UB 交互UB 层tile 数据先搬入本地 buffer作为计算的 staging 区Register/VF 层RegBase 在__simd_vf__函数内使用RegTensor执行寄存器级计算。MemBase 与 RegBase 的核心区别不在 GM/UB 搬运而在 UB 内的计算方式层级MemBase 写法RegBase 写法GM/UB 搬运DataCopyPad(GlobalTensor - LocalTensor)保持相同UB stagingLocalTensor、TQue、TPipe保持相同计算对象LocalTensorTRegTensorT计算入口普通__aicore__成员函数__simd_vf__UB 到寄存器隐含在标准 Vector API 中显式Reg::LoadAlign寄存器到 UB隐含在标准 Vector API 中显式Reg::StoreAligntail 控制count、repeat、mask 参数MaskReg因此平滑迁移的原则是保留 MemBase 已验证过的数据流只替换 Vector 计算部分。2. 算子计算语义RMSNorm RoPE KV Cache该算子把kv最后一维拆成两段分别处理kv[..., :Dv] - RMSNorm - v_out - v_cache kv[..., Dv:] - RoPE - k_out - k_cache2.1 RMSNormRMSNorm 的计算过程为mean_square mean(x * x) rms sqrt(mean_square epsilon) v_out (x / rms) * gamma2.2 RoPERoPE 部分把交错存储的复数拆成 real/imag 两路再用两段 cos/sin 完成复数旋转real rope_x[..., 0::2] imag rope_x[..., 1::2] k_out_first real * cos_first - imag * sin_first k_out_second imag * cos_second real * sin_second2.3 数据生成与 golden 佐证上述语义与 golden 生成脚本 scripts/gen_data.py 中的build_golden完全一致rms_x取kv[..., :dv]rope_x取kv[..., dv:]k_out concat(real, imag) * cos concat(-imag, real) * sin其中part1 concat(real, imag)、part2 concat(-imag, real)恰好对应两段输出k_out_first与k_out_second。脚本还会根据index把k_out/v_out写回k_cache/v_cachecache_idx 0时写回与内核CopyOutK/CopyOutV的 cache 写回逻辑一一对应。当前规格下Dv512、Dk128。RMSNorm 每行需要处理 512 个元素包含平方、归约、开方、除法、乘gamma、BF16/FP32 转换等操作是主要优化热点RoPE 每行处理 128 个元素计算链较短但也适合用 VF 避免中间 UB 临时张量。3. 快速编译与运行在cann-samples仓库根目录执行cmake -S . -B build -DNPU_ARCHdav-3510 cmake --build build --target kv_rms_norm_rope_cache_story构建完成后会生成两个可执行文件build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_membase_full_load build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_regbase_full_load分别运行./build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_membase_full_load ./build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_regbase_full_load每个可执行文件会自动生成输入数据并执行 golden 精度校验预期最后输出包含PASS从构建脚本 CMakeLists.txt 可以看到两个.asc文件各自生成独立可执行目标目标名规则为kv_rms_norm_rope_cache_${子目录}_${文件名}并且以-O3优化级别编译、以SOURCE_DIR宏注入源码目录用于运行时定位gen_data.py。运行时的完整 host 流程封装在 sample_common.h 的RunSample中aclInit→ 创建 stream → 调用GenerateDatapython3 gen_data.py --batch 8 --seq 128 --dv 512 --dk 128→ 读取input/*.bin与output/*_golden.bin→aclrtMallocaclrtMemcpy搬入设备 → 以blockNum, nullptr, stream方式启动内核 → 搬回结果后对k_cache、v_cache、k_out、v_out四路输出逐一调用CompareBf16比对全部误差在阈值内才打印PASS。tiling 计算位于 BuildTiling通过platform_ascendc::PlatformAscendCManager获取 AIV core 数totalRows B * N * SblockFactor ceil(totalRows / coreNum)blockNum ceil(totalRows / blockFactor)ubFactor min(SAMPLE_UB_FACTOR, blockFactor)并将epsilon与reciprocal 1/Dv一并写入 tiling 结构体传给内核。4. MemBase 版本以 LocalTensor 为中心的计算链MemBase 版本位于 membase/full_load.asc整体采用Init - Process - ProcessTile的结构。4.1 外层数据流Init完成 GM 地址绑定SetGlobalBuffer、block 行数计算、TQue/TBuf初始化Process先预加载并转换共享gammaBF16DataCopy后整段Cast成 FP32随后按ubFactor循环处理 tileProcessTile串联当前 tile 的搬入、RoPE、RMSNorm 和写回。完整数据流如下Process Load gamma(BF16) - Cast gamma(FP32) for each tile: CopyRopeAndX kv[..., :Dv] - xLocal kv[..., Dv:] - ropeLocal Load cos/sin Rope ropeLocal cos/sin - k outLocal CopyOutK k outLocal - k_out k outLocal - k_cache[index] RmsNorm xLocal gammaFp32 - v outLocal CopyOutV v outLocal - v_out v outLocal - v_cache[index]在 GM/UB 搬运层CopyRopeAndX使用两段DataCopyPad源码以DataCopyExtParams指定每行dk/dv个元素的跨步搬移把后Dk维搬到ropeLocal、前Dv维搬到xLocal从而让 RoPE 和 RMSNorm 在 UB 中独立消费各自的数据段。cos、sin也按 tile 搬入cosSinQueue_与当前ropeLocal对齐。CopyOutK/CopyOutV源码按行逐条DataCopyPad写回 output并通过Mutex::LockPIPE_MTE3保护对indexGm_的读取cacheOffset 0时再按batch * cacheLength cacheOffset计算 cache 地址写回。4.2 RoPE 的 LocalTensor 实现MemBase 的Rope源码仍以LocalTensor为核心先将cos/sinCast成 FP32再通过两次GatherMask从交错 rope 数据中分别抽取 real/imag{1, 1, NUM_EIGHT, 0}步长随后执行两组乘加得到y0/y1Add合并后Cast回 BF16。这个流程直观、易验证但会在cosFp32、sinFp32、y0、y1、realFp32、imagFp32等 UB 临时张量之间产生多次读写。4.3 RMSNorm 的 LocalTensor 实现RMSNorm 是 MemBase 版本的主要热点。RmsNorm源码先将xLocal从 BF16 cast 到 FP32再执行平方、分段Add归约利用{1,1,1,NUM_EIGHT,NUM_SIXTEEN,NUM_SIXTEEN}repeat/stride 参数把 512 元素折半累加、WholeReduceSum得到行和、乘1/Dv、加epsilon、Sqrt、Brcb广播、逐行Div、乘gammaFp32最后Cast回 BF16。该链路依赖多个 UB 临时区xFp32 - square - rowSum - broadcast rms - normalized xFp32 - outLocal这些临时区全部来自wsBuffer_其分配见 Initpipe_.InitBuffer(wsBuffer_, ubFactor_ * (tiling_-dv * NUM_THREE tiling_-dk * NUM_EIGHT) * sizeof(float));即每行需要Dv*3 Dk*8个 float 的临时空间。MemBase 版本的仿真打点图如下图中展示的是一次循环的 Vector 计算内容左边是 ROPE右边是 RmsNorm所有的 Vector 指令都是一个单独的 VF整体看 Vector 指令较为细碎——这正是引入 RegBase 融合计算链的动机。5. 从 MemBase 到 RegBase 的平滑迁移边界从 MemBase 迁移到 RegBase 的整体原则有在一个完整的 VF 功能内应融尽融。理论上 VF 内融合的指令越多UB 和寄存器交互的次数越少性能越好。VF 内尽量少使用 Reduce 类、非对齐访存、Interleave/Deinterleave 类单发指令。VF 内尽量少使用 MembarMembar 会导致 VF 内流水中断使整体 IPC 不高。迁移不是推倒重写。当前 RegBase 版本复用了 MemBase 中已经验证过的外层结构Init中的 GM tensor 绑定方式保持一致BuildTiling、blockFactor、ubFactor、blockNum保持一致CopyRopeAndX中的 GM - UB 搬运保持DataCopyPadCopyOutK/CopyOutV中的 cache 写回和 output 写回保持一致inQueue、cosSinQueue、outQueue的 double bufferBUFFER_NUM2数据流保持一致输出仍走同一套 golden 校验比较k_cache、v_cache、k_out、v_out。真正变化的是RmsNorm和Rope的 compute bodyMemBase: LocalTensor - Cast/Mul/Add/WholeReduceSum/Div/Mul/Cast - LocalTensor RegBase: __ubuf__ pointer - Reg::LoadAlign - RegTensor RegTensor compute chain Reg::StoreAlign - __ubuf__ pointer这种边界划分的好处是外层索引、cache offset、batch 切分和 GM 地址计算不变降低迁移风险原有精度 golden 可以直接复用MemBase 可以作为 RegBase 的对照基线出现错误时可以快速判断问题在 GM/UB 数据流还是 VF compute bodyRegBase 的收益集中在热计算链避免为追求寄存器化引入额外 pipeline 复杂度。6. VF 方案设计RegTensor 寄存器计算RegBase 版本位于 regbase/full_load.asc。普通的RmsNorm/Rope包装函数只做一件事通过LocalTensor.GetPhyAddr()把 UB 地址转成__ubuf__指针再调用独立的__simd_vf__函数RmsNorm 包装、Rope 包装。6.1 RMSNormVF两段式 VF 循环RmsNormVF源码把每一行的 RMSNorm 拆成两个 VF 循环。第一段循环计算平方和for each row: reduceSum 0 for i in Dv chunks: xB16 LoadAlign(DIST_UNPACK_B16) xFp32 Cast(BF16 - FP32) square xFp32 * xFp32 chunk ReduceSum(square) reduceSum chunk第二段循环归一化并乘gammarms sqrt(reduceSum / Dv epsilon) for i in Dv chunks: xB16 LoadAlign(DIST_UNPACK_B16) gammaB16 LoadAlign(DIST_UNPACK_B16) xFp32 Cast(BF16 - FP32) gammaFp32 Cast(BF16 - FP32) norm (xFp32 / rms) * gammaFp32 outB16 Cast(FP32 - BF16) StoreAlign(DIST_PACK_B32)对应到源码chunk 粒度由VL_FP32_SIZE 256 / sizeof(float) 64决定dvLoop_ (dv VL_FP32_SIZE - 1) / VL_FP32_SIZE即Dv512时每行 8 个 chunk。第一段用ReduceSum(reduceSum, reduceSum, fullMask)把 64-lane 的行内和归约到单一值随后Muls乘reciprocal、Adds加epsilon、Sqrt开方、Div(invRms, one, rmsValue)取倒数再用DuplicateHighLowPart::LOWEST把标量广播成invRmsBrc供整行乘除使用第二段按 chunk 同时LoadAlignx 与 gammaMul(x, invRmsBrc)与Mul(norm, gammaFp32)两级乘后Cast回 BF16 并StoreAlign。设计要点BF16 输入先升到 FP32 做平方、归约、除法和乘法保证 RMSNorm 中间计算精度reduceSum、rmsValue、invRms、invRmsBrc、norm等短生命周期中间量全部保留在RegTensor中gamma不再预先 cast 成一整段 FP32 UB 临时张量而是在 VF 中按 chunk load BF16 并 cast 到 FP32通过UpdateMaskfloat(remaining)控制最后一个 chunk 的有效 lane避免把 tail padding 当有效数据。6.2 RopeVF寄存器内复数旋转RopeVF源码的核心思路是直接从 BF16 rope 数据中拆出 real/imag并用两段 cos/sin 完成复数旋转real, imag LoadAlign(DIST_DINTLV_B16, rope) realFp32 Cast Interleave imagFp32 Cast Interleave out_first realFp32 * cos_first - imagFp32 * sin_first out_second imagFp32 * cos_second realFp32 * sin_second StoreAlign(out_first) StoreAlign(out_second)从源码实现看real/imag 的拆分是通过两次带不同RegLayout的 Cast 完成的CAST_B16_TO_B32RegLayout::ZERO取 even lane得到realFp32CAST_B16_TO_B32_ODDRegLayout::ONE取 odd lane得到imagFp32cast trait 定义效果等价于对交错数据做 deinterleave。随后Mul/Sub/Add全部在寄存器内串联两次StoreAlign分别写回outFirst与outSecond。设计要点real/imag、cos/sin、mul/result 都在寄存器内串联无 UB 中间张量RoPE 的pairCount Dk / 2 64mask UpdateMaskfloat(pairCount)固定控制 64 个 FP32 lane输出分成前后两段写回包装函数中以out pairCount作为第二段起始地址与 golden 中concat(real, imag)的布局保持一致。7. 理论分析与实测验证7.1 UB 使用量减少MemBase RMSNorm 需要多个 UB 临时张量xFp32 - square - rowSum - broadcast rms - xFp32 norm - outLocal并且标准 Vector API 之间通常会形成 UB load/store 边界。RegBase 把平方、归约、除法、乘 gamma 等中间状态尽量保存在寄存器中只在 VF 入口 load 输入、在 VF 末尾 store 输出。理论收益来自减少中间结果落 UB减少多段 Vector API 之间的 UB 往返降低 UB 临时 buffer 占用。对比两个版本的 buffer 分配MemBase: gammaQueue: BF16 gamma FP32 gamma wsBuffer_: rows * (Dv * 3 Dk * 8) * sizeof(float) RegBase: gammaQueue: BF16 gamma no wsBuffer_RegBase 的 Init 中已无TBufTPosition::VECCALC wsBuffer_成员这是最直接的 UB 占用优化。7.2 精度路径保持一致样例输入和输出是 BF16但 RMSNorm 的关键中间计算必须使用 FP32BF16 load - FP32 compute - BF16 storeRegBase 版本通过CAST_B16_TO_B32与CAST_B32_TO_B16显式表达转换路径其中回写 cast 使用RoundMode::CAST_RINTround to nearest。实测中MemBase 和 RegBase 均通过同一套 golden 校验MemBase: k_cache/v_cache/k_out/v_out 全部 PASS RegBase: k_cache/v_cache/k_out/v_out 全部 PASS其中 RegBasev_cache/v_out的最大 diff 为0.0078125远低于 sample_common.h 中定义的 BF16 比对阈值6e-2。8. 实际调优路径8.1 保持 full-load 外层数据流当前样例一次 tile 处理ubFactor8行运行时实测打印blockNum 54, blockFactor 19, ubFactor 8这与 BuildTiling 的推导吻合总行数8 * 1 * 128 1024按 AIV core 数向上取整得到每个 block 约 19 行UB 内一次搬入最多 8 行。这个切分对 MemBase 和 RegBase 共同适用因此调优时先不改外层只验证 compute body 的收益和正确性。8.2 减少 UB 使用MemBase 版本需要wsBuffer_存放 RMSNorm 和 RoPE 的 FP32 中间结果RegBase 版本删除了这个大块 VECCALC改为在 VF 中使用RegTensor临时变量。对比MemBase: gammaQueue: BF16 gamma FP32 gamma wsBuffer_: rows * (Dv * 3 Dk * 8) * sizeof(float) RegBase: gammaQueue: BF16 gamma no wsBuffer_这是最直接的 UB 占用优化为 double buffer 深度或更大的ubFactor留出了容量空间。8.3 gamma 处理整段预转换改为 VF 内按需转换MemBase 版本在主循环前把整段gamma从 BF16 cast 成 FP32并保存在 UB 中供后续每行复用RegBase 版本保留 BF16gammaLocal在RmsNormVF的第二段循环中按 chunk load 并 cast。这样做的取舍优点减少一份完整 FP32 gamma UB buffer优点gamma 与 x 的 load/cast 粒度一致寄存器链更紧凑代价每行都会重新 cast gamma chunk。当前样例的主要目标是演示 RegBase VF 迁移和减少 UB 中间张量如果后续 profiler 证明 gamma cast 成本突出可以考虑引入小块 gamma 预处理或更细粒度复用策略但不能破坏 UB 容量和 double buffer 数据流。8.4 切换验证顺序建议按以下顺序调优避免一次修改多个变量后无法归因固定输入数据和 shape只替换 compute body确认精度 PASS对比 MemBase 和 RegBase 的输出 diff先确认 BF16 阈值内一致使用 profiler 观察 UB load/store、Vector 指令、pipeline stall若 UB 压力仍高继续检查 VF 内是否还有可消除的 store/load若算术指令成为瓶颈再检查 reduce 组织、gamma cast 复用和 div/sqrt 链路若搬运成为瓶颈再调整ubFactor、double buffer 深度或 GM copy 形态。当前已完成的功能验证kv_rms_norm_rope_cache_membase_full_load: PASS kv_rms_norm_rope_cache_regbase_full_load: PASS9. RegBase 版本运行效果RegBase 版本位于 regbase/full_load.asc。仿真打点图中左侧是 rope已经融合成 1 个 VF右侧是 RmsNorm融合成 1 个 VF。选择 Rope 中 RVEC_EX 部分的指令从下侧的简易统计可以估算出 IPC 在 1.36 左右选择 RmsNorm 中 RVEC_EX 部分的指令估算出 IPC 在 1.15 左右。RmsNorm 的 IPC 比 Rope 低是因为其中有单发指令而且代码中第一个循环体前后有依赖所以整体 IPC 比 Rope 低需要说明的是上述 IPC 数值来自仿真打点图的简易统计估算用于对比两个 VF 内部的指令密度差异实际性能收益应结合 profiler 数据进一步确认。10. 迁移实践清单从 MemBase 切到 RegBase 时可以按以下清单执行先保留 MemBase 的 host、tiling、queue、CopyIn、CopyOut找出最热的LocalTensor计算链本样例是 RMSNorm其次是 RoPE把计算链拆成独立__simd_vf__函数在普通__aicore__函数中只做LocalTensor.GetPhyAddr()到__ubuf__指针的转换VF 内只使用RegTensor、MaskReg、Reg::Load*、Reg::Store*和Reg::*compute APIBF16 路径显式写清BF16 - FP32 - BF16的 cast trait 和 round modetail 只用MaskReg控制不依赖 UB padding 参与数学每次优化后都运行同一套 golden 校验确认k_cache、v_cache、k_out、v_out全部 PASS。一句话总结RegBase 的平滑迁移路线是外层不动、热链下沉、寄存器串联、同源校验。本样例正是沿着这条路线把 MemBase 中以LocalTensor为中心的计算链切换为 VF 中以RegTensor为中心的寄存器计算链——先通过同一套 golden 校验守住精度底线再借助仿真打点观察 VF 融合效果最终在减少 UB 中间张量的同时让热计算链在寄存器内连续执行。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN 算子性能调优实战：KvRmsNormRopeCache 从 MemBase 到 RegBase 的平滑迁移与寄存器级 VF 计算
CANN 算子性能调优实战KvRmsNormRopeCache 从 MemBase 到 RegBase 的平滑迁移与寄存器级 VF 计算【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples本文基于 cann-samples 仓库中的KvRmsNormRopeCache实战样例完整讲解一个融合了 RMSNorm、RoPE 与 KV Cache 写回的 AscendC 算子如何从 MemBase以LocalTensor为中心的 UB 计算平滑迁移到 RegBase以RegTensor为中心的寄存器级 VF 计算。读者将掌握 MemBase 与 RegBase 的差异边界、VFVector Function设计方法、Reg::LoadAlign/Reg::StoreAlign的显式寄存器搬运范式以及一套外层数据流不动、热计算链下沉寄存器的可复用调优路线。1. 样例定位与硬件编程模型该样例位于Samples/2_Performance/kv_rms_norm_rope_cache_story/目录结构如下MemBase 版本membase/full_load.ascRegBase 版本regbase/full_load.asc公共 host 侧、tiling、数据生成与精度校验include/sample_common.h、scripts/gen_data.py构建入口CMakeLists.txt1.1 目标硬件与数据规格样例面向 Ascend 950PR/950DT编译时通过NPU_ARCHdav-3510指定架构固定走 AIV Vector kernel数据类型为 BF16。典型 shape 如下表所示张量Shape说明kv[B, N, S, (dv dk)]最后一维前dv段给 RMSNorm后dk段给 RoPEgamma[dv]RMSNorm 缩放参数cos/sin[B, N, S, dk]RoPE 旋转角三角函数k_cache[B, N, S, dk]K 缓存按 index 写回v_cache[B, N, S, dv]V 缓存按 index 写回样例固定参数定义在 sample_common.hSAMPLE_BATCH8、SAMPLE_NUM_HEAD1、SAMPLE_SEQ128、SAMPLE_DV512、SAMPLE_DK128、SAMPLE_UB_FACTOR8、SAMPLE_EPSILON1e-5f、BF16 比对阈值BF16_COMPARE_TOL6e-2f。1.2 三层存储/计算模型从性能角度看算子运行在三层存储/计算模型上GM 层输入、输出和 cache 位于全局内存通过DataCopy/DataCopyPad与 UB 交互UB 层tile 数据先搬入本地 buffer作为计算的 staging 区Register/VF 层RegBase 在__simd_vf__函数内使用RegTensor执行寄存器级计算。MemBase 与 RegBase 的核心区别不在 GM/UB 搬运而在 UB 内的计算方式层级MemBase 写法RegBase 写法GM/UB 搬运DataCopyPad(GlobalTensor - LocalTensor)保持相同UB stagingLocalTensor、TQue、TPipe保持相同计算对象LocalTensorTRegTensorT计算入口普通__aicore__成员函数__simd_vf__UB 到寄存器隐含在标准 Vector API 中显式Reg::LoadAlign寄存器到 UB隐含在标准 Vector API 中显式Reg::StoreAligntail 控制count、repeat、mask 参数MaskReg因此平滑迁移的原则是保留 MemBase 已验证过的数据流只替换 Vector 计算部分。2. 算子计算语义RMSNorm RoPE KV Cache该算子把kv最后一维拆成两段分别处理kv[..., :Dv] - RMSNorm - v_out - v_cache kv[..., Dv:] - RoPE - k_out - k_cache2.1 RMSNormRMSNorm 的计算过程为mean_square mean(x * x) rms sqrt(mean_square epsilon) v_out (x / rms) * gamma2.2 RoPERoPE 部分把交错存储的复数拆成 real/imag 两路再用两段 cos/sin 完成复数旋转real rope_x[..., 0::2] imag rope_x[..., 1::2] k_out_first real * cos_first - imag * sin_first k_out_second imag * cos_second real * sin_second2.3 数据生成与 golden 佐证上述语义与 golden 生成脚本 scripts/gen_data.py 中的build_golden完全一致rms_x取kv[..., :dv]rope_x取kv[..., dv:]k_out concat(real, imag) * cos concat(-imag, real) * sin其中part1 concat(real, imag)、part2 concat(-imag, real)恰好对应两段输出k_out_first与k_out_second。脚本还会根据index把k_out/v_out写回k_cache/v_cachecache_idx 0时写回与内核CopyOutK/CopyOutV的 cache 写回逻辑一一对应。当前规格下Dv512、Dk128。RMSNorm 每行需要处理 512 个元素包含平方、归约、开方、除法、乘gamma、BF16/FP32 转换等操作是主要优化热点RoPE 每行处理 128 个元素计算链较短但也适合用 VF 避免中间 UB 临时张量。3. 快速编译与运行在cann-samples仓库根目录执行cmake -S . -B build -DNPU_ARCHdav-3510 cmake --build build --target kv_rms_norm_rope_cache_story构建完成后会生成两个可执行文件build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_membase_full_load build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_regbase_full_load分别运行./build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_membase_full_load ./build/Samples/2_Performance/kv_rms_norm_rope_cache_story/kv_rms_norm_rope_cache_regbase_full_load每个可执行文件会自动生成输入数据并执行 golden 精度校验预期最后输出包含PASS从构建脚本 CMakeLists.txt 可以看到两个.asc文件各自生成独立可执行目标目标名规则为kv_rms_norm_rope_cache_${子目录}_${文件名}并且以-O3优化级别编译、以SOURCE_DIR宏注入源码目录用于运行时定位gen_data.py。运行时的完整 host 流程封装在 sample_common.h 的RunSample中aclInit→ 创建 stream → 调用GenerateDatapython3 gen_data.py --batch 8 --seq 128 --dv 512 --dk 128→ 读取input/*.bin与output/*_golden.bin→aclrtMallocaclrtMemcpy搬入设备 → 以blockNum, nullptr, stream方式启动内核 → 搬回结果后对k_cache、v_cache、k_out、v_out四路输出逐一调用CompareBf16比对全部误差在阈值内才打印PASS。tiling 计算位于 BuildTiling通过platform_ascendc::PlatformAscendCManager获取 AIV core 数totalRows B * N * SblockFactor ceil(totalRows / coreNum)blockNum ceil(totalRows / blockFactor)ubFactor min(SAMPLE_UB_FACTOR, blockFactor)并将epsilon与reciprocal 1/Dv一并写入 tiling 结构体传给内核。4. MemBase 版本以 LocalTensor 为中心的计算链MemBase 版本位于 membase/full_load.asc整体采用Init - Process - ProcessTile的结构。4.1 外层数据流Init完成 GM 地址绑定SetGlobalBuffer、block 行数计算、TQue/TBuf初始化Process先预加载并转换共享gammaBF16DataCopy后整段Cast成 FP32随后按ubFactor循环处理 tileProcessTile串联当前 tile 的搬入、RoPE、RMSNorm 和写回。完整数据流如下Process Load gamma(BF16) - Cast gamma(FP32) for each tile: CopyRopeAndX kv[..., :Dv] - xLocal kv[..., Dv:] - ropeLocal Load cos/sin Rope ropeLocal cos/sin - k outLocal CopyOutK k outLocal - k_out k outLocal - k_cache[index] RmsNorm xLocal gammaFp32 - v outLocal CopyOutV v outLocal - v_out v outLocal - v_cache[index]在 GM/UB 搬运层CopyRopeAndX使用两段DataCopyPad源码以DataCopyExtParams指定每行dk/dv个元素的跨步搬移把后Dk维搬到ropeLocal、前Dv维搬到xLocal从而让 RoPE 和 RMSNorm 在 UB 中独立消费各自的数据段。cos、sin也按 tile 搬入cosSinQueue_与当前ropeLocal对齐。CopyOutK/CopyOutV源码按行逐条DataCopyPad写回 output并通过Mutex::LockPIPE_MTE3保护对indexGm_的读取cacheOffset 0时再按batch * cacheLength cacheOffset计算 cache 地址写回。4.2 RoPE 的 LocalTensor 实现MemBase 的Rope源码仍以LocalTensor为核心先将cos/sinCast成 FP32再通过两次GatherMask从交错 rope 数据中分别抽取 real/imag{1, 1, NUM_EIGHT, 0}步长随后执行两组乘加得到y0/y1Add合并后Cast回 BF16。这个流程直观、易验证但会在cosFp32、sinFp32、y0、y1、realFp32、imagFp32等 UB 临时张量之间产生多次读写。4.3 RMSNorm 的 LocalTensor 实现RMSNorm 是 MemBase 版本的主要热点。RmsNorm源码先将xLocal从 BF16 cast 到 FP32再执行平方、分段Add归约利用{1,1,1,NUM_EIGHT,NUM_SIXTEEN,NUM_SIXTEEN}repeat/stride 参数把 512 元素折半累加、WholeReduceSum得到行和、乘1/Dv、加epsilon、Sqrt、Brcb广播、逐行Div、乘gammaFp32最后Cast回 BF16。该链路依赖多个 UB 临时区xFp32 - square - rowSum - broadcast rms - normalized xFp32 - outLocal这些临时区全部来自wsBuffer_其分配见 Initpipe_.InitBuffer(wsBuffer_, ubFactor_ * (tiling_-dv * NUM_THREE tiling_-dk * NUM_EIGHT) * sizeof(float));即每行需要Dv*3 Dk*8个 float 的临时空间。MemBase 版本的仿真打点图如下图中展示的是一次循环的 Vector 计算内容左边是 ROPE右边是 RmsNorm所有的 Vector 指令都是一个单独的 VF整体看 Vector 指令较为细碎——这正是引入 RegBase 融合计算链的动机。5. 从 MemBase 到 RegBase 的平滑迁移边界从 MemBase 迁移到 RegBase 的整体原则有在一个完整的 VF 功能内应融尽融。理论上 VF 内融合的指令越多UB 和寄存器交互的次数越少性能越好。VF 内尽量少使用 Reduce 类、非对齐访存、Interleave/Deinterleave 类单发指令。VF 内尽量少使用 MembarMembar 会导致 VF 内流水中断使整体 IPC 不高。迁移不是推倒重写。当前 RegBase 版本复用了 MemBase 中已经验证过的外层结构Init中的 GM tensor 绑定方式保持一致BuildTiling、blockFactor、ubFactor、blockNum保持一致CopyRopeAndX中的 GM - UB 搬运保持DataCopyPadCopyOutK/CopyOutV中的 cache 写回和 output 写回保持一致inQueue、cosSinQueue、outQueue的 double bufferBUFFER_NUM2数据流保持一致输出仍走同一套 golden 校验比较k_cache、v_cache、k_out、v_out。真正变化的是RmsNorm和Rope的 compute bodyMemBase: LocalTensor - Cast/Mul/Add/WholeReduceSum/Div/Mul/Cast - LocalTensor RegBase: __ubuf__ pointer - Reg::LoadAlign - RegTensor RegTensor compute chain Reg::StoreAlign - __ubuf__ pointer这种边界划分的好处是外层索引、cache offset、batch 切分和 GM 地址计算不变降低迁移风险原有精度 golden 可以直接复用MemBase 可以作为 RegBase 的对照基线出现错误时可以快速判断问题在 GM/UB 数据流还是 VF compute bodyRegBase 的收益集中在热计算链避免为追求寄存器化引入额外 pipeline 复杂度。6. VF 方案设计RegTensor 寄存器计算RegBase 版本位于 regbase/full_load.asc。普通的RmsNorm/Rope包装函数只做一件事通过LocalTensor.GetPhyAddr()把 UB 地址转成__ubuf__指针再调用独立的__simd_vf__函数RmsNorm 包装、Rope 包装。6.1 RMSNormVF两段式 VF 循环RmsNormVF源码把每一行的 RMSNorm 拆成两个 VF 循环。第一段循环计算平方和for each row: reduceSum 0 for i in Dv chunks: xB16 LoadAlign(DIST_UNPACK_B16) xFp32 Cast(BF16 - FP32) square xFp32 * xFp32 chunk ReduceSum(square) reduceSum chunk第二段循环归一化并乘gammarms sqrt(reduceSum / Dv epsilon) for i in Dv chunks: xB16 LoadAlign(DIST_UNPACK_B16) gammaB16 LoadAlign(DIST_UNPACK_B16) xFp32 Cast(BF16 - FP32) gammaFp32 Cast(BF16 - FP32) norm (xFp32 / rms) * gammaFp32 outB16 Cast(FP32 - BF16) StoreAlign(DIST_PACK_B32)对应到源码chunk 粒度由VL_FP32_SIZE 256 / sizeof(float) 64决定dvLoop_ (dv VL_FP32_SIZE - 1) / VL_FP32_SIZE即Dv512时每行 8 个 chunk。第一段用ReduceSum(reduceSum, reduceSum, fullMask)把 64-lane 的行内和归约到单一值随后Muls乘reciprocal、Adds加epsilon、Sqrt开方、Div(invRms, one, rmsValue)取倒数再用DuplicateHighLowPart::LOWEST把标量广播成invRmsBrc供整行乘除使用第二段按 chunk 同时LoadAlignx 与 gammaMul(x, invRmsBrc)与Mul(norm, gammaFp32)两级乘后Cast回 BF16 并StoreAlign。设计要点BF16 输入先升到 FP32 做平方、归约、除法和乘法保证 RMSNorm 中间计算精度reduceSum、rmsValue、invRms、invRmsBrc、norm等短生命周期中间量全部保留在RegTensor中gamma不再预先 cast 成一整段 FP32 UB 临时张量而是在 VF 中按 chunk load BF16 并 cast 到 FP32通过UpdateMaskfloat(remaining)控制最后一个 chunk 的有效 lane避免把 tail padding 当有效数据。6.2 RopeVF寄存器内复数旋转RopeVF源码的核心思路是直接从 BF16 rope 数据中拆出 real/imag并用两段 cos/sin 完成复数旋转real, imag LoadAlign(DIST_DINTLV_B16, rope) realFp32 Cast Interleave imagFp32 Cast Interleave out_first realFp32 * cos_first - imagFp32 * sin_first out_second imagFp32 * cos_second realFp32 * sin_second StoreAlign(out_first) StoreAlign(out_second)从源码实现看real/imag 的拆分是通过两次带不同RegLayout的 Cast 完成的CAST_B16_TO_B32RegLayout::ZERO取 even lane得到realFp32CAST_B16_TO_B32_ODDRegLayout::ONE取 odd lane得到imagFp32cast trait 定义效果等价于对交错数据做 deinterleave。随后Mul/Sub/Add全部在寄存器内串联两次StoreAlign分别写回outFirst与outSecond。设计要点real/imag、cos/sin、mul/result 都在寄存器内串联无 UB 中间张量RoPE 的pairCount Dk / 2 64mask UpdateMaskfloat(pairCount)固定控制 64 个 FP32 lane输出分成前后两段写回包装函数中以out pairCount作为第二段起始地址与 golden 中concat(real, imag)的布局保持一致。7. 理论分析与实测验证7.1 UB 使用量减少MemBase RMSNorm 需要多个 UB 临时张量xFp32 - square - rowSum - broadcast rms - xFp32 norm - outLocal并且标准 Vector API 之间通常会形成 UB load/store 边界。RegBase 把平方、归约、除法、乘 gamma 等中间状态尽量保存在寄存器中只在 VF 入口 load 输入、在 VF 末尾 store 输出。理论收益来自减少中间结果落 UB减少多段 Vector API 之间的 UB 往返降低 UB 临时 buffer 占用。对比两个版本的 buffer 分配MemBase: gammaQueue: BF16 gamma FP32 gamma wsBuffer_: rows * (Dv * 3 Dk * 8) * sizeof(float) RegBase: gammaQueue: BF16 gamma no wsBuffer_RegBase 的 Init 中已无TBufTPosition::VECCALC wsBuffer_成员这是最直接的 UB 占用优化。7.2 精度路径保持一致样例输入和输出是 BF16但 RMSNorm 的关键中间计算必须使用 FP32BF16 load - FP32 compute - BF16 storeRegBase 版本通过CAST_B16_TO_B32与CAST_B32_TO_B16显式表达转换路径其中回写 cast 使用RoundMode::CAST_RINTround to nearest。实测中MemBase 和 RegBase 均通过同一套 golden 校验MemBase: k_cache/v_cache/k_out/v_out 全部 PASS RegBase: k_cache/v_cache/k_out/v_out 全部 PASS其中 RegBasev_cache/v_out的最大 diff 为0.0078125远低于 sample_common.h 中定义的 BF16 比对阈值6e-2。8. 实际调优路径8.1 保持 full-load 外层数据流当前样例一次 tile 处理ubFactor8行运行时实测打印blockNum 54, blockFactor 19, ubFactor 8这与 BuildTiling 的推导吻合总行数8 * 1 * 128 1024按 AIV core 数向上取整得到每个 block 约 19 行UB 内一次搬入最多 8 行。这个切分对 MemBase 和 RegBase 共同适用因此调优时先不改外层只验证 compute body 的收益和正确性。8.2 减少 UB 使用MemBase 版本需要wsBuffer_存放 RMSNorm 和 RoPE 的 FP32 中间结果RegBase 版本删除了这个大块 VECCALC改为在 VF 中使用RegTensor临时变量。对比MemBase: gammaQueue: BF16 gamma FP32 gamma wsBuffer_: rows * (Dv * 3 Dk * 8) * sizeof(float) RegBase: gammaQueue: BF16 gamma no wsBuffer_这是最直接的 UB 占用优化为 double buffer 深度或更大的ubFactor留出了容量空间。8.3 gamma 处理整段预转换改为 VF 内按需转换MemBase 版本在主循环前把整段gamma从 BF16 cast 成 FP32并保存在 UB 中供后续每行复用RegBase 版本保留 BF16gammaLocal在RmsNormVF的第二段循环中按 chunk load 并 cast。这样做的取舍优点减少一份完整 FP32 gamma UB buffer优点gamma 与 x 的 load/cast 粒度一致寄存器链更紧凑代价每行都会重新 cast gamma chunk。当前样例的主要目标是演示 RegBase VF 迁移和减少 UB 中间张量如果后续 profiler 证明 gamma cast 成本突出可以考虑引入小块 gamma 预处理或更细粒度复用策略但不能破坏 UB 容量和 double buffer 数据流。8.4 切换验证顺序建议按以下顺序调优避免一次修改多个变量后无法归因固定输入数据和 shape只替换 compute body确认精度 PASS对比 MemBase 和 RegBase 的输出 diff先确认 BF16 阈值内一致使用 profiler 观察 UB load/store、Vector 指令、pipeline stall若 UB 压力仍高继续检查 VF 内是否还有可消除的 store/load若算术指令成为瓶颈再检查 reduce 组织、gamma cast 复用和 div/sqrt 链路若搬运成为瓶颈再调整ubFactor、double buffer 深度或 GM copy 形态。当前已完成的功能验证kv_rms_norm_rope_cache_membase_full_load: PASS kv_rms_norm_rope_cache_regbase_full_load: PASS9. RegBase 版本运行效果RegBase 版本位于 regbase/full_load.asc。仿真打点图中左侧是 rope已经融合成 1 个 VF右侧是 RmsNorm融合成 1 个 VF。选择 Rope 中 RVEC_EX 部分的指令从下侧的简易统计可以估算出 IPC 在 1.36 左右选择 RmsNorm 中 RVEC_EX 部分的指令估算出 IPC 在 1.15 左右。RmsNorm 的 IPC 比 Rope 低是因为其中有单发指令而且代码中第一个循环体前后有依赖所以整体 IPC 比 Rope 低需要说明的是上述 IPC 数值来自仿真打点图的简易统计估算用于对比两个 VF 内部的指令密度差异实际性能收益应结合 profiler 数据进一步确认。10. 迁移实践清单从 MemBase 切到 RegBase 时可以按以下清单执行先保留 MemBase 的 host、tiling、queue、CopyIn、CopyOut找出最热的LocalTensor计算链本样例是 RMSNorm其次是 RoPE把计算链拆成独立__simd_vf__函数在普通__aicore__函数中只做LocalTensor.GetPhyAddr()到__ubuf__指针的转换VF 内只使用RegTensor、MaskReg、Reg::Load*、Reg::Store*和Reg::*compute APIBF16 路径显式写清BF16 - FP32 - BF16的 cast trait 和 round modetail 只用MaskReg控制不依赖 UB padding 参与数学每次优化后都运行同一套 golden 校验确认k_cache、v_cache、k_out、v_out全部 PASS。一句话总结RegBase 的平滑迁移路线是外层不动、热链下沉、寄存器串联、同源校验。本样例正是沿着这条路线把 MemBase 中以LocalTensor为中心的计算链切换为 VF 中以RegTensor为中心的寄存器计算链——先通过同一套 golden 校验守住精度底线再借助仿真打点观察 VF 融合效果最终在减少 UB 中间张量的同时让热计算链在寄存器内连续执行。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
