Summary

Ports the AOT sgl_per_token_group_quant_8bit_v2 per-token-group FP8 quantization kernel (sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu) into the jit_kernel module with full feature parity, and routes the enable_v2 path of sglang_per_token_group_quant_fp8 to the JIT version. The AOT kernel stays in the
tree as a dormant fallback.

Motivation: moving this kernel to jit_kernel removes the AOT wheel rebuild from the iteration loop (first-use JIT compilation), keeps it next to the other JIT quant kernels, and lets us tune it without a sgl-kernel release.

The JIT kernel is a near-verbatim port (same 256-bit vectorized loads, 8 threads / 128-group, PDL, NaiveScheduler + MaskedLayoutScheduler, ue8m0 / float scales, fp8 / int8 output, fused silu+mul). It is bit-exact with the AOT kernel.

What changed

New

• python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh — the ported kernel + launcher.
• python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py — Python wrapper.
• python/sglang/jit_kernel/tests/test_per_token_group_quant_8bit_v2.py — bit-exact tests vs the AOT op.
• python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit_v2.py — JIT vs AOT benchmark.

Modified

• python/sglang/srt/layers/quantization/fp8_kernel.py — import the JIT v2 wrapper and call it from the enable_v2 branch of sglang_per_token_group_quant_fp8.

Implementation notes

• All shape-derived scalars are computed in the Python wrapper and passed to the kernel, so the C++ side only touches TensorView::data_ptr() / device().
• --use_fast_math is required in the JIT build to match the AOT kernel (which is compiled with -use_fast_math); without it the FP8 scale division rounds slightly differently and produces scattered 1-ulp FP8 differences.
• PDL is wired through the JIT device::PDLWaitPrimary / PDLTriggerSecondary helpers.

Validation

Hardware: NVIDIA H20-3e (SM90), CUDA 13.

Correctness — bit-exact vs the AOT kernel (104 passed)

torch.equal on both the FP8 codes and the scales, against
sgl_per_token_group_quant_8bit(..., enable_v2=True):

• 96 cases (NaiveScheduler): vanilla + fused silu+mul × bf16 / fp16 × num_tokens {1,7,64,333} × hidden {128,2048,4096} × float / ue8m0 scales.
• 8 cases (MaskedLayoutScheduler): ue8m0 + silu + masked 3D [num_experts, tokens_pad, hidden*2] × experts {2,5} × hidden {2048,4096} × tokens_pad {128,384}.

Benchmark — JIT v2 vs AOT v2

group_size=128, hidden=4096, bf16, CUDA-graph timing:

JIT v2 matches the AOT kernel at large batch and is slightly faster at small batch (no per-call output allocation on the measured path).

Note: This kernel occupies decode GPU 45.7% in Qwen/Qwen3.5-122B-A10B-FP8.

End-to-end — Qwen/Qwen3.5-122B-A10B-FP8, TP4, H20-3e

Served the model with the wired code (enable_v2 → JIT v2):

• CUDA graph capture succeeds (70 s, clean) — confirms the JIT kernel is CUDA-graph-compatible (compiles during warmup, captures cleanly).

• GSM8K (200 questions, 5-shot): accuracy 0.970, invalid 0.000 — identical to the AOT baseline (0.970), so the dispatch change is accuracy-neutral.

  Accuracy: 0.970
  Invalid: 0.000
  Latency: 21.618 s
  Output throughput: 1540.048 token/s
  

Notes

• The AOT sgl_per_token_group_quant_8bit remains available as a fallback (just not the default path).
• The enable_v2 group-size gate ({16, 32, 64, 128}) is unchanged.

CI States

Latest PR Test (Base): ✅ Run #27405353245
Latest PR Test (Extra): ✅ Run #27405352943