Flash Grouped Per-Token FP8 量化

DeepSeek-V3 采用了低精度 (FP8) 进行训练，因此大部分 Linear (GEMM) 操作是在 FP8 精度下进行的。使用 FP8 精度 GEMM 的优势在于：在 GPU 上其耗时只有 bfloat16 的一半。这听起来很美好，但忽略了其他操作（如 attention）仍旧在 bfloat16 下进行，所以在进行 GEMM 之前，需要先将 input 转换为 FP8。在实际训练中，这部分如果没有高性能的 kernel 实现，往往会严重影响性能。在最初对 DeepSeek-V3 进行后训练 (post-training) 的过程中，我们发现使用 PyTorch 实现的 naive 版本进行 FP8 训练，甚至比直接使用 bfloat16 训练还要慢。本文将主要探讨如何解决这个问题。

1. 背景

作为背景知识的一部分，我们首先介绍下 DeepSeek-V3 为了进行 FP8 训练所采用的量化方式：

1. 针对 input (activation) 的 Per-Token Quantization (1x128 tile)
2. 针对 weight 的 Per-Block Quantization (128x128 tile)

Per-Token Quantization 主要是对 input 进行量化：对于每个 token，沿着 channel 维度每 128 个元素计算最大值作为 scale，利用该 scale 对这 128 个元素进行缩放，随后转换为 FP8 精度。Python 代码如下：

## per token cast to fp8
def per_token_cast_to_fp8(
    x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    assert x.dim() == 2
    m, n = x.shape
    # n is divisible by 128 for simplicity; otherwise padding is needed
    x_view = x.view(m, -1, 128) # (m, n//128, 128)
    x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) # (m, n//128)
    fp8_data = x_view * (448.0 / x_amax.unsqueeze(2))
    fp8_data = fp8_data.to(torch.float8_e4m3fn)
    scale = (x_amax / 448.0).view(m, -1)
    return fp8_data, scale

当然，Per-Token Quantization 的整体计算流程并不复杂，目前已有许多开源的高性能实现，因此在本文中只做简单介绍。同时，本文不涉及 weight 的量化，因为 weight 的量化在初始化阶段即已完成，属于一次性开销。

本文想要着重探讨的是如何实现 Grouped Per-Token Quantization，它是 Per-Token Quantization 的一种特例。要引出 Grouped Per-Token Quantization，首先需要了解下 DeepSeek-V3 中 GEMM 的两种类型：Normal GEMM (比如 DeepSeekV3Attention 中的 Linear) 和 Group GEMM (DeepSeekV3MLP 中的 Linear)，因为 GEMM 类型和对 input 的量化类型是有对应关系的。Normal GEMM 就是普通的矩阵乘法，Group GEMM 则是多了 group 维度，input 按照不同的 group 进行区分，去和不同的 weight 进行 Normal GEMM， 相当于多了一层 for 循环。MoE 模型往往有多个 experts，对应不同的 weight，每个 token 在进行 forward 时会激活多个 expert，这里的 group 指的就是多个 expert 场景下单个 expert 分到的需要计算的 token 组成一个 group。

Figure 1 展示了 DeepSeek-V3 的 MoE 中 Group GEMM 的计算流程，一次完整的 Group GEMM 前向 (forward) 和反向 (backward) 计算包含三次 GEMM 计算：

# x.shape = (num_groups, m, k)
# w.shape = (num_groups, n, k)
# 在 DeepGEMM中，对 group 的区分在 kernel 里完成，在外面看起来就和普通的 GEMM 一样。

# note 1: 对于 GEMM，这里延续了 PyTorch Linear 的格式，也就是 w 的 shape 并非 (..., k, n)，而是(..., n, k)
# note 2: _nt 代表第一个 input(A)：Non-Transposed， 第二个 input(B)：Transposed，需要注意`理论计算`和`gemm api 调用`的区别。

# 1. forward
# o = x @ w.T
o = group_gemm_nt(x, w)

