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

[zzjc1234]

Motivation

mhc_fused_post_pre is the mHC layer boundary DeepSeek-V4 runs between every pair of layers. Above fma_token_threshold = 32 tokens it splits into mhc_post + a pre-norm GEMM, and when SGLANG_OPT_DEEPGEMM_HC_PRENORM is off it takes a fallback whose comment reads:

else:
    # Fallback mirrors mhc_pre when DeepGEMM prenorm is disabled.

It does not mirror mhc_pre. In the same file and the same configuration, mhc_pre runs mhc_pre_gemm_sqrsum_splitk_kernel for any batch up to 2048 tokens, while this fallback calls _mhc_pre_gemm_sqrsum_dispatch() — the non-split kernel, whose grid is T.Kernel(T.ceildiv(num_tokens, token_block)) with token_block=32. At a decode batch of 40 that is 2 thread blocks, so 2 SMs stream the whole 1.5 MiB fp32 fn matrix while the rest of the device idles.

Measured standalone at DeepSeek-V4-Flash shapes (hidden_size=4096, hc_mult=4, so hc_hidden_size=16384), 40 tokens:

	RTX PRO 6000 Blackwell (sm_120)	RTX 4090 (sm_89)
mhc_fused_post_pre, plain fallback	84.1 us	103.5 us
mhc_fused_post_pre, mhc_pre's split-K	12.4 us	12.5 us

The op itself only has ~12 us of work in it: at 40 tokens the input is 40 x 4 x 4096 bf16 (1.3 MB) and the weights are 1.5 MiB, so the whole boundary is a few microseconds of traffic plus launch overhead. The remaining 72 us is one kernel occupying two SMs. Over the 43 layers of DeepSeek-V4-Flash that is 3.1 ms per decode step on sm_120 that no configuration needs to pay.

Who reaches this. A DeepSeek-V4 decode batch above 32 with the DeepGEMM pre-norm path off, i.e. any of:

• SGLANG_OPT_DEEPGEMM_HC_PRENORM=0 set directly;
• SGLANG_ENABLE_JIT_DEEPGEMM=0, which is exactly the configuration in [Bug] NameError: name 'deep_gemm' is not defined in tf32_hc_prenorm_gemm when SGLANG_ENABLE_JIT_DEEPGEMM=0 #29738 (NameError: name 'deep_gemm' is not defined in tf32_hc_prenorm_gemm, being fixed in Fix DeepGEMM import for independently enabled MHC prenorm #38442) — turning the pre-norm flag off is the workaround people land on;
• a DeepGEMM build without support for the running arch.

Batch <= 32 takes the FMA path and never sees it, which is why a batch-size-dependent 6x on a layer-boundary op has gone unreported.

The two have differed since mhc_fused_post_pre was introduced in #25976 — mhc_pre had already moved to the split-K kernel in #23882 three weeks earlier. It cost nothing at the time, because the fused path was opt-in; #34019 and #35214 made it the default in August, and today's main (ccfa120dae) still takes the plain kernel here.

Root cause

mhc_pre and mhc_fused_post_pre each open-code the same "no DeepGEMM" decision, and only mhc_pre was updated when the split-K kernel landed.

mhc_pre:

if num_tokens <= 2048:
    ...
    kernel_0, _ = mhc_pre_gemm_sqrsum_splitk_kernel(..., split_k=n_splits_pre, ...)
    gemm_last_dim = 32                 # partials, stage_1 folded into big_fuse
    big_fuse_n_splits = n_splits_pre   # 32

mhc_fused_post_pre:

n_splits = 1
_mhc_pre_gemm_sqrsum_dispatch()(...)   # grid = ceil(num_tokens / 32)

The two kernels compute the same thing; they differ only in how the hc_hidden_size = 16384 reduction is parallelized. The split-K kernel launches (ceil(num_tokens / 32), split_k) blocks with split_k = 32, so each block reads 1/32 of fn and the reduction over splits is folded into mhc_pre_big_fuse, which already takes an n_splits argument. The plain kernel gives one block per 32 tokens and no split at all, so its occupancy is set entirely by the batch size.

CUDA profile of the boundary at 40 tokens on sm_120 (us per call):

before:  76.6  mhc_pre_gemm_sqrsum_tilelang_kernel
          4.0  mhc_pre_big_fuse_with_norm_tilelang_kernel
          3.1  mhc_post_tilelang_kernel

