DeepSpeed FP6：以 FP6 量化与 TC-FPx 全栈 GPU Kernel 为核心的 LLM 推理服务加速

DeepSpeed FP6：以 FP6 量化与 TC-FPx 全栈 GPU Kernel 为核心的 LLM 推理服务加速

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

DeepSpeed-FP6 是 DeepSpeed 推理生态中面向大型语言模型（LLMs）的权重量化服务方案：它以 6 位浮点（FP6）作为权重精度，通过第一个具备 Tensor Core 支持的不规则位宽浮点权重 GPU 系统方案 TC-FPx 打破硬件位宽限制，并集成到 DeepSpeed-FastGen/MII 的 v2 推理引擎中实现“运行时即时量化 + 低显存高吞吐”的 LLM 部署。读完本篇，你将理解 FP6 相对 INT4/FP16 的取舍逻辑、TC-FPx 的 4+2 位平面拆分与位级打包实现、wf6af16量化模式在仓库源码中的完整调用链（配置解析 → 权重预处理 → CUDA kernel → 门控/激活融合），以及 LLaMA-2-70B 级别模型的实际使用方式与适用边界。本文主体内容基于官方中文博客 README-Chinese.md，并结合当前仓库的源码实现进行纵深扩充。

1. 为什么选择 6 位浮点（FP6）

LLMs 领域迅猛发展，模型量化是提升推理服务性能的关键技术之一，目标是在提高计算效率、压缩存储的同时保持模型质量。DeepSpeed-FP6 的选型逻辑建立在两个研究结论之上：

INT4 的挑战。在 ZeroQuant(4+2)（arXiv:2312.08583）的研究中，团队探索了 INT4 量化技术（如 GPTQ 算法）在大语言模型中的表现。这些技术虽然能显著减小模型体积与参数存储量，但由于过拟合问题，在更一般的许多任务中往往表现不佳——包括代码生成和摘要等更多生成任务。因此在保持 INT4 级别内存效率的同时，需要一种精度上更稳健的方案。

FP6 的突破。对不同量化方法的探索最终指向了 FP6 精度标准。尽管 FP6 这种“非 2 的幂”的不规则位宽在当前 AI 硬件上缺乏高效支持（这一点由 TC-FPx 解决），该格式在各种任务上的性能和灵活性均表现出色：使用 FP6 量化的 StarCoder-15B 在代码生成上达到了与 FP16 模型相当的结果，较小的模型（如 BART-406M）在摘要任务上达到了标准 FP16 性能水平。为提高 FP6 在主流 AI 硬件上的执行效率，团队进一步提出了 4+2 位平面的 FP6 GPU kernel 方案，使 FP6 成为兼顾效率与有效性的量化途径。

如果要引用 DeepSpeed-FP6，官方建议引用以下两篇 arXiv 报告——ZeroQuant(4+2) 和 FP6-LLM：

@article{wu2023zeroquant, title={Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks}, author={Wu, Xiaoxia and Xia, Haojun and Youn, Stephen and Zheng, Zhen and Chen, Shiyang and Bakhtiari, Arash and Wyatt, Michael and Aminabadi, Reza Yazdani and He, Yuxiong and Ruwase, Olatunji and Song, Leon and others}, journal={arXiv preprint arXiv:2312.08583}, year={2023} } @article{xia2024fp6, title={FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design}, author={Xia, Haojun and Zheng, Zhen and Wu, Xiaoxia and Chen, Shiyang and Yao, Zhewei and Youn, Stephen and Bakhtiari, Arash and Wyatt, Michael and Zhuang, Donglin and Zhou, Zhongzhu and others}, journal={arXiv preprint arXiv:2401.14112}, year={2024} }

2. FP6 的系统支持：TC-FPx 全栈 GPU Kernel

2.1 设计思想：SIMT 反量化 + Tensor Core 计算

FP6 量化的核心难点在于缺乏针对这种不规则位宽的高效 GPU Kernel。FP6-LLM（arXiv:2401.14112）给出了 TC-FPx——第一个具有 Tensor Core 支持的、覆盖 FP6 及其他不规则量化位宽（6 位、5 位、3 位等）浮点权重的全栈 GPU 系统设计方案，用于缓解 LLM 推理期间的“内存墙”问题。TC-FPx 打破了底层 GPU 硬件的限制，使 GPU 可以支持涉及任意位宽模型权重的矩阵乘法计算。其分工是：

