内存优化

本仓库很多价值都来自于减少不必要的内存搬运。

为什么融合有效

很多 Transformer 子图的瓶颈并不只是算力，还来自中间张量频繁写回再读回 HBM。

融合的核心价值，就是让中间值尽量留在更靠近计算的位置，而不是先写回全局显存再被下一个算子读走。

fused_rmsnorm_rope

最大的收益是：归一化结果不需要先落成单独的全局张量，再去做 RoPE。

实际效果：

• 更少的 HBM 读写，
• 比拆成两个操作更少的 launch 开销，
• 整体更偏向带宽优化问题。

fused_gated_mlp

输入 tile 会在一次 kernel 路径里复用到两条投影中，随后完成激活与相乘，再写出最终结果。

这减少了如果拆成多个 operator 时会出现的中间张量搬运。

FP8 路径

fp8_gemm 还通过使用单字节量化矩阵来降低内存压力。

它同时影响：

• 存储占用，
• 每次矩阵读取所需搬运的字节数。

代价是会引入受控的量化误差，这也是为什么 FP8 指南强调一定要和 baseline 对比。

实践中需要关注的点

• tensor 是否 contiguous，
• 形状是否适合当前 kernel 的分块策略，
• 真实 workload 的瓶颈是否确实位于这些融合片段，而不是位于外围模型代码。

结论

当被减少的内存往返本来就占了显著成本时，本仓库的收益最大。也就是说，它最适合放在那些原本会频繁把中间结果写回 HBM 的热点路径上。