after:    4.3  mhc_pre_big_fuse_with_norm_tilelang_kernel
          4.1  mhc_pre_gemm_sqrsum_splitk_stage_0_kernel
          3.0  mhc_post_tilelang_kernel

Nothing else moves: the split-K partials cost big_fuse 0.3 us more, and the 76.6 us kernel becomes a 4.1 us one.

Modifications

Extract the split-K launch into _mhc_pre_gemm_sqrsum_splitk, shared by both call sites, so the branch cannot drift again — the duplication is what caused this. The kernel selection constants come with it:

MHC_PRE_SPLITK_MAX_TOKENS = 2048
MHC_PRE_SPLITK_HIDDEN_BLOCK = {16384: 256, 28672: 128}

mhc_pre keeps its exact behaviour; its inline if hc_hidden_size == 16384 / elif 28672 / else raise moves into the helper unchanged. mhc_fused_post_pre takes the helper on the same condition, tracks the resulting last GEMM dim in gemm_last_dim (32 for split-K partials, hc_mult3 otherwise) the way mhc_pre already does, and gains the same n_splits_pre: int = 32 parameter.

One deliberate asymmetry: for an hc_hidden_size the split-K kernel is not specialized for, mhc_pre raises NotImplementedError but mhc_fused_post_pre has always accepted any size through the plain kernel. The new branch is therefore an elif guarded on hc_hidden_size in MHC_PRE_SPLITK_HIDDEN_BLOCK, and unqualified shapes keep the plain kernel rather than starting to raise. Batches above MHC_PRE_SPLITK_MAX_TOKENS also keep it, matching mhc_pre and the measurements below.

No new JIT compilation: mhc_pre_gemm_sqrsum_splitk_kernel is functools.cached on (hc_mult3, hc_hidden_size, split_k, token_block, hidden_block) and the fallback passes the arguments mhc_pre already passes, so prewarm_mhc_pre covers it.

test_mhc_fused_post_pre_no_deepgemm_matches_mhc_pre asserts the branch, not just the numbers. With SGLANG_OPT_DEEPGEMM_HC_PRENORM overridden off it counts _mhc_pre_gemm_sqrsum_dispatch calls across an mhc_post + mhc_pre reference and an mhc_fused_post_pre call at 40 and 64 tokens, and requires zero from both. This is the assertion that fails on the pre-fix code — a numeric-only test would not, because the plain kernel is correct, just slow:

E  AssertionError: mhc_fused_post_pre took the plain pre-norm GEMM where mhc_pre uses split-K
E  assert 1 == 0

It then checks the outputs against the unfused sequence at the tolerances the neighbouring test uses. The TP-group bypass the existing test needed is factored into _bypass_tp_group and shared.

Accuracy Tests

Switching a reduction from one accumulator to 32 partials changes fp32 summation order, so the honest question is not whether the two kernels agree bit-for-bit but which is closer to the truth. Against a float64 reference of the same mhc_post -> mhc_pre sequence (both kernels consume a bit-identical bf16 residual_cur, so this isolates the GEMM), mean |diff| / mean |ref| on sm_120:

tokens	post plain	post split-K	comb plain	comb split-K	layer_input plain	layer_input split-K
40	5.855e-05	5.530e-05	8.404e-05	7.905e-05	1.413e-03	1.413e-03
64	5.648e-05	5.331e-05	8.261e-05	7.761e-05	1.413e-03	1.413e-03
512	5.473e-05	5.158e-05	7.878e-05	7.422e-05	1.410e-03	1.410e-03
2048	5.342e-05	5.021e-05	8.050e-05	7.575e-05	1.410e-03	1.410e-03

The split-K path is consistently more accurate on post and comb — pairwise summation over 32 partials loses less than one long accumulator — and identical on layer_input, whose error is bf16 output rounding, not the GEMM. residual is bit-identical (it comes from mhc_post, untouched). Direct plain-vs-split-K deltas are correspondingly small: post <= 1.7e-05 and comb <= 7.8e-06 absolute at 40 tokens, and 327 of 163840 layer_input elements differ by more than one bf16 ulp.

Unit tests, RTX PRO 6000 Blackwell (sm_120), test/registered/kernels/ops/layernorm/test_mhc_kernels.py:

repo defaults (DeepGEMM pre-norm on):        28 passed
SGLANG_OPT_DEEPGEMM_HC_PRENORM=0:            28 passed

28 = the 24 existing cases plus the 4 new ones. Both configurations were run because the changed branch is only reachable in the second.

No model eval was run. The DeepGEMM path — the default, and what CI exercises — is untouched, and on the fallback the change is a reduction-order difference that measures closer to a float64 reference than what it replaces.

Speed Tests and Profiling

mhc_fused_post_pre alone, CUDA-graph replay, 200 replays, hidden_size=4096 / hc_mult=4 / norm_weight set, SGLANG_OPT_DEEPGEMM_HC_PRENORM=0. "before" is the plain fallback, "after" is split-K; A/B in one process so the two share compilation and allocator state.

RTX PRO 6000 Blackwell (sm_120):

tokens	before	after	speedup
24 (FMA path, unchanged)	16.5 us	16.5 us	1.00x
33	84.1 us	12.4 us	6.80x
40	84.1 us	12.4 us	6.80x
64	84.1 us	12.4 us	6.80x
128	85.2 us	12.4 us	6.89x
512	94.3 us	30.8 us	3.06x
2048	186.0 us	158.7 us	1.17x

RTX 4090 (sm_89):

tokens	before	after	speedup
24 (FMA path, unchanged)	14.4 us	14.4 us	1.00x
33	103.5 us	12.5 us	8.30x
40	103.5 us	12.5 us	8.28x
64	104.0 us	12.8 us	8.11x
128	104.7 us	14.5 us	7.24x
512	109.3 us	42.0 us	2.60x
2048	272.2 us	269.1 us	1.01x

The shape of the curve is the diagnosis. Below 32 tokens nothing changes, because that is the FMA path. From 33 to 128 the plain kernel is flat at ~84 us regardless of batch — it is not doing more work, it is doing the same work on 2 to 4 SMs — while split-K is flat at 12.4 us. By 512 tokens the plain grid is wide enough to matter and the gap narrows; by 2048, MHC_PRE_SPLITK_MAX_TOKENS, they converge, which is why both mhc_pre and this fallback stop using split-K there.

For DeepSeek-V4-Flash (43 layers) at a decode batch of 40 on sm_120, 43 x 71.7 us = 3.1 ms per decode step returned to the model. Standalone the kernel has the device to itself; in a real forward pass it is competing for SMs with the rest of the layer, so this is a floor rather than an estimate. No full-model end-to-end benchmark is claimed here — the numbers above are main plus this patch and are exactly what the attached script reproduces.

Repro (both tables and the accuracy table come from this script):

# /tmp/bench_mhc_fused_post_pre.py, attached below
SGLANG_OPT_DEEPGEMM_HC_PRENORM=0 python3 /tmp/bench_mhc_fused_post_pre.py
BENCH_ACCURACY=1 SGLANG_OPT_DEEPGEMM_HC_PRENORM=0 python3 /tmp/bench_mhc_fused_post_pre.py 40 64 512 2048
BENCH_KERNELS=1  SGLANG_OPT_DEEPGEMM_HC_PRENORM=0 python3 /tmp/bench_mhc_fused_post_pre.py 40

bench_mhc_fused_post_pre.py

"""Repro for the mhc_fused_post_pre pre-norm GEMM fallback.

Times the layer-boundary op at DeepSeek-V4-Flash decode shapes with the
DeepGEMM pre-norm path disabled, which is the configuration that reaches the
fallback branch.  On a patched tree it also A/Bs against the old plain-GEMM
fallback by emptying MHC_PRE_SPLITK_HIDDEN_BLOCK, and reports the output delta.

    SGLANG_OPT_DEEPGEMM_HC_PRENORM=0 python bench_mhc_fused_post_pre.py
"""

import contextlib
import os
import sys

import torch

os.environ.setdefault("SGLANG_OPT_DEEPGEMM_HC_PRENORM", "0")

import sglang.kernels.ops.layernorm.mhc as mhc  # noqa: E402

# Single-process kernel bench: no TP group, no symmetric-memory pool.
mhc.use_symmetric_memory = lambda *a, **k: contextlib.nullcontext()
mhc.get_tp_group = lambda: None
mhc.is_allocation_symmetric = lambda: False
mhc.is_dsa_prefill_cp_round_robin_split = lambda: False