# 2. input backward (Dgrad)
# dx = do @ w
dx = group_gemm_nt(do, w.transpose(1,2)) # 注意：w 需要转置才是 group_gemm_nt 需要的格式

# 3. weight backward (Wgrad)
# dw = do.T @ x
dw = group_gemm_nt(do.transpose(1,2), x.transpose(1,2)) # do 和 x 均需要转置

仔细观察上述代码，我们可以发现，forward 阶段和 input backward 的计算其实和普通的 Per-Token Quantization 并无区别，因为对于 shape 为 (num_groups, m, n) 的 input，其实 dim-0 并无影响，也无需感知，和上面的 Per-Token Quantization 一样计算即可。

对于 weight backward (Wgrad)，情况则有所不同：

• GEMM 的两个输入（do 和 x）都需要进行转置。
• 转置之后，沿着原本的 m-dim 进行 scaling (考虑上述代码中的 x，forward 时是沿着 k-dim 进行 scaling)。

这就是我们想要探讨的 Grouped Per-Token Quantization，该操作是 “transpose + quantize” 的组合，其 naive 的 PyTorch 实现如下：

def grouped_fp8_cast(input, m_indices_counts):
    """
    naive 实现：先转置，再分块量化
    Args:
        input: (total_tokens, hidden)
        m_indices_counts: List of token counts per group
    """
    m = input.shape[0] # total_tokens
    total_scale_factors = sum(ceil_div(m_size, 128) for m_size in m_indices_counts)

    # 1. 显式转置
    # input.T -> (hidden, total_tokens)
    input_t = input.T.contiguous()

    n, m = input_t.shape # n=hidden, m=total_tokens
    output_t_fp8_flat = torch.empty((m * n,), device="cuda", dtype=torch.float8_e4m3fn)
    # output buffer for scales: (total_blocks, hidden)
    output_t_scales = torch.empty((total_scale_factors, n), device="cuda", dtype=torch.float)

    lhs_offset, scale_offset = 0, 0

    # 2. 循环处理每个 group
    for m_size in m_indices_counts.cpu().tolist():
        # slice 当前 group
        # input_t_chunk: (hidden, m_size)
        input_t_chunk = input_t[:, lhs_offset : lhs_offset + m_size]

        # 调用标准量化 kernel
        # note: per_token_cast_to_fp8 quantizes along the LAST dim (m_size).
        # result scales: (hidden, m_size//128)
        input_t_fp8_chunk, input_t_scale_chunk = per_token_cast_to_fp8(input_t_chunk)

        # 拷贝回结果 buffer
        output_t_fp8_flat[lhs_offset * n : (lhs_offset + m_size) * n].copy_(
            input_t_fp8_chunk.flatten()
        )

        num_scales = ceil_div(m_size, 128)
        # transpose scales to (m_size//128, hidden) for continuous storage per block-row
        output_t_scales[scale_offset : scale_offset + num_scales].copy_(
            input_t_scale_chunk.T
        )

        lhs_offset += m_size
        scale_offset += num_scales

    return output_t_fp8_flat, output_t_scales

我们解释一下上面 Grouped Per-Token Quantization 的 Python 代码逻辑，这可以拆分为两个步骤来理解：

1. Transpose: for-loop 之前，对整个 input 进行转置得到 input_t，并构造 output 和 scale。
2. Per-Token Quantization: 对 input_t 按 group 进行量化。

计算流程如 Figure 2 和 Figure 3 所示。

最终产生的 output_t_fp8_flat 和 output_t_scales 的格式可能会让人觉得困惑。具体为什么这么做需要结合 DeepGEMM 的实现，简单概括来说就是 Wgrad kernel 需要这种数据格式。此外：

• 将 output 进行 flatten 是因为每个 group 的 output tensor 大小不一样，flatten 之后可以拼接到一起，否则就需要构造多个 output tensor。
• scale 进行了转置，这也是 DeepGEMM kernel 的要求。如果不在这里做，DeepGEMM kernel 在调用接口时也会自己做一次（这将严重影响性能）。

