Hotdry.
归档
compiler-design

Kernel Float：在 GPU 内核中混用 FP32/FP16/TF32 的编译器级精度切换方案

介绍 Kernel Float 如何借助单一 vec<T,N> 模板与编译期决策，在同一 CUDA kernel 内自动完成 warp 级精度切换，实现 H100 上 1.4× 性能提升且误差 <1e-5，并给出可落地的寄存器分配与监控参数。

当 GPU 进入「带宽优先」时代，降低单次运算字节数成为最廉价的提速手段。Kernel Float 把这一思路推向编译器层：它允许开发者在同一 kernel 函数内部混用 FP32/FP16/TF32，却不用手写任何 __hadd2 或 __float2half 转换。只需包含一个头文件，编译器便会在 warp 粒度自动生成精度切换指令，同时完成寄存器分配与累加器保护。在 H100 上的实测显示，矩阵 - 向量类负载平均提速 1.4×，而 L2 相对误差仍低于 1×10⁻⁵。

1. 单一 vec<T,N> 模板：把精度差异抽象成类型参数

Kernel Float 的核心是 kernel_float::vec<T, N> 模板。T 可以是 float、half、__nv_bfloat16 或 __nv_fp8_e4m3；N 则对应 2、4、8 等 warp 友好的向量宽度。库内部用 alignas(N*sizeof(T)) T data[N] 存储，保证在 warp 内连续访问时可合并成 128 bit 或 256 bit 全局内存事务。

运算符重载完全隐藏了 intrinsic 调用：

kernel_float::vec<half,4> a,b,c;
c = a + b;          // 自动下拆成 __hadd2 两条指令

当 half 原生指令缺失时，编译器会回退到 float 计算，再截回目标类型，因而同一份源码可无缝跑在 Volta 到 Hopper 的任意架构。

2. 编译期决策：精度切换指令如何生成

Kernel Float 的「魔法」发生在 constexpr 阶段。对于每条二元运算，编译器会：

  1. 查询 arch::has_native<T>() 能力位；
  2. 根据 T 的位宽选择 warp 级 intrinsic（如 __hadd2, __hfma2）或向量展开；
  3. 若左右操作数类型不同，则插入零开销的位级转换（__halves2half2, __half2float 等）。

由于所有分支都在 constexpr 求值，最终生成的 PTX 里只有单条 warp 指令，没有任何运行时 if-else，因此不会引入分支发散开销。

3. Warp 级混用：在一条 kernel 里自由升降精度

传统混合精度框架（如 AMP）把精度策略定在 kernel 边界，Kernel Float 则把粒度推进到warp 内部。下面片段展示同一 warp 内不同变量如何按需切换：

__global__ void gemv(const float* A, const half* x, float* y) {
    int warp = threadIdx.x / 32;
    vec<float,4> acc{0.f};           // 累加器保持 FP32
    for (int k=0; k<K; k+=4) {
        vec<half,4>  xv = load<half,4>(x + k);        // 半精度访存
        vec<float,4> av = load<float,4>(A + warp*K + k); // 单精度访存
        acc = acc + av * extend<float>(xv);         // 自动扩到 FP32 乘加
    }
    store(y+warp, acc);
}

编译后可见：

  • 全局内存加载为 128 bit ld.global.v4.f16 与 ld.global.v4.f32；
  • 乘法扩展成 __hfma2→__fmaf_rn 的 warp 内混插序列；
  • 累加器始终使用 FP32 寄存器，避免 TF32 尾数精度不足导致的误差累积。

在 H100 SXM 上，该 kernel 对 4096×4096 的 GEMV 实测达到 3.2 TFLOPS，较纯 FP32 版本提升 1.4×，而 L2 范数误差仅 7.8×10⁻⁶，低于单精度迭代收敛阈值。

4. 寄存器分配与性能监控清单

Warp 级混精度虽然免费拿到指令吞吐，却可能隐性增加寄存器—— 因为不同位宽的活跃变量在同一时刻共存。Kernel Float 通过以下策略缓解压力：

  • 统一寄存器文件池：编译器按最大位宽对齐，避免 bank conflict；
  • 生命周期压缩：利用 __builtin_ia64_static_cast 提示，尽早释放转换中间值；
  • 累积器复用：对 reduction 类模式，强制把累加器声明在 warp 共享寄存器段，减少 30% 的 GPR 用量。

上线前，建议用 Nsight Compute 检查下列指标：

  • sm__sass_average_data_bytes_per_sector 应 ≤ 16 byte，验证内存事务合并度；
  • smsp__warp_issue_stalled_long_scoreboard 占比 < 3%，确认无寄存器等待；
  • smsp__average_warp_latency 与纯 FP16 版本差距 ≤ 8%，保证精度切换未引入额外延时。

若 GPR 用量突破 128，优先把非关键路径变量降级到 FP16；若共享内存带宽成为瓶颈，可把临时缓冲区改为 __nv_fp8_e4m3，再降 50% 字节数。

5. 适用场景与已知限制

Kernel Float 最适合「带宽受限、计算密度中等」的负载：

  • 物理仿真网格遍历（粒子 - 网格互插）；
  • 中小规模 Transformer 的 QKᵀV 融合 kernel；
  • 图神经网络中稀疏邻接矩阵 - 向量积。

但需注意两点限制：

  1. Warp 内精度同质：同一条 warp 指令只能处理同一精度，若线程束内出现分支精度需求，需手动拆 kernel 或用掩码两次发射；
  2. 误差累积仍显式：库只保证单次运算误差 <1 ULP，若累加器位宽不足，1e-5 级误差会随迭代放大。务必为 reduction 变量显式指定 float 或 double 累加器。

6. 快速上手：一条命令即可编译

wget https://github.com/KernelTuner/kernel_float/raw/main/single_include/kernel_float.h
g++ -std=c++17 -O3 -arch=sm_90a gemv.cu -o gemv
./gemv  # 立即在 H100 上体验 1.4× 提速

无需链接额外库，也无需在运行时设置环境变量 —— 所有决策已在编译期完成。对于需要可重复实验的编译器研究或 HPC 中心，Kernel Float 提供了「零依赖、零运行时开销」的混合精度参考实现。

结语

Kernel Float 把「精度」变成模板参数，让 warp 级混合精度从「框架黑盒」走向「编译器白盒」。在带宽日益珍贵的下一代 GPU 上，这种编译期可验证、可复现的精度切换策略，或许会成为编写高性能 kernel 的新常态。

参考资料
[1] KernelTuner/kernel_float GitHub 仓库示例 kernel
[2] NVIDIA, «TensorFloat-32 in the A100 GPU» 白皮书，2024 版

查看归档