NeFut Logo NeFut
EN 管理员登录

[AI学术] FuTransformer:在 Hopper 上实现 FP8 预填充和 INT4 解码的 Llama 风格推理

发布于:2026-10-20 22:00 最后更新:2026-10-22 01:04
#AI #LLM #Compiler

Transformer 预填充在每个 token 上复用权重,主要受矩阵乘法和激活数据搬移支配;单 token 解码复用少,受权重和 KV-cache 读取主导。我们推出 FUTRANSFORMER,在 NVIDIA Hopper GPU 上的 CUDA 实现,遵循 CODA 的 GEMM-plus-epilogue 思路,并加入以下优化:

在一块 H100 上,预填充比编译的 PyTorch FP8 快至 1.48 倍,算力达 1.29 PFLOP/s。对于 8B 模型的 gate/up 投影使用 SwiGLU,BF16 路径比 CODA 快 1.13 倍,FP8 路径比 BF16 CODA 快 2.02-2.18 倍。跨 8B 与 70B、批次 1-32,INT4 解码比对应的 PyTorch BF16 与 INT4 实现快 1.3-2.4 倍,分别比编译的 CUDA-graph BF16 快 2.1-3.3 倍,比 PyTorch INT4 投影内核快 1.3-6.3 倍,并且比我们的 FP8 路径快 31-51%。完整的 Llama-3.1-8B 解码在 batch-1 时每步仅需 2.8 ms,而编译的 transformers 实现需 8.1 ms。

在 WikiText-2 上使用 INT4 权重路径解码时,困惑度从 6.24 上升到 7.23(四舍五入权重)或 6.99(GPTQ 同格式),旋转的 INT4 缓存额外增加 0.20。我们的内核在困惑度上与量化参考计算的差异不超过 0.001。

点评

原文链接: https://arxiv.org/abs/2610.23796

[h] 返回首页