From e4eea7ce2ffab90d6eef2e79cbb39e4d14577f7c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 21 Jul 2026 20:05:17 +0800 Subject: [PATCH] Optimize LongCat-Flash router GEMM with the HPC-Ops bf16xfp32 kernel (#30247) Co-authored-by: Claude Fable 5 Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com> --- python/sglang/jit_kernel/dsv4/gemm.py | 121 +++++++++++++++++- python/sglang/srt/models/longcat_flash.py | 31 ++++- .../gemm/test_linear_bf16_fp32_hpc.py | 75 +++++++++++ .../test_longcat_flash_router_hpc_gemm.py | 116 +++++++++++++++++ 4 files changed, 336 insertions(+), 7 deletions(-) create mode 100644 test/registered/gemm/test_linear_bf16_fp32_hpc.py create mode 100644 test/registered/unit/models/test_longcat_flash_router_hpc_gemm.py diff --git a/python/sglang/jit_kernel/dsv4/gemm.py b/python/sglang/jit_kernel/dsv4/gemm.py index da60eccec..c02a7c7a5 100644 --- a/python/sglang/jit_kernel/dsv4/gemm.py +++ b/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: - if _use_aiter: +@functools.cache +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) diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index eeb289c8c..c0662e1c0 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/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_scaling=config.rope_scaling, + rope_theta=( + 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 diff --git a/test/registered/gemm/test_linear_bf16_fp32_hpc.py b/test/registered/gemm/test_linear_bf16_fp32_hpc.py new file mode 100644 index 000000000..1e716220f --- /dev/null +++ b/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() diff --git a/test/registered/unit/models/test_longcat_flash_router_hpc_gemm.py b/test/registered/unit/models/test_longcat_flash_router_hpc_gemm.py new file mode 100644 index 000000000..834cff921 --- /dev/null +++ b/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()