vLLM FP8 KV-Cache 和 Attention 量化更新

vLLM FP8 KV-Cache 和 Attention 量化更新

vLLM 已针对 FP8 KV-cache 和 attention 量化实施了关键修复和优化,能够显著降低长上下文 LLM 服务中的内存占用和逐 token 延迟 (ITL)。通过使用 --kv-cache-dtype fp8 标志,用户可以将 KV-cache 存储减半,并在 FP8 (e4m3) 下运行 attention 计算,这对于受内存限制的解码工作负载特别有效。

技术改进与错误修复

vLLM 的最新更新解决了在 Hopper 和 Blackwell GPU 上进行压力测试时发现的关键准确性和性能退化问题。

用于提高准确性的两级累加 (Two-Level Accumulation)

为了解决 Hopper GPU 上严重的准确性退化问题——在 128k needle-in-a-haystack 任务中,准确率从 91% (BF16) 降至 13% (FP8)——vLLM 引入了两级累加策略。这种方法将部分累加结果写入实际的 FP32 寄存器,以减轻当收缩维度达到 100k 或更多时 Tensor Cores 中的精度损失。该修复将准确率恢复到了 89%。

混合 Attention 和层跳过 (Layer Skipping)

对于具有混合 attention 架构的模型(例如 GPT-OSS),其中某些层使用小窗口(例如 128 tokens)的滑动窗口 attention,FP8 量化的开销往往超过了内存收益。vLLM 现在包含了 --kv-cache-dtype-skip-layers sliding_window 标志,允许这些特定层保持在 BF16 模式以提高解码速度。

Kernel 和融合优化

  • Per-Head Scales: Flash Attention 3 (FA3) kernel 现在支持 FP8 量化的 scale 数组,每个 scale 对应一个 KV-head。
  • Query Quantization Fusion: Query 量化已移至 torch.compile 可以融合的 torch 实现中,消除了固定的逐 token 开销。
  • Tiling 配置: 针对 head_dim = 64head_dim = 128 的 prefill tiling 进行了调优,以减少由两级累加引起的寄存器溢出 (register spills)。

性能基准测试

将 KV-cache 从 BF16 量化为 FP8 可以使每个 attention 步骤的内存带宽减半,这直接降低了 ITL 斜率(即延迟随输入长度增加的速率)。

单请求延迟

在 H100 GPU 上使用 Llama-3.1-8B 时,ITL 斜率从 4.37e-05 降至 2.37e-05 ms/token(降低了 54%),使解码盈亏平衡点降至约 7k tokens。对于 gpt-oss-20b,使用 skip-SW 变体可将 ITL 斜率降低至 BF16 基准的 71%。

高负载下的吞吐量

在高负载场景下(并发度 8,约 20k 输入 tokens),FP8 带来了以下收益:

  • Llama-3.1-8B: 输出吞吐量提高 14.9%,总运行时间缩短 13.0%。
  • gpt-oss-20b: 输出吞吐量提高 4.8%(使用 skip-SW 变体)。

特定架构的结果

  • Blackwell (B200): 使用 FlashInfer 后端时,Llama-3.1-8B 的 ITL 斜率降至 BF16 的 54%,盈亏平衡点约为 4k tokens。Blackwell 不需要两级累加,因为 Hopper 中发现的精度问题在 Blackwell 上并不存在。
  • 具有大 Head Dimensions 的 Hopper (H100): 对于 head_dim = 256 的模型(例如 gemma-4-E2B),由于两级累加增加了寄存器压力,prefill 性能 (TTFT) 比 BF16 慢(在长上下文下慢约 1.6 倍)。

准确性与验证

评估是使用未校准的量化 scale (scale = 1.0) 进行的,以建立性能下限。

推理与长上下文任务

  • 推理: 在 Qwen3-30B-A3B-Thinking-2507 和 Qwen3.5-27B 上,FP8 KV-cache 和 attention 量化导致的准确性损失微乎其微,通常在 0 到 2 个点之间。
  • 长上下文 (MRCR): Llama-3.3-70B-Instruct 恢复了基准 AUC@128k 的 97-98%。Qwen3.5-27B 完全恢复了聚合 AUC@1M 指标,证明了即使在极端上下文长度下也具有稳定性。

何时使用校准

虽然未校准的 FP8 通常足够,但某些模型会出现系统性退化。例如,使用 FlashMLA 后端的 Kimi-K2.5 在不同序列长度下表现出持续的准确性下降。在这种情况下,vLLM 支持通过 llm-compressor 或逐 attention-head 量化 scale 进行 scale 校准。

使用建议摘要

场景 建议
解码密集型、受内存限制的工作负载 使用 --kv-cache-dtype fp8
混合 attention 模型 使用 --kv-cache-dtype fp8 --kv-cache-dtype-skip-layers sliding_window
短上下文 (< 7k tokens) 坚持使用 BF16 以避免 FP8 带来的微小常数开销
head_dim = 256 且预填充优先 坚持使用 BF16 或禁用两级累加(需要进行准确性验证)
持续的准确性损失 (< 95%) 使用 llm-compressor 进行校准

Sources