[mmangkad]

Motivation

The DSV4 FP8 wo_a path in DeepSeek-V4 needs per-token-group scales in the layout consumed by deep_gemm.fp8_einsum("bhr,hdr->bhd", ...). This PR adds a DSV4-specific fused JIT quantization path that produces the FP8 activation tensor and writes the final UE8M0-rounded scale tensor directly in the fp8_einsum layout while preserving the existing FP8 activation quantization behavior.

This fuses the scale rounding and DSV4 scale-layout write into the quantization kernel, avoiding a separate post-quant scale conversion/layout step. This PR only changes the DSV4 FP8 wo_a path and is based on the JIT v2 per-token-group quantization infrastructure from #27449.

Modifications

• Added PerTokenGroupQuantFp8Dsv4Kernel to the JIT v2 per-token-group quantization kernel. It reuses the JIT v2 FP8 quantization math and FP8 conversion path, and fuses UE8M0 scale rounding plus DSV4 scale-layout writes into the same kernel.
• Added Python JIT wrapper per_token_group_quant_fp8_dsv4.
• Added SGLang helper sglang_per_token_group_quant_fp8_dsv4_ue8m0.
• Routed the DeepSeek-V4 DSV4 FP8 wo_a path through the new helper.
• Added a registered test covering BF16/FP16 inputs, empty/non-empty inputs, exact FP8 bytes, and exact UE8M0 DSV4 scale layout.

Accuracy Tests

GSM8K

Command:

sgl-eval run gsm8k \
  --base-url http://localhost:30000/v1 \
  --num-threads 256

Result:

benchmark	examples	score	latency	output throughput
GSM8K	1319	96.51%	92.35s	1569.66 tok/s

Kernel Exactness Test

The unit test compares the new DSV4 path against the previous tensor behavior (sglang_per_token_group_quant_fp8 + deep_gemm.ceil_to_ue8m0 + DSV4 layout transpose). It checks exact FP8 bytes, UE8M0 scale values, and scale tensor layout for BF16/FP16 inputs, including empty inputs.

python -m pytest test/registered/jit/test_per_token_group_quant_8bit_v2.py -k dsv4 -q
10 passed, 104 deselected

Speed Tests and Profiling

Benchmarks were run on one GPU from a single GB300 node with deepseek-ai/DeepSeek-V4-Flash.

Bench One Batch, 1k/1k

Command:

CUDA_VISIBLE_DEVICES=0 SGLANG_JIT_DEEPGEMM_FAST_WARMUP=1 python -m sglang.bench_one_batch \
  --model-path deepseek-ai/DeepSeek-V4-Flash \
  --trust-remote-code \
  --moe-runner-backend flashinfer_mxfp4 \
  --disable-flashinfer-autotune \
  --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 46}' \
  --batch-size 1 8 \
  --input-len 1024 \
  --output-len 1024

Prefill latency:

batch size	main	this PR	delta
1	0.13613s	0.13496s	0.9% faster
8	0.52655s	0.52408s	0.5% faster

Decode latency:

batch size	main	this PR	delta
1	0.007508s	0.006957s	7.3% faster
8	0.010594s	0.009316s	12.1% faster

Overall throughput:

batch size	main	this PR	delta
1	261.95 tok/s	282.32 tok/s	+7.8%
8	1443.15 tok/s	1630.04 tok/s	+13.0%

Bench Serving, 1k/1k

Server:

CUDA_VISIBLE_DEVICES=0 SGLANG_JIT_DEEPGEMM_FAST_WARMUP=1 sglang serve \
  --trust-remote-code \
  --model-path deepseek-ai/DeepSeek-V4-Flash \
  --moe-runner-backend flashinfer_mxfp4 \
  --disable-flashinfer-autotune \
  --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 46}'

Benchmark:

python -m sglang.bench_serving \
  --backend sglang \
  --model deepseek-ai/DeepSeek-V4-Flash \
  --dataset-name random \
  --random-input-len 1024 \
  --random-output-len 1024 \
  --random-range-ratio 1.0 \
  --num-prompts $c \
  --request-rate inf \
  --max-concurrency $c \
  --seed 1 \
  --disable-tqdm \
  --flush-cache

Output throughput:

concurrency	main	this PR	delta
1	138.54 tok/s	149.94 tok/s	+8.2%
8	722.61 tok/s	788.43 tok/s	+9.1%

TTFT:

concurrency	main	this PR	delta
1	172.09ms	156.19ms	9.2% faster
8	755.62ms	660.83ms	12.5% faster

TPOT:

concurrency	main	this PR	delta
1	7.063ms	6.499ms	8.0% faster
8	10.333ms	9.502ms	8.0% faster

---

CI States

Latest PR Test (Base): ❌ Run #27456441998
Latest PR Test (Extra): ❌ Run #27456441910

[b8zhong]

#26766 (the author is rebasing when the JIT path is merged)

[mmangkad]

#26766 (the author is rebasing when the JIT path is merged)

@b8zhong ah, missed that one, closing this as duplicate. The only extra bit here was writing the UE8M0 scales directly in the DSV4 fp8_einsum layout.

[b8zhong]

@mmangkad I'll add you as a co-autho r