[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(