Google 的 QAT Gemma 4 26B-A4B 在单块 TPU v6e 上：KV 缓存提升 15.6 倍，吞吐量达 FP8 的 1.9 倍

📌 One-Sentence Summary

演示了如何将 Google 的 QAT Gemma 4 26B 无损重打包为 W4A16 格式，在 TPU 上运行 vLLM，实现 KV 缓存提升 15.6 倍、吞吐量达 FP8 的 1.9 倍，且精度损失可忽略不计。

📝 Summary

本文提供了一份分步指南，介绍如何在单块 Google Cloud TPU v6e 芯片上使用 vLLM 部署经过量化感知训练（QAT）的 Gemma 4 26B-A4B 模型。针对 bf16 模型的内存限制，作者通过自定义重打包脚本将 Google 的「未量化」QAT 导出转换为 Compressed-Tensors W4A16 检查点，同时保留原始的 4 位网格结构。作者还为 vLLM 的 TPU 后端贡献了新方法以支持这些权重。基准测试结果显示，W4A16 构建仅使用 17.43 GiB HBM，而 RedHat 的 FP8 构建使用 27.99 GiB，KV 缓存容量从 3，456 个 token 提升至 53，888 个，输出吞吐量从每秒 668 个 token 提升至 1，283 个。在包含 3，880 条记录的测试集上，两种量化格式的精度差异在统计上不显著（-0.7 个百分点）。该方案具有可移植性，同一检查点无需补丁即可在 NVIDIA GPU 上运行。

💡 Main Points

QAT 权重的无损重打包实现了 TPU 上高效的 W4A16 推理。

标准量化规则不适用于 QAT 导出，因为它们假设了不同的缩放网格。作者开发了一种方法，从 bf16 张量中恢复精确的步长，从而在不改变底层 4 位级别的情况下转换为 Compressed-Tensors W4A16，保留了训练带来的优势。

W4A16 在单芯片 TPU 上的内存效率和吞吐量显著优于 FP8。

通过降低权重精度同时保持激活值为 bf16，模型占用从约 28 GiB 降至约 17 GiB。这释放了大量 HBM 用于 KV 缓存，容量提升 15.6 倍，token 生成速度相比 FP8 基线提升近 2 倍。

在此场景下，激进量化带来的精度下降在统计上可忽略不计。

分类任务的对比基准测试显示，W4A16 构建的性能与 FP8 构建相差在 0.7 个百分点以内。置信区间跨越零，表明性能差异不显著，验证了以大幅资源节省换取微小精度损失的合理性。

优化后的检查点保持了跨平台兼容性。

虽然 TPU 实现需要对 vLLM 后端进行特定补丁，但生成的 W4A16 检查点可在 NVIDIA GPU（如 L4）上的标准 vLLM 安装中正确加载和运行，展示了 compressed-tensors 格式的可移植性。

💬 Key Quotes

QAT 26B 在单块 v6e 芯片上以 17.43 GiB HBM 运行，拥有 53，888 个 token 的 KV 缓存和每秒 1，283 个输出 token。

RedHat 的 FP8 构建使用 27.99 GiB，容纳 3，456 个 token，每秒服务 668 个。

在包含 3，880 条记录的分类测试集上，两者的差距在 1 个百分点以内。

路径是对 Google「未量化」QAT 导出的无损重打包，以及 vLLM TPU 后端中的一个新方法；同一个重打包检查点在 NVIDIA L4 上的 vLLM 0.30.0 中无需补丁即可加载。

📊 Article Meta

AI Screening： 86

Source： DEV Community: machinelearning

Author： xbill

Category： 人工智能

Language： 英文

Read Time： 10 min

Word Count： 2415

Tags：

AI 与智能应用 , 模型训练与推理 , AI 工程 , KV Cache优化 , AI 硬件与芯片