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

python/sglang/kernels/ops/layernorm/mhc.py

@@ -735,6 +735,49 @@ def mhc_pre_gemm_sqrsum_splitk_stage_1(
     )
 
 
+# Token counts above this keep the plain (non split-K) pre-norm GEMM: with that
+# many tokens the ceil(num_tokens / 32) grid already fills the device.
+MHC_PRE_SPLITK_MAX_TOKENS = 2048
+# hc_hidden_size -> hidden_block for mhc_pre_gemm_sqrsum_splitk_kernel.
+MHC_PRE_SPLITK_HIDDEN_BLOCK = {16384: 256, 28672: 128}
+
+
+def _mhc_pre_gemm_sqrsum_splitk(
+    x: torch.Tensor,
+    fn: torch.Tensor,
+    hc_mult3: int,
+    hc_hidden_size: int,
+    split_k: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Split-K pre-norm GEMM and RMS square sums, left as per-split partials.
+
+    Shared by mhc_pre and mhc_fused_post_pre so the two cannot drift apart. The
+    stage_1 reduction is folded into mhc_pre_big_fuse, which is why the caller
+    passes ``split_k`` as its ``n_splits`` and ``32`` as its last GEMM dim.
+    """
+    hidden_block = MHC_PRE_SPLITK_HIDDEN_BLOCK.get(hc_hidden_size)
+    if hidden_block is None:
+        raise NotImplementedError(
+            f"mhc_pre splitk kernel only supports hc_hidden_size in "
+            f"{sorted(MHC_PRE_SPLITK_HIDDEN_BLOCK)}, got {hc_hidden_size}"
+        )
+    kernel_0, _ = mhc_pre_gemm_sqrsum_splitk_kernel(
+        hc_mult3,
+        hc_hidden_size,
+        split_k=split_k,
+        token_block=32,
+        hidden_block=hidden_block,
+    )
+    partial_out = torch.empty(
+        split_k, x.shape[0], 32, dtype=torch.float32, device=x.device
+    )
+    partial_sqrsum = torch.empty(
+        split_k, x.shape[0], dtype=torch.float32, device=x.device
+    )
+    kernel_0(x, fn, partial_out, partial_sqrsum)
+    return partial_out, partial_sqrsum
+
+
 def _compute_num_split_for_mhc_pre(num_tokens: int, hc_hidden_size: int) -> int:
     block_m, block_k = 64, 64
     grid_size = (num_tokens + block_m - 1) // block_m
@@ -1053,43 +1096,16 @@ def mhc_pre(
         gemm_last_dim = hc_mult3
         big_fuse_n_splits = n_splits
     else:
-        if num_tokens <= 2048:
+        if num_tokens <= MHC_PRE_SPLITK_MAX_TOKENS:
             assert n_splits == 1
-            if hc_hidden_size == 16384:
-                hidden_block = 256
-            elif hc_hidden_size == 28672:
-                hidden_block = 128
-            else:
-                raise NotImplementedError(
-                    f"mhc_pre splitk kernel only supports hc_hidden_size in {{16384, 28672}}, "
-                    f"got {hc_hidden_size}"
-                )
-            kernel_0, _ = mhc_pre_gemm_sqrsum_splitk_kernel(
-                hc_mult3,
-                hc_hidden_size,
+            # Stage_1 reduction is folded into big_fuse below; skip launching it.
+            gemm_out_mul, gemm_out_sqrsum = _mhc_pre_gemm_sqrsum_splitk(
+                x=residual_flat.view(num_tokens, hc_hidden_size),
+                fn=fn_flat,
+                hc_mult3=hc_mult3,
+                hc_hidden_size=hc_hidden_size,
                 split_k=n_splits_pre,
-                token_block=32,
-                hidden_block=hidden_block,
-            )
-            partial_out = torch.empty(
-                n_splits_pre,
-                num_tokens,
-                32,
-                dtype=torch.float32,
-                device=residual.device,
-            )
-            partial_sqrsum = torch.empty(
-                n_splits_pre, num_tokens, dtype=torch.float32, device=residual.device
             )
-            kernel_0(
-                residual_flat.view(num_tokens, hc_hidden_size),
-                fn_flat,
-                partial_out,
-                partial_sqrsum,
-            )
-            # Stage_1 reduction is folded into big_fuse below; skip launching it.
-            gemm_out_mul = partial_out
-            gemm_out_sqrsum = partial_sqrsum
             gemm_last_dim = 32
             big_fuse_n_splits = n_splits_pre
         else:
@@ -1520,6 +1536,7 @@ def mhc_fused_post_pre(
     sinkhorn_repeat: int,
     n_splits: int = 1,
     tile_n: int = 1,
+    n_splits_pre: int = 32,
     *,
     norm_weight: torch.Tensor | None = None,
     norm_eps: float | None = None,
@@ -1611,6 +1628,7 @@ def mhc_fused_post_pre(
         device=residual.device,
     )
     residual_cur = torch.empty_like(residual_flat)
+    gemm_last_dim = hc_mult3
 
     if num_tokens <= fma_token_threshold:
         # Small-batch path: one TileLang launch computes hc_post, the bf16
@@ -1653,8 +1671,28 @@ def mhc_fused_post_pre(
                 gemm_out_sqrsum,
                 num_splits=n_splits,
             )
+        elif (
+            num_tokens <= MHC_PRE_SPLITK_MAX_TOKENS
+            and hc_hidden_size in MHC_PRE_SPLITK_HIDDEN_BLOCK
+        ):
+            # Fallback mirrors mhc_pre when DeepGEMM prenorm is disabled: same
+            # split-K kernel, same folded stage_1 reduction. The plain kernel
+            # below launches only ceil(num_tokens / 32) blocks, so a decode
+            # batch just past fma_token_threshold streams the whole fn matrix
+            # through two SMs.
+            n_splits = n_splits_pre
+            gemm_out_mul, gemm_out_sqrsum = _mhc_pre_gemm_sqrsum_splitk(
+                x=residual_cur.view(num_tokens, hc_hidden_size),
+                fn=fn,
+                hc_mult3=hc_mult3,
+                hc_hidden_size=hc_hidden_size,
+                split_k=n_splits_pre,
+            )
+            gemm_last_dim = 32
         else:
-            # Fallback mirrors mhc_pre when DeepGEMM prenorm is disabled.
+            # hc_hidden_size the split-K kernel is not specialized for, or a
+            # batch large enough that ceil(num_tokens / 32) already fills the
+            # device: plain GEMM, as in mhc_pre above MHC_PRE_SPLITK_MAX_TOKENS.
             n_splits = 1
             gemm_out_mul_2d = torch.empty(
                 num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
@@ -1727,7 +1765,7 @@ def mhc_fused_post_pre(
             norm_eps,
             n_splits,
             hc_mult,
-            hc_mult3,
+            gemm_last_dim,
         )
     else:
         # Same mhc_pre finalization without the model-layer RMSNorm.
@@ -1748,7 +1786,7 @@ def mhc_fused_post_pre(
             sinkhorn_repeat,
             n_splits,
             hc_mult,
-            hc_mult3,
+            gemm_last_dim,
         )
 
     return (

test/registered/kernels/ops/layernorm/test_mhc_kernels.py

@@ -5,11 +5,26 @@
 
 import sglang.kernels.ops.layernorm.mhc as mhc
 from sglang.kernels.ops.layernorm.mhc import mhc_fused_post_pre, mhc_post, mhc_pre
+from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
 
 
+def _bypass_tp_group(monkeypatch):
+    """These are single-process kernel unit tests with no TP group initialized.
+
+    mhc_pre / mhc_fused_post_pre allocate the MoE input in the symmetric-memory
+    pool via use_symmetric_memory(get_tp_group(), ...); bypass that path so the
+    kernel runs with a plain torch.empty allocation. Mirrors the workaround in
+    test_mxfp4_sm90_cutlass.py for the same TP-group-not-initialized case.
+    """
+    monkeypatch.setattr(mhc, "is_dsa_prefill_cp_round_robin_split", lambda: False)
+    monkeypatch.setattr(mhc, "use_symmetric_memory", lambda *a, **kw: nullcontext())
+    monkeypatch.setattr(mhc, "is_allocation_symmetric", lambda: False)
+    monkeypatch.setattr(mhc, "get_tp_group", lambda: None)
+
+
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("num_tokens", [0, 1, 8, 17, 32, 64])
 @pytest.mark.parametrize("use_norm", [False, True])
@@ -19,15 +34,7 @@ def test_mhc_fused_post_pre_matches_unfused(
     if not torch.cuda.is_available():
         pytest.skip("CUDA is required for TileLang mHC kernels")
 
-    monkeypatch.setattr(mhc, "is_dsa_prefill_cp_round_robin_split", lambda: False)
-    # This is a single-process kernel unit test with no TP group initialized.
-    # mhc_pre / mhc_fused_post_pre allocate the MoE input in the symmetric-memory
-    # pool via use_symmetric_memory(get_tp_group(), ...); bypass that path so the
-    # kernel runs with a plain torch.empty allocation. Mirrors the workaround in
-    # test_mxfp4_sm90_cutlass.py for the same TP-group-not-initialized case.
-    monkeypatch.setattr(mhc, "use_symmetric_memory", lambda *a, **kw: nullcontext())
-    monkeypatch.setattr(mhc, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(mhc, "get_tp_group", lambda: None)
+    _bypass_tp_group(monkeypatch)
     torch.manual_seed(0)
     device = torch.device("cuda")
     hc_mult = 4
@@ -124,6 +131,101 @@ def test_mhc_fused_post_pre_matches_unfused(
     torch.testing.assert_close(layer_out, layer_ref, atol=layer_atol, rtol=layer_rtol)
 
 
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("num_tokens", [40, 64])
+def test_mhc_fused_post_pre_no_deepgemm_matches_mhc_pre(
+    monkeypatch, hidden_size, num_tokens
+):
+    """Without DeepGEMM the fused boundary must run mhc_pre's pre-norm GEMM.
+
+    num_tokens is above mhc_fused_post_pre's FMA threshold and below mhc_pre's
+    split-K token limit, so both paths belong on the split-K kernel. Taking the
+    plain kernel instead is correct but launches only ceil(num_tokens / 32)
+    blocks, which is a large slowdown at decode batch sizes, so assert on the
+    branch rather than on the outputs alone.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for TileLang mHC kernels")
+
+    _bypass_tp_group(monkeypatch)
+    plain_gemm_calls = 0
+    real_dispatch = mhc._mhc_pre_gemm_sqrsum_dispatch
+
+    def counting_dispatch():
+        nonlocal plain_gemm_calls
+        plain_gemm_calls += 1
+        return real_dispatch()
+
+    monkeypatch.setattr(mhc, "_mhc_pre_gemm_sqrsum_dispatch", counting_dispatch)
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    hc_mult = 4
+    hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
+    hc_hidden_size = hc_mult * hidden_size
+    assert hc_hidden_size in mhc.MHC_PRE_SPLITK_HIDDEN_BLOCK
+    assert num_tokens <= mhc.MHC_PRE_SPLITK_MAX_TOKENS
+
+    x = torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
+    residual = (
+        torch.randn(
+            num_tokens, hc_mult, hidden_size, device=device, dtype=torch.bfloat16
+        )
+        * 0.1
+    )
+    post_prev = torch.rand(num_tokens, hc_mult, 1, device=device, dtype=torch.float32)
+    comb_prev = (
+        torch.rand(num_tokens, hc_mult, hc_mult, device=device, dtype=torch.float32)
+        * 0.25
+    )
+    fn = (
+        torch.randn(hc_mult3, hc_hidden_size, device=device, dtype=torch.float32) * 0.01
+    )
+    hc_scale = torch.tensor([0.5, 0.25, 0.25], device=device, dtype=torch.float32)
+    hc_base = torch.zeros(hc_mult3, device=device, dtype=torch.float32)
+    rms_eps = hc_eps = 1e-6
+    sinkhorn_repeat = 2
+
+    with envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.override(False):
+        residual_ref = mhc_post(x, residual, post_prev, comb_prev)
+        post_ref, comb_ref, layer_ref = mhc_pre(
+            residual_ref,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_eps,
+            hc_eps,
+            2.0,
+            sinkhorn_repeat,
+        )
+        assert plain_gemm_calls == 0, "mhc_pre took the plain pre-norm GEMM"
+
+        residual_out, post_out, comb_out, layer_out = mhc_fused_post_pre(
+            x,
+            residual,
+            post_prev,
+            comb_prev,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_eps,
+            hc_eps,
+            2.0,
+            sinkhorn_repeat,
+        )
+    torch.cuda.synchronize()
+    assert plain_gemm_calls == 0, (
+        "mhc_fused_post_pre took the plain pre-norm GEMM where mhc_pre uses split-K"
+    )
+
+    torch.testing.assert_close(residual_out, residual_ref, atol=0, rtol=0)
+    torch.testing.assert_close(post_out, post_ref, atol=1e-3, rtol=1e-3)
+    torch.testing.assert_close(comb_out, comb_ref, atol=1e-3, rtol=1e-3)
+    torch.testing.assert_close(layer_out, layer_ref, atol=2e-3, rtol=2e-3)
+
+
 if __name__ == "__main__":
     import sys