[DeepSeek-V4] Add mhc_fused_post_pre kernel by JoeLee314 · Pull Request #25976 · sgl-project/sglang

JoeLee314 · 2026-05-21T09:01:12Z

[DeepSeek-V4] Add fused mHC post/pre kernel

Motivation

This PR optimizes the DeepSeek-V4 mHC path by fusing the boundary between one mHC post step and the next mHC pre step.

The approach is adapted from TRTLLM. The main idea is to avoid launching separate kernels for the latency-sensitive decode path:

mhc_post -> pre-norm GEMM -> mhc_pre finalization

The fused path keeps the existing numerically sensitive mhc_pre_big_fuse(_with_norm) finalization stage, but removes the separate post/pre scheduling boundary.

For small token batches, a TileLang scalar-FMA kernel computes the following work in one launch:

previous hc_post + bf16 residual materialization + pre-norm GEMM partials + RMS square-sum partials

For larger token batches, the implementation keeps the existing high-throughput TileLang mhc_post + DeepGEMM decomposition.

The optimization is guarded by a new env var:

SGLANG_OPT_FUSE_MHC_POST_PRE=1

It intentionally depends on the existing TileLang mHC switches:

SGLANG_OPT_USE_TILELANG_MHC_PRE=1
SGLANG_OPT_USE_TILELANG_MHC_POST=1

because the fused implementation directly reuses the TileLang mHC tensor layouts and the mhc_pre_big_fuse(_with_norm) finalization kernels.

Performance Test

total_throughput increased by 3.35%

Performance summary

Accuracy / Correctness Tests

GSM8K score: 0.975

Modifications

`python/sglang/srt/environ.py`

Add:

SGLANG_OPT_FUSE_MHC_POST_PRE = EnvBool(False)

`python/sglang/srt/layers/mhc.py`

Add mhc_fused_post_pre_fma_tilelang, a small-token fused TileLang kernel.
Add mhc_fused_post_pre, which dispatches:
- small batches to the fused FMA kernel;
- larger batches to the existing TileLang mhc_post + DeepGEMM path.
Preserve the existing mhc_pre_big_fuse(_with_norm) finalization stage.
Handle zero-token DP/EP ranks with correctly typed empty tensors.

`python/sglang/srt/models/deepseek_v4.py`

Add fused mHC enable check cached in __init__.
Fuse:
- cross-layer previous FFN hc_post + next attention hc_pre;
- within-layer attention hc_post + FFN hc_pre.
Defer the final FFN hc_post and close it at model tail.
Cache bf16 RMSNorm weights for the fused path to avoid per-forward casts.

`python/sglang/srt/models/deepseek_v4_nextn.py`

Adapt NextN decoder usage to the decoder layer returning deferred fused state.

`tests/kernels/test_mhc_kernels.py`

Add coverage for mhc_fused_post_pre vs the unfused reference sequence.

Checklist

Added fused mHC post/pre boundary optimization.
Added env guard SGLANG_OPT_FUSE_MHC_POST_PRE.
Kept the fused path dependent on TileLang mHC pre/post kernels.
Added kernel unit test for fused vs unfused mHC.
Verified operator-level precision.
Verified short deterministic sanity prompts.
Verified GSM8K-200 end-to-end accuracy.
Verified bench_serving performance improvement on DeepSeek-V4-Flash.

CI States

Latest PR Test (Base): ❌ Run #26549213709
Latest PR Test (Extra): ❌ Run #26549213622

gemini-code-assist

Code Review

This pull request introduces kernel fusion for the multi-head latent attention (mHC) post-mapping and pre-norm GEMM stages in DeepSeek-V4 models. It adds a new TileLang JIT kernel optimized for small token batches and a unified wrapper, mhc_fused_post_pre, which supports both within-layer and cross-layer fusion. The review feedback highlights several code quality improvements, including the removal of unused variables, consolidation of redundant local imports, and the clarification of magic numbers in the model implementation.

yhyang201 · 2026-05-26T08:20:29Z

+    monkeypatch.setattr(mhc, "is_dsa_prefill_cp_round_robin_split", lambda: False)
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    hc_mult = 2


Could we change this to hc_mult=4 to match 4's actual config?

yhyang201 · 2026-05-26T08:25:12Z

+from sglang.srt.layers.mhc import mhc_fused_post_pre, mhc_post, mhc_pre
+
+
+@pytest.mark.parametrize("num_tokens", [0, 1, 8, 17, 32])


Can we add a batch size of 64 to test the correctness of the original (unfused large-batch) path?

yhyang201 · 2026-05-26T08:44:49Z

+            hc_mult,
+            hidden_size,
+            hc_mult3,
+            tile_n=tile_n,


It looks like mhc_fused_post_pre_fma_tilelang doesn’t have the parameter tile_n?

You’re right! tile_n should be passed to tile_mix_outputs, I missed that during the refactor.
Thanks for catching it, fixed now.

yhyang201 · 2026-05-26T10:45:42Z

+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    hc_mult = 4
+    hidden_size = 4096


can we parameterize hidden_size and add 7168 (v4pro)?

yhyang201 · 2026-05-26T10:46:36Z

+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
+)
+def mhc_fused_post_pre_fma_tilelang(


Can we add warmup for the fused kernel in prewarm_mhc_token_counts?

yhyang201 · 2026-05-26T10:46:48Z

/tag-and-rerun-ci

Co-authored-by: Qichao Li <liqichao@baidu.com>

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: trigger the same prewarm 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. The in-forward trigger stays as a fallback for entry points that bypass load_weights. 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).

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).

