[DSV4.1] Big fused wo_a quant

[DarkSharpness]

Motivation

On small DeepSeek-V4.1 TP4 batches, inverse RoPE, WO-A projection, and MXFP8 quantization require separate kernel launches. This PR fuses the sequence to reduce launch and intermediate-memory overhead.

Changes

• Add the SM100/SM103 fused inverse-RoPE + WO-A kernel, with BF16 or FlashInfer-swizzled MXFP8 output, for up to 32 token rows and the supported (groups, rank, K) = (2, 1024, 4096) shape.
• Dispatch it for eligible DSV4.1 layers; preserve main's model-type gating, NPU WO-A path, and the existing fallback for other shapes/backends.
• Add the JIT CUDA-driver link support and mbarrier helper required by the kernel.
• Keep 18 focused correctness cases in test/registered/kernels/: BF16/MXFP8 output, padded inputs, token-count boundaries, scale padding, and int32/int64 positions. Validate scales in bulk and update the WO-A imports in existing tests.

Validation

Rebased onto main at a6cf05817f11d22023fd951a76255ef50fb09f49, after #38798 and #40039 landed. All local pre-commit checks on changed files passed. Full suite-registration validation passed for all 2,489 registrations; the experimental benchmark with an invalid B200 suite has been removed, along with its unrelated global benchmark changes. Vectorized scale indexing was checked against all 8,192 original offsets. CI is rerunning on the updated head. GPU tests and benchmarks have not yet been rerun on this new main-based revision.

Historical measurements on e919e8be784473d2fb7ea10d10f45c23af658f68 and its direct parent:

• 4×GB300, TP4/EP1, concurrency 1, 4096 input / 1024 output, DSPARK block size 5, simulated acceptance length 5.5: median client decode throughput 1167.22 → 1202.65 tokens/s (+3.03%). The 18 added kernel correctness tests passed. Full setup and results.
• Real-acceptance accuracy checks covered GSM8K, AIME 2026, and GPQA, with repeated GSM8K runs. These are not a claim of exact numerical/statistical equivalence. Accuracy results and limitations.

---

CI States

Latest PR Test (Base): ⏳ Run #35437055876
Latest PR Test (Extra): ⏳ Run #35437055726
Latest PR Test (AMD ROCm 10): ⏳ Run #35437055858

[BBuf]

Tested this PR on 4× GB300 (SM103). At request concurrency 1, median client-side decode throughput improved from 1167.22 to 1202.65 tokens/s (+3.03%).

• Baseline: 9a77971b903360f6700c797bb4a2b5a62e8c391f (the direct parent).
• Candidate: e919e8be784473d2fb7ea10d10f45c23af658f68.
• DeepSeek-V4.1-Flash, TP4 / EP1, input 4096 / output 1024 tokens, DSPARK block size 5, simulated acceptance length 5.5 (SGLANG_SIMULATE_ACC_METHOD=match-expected). Identical deployment arguments, weights, input and software environment for both revisions; no source overlays. The separate router PDL change in [DSV4.1] Avoid PDL for small target routers on SM103 #39941 was not included.
• PyTorch 2.13.0+cu130, Triton 3.7.1, FlashInfer 0.6.18, sglang-kernel 0.4.7.

Each revision had two fresh server starts, with one warmup and 20 measured requests per start. Execution order was baseline → PR → PR → baseline. Prefix caches were flushed before each request; all 80 measured requests had zero cached tokens. No samples were discarded.

Version	First start median	Second start median	Pooled median (40 requests)
Baseline	1164.29	1170.13	1167.22
This PR	1205.98	1197.09	1202.65

Units: tokens/s. Throughput is measured at the streaming client as (completion_tokens - first_chunk_tokens) / (last_chunk_time - first_chunk_time), excluding TTFT and the first output chunk. Profiling was disabled during these measurements. The per-request acceptance-length sequences matched across all four starts, with median 5.505376.

The isolated inverse-RoPE → WO-A partial GEMM → reduce/quant chain also shows a consistent improvement:

Token rows T	Original three-kernel chain	Fused kernel	Latency reduction
2	8.959 µs	4.964 µs	44.59%
6	9.208 µs	4.964 µs	46.08%
8	9.226 µs	4.983 µs	45.99%

These are medians of three sweep medians, alternating measurement order, using the same benchmark helper/L2 marker for both implementations and rotating weight buffers beyond L2 capacity (30 warmups and 2000 replays per sweep).

Separate model traces confirm the fused path is active: in each rank0 trace covering 20 decode steps, 800 instances of the original three-kernel chain were replaced by 800 wo_a_fused_kernel calls, eliminating 1600 kernel invocations. The remaining 60 instances of the original non-quantized path were retained.

Validation: all 18 added kernel tests passed. This run used a synthetic token prompt and simulated acceptance for performance measurement; GSM8K/AIME/GPQA were not run, so this does not establish task-level accuracy equivalence. BS32/64 and the combination with #39941 were not measured here.

[BBuf]

Accuracy follow-up on e919e8be784473d2fb7ea10d10f45c23af658f68, compared with its parent 9a77971b903360f6700c797bb4a2b5a62e8c391f.