2. 性能分析

现在，我们来分析 grouped_fp8_cast 的性能瓶颈，根据上述的代码，可以发现：

1. 显式转置: input.T.contiguous() 强制产生了一次完整的 global memory 读写（$2\times$ bytes），这对于 memory-bound 的 kernel 而言是巨大的开销。
2. 多次 kernel launch 以及负载不均: for-loop 实现导致每个 group 需要进行一次 kernel launch。负载不均，考虑 m=512 和 m=8192，均需执行一次 kernel，对于较小的 m，kernel launch 的耗时占比较大，同时也无法使得 GPU 保持较高的利用率。

input_shape = (24064, 7168) # bfloat16
m_indices = [5120, 512, 2560, 2048, 2048, 512, 3072, 8192]

我们选取训练过程中实际需要的一组配置，根据这个配置构造 input ，对 PyTorch 版本的实现进行耗时测试：

• Naive PyTorch 耗时: 4.3021ms。这个数字甚至超过了随后进行的 Group GEMM 计算本身的耗时！
• torch.compile 优化后: 1.5227ms。虽然性能有所提升，但作为 element-wise 操作，其速度仍然不够理想。

对于 naive 的 PyTorch 代码，input.T.contiguous() 的耗时大致是 1.3147ms，占比大约是 32%。

再观察 PyTorch Profiler 导出的信息，可以看到有大量细碎的 kernel 和一些 bubble，也说明整体的执行非常低效。

其实对于这类 element-wise 的操作，有一个简单和直观的方法来判断其性能是否足够好，那就是看整体的 throughput。简单来说，如果是纯内存拷贝，也就是 kernel 只负责把数据从 global memory 拷贝到 register，再写回，什么计算也不做。那么根据 H100 的带宽 (HBM Bandwidth ~3.3 TB/s)，对于上述配置：

• Input: $24064\times 7168\times 2$ bytes (BF16) = 329 MB
• Output: $24064\times 7168\times 1$ bytes (FP8) = 164.5 MB

理论耗时应在 0.15ms 左右。目前有效带宽利用率严重不足，即使是 torch.compile 优化后的版本，利用率也仅为 10%，说明存在巨大的优化空间。

3. 优化实现

根据上述的分析，核心思路非常直观：算子融合 (Kernel Fusion)。我们希望在一个 kernel 内完成数据读取、隐式转置、量化计算和写回，解决 grouped_fp8_cast 存在的性能瓶颈。

3.1 Persistent Kernel 与合并扫描

为了解决多次 kernel launch 以及负载不均的问题，我们采用 Persistent Kernel 的模式，所有的 thread block 协作处理整个输入直至整个任务完成，而不是为每个 group 单独启动一次 kernel。

如 Figure 5 所示，我们将 input 整体划分为 $128\times \text{TILE\_COL\_SIZE}$ 的 tile，启动 (launch) 一定数量的 thread block 去计算每个 tile，算完之后继续向后扫描，查看是否还有需要计算的 tile，直到结束。

由于 group 大小不一，在 thread block 内部需要一种方法让任意一个 tile 知道自己属于哪个 group（以便获取对应的 output 写入位置）。

解决思路：前缀和 + 二分查找

1. 预先对 group_m = $[m_0,m_1,...,m_g]$ 求前缀和，得到 accum_m = $[a_0,a_1,...,a_g]$。
2. 每个 thread block 根据自己当前计算的全局行号 (global_row_idx)，在 accum_m 数组中进行二分查找，定位 group ID。

根据上述的讨论，主体代码框架如下面的代码所示：

__global__ void per_token_cast_to_fp8_grouped_kernel(...) {
    // persistent kernel Loop
    int32_t tile_id = blockIdx.x;
    while (tile_id < total_tiles) {
        // ... 计算 global coordinates ...

        // 定位当前 tile 属于哪个 Group
        int32_t group_id = binary_group_id_locate(global_row_idx, group_offsets);

        // 1. load data (bfloat16) from global memory
        // 2. compute scale & quantize
        // 3. store data (FP8) to global memory

        tile_id += gridDim.x; // Move to next tile (Grid-Stride Loop)
    }
}

在 Expert Parallelism 场景下，group 数量很少（通常 < 128），将这些元数据加载到 shared memory 中进行二分查找几乎没有开销。

3.2 利用 Shared Memory

由于我们不希望额外进行一次转置，所以必须在 kernel 内部完成隐式的转置操作。如 Figure 5 所示，我们让 thread block 内部的每个 warp 负责处理 tile 中的一列，计算完 scale 和 FP8 结果后写回。这么做主要基于两个考量：

1. FP8 的量化正好就是沿着一列进行的，可以利用 warp reduce。
2. 一个 warp 完成计算后写入一个连续的 FP8 数据块 (coalesced access)，可以充分利用 HBM 带宽。

问题： 写回是连续的，但如果按照同样的方式读取，则是非连续的（因为 input 和 output 之间存在一个转置的过程）。

我们的思路是借助 shared memory：

1. Global memory (row-major) $\to$ shared memory (row-major)。这保证了 global memory 读取是连续的。
2. Shared memory $\to$ registers。warp 竖着读 shared memory 进行计算，但这又引出了下面的问题。

Bank Conflict： Shared memory 每 32 bits (4 bytes) 构成一个 bank，一共有 32 个不同的 bank (128 bytes)。

byte addr           bank
0-3                 0
4-7                 1
...
124-127             31
128-131             0
...

当一个 warp 的 32 个 threads 同时访问不同的 bank 时，只需一次读写即可完成，否则将串行化。 Warp 竖着读取 shared memory 时，bank conflict 几乎无法避免，最坏的情况下 32 个 threads 竖着读会访问到同一个 bank。 如 Figure 6 左边所示，假设 shared memory 一行有 32 个 bfloat16 数据，这导致了 16 路 Bank Conflict (16-way Bank Conflicts)（thread-0, thread-2, ..., thread-30 均命中 bank-0），极大降低了 shared memory 的带宽。

解决方法：Shared Memory Swizzling 在将数据写入 shared memory 时对 index 进行 swizzle，使每个 warp 竖向读取时也能命中不同的 bank，如 Figure 6 右边所示。

// 简易的 swizzle 实现
template<int32_t OuterSize, int32_t InnerSize, int32_t base_mask>
struct IndexSwizzler {
    static constexpr int32_t outer_bits = static_log2(OuterSize);
    static constexpr int32_t inner_bits = static_log2(InnerSize);
    static constexpr int32_t shiftr_bits = outer_bits - inner_bits;
    static constexpr int32_t offset_mask = base_mask << outer_bits;

    __device__ __forceinline__ int32_t swizzle(int32_t offset) const {
        return offset ^ ((offset & offset_mask) >> shiftr_bits);
    }
};

4. 最终结果

完成整体实现后，我们重新测试了之前的配置：

input_shape = (24064, 7168) # bfloat16
m_indices = [5120, 512, 2560, 2048, 2048, 512, 3072, 8192]

• 优化后耗时: 0.189ms
• 相比 Naive (4.3ms): 22.7x 加速
• 相比 torch.compile (1.52ms): 8x 加速
• Throughput: 2735 GB/s

在 H100 上，这一性能已经比较接近 HBM 的理论带宽。耗时也和常规的 per_token_cast_to_fp8（没有 transpose 也不用区分 group）基本一致。下面是一些测试数据：

参考

• DeepSeek-V3 Technical Report
• GitHub - deepseek-ai/DeepGEMM

@misc{2025-grouped-fp8,
  title   = {Flash Grouped Per-Token FP8 Quantization},
  author  = {Zhen Jiang in collaboration with others at Dawning Road},
  year    = {2025},
  url     = {/blog/grouped-fp8/},
  note    = {Blog post},
  urldate = {2026-10-24},
}