DeepSeek-V3 FP8 量化原理与工程实现

量化精度与粒度

• 量化精度：FP8
• 权重：block-wise 量化，每个 block 的 shape 是 [128, 128]，静态离线量化
• 激活：per-token-group 量化，动态在线量化

(1) 对于激活值，我们以 1x128 的组为基础对元素进行分组和缩放（每个 token 每 128 个通道）； (2) 对于权重，我们以 128x128 的块为基础对元素进行分组和缩放（每 128 个输入通道每 128 个输出通道）。

权重量化（block-wise）

假设权重 B 的 shape 为：[hidden_dim, out_dim]

1. 分块方式：
   • 在 hidden_dim 维度上每 128 个输入特征一组
   • 在 out_dim 维度上每 128 个输出特征一组
2. 量化缩放因子（scales）：
   • Bs 的 shape: [hidden_dim//128, out_dim//128]
   • 每个权重块使用独立的 scale
   • 静态量化：权重量化是离线预计算好的

激活量化（per-token-group）

假设激活 A 输入的 shape 为：[batch_size x seq_len, hidden_dim]

1. 分块方式：
   • 对于每一个 token，在 hidden_dim 维度上每 128 个通道的激活值分为一组，并为这一组计算一个单独的缩放因子。
2. 量化缩放因子：
   • As 的 shape: [batch_size x seq_len, hidden_dim//128]
   • 每个块使用独立的 scale 进行量化
   • 动态量化：在推理过程中，实时对激活进行量化
   • 不对 token 维度分块

FP8-GEMM 工程实现

下面主要针对 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

GPU 计算流程

背景：

• 下溢和精度损失： 使用 FP8 等低精度格式进行 GEMM 运算时，中间结果的累加容易出现下溢，导致精度损失。传统的做法是使用 FP32 进行累加，以保持精度。

下溢指的是计算结果的绝对值非常小，小于浮点数所能表示的最小正数（非零）。换句话说，计算结果太接近于零，以至于计算机无法用当前的浮点数格式精确地表示它，通常会被近似为零。

DeepSeek-V3 的方案： 所有 FP8 张量都采用 E4M3 格式 (4 位指数和 3 位尾数)，以获得更高的精度。

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

每当 Tensor Core 累加了 128 个 FP8 结果后，就会将这些结果转换（或缩放）到 FP32 精度，然后在 CUDA Cores 的 FP32 寄存器中进行累加。

详细步骤：

1. Tensor Core 以 FP8 精度高效地执行大量的矩阵乘法和累加（MMA）操作。使用低精度累加器存储中间结果。
2. 每累加 128 个元素（Nc = 128），就将这些 FP8 累加结果转换到 FP32 精度。
3. 在 CUDA Cores 的 FP32 寄存器中进行高精度的累加，最终结果经过 Scaling Factor 缩放，也就是反量化。
4. 重复步骤 1-3，直到完成所有的矩阵乘法和累加操作。

Python Native 实现

核心代码示例：

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 浮点结果。

Triton 实现

代码参考 sglang 中的实现：

1. 函数接口

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

2. Triton 算子配置

# 尝试加载之前通过 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 算子。

3. 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

CUTLASS 实现

先了解一下几种量化缩放的术语（和量化粒度有关）：

• 张量级缩放（Tensor-wise Scaling）： 每个张量使用单个缩放因子，在尾声（epilogue）中应用。
• 行级缩放（Row-wise Scaling）： 使用一个行向量进行缩放，对于操作数 A 的维度为 Mx1，对于操作数 B 的维度为 1xN，避免沿归约维度进行缩放。这也可以在尾声中使用 EpilogueVisitorTree 来处理。
• 分块缩放（Block-wise Scaling）： 引入一个 2D 缩放张量，每个 CTA 块分配一个缩放值。由于此缩放涉及归约维度 (M, N, K)，因此必须在主循环中应用，这会影响性能。
• 分组缩放（Group-wise Scaling）： 使用一个 2D 缩放张量，每个 CTA 块有多个缩放值。缩放粒度独立于 CTA 块配置，为将来的实现提供了更大的灵活性。

关于 FP8-block-wise 量化有先后两个 PR，第一个 PR 先支持了Blockwise Scaling，第二个 PR 在第一个的基础上支持了 Groupwise Scaling，下面依次介绍。

分块缩放

第一个 PR 实现了 CUTLASS F8 GEMM 的分块缩放（Blockwise Scaling），通过共享内存暂存缩放张量，并为将来支持分组缩放做准备。

上面这张图表示了分块缩放：

• CTA Tile (128x128x128): 这表示一个 CUDA 线程块（CTA）处理的数据块的大小。在这个例子中，每个 CTA 处理一个 128x128 的输出块，其中 K 维度（归约维度）的大小也是 128。这个 128x128x128 是逻辑上的，实际的计算可能根据 warp size 进一步划分。
• A (640x512): 输入矩阵 A 的维度是 640x512 (M x K)。
• B (512x384): 输入矩阵 B 的维度是 512x384 (K x N)。
• Scale A (5x4): A 的缩放因子矩阵，维度是 5x4。这里的 5 对应 M 维度被划分成的块数（640 / 128 = 5，向上取整），4 对应 K 维度被划分成的块数（512 / 128 = 4）。每个元素对应 A 的一个 128x128 的块的缩放因子。
• Scale B (4x3): B 的缩放因子矩阵，维度是 4x3。这里的 4 对应 K 维度被划分成的块数（512 / 128 = 4），3 对应 N 维度被划分成的块数（384 / 128 = 3）。每个元素对应 B 的一个 128x128 的块的缩放因子。
• 输出矩阵 C (绿色部分): 输出矩阵 C 的维度是 640x384 (M x N)。它被划分成若干个 128x128 的块，每个块由一个 CTA 计算。

分组缩放

第二个 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，有两个利器可以使用：

1. CuTe：大大简化了 CUTLASS 中复杂数据布局和线程组织的管理。
2. EVT：用于在 GEMM 的尾声阶段融合各种后处理操作的框架，以提高性能和效率。参考论文：EVT: Accelerating Deep Learning Training with Epilogue Visitor Tree。

总结

本文详细介绍了 FP8 block-wise 量化的原理以及推理的工程实现。很多时候想到一个好的量化算法并不难，难的是和硬件特性结合起来，在保持精度的前提下发挥量化的最大性能。

DeepSeek-v3 这类 MOE 模型在推理中，有需要用到 Grouped-GEMM 的场景，后续可进一步分析这一块。随着大模型应用日渐普遍，大模型的训练和推理工程基础建设如火如荼，AI Infra 在大模型这波浪潮中将扮演重要的角色。掌握底层量化原理与算子实现能力，对于构建高性能大模型系统至关重要。