• Tensor Cores负责矩阵乘法的密集计算；
• SIMT cores在运行时负责权重反量化，将模型权重反量化为 FP16 类型，再由 Tensor Core 基于 FP16 进行计算。

它包含三项关键创新：

1. 运行前比特层级的数据排布转换：解决权重具有不规则位宽时不友好的内存访问问题，实现 GPU 内存的最优访问；
2. 运行时的高效 SIMT 计算：最小化权重反量化的运行时开销；
3. 全栈的高效流水线设计：对 SIMT 计算、Tensor Core 计算和 GPU 内存访问进行高效调度，最大程度提升性能。

其中“4+2”的含义是：每个 6 位权重被拆分为一个 2 位平面和一个 4 位平面分别打包存储（2+4=6 bit），使存储访问对齐到字节粒度，同时在 kernel 内按位平面重组出完整的 FP6 权重。这一点可以直接在当前仓库的 CUDA 源码中得到印证。

2.2 仓库源码印证：从 FP6 常量到 4+2 位平面打包

当前仓库中该 kernel 的实现位于 DeepSpeed Inference v2（DeepSpeed-FastGen 的 ragged 推理引擎）路径下，关键文件为 linear_kernels.cpp。

FP6 数值格式常量。cast_fp16_fp6函数定义了 FP6 的位结构：指数 3 位、尾数 2 位（exponent_nbits_fp6 = 3、mantissa_nbits_fp6 = 2），值域上界absmax_fp6 = 28、最小非零绝对值absmin_nonzero_fp6 = 0.0625，并且注释明确说明“指数 '111' 被当作普通值而非 NaN/inf 处理，这与 qtorch 的行为一致”——这正是 FP6 与量化库 qtorch 对齐的实现细节（见 linear_kernels.cpp#L38-L106）。该值域常量 28 也出现在仓库另一处 FP 量化器实现 fp_quantize.cpp（q_bits == 6 ? 28.0）中，两处保持一致。

运行前的位级打包。weight_prepacking_fp16_to_fp6将 FP16 权重重排为连续存储的 6 位权重：每 4 个 16 位元素打包成 3 个字节（4×6bit = 3×8bit），即输出体积为M*K*6/8字节（见 linear_kernels.cpp#L116-L134）。随后preprocess_weight调用的weight_matrix_prepacking把连续 6 位位流拆分为两个位平面张量（见 linear_kernels.cpp#L194-L224）：

• weight_2bit：形状M*K*2/8字节（2 位平面）；
• weight_4bit：形状M*K*4/8字节（4 位平面）。

两者合计恰为每权重 6 bit，这就是博客中“4+2”方案的存储形态；代码还特意clone()以切断与原张量的底层内存共享，避免原始权重释放后出现悬空引用。

推理阶段的 kernel 入口。cuda_wf6af16_linear（linear_kernels.cpp#L153-L183）在 CUDA 流上启动fp6_linear_kernel，接收weights_2bit、weights_4bit、scales（形状[M]，每输出通道一个 scale）、hidden_states（FP16 激活）与一个M*N*split_k的 FP32 workspace，完成“反量化（SIMT）+ 位平面重组 + Tensor Core GEMM + split-K 归约”的融合计算。

2.3 Kernel 的封装与 Split-K 调度策略

Python 侧封装在 cuda_linear.py 中：

• 形状约束：out_channels % 256 != 0或in_channels % 64 != 0会直接抛出ValueError，且不支持 batched-matmul（cuda_linear.py#L181-L182）；
• Split-K 查表：由于 decoder 阶段 GEMM 的矩阵形状狭长（token 数少、通道数多），kernel 采用 split-K 提升 SM 利用率。仓库内置了一张在A100-80G 上实测调优的split_k_map：以 64 个 token 为一个区间，按out_channels（覆盖 3072、4096、5120、6144、8192、10240、14336、28672、57344 等常见 LLM 通道数）查表得到 split-K 值；超出 768 token 或查不到的形状退化为split_k = 1并记录告警（cuda_linear.py#L29-L197）。这张表本身就是“TC-FPx 仅支持 Ampere、在 A100 上验证”这一限制在工程层面的具体体现。

2.4 性能结论

• 在 NVIDIA A100 上进行受参数矩阵访存瓶颈限制（decoder 矩阵形状狭长）的 GEMM 时，FP6 kernel 平均处理速度比 FP16 cuBLAS 基准提高2.1 倍；
• 借助 FP6 量化，LLaMA-70B 模型能够在单个 A100 GPU上运行，在 batch 小于 32 的 LLM 推理任务中，归一化吞吐比 FP16 基准高1.69 到 2.65 倍；
• 适用边界：TC-FPx 目前仅支持NVIDIA Ampere GPU，且仅在 A100 上进行了测试和验证。

