[KongchangAI]
· 2 min read· 1,427 words

FP8 Prefill + INT4 Decode: New Optimization for Llama Inference on Hopper

FP8 Prefill + INT4 Decode: New Optimization for Llama Inference on Hopper

FuTransformer designs separate FP8 and INT4 quantization paths for Prefill and Decode on Hopper, balancing compute and bandwidth bottlenecks.

FuTransformer is a CUDA inference implementation for Llama-style Transformers on NVIDIA Hopper GPUs. It recognizes the asymmetric computational characteristics of Prefill and Decode and adopts different quantization strategies: FP8 prefill with GEMM-epilogue fusion, RMSNorm gain folding, delayed scaling, and direct KV-cache writes reaches 1.29 PFLOP/s on a single H100, up to 1.48x faster than compiled PyTorch FP8; INT4 decode uses INT4 weights and KV-cache with dequantization results fed directly into Tensor Core MMA operands and a block-diagonal Hadamard transform to mitigate outlier-channel error, achieving 1.3-2.4x speedup over the fastest PyTorch scheme at 8B/70B shapes. Accuracy-wise, INT4 weights raise perplexity from 6.24 to 7.23 (RTN) or 6.99 (GPTQ), with the rotated cache adding another 0.20. The authors stress that quantization strategy should follow the phase bottleneck rather than a unified format, though results are limited to single-card H100 experiments.

Why do Prefill and Decode need different quantization strategies?

The two phases of Transformer inference have fundamentally different computational characteristics. In the Prefill phase, the entire prompt sequence is processed at once, so weights can be reused across many tokens, and the performance bottleneck is mainly matrix multiplication and data movement of activations. In contrast, the single-token Decode phase has almost no weight reuse, and the bottleneck shifts to weight reads and KV-cache reads.

This asymmetry means that covering both phases with a single quantization scheme often forces a compromise on one side. FuTransformer is exactly a CUDA implementation proposed to address this contradiction. It designs separate FP8 prefill paths and INT4 decode paths for Llama-style Transformer blocks on NVIDIA Hopper GPUs.

FuTransformer architecture overview

KV-cache is the mechanism in autoregressive decoding that caches the Key and Value tensors of historical tokens, avoiding recomputation of attention over the entire prefix at each step. Its size grows linearly with batch size, sequence length, number of layers, and hidden dimension, and in long-context scenarios it can account for the main portion of GPU memory. In the Decode phase, every generated token requires reading the entire KV-cache, so its bandwidth overhead is proportional to sequence length. This is exactly why the article lists KV-cache reads as a Decode bottleneck and designs a separate INT4 cache copy for it.

Tensor Core fragment layout refers to the fact that matrix multiply-add instructions on NVIDIA GPUs (such as MMA) have specific requirements for how operands are arranged in registers and shared memory. Data must be placed according to a specific thread-element correspondence. If the dequantization result does not satisfy this layout, it must first be written to shared memory and then rearranged, introducing extra movement. The article emphasizes that weights and keys are stored in fragment layout precisely so that the dequantization output can directly serve as MMA operands, eliminating this rearrangement and write-back.

FP8 Prefill: fusing GEMM and epilogue all the way

FuTransformer follows CODA's GEMM-plus-epilogue approach, but applies it to inference scenarios and makes several targeted modifications: writing results directly to KV-cache locations, folding fixed RMSNorm gains into weights, using delayed scaling to directly output FP8, and supporting prefill without padding.

The common goal of these changes is to reduce round-trip movement of intermediate data. Folding RMSNorm gains into weights means one fewer elementwise multiply-add during inference; delayed scaling eliminates an additional quantization step for FP8 output. On a single H100, this prefill path is up to 1.48x faster than compiled PyTorch FP8, reaching 1.29 PFLOP/s.

On the gate/up projection with SwiGLU at 8B shapes, its BF16 path is 1.13x faster than CODA, while the FP8 path is 2.02 to 2.18x faster than the evaluated BF16 CODA path.

INT4 Decode: letting dequantization feed directly into MMA operands

The Decode path uses INT4 projection weights and an INT4 KV-cache copy. The key design is that both weights and keys are stored according to the Tensor Core fragment layout, so after dequantization they can be used directly as MMA operands without generating a full intermediate tile. This eliminates an extra memory write common in INT4 schemes.

Another detail is precision compensation for outlier channels. The paper uses a block-diagonal Hadamard transform to reduce cache quantization error and mitigate the interference of outlier channels with INT4 quantization.

Across 8B and 70B shapes and batch sizes 1 to 32, INT4 decoding is 1.3 to 2.4x faster than the faster of PyTorch's BF16 and INT4 blocks; 2.1 to 3.3x faster than compiled and CUDA-graphed BF16; and 1.3 to 6.3x faster than the same block using PyTorch's INT4 projection kernel, which is designed for single-query rows and even loses to BF16 at batch 32. Compared with its own FP8 path, INT4 decoding is also 31% to 51% faster.

