Repository navigation
[Perf] per_token_group_quant_8bit: multiply by reciprocal instead of dividing per element - #61343
Merged
Merged
Conversation
gau-nernst
force-pushed
the
gau/fp8-quant-fastdiv
branch
from
October 13, 2026 05:44
b02d17b to
c04f719
Compare
…dividing per element Both the smem and register-resident schedules now scale with v * (1/y_s), computing one div.rn per 128-element group instead of one per element. Quantization is division-throughput-bound at large M; this is ~1.25x on GB300 (e.g. 101 -> 76 us at M=32768, K=4096; 2.34 -> 1.95 us at M=1). Numerics: scales are unchanged (bit-exact). The quantized values can differ by 1 fp8 ulp on ~0.05% of elements (reciprocal rounding vs division) -- the same class of numerics the UE8M0 path in this kernel and SGLang's production quant kernel already use. The two schedules use identical math so the bitwise register-vs-smem invariant (test_per_token_group_quant_fp8_register_ matches_smem) is preserved; UE8M0 behavior is unchanged (reciprocal of a power of two is exact). Co-authored-by: Kimi Code Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
gau-nernst
force-pushed
the
gau/fp8-quant-fastdiv
branch
from
October 13, 2026 05:46
c04f719 to
1aca0a2
Compare
gau-nernst
marked this pull request as ready for review
October 13, 2026 05:48
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94388 for commit |
Collaborator
CI selector (shadow): 41 test steps (80 jobs) instead of 195 (271 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Today's rules run everything for this diff. Selector would run (41)
Would skip (today's rules run them) (154)
Would add (today's rules do not run them) (0)none AMD mirrors: would skip (0)none AMD mirrors: would add (0)none 2 changed files · base |
zyongye
approved these changes
Oct 13, 2026
The fp8 group path divided per element while per_token_group_quant_8bit now multiplies by the reciprocal; the divergence broke test_fused_quant_layernorm's strict outlier budget and mixes numerics when a model uses both paths. Align on reciprocal-multiply via ScaledQuant<..., is_scale_inverted=true> (also drops 4 div.rn per vec4). Co-authored-by: Kimi Code Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
gau-nernst
force-pushed
the
gau/fp8-quant-fastdiv
branch
from
October 13, 2026 08:27
6669cac to
3995312
Compare
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94421 for commit |
AndreasKaratzas
pushed a commit
to AndreasKaratzas/vllm
that referenced
this pull request
Oct 16, 2026
…dividing per element (vllm-project#61343) Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>
michael-gpu
pushed a commit
to practical-inference/vllm-ampere-extended
that referenced
this pull request
Oct 20, 2026
…dividing per element (vllm-project#61343) Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Closed
3 of 5 tasks
faresobeid
added a commit
to PrimeIntellect-ai/prime-rl
that referenced
this pull request
Nov 3, 2026
…LM 0.33 vllm-project/vllm#61343 quantizes per-token groups as x * (1/s) instead of x / s, so the trainer's Triton casts (meant to be bit-identical to vLLM's) differed in 0.03-0.08% of fp8 values. Multiply by the exact reciprocal (div_rn(1, s)) on the non-ue8m0 path; test_fp8_utils passes again (25/25 on an RTX PRO 6000). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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.
Overview
Speed up FP8 per-token-group quantization by computing the per-group scale reciprocal once and multiplying, instead of dividing per element. This covers both CUDA group-quant kernels:
per_token_group_quant_8bit(smem and register-resident schedules) and the fused RMSNorm+group-quant kernel (layernorm_utils.cuh), so that all CUDA fp8 group-quant paths share the same reciprocal-multiply numerics. ~1.2-1.3x kernel speedup on GB300 across M=1..32768; ~1-2% e2e TPOT improvement on Qwen3.5-397B-A17B-FP8 (TP4).Claims
per_token_group_fp8_quantCUDA kernel on GB300 across M=1..32768 (details below).Validation
Correctness
pytest tests/kernels/quantization/test_per_token_group_quant.py: 848 passed, includingtest_per_token_group_quant_fp8_register_matches_smem(bitwise register-vs-smem equality; both schedules are changed identically, so the invariant is preserved).Microbenchmarks
GB300, K=4096, group=128, bf16 -> fp8e4m3 + fp32 scales. FlashInfer CUPTI with CUDA graph and cold L2; median of interleaved rounds. Both columns are standalone same-toolchain builds (nvcc 13.1) of the register-resident kernel; repro harness:
benchmarks/kernels/benchmark_per_token_group_quant.py.The kernel is memory-bound (~2 FLOPs/element; TFLOPS is not the meaningful metric here). This PR lifts achieved bandwidth from 4.0 to 5.2 TB/s at M=32768. A flat-index Triton variant we prototyped reaches 6.7 TB/s at M=32768 but loses at small M; modifying the existing CUDA kernel is the lower-risk change.
E2E perf
Qwen3.5-397B-A17B-FP8, TP4, 8k input / 1k output, vllm bench serve; median TPOT. main (
c077dc663d) and the branch (1aca0a263f) ran sequentially on the same GB300 node:Request throughput +1.2% at concurrency 64. Note: the e2e delta includes a toolchain effect (branch built from source with local nvcc 13.1; baseline is the nvcc 13.0 wheel, whose stock quant kernel is mis-scheduled -- see Details), so it mildly overstates the pure source-change effect. The same-toolchain kernel numbers above are the honest claim.
Accuracy
Qwen3.5-397B-A17B-FP8, TP4, identical server flags on the same GB300 node:
Mixed and within noise (binomial std error ~1.06pp per arm at n=1319), but GSM8K drifts down directionally; consistent with the ~0.05% of elements whose fp8 rounding flips under reciprocal-multiply.
Details
div.rn.f32is a multi-instruction slow sequence, and the kernel issued it once per element when one division per group suffices. The kernel's own UE8M0 path already uses reciprocal-multiply (a power-of-two scale makes the reciprocal exact); this PR applies the same transform to the fp32-scale path._per_token_group_quant_fp8), and SGLang's production quant kernel.layernorm_utils.cuh, the target ofRMSNormQuantFusionPassfor block-FP8 models, on both CUDA and ROCm via hipify) is switched to the same reciprocal-multiply. Without this, fused and unfused paths diverge by 1 fp8 ulp on ~0.05% of elements (caught bytest_fused_quant_layernormandtest_mamba_prefix_cache). The fused kernel also drops 4 div.rn per vec4. Per-token fp8 keeps exact per-element division for FBGemm parity; int8 paths are unchanged.Toolchain note: nvcc 13.0 x sm_100 mis-schedules the stock kernel
All official artifacts (release pipeline nvcc 13.0.2, Dockerfile CUDA_VERSION=13.0.3) target
... 10.0 ...(no sm_103). nvcc 13.0 targeting sm_100 mis-schedules this kernel's div.rn slow paths: 864 instructions / 3095 static stall cycles vs 568 / 1655 under nvcc 13.1; the shipped wheel runs the stock kernel at 167.6us vs 101.2us for a 13.1 rebuild at M=32768 (sm_100 SASS loads natively on sm_103; verified with CUDA_DISABLE_PTX_JIT=1). This PR's variant is insensitive to the compiler (78.3us on nvcc 13.0 vs 78.1us on 13.1) because it removes the per-element div.rn entirely. Worth a separate build-toolchain discussion; out of scope here.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing
AI assistance: Kimi Code helped develop, benchmark, and validate this change.