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>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Fable 5
Halcyon
Claude Fable 5
Halcyon
parent
303896a475
commit
e4eea7ce2f
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user