Skip to content

[Perf] Halve the non-finite sanitization overhead in per_token_group_quant - #32296

Merged
BBuf merged 2 commits into
mainfrom
perf/ptgq-single-sided-clamp
Jul 25, 2026
Merged

BBuf merged 2 commits into
mainfrom
perf/ptgq-single-sided-clamp

Conversation

@DarkSharpness

@DarkSharpness DarkSharpness commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

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. 32 FMNMX (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:

  • NaN / +inf: IEEE fminf / __hmin2 return the non-NaN operand, so a single upper min(v, 448) already maps them to +448.
  • -inf: passes through the min, but the fp8 cvt.satfinite conversion 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).
  • ue8m0 half path: __hmin2(__hmax2(x, lo2), hi2)__hmin2(x, max_clip2).
  • Both flat and masked kernels: move PDLWaitPrimary below the pure index arithmetic. This is codegen-neutral today (ptxas schedules the ACQBULK identically 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.
  • Updated the sanitization comments to describe the min + SATFINITE split.
  • bench_per_token_group_quant.py: drop the graph_clone_args=(1,) override — fall back to do_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 eps FMNMX are excluded):

arch path pre-#32188 main (#32188) this PR
sm_100a ue8m0 128 insns / 0 clamp 144 / 16 HMNMX2 136 / 8 HMNMX2
sm_100a fp32 scale 184 / 0 216 / 32 FMNMX 200 / 16 FMNMX
sm_90a ue8m0 192 / 0 216 / 16 HMNMX2 200 / 8 HMNMX2
sm_90a fp32 scale 192 / 0 224 / 32 FMNMX 208 / 16 FMNMX

Benchmark (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_v2 baseline 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

…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>
@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.

@github-actions github-actions Bot added quant LLM Quantization jit-kernel labels Jul 24, 2026
…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

BBuf commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

/rerun-test registered/ep/test_deepep_small.py

@github-actions

github-actions Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/ep/test_deepep_small.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/ep/test_deepep_small.py

@BBuf
BBuf merged commit 9402012 into main Jul 25, 2026
101 of 117 checks passed
@BBuf
BBuf deleted the perf/ptgq-single-sided-clamp branch July 25, 2026 00:17
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

jit-kernel quant LLM Quantization

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants