SGLang W8A8 块量化 Triton 内核自动调优实战：从 DeepSeek-V3 到自定义模型的 tuning_block_wise_kernel 全指南

SGLang W8A8 块量化 Triton 内核自动调优实战从 DeepSeek-V3 到自定义模型的 tuning_block_wise_kernel 全指南【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglangW8A8 块量化block-wise quantization是 SGLang 在 FP8/INT8 低比特推理中支撑高性能 GEMM 的核心手段之一。本指南围绕 benchmark/kernels/quantization/README.md 展开完整讲解 Triton FP8/INT8 块量化内核的自动调优工具 tuning_block_wise_kernel.py 的适用场景、命令行用法、N/K 维度计算以及调优结果的落盘与运行时加载机制。读完本文你将能独立为 DeepSeek-V3 或任意自定义模型如 Qwen3-VL生成针对本机 GPU 的最优内核配置并理解这些配置如何被 SGLang 运行时消费。一、什么时候用 Triton 块量化内核什么时候用 DeepGEMMSGLang 的 W8A8 块量化 GEMM 存在两套主流实现Triton 内核与 DeepGEMM。调优之前必须先明确你当前需要优化的是哪一条路径因为二者适用的硬件与输出数据类型不同。优先使用 Triton FP8 块量化内核的场景输出 dtype 不是bfloat16例如float16、float32。DeepGEMM 的w8a8_block_fp8_matmul_deepgemm在 fp8_kernel.py 中强制断言C.dtype torch.bfloat16不满足即无法运行DeepGEMM 被禁用环境变量SGLANG_ENABLE_JIT_DEEPGEMM0GPU 计算能力低于 SM90DeepGEMM 要求 SM90即 Hopper 及以上的 H100/H200 等需要跨平台兼容性。Triton 内核同时支持 NVIDIA 与 AMD GPU在 ROCm 平台上如 MI300X/MI325XFP8 数据类型会自动切换为torch.float8_e4m3fnuz见 tuning_block_wise_kernel.py。优先使用 DeepGEMM 的场景输出 dtype 为bfloat16且 DeepGEMM 已启用运行在计算能力 ≥ SM90 的 NVIDIA GPU 上如 H100、H200生产负载追求极致性能——DeepGEMM 针对 Hopper 架构做了深度优化。注意DeepGEMM 要求 CUDA compute capability ≥ 9.0SM90是专门面向 NVIDIA Hopper GPUH100/H200优化的实现。从源码结构看SGLang 的自动选择逻辑位于 fp8_utils.py 的_dispatch_auto_backend中优先 DeepGEMM其次按 Blackwell FlashInfer、SM120 CUTLASS、AMD AITER、NPU 的顺序尝试最终回退到triton_w8a8_block_fp8_lineardeepgemm_w8a8_block_fp8_linear_with_fallback中dtype_supported output_dtype torch.bfloat16与shape_supported weight.shape[0] % 64 0 and weight.shape[1] % 128 0任一不满足时也会自动回退到 Triton 路径。因此Triton 内核在非 Hopper、非 bf16 输出、跨平台三类场景下几乎是唯一选择也是调优脚本的重点对象。二、快速开始调优脚本 tuning_block_wise_kernel.py 会针对给定的权重形状 (N, K) 与一组 batch size在多块 GPU 上并行穷举 Triton 内核的 tile 配置并选出每档 batch 的最优参数。默认模式DeepSeek-V3python benchmark/kernels/quantization/tuning_block_wise_kernel.py --tp-size 8未指定--N/--K时脚本调用get_weight_shapes(tp_size)tuning_block_wise_kernel.py加载 DeepSeek-V3 的预设权重形状集合并按张量并行度切分可切分的维度N 可 TP 与 K 可 TP 的形状会分别除以tp_size。自定义模型指定 N 和 Kpython benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 25600前提条件机器上至少有一块可用 GPU脚本在main中通过get_device_count()检查为 0 时直接抛出RuntimeError同时需要安装sglang导入sglang.kernels.ops.quantization.fp8_kernel与triton、torch、tqdm。脚本要求--N与--K必须同时给出或同时不给否则 argparse 会报错终止。三、命令行参数详解脚本完整参数定义见 tuning_block_wise_kernel.py参数说明默认值--tp-size/-tp张量并行度仅用于 DeepSeek-V3 预设形状指定--N/--K时被忽略8--N权重矩阵输出维度output dim即列数None须与--K同时给出--K权重矩阵输入维度input dim即行数None须与--N同时给出--input-type量化类型fp8或int8fp8--out-dtype输出数据类型可选float32/float16/bfloat16/halffloat16--block-n权重量化块的行粒度block_n128--block-k权重量化块的列粒度block_k128--batch-sizes自定义 batch size 列表可变参数不传则使用内置 18 档见下--save-path最优配置的保存目录python/sglang/kernels/ops/quantization/configs补充说明几点影响实际效果的关键行为默认 batch sizes为[1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536, 2048, 3072, 4096]tuning_block_wise_kernel.py覆盖了从单请求到大规模并发推理的典型负载候选搜索空间由get_configs_compute_bound()生成tuning_block_wise_kernel.pyNVIDIA 上遍历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}、num_stages ∈ {2,3,4,5}的组合ROCm 平台使用独立的get_rocm_configs_compute_bound()配置集。搜索前还会按block_k % BLOCK_SIZE_K 0过滤掉与量化块粒度不匹配的候选对每个候选配置脚本先做 5 次 warmup触发 JIT 编译再用 CUDA/HIP 事件计时 10 次取平均tuning_block_wise_kernel.py编译失败OutOfResources的配置会被跳过多 GPU 并行batch sizes 被打乱后均分到所有 GPU每个 GPU 处理全部权重形状跨进程通过multiprocessing的spawn上下文与Manager().Lock()保证并发写 JSON 不冲突tuning_block_wise_kernel.py脚本还内置了对 AMD GPU 的适配当 FP8 内核的网格规模小于等于设备计算单元数时会切换到手写展开 4 次的_w8a8_block_fp8_matmul_unrolledx4内核以获得更优性能tuning_block_wise_kernel.py。四、如何计算 N 和 K对线性层y xW^T其中x形状为 (M, K)、W形状为 (N, K)N输出特征数权重矩阵的输出维度K输入特征数权重矩阵的输入维度。若不清楚模型各层维度可从模型配置hidden_size、intermediate_size、注意力头数等手工推导。示例Qwen3-VL-32B设hidden_size5120、intermediate_size25600、num_heads64、num_kv_heads8、head_dim128TP1# QKV projection: Q(8192) K(1024) V(1024) 10240 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 10240 --K 5120 # MLP gateup (SwiGLU): 2 * intermediate_size 51200 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 51200 --K 5120 # MLP down projection python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 25600 # O projection (if separate from QKV) python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 8192其中 Q 头数为 64Q 维度 64 × 128 8192KV 头数为 8K/V 维度各 8 × 128 1024。若 TP8N 维被切分K 维保持不变# QKV projection: Q(8192) K(1024) V(1024) 10240 / TP8 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 1280 --K 5120 # MLP gateup (SwiGLU): 2 * intermediate_size 51200 / TP8 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 6400 --K 5120 # MLP down projection python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 3200 # O projection (if separate from QKV) python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 1024五、调优输出配置文件格式与运行时加载5.1 输出文件命名调优完成后脚本为每个 (N, K, block_shape, input_type) 组合生成一个 JSON 配置文件写入--save-path指定的目录。命名格式为N{N},K{K},device_name{DEVICE},dtypefp8_w8a8,block_shape[128,128].json例如当前仓库 python/sglang/kernels/ops/quantization/configs 目录下已收录 160 个面向不同 GPU 型号的预调优配置文件名如N1280,K5120,device_nameNVIDIA_B200,dtypefp8_w8a8,block_shape[128, 128].json。device_name由get_device_name().replace( , _)得到因此同一形状在不同 GPUA100、H100、H200、B200、MI300X 等上的最优配置互不干扰。5.2 JSON 内容结构文件内容是一个以 batch size字符串形式的 M为键、内核 tile 参数为值的映射把每档 batch size 映射到最优内核配置{ 1: {BLOCK_SIZE_M: 16, BLOCK_SIZE_N: 64, BLOCK_SIZE_K: 128, ...}, 2048: {BLOCK_SIZE_M: 128, BLOCK_SIZE_N: 128, BLOCK_SIZE_K: 128, ...} }仓库中真实文件示例NVIDIA B200、N1280/K5120为{ 2048: { BLOCK_SIZE_M: 64, BLOCK_SIZE_N: 128, BLOCK_SIZE_K: 128, GROUP_SIZE_M: 1, num_warps: 4, num_stages: 4 }, 3072: { BLOCK_SIZE_M: 64, BLOCK_SIZE_N: 128, BLOCK_SIZE_K: 128, GROUP_SIZE_M: 64, num_warps: 4, num_stages: 4 }, 4096: { BLOCK_SIZE_M: 64, BLOCK_SIZE_N: 128, BLOCK_SIZE_K: 128, GROUP_SIZE_M: 64, num_warps: 3, num_stages: 3 } }可见不同 batch size 下最优的GROUP_SIZE_M、num_stages会有差异——这正是按 batch 分档调优的意义所在。若 JSON 文件已存在脚本会读取旧配置、用新结果覆盖/合并再写回按 batch size 排序后整体落盘见 tuning_block_wise_kernel.py便于多次增量调优。5.3 运行时的加载与使用Triton 内核入口w8a8_block_fp8_matmul_triton在每次 GEMM 时调用get_w8a8_block_fp8_configs(N, K, block_n, block_k)fp8_kernel.py按N{N},K{K},device_name{DEVICE},dtypefp8_w8a8,block_shape[{block_n}, {block_k}].json拼接文件名在当前configs目录与fp8_kernel.py同级的configs/子目录查找配置文件找到则读取为{int(key): val}映射并对其中BLOCK_SIZE_K block_k且不满足整除关系的条目做BLOCK_SIZE_K钳位修正因为 scale stepping 要求BLOCK_SIZE_K block_k随后返回未找到则打印告警 Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! 并返回None运行时回退到默认配置{BLOCK_SIZE_M: 64, BLOCK_SIZE_N: block_n, BLOCK_SIZE_K: block_k, GROUP_SIZE_M: 32, num_warps: 4, num_stages: 3}fp8_kernel.py拿到配置映射后对当前实际 Mbatch size取min(configs.keys(), keylambda x: abs(x - M))即选择距离最近的一档 batch 配置fp8_kernel.py。注意路径细节README 中描述的保存目录为python/sglang/srt/layers/quantization/configs/而当前仓库脚本源码中--save-path的实际默认值以及运行时读取目录均为python/sglang/kernels/ops/quantization/configs见 tuning_block_wise_kernel.py 与 fp8_kernel.py。因此请使用脚本的默认保存路径确保生成的配置能被运行时直接发现若手动指定--save-path请保证最终文件落在内核模块的configs/目录下否则将走默认配置并出现性能告警。另外get_w8a8_block_fp8_configs在torch.compile图编译期间会直接返回None跳过文件查找避免 host 端非张量操作此时同样使用默认配置。六、调优背后的原理块量化 GEMM 的输入输出约定理解调优脚本内部逻辑有助于正确使用它。tune()函数tuning_block_wise_kernel.py的核心流程是用torch.rand生成服从均匀分布的 FP32 激活 A (M, K) 与权重 B (N, K)再 clamp 到类型上下界后量化为 FP8NVIDIA 用float8_e4m3fnROCm 用float8_e4m3fnuz或 INT8依据block_n/block_k计算 tile 数n_tiles ceil(N / block_n)、k_tiles ceil(K / block_k)生成形状为 (M, k_tiles) 的激活 scaleAs与形状为 (n_tiles, k_tiles) 的权重 scaleBs放大系数 1e-2。这验证了块量化 scale 张量的布局激活按 token 组粒度、权重按块粒度对搜索空间中的每个配置调用benchmark_config计时记录最优结果。在 fp8_utils.py 的triton_w8a8_block_fp8_linear生产路径中真实激活会先经per_token_group_quant_fp8(input_2d, block_size[1], ...)做 per-token-group 量化组大小即block_k再交给w8a8_block_fp8_matmul_triton执行。这印证了调优脚本中 scale 形状与量化粒度的设定与运行时完全一致——调优结果的可用性由此得到保证。七、验证与进阶建议验证调优结果SGLang 的量化模块测试如 test/registered/kernels 与 test/registered/unit 下的相关用例覆盖了 W8A8 块量化内核的正确性为自定义模型调优后建议先在离线引擎上对比 FP8 与 BF16 的生成质量再投入线上。多模型增量调优脚本支持合并写入可先跑 DeepSeek-V3 默认形状再追加自定义模型的 (N, K)同一 GPU 同一形状的配置会被保留无需清空目录。AMD 平台ROCm 上input_typefp8会使用e4m3fnuz格式搜索空间与内核选择含 unrolledx4 变体均由_is_hip自动适配无需额外参数。批量负载适配若你的线上 batch 分布集中在某一区间如长序列推理下 M 偏小可通过--batch-sizes 1 2 4 8 16 32 64自定义调优档位缩小搜索范围、加快调优速度同时让最近邻查找更贴合实际负载。总而言之tuning_block_wise_kernel.py是 SGLang 为 W8A8 块量化 Triton 内核准备的一键式硬件适配工具它把内核 tile 搜索、批量计时、多 GPU 并行与 JSON 落盘整合为一条命令产出的配置由 fp8_kernel.py 在每次 GEMM 时按最近邻 batch 自动加载。结合本文的 N/K 推导方法与 DeepGEMM 选型判断你可以为自己的 GPU 与模型组合获得接近硬件上限的 Triton 块量化 GEMM 性能。【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考