DeepSeek-V3 FP8 量化原理与工程实现
DeepSeek-V3 横空出世,训练和推理成本极低,一个重要的原因就是采用了 FP8 进行训练和推理。本文结合实践分析其中的原理、配置及工程落地细节。
DeepSeek-V3 采用 FP8 量化技术显著降低了训练和推理成本。文章深入解析了 Group/Block-wise 量化原理,解释了如何通过分块量化解决特征异常值问题并平衡精度与效率。详细介绍了 DeepSeek-V3 的量化配置,包括权重的静态 Block-wise 量化和激活的动态 Per-token-group 量化。重点讨论了 FP8-GEMM 的工程实现,涵盖 GPU 计算流程中的下溢处理、Python Native 实现、Triton 算子配置及 CUTLASS 库中的分块与分组缩放策略。最后总结了硬件支持(如 Hopper 架构 TMA)与算子优化的关键点。

DeepSeek-V3 横空出世,训练和推理成本极低,一个重要的原因就是采用了 FP8 进行训练和推理。本文结合实践分析其中的原理、配置及工程落地细节。
分块量化(Block-wise Quantization),也称为分组量化(Per-group Quantization),是一种细粒度量化方法。
特征异常值是指在特征分布中远离大部分数据的极端值。这些异常值对量化尤其具有挑战性,因为如果使用全局的量化参数(例如最大值),则这些异常值可能会导致大部分数据的量化精度下降。
细粒度量化的核心思想是使用更精细的量化粒度,即对输入和权重的不同部分使用不同的缩放因子。这样可以更好地适应数据的局部特征,减少异常值的影响。

分块量化将张量分割成更小的块或组,并为每个块分配独立的量化参数(缩放因子 s 和零点 z)。
如上图所示,矩阵被分割成多个小块,每个小块使用不同的颜色进行标注,对应不同的量化参数。
优点:
缺点:
总之 Block-wise 量化是对矩阵分组,每一组有独立的量化参数,可以更好的控制精度损失。
首先看 DeepSeek-V3 FP8 版本的模型配置:
{
"quantization_config": {
"activation_scheme": "dynamic",
"fmt": "e4m3",
"quant_method": "fp8",
"weight_block_size": [
128,
128
]
}
}
[128, 128],静态离线量化
(1) 对于激活值,我们以 1x128 的组为基础对元素进行分组和缩放(每个 token 每 128 个通道); (2) 对于权重,我们以 128x128 的块为基础对元素进行分组和缩放(每 128 个输入通道每 128 个输出通道)。
假设权重 B 的 shape 为:[hidden_dim, out_dim]
hidden_dim 维度上每 128 个输入特征一组out_dim 维度上每 128 个输出特征一组[hidden_dim//128, out_dim//128]假设激活 A 输入的 shape 为:[batch_size x seq_len, hidden_dim]
hidden_dim 维度上每 128 个通道的激活值分为一组,并为这一组计算一个单独的缩放因子。[batch_size x seq_len, hidden_dim//128]下面主要针对 FP8 GEMM 的工程实现讨论。理解了上面的权重和激活量化原理,那么下面来看如何进行两个 FP8 量化矩阵的乘法运算。
经过量化,我们得到了下面参数:
// inputs
// A [M, K] fp8 (按行分组量化,每组对应一个 As 元素)
// B [N, K] fp8 (按块量化,块大小为 [block_k, block_n],每个块对应一个 Bs 元素)
// As [M, K/block_k] fp32 (A 的每行(或每组)的量化比例因子)
// Bs [K/block_k, N/block_n] fp32 (B 的每个块的量化比例因子)
// outputs
// mat [M, N] fp32
背景:
下溢指的是计算结果的绝对值非常小,小于浮点数所能表示的最小正数(非零)。换句话说,计算结果太接近于零,以至于计算机无法用当前的浮点数格式精确地表示它,通常会被近似为零。
DeepSeek-V3 的方案: 所有 FP8 张量都采用 E4M3 格式 (4 位指数和 3 位尾数),以获得更高的精度。

计算过程:
以 NC = 128 个元素 MMA 的间隔转移到 CUDA Cores 进行高精度累加。

每当 Tensor Core 累加了 128 个 FP8 结果后,就会将这些结果转换(或缩放)到 FP32 精度,然后在 CUDA Cores 的 FP32 寄存器中进行累加。
详细步骤:
核心代码示例:
def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.float16):
"""This function performs matrix multiplication with block-wise quantization using native torch.
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
The output is returned in the specified `output_dtype`.
"""
M, K = A.shape
N, K_B = B.shape
block_k, block_n = block_size
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
assert n_tiles == Bs.shape[0]
assert k_tiles == Bs.shape[1]
C_shape = (M, N)
C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
B_tiles = [
[
B[
j * block_n : min((j + 1) * block_n, N),
i * block_k : min((i + 1) * block_k, K),
]
for i in range(k_tiles)
]
for j in range(n_tiles)
]
C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
for i in range(k_tiles):
for j in range(n_tiles):
a = A_tiles[i] # [M, 128]
b = B_tiles[j][i] # [128, 128]
c = C_tiles[j] # [M, 128]
s = As_tiles[i] * Bs[j][i] #[M, 1]
c[:, :] += torch.matmul(a, b.t()) * s
C = C.reshape(C_shape).to(output_dtype)
return C
可以结合上面对矩阵乘法的注释来理解分块矩阵乘法的过程:进行矩阵乘法的时候,先对矩阵 A 和 B 依照各自的量化粒度分块,在分块的粒度上进行矩阵乘法运算,然后再乘以量化因子进行反量化,得到分块的 FP32 浮点结果。
代码参考 sglang 中的实现:
def w8a8_block_fp8_matmul(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
block_size: List[int],
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
"""This function performs matrix multiplication with block-wise quantization.
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
The output is returned in the specified `output_dtype`.
Args:
A: The input tensor, e.g., activation.
B: The input tensor, e.g., weight.
As: The per-token-group quantization scale for `A`.
Bs: The per-block quantization scale for `B`.
block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128].
output_dtype: The dtype of the returned tensor.
Returns:
torch.Tensor: The result of matmul.
"""
pass
# 尝试加载之前通过 tuning 方式获得的最佳配置信息。
configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1])
if configs:
# If an optimal configuration map has been found, look up the
# optimal config
config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
else:
# Default config
# Block-wise quant: BLOCK_SIZE_K must be divisable by block_size[1]
config = {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": block_size[0],
"BLOCK_SIZE_K": block_size[1],
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
}
可以通过对 Triton 算子进行 tuning 来得到最优的 kernel 配置,接着调用 Triton 算子。
Triton 的代码介于 PyTorch 和 CUDA 代码之间,它提供了一种比手写 CUDA 算子更高层次的抽象,方便开发。
核心计算流程如下,注意累加器 accumulator 是 float32 精度的。
@triton.jit
def _w8a8_block_fp8_matmul(
# Pointers to inputs and output
# ... parameters ...
):
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
k_start = k * BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
先了解一下几种量化缩放的术语(和量化粒度有关):
关于 FP8-block-wise 量化有先后两个 PR,第一个 PR 先支持了Blockwise Scaling,第二个 PR 在第一个的基础上支持了 Groupwise Scaling,下面依次介绍。
第一个 PR 实现了 CUTLASS F8 GEMM 的分块缩放(Blockwise Scaling),通过共享内存暂存缩放张量,并为将来支持分组缩放做准备。

