FuTransformer: FP8 prefill and INT4 decode for Llama-style inference on Hopper

Wei Wang

LLM MQIA AI Systems and Hardware Other AI Sys & HW

2026年10月16日

Transformer 的预填充（prefill）阶段在不同 token 之间复用权重，其计算主要由矩阵乘法及激活值的数据搬移所主导；而单 token 解码（decode）阶段权重复用有限，主要开销在于权重读取与 KV 缓存（KV-cache）读取。我们提出 FUTRANSFORMER——一种面向 NVIDIA Hopper 架构 GPU 的、基于 Llama 架构的 Transformer 块 CUDA 实现。该实现借鉴了 CODA 提出的“GEMM 加后处理”（GEMM-plus-epilogue）推理优化范式，并进一步引入多项关键改进：直接将 KV 缓存写入目标位置（direct KV-cache placement）、将固定的 RMSNorm 归一化增益（gain）融合进权重中（folding RMSNorm gain into weights）、以 FP8 格式直接输出结果并采用延迟缩放（delayed scaling）机制，以及支持无需填充（padding-free）的预填充操作。在解码阶段，我们采用 INT4 量化权重进行投影运算，并将 KV 缓存也以 INT4 格式存储和复制。所有权重与键（keys）均按 Tensor Core 片段（fragment）布局组织，从而支持直接反量化至矩阵乘加（MMA）运算单元的输入寄存器，无需构建完整的中间 tile。此外，我们引入分块对角（block-diagonal）哈达玛（Hadamard）变换，以缓解因通道离群值（outlier channels）导致的缓存量化误差。在单颗 H100 GPU 上，FUTRANSFORMER 的预填充吞吐量最高可达编译版 PyTorch FP8 实现的 1.48 倍，峰值算力达 1.29 PFLOP/s。对于 8B 规模模型中采用 SwiGLU 激活函数的门控/上投影（gate/up projection）层，我们的 BF16 实现比 CODA 快 1.13 倍；而 FP8 实现则比 CODA 所评估的 BF16 路径快 2.02–2.18 倍。在 8B 与 70B 两种模型规模、批大小（batch size）为 1–32 的全范围测试中，INT4 解码性能相较 PyTorch 在各配置下各自最快的 BF16 或 INT4 实现提升 1.3–2.4 倍：具体而言，相比经编译优化且启用 CUDA Graph 的 BF16 实现，提速 2.1–3.3 倍；相比在同一架构下使用 PyTorch 自带 INT4 投影核（kernel）的实现（该核仅针对单个查询行设计，在批大小为 32 时性能反低于 BF16），提速达 1.3–6.3 倍。同时，INT4 解码路径亦比我们自己的 FP8 路径快 31%–51%。在批大小为 1 的条件下，完整执行 Llama-3.1-8B 模型的一次解码步（per step）仅需 2.8 毫秒，远优于编译优化后的 transformers 库实现（8.1 毫秒）。在 WikiText-2 数据集上全程使用 INT4 权重路径进行逐 token 解码时，困惑度（perplexity）从基线 6.24 分别上升至 7.23（采用最近邻舍入 quantization，round-to-nearest）或 6.99（采用 GPTQ 量化，保持相同格式）；若额外启用旋转（rotated）INT4 KV 缓存，则困惑度再增加 0.20。所有核心算子均能精确匹配量化参考实现，二者困惑度差异控制在 0.001 以内。

Transformer prefill reuses weights across tokens and is dominated by matrix multiplication and activation data movement; single-token decode has limited reuse and is dominated by weight and KV-cache reads. We present FUTRANSFORMER, a CUDA implementation of a Llama-style Transformer block for NVIDIA Hopper GPUs. It applies CODA's GEMM-plus-epilogue approach to inference, adding direct KV-cache placement, folding the fixed RMSNorm gain into weights, writing FP8 output directly with delayed scaling, and supporting prefill without padding. Decode uses INT4 projection weights and an INT4 KV-cache copy. Weights and keys follow Tensor Core fragment layouts for direct dequantization into MMA operands without full intermediate tiles. A block-diagonal Hadamard transform reduces cache quantization error from outlier channels. On one H100, prefill is up to 1.48x faster than compiled PyTorch FP8 and reaches 1.29 PFLOP/s. For gate/up projection with SwiGLU at the 8B shape, our BF16 path is 1.13x faster than CODA. Our FP8 path is 2.02-2.18x faster than the evaluated BF16 CODA path. Across 8B and 70B shapes at batches 1-32, INT4 decode is 1.3-2.4x faster than the faster of PyTorch's BF16 and INT4 blocks at each point: 2.1-3.3x faster than compiled, CUDA-graphed BF16 and 1.3-6.3x faster than the same block using PyTorch's INT4 projection kernel, which targets a single query row and loses to BF16 at batch 32. It is also 31-51% faster than our FP8 path. Complete Llama-3.1-8B decode takes 2.8 ms per step at batch one, versus 8.1 ms for the compiled transformers implementation. Decoding every WikiText-2 token through the INT4 weight path raises perplexity from 6.24 to 7.23 with round-to-nearest weights or 6.99 with GPTQ in the same format; the rotated INT4 cache adds 0.20. The kernels match a quantized reference computation within 0.001 perplexity.

许愿