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>

python/sglang/jit_kernel/dsv4/gemm.py

@@ -1,7 +1,10 @@
+import functools
+import importlib.util
+from typing import Optional
+
 import torch
 
 from sglang.srt.environ import envs
-from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.utils import get_bool_env_var, is_hip
 
 _is_hip = is_hip()
@@ -11,14 +14,122 @@ if _use_aiter:
     from aiter.tuned_gemm import tgemm
 
 _linear_bf16_fp32_algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get()
+_HPC_GEMM_WEIGHT_CACHE_ATTR = "_sglang_bf16xfp32_weight_cache"
+# The HPC-Ops bf16xfp32 GEMM consumes the fp32 weight decomposed into two
+# bf16 halves: w_high = w.bf16 and w_low = ((w - w_high) / scale).bf16 with
+# scale = 1/256, so that w ~= w_high + scale * w_low.
+_HPC_GEMM_WEIGHT_SCALE = 1.0 / 256.0
 
 
-def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+@functools.cache
-    if _use_aiter:
+def _hpc_gemm_bf16xfp32_available() -> bool:
+    """HPC-Ops (https://github.com/Tencent/hpc-ops) ships sm90a kernels."""
+    if importlib.util.find_spec("hpc") is None:
+        return False
+    if not torch.cuda.is_available():
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major == 9
+
+
+def _can_use_hpc_gemm_bf16xfp32(
+    x: torch.Tensor, y: torch.Tensor, *, min_m: int = 8
+) -> bool:
+    if x.dim() != 2 or y.dim() != 2 or x.shape[1] != y.shape[1]:
+        return False
+    if x.shape[0] < min_m:
+        return False
+    if not (x.is_cuda and y.is_cuda):
+        return False
+    if x.dtype != torch.bfloat16 or y.dtype != torch.float32:
+        return False
+    if not (x.is_contiguous() and y.is_contiguous()):
+        return False
+    if y.shape[0] % 64 != 0:
+        return False
+    return _hpc_gemm_bf16xfp32_available()
+
+
+def _get_bf16xfp32_weight_split(
+    y: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split the fp32 weight for the HPC-Ops kernel and cache the result
+    (plus the split-K flag workspace, which the kernel leaves zeroed) on the
+    weight tensor."""
+    import hpc
+
+    cache_key = (
+        y.data_ptr(),
+        y._version,
+        tuple(y.shape),
+        tuple(y.stride()),
+        y.device.index,
+        y.dtype,
+    )
+    cache = getattr(y, _HPC_GEMM_WEIGHT_CACHE_ATTR, None)
+    if cache is not None and cache[0] == cache_key:
+        return cache[1], cache[2], cache[3]
+
+    with torch.no_grad():
+        w_high = y.to(torch.bfloat16)
+        w_low = ((y - w_high.float()) / _HPC_GEMM_WEIGHT_SCALE).to(torch.bfloat16)
+    split_flag = hpc.get_gemm_bf16xfp32_workspace(y.shape[0])
+    setattr(y, _HPC_GEMM_WEIGHT_CACHE_ATTR, (cache_key, w_high, w_low, split_flag))
+    return w_high, w_low, split_flag
+
+
+def _linear_bf16_fp32_cublas(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    if x.is_cuda and x.dtype == torch.bfloat16 and y.dtype == torch.bfloat16:
+        return torch.mm(x, y.t(), out_dtype=torch.float32)
+    return torch.mm(x.float(), y.float().t())
+
+
+def _linear_bf16_fp32_hpc(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    *,
+    min_m: int = 8,
+) -> Optional[torch.Tensor]:
+    if not _can_use_hpc_gemm_bf16xfp32(x, y, min_m=min_m):
+        return None
+
+    import hpc
+
+    w_high, w_low, split_flag = _get_bf16xfp32_weight_split(y)
+    return hpc.gemm_bf16xfp32(
+        x,
+        w_high,
+        w_low,
+        _HPC_GEMM_WEIGHT_SCALE,
+        use_fp32_output=True,
+        use_splitk=True,
+        split_flag=split_flag,
+    )
+
+
+def linear_bf16_fp32(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    *,
+    hpc_kernel_min_m: Optional[int] = None,
+) -> torch.Tensor:
+    if _use_aiter and y.dtype == torch.bfloat16:
         return tgemm.mm(x, y, otype=x.dtype).float()
-    elif _linear_bf16_fp32_algo == "deep_gemm":
+    elif hpc_kernel_min_m is not None:
+        output = _linear_bf16_fp32_hpc(x, y, min_m=hpc_kernel_min_m)
+        if output is not None:
+            return output
+        return _linear_bf16_fp32_cublas(x, y)
+    elif _linear_bf16_fp32_algo == "hpc":
+        output = _linear_bf16_fp32_hpc(x, y)
+        if output is not None:
+            return output
+        return _linear_bf16_fp32_cublas(x, y)
+    elif _linear_bf16_fp32_algo == "deep_gemm" and y.dtype == torch.bfloat16:
+        from sglang.srt.layers import deep_gemm_wrapper
+
         z = torch.empty(x.size(0), y.size(0), dtype=torch.float32, device=x.device)
         deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, y, z)
         return z
     else:
-        return torch.mm(x, y.t(), out_dtype=torch.float32)
+        return _linear_bf16_fp32_cublas(x, y)

python/sglang/srt/models/longcat_flash.py

@@ -37,6 +37,7 @@ from typing import Iterable, List, Optional, Tuple
 import torch
 from torch import nn
 
+from sglang.jit_kernel.dsv4 import linear_bf16_fp32
 from sglang.kernels.ops.moe.ep_moe_kernels import zero_experts_compute_triton
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.configs import LongcatFlashConfig
@@ -122,6 +123,15 @@ else:
 
 logger = logging.getLogger(__name__)
 
+# Minimum m (num_tokens) from which the JIT bf16xfp32 router GEMM beats
+# cublas, benchmarked per router shape (hidden_size, n_routed_experts) on H200.
+_LONGCAT_FLASH_ROUTER_HPC_GEMM_MIN_M = {
+    # LongCat-Flash-Chat-FP8: 6144 hidden size, 512 routed experts + 256 zero experts.
+    (6144, 768): 64,
+    # LongCat-Flash-Lite-FP8: 3072 hidden size, 256 routed experts + 128 zero experts.
+    (3072, 384): 128,
+}
+
 
 def _scmoe_align_rows(t, target):
     """Align a [rows,H] tensor to `target` rows across the attn-tp group:
@@ -207,8 +217,21 @@ class LongcatFlashRouter(nn.Module):
         self.e_score_correction_bias = nn.Parameter(
             torch.zeros((self.n_routed_experts), dtype=rounter_params_dtype)
         )
+        self.hpc_kernel_min_m = _LONGCAT_FLASH_ROUTER_HPC_GEMM_MIN_M.get(
+            (config.hidden_size, self.n_routed_experts)
+        )
 
     def forward(self, hidden_states):
+        if (
+            self.hpc_kernel_min_m is not None
+            and self.rounter_params_dtype == torch.float32
+            and self.classifier.bias is None
+        ):
+            return linear_bf16_fp32(
+                hidden_states,
+                self.classifier.weight,
+                hpc_kernel_min_m=self.hpc_kernel_min_m,
+            )
         logits, _ = self.classifier(hidden_states.to(self.rounter_params_dtype))
         return logits
 
@@ -349,8 +372,12 @@ class LongcatFlashDecoderLayer(nn.Module):
                     v_head_dim=config.v_head_dim,
                     q_lora_rank=config.q_lora_rank,
                     kv_lora_rank=config.kv_lora_rank,
-                    rope_theta=config.rope_theta,
+                    rope_theta=(
-                    rope_scaling=config.rope_scaling,
+                        config.rope_parameters["rope_theta"]
+                        if "rope_theta" in getattr(config, "rope_parameters", {})
+                        else config.rope_theta
+                    ),
+                    rope_scaling=getattr(config, "rope_scaling", None),
                     max_position_embeddings=config.max_position_embeddings,
                     quant_config=(
                         None

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