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

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