Same 4×GB300, TP4/EP1, DSPARK block size 5, real speculative acceptance (no simulated acceptance), temperature 0, and identical frozen prompts/evaluator settings. The full suites used concurrency 4, exercising the fused path with 24 verify rows.

Dataset	Baseline	PR	Output-limit hits (base → PR)
GSM8K, 1,314 held-out questions, 5-shot	1,282/1,314 (97.56%)	1,278/1,314 (97.26%)	1 → 1
AIME 2026, 30 questions	28/30 (93.33%)	28/30 (93.33%)	2 → 2
GPQA Diamond, 198 questions	167/198 (84.34%)	174/198 (87.88%)	16 → 12

There were no request errors. GSM8K excludes the five examples used as demonstrations. Output-limit hits remain in the scored results.

For the small GSM8K difference, fresh-server repeats produced 1,281/1,314 on baseline and 1,279/1,314 on the PR. Thus baseline scored 97.49–97.56% and the PR 97.26–97.34%; five of the seven initially lost questions recovered in the PR repeat. There is question-level run-to-run variation, so I would not treat the initial four-question difference alone as evidence of a kernel accuracy bug. These runs do not establish exact numerical or statistical equivalence. The GPQA increase also should not be interpreted as a model-quality improvement: output-limit outcomes contribute to the difference.

Additional checks:

• All 18 existing fused WO-A kernel tests passed. They check BF16 results against the FP64 reference and exact quantized values/scales against the corresponding BF16 quantization reference; this is not a claim of bitwise equality to the old GEMM path.
• Fixed concurrency-1 subsets: GSM8K 64/64 → 63/64, AIME 7/8 → 7/8, GPQA 26/32 → 28/32. These are small supplemental subsets, not full BS=1 accuracy suites.
• A real-acceptance, four-rank trace confirmed that the candidate exercised the fused kernel (800 fused WO-A calls on rank 0 over the 20-step diagnostic).

Both source snapshots were hash-checked after evaluation. The evaluator was pinned to SGL-Eval 01a94430ea242033b1051687104801f4f3a3536d, and the prompts, answers, sample IDs, and evaluation settings matched across arms.

[BBuf]

/rerun-test test/registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py

[github-actions]

Results for /rerun-test test/registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py:

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

cd test/ && python3 registered/quant/test_llama8b_nvfp4_kv_cache_sm120.py

[BBuf]

Base NVIDIA CI has a PR-related unit-test failure: test_deepseek_v4_unified_fp8_q_pair.py constructs _Harness without calling MQALayer.__init__, while this PR makes forward() read the new self.use_fused_wo_a field. This deterministically raises an AttributeError because _Harness has no use_fused_wo_a attribute.

Could you initialize self.use_fused_wo_a = False in that test harness? I did not rerun this failure because it is directly related to the PR and is not flaky.

[BBuf]

/rerun-test test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py

[github-actions]

Results for /rerun-test test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py

[BBuf]

Fixed the CPU test regression in 280d84c. The test harness bypasses MQALayer.init, so it now initializes use_fused_wo_a=False to match the production default and keep this unified-KV test on the intended non-fused path.

Validation:

• Targeted CPU rerun passed: https://github.com/sgl-project/sglang/actions/runs/35419009188
• The same test also passed in the full Base CPU shard: https://github.com/sgl-project/sglang/actions/runs/35418955183/job/105832992563

The remaining Base failure is an unrelated current-main test regression (test_layer_communicator_fusion_gate.py mocks get_tp_group after that symbol was removed). Extra jobs are being fast-failed by the unrelated current-main ruff-format failure in logits_processor.py; no #39957 GPU test reported a real failure on this head.

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_wo_a_fused.py test/registered/kernels/ops/attention/test_fp8_wo_a.py test/registered/kernels/ops/layernorm/test_mxfp8_epilogue.py test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/attention/test_wo_a_fused.py test/registered/kernels/ops/attention/test_fp8_wo_a.py test/registered/kernels/ops/layernorm/test_mxfp8_epilogue.py test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py:

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

cd test/ && python3 registered/kernels/ops/attention/test_wo_a_fused.py
cd test/ && python3 registered/kernels/ops/attention/test_fp8_wo_a.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mxfp8_epilogue.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

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

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

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

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

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_wo_a_fused.py test/registered/kernels/ops/attention/test_fp8_wo_a.py test/registered/kernels/ops/layernorm/test_mxfp8_epilogue.py test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py test/registered/unit/managers/test_disagg_idle_step_counters.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/attention/test_wo_a_fused.py test/registered/kernels/ops/attention/test_fp8_wo_a.py test/registered/kernels/ops/layernorm/test_mxfp8_epilogue.py test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py test/registered/unit/managers/test_disagg_idle_step_counters.py:

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

cd test/ && python3 registered/kernels/ops/attention/test_wo_a_fused.py
cd test/ && python3 registered/kernels/ops/attention/test_fp8_wo_a.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mxfp8_epilogue.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

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

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

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

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

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

cd test/ && python3 registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py
cd test/ && python3 registered/unit/managers/test_disagg_idle_step_counters.py