# DeepSeek-V4-Flash: hidden 4096, hc_mult 4 -> hc_hidden_size 16384.
HIDDEN, HC_MULT, SINKHORN, EPS = 4096, 4, 20, 1e-6
HC_MULT3 = HC_MULT * 2 + HC_MULT * HC_MULT
HAS_SPLITK = hasattr(mhc, "MHC_PRE_SPLITK_HIDDEN_BLOCK")


def make_inputs(num_tokens):
    torch.manual_seed(0)
    d = "cuda"
    return dict(
        x=torch.randn(num_tokens, HIDDEN, device=d, dtype=torch.bfloat16) * 0.1,
        residual=torch.randn(
            num_tokens, HC_MULT, HIDDEN, device=d, dtype=torch.bfloat16
        )
        * 0.1,
        post_layer_mix=torch.rand(num_tokens, HC_MULT, 1, device=d),
        comb_res_mix=torch.rand(num_tokens, HC_MULT, HC_MULT, device=d) / HC_MULT,
        fn=torch.randn(HC_MULT3, HC_MULT * HIDDEN, device=d) * 0.01,
        hc_scale=torch.tensor([0.5, 0.25, 0.25], device=d),
        hc_base=torch.randn(HC_MULT3, device=d) * 0.1,
        norm_weight=torch.ones(HIDDEN, device=d, dtype=torch.bfloat16),
    )


def call(t):
    return mhc.mhc_fused_post_pre(
        t["x"],
        t["residual"],
        t["post_layer_mix"],
        t["comb_res_mix"],
        t["fn"],
        t["hc_scale"],
        t["hc_base"],
        EPS,
        EPS,
        EPS,
        2.0,
        SINKHORN,
        norm_weight=t["norm_weight"],
        norm_eps=EPS,
    )


def time_us(t, reps=200):
    for _ in range(5):
        call(t)
    torch.cuda.synchronize()
    g, s = torch.cuda.CUDAGraph(), torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(2):
            call(t)
        with torch.cuda.graph(g, stream=s):
            call(t)
    torch.cuda.current_stream().wait_stream(s)
    torch.cuda.synchronize()
    a, b = (torch.cuda.Event(enable_timing=True) for _ in range(2))
    a.record()
    for _ in range(reps):
        g.replay()
    b.record()
    torch.cuda.synchronize()
    return a.elapsed_time(b) / reps * 1e3


@contextlib.contextmanager
def force_plain_gemm():
    """Route the fallback to the pre-fix plain kernel."""
    saved = mhc.MHC_PRE_SPLITK_HIDDEN_BLOCK
    mhc.MHC_PRE_SPLITK_HIDDEN_BLOCK = {}
    try:
        yield
    finally:
        mhc.MHC_PRE_SPLITK_HIDDEN_BLOCK = saved


def fp64_reference(t):
    """mhc_post -> mhc_pre in float64, without the fused RMSNorm.

    Both kernels consume the same bf16 residual_cur, so this isolates what the
    pre-norm GEMM and its reduction contribute to the error.
    """
    x, res = t["x"].double(), t["residual"].double()
    residual_cur = (
        t["post_layer_mix"].double() * x.unsqueeze(1)
        + (t["comb_res_mix"].double().unsqueeze(-1) * res.unsqueeze(2)).sum(dim=1)
    ).bfloat16()  # the kernels round here too; residual_cur is bit-identical

    n = HC_MULT
    flat = residual_cur.view(residual_cur.shape[0], -1).double()
    rsqrt = torch.rsqrt(flat.square().mean(-1, keepdim=True) + EPS)
    mixes = (flat @ t["fn"].double().T) * rsqrt
    scale, base = t["hc_scale"].double(), t["hc_base"].double()
    pre = torch.sigmoid(mixes[:, :n] * scale[0] + base[:n]) + EPS
    post = 2.0 * torch.sigmoid(mixes[:, n : 2 * n] * scale[1] + base[n : 2 * n])
    comb = mixes[:, 2 * n :].view(-1, n, n) * scale[2] + base[2 * n :].view(n, n)
    comb = comb.softmax(-1) + EPS
    comb = comb / (comb.sum(-2, keepdim=True) + EPS)
    for _ in range(SINKHORN - 1):
        comb = comb / (comb.sum(-1, keepdim=True) + EPS)
        comb = comb / (comb.sum(-2, keepdim=True) + EPS)
    layer_input = (pre.unsqueeze(-1) * residual_cur.double()).sum(dim=1)
    return post.unsqueeze(-1), comb, layer_input


