[zcnrex]

Summary

Bumps sgl-deep-gemm to 0.2.6 and adopts the two 0.2.6 changes that need a caller-side change. The other 0.2.6 changes are DeepGEMM-internal and on by default: Mega MoE block_m 64 for decode loads, the mega_mhc split-barrier fix under CUDA graph replay, the K=64 MXF8F6F4 dense MMA on CC 10.7, and the SIMT dispatch pull for FP4 activations.

DeepGEMM producers write FlashInfer's 128x4 scale layout (sgl-project/DeepGEMM#136)

On CC 10.7, two DSv4 linears take their input as an MXFP8 pair written by a DeepGEMM kernel. wo_b gets it from the wo_a fp8_einsum epilogue, and wqkv_a from Mega mHC. 0.2.6 lets both producers write the scales directly in FlashInfer's 128x4 layout: pass a 4-D int32 scale [ceil(rows/128), K/128, 32, 4], and its flat uint8 view is FlashInfer's a_descale. The pair then runs on the FlashInfer MXFP8 GEMM with no conversion step.

• A producer writes the FlashInfer layout whenever its consumer is a FlashInfer layer (the only case in which a pair is produced), outside batch-invariant and deterministic modes. There is no per-shape or per-row threshold: FlashInfer beats or ties DeepGEMM at every measured row count (table below). The scale converter added earlier in this PR is removed.
• mxfp8_input.py
  • deepgemm_pair_takes_flashinfer is the route check. batch_invariant_or_deterministic moves here from fp8.py (with top-level imports), and Fp8LinearMethod.apply uses it for pin_tactic.
  • deepgemm_mxfp8_linear_input hands a filled 4-D pair on as Mxfp8SwizzledInput. Fp8LinearMethod.apply already sends that type to FlashInfer.
  • zero_deepgemm_mxfp8_scale_rows zeroes padded rows in either layout.
• deepseek_v4.py / deepseek_v4_mhc.py: the producers allocate the 4-D scale (flashinfer= on the allocation) when the route applies.
  • In the breakable-CUDA-graph path, the live rows' 4-D scales are a prefix of the padded ones, so they need no restride.
• A DeepGEMM-packed int32 pair (batch-invariant or deterministic mode) still runs on DeepGEMM, as on main.

Current MQA logits API for bf16 index weights (sgl-project/DeepGEMM#134)

The 0.2.3-style call (positional clean_logits plus logits_dtype) is routed by sgl-deep-gemm to fp8_fp4_mqa_logits_legacy. With bf16 weights, flat_index_logits_tiles now calls the current API, which reaches the kernel #134 sped up. fp32 weights keep the legacy call. The logits are bitwise identical.

Decode

• The wo_a einsum producer also runs for decode and verify batches up to 2048 rows. So wo_b takes the FlashInfer route there too, and the kernel table below covers decode-sized row counts.
• wqkv_a's Mega mHC producer is prefill-only, and the paged MQA logits call is unchanged, so neither route affects decode.
• Accuracy in decode is covered by the GSM8K runs below. Decode throughput is in the decode table below.

Performance

Kernel level, CC 10.7. Producer + FlashInfer GEMM on 4-D scales, divided by producer + DeepGEMM GEMM on packed scales, both on 0.2.6. Each cell is one run or two runs (run1 / run2) on swapped GPUs. Lower is better.

rows	wo_b	wqkv_a
1	0.92 / 0.87	0.93 / 0.93
8	0.91 / 0.90	0.94 / 1.03
16	0.88 / 0.87	0.96 / 0.96
32	0.87	0.89
64	0.85 / 0.86	0.96 / 0.95
128	0.93 / 0.92	0.90 / 0.91
256	0.88 / 0.88	0.95 / 0.93
512	0.92 / 0.93	0.91 / 0.91
1024	0.92 / 0.96	0.97 / 0.97
2048	0.94 / 0.98	0.98 / 0.96
4096	0.94 / 0.92	0.96 / 0.97
8192	0.96 / 0.93	0.99 / 1.00
16384	0.91 / 0.94	0.99 / 1.00

• At 1, 129, 1000, 4100 and 16384 rows, the 4-D scales and fp8 data are bitwise equal to the packed-scale path followed by a conversion. So is the FlashInfer output, including with NaN-filled pad scales.
• DeepGEMM's best case is a tie: wqkv_a at 8K–16K rows (0.99–1.00). On CC 10.3, an earlier measurement of the same producer + GEMM pair puts FlashInfer at 0.89–0.99, and DeepGEMM ahead by at most 1% (wqkv_a at 8K–16K). So no shape keeps a DeepGEMM route.

#134 (CC 10.3, index logits per prefill chunk, H32, 131K context, 4×4096 new tokens): −4.3% with the new call. The old call is unchanged. Logits are bitwise identical.

Server, CC 10.7, DeepSeek-V4.1-Flash, DEP4 (--dsv4-attn-backend flashmla_mega, SGLANG_OPT_DSV41_MEGA_MHC_PREFILL=1, DSPARK). Each arm ran once on each of two nodes.

arm	GSM8K-200	prefill TTFT, 16 × 16K prompts (s)
main + 0.2.5.post3	0.990 / 0.995	0.384 / 0.386
main + 0.2.6	0.985 / 0.995	0.380 / 0.381
this PR + 0.2.6	0.995 / 0.985	0.377 / 0.390
this PR + 0.2.6, no routing table (head)	0.985 / 0.995	0.376 / 0.378

Accuracy is unchanged. At 16K-token prompts, prefill time is within run-to-run noise: these GEMMs are a small share of the step.

Decode, same server setup. Steady-state generation throughput summed over the 4 DP ranks, in tok/s. 128K-token context, simulated acceptance length 3.51, decode CUDA graphs up to 128. One server per cell; each arm ran once on each of two nodes.

arm	32 requests	128 requests
main + 0.2.5.post3	14000 / 14340	40941 / 40551
main + 0.2.6	14641 / 14549	43417 / 43455
this PR + 0.2.6	15073 / 14357	44068 / 44249

• At 128 requests, 0.2.6 alone is +6.6%. The wo_b route adds +1.7% on top (+1.5% and +1.8% on the two nodes).
• At 32 requests, 0.2.6 is +3.0%. The wo_b route adds +0.8%, within noise: +3.0% on one node and −1.3% on the other.

A profiled prefill step on this PR (rank 0) shows both routes taken:

• wo_b and wqkv_a run on FlashInfer's block-scaled dense GEMM.
• The only DeepGEMM dense GEMM left is the wo_a einsum that produces wo_b's input.
• Index logits run sm100_mqa_logits, not the legacy kernel.

Tests

• test_fp8_blockwise_linear_backends.py::test_deepgemm_producer_flashinfer_layout:

  • the 4-D scale's flat view equals flashinfer.block_scale_interleave at 200 and 256 rows;
  • the FlashInfer route matches the DeepGEMM route and the fp32 reference;
  • a spy checks which route runs;
  • the pad-row zeroing matches 2-D zeroing followed by a relayout.

  CC 10.3 with sgl-deep-gemm 0.2.6: the file gives 17 passed and 7 skipped.

• Indexer tests (test_fp4_indexer, test_dense_prefill_indexer, test_dsv41_prefill_sparse_indexer, test_dsv41_indexer_tp_shard) on CC 10.3 with 0.2.6: 117 passed, 2 skipped.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #41934628566
Latest PR Test (Extra): ❌ Run #41934628353
Latest PR Test (AMD ROCm 10): ❌ Run #41934628246

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[zcnrex]

/rerun-test registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py registered/kernels/ops/attention/dsv4/test_mega_attention.py registered/kernels/ops/layernorm/test_mhc_kernels.py registered/unit/models/test_deepseek_v4_boundary.py registered/unit/models/test_deepseek_v4_pipeline.py registered/e2e/model_executor/test_decoder_tail_moe_bound.py registered/kernels/ops/attention/test_dsv41_prefill_sparse_indexer.py registered/kernels/ops/attention/test_dense_prefill_indexer.py registered/kernels/ops/attention/test_fp4_indexer.py registered/unit/layers/test_dsv41_indexer_tp_shard.py registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py

[github-actions]

🚀 4-gpu-b200 (4 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py
cd test/ && python3 registered/kernels/ops/attention/test_dsv41_prefill_sparse_indexer.py
cd test/ && python3 registered/kernels/ops/attention/test_dense_prefill_indexer.py

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/e2e/model_executor/test_decoder_tail_moe_bound.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py
cd test/ && python3 registered/kernels/ops/attention/test_fp4_indexer.py

🚀 ubuntu-latest (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/models/test_deepseek_v4_boundary.py
cd test/ && python3 registered/unit/models/test_deepseek_v4_pipeline.py
cd test/ && python3 registered/unit/layers/test_dsv41_indexer_tp_shard.py

🚀 2-gpu-gb300 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py