3. 使用 FP6 服务 LLMs：集成到 DeepSpeed-FastGen/MII

FP6 量化 kernel 已成功集成到 DeepSpeed-FastGen（当前仓库中的 v2 推理引擎），实现了运行时的即时量化。通过统一配置选项，用户即可高效量化和部署大型语言模型，输入可以是 HuggingFace 模型名称或本地 checkpoint 目录。加载 checkpoint 时，系统对每个线性层施加 FP6 round-to-nearest 量化，并对量化权重执行比特层级的数据排布转换；转换后的张量成为推理时使用的模型权重，原始 FP16 权重随即被丢弃以释放显存；推理阶段则由 FP6 kernel 直接利用这些 6 位权重计算。

这一流程在仓库源码中对应清晰的调用链：

1. 配置解析：v2 引擎配置中的QuantizationConfig只定义了一个字段quantization_mode，当前唯一支持的值即'wf6af16'（FP6 权重 + FP16 激活的权重-only 量化），见 config_v2.py#L20-L27；
2. 硬件门禁与模块选择：instantiate_linear在quantization_mode == "wf6af16"时依次检查——必须可用 CUDA、非 ROCm 构建、且torch.cuda.get_device_properties(0).major == 8（即 Ampere 架构），任一不满足即抛出明确错误；随后从DSLinearRegistry实例化quantized_wf6af16_linear模块，见 heuristics.py#L89-L107；
3. 线性层实现：QuantizedWf6Af16Linear（quantized_linear.py#L95-L205）负责权重转换与前向计算：
   • supports_config要求输入/输出 dtype 必须为 FP16（FP6 数据项以 fp16 张量集合的形式打包存储，例如 8 个 fp6 项存于 3 个 fp16 槽位），并验证所请求激活（含门控激活）kernel 的可用性；
   • transform_param对每个线性层权重执行fp_quantize(param, num_bits=6, exp_bits=3)得到“fp16 容器中的 fp6 值 + 每输出通道 scale”，再调用preprocess_weight拆成 2bit/4bit 两个位平面，bias 等一维参数原样透传；
   • forward区分门控线性层（如 LLaMA 风格的 gate+up 融合，输出通道数翻倍、使用预分配 double buffer，一次 kernel 调用产出两路结果后再接门控激活）与普通线性层（matmul 后接 bias + 激活），输出写入预分配缓冲，避免每步重复分配。

量化函数细节。同文件中的fp_quantize（quantized_linear.py#L25-L92）值得展开：它依赖 qtorch 的float_quantize（未安装会明确提示pip install qtorch），固定采用 6 位、3 位指数的默认配置，量化范围q_range = 28；按输出通道（默认group_size为最后一维）计算scale = max(|x|)/28，并将零 scale 置 1 防止除零；缩放后的张量经 qtorch 做 round-to-nearest 浮点量化，最后以 FP16 容器回传，反量化值即quantized_fake_fp6 * scales。目前仅支持按最后一维的分组方式，其他group_size会抛出NotImplementedError。

仓库的量化算子测试 test_fp_quant.py 进一步提供了交叉验证：DeepSpeed 的 FP 量化器与 qtorch 的float_quantize在相同exp_bits/man_bits/group_size下逐项对比，要求两者平均量化误差差值小于 0.0004，覆盖了整张量量化与随机行索引子集两种场景，是 FP6 量化正确性的直接回归依据。

3.1 端到端服务性能

在两块 A100-80G GPU 上对 LLaMA-2-70B 使用 FP6 量化进行服务评估，相比 FP16 基线实现1.5 倍推理延迟降低与3.5 倍推理吞吐提升。FP6 量化带来三类关键收益：

1. 更少 GPU 部署大模型：LLaMA-70B 以 FP6 形式可在单个 A100-80G 上运行，而 FP16 至少需要两个 GPU；
2. 加速访存受限的线性层：小 batch 下 decode 阶段的线性层计算以内存访问为瓶颈，位宽减半的权重直接减少搬运量；
3. 降低权重显存占用：允许同时服务更多查询，提高服务吞吐。

系统在长序列生成场景下效率尤为突出：对于超过提示长度的生成长度，性能优势显著，且随生成序列延长，FP6 与 FP16 的差距进一步扩大。博客给出的测试设置（DeepSpeed-MII，128 个请求、32 个客户端、2xA100-80G）中，尝试了 128/256/512 等不同请求规模，加速效果相似。解码变长导致访存瓶颈增强的原因有二：