Co-authored-by: Qichao Li <liqichao@baidu.com>

JoeLee314 requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Ying1123, ch-wan, ispobock and merrymercy as code owners May 21, 2026 09:01

github-actions Bot added the deepseek label May 21, 2026

gemini-code-assist Bot reviewed May 21, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/mhc.py Outdated

Comment thread python/sglang/srt/layers/mhc.py Outdated

Comment thread python/sglang/srt/models/deepseek_v4.py Outdated

Comment thread python/sglang/srt/models/deepseek_v4.py Outdated

JoeLee314 force-pushed the add_fused_mhc_post_pre branch 3 times, most recently from a9f629e to a7c2a80 Compare May 21, 2026 14:11

amd-bot mentioned this pull request May 22, 2026

[CI Monitor] Daily Report - 2026-05-22 bingxche/sglang-ci-bot#80

Open

JoeLee314 force-pushed the add_fused_mhc_post_pre branch 2 times, most recently from 3cb8300 to a9f9e28 Compare May 22, 2026 10:36

JoeLee314 changed the title ~~Add mhc_fused_post_pre kernel~~ [DeepSeek-V4] Add mhc_fused_post_pre kernel May 22, 2026

JoeLee314 force-pushed the add_fused_mhc_post_pre branch from a9f9e28 to 2c52dfd Compare May 22, 2026 15:43

This was referenced May 23, 2026

[CI Monitor] Daily Report - 2026-05-23 bingxche/sglang-ci-bot#81

Open

[CI Monitor] Daily Report - 2026-05-26 bingxche/sglang-ci-bot#84

Open

JoeLee314 force-pushed the add_fused_mhc_post_pre branch 2 times, most recently from 846c8b6 to 2abba97 Compare May 26, 2026 07:01

ch-wan assigned yhyang201 May 26, 2026

yhyang201 reviewed May 26, 2026

View reviewed changes

JoeLee314 force-pushed the add_fused_mhc_post_pre branch 3 times, most recently from 5beed4e to a85f425 Compare May 26, 2026 09:36

yhyang201 reviewed May 26, 2026

View reviewed changes

yhyang201 added the run-ci label May 26, 2026

JoeLee314 force-pushed the add_fused_mhc_post_pre branch 2 times, most recently from 6c8bb20 to 053d52a Compare May 27, 2026 12:03

yhyang201 approved these changes May 27, 2026

View reviewed changes

Add mhc_fused_post_pre kernel

48c47c5

JoeLee314 force-pushed the add_fused_mhc_post_pre branch from 053d52a to 48c47c5 Compare May 28, 2026 01:34

ch-wan merged commit 7c5708c into sgl-project:main May 30, 2026
179 of 195 checks passed

yhyang201 mentioned this pull request May 30, 2026

Remove dead prewarm_mhc_token_count_buckets #26787

Open

mqhc2020 pushed a commit to mqhc2020/sglang that referenced this pull request Jun 2, 2026

[DeepSeek-V4] Add mhc_fused_post_pre kernel (sgl-project#25976)

b4957ba

Co-authored-by: Qichao Li <liqichao@baidu.com>

hanming-lu pushed a commit that referenced this pull request Jun 3, 2026

[DeepSeek-V4] Add mhc_fused_post_pre kernel (#25976)

2e25a62

Co-authored-by: Qichao Li <liqichao@baidu.com>

DarkSharpness mentioned this pull request Jun 10, 2026

[CI] Move misplaced mhc kernel test into test/registered/kernels #27781

Merged

3 tasks

YAMY1234 mentioned this pull request Jul 3, 2026

[dsv4] Trigger MHC prenorm prewarm at weight-load time with rank sync #29988

Merged

5 tasks

ormandj mentioned this pull request Jul 12, 2026

[SM120] Allow fused MHC opt-in with standalone TileLang pre disabled #30954

Merged

4 tasks

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[DeepSeek-V4] Add mhc_fused_post_pre kernel (sgl-project#25976)

78755ca

Co-authored-by: Qichao Li <liqichao@baidu.com>

zzjc1234 mentioned this pull request Sep 8, 2026

[DSV4] Use mhc_pre's split-K pre-norm GEMM in the mhc_fused_post_pre fallback #38475

Open

5 tasks

jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026

[DeepSeek-V4] Add mhc_fused_post_pre kernel (sgl-project#25976)

d1bb06c

Co-authored-by: Qichao Li <liqichao@baidu.com>

AndyLi429 mentioned this pull request Sep 10, 2026

DSV4 optimization AndyLi429/sglang#7

Open

		from sglang.srt.layers.mhc import mhc_fused_post_pre, mhc_post, mhc_pre


		@pytest.mark.parametrize("num_tokens", [0, 1, 8, 17, 32])

Uh oh!

Conversation

JoeLee314 commented May 21, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

[DeepSeek-V4] Add fused mHC post/pre kernel

Motivation

Performance Test

Performance summary

Accuracy / Correctness Tests

Modifications

python/sglang/srt/environ.py

python/sglang/srt/layers/mhc.py

python/sglang/srt/models/deepseek_v4.py

python/sglang/srt/models/deepseek_v4_nextn.py

tests/kernels/test_mhc_kernels.py

Checklist

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yhyang201 commented May 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

JoeLee314 commented May 21, 2026 •

edited by github-actions Bot

Loading

`python/sglang/srt/environ.py`

`python/sglang/srt/layers/mhc.py`

`python/sglang/srt/models/deepseek_v4.py`

`python/sglang/srt/models/deepseek_v4_nextn.py`

`tests/kernels/test_mhc_kernels.py`