dsv4: use fused hc head in nextn

Author: bhaktatejas922

python/sglang/srt/models/deepseek_v4_nextn.py

@@ -117,6 +117,17 @@ def hc_head(
         hc_scale: torch.Tensor,
         hc_base: torch.Tensor,
     ):
+        if x.numel() > 0:
+            from sglang.srt.layers.mhc_head import fused_hc_head
+
+            return fused_hc_head(
+                x.contiguous(),
+                hc_fn,
+                hc_scale,
+                hc_base,
+                norm_eps=self.rms_norm_eps,
+                hc_eps=self.hc_eps,
+            )
         shape, dtype = x.size(), x.dtype
         x = x.flatten(1).float()
         rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.rms_norm_eps)

test/registered/layers/test_dsv4_mhc_head.py

@@ -0,0 +1,71 @@
+import unittest
+
+import torch
+
+from sglang.srt.models.deepseek_v4_nextn import DeepseekV4ModelNextN
+from sglang.test.ci.ci_register import register_cuda_ci
+
+
+register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-large")
+
+
+def _reference_hc_head(
+    x: torch.Tensor,
+    hc_fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    norm_eps: float,
+    hc_eps: float,
+) -> torch.Tensor:
+    shape, dtype = x.size(), x.dtype
+    flat = x.flatten(1).float()
+    rsqrt = torch.rsqrt(flat.square().mean(-1, keepdim=True) + norm_eps)
+    mixes = torch.nn.functional.linear(flat, hc_fn) * rsqrt
+    pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
+    out = torch.sum(pre.unsqueeze(-1) * flat.view(shape), dim=1)
+    return out.to(dtype)
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestDeepseekV4NextNHcHead(unittest.TestCase):
+    def _run_case(self, tokens: int, hc_mult: int, hidden_size: int) -> None:
+        torch.manual_seed(1234 + tokens + hidden_size)
+        device = "cuda"
+        dtype = torch.bfloat16
+        norm_eps = 1.0e-6
+        hc_eps = 1.0e-3
+
+        model = object.__new__(DeepseekV4ModelNextN)
+        model.rms_norm_eps = norm_eps
+        model.hc_eps = hc_eps
+
+        x = torch.randn(
+            tokens, hc_mult, hidden_size, device=device, dtype=torch.float32
+        ).to(dtype)
+        hc_fn = torch.randn(
+            hc_mult,
+            hc_mult * hidden_size,
+            device=device,
+            dtype=torch.float32,
+        ) * 0.02
+        hc_scale = torch.randn(1, device=device, dtype=torch.float32)
+        hc_base = torch.randn(hc_mult, device=device, dtype=torch.float32)
+
+        expected = _reference_hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps)
+        actual = DeepseekV4ModelNextN.hc_head(model, x, hc_fn, hc_scale, hc_base)
+
+        self.assertEqual(actual.shape, (tokens, hidden_size))
+        self.assertEqual(actual.dtype, dtype)
+        torch.testing.assert_close(
+            actual.float(), expected.float(), rtol=3.0e-2, atol=3.0e-2
+        )
+
+    def test_nextn_hc_head_uses_fused_kernel_at_dsv4_shape(self):
+        self._run_case(tokens=16, hc_mult=4, hidden_size=7168)
+
+    def test_nextn_hc_head_handles_empty_batch(self):
+        self._run_case(tokens=0, hc_mult=4, hidden_size=256)
+
+
+if __name__ == "__main__":
+    unittest.main()