def accuracy(counts):
    """Error of each fallback against a float64 reference (no RMSNorm)."""
    print("Error vs float64 reference (mean |diff| / mean |ref|):")
    for n in counts:
        t = make_inputs(n)
        with torch.inference_mode():
            new = mhc.mhc_fused_post_pre(
                t["x"], t["residual"], t["post_layer_mix"], t["comb_res_mix"],
                t["fn"], t["hc_scale"], t["hc_base"], EPS, EPS, EPS, 2.0, SINKHORN,
            )
            with force_plain_gemm():
                ref = mhc.mhc_fused_post_pre(
                    t["x"], t["residual"], t["post_layer_mix"], t["comb_res_mix"],
                    t["fn"], t["hc_scale"], t["hc_base"], EPS, EPS, EPS, 2.0, SINKHORN,
                )
        gold = fp64_reference(t)
        out = [f"tokens={n:5d}"]
        for name, g, a, b in zip(("post", "comb", "layer_input"), gold, ref[1:], new[1:]):
            rel = lambda v: ((v.double() - g).abs().mean() / g.abs().mean()).item()
            out.append(f"{name}: plain {rel(a):.3e}  split-K {rel(b):.3e}")
        print("  " + "   ".join(out))


def kernels(t, tag):
    from torch.profiler import ProfilerActivity, profile

    for _ in range(3):
        call(t)
    torch.cuda.synchronize()
    with profile(activities=[ProfilerActivity.CUDA]) as prof:
        for _ in range(20):
            call(t)
        torch.cuda.synchronize()
    rows = sorted(
        (
            (e.key[:60], e.device_time_total / 20)
            for e in prof.key_averages()
            if e.device_time_total > 0 and "Memcpy" not in e.key
        ),
        key=lambda r: -r[1],
    )
    print(f"  [{tag}] top CUDA kernels, us/call:")
    for name, us in rows[:4]:
        print(f"      {us:8.1f}  {name}")


def main():
    counts = [int(a) for a in sys.argv[1:]] or [24, 33, 40, 64, 128, 512, 2048]
    print(torch.cuda.get_device_name(0), "| split-K fallback:", HAS_SPLITK)
    if HAS_SPLITK and os.environ.get("BENCH_ACCURACY"):
        accuracy(counts)
    names = ("residual", "post", "comb", "layer_input")
    for n in counts:
        t = make_inputs(n)
        after = time_us(t)
        if not HAS_SPLITK:
            print(f"tokens={n:5d}  main {after:8.1f} us")
            continue
        new = [o.clone() for o in call(t)]
        with force_plain_gemm():
            before = time_us(t)
            ref = [o.clone() for o in call(t)]
        deltas = []
        for name, a, b in zip(names, ref, new):
            af, bf = a.float(), b.float()
            d = (af - bf).abs()
            if a.dtype == torch.bfloat16:
                # bf16 keeps 8 significand bits, so one ulp is a 2**-8
                # relative step; count elements that differ by more than that.
                over = (d > 2.0**-8 * af.abs()).sum().item()
                deltas.append(f"{name} {over}/{a.numel()} over 1 bf16 ulp")
            else:
                deltas.append(f"{name} max {d.max().item():.1e} abs")
        print(
            f"tokens={n:5d}  before {before:8.1f} us   after {after:8.1f} us"
            f"   {before / after:5.2f}x   " + ", ".join(deltas)
        )
        if os.environ.get("BENCH_KERNELS"):
            with force_plain_gemm():
                kernels(t, "before")
            kernels(t, "after")


main()

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

Happy to open a tracking issue if maintainers prefer one. If the benchmark script is worth keeping in-tree I can add it under a benchmark directory instead of inlining it here.

---

CI States

Latest PR Test (Base): ❌ Run #34210698260
Latest PR Test (Extra): ❌ Run #34210697766
Latest PR Test (AMD ROCm 7.2): ❌ Run #34210698066