feat(quant): add triton fp8 w8a8 per-tensor quantization method#1404
Open
sufubao wants to merge 12 commits into
Open
feat(quant): add triton fp8 w8a8 per-tensor quantization method#1404sufubao wants to merge 12 commits into
sufubao wants to merge 12 commits into
Conversation
Add TritonFP8w8a8PerTensorQuantizationMethod (triton-fp8w8a8-pertensor/-pt) with deferred staging for fused multi-split and per-expert weights. Support per-tensor weight scale in the grouped MoE matmul (WEIGHT_SCALE_PER_TENSOR) and the per-token scaled_mm kernel (B_SCALE_IS_TENSOR).
TritonFP8w8a8PerTensorQuantizationMethod.apply() quantized the activation with vllm's scaled_fp8_quant, which is only bound under `if HAS_VLLM`. On images without vllm this raised `NameError: name 'scaled_fp8_quant' is not defined` at first prefill, crashing model init (router saw only EOFError). Use the pure-triton per-token kernel instead. Call lightllm_per_token_group_quant_fp8 directly rather than the per_token_group_quant_fp8 wrapper: the wrapper dispatches to the sgl kernel when present, which rejects group_size == k (per-token). group_size == k matches vllm's use_per_token_if_dynamic semantics. Verified on H200: per-token scale exact vs max(|x|)/448; mm cosine 0.9993 vs bf16 ref (> repo's 0.99 gate) across M=1/16/47/128/7 and non-aligned K/N.
The triton scaled_mm_per_token kernel underperforms on the attention-projection shapes (qkv N=8192, o N=5120,K=6144) at decode batch sizes: it runs ~0.4-0.6x of vLLM's Cutlass fp8 scaled-mm and even slower than bf16, because its grid (cdiv(M,BM)*cdiv(N,BN), no split-K) can't fill the GPU at small N + small M. A split-K prototype confirmed the diagnosis but never matched Cutlass, which holds a ~2x structural lead. This capped end-to-end fp8 speedup at ~1.10x over bf16 vs vLLM's ~1.20x on Qwen3.5-27B / H200. Select the fp8 GEMM backend via LIGHTLLM_FP8_GEMM (auto|cutlass|sgl|triton); auto prefers vllm cutlass_scaled_mm, then sgl_kernel.fp8_scaled_mm, then the triton kernel, so images without a Cutlass-class kernel fall back cleanly. Lifts end-to-end fp8 throughput to ~1.19x over bf16 (matches vLLM); GSM8K lossless (cutlass 0.850 / sgl 0.860 / triton 0.825 / bf16 0.860). sgl_kernel rejects a per-tensor weight scale, so the scalar is expanded to per-channel [N] and cached on the weight_pack.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
shihaobai
force-pushed
the
support_fp8_pt
branch
from
July 24, 2026 15:33
847dcef to
7157d5b
Compare
only extend scaled_mm_per_token static key for the per-tensor case so existing per-channel tuned configs still match; validate explicit LIGHTLLM_FP8_GEMM=cutlass|sgl against availability at import instead of crashing at first inference.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Add
triton-fp8w8a8-pertensor(aliasesfp8w8a8-pertensor,triton-fp8w8a8-pt,fp8w8a8-pt) — an FP8 w8a8 quantization method with no vLLM dependency: weight is per-tensor FP8, activation is per-token FP8, both via Triton kernels. The fp8 GEMM dispatches to the best available backend.新增
triton-fp8w8a8-pertensor(别名fp8w8a8-pertensor/triton-fp8w8a8-pt/fp8w8a8-pt):不依赖 vLLM 的 FP8 w8a8 量化方法,权重 per-tensor、激活 per-token,全部走 Triton kernel;fp8 GEMM 自动选择最优后端。Why
The existing
vllm-fp8w8a8*methods require vLLM. On images without vLLM (or when a pure-Triton path is preferred) there was no fp8 w8a8 option. This fills that gap and reaches vLLM-comparable end-to-end fp8 throughput by routing the GEMM to a Cutlass-class kernel when available.现有
vllm-fp8w8a8*强依赖 vLLM;无 vLLM 的镜像(或希望走纯 Triton 路径)下缺少 fp8 w8a8 方案。本方法补上该能力,并在存在 Cutlass 级 kernel 时把 GEMM 路由过去,端到端吞吐与 vLLM 持平。Changes
lightllm/common/quantization/w8a8.pyTritonFP8w8a8PerTensorQuantizationMethod: per-tensor weight quant (_fp8_per_tensor_quant), per-token activation quant via the pure-Tritonlightllm_per_token_group_quant_fp8(not the wrapper, which would dispatch to the sgl kernel that rejectsgroup_size == k)._try_finalize_deferred_weight).LIGHTLLM_FP8_GEMM).lightllm/common/basemodel/triton_kernel/quantization/scaled_mm_per_token_kernel.py:B_SCALE_IS_TENSORpath so the per-token scaled_mm accepts a per-tensor weight scale[1], not only per-channel[N].lightllm/common/basemodel/layer_infer/fused_moe/grouped_fused_moe.py:WEIGHT_SCALE_PER_TENSORsupport in the fused MoE matmul.lightllm/server/api_cli.py: listtriton-fp8w8a8-pertensorin the--quant_methodhelp.Usage
autopickscutlass(vLLMcutlass_scaled_mm) →sgl(sgl_kernel.fp8_scaled_mm) →triton, by availability, so images without a Cutlass-class kernel fall back to Triton cleanly.sgl_kernelrejects a per-tensor weight scale, so the scalar is expanded to per-channel[N]once and cached on theweight_pack.Validation (H200, Qwen3.5-27B)
max(|x|)/448; matmul cosine 0.9993 vs bf16 ref (above the repo's 0.99 gate) across M = 1/16/47/128/7 and non-aligned K/N.