diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index eb782c22744b..35a62ebe3a2a 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 3c883cd2f3ba..9492090de938 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1988,6 +1988,8 @@ def hc_pre_torch_impl(x, hc_fn): 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 @@ def hc_pre_torch_impl(x, hc_fn): 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( diff --git a/test/registered/kernels/ops/layernorm/test_hc_combine.py b/test/registered/kernels/ops/layernorm/test_hc_combine.py new file mode 100644 index 000000000000..d286b03a015d --- /dev/null +++ b/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