Skip to content

feat(quant): add triton fp8 w8a8 per-tensor quantization method#1404

Open
sufubao wants to merge 12 commits into
ModelTC:mainfrom
sufubao:support_fp8_pt
Open

feat(quant): add triton fp8 w8a8 per-tensor quantization method#1404
sufubao wants to merge 12 commits into
ModelTC:mainfrom
sufubao:support_fp8_pt

Conversation

@sufubao

@sufubao sufubao commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

What

Add triton-fp8w8a8-pertensor (aliases fp8w8a8-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.py
    • New TritonFP8w8a8PerTensorQuantizationMethod: per-tensor weight quant (_fp8_per_tensor_quant), per-token activation quant via the pure-Triton lightllm_per_token_group_quant_fp8 (not the wrapper, which would dispatch to the sgl kernel that rejects group_size == k).
    • Deferred staging for fused multi-split weights and per-expert MoE weights (_try_finalize_deferred_weight).
    • fp8 GEMM backend selector (LIGHTLLM_FP8_GEMM).
  • lightllm/common/basemodel/triton_kernel/quantization/scaled_mm_per_token_kernel.py: B_SCALE_IS_TENSOR path 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_TENSOR support in the fused MoE matmul.
  • lightllm/server/api_cli.py: list triton-fp8w8a8-pertensor in the --quant_method help.

Usage

--quant_method triton-fp8w8a8-pertensor
# optional GEMM backend, default auto
export LIGHTLLM_FP8_GEMM=auto   # auto | cutlass | sgl | triton

auto picks cutlass (vLLM cutlass_scaled_mm) → sgl (sgl_kernel.fp8_scaled_mm) → triton, by availability, so images without a Cutlass-class kernel fall back to Triton cleanly. sgl_kernel rejects a per-tensor weight scale, so the scalar is expanded to per-channel [N] once and cached on the weight_pack.

Validation (H200, Qwen3.5-27B)

  • End-to-end fp8 throughput ≈ 1.19x over bf16 (matches vLLM ~1.20x; the pure-triton GEMM alone capped at ~1.10x because its grid underfills the GPU on the attention-projection shapes at decode batch sizes).
  • GSM8K accuracy by backend: cutlass 0.850 / sgl 0.860 / triton 0.825 / bf16 0.860 (sgl matches bf16; cutlass/triton within ~0.01–0.035).
  • Correctness: per-token scale exact vs 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.

sufubao added 3 commits July 24, 2026 22:42
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.
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

shihaobai and others added 9 commits July 24, 2026 23:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants