[JIT] per_token_group_quant: fix the fp16 UE8M0 multiplier overflow, 1.23x on the fused EP-MoE path, unify the scale allocation - #33533
Open
DarkSharpness wants to merge 4 commits into
Conversation
…drop __frcp_rn from silu Two independent fixes to the trait-driven quant kernel, both traced back to sgl-project#30924 and both verified on H200 (sm_90a). 1. fp16 + UE8M0 quantized whole groups to +-448. The power-of-two quant multiplier was narrowed to the input dtype so the scaling could be one packed __hmul2. That multiplier reaches kMaxValue / eps = 4.5e12, which bfloat16 holds (it has fp32's exponent range) but float16 does not (max 65504). Any fp16 group whose absmax fell below 448/65504 = 6.8e-3 therefore got an `inf` multiplier, and every element in the group quantized to +-448 -- an all-zero row went through `0 * inf` to NaN, which the sanitizing clamp also turned into +448. Measured 4095/4096 codes wrong, 100% of them +-448, silently. The condition is now derived at compile time from DTypeTrait<T>::kFloatMax and a newly-named kAmaxFloor, keeping the packed half multiply where it is provably safe and scaling in fp32 otherwise. The bf16 production path is codegen-identical: SASS for the bf16/ue8m0/g128 col-major flat kernel is byte-for-byte unchanged at 200 instructions. fp16 + ue8m0 goes 128 -> 160 instructions, the cost of correctness on a path no bf16 model uses. Regression test verified to fail before the fix (3 fp16 cases red, 3 bf16 green) and pass after. 2. details::silu used __frcp_rn where the v2 kernel used a plain divide. Under --use_fast_math `a / b` is MUFU.RCP + FMUL (2 instructions); __frcp_rn is a Newton refinement (MUFU.RCP + 2 FFMA + FADD) behind a branch into an out-of-line special-value fixup (~10 instructions + divergence). silu runs per ELEMENT, not once per group like the quant multiplier, so this dominated the fused path: the bf16 fused masked kernel goes 608 -> 384 SASS instructions and BSSY/CALL 35 -> 2, worth 1.23x geomean (1.13-1.26x) on graph-captured EP-MoE decode shapes. It also restores bit-exactness with the AOT v2 op, which uses the same divide. Non-fused kernels are unaffected (silu is only instantiated for kFuseSiluAndMul). Note sgl-project#32616 made exactly this substitution back for the quant multiplier but missed silu; the multiplier is once per group, so it was the smaller half. Verified: 781 passed across the quantization + moe kernel suites; the current kernel is bitwise equal to the pre-sgl-project#30924 v2 kernel across 26 combinations (fp32/ue8m0/int8 x row/col-major x bf16/fp16 x plain/fused/masked).
…scatter fusion minimax/per_token_quant_ue8m0.cuh held two independent kernels. The non-scatter one was a pure duplicate of the shared per_token_group_quant row-packed UE8M0 path -- measured byte-identical across 8 shape x group-size combinations -- and had no production caller; only tests used it, as the reference for the scatter variant. Delete it and point those tests at the shared kernel. That also strengthens the scatter test: its reference is no longer a sibling kernel in the same file (where a shared mistake would cancel out) but the kernel that has its own bit-exact coverage. The scatter variant stays. What it adds over the shared kernel is the output mapping -- one source token replicated to its topk destination rows, with the scale byte-scattered into the MN-major grouped-GEMM buffer -- which eliminates a fill_gateup_input_triton_kernel launch and the intermediate buffers. Its module comment now says that, instead of referring to the deleted kernel. Verified: the edited .cuh still compiles for the scatter instantiation; 781 passed across the quantization + moe kernel suites. test_minimax_quant_scatter itself is SM100-gated and skips on H200, so it still needs a B200 run.
…entry points
On CUDA `eps` had already stopped being a knob: the JIT kernel bakes the
group-absmax floor in at compile time (QuantTrait::kAmaxFloor), so
_run_per_token_group_quant_8bit_kernel could only assert `eps == 1e-10` or raise,
and sglang_per_token_group_quant_int8 carried the same assert. Its only remaining
jobs were being forwarded to the MUSA AOT op and to the deprecated v2 JIT kernel,
both of which still take it as a real runtime argument.
An AST sweep of every call site repo-wide (python/sglang, test, benchmark; 59
calls across the whole quant family, resolving positional as well as keyword
arguments) found no caller anywhere that passes a non-default eps: the 13 that
pass it explicitly are 4 internal forwards inside fp8_kernel.py plus 9 v2
tests/benchmarks passing the literal 1e-10 because v2's parameter has no default.
Remove it from the five entry points that reach the JIT kernel
(sglang_per_token_group_quant_fp8 / _row_padded / _ue8m0 /
sglang_per_token_group_quant_8bit / sglang_per_token_group_quant_int8) and from
the dispatcher, and introduce PER_TOKEN_GROUP_QUANT_EPS for the two internal
forwards that still need the value. Behaviour is unchanged.
Left alone deliberately: the Triton per_token_group_quant_{fp8,int8,8bit} family,
per_token_group_quant_8bit_v2 and the AOT sgl_per_token_group_quant_8bit (faithful
ports whose eps is a live runtime argument), and the MLA quants that use 1e-12.
Removing a positional parameter shifts the ones after it, so that was checked
separately: no call site passes more than two positional arguments (x,
group_size), well before eps' old index. One caller could not be seen by AST at
all -- the AOT test built a single `execute_kwargs` dict containing `eps=1e-10`
and `**`-unpacked it into both the Triton reference and the sglang entry point.
Confirmed by experiment that this really did break (TypeError: unexpected keyword
argument 'eps'); it now keeps eps only for the Triton call. A repo-wide scan for
`**`-unpacked calls into any of the six reshaped functions found that this was
the only such site.
Verified: the AOT suite that owns that caller passes in full (1760 passed, 1872
skipped -- the skips are the Blackwell-only ue8m0 configs); the five eps-relevant
kernel test files pass (5 passed, 2 skipped, 160 subtests); all eight reshaped
entry points were smoke-called on device; and test/manual/dsv4/test_wo_a_fp8_sm90.py
passes end to end on sm90.
…dup the entry points
Four places decided what a group quant allocates, and none of them wrote down
which scale format they wanted -- it fell out of whichever allocator you happened
to call:
create_per_token_group_quant_fp8_output_scale row-major UE8M0 -> fp32 2^(e-127)
per_token_group_quant._allocate_outputs row-major UE8M0 -> int32 packed
sglang_per_token_group_quant_fp8_ue8m0 inlined the int32 TMA branch
sglang_per_token_group_quant_fp8_row_padded inlined the col-major fp32 branch
The first two disagree for identical flags. That is the root cause behind several
recent surprises: test_fp8_wo_a depends on the fp32 flavor purely through a
defaulted column_major_scales, and the deprecated v2 fallback exists because
"layout is row-major" was being used to imply "format is fp32 pow-2".
Collapse them into quant_format.create_group_quant_{scale,outputs}. Being a leaf
module matters: fp8_kernel imports per_token_group_quant at module level, so
per_token_group_quant could only reach the allocation through a lazy import to
break the cycle. Both sides -- and dsv4/moe.py -- now import it normally, which is
why is_hip() is inlined as `torch.version.hip is not None` (exactly what
sglang.srt.utils.is_hip does) and ceil_align as two lines.
Unifying forces the ambiguity to be named, hence pack_ue8m0 (read only under
scale_ue8m0; without pow-2 rounding there is no single byte to pack). Rounding,
storage and layout are now three parameters instead of two booleans plus whichever
function you reached for. Behaviour is unchanged: the old create_ contract was
exactly "packed iff column-major", so every migrated site passes pack_ue8m0 equal
to its column_major_scales expression, and per_token_group_quant's auto-allocation
passes pack_ue8m0=True to keep _allocate_outputs' int32 row-packed result.
The module also owns is_fp8_fnuz / fp8_dtype / fp8_max / fp8_min, moved out of
fp8_kernel and re-exported from there so the ~60 existing callers are untouched.
That is not bookkeeping: it fixes a real bug this refactor would otherwise have
introduced. An earlier draft hardcoded torch.float8_e4m3fn in dsv4/moe.py on the
theory that the kernel is CUDA-only. It is not -- load_jit strips --use_fast_math
under torch.version.hip, utils.cuh maps fp8_e4m3_t to uint8_t on ROCm, and
type.cuh's kFP8E4M3Max is 224 under HIP_FP8_TYPE_FNUZ, which
silu_and_mul_masked_post_quant.cuh consumes via math::FP8_E4M3_MAX. So on gfx94x
the kernel writes fnuz codes and a hardcoded e4m3fn buffer would silently
disagree with them. Allocation now goes through fp8_dtype, which follows the same
split.
Entry-point dedup on top of that:
- Delete sglang_per_token_group_quant_fp8_ue8m0. No callers anywhere (not in
__all__, not in docs), and step-for-step identical to
sglang_per_token_group_quant_fp8 with column_major_scales / scale_tma_aligned /
scale_ue8m0 all True -- the whole-row shortcut cannot fire once scale_ue8m0 is
set.
- sglang_per_token_group_quant_fp8_row_padded keeps its reason to exist (row
padding + tail zeroing, which the general entry point cannot express) but stops
hand-rolling the col-major fp32 buffer.
- silu_and_mul_contig_post_quant allocates its own outputs, so deep_gemm.py and
deepseek_v2.py no longer hand-allocate a codes buffer plus a scale buffer each
(~10 duplicated lines apiece). Its output/output_scale parameters are gone
rather than optional, since no caller passed them. That removes the last two
production callers of the allocator from outside sglang.kernels.ops: scale
layout is knowledge the kernel wrappers own, not something model code
re-derives.
- Five call sites that allocated codes and scales side by side now make one call.
The three that fed torch.zeros keep zeroing explicitly -- test_8bit_v2's _alloc
documents why (unwritten padding has to compare equal). _alloc_scale in
test_per_token_group_quant stays on the scale-only entry point: it is the one
caller that genuinely wants just a scale, because its tests build codes buffers
of differing dtypes themselves.
Also shortened the name -- the old one said "fp8" though the int8 path uses it for
its fp32 scales.
Verified on H200: create_group_quant_scale reproduces both old allocators exactly
(shape/stride/dtype/storage_offset over 64 flag combinations including the assert
cases, plus the int32 row-packed branch over 10 shapes), and
create_group_quant_outputs' scale matches it over the same 64 with correct codes
shape/dtype; auto-allocating silu_and_mul_contig_post_quant is bit-identical to
the hand-allocated form over 5 shape/flag combinations; no import cycles, and
per_token_group_quant is still the package's function rather than the shadowing
submodule. Failure set across the quantization + quant + moe + attention + AOT
suites is identical to main (5 failed / 19 errors, all model- or
library-dependent), with 2572 passed / 1979 skipped unchanged from before this
refactor.
DarkSharpness
requested review from
BBuf,
Edwardf0t1,
FlamingoPg,
Fridge003,
HaiShaw,
HydraQYH,
Ying1123,
celve,
ch-wan,
fzyzcjy,
ispobock,
merrymercy,
yizhang2077 and
yuan-luo
as code owners
August 4, 2026 10:08
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
DarkSharpness
force-pushed
the
ptgq-quant-fixes
branch
from
August 4, 2026 11:24
34b0e9b to
f5fff06
Compare
41 tasks
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
2 participants
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.
Motivation
Two silent-correctness bugs in the trait-driven
per_token_group_quantkernel, a1.23x left on the table in the EP-MoE fused path, and the duplication that let all
of it hide.
Both bugs trace back to #30924. They surfaced while trying to reproduce the
GLM-5.2 H20 decode regression (#32582): #32616 reverted that PR's
kMaxValue * __frcp_rn(amax)back to a plain divide, noting "the exact mechanismis still unclear".
The mechanism is fp8 round-to-nearest tie-breaking, and the reverted code was not
the accuracy problem. Under
--use_fast_math,448.0f / amaxgives a multiplier1 ULP high, so products land just above an fp8 tie (168.0000153 → rounds up to
176);
__frcp_rngives the correctly-rounded multiplier, so products land on thetie (168.0) and round half-to-even (160). Only ~0.01% of codes on randn bf16
differ, and the mean relative dequant error is identical to six digits. So #32616
restored bit-exactness with the legacy v2 kernel, not accuracy — and it was right
to, because
__frcp_rnis also ~5x the instruction cost: 2 SASS instructions(
MUFU.RCP+FMUL.FTZ) versus a Newton refinement behind a branch into anout-of-line special-value fixup. This PR does not re-land it.
But #32616 only fixed the quant multiplier. #30924 made the same substitution in
details::silu, which runs per element rather than once per group.Modifications
1. fp16 + UE8M0 quantized whole groups to ±448
The power-of-two quant multiplier was narrowed to the input dtype so the scaling
could be one packed
__hmul2. That multiplier reacheskMaxValue / eps= 4.5e12,which bfloat16 holds (it has fp32's exponent range) but float16 does not (max
65504). Any fp16 group whose absmax fell below 448/65504 = 6.8e-3 therefore got an
infmultiplier and every element in the group quantized to ±448 — an all-zero rowwent through
0 * infto NaN, which the sanitizing clamp also turned into +448.Measured 4095/4096 codes wrong, 100% of them ±448, silently. The old v2 kernel
handles the same input correctly. Not covered before because the fp16 UE8M0 cases
all used unit-variance
randn, whose absmax sits ~400x above the threshold.The condition is now derived at compile time from
DTypeTrait<T>::kFloatMaxand anewly-named
kAmaxFloor, keeping the packed half multiply where it is provablysafe and scaling in fp32 otherwise. The bf16 production path is
codegen-identical — SASS for the bf16/ue8m0/g128 col-major flat kernel is
byte-for-byte unchanged at 200 instructions. fp16 + ue8m0 goes 128 → 160
instructions, the cost of correctness on a path no bf16 model uses.
2.
details::siluused__frcp_rnwhere v2 used a plain divide — 1.23xBecause silu runs per element, this dominated the fused path. Switching it back on
the non-Blackwell branch: the bf16 fused masked kernel goes 608 → 384 SASS
instructions,
BSSY/CALL35 → 2, worth 1.23x geomean (1.13–1.26x) ongraph-captured EP-MoE decode shapes. It also restores bit-exactness with the AOT v2
op, which uses the same divide. Non-fused kernels are unaffected (silu is only
instantiated for
kFuseSiluAndMul).3. Dedup: one allocator, one leaf module, three fewer entry points
Four places decided what a group quant allocates, and none wrote down which scale
format they wanted — it fell out of whichever allocator you happened to call:
create_per_token_group_quant_fp8_output_scale2^(e-127)per_token_group_quant._allocate_outputssglang_per_token_group_quant_fp8_ue8m0sglang_per_token_group_quant_fp8_row_paddedThe first two disagree for identical flags. That is the root cause behind
several surprises here:
test_fp8_wo_adepends on the fp32 flavor purely through adefaulted
column_major_scales, and the deprecated v2 fallback exists because"layout is row-major" was being used to imply "format is fp32 pow-2".
quant_format.create_group_quant_{scale,outputs}. Being a leafmodule matters:
fp8_kernelimportsper_token_group_quantat module level, soper_token_group_quantcould only reach the allocation through a lazy import tobreak the cycle. Both sides now import it normally.
pack_ue8m0. Rounding, storageand layout are three parameters instead of two booleans plus whichever function
you reached for. Behaviour is unchanged: the old contract was exactly "packed
iff column-major", so every migrated site passes
pack_ue8m0equal to itscolumn_major_scalesexpression.is_fp8_fnuz/fp8_dtype/fp8_max/fp8_minmoved into the same leaf andre-exported from
fp8_kernel, so the ~60 existing callers are untouched. Thisalso fixes a bug the refactor would otherwise have introduced: an earlier
draft hardcoded
torch.float8_e4m3fnon the theory that the dsv4 kernel isCUDA-only. It is not —
load_jitstrips--use_fast_mathundertorch.version.hip,utils.cuhmapsfp8_e4m3_ttouint8_ton ROCm, andtype.cuh'skFP8E4M3Maxis 224 underHIP_FP8_TYPE_FNUZ, which the kernelconsumes via
math::FP8_E4M3_MAX. So on gfx94x it writes fnuz codes and ahardcoded e4m3fn buffer would silently disagree.
sglang_per_token_group_quant_fp8_ue8m0— no callers anywhere, andstep-for-step identical to
sglang_per_token_group_quant_fp8withcolumn_major_scales/scale_tma_aligned/scale_ue8m0all True.minimax/per_token_quant_ue8m0kernel —measured byte-identical to the shared row-packed UE8M0 path across 8 shape ×
group-size combinations, and used only as the scatter variant's test reference.
That reference is now the shared kernel, which has its own bit-exact coverage.
silu_and_mul_contig_post_quantallocates its own outputs, sodeep_gemm.pyanddeepseek_v2.pyno longer hand-roll a codes buffer plus a scale buffer each.Zero production callers of the allocator now live outside
sglang.kernels.ops— scale layout is knowledge the kernel wrappers own.4. Dropped the inert
epsparameterOn CUDA it had already stopped being a knob: the kernel bakes the absmax floor in
at compile time, so the dispatcher could only assert
eps == 1e-10or raise. AnAST sweep of all 59 call sites repo-wide (resolving positional as well as keyword
arguments) found no caller anywhere that passes a non-default eps. Removed from
the entry points;
PER_TOKEN_GROUP_QUANT_EPSremains for the MUSA AOT op and theTriton reference, whose
epsis a live runtime argument.Accuracy Tests
All on H200 (sm_90a). Bit-exactness is against the pre-#30924
per_token_group_quant_8bit_v2kernel.row/col-major × bf16/fp16 × plain / fused-silu / masked.
green — exactly the expected signature.
dtype / storage_offset over 64 flag combinations including the assert cases, plus
the int32 row-packed branch over 10 shapes.
silu_and_mul_contig_post_quantis bit-identical to thehand-allocated form over 5 shape/flag combinations.
test/registered/kernels/ops/quantization/,test/registered/quant/,test/registered/kernels/ops/moe/test_minimax_quant_scatter.py,test_kimi_k3_prerequisite_ops.py,test_fp8_wo_a.py, and the AOTtest_per_token_group_quant_8bit.py(3632 cases). Failure set is identical tomain(24 entries, all model- or library-dependent:auto-roundnotinstalled, GGUF/AWQ/kimi/w8a8/gptqmodel need weights and a server).
Speed Tests
Graph-captured (
torch.cuda.CUDAGraph) EP-MoE decode shapes, bf16, group 128,col-packed UE8M0, fused silu + masked — the
--moe-a2a-backend deepeppath.Measuring uncaptured would have been meaningless:
per_token_group_quantcarries~8us of Python dispatch against a 2–5us kernel.
geomean 1.227x. Rep spread 0.1–0.3%. Non-fused kernels unchanged (200
instructions, byte-identical SASS).
Blackwell verification
Two tests are gated on
get_device_sm() >= 100and only skip on H200. Both pass onB200:
test/registered/kernels/ops/attention/test_fp8_wo_a.pytest/registered/kernels/ops/moe/test_minimax_quant_scatter.pyThe first is the one that made me leave the v2 fallback in place: its
_flat_referenceasks forscale_ue8m0=Truewithcolumn_major_scalesdefaulted,which allocates fp32 and so needs the fp32-pow-2 format the shared kernel still
cannot produce. Adding that format (and retiring the fallback) is prepared but
deliberately not in this PR — it is a separate change with its own tradeoff.
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): 🚫 Run #30904663212
Latest PR Test (Extra): ❌ Run #30904668323