Optimize LongCat-Flash router GEMM with the HPC-Ops bf16xfp32 kernel (#30247)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com>

test/registered/gemm/test_linear_bf16_fp32_hpc.py

@@ -0,0 +1,75 @@
+"""Numerical tests for the HPC-Ops bf16xfp32 router GEMM path.
+
+Validates sglang.jit_kernel.dsv4.linear_bf16_fp32's HPC-Ops branch against
+the fp32 reference on the LongCat-Flash router shapes. Skipped when HPC-Ops
+(https://github.com/Tencent/hpc-ops) is not installed or the GPU is not
+Hopper (the kernels ship sm90a only).
+"""
+
+import unittest
+
+import torch
+
+from sglang.jit_kernel.dsv4.gemm import (
+    _hpc_gemm_bf16xfp32_available,
+    _linear_bf16_fp32_hpc,
+    linear_bf16_fp32,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large")
+
+# (hidden_size, n_routed_experts + zero experts) for LongCat-Flash Chat / Lite.
+_ROUTER_SHAPES = ((6144, 768), (3072, 384))
+
+
+@unittest.skipUnless(
+    _hpc_gemm_bf16xfp32_available(),
+    "requires HPC-Ops (https://github.com/Tencent/hpc-ops) and a Hopper GPU",
+)
+class TestLinearBf16Fp32Hpc(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        torch.manual_seed(0)
+
+    def test_matches_fp32_reference(self):
+        for k, n in _ROUTER_SHAPES:
+            for m in (8, 64, 512):
+                with self.subTest(m=m, k=k, n=n):
+                    x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
+                    w = torch.randn(n, k, dtype=torch.float32, device="cuda")
+                    out = _linear_bf16_fp32_hpc(x, w)
+                    self.assertIsNotNone(out)
+                    self.assertEqual(out.dtype, torch.float32)
+                    ref = torch.mm(x.float(), w.t())
+                    torch.testing.assert_close(out, ref, rtol=0.08, atol=0.01)
+
+    def test_min_m_dispatch(self):
+        k, n = _ROUTER_SHAPES[0]
+        w = torch.randn(n, k, dtype=torch.float32, device="cuda")
+        below = torch.randn(4, k, dtype=torch.bfloat16, device="cuda")
+        self.assertIsNone(_linear_bf16_fp32_hpc(below, w, min_m=8))
+        # The public entry falls back to cublas below min_m and still
+        # returns the correct fp32 result.
+        out = linear_bf16_fp32(below, w, hpc_kernel_min_m=8)
+        torch.testing.assert_close(
+            out, torch.mm(below.float(), w.t()), rtol=0.08, atol=0.01
+        )
+
+    def test_weight_split_cache_reused(self):
+        k, n = _ROUTER_SHAPES[1]
+        x = torch.randn(16, k, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(n, k, dtype=torch.float32, device="cuda")
+        out1 = _linear_bf16_fp32_hpc(x, w)
+        cache = getattr(w, "_sglang_bf16xfp32_weight_cache")
+        out2 = _linear_bf16_fp32_hpc(x, w)
+        self.assertIs(getattr(w, "_sglang_bf16xfp32_weight_cache"), cache)
+        torch.testing.assert_close(out1, out2)
+        # The kernel leaves the cached split-K workspace zeroed.
+        self.assertTrue((cache[3] == 0).all().item())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/models/test_longcat_flash_router_hpc_gemm.py

@@ -0,0 +1,116 @@
+"""Unit tests for LongCat-Flash router GEMM dispatch to the HPC-Ops bf16xfp32 kernel."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.models.longcat_flash import LongcatFlashRouter  # noqa: E402
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _longcat_config(hidden_size, n_routed_experts, *, router_bias=False):
+    return SimpleNamespace(
+        hidden_size=hidden_size,
+        n_routed_experts=n_routed_experts,
+        router_bias=router_bias,
+    )
+
+
+class TestLongcatFlashRouterHpcGemm(CustomTestCase):
+    def _assert_dispatches_to_hpc_gemm(
+        self,
+        *,
+        hidden_size,
+        n_routed_experts,
+        zero_expert_num,
+        expected_min_m,
+    ):
+        router = LongcatFlashRouter(
+            _longcat_config(hidden_size, n_routed_experts),
+            zero_expert_num=zero_expert_num,
+            rounter_params_dtype=torch.float32,
+        )
+        hidden_states = torch.randn((4, hidden_size), dtype=torch.bfloat16)
+        expected = torch.randn(
+            (4, n_routed_experts + zero_expert_num), dtype=torch.float32
+        )
+
+        with patch(
+            "sglang.srt.models.longcat_flash.linear_bf16_fp32",
+            return_value=expected,
+        ) as mock_linear:
+            out = router(hidden_states)
+
+        self.assertIs(out, expected)
+        mock_linear.assert_called_once()
+        args, kwargs = mock_linear.call_args
+        self.assertIs(args[0], hidden_states)
+        self.assertIs(args[1], router.classifier.weight)
+        self.assertEqual(kwargs["hpc_kernel_min_m"], expected_min_m)
+
+    def _assert_uses_classifier(self, router, hidden_size):
+        hidden_states = torch.randn((4, hidden_size), dtype=torch.bfloat16)
+        expected = torch.randn((4, router.n_routed_experts), dtype=torch.float32)
+
+        with (
+            patch(
+                "sglang.srt.models.longcat_flash.linear_bf16_fp32",
+                side_effect=AssertionError("unexpected hpc kernel dispatch"),
+            ),
+            patch.object(
+                router.classifier,
+                "forward",
+                return_value=(expected, None),
+            ) as mock_classifier,
+        ):
+            out = router(hidden_states)
+
+        self.assertIs(out, expected)
+        mock_classifier.assert_called_once()
+        self.assertEqual(mock_classifier.call_args.args[0].dtype, torch.float32)
+
+    def test_chat_shape_dispatches_with_benchmark_guard(self):
+        self._assert_dispatches_to_hpc_gemm(
+            hidden_size=6144,
+            n_routed_experts=512,
+            zero_expert_num=256,
+            expected_min_m=64,
+        )
+
+    def test_lite_shape_dispatches_with_benchmark_guard(self):
+        self._assert_dispatches_to_hpc_gemm(
+            hidden_size=3072,
+            n_routed_experts=256,
+            zero_expert_num=128,
+            expected_min_m=128,
+        )
+
+    def test_unbenchmarked_shape_uses_classifier(self):
+        router = LongcatFlashRouter(
+            _longcat_config(4096, 256),
+            zero_expert_num=128,
+            rounter_params_dtype=torch.float32,
+        )
+
+        self._assert_uses_classifier(router, hidden_size=4096)
+
+    def test_router_bias_uses_classifier(self):
+        router = LongcatFlashRouter(
+            _longcat_config(6144, 512, router_bias=True),
+            zero_expert_num=256,
+            rounter_params_dtype=torch.float32,
+        )
+
+        self._assert_uses_classifier(router, hidden_size=6144)
+
+
+if __name__ == "__main__":
+    unittest.main()