[DSV4] hc-prenorm: fuse the combine step into a Triton kernel (#35118)

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

@@ -1789,3 +1789,53 @@ def npu_hc_pre(
     # not fold input_layernorm. Return norm_fused=False so the caller
     # applies the layernorm itself, matching the deepgemm/torch paths.
     return y.to(dtype), post, comb, False
+
+
+@triton.jit
+def _hc_combine_kernel(
+    x_ptr,
+    pre_ptr,
+    y_ptr,
+    H,
+    x_stride_m,
+    pre_stride_m,
+    pre_stride_k,
+    y_stride_m,
+    HC: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+    acc = tl.zeros([BLOCK_H], dtype=tl.float32)
+    for k in tl.static_range(HC):
+        pk = tl.load(pre_ptr + pid_m * pre_stride_m + k * pre_stride_k).to(tl.float32)
+        xv = tl.load(
+            x_ptr + pid_m * x_stride_m + k * H + offs_h, mask=mask, other=0.0
+        ).to(tl.float32)
+        acc += pk * xv
+    tl.store(y_ptr + pid_m * y_stride_m + offs_h, acc, mask=mask)
+
+
+def hc_combine(
+    x_flat: torch.Tensor, pre: torch.Tensor, hc: int, out_dtype: torch.dtype
+) -> torch.Tensor:
+    """Fused y[m, h] = sum_k pre[m, k] * x_flat[m, k*H + h]."""
+    m = x_flat.shape[0]
+    h = x_flat.shape[1] // hc
+    y = torch.empty((m, h), dtype=out_dtype, device=x_flat.device)
+    block_h = 1024
+    _hc_combine_kernel[(m, triton.cdiv(h, block_h))](
+        x_flat,
+        pre,
+        y,
+        h,
+        x_flat.stride(0),
+        pre.stride(0),
+        pre.stride(1),
+        y.stride(0),
+        HC=hc,
+        BLOCK_H=block_h,
+    )
+    return y

python/sglang/srt/models/deepseek_v4.py

@@ -1988,6 +1988,8 @@ class DeepseekV4DecoderLayer(nn.Module):
             self.hc_sinkhorn_iters,
             self.hc_eps,
         )
+        from sglang.kernels.ops.layernorm.mhc import hc_combine
+
         # y is the post-norm activation fed into the MoE. Allocate it in the
         # symmetric memory pool so the downstream all-reduce uses the low-latency
         # NCCL symmetric path: the Triton inplace MoE runner writes the expert
@@ -1997,7 +1999,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         with use_symmetric_memory(
             get_tp_group(), disabled=not is_allocation_symmetric()
         ):
-            y = (pre.squeeze(1).unsqueeze(-1) * x_flat.view(shape)).sum(dim=1).to(dtype)
+            y = hc_combine(x_flat, pre.squeeze(1), self.hc_mult, dtype)
         return y, post.squeeze(1), comb.squeeze(1), False
 
     def hc_post(

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

@@ -0,0 +1,53 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.layernorm.mhc import hc_combine
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.cuda is None,
+    reason="hc_combine requires CUDA",
+)
+
+
+def _reference(
+    x_flat: torch.Tensor, pre: torch.Tensor, hc: int, out_dtype: torch.dtype
+) -> torch.Tensor:
+    m, h = x_flat.shape[0], x_flat.shape[1] // hc
+    return (pre.unsqueeze(-1) * x_flat.view(m, hc, h)).sum(dim=1).to(out_dtype)
+
+
+@pytest.mark.parametrize("m", [1, 7, 64, 1024, 2048])
+@pytest.mark.parametrize("hc,h", [(4, 128), (4, 512), (8, 256)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_hc_combine_matches_reference(m: int, hc: int, h: int, dtype: torch.dtype):
+    torch.manual_seed(0)
+    x_flat = torch.randn(m, hc * h, device="cuda", dtype=dtype)
+    pre = torch.randn(m, hc, device="cuda", dtype=dtype).contiguous()
+
+    got = hc_combine(x_flat, pre, hc, dtype)
+    ref = _reference(x_flat, pre, hc, dtype)
+
+    assert got.shape == (m, h)
+    assert got.dtype == dtype
+    # hc_combine accumulates in fp32 while the reference accumulates in `dtype`,
+    # so compare against the fp32 result the kernel is approximating.
+    ref_fp32 = _reference(x_flat.float(), pre.float(), hc, torch.float32)
+    scale = ref_fp32.abs().max().clamp(min=1e-6)
+    assert (got.float() - ref_fp32).abs().max() / scale < 5e-2
+    assert (ref.float() - got.float()).abs().max() / scale < 5e-2
+
+
+def test_hc_combine_strided_pre():
+    torch.manual_seed(0)
+    m, hc, h = 32, 4, 128
+    x_flat = torch.randn(m, hc * h, device="cuda", dtype=torch.bfloat16)
+    # A non-contiguous view of pre, to check the kernel honours pre's strides.
+    pre = torch.randn(m, hc, 2, device="cuda", dtype=torch.bfloat16)[..., 0]
+    assert not pre.is_contiguous()
+    got = hc_combine(x_flat, pre, hc, torch.bfloat16)
+    ref = _reference(x_flat, pre, hc, torch.bfloat16)
+    scale = ref.float().abs().max().clamp(min=1e-6)
+    assert (got.float() - ref.float()).abs().max() / scale < 5e-2