• KV cache 显存随序列增长，可容纳的 batch 变小，线性层从计算受限转变为参数访存受限；
• DeepSpeed-FastGen 的 prefill-decoding-mixed-batch 技术下，当 decoding 较长时，参与 mixed-batching 的 prefill 切块相对不足，纯 decoding batch 的出现频率增加，进一步加剧访存瓶颈——这恰好是权重-only 量化 kernel 最能发挥优势的区间。

3.2 当前限制

• 当 GEMM 因 batch 较大或显存充足而转为Tensor Core 计算受限时，权重-only 量化 kernel 可能无法保持对 cuBLAS 等厂商优化库的延迟优势，但低显存占用仍是关键优势；
• 目前支持范围限于非 MoE（Non-MoE）稠密结构，MoE 支持正在进行中；
• 当前仅兼容FP16 输入模型，因为该 FP6 kernel 仅处理 FP16 激活（源码层面即体现为supported_dtypes = [DtypeEnum.fp16]与input_dtype != torch.float16时拒绝配置）。

4. 如何开始

DeepSpeed-FP6 的量化与推理体验简单直接。以 LLaMA-2-70B 为例：

import mii pipe = mii.pipeline("NousResearch/Llama-2-70b-hf", quantization_mode='wf6af16') response = pipe(["DeepSpeed is", "Seattle is"], max_new_tokens=128) print(response)

环境依赖（官方博客给出的安装命令）：

pip install deepspeed-mii pip install qtorch

其中 qtorch 是 FP6 量化的运行时依赖，当前仓库的推理依赖清单 requirements-inf.txt 中已包含qtorch（开发依赖 requirements-dev.txt 固定为qtorch==0.3.0）。CUDA 侧的 FP6 算子通过InferenceCoreBuilder构建加载（quantized_linear.py#L154-L156 中的InferenceCoreBuilder().load()），算子构建器定义见 inference_core_ops.py。

需要说明的适用前提：quantization_mode='wf6af16'仅在当前 GPU 为NVIDIA Ampere 架构（compute capability major == 8）且运行于原生 CUDA（非 ROCm）的 PyTorch 时可用，否则引擎会抛出明确的ValueError；模型线性层的out_channels/in_channels还需满足 256/64 的整除约束。进行基准测试可参考 DeepSpeedExamples 仓库中的benchmarks/inference/mii/run_fp6.sh脚本；FP6 的独立 kernel 由悉尼大学发布（FP6-LLM 项目）。

5. 软件改进与演进方向

DeepSpeed-FP6 目前仅支持线性 GEMM，团队期待未来支持 MoE GEMM，并将持续依据社区反馈改进。DeepSpeed-FP6 是更大 DeepSpeed 生态的一部分，覆盖一系列深度学习系统与建模技术。社区贡献方面，欢迎报告问题、提交 PR 并参与讨论，贡献流程可参考仓库内的 CONTRIBUTING.md。

6. 致谢与贡献

DeepSpeed-FP6 感谢悉尼大学和罗格斯大学的合作，以及开源量化库 qtorch 的支持。

贡献者：Xiaoxia Wu*、Zhen Zheng*、Haojun Xia*（平等贡献）、Arash Bakhtiari、Michael Wyatt、Shiyang Chen、Stephen Youn、Reza Yazdani Aminabadi、Yuxiong He、Olatunji Ruwase、Zhewei Yao、Leon Song（项目负责人）；机构为微软（1）、悉尼大学（2）、罗格斯大学（3）。

相关文献：

1. ZeroQuant(4+2): Redefining LLMs Quantization with a New FP6-Centric Strategy for Diverse Generative Tasks. arXiv:2312.08583
2. FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design. arXiv:2401.14112
3. FP6-LLM kernel 独立发布（悉尼大学 fp6_llm 项目）

7. 小结：从博客到源码的一张完整地图

从“为什么选 6 位”到“6 位如何在 GPU 上跑得快”，再到“如何在 MII 引擎里一行配置启用”，DeepSpeed-FP6 构成了一条完整的、有源码可查的技术链路。对于以 decode 为主的中小 batch 在线服务场景，它是当前仓库内降低 LLM 显存占用与推理延迟的成熟量化路径；而对于大 batch 计算受限场景与 MoE 模型，则需要等待后续版本的能力扩展。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed