vLLM 量化 Kernel 配置机制:W8A8 Block FP8 Triton 内核的调优、选择与回退原理

vLLM 为块量化 W8A8(FP8)GEMM 内核维护了一套“设备 + 形状”维度的预调优 Triton 配置库,运行时按 GEMM 形状与批大小自动命中最优配置,未命中时回退默认配置并告警。本文基于 configs/README.md 所指的配置目录,完整讲解这些 JSON 配置文件的命名规范、字段含义、用 benchmark_w8a8_block_fp8.py 自行生成配置的全流程,以及 vLLM 运行时加载与回退的源码机制,帮助你在自己的 GPU 上复现调优并理解“为什么这块卡上推理更快”。

一、什么是 Quantization Kernel Config

配置文件目录 vllm/model_executor/layers/quantization/utils/configs/ 中的 README 只有两句话,但指向了核心事实:

Use scripts under benchmarks/kernels/ to generate these config files.

即:该目录下所有 JSON 文件均由仓库 benchmarks/kernels/ 下的调优脚本在真实 GPU 上测量生成,而非人工拍脑袋设定。当前目录包含 219 个 JSON 文件,按“矩阵形状 × 设备 × 数据类型 × 量化块”四维命名,例如:

N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json
N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128,128].json
N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128,128].json
N=1536,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128,128].json

文件名各段含义:

字段	含义	示例取值
N=...	权重矩阵的输出维度(行)	1536、7168、24576
K=...	权重矩阵的输入维度(列)	512、7168、16384
device_name=...	规范化后的设备名(见下文)	NVIDIA_H100_80GB_HBM3、AMD_Instinct_MI300X
dtype=..._w8a8	内核数据类型	fp8_w8a8、int8_w8a8
block_shape=[n,k]	权重块量化粒度	[128,128]

其中 N/K 组合(如 N=7168, K=1536、N=12288/24576, K=7168)对应 DeepSeek-V3/R1 系列模型中典型的注意力投影与 MoE 专家投影形状;覆盖的设备横跨 NVIDIA(H100、H200、H20、B200、A100、A800、L20、L20Y、L40S)与 AMD(MI300X、MI325X、MI325_OAM)。目录中同时存在 int8_w8a8 条目(主要是 A100/A800 上无 FP8 Tensor Core 的块量化路径),说明同一套“形状 → Triton 配置”机制被复用于不同数据类型。

设备名的规范化由 get_device_name_as_file_name 完成:它把 torch.cuda.get_device_name() 返回的字符串去除空格,保证文件名在不同驱动/系统版本下保持稳定。

二、配置文件内部结构:一个“批大小网格 → Triton 启动参数”的映射

以一个真实文件 N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json 为例,其结构为:

{
    "1":    { "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128,
              "GROUP_SIZE_M": 32, "num_warps": 4, "num_stages": 4 },
    "2":    { "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 128,
              "GROUP_SIZE_M": 1,   "num_warps": 4, "num_stages": 5 },
    ...
    "512":  { "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 128,
              "GROUP_SIZE_M": 1,   "num_warps": 4, "num_stages": 3 },
    "1024": { "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128,
              "GROUP_SIZE_M": 16,  "num_warps": 4, "num_stages": 3 }
}

• 顶层 key 是字符串化的批大小 M(请求 token 数),取值来自一个非均匀网格:1、2、4、8、16、24、32、48、64、96、128、256、512、1024、1536、2048、3072、4096。网格在 decode 小批区域采样密集(1~128 逐个覆盖),在 prefill 大批区域采样稀疏——这正是 LLM 服务中最常见的两种工作负载形态。
• value 是该 M 下测得最快的 Triton 内核启动配置,字段含义:
  • BLOCK_SIZE_M / BLOCK_SIZE_N / BLOCK_SIZE_K:Triton GEMM 的 tile 大小;
  • GROUP_SIZE_M:grouped (swizzle) 调度时 M 方向的分组数,影响 L2 复用;
  • num_warps:每个 program 的 warp 数;
  • num_stages:软件流水级数,决定寄存器/SRAM 占用。

可以观察到,小 M 下 BLOCK_SIZE_N 常取 32 以减小计算冗余,大 M 下回到 64;num_stages 则随 tile 大小在 3~5 之间变化,以在寄存器压力和流水深度之间取得平衡。这些差异正是逐设备、逐形状实测出来的结果,无法靠通用规则推导。

三、如何生成配置文件:benchmark_w8a8_block_fp8.py 全流程

生成入口是 benchmarks/kernels/benchmark_w8a8_block_fp8.py(注释中说明其改编自 SGLang 的块量化内核调优脚本)。脚本头部的官方用法:

# Tune triton w8a8 block fp8 for DeepSeek-V3/DeepSeek-R1:
python3 benchmark_w8a8_block_fp8.py --tp-size 8 --input-type fp8
# 然后把输出拷贝到 vllm/model_executor/layers/quantization/utils/configs

3.1 命令行参数

参数	默认值	说明
--tp-size / -tp	8	张量并行度,用于切分权重形状得到每卡实际的 N/K
--input-type	fp8	输入数据类型,当前仅支持 fp8(float8_e4m3fn)
--out-dtype	float16	输出 dtype,可选 float32/float16/bfloat16/half
--block-n / --block-k	128 / 128	权重块量化粒度,与文件名 block_shape 对应
--batch-size	全部 18 个网格点	只测单个批大小;指定后仅用 1 张 GPU
--save-path	./	JSON 输出目录,需再拷贝进包内 configs 目录

3.2 形状集与张量并行切分

get_weight_shapes(tp_size) 内置了 DeepSeek-V3 的三类形状:不可 TP 的(如 (512+64, 7168)、(2112, 7168)、(128*256, 7168) 专家形状)、N 维度可 TP 的(如 (24576, 1536)、(12288, 7168))、K 维度可 TP 的(如 (7168, 18432))。TP=8 时,例如 (128*(128+128), 512) 会切成 (4096, 512)。如果你的目标模型不同(注释中明确提示 “Modify them, if you tune for another different model”),需要按目标模型的每卡投影矩阵形状替换该表。

3.3 搜索空间与测量方法

get_configs_compute_bound() 生成约 1280 个候选配置:num_stages ∈ {2,3,4,5} × BLOCK_SIZE_M ∈ {16,32,64,128,256} × BLOCK_SIZE_K ∈ {64,128} × BLOCK_SIZE_N ∈ {32,64,128,256} × num_warps ∈ {4,8} × GROUP_SIZE_M ∈ {1,16,32,64},再按 block_k % BLOCK_SIZE_K == 0 过滤(块量化要求 K-tile 整除量化块)。

对每个 (M, 配置) 组合,benchmark_config 用 CUDA Event 计时:5 次 warmup(顺带触发 Triton JIT 编译)后取 10 次迭代的平均微秒数;编译期资源不足的候选(OutOfResources)被直接跳过。最终每个 M 保留延迟最低的一组配置,写入 JSON。

3.4 多 GPU 并行策略

main 会探测本机全部 GPU,把默认 18 个批大小按 distribute_batch_sizes 均分到各 GPU,通过 multiprocessing spawn 池并行调优——每卡负责一部分 M,但对所有权重形状都要跑完,因此多卡能线性缩短总时长。单个 --batch-size 场景则回落到单卡串行。

跑完后,脚本打印 Writing best config to ...,把 JSON 拷入 vllm/model_executor/layers/quantization/utils/configs/ 即可被运行时命中(该目录的 *.json 已在 setup.py 的包数据清单中,随 wheel 一起分发)。

四、运行时机制:自动命中、最近邻选择与默认回退

4.1 配置查找:get_w8a8_block_fp8_configs

加载逻辑在 fp8_utils.py:

@functools.lru_cache
def get_w8a8_block_fp8_configs(N, K, block_n, block_k):
    device_name = get_device_name_as_file_name()
    json_file_name = (
        f"N={N},K={K},device_name={device_name},"
        f"dtype=fp8_w8a8,block_shape=[{block_n},{block_k}].json"
    )
    config_file_path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
    if os.path.exists(config_file_path):
        return {int(key): val for key, val in json.load(open(config_file_path)).items()}
    logger.warning(
        "Using default W8A8 Block FP8 kernel config. Performance might "
        "be sub-optimal! Config file not found at %s", config_file_path)
    return None

要点:

1. 文件名即索引:查找时与生成脚本 save_configs 使用同一套命名模板,生成/加载两端天然对齐;
2. lru_cache:同一 (N, K, block) 组合只读盘一次,字典常驻内存;
3. 未命中不报错而是告警回退:日志形如 Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! Config file not found at ...——benchmarks/kernels/deepgemm/README.md 中就保留了这样一段真实运行日志,展示 H100 上部分形状命中、部分形状回退的混合状态。

4.2 内核调用:按最近邻批大小取配置

w8a8_triton_block_scaled_mm(fp8_utils.py)是块量化 FP8 GEMM 的统一入口(DeepSeek-V3/R1 等块量化模型的线性层经 FP8 量化路径最终走到这里)。它在拿到 M 后:

configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1])
if configs:
    config = configs[min(configs.keys(), key=lambda x: abs(x - M))]

即:在配置的 M 网格中做最近邻匹配——M=37 会命中“32”那行配置,M=2050 命中“2048”。这就是配置 key 采用非均匀网格的原因:小 M 段密集保证 decode 精度,大 M 段稀疏且误差相对可忽略。若 configs 为 None(未命中文件),则使用函数内置的默认 Triton 配置运行,功能正确但性能可能次优。

五、同构机制的延伸:MoE 内核配置与用户自定义目录

同样的“形状 → 逐批大小 Triton 配置”思想被 fused_moe.py 的 get_moe_configs 复用,但有三个值得注意的差异:

• 文件名多一个专家维度:E={E},N={N},device_name=...{dtype}{block_shape}.json,放在 fused_moe/configs 目录,由 benchmark_moe.py 类脚本生成;
• 支持用户目录覆盖:优先查找环境变量 VLLM_TUNED_CONFIG_FOLDER 指向的目录(“note that we prioritize user defined config”),再回落到包内目录。也就是说,你不必改仓库,只需在自己的目录放同名 JSON 并设置该环境变量即可让 vLLM 使用自定义调优结果;
• 批不变模式禁用调优配置:当 VLLM_BATCH_INVARIANT=1 时直接返回 None 使用默认配置,保证不同批大小下输出比特级一致(见 batch_invariance 特性文档)。

对比之下,get_w8a8_block_fp8_configs 目前只查包内 configs 目录;如果部署环境需要覆盖 W8A8 配置,可参照 MoE 路径的做法理解这一机制的差异(从源码结构看,这是两条独立实现)。

六、实践建议

1. 先查日志再调优:启动服务时若看到 Using default W8A8 Block FP8 kernel config... 告警,说明你的“卡型 × 形状 × dtype”组合尚无预调优配置;先确认该形状是否真的出现在目标模型里,再决定是否调优。
2. 在自己的设备上完整复现:按 --tp-size 指定生产环境的 TP 度调优,因为 N/K 会随 TP 切分而改变;跑完把 JSON 放入包内目录(或通过部署层挂载/补丁方式分发),文件名必须与运行时命名模板逐字符一致。
3. 调优对象要匹配:内置权重形状面向 DeepSeek-V3/R1;服务其他模型时替换 get_weight_shapes 中每卡实际投影形状,并可用 --batch-size 对个别热点形状做补测。
4. 理解回退是安全设计:缺配置只降级性能、不影响正确性,因此可以把“全形状覆盖”作为持续集成任务(新卡型出货后批量跑一遍),而不必担心线上风险。

参考

• 配置目录与命名总览:vllm/model_executor/layers/quantization/utils/configs/(219 个 JSON,README 见 configs/README.md)
• 生成脚本:benchmarks/kernels/benchmark_w8a8_block_fp8.py
• 运行时查找与内核入口:fp8_utils.py
• MoE 配置查找与用户目录覆盖:fused_moe.py
• 打包分发配置:setup.py