[DeepSeek-V4.1] Optimize DSpark verify and MoE kernels on Blackwell - #38879
Merged
Merged
Conversation
Overlap mHC statistics and routed input quantization, write WO-A output in token-major layout, fuse candidate masking and route packing, and fuse MoE finalize/shared add with the custom push all-reduce on supported small TP4 batches. Preserve unfused fallbacks and phase synchronization. Co-authored-by: DarkSharpness <2040703891@qq.com>
BBuf
marked this pull request as ready for review
September 10, 2026 13:49
BBuf
requested review from
Alisehen,
AniZpZ,
DarkSharpness,
Edwardf0t1,
FlamingoPg,
Fridge003,
HaiShaw,
HydraQYH,
OrangeRedeng,
Qiaolin-Yu,
Ying1123,
b8zhong,
celve,
ch-wan,
fzyzcjy,
hebiao064,
ispobock,
merrymercy,
mmangkad and
yuan-luo
as code owners
September 10, 2026 13:49
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
DSpark target verification and draft steps miss several decode fast paths. MoE routing, input quantization, and finalize also leave work on the critical path. This carries the compatible kernel changes onto
dsv4.1.Modifications
The updated public top-k implementation is retained. This PR contains production changes only; standalone validation harnesses are kept outside the source diff.
Co-authored with @DarkSharpness (Ziyi Xu), who contributed the MoE router and input prequantization, finalize/shared-add/all-reduce fusion, and custom all-reduce backend selection.
Speed Tests
Base: public
dsv4.1at7bdebdab7db4befb71c64ae0d6f0eb37fe7d8402. Candidate:1b742acd2a49ebd7acd017032875552099d12391.4x B300 SXM6 AC, TP4/EP4, PyTorch 2.13.0+cu130, FlashInfer 0.6.18, Triton 3.7.1,
sglang-kernel==0.4.6.post1,sgl-deep-gemm==0.1.7, CUTLASS DSL 4.6.2, nvcc 13.0.88, driver 580.126.20. DSpark block size 5 with real acceptance. Both arms use the same checkpoint, package environment, fixed KV pool sizes, and hashed input IDs.Two independent server launches per arm, with one warmup and three measured repetitions each. Order: baseline A, candidate A, candidate B, baseline B. Medians below use all six measured repetitions.
Median BS1 acceptance length: 5.658 → 5.818; BS64: 5.407 → 5.479. No simulated acceptance. BS1 acceptance/TPS cycle proxy: 9.897 → 7.603 ms. This proxy is separate from the GPU trace timings below.
BS64 decode samples require consecutive log intervals with exactly 64 running requests and CUDA graphs active, excluding intervals across prefill. The full-batch metric includes prefill and the draining tail. BS1 excludes the first SSE event and counts only the remaining output tokens. These metrics should not be compared interchangeably.
Per-server measurements
Accuracy Tests
All requested samples completed with zero request errors. GSM8K has no empty or truncated responses. AIME budget truncations are counted in the score, not dropped. The repeated AIME score is correct samples / 480, not pass@16.
GSM8K has 14 correct-to-incorrect and 12 incorrect-to-correct changes under the unchanged prompt and scorer. Generated text is not bitwise identical. These results and the kernel checks do not establish exact model equivalence for every input.
GSM8K uses the first five test rows as demonstrations and excludes them from evaluation. AIME uses
sgl-eval==0.1.0, the bundled MathArena prompt and NeMo-Skills revision645cf567ff08c0ae9cc3fc8e1edbb975b3067816, thinking mode, top-p 0.95, and max tokens 65,536. Serial AIME uses seed 0; repeated AIME leaves the request seed unset.Kernel validation and profiling
GPU-only TP0 traces, BS1 with 4096 input tokens, 20 target graph replays and 20 draft graph replays per arm. Each draft graph contains the three draft stages. Trace timings are reported separately from unprofiled serving throughput.
The following launch counts cover the complete 20-cycle trace:
mHC statistics overlap with other kernel families on different CUDA streams increases from 0.18% to 93.50% of its kernel duration. This measures time overlap in the trace; it is not an occupancy metric or a guarantee that the overlapped work has zero cost. Graph spans above measure the net effect.
Server command (both arms)
After the server is ready, call
POST /freeze_gc. Flush the request cache before each measured repetition. Both arms use the identical launch command; the runtime selects the compatible kernel paths.Checkouts used for the comparison
From a clone of
sgl-project/sglang:Run each arm from its own worktree with the same installed dependency versions. Start with inherited
SGLANG_*overrides removed; the measured controllers clear those variables. The same checkpoint directory is used for both arms.Benchmark and GSM8K client commands and harnesses
Save the three Python blocks below beside
gsm8k-test.jsonl. SetMODEL_PATHto the same checkpoint used by the server andSERVER_LOGto its log file. Keep the serving checkout first onPYTHONPATH; the preinstalled editable SGLang package is not the source under test. The benchmark scripts differ from the measured copies only in making the tokenizer path configurable.GSM8K uses the public
openai/grade-school-mathgrade_school_math/data/test.jsonlfile, SHA2563730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14. The first five rows are demonstrations and are excluded from the 1,314 evaluated questions.bench_bs1.pybench_bs64.pygsm_eval.pyAIME 2026 commands
The run uses
sgl-eval==0.1.0with its bundled NeMo-Skills data/prompt (645cf567ff08c0ae9cc3fc8e1edbb975b3067816). The evaluator is installed outside the serving environment. Both arms use the same files and sampling settings.The repeated lane leaves the request seed unset. Accuracy is the fraction of correct samples, not pass@16. Errors and token-budget truncations are counted separately.
CI
GPU CI is blocked before tests by the global requirement to include
maincommit3700c4ee26a1, which also rejects thisdsv4.1-based PR (job log). The B300 results above were collected directly against the stated base and candidate.CI States
Latest PR Test (Base): ❌ Run #34478244464
Latest PR Test (Extra): ❌ Run #34478243976
Latest PR Test (AMD ROCm 10): ❌ Run #34478244607