Full Llama-3.1-8B takes 2.8 milliseconds per decode step at batch 1, compared with 8.1 milliseconds for the compiled transformers implementation.

Hadamard transform is an orthogonal linear transform that evenly distributes the energy of each component in a vector. During quantization, if some channels have outliers with magnitudes much larger than others, the limited dynamic range of INT4 is occupied by these channels, increasing quantization error for the remaining channels. The block-diagonal Hadamard transform applies rotations to weights or activations in blocks, spreading outlier energy across dimensions within the block and thereby reducing overall quantization error. The article uses it for the KV-cache at the cost of about a 0.20 increase in perplexity, in exchange for lower cache quantization error; the transform must be applied correspondingly before and after inference, so it also introduces a small amount of extra computation.

Accuracy cost: perplexity from 6.24 to 7.23

Quantization acceleration is not free. Decoding every token of WikiText-2 through the INT4 weight path raises perplexity from 6.24 to 7.23 (round-to-nearest weights) or 6.99 (GPTQ in the same format); the rotated INT4 cache adds another 0.20.

These numbers provide a clear trade-off reference: GPTQ is clearly better than naive round-to-nearest in this format, and the extra 0.20 perplexity from the Hadamard rotation buys lower cache quantization error. The perplexity difference between the kernel implementation and the quantization reference computation is within 0.001, indicating that the acceleration does not introduce additional numerical deviation.

Perplexity is a common metric for language model evaluation, defined as the exponential of the average negative log-likelihood per token that the model assigns to test text. Lower values indicate more accurate predictions. It is affected by tokenization, test set, and context length, so it is only meaningful to compare horizontally under the same setting. The increase from 6.24 to 7.23 in the article is about 16%, corresponding to the quantization noise introduced by the INT4 weight path; GPTQ's 6.99 shows that using calibration data to optimize rounding can recover part of the loss. The difference of less than 0.001 between the kernel and the reference implementation indicates that the accelerated implementation itself does not introduce additional numerical deviation.

GPTQ is a post-training quantization method that uses a small amount of calibration data to minimize quantization error layer by layer. It uses approximate second-order information (Hessian) to decide the weight rounding direction, minimizing the deviation between quantized output and full-precision output. Compared with naive round-to-nearest (direct nearest rounding), GPTQ typically significantly reduces accuracy loss at very low bit widths (such as INT4), which is exactly why GPTQ's perplexity (6.99) is clearly better than RTN (7.23) in the same INT4 format in the article. Its cost is a one-time offline calibration process, and different models need to be recalibrated.

Practical significance

The value of FuTransformer lies not in a breakthrough in any single metric, but in demonstrating the engineering feasibility of phase-specific quantization on Hopper: prefill uses FP8 to saturate compute, decode uses INT4 to reduce bandwidth pressure, and both share the same fragment-layout dequantization logic. For teams deploying Llama-style models, this path suggests a direction: quantization strategy should follow the bottleneck of the inference phase rather than pursuing a unified format across the entire pipeline.

It should be noted that the above acceleration data all come from a single-card H100 experimental environment, and performance under multi-GPU parallelism, longer contexts, and different batch distributions still needs further verification.

One-line Summary

FuTransformer designs separate FP8 and INT4 quantization paths for Prefill and Decode on Hopper, balancing compute and bandwidth bottlenecks.

Paragraph Summary

FuTransformer is a CUDA inference implementation for Llama-style Transformers on NVIDIA Hopper GPUs. Its core idea is to recognize the asymmetric computational characteristics of Prefill and Decode and adopt different quantization strategies for each. In Prefill, weight reuse is abundant and the bottleneck is matrix multiplication, so it uses an FP8 path, reducing data movement through GEMM-epilogue fusion, RMSNorm gain folding, delayed scaling, and direct writes to KV-cache, reaching 1.29 PFLOP/s on a single H100 and up to 1.48x faster than compiled PyTorch FP8. In Decode, there is almost no weight reuse and the bottleneck shifts to weight and KV-cache reads, so it uses INT4 weights and an INT4 KV-cache, and lets dequantization results serve directly as Tensor Core MMA operands, with a block-diagonal Hadamard transform to mitigate outlier-channel error. At 8B/70B shapes it is 1.3 to 2.4x faster than the fastest PyTorch scheme. In terms of accuracy, INT4 weights raise perplexity from 6.24 to 7.23 (RTN) or 6.99 (GPTQ), and the rotated cache adds another 0.20. The authors emphasize that quantization strategy should follow the bottleneck of the inference phase rather than pursuing a unified format across the entire pipeline, but the conclusions are currently limited to a single-card H100 experimental environment.

#FP8 prefill#INT4 decode#Hopper GPU#Llama inference#quantization#FuTransformer
Share:

Related articles