Repository navigation
[dsv4] Prewarm MHC prenorm kernel at startup - #27986
Merged
Fridge003 merged 3 commits intoJun 15, 2026
Merged
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
YAMY1234
changed the title
sgl-project#26238 routed MHC-prenorm warmup through the DeepGEMM wrapper, which only compiles the GEMM. The TileLang `mhc_pre_big_fuse_with_norm` kernel runs after the GEMM in mhc_pre() regardless of backend, is specialized per `n_splits` bucket, and was then compiled lazily on the first prefill/decode -- costing first-run throughput. (`num_tokens` is a dynamic dim, so only `n_splits` keys the kernel.) On the first non-capturing mhc_pre() call (the pre-capture warmup run), replay mhc_pre once per distinct n_splits bucket over [1, chunked_prefill_size] with the call's real weights, so the compiled variants match the runtime exactly. Guarded by torch.cuda.is_current_stream_capturing() (never recorded into a CUDA graph) and a one-shot flag. Gated by SGLANG_DSV4_MHC_PREWARM (default on). Re-adds get_mhc_pre_token_count_representatives (removed by sgl-project#26238). Self-contained in mhc.py + one env var -- no changes to model_runner or the model definition. Note replaying mhc_pre also warms the DeepGEMM prenorm GEMM (tf32_hc_prenorm_gemm) per n_splits bucket, so it subsumes sgl-project#27729's warmup. Faithful repro (Shu-Hao's exact harness, B300 TP4 conc32 1k1k EAGLE, two launches): first run 1016->1260 tok/s (+24%) vs no-warmup baseline; run1/run2 gap 234->25. One-time startup warmup ~145s.
…p prewarm The previous commit prewarms mhc_pre() per n_splits bucket at startup, and mhc_pre internally calls deep_gemm.tf32_hc_prenorm_gemm -- so the dedicated DeepGEMM prenorm-GEMM warmup added in sgl-project#26238 (the deep_gemm_execution_hook routing + _TF32HcPrenormWarmupExecutor) now compiles the same per-n_splits kernels a second time. Remove it: tf32_hc_prenorm_gemm() becomes a thin pass-through to deep_gemm, and the TF32_HC_PRENORM_GEMM executor / enum member / registration / memory estimate are dropped. The runtime path (mhc_pre -> deepgemm TF32 prenorm GEMM) is unchanged; the GEMM is now warmed solely by the startup prewarm replay (verified: that replay already emits "Try DeepGEMM JIT Compiling <TF32_HC_PRENORM_GEMM> num_groups=..." for every bucket at startup).
YAMY1234
force-pushed
the
dsv4-mhc-prenorm-startup-prewarm-v2
branch
from
June 12, 2026 01:11
b0ec97f to
9736a78
Compare
YAMY1234
marked this pull request as ready for review
June 12, 2026 01:11
YAMY1234
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Ying1123,
ch-wan,
ispobock and
merrymercy
as code owners
June 12, 2026 01:11
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Fridge003
reviewed
Jun 12, 2026
Per review on sgl-project#27986: - Split mhc_pre into a thin one-shot guard + dispatch wrapper and a pure _mhc_pre_impl. The startup prewarm now replays _mhc_pre_impl (not mhc_pre), so there is no mhc_pre <-> warmup recursion -- warm-once no longer depends on flag-ordering. - Pass arguments explicitly instead of locals()/**kwargs forwarding. - Document that compiled kernels persist in the on-disk JIT cache, so the warmup cost is paid only on a cold cache; later runs hit the cache.
YAMY1234
force-pushed
the
dsv4-mhc-prenorm-startup-prewarm-v2
branch
from
June 13, 2026 15:47
8f3486e to
143d18d
Compare
Fridge003
approved these changes
Jun 14, 2026
Collaborator
|
/rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_h200.py test/registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py test/registered/disaggregation/test_disaggregation_dsv4.py |
Contributor
|
Results for 🚀 🚀 |
Merged
3 of 5 tasks
YAMY1234
added a commit
to YAMY1234/sglang
that referenced
this pull request
Jul 3, 2026
The one-shot MHC prenorm prewarm (sgl-project#27986) triggers lazily inside the first forward that carries tokens, per rank, uncoordinated. On wide-EP disagg prefill (GB200 1p1d dep8) that is fatal: idle EP peers still launch deep_gemm.fp8_fp4_mega_moe every layer and wait for the compiling rank inside its NVLink barrier, which device-traps after 180 s ("DeepGEMM NVLink barrier timeout ... signal=7, target=8") and poisons the context (CUDA error 719 on later launches). Measured on GB200: the 23-bucket burst is 159-166 s on a cold JIT cache, so any >15 s trigger stagger between ranks blows the barrier budget. Fix: drive prewarm_mhc_pre() from the tail of load_weights on every rank — before the memory pool, CUDA-graph capture, and any forward — then barrier on the TP group so no rank proceeds while a peer still compiles, and empty_cache so the transient prewarm buffers don't skew KV-pool sizing. With the trigger out of the forward path, the in-forward gate and the mhc_pre/_mhc_pre_impl split (which existed only so the prewarm replay would not re-enter mhc_pre) lose their reason to exist: remove them and make mhc_pre the single implementation again. Also drop the orphaned prewarm_mhc_token_counts/_buckets pair (added by sgl-project#25810, removed by sgl-project#26238, resurrected without a caller by the sgl-project#25976 merge).
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
3 of 5 tasks
4 of 5 tasks
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
On DeepSeek-V4, the first benchmark run on a fresh server is much slower than later runs (~1016 vs ~1250 tok/s; B300, TP4 / conc 32 / 1k1k / EAGLE). This persists even though the benchmark already runs a
2×concurrencywarmup before the timed run: that warmup uses a single fixed-shape input, so it only compiles a few of the MHC prenorm kernel's per-n_splitsvariants; the timed run (random lengths) hits the rest and JIT-compiles them mid-run. A single warmup pass does not fix it — everyn_splitsbucket must be compiled at startup.Two prior changes are already merged: #25810 added a model-side prenorm warmup (later reverted — the logic was spread across the model and
ModelRunner), and #26238 routed MHC prenorm through the DeepGEMM wrapper and warmed the prenorm GEMM (tf32_hc_prenorm_gemm) there. However that turned out not to be sufficient, so #27729 tried to explicitly sweep alln_splitsbuckets of that GEMM — yet the first-run gap remains (see Speed Tests). The reason:mhc_pre()runs a second, TileLang kernel after the GEMM (mhc_pre_big_fuse_with_norm, specialized pern_splits) which is not a DeepGEMM kernel, so DeepGEMM-side warmup never reaches it.This PR instead warms the actual
mhc_pre()path at startup. Replayingmhc_pre()compiles the TileLang kernel and, via the same call, the DeepGEMM GEMM — so it covers what #26238/#27729 warm plus the kernel they miss, superseding them. The second commit removes #26238's now-redundant DeepGEMM prenorm warmup.Modifications
layers/mhc.py,environ.py): on the first non-capturingmhc_pre()call, compile the kernel for everyn_splitsbucket over[1, chunked_prefill_size]by replayingmhc_pre()with the call's real weights; skipped during CUDA graph capture. Gated bySGLANG_DSV4_MHC_PREWARM(default on). One-time cost ~145 s, absorbed into the capture phase.deep_gemm_wrapper/): drop_TF32HcPrenormWarmupExecutorand its hook; the runtime GEMM path is unchanged.Accuracy Tests
Warmup-only change: it compiles the same kernels the runtime already uses, just at startup. No change to model outputs.
Speed Tests and Profiling
B300, DeepSeek-V4-Pro, TP4, conc 32, ISL=OSL=1024, EAGLE,
flashinfer_mxfp4. Standard2×concwarmup +10×conctimed run, over two separate server launches sharing the on-disk JIT cache —run1= a fresh server's first run (cold disk),run2= warm. Output token throughput (tok/s).Does the DeepGEMM-side warmup (#27729) fix it, and is it still needed with this PR?
#27729 alone lifts the first run +11% but leaves the gap; this PR lifts it +24% (1016→1260) and nearly closes the run1↔run2 gap. This PR ± #27729 is within noise (1260 vs 1269) → #27729 is redundant once this lands.
This PR's two commits together (i.e. removing #26238's warmup):
Removing #26238's DeepGEMM prenorm warmup leaves run1 unchanged with a clean startup and zero lazy kernel compilation during the timed run — confirming the startup replay covers both the TileLang kernel and the DeepGEMM GEMM.
Checklist
CI States
Latest PR Test (Base): ❌ Run #27471417680
Latest PR Test (Extra): ⏳ Run #27574804670