[Perf] Halve the non-finite sanitization overhead in per_token_group_quant - #32296
Merged
Merged
Conversation
…quant #32188 sanitized non-finite quant inputs with a two-sided clamp (fminf(fmaxf(v, -448), 448) / __hmin2(__hmax2(...))), adding 16 HMNMX2 (ue8m0 path) resp. 32 FMNMX (fp32-scale path) per bf16 group-of-128 on top of the bare-SATFINITE kernel. The lower clamp is redundant: IEEE min already returns the non-NaN operand (NaN / +inf -> +448), and the fp8 SATFINITE conversion saturates -inf to -448 on its own. Keep only the upper min, halving the added instructions, and hoist the PDL wait below the pure index arithmetic (codegen-neutral today; keeps the source honest about what the wait actually orders). Per-kernel SASS (bf16, gs=128, flat, sm_100a): ue8m0 144 -> 136 insns (20 -> 12 HMNMX2), fp32-scale 216 -> 200 (33 -> 17 FMNMX); sm_90a moves the same way. B200 microbench vs main: geomean 1.009x over 126 configs, up to 1.03x in the latency-bound small-batch range, no case slower. NaN now quantizes to +448 instead of -448 -- still finite, which is the property the sanitization tests pin. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DarkSharpness
requested review from
BBuf,
HydraQYH,
celve and
yuan-luo
as code owners
July 24, 2026 07:57
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
…uant Fall back to do_bench's default of cloning every input arg per graph iteration instead of only the activation, so the quant outputs also rotate buffers between iterations. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
BBuf
approved these changes
Jul 24, 2026
Collaborator
|
/rerun-test registered/ep/test_deepep_small.py |
Contributor
|
Results for 🚀 |
5 tasks
This was referenced Jul 25, 2026
whybeyoung
pushed a commit
to stepinto/sglang
that referenced
this pull request
Jul 28, 2026
…nel (PR sgl-project#30924) Extends 4f45d01: my earlier fp8_kernel.py revert only covered callers that go through _run_per_token_group_quant_8bit_kernel. The EP-MoE fused-silu-mul-quant default plain-silu path in _varlen_deep_gemm_silu_mul_quant bypasses that dispatcher and calls the unified JIT per_token_group_quant() directly, so it still hit the two PR sgl-project#30924 Hopper bugs (scale-tensor use-after-free from returning non-owning views for already-aligned scales; missing non-finite quant sanitization -- NaNs propagate to fp8 NaN codes instead of saturating). Observed symptom on GLM 5.2 NVFP4 MTP on H20 with --enable-ep-moe: - After 4f45d01, accept rate improved from ~0.20 -> ~0.32 (still low; MTP target is >0.5). Log evidence: DP0-DP7 EP0-EP7 all reporting accept_len 2.13-3.18, accept_rate 0.23-0.44. - The plain-silu masked EP-MoE fused-quant path was still on the buggy new kernel. Fix: allocate output_q and output_s inline using create_per_token_group_quant_fp8_output_scale (same layouts deep_gemm consumes: packed-int32 col-major for UE8M0, TMA-aligned col-major fp32 otherwise -- caller's get_mn_major transform still short-circuits) and call per_token_group_quant_8bit_v2 with fuse_silu_and_mul + masked_m. Notes: - v2 has full feature parity for (fuse_silu_and_mul + masked_m + scale_ue8m0 + col-major output_s), verified in the v2 custom op. - expected_m (a PR sgl-project#30924 grid-bound optimization) is dropped; v2 uses tokens_pad as the grid bound -- numerically identical, marginally slower. - gemm1_alpha and swiglu_limit/swizzle branches are untouched (they already use legacy silu_and_mul_masked_post_quant kernels, not the new JIT unified path). - Upstream long-term fix: PR sgl-project#32188 + PR sgl-project#32296. This is a debug revert until those are rebased in.
jinzhenfan
pushed a commit
to jinzhenfan/sglang
that referenced
this pull request
Jul 29, 2026
…quant (sgl-project#32296) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
jakki-amd
pushed a commit
to jakki-amd/sglang
that referenced
this pull request
Sep 9, 2026
…quant (sgl-project#32296) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Atituiset
pushed a commit
to Atituiset/sglang
that referenced
this pull request
Sep 10, 2026
…quant (sgl-project#32296) Co-authored-by: Claude Fable 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.
Motivation
#32188 fixed the H100 deepep TBO CI break by sanitizing non-finite quant inputs with a two-sided clamp before the fp8 conversion. That fix is correct, but it costs more than it needs to: the trait-driven quant kernel is latency-bound at decode batch sizes, and the clamp added 16
HMNMX2(ue8m0 path) resp. 32FMNMX(fp32-scale path) per bf16 group-of-128 — a 10–17% instruction-count increase on a kernel that runs twice per layer.The lower clamp is redundant:
fminf/__hmin2return the non-NaN operand, so a single uppermin(v, 448)already maps them to +448.cvt.satfiniteconversion saturates it to -448 on its own.So a single-sided min preserves the sanitization contract (non-finite inputs never reach an fp8 NaN code) at half the added cost. The only behavior delta: NaN now quantizes to +448 instead of -448 — both finite, which is the property the #32188 tests pin.
Modifications
WeightTrait<fp8_e4m3_t>::quant(fp32-scale path):fminf(fmaxf(v, -448), 448)→fminf(v, 448).__hmin2(__hmax2(x, lo2), hi2)→__hmin2(x, max_clip2).PDLWaitPrimarybelow the pure index arithmetic. This is codegen-neutral today (ptxas schedules theACQBULKidentically on sm_90a and sm_100a since the wait only orders memory ops), but keeps the source honest about what the wait actually protects: the first dependent global read (in.load/masked_m[...]), not the arithmetic.bench_per_token_group_quant.py: drop thegraph_clone_args=(1,)override — fall back todo_bench's default of cloning every input arg per graph iteration, so the quant outputs also rotate buffers between iterations.SASS evidence (bf16, gs=128, flat kernel,
--use_fast_math)Instruction counts of the whole kernel; clamp column counts only the sanitization min/max (the 4 amax-reduction
HMNMX2/ 1 epsFMNMXare excluded):HMNMX2HMNMX2FMNMXFMNMXHMNMX2HMNMX2FMNMXFMNMXBenchmark (B200,
bench_per_token_group_quant.py, this PR vs main)126 configs (gs 32/64/128 × row/col fp32 / col ue8m0 × 1–8192 tokens): geomean 1.009x, max 1.029x, min 0.993x, no config slower than noise. The win concentrates in the latency-bound 1–512 token range (1–3%); large batches are bandwidth-bound and unchanged. The
jit_v2baseline column drifted ≤1% between the two runs, so the comparison is clean.Test
test/registered/kernels/ops/quantization/test_per_token_group_quant.py: 105/105 pass on B200, including the 12 non-finite sanitization cases (NaN / ±inf × ue8m0/fp32 × flat/masked) added by #32188.🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #30077678531
Latest PR Test (Extra): ❌ Run #30077678071