【训练营】BF16 Elementwise Kernel 优化

[WjhisTy]

一、 问题描述

在目前的 InfiniTrain 框架中，BF16（nv_bfloat16）类型的 Elementwise 算子（如 UnaryForwardKernel 等）默认采用了最基础的标量访存模式。
由于 BF16 单个元素仅占 2 个字节，当 GPU 线程逐个读取数据时，完全无法有效利用现代 GPU 128-bit（16 字节）的显存事务宽度，导致极大的带宽浪费。这种 访存瓶颈 问题在进行大规模张量运算时，会显著拖慢底层算子的执行效率。

二、 当前（优化前）的性能记录和分析

• 环境：NVIDIA 服务器
• 端到端基准测试：在未优化的原生代码下运行 LLaMA-3.2 1B 模型，受限于公共节点的显存带宽竞争和底层标量读取的低效，端到端吞吐量曾跌至约 625 tokens/sec（单步耗时高达 410 ms）。
• 理论分析：优化前，每个线程发起一次 load/store 指令仅搬运 2 bytes 数据，导致全局显存带宽利用率极低。

三、 优化方法与开发历程记录

1. 优化历程与踩坑记录
在性能剖析阶段，我经历了极为曲折但收获颇丰的排查过程：

• 尝试使用 nsys 和 ncu 时，遭遇了云容器共享存储空间爆满（/tmp 目录 No space left on device）的问题，通过重定向 TMPDIR 成功绕过了编译器的 IO 限制。
• 之后在开放ncu的时候，使用 ncu -k 精准狙击抓取到了目标算子的底层硬件计数器数据。

2. 核心优化方案：128-bit 向量化访存
为了解决带宽浪费，我重构了 infini_train/src/kernels/cuda/elementwise.cu 中的核心逻辑：

• 构建向量化数据结构：引入了 AlignedVector<T, kVecSize>，通过 alignas(sizeof(T) * VecSize) 强制 16 字节对齐。针对 BF16，kVecSize 设为 8，使得一个线程能一次性吞吐 8 个元素（16 bytes）。
• 安全对齐校验：增加了 IsAligned 运行时指针校验，确保仅在输入/输出指针严格对齐时开启向量化模式，否则安全回退至标量处理。
• 循环展开批量处理：在主循环中使用 #pragma unroll，促使编译器在底层发射 LDG.E.128 和 STG.E.128 宽指令。
• 严谨的尾部处理：精确计算 tail_start，将无法被 8 整除的剩余元素交由标量安全处理，保证精度下的结果绝对正确。

四、 最终优化后的性能记录和分析

• 端到端性能表现：应用向量化访存后，LLaMA-3.2 1B 的端到端吞吐量稳定在 ~1410 tokens/sec（单步耗时降至 ~181 ms）。
• 分析：根据阿姆达尔定律，虽然 Elementwise 算子在总执行时间中占比不如 MatMul 大，但在微观层面，该核函数自身的显存吞吐效率得到了结构性的倍数提升，显著降低了该算子自身的耗时。

五、 NCU 使用和分析

为深度剖析瓶颈，我使用 ncu -o bf16_unary_report -f -k "UnaryForwardKernel" -c 2 成功抓取了该算子的 Profiling 报告。
报告数据发现：

• 目标调用（如 MulScalarForward / PowForward）的执行时间极短（Duration 分别为 3.46 μs 和 4.26 μs）。
• 报告显示 Memory Throughput 和 SM Throughput 均不足 1%。
• Launch Statistics 指出，该 Kernel 启动的 Grid Size 仅为 1（Threads = 256）。
  深度分析：
  在配备 108 个 SM 的顶级 GPU 上，处理极小尺寸（≤256 元素）时会引发严重的“尾部效应”。NCU 警告指出 "only 0.0 full waves across all SMs"，证明当前框架中部分 Elementwise 算子处于 Kernel Launch Overhead Bound（内核启动延迟瓶颈），这也为下一步的优化指明了方向。

六、 未来可继续提升的地方

基于 NCU 的诊断分析，未来的优化方向应聚焦于：

1. 算子融合 (Kernel Fusion)：鉴于部分 Elementwise 算子存在启动开销远大于计算时间的现象，应将连续的操作融合成单一 Kernel，以彻底消除 Launch Overhead 和重复读写全局显存的开销。为此，我在 elementwise.cu 底部硬编码实现了一个 FusedAddSigmoidKernel 作为概念验证 (PoC)。
2. 将 128-bit 向量化访存全面覆盖至 BinaryForwardKernel 及反向传播逻辑中。