[Kernel] Expand FP8 support to Ampere GPUs using FP8 Marlin

[mgoin]

This work expands FP8 support in vLLM from GPUs with hardware FP8 support (Hopper and Ada Lovelace) to GPUs without native support (currently Ampere) by introducing FP8 Marlin - a fast fused dequantization kernel for FP8 to BF16/FP16 conversion.

Key features:

• Enables FP8 quantization on a wider range of GPUs (SM 8.0 and 8.7, Ampere)
• Improves performance up to 2x in memory-bound scenarios
• Maintains accuracy comparable to FP16 baselines
• Reduces weight memory usage by 2x, allowing larger batches
• Simple to use - just specify quantization="fp8" at runtime or use pre-quantized FP8 checkpoints

Implementation details:

• Based on existing 8-bit integer support in GPTQ Marlin kernel
• Packs FP8 weights into int32 doublewords (GPTQ format) and then permutes weights into Marlin format
• Efficient 4xFP8 to 4xFP16/BF16 dequantization using bit arithmetic and SIMT operations

End-to-end performance and accuracy results:



Individual layer sweeps:

As shown in the graphs, FP8 Marlin can provide significant speedups with minimal accuracy impact. Performance gains are higher on GPUs with less memory bandwidth (A10, RTX 3090) and for larger models.

Notes:

• This weight-only approach differs slightly from the existing W8A8 FP8 quantization, offering higher accuracy because the activations have no need to be quantized
• Currently expanding scales to be channelwise; future work will revert to per-tensor scales
• This does not include support for MoE models.

Testing:

• Tested on H100, A100, and A10 GPUs

This enhancement enables more users to benefit from FP8 quantization without hardware restrictions, improving vLLM's performance and efficiency across a broader range of setups!

[robertgshaw2-neuralmagic]

This is an awesome feature!

[comaniac]

Overall LGTM. Thanks!

[fxmarty]

@mgoin awesome feature! I suppose that the perf benchmark was run with cuda graph enabled? Out of curiosity, did you run it without cuda graph?

As this kernel has been integrated in TGI as well, it appears having CUDA graph enabled is rather critical so as to get speedups in the decoding (which I don't really explain to myself - but haven't profiled). In the prefill, as cuda graphs are never used for long enough seqlens, I do get a slight slowdown.

I did not benchmark on vllm, but I suppose the trend is similar. Probably depends on the gpu/tp config/model as well.

related: huggingface/optimum-quanto#241 (comment)

[mgoin]

Glad you're enjoying it @fxmarty. Thanks for sharing your analysis. My end-to-end benchmarks were all done with cuda graphs enabled as this is the default in vLLM. Note that it is expected to see a slight slow-down at prefill (M>256), we trade this off to see the improvements at decode.

I'm curious, have you seen the same difference for marlin int8 or int4? Aside from this, I think there could be additional tuning for A100 problem shapes.

[HPC4AI]

Hello, I noticed that you used the dequant_8bit function to dequantize FP8 data to FP16 data, but I'm not clear on the underlying principle. Could you please also provide the code for quantizing FP16 to FP8? Thanks.

[AllenDou]

Hello, I noticed that you used the dequant_8bit function to dequantize FP8 data to FP16 data, but I'm not clear on the underlying principle. Could you please also provide the code for quantizing FP16 to FP8? Thanks.

vllm/csrc/quantization/fp8/fp8_marlin.cu

Line 164 in 2b0879b

__device__ inline typename ScalarType<nv_bfloat16>::FragB

https://github.com/IST-DASLab/marlin/blob/1f25790bdd49fba53106164a24666dade68d7c90/marlin/marlin_cuda_kernel.cu#L131

https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h