上面这张图表示了分块缩放:
第二个 PR 在第一个 PR(添加了分块缩放策略)的基础上,进一步添加了针对 A 张量 M 维度的分组缩放策略。
沿 M 维度的缩放粒度与 CTA 块配置无关,但是,沿 N 和 K 维度的缩放粒度仍然是分块的(即每个 CTA 块一个缩放值)。
所以到了这一步,基于这个 PR 我们才能实现与前面 Pytorch 和 Triton 代码功能相同的 kernel。
作者具体使用了 CUTLASS 3.0 新 API 在 Hopper 架构上进行分组缩放 FP8 GEMM 运算。
NVIDIA Hopper 架构引入了新的 tensor core 指令集 (GMMA),比 Ampere 的 tensor core 指令更高效。
Hopper 架构包含新的 Tensor Memory Accelerator (TMA) 单元,可以在全局内存和共享内存之间高效地传输大型数据块。TMA 还支持线程块之间异步拷贝。
使用了 Warp Specialized 内核设计。
CUTLASS 中 FP8 E4M3 使用 cutlass::float_e4m3_t 来表示,代码比较长,在这里不详细分析了。
作者在 PR 里给的 example 只是告诉了如何使用这个 GEMM,基于特定场景还需要做定制化开发优化。
如果要基于 CUTLASS 3.0 开发 kernel,有两个利器可以使用:
本文详细介绍了 FP8 block-wise 量化的原理以及推理的工程实现。很多时候想到一个好的量化算法并不难,难的是和硬件特性结合起来,在保持精度的前提下发挥量化的最大性能。
DeepSeek-v3 这类 MOE 模型在推理中,有需要用到 Grouped-GEMM 的场景,后续可进一步分析这一块。随着大模型应用日渐普遍,大模型的训练和推理工程基础建设如火如荼,AI Infra 在大模型这波浪潮中将扮演重要的角色。掌握底层量化原理与算子实现能力,对于构建高性能大模型系统至关重要。

微信公众号「极客日志」,在微信中扫描左侧二维码关注。展示文案:极客日志 zeeklog
使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具,Base64 字符串编码/解码在线工具,online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具,Base64 文件转换器在线工具,online