FuTransformer：Hopper 上 Llama
来自 arXiv · 期号 2026-10-17
摘要
Transformer 的 prefill 阶段在 token 间复用权重，瓶颈是矩阵乘法和激活数据搬运；单 token decode 几乎没有复用，瓶颈是权重和 KV-cache 读取。 FuTransformer 是一套针对 Hopper GPU 的 Llama 风格 Transformer 块的 CUDA 实现。它把 CODA 的 GEMM-plus-epilogue 思路用到推理上：KV-cache 直接放好、RMSNorm 增益折叠进权重、FP8 输出用延迟缩放直接写、prefill 不用 padding。Decode 用 INT4 投影权重和 INT4 KV-cache 副本，权重和 key 按 Tensor Core 分片布局排好，反量化后直接进 MMA 操作数，不用拼完整中间块。一个块对角 Hadamard 变换减少离群通道带来的缓存量化误差。 在一张 H100 上，prefill 比编译后的 PyTorch FP8 快最多 1.48 倍，达到 1.29 PFLOP/s。8B 和 70B 规模、batch 1-32 下，INT4 decode 比 PyTorch 的 BF16 和 INT4 块中更快的那一个快 1.3-2.4 倍。完整跑一遍 Llama-3.1-8B decode，每步 2.8 ms，而编译版 transformers 实现是 8.1 ms。INT4 权重路径在 WikiText-2 上的困惑度从 6.24 升到 7.23（四舍五入）或 6.99（GPTQ）。
深读
导读 这篇论文介绍了一套名为 FuTransformer 的 CUDA 实现，专门针对 NVIDIA Hopper GPU 优化 Llama 风格 Transformer 块的推理性能。它通过将 CODA 的 GEMM-plus-epilogue 思路应用于推理，并引入 INT4 权重和 KV-cache 量化，在 prefill 和 decode 阶段均取得了显著的加速。读完可以了解当前 GPU 推理优化的前沿方向，以及量化与算子融合如何在实际硬件上落地。 关键信息 - Prefill 阶段在 H100 上最高可达 PyTorch FP8 编译版本的 1.48 倍加速，算力达到 1.29 PFLOP/s。 - 针对 8B 和 70B 模型规模、batch 1–32 的 INT4 decode 速度，比 PyTorch 的 BF16 和 INT4 块中较快者快 1.3–2.4 倍。 - 完整的 Llama-3.1-8B decode 在 batch 1 时每步仅需 2.8 ms，而 PyTorch 编译版本需要 8.1 ms。 - 使用 round-to-nearest 权重的 INT4 路径在 WikiText-2 上困惑度从 6.24 升至 7.23；使用 GPTQ 格式则为 6.99；旋转后的 INT4 cache 额外增加 0.20。 逐节讲解 问题背景与设计思路 Transformer 推理的两个阶段瓶颈不同：prefill 阶段在多个 token 间复用权重，主要受矩阵乘法和激活数据搬运限制；单 token decode 阶段几乎没有复用，瓶颈是权重和 KV-cache 的读取。FuTransformer 将 CODA 的 GEMM-plus-epilogue 方法（即把归一化、激活函数等后处理步骤融合进矩阵乘法的 epilogue 中）应用到推理场景，并加入多项针对 Hopper GPU 的优化。 核心优化技术 具体优化包括：直接放置 KV-cache 到目标位置、将固定的 RMSNorm 增益折叠进权重、使用延迟缩放直接写出 FP8 输出、以及支持无 padding 的 prefill。Decode 阶段使用 INT4 投影权重和 INT4 KV-cache 副本，权重和 key 按 Tensor Core 分片布局排列，反量化后直接作为 MMA 操作数，无需拼出完整的中间块。此外，一个块对角 Hadamard 变换被用来减少离群通道带来的缓存量化误差。 性能评测与精度验证 在单张 H100 上，FuTransformer 的 FP8 prefill 路径比 PyTorch 编译的 FP8 快最多 1.48 倍，达到 1.29 PFLOP/s。在 8B 模型规模的 gate/up 投影（SwiGLU）上，BF16 路径比 CODA 快 1.1…