[Perf] beam: Add optional cuTile BF16x3 router GEMM

python/sglang/kernels/fused_op.py

@@ -106,6 +106,7 @@
     KernelBackend.JIT: "forward_jit",
     KernelBackend.AOT: "forward_aot",
     KernelBackend.CUTE_DSL: "forward_cute_dsl",
+    KernelBackend.CUTILE: "forward_cutile",
     KernelBackend.FLYDSL: "forward_flydsl",
     KernelBackend.HELION: "forward_helion",
     KernelBackend.KDA: "forward_kda",
@@ -421,6 +422,9 @@ def forward_aot(self, *args, **kwargs):
     def forward_cute_dsl(self, *args, **kwargs):
         raise NotImplementedError(f"{self._op_label()}: no cute_dsl backend")
 
+    def forward_cutile(self, *args, **kwargs):
+        raise NotImplementedError(f"{self._op_label()}: no cutile backend")
+
     def forward_flydsl(self, *args, **kwargs):
         raise NotImplementedError(f"{self._op_label()}: no flydsl backend")
 

python/sglang/kernels/ops/gemm/__init__.py

@@ -525,13 +525,6 @@ def kimi_k3_tiny_gemm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
         capabilities=CUDA_ONLY,
     )
 )
-register_kernel(
-    KernelSpec(
-        op="gemm.split_fp32_to_bf16x3",
-        backend=KernelBackend.TRITON,
-        target="sglang.kernels.ops.gemm.router_gemv:split_fp32_to_bf16x3",
-    )
-)
 register_kernel(
     KernelSpec(
         op="gemm.router_gemm_bf16x3_cublas",
@@ -540,3 +533,14 @@ def kimi_k3_tiny_gemm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
         capabilities=CUDA_ONLY,
     )
 )
+# Launch configs are tuned for Hopper and SM10x only.
+register_kernel(
+    KernelSpec(
+        op="gemm.router_gemm_bf16x3",
+        backend=KernelBackend.CUTILE,
+        target="sglang.kernels.ops.gemm.router_gemm_cutile:router_gemm_bf16x3_cutile",
+        capabilities=frozenset(
+            {CapabilityRequirement.cuda(min_sm=(9, 0), max_sm=(10, 9))}
+        ),
+    )
+)

python/sglang/kernels/ops/gemm/router_gemm_cutile.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+"""cuTile router GEMM: BF16 activations × FP32 weights → FP32 logits.
+
+The caller splits each FP32 weight into three BF16 components before serving.
+Main-product K tiles are combined with explicit FP32 rounding for accuracy.
+Split-K configurations use a separate reduction kernel.
+"""
+
+import functools
+
+import cuda.tile as ct
+import torch
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+@ct.kernel(num_ctas=1)
+def _router_gemm_bf16x3_kernel(
+    X,
+    W,
+    WR,
+    Y,
+    BM: ct.Constant[int],
+    BN: ct.Constant[int],
+    BK: ct.Constant[int],
+    SPLIT_K: ct.Constant[int],
+    EXPERT_MAJOR: ct.Constant[bool],
+    PDL: ct.Constant[bool],
+):
+    col, row, part = ct.bid(0), ct.bid(1), ct.bid(2)
+    if PDL:
+        ct.grid_dependency_control_launch_dependents()
+        ct.grid_dependency_control_wait()
+    steps = ct.cdiv(ct.cdiv(X.shape[1], BK), SPLIT_K)
+    if EXPERT_MAJOR:
+        total = ct.full((4 * BN, BM), 0, ct.float32)
+    else:
+        total = ct.full((BM, BN), 0, ct.float32)
+        residual = ct.full((BM, 2 * BN), 0, ct.float32)
+    for step in range(steps):
+        k = part * steps + step
+        x = ct.load(
+            X,
+            (row, k),
+            (BM, BK),
+            padding_mode=ct.PaddingMode.ZERO,
+            latency=4 if EXPERT_MAJOR or BK == 128 else 8,
+        )
+        if EXPERT_MAJOR:
+            w = ct.load(
+                W, (0, col, k), (4, BN, BK), padding_mode=ct.PaddingMode.ZERO, latency=4
+            )
+            local = ct.mma(
+                ct.reshape(w, (4 * BN, BK)),
+                ct.transpose(x),
+                ct.full((4 * BN, BM), 0, ct.float32),
+            )
+        else:
+            w0 = ct.load(
+                W,
+                (0, col, k),
+                (1, BN, BK),
+                padding_mode=ct.PaddingMode.ZERO,
+                latency=4 if BK == 128 else 8,
+            )
+            wr = ct.load(
+                WR,
+                (0, col, k),
+                (2, BN, BK),
+                padding_mode=ct.PaddingMode.ZERO,
+                latency=4 if BK == 128 else 8,
+            )
+            local = ct.mma(
+                x,
+                ct.transpose(ct.reshape(w0, (BN, BK))),
+                ct.full((BM, BN), 0, ct.float32),
+            )
+            residual = ct.mma(x, ct.transpose(ct.reshape(wr, (2 * BN, BK))), residual)
+        # Round each main-product tile before adding the next tile.
+        total = ct.add(total, local, rounding_mode=ct.RoundingMode.RN)
+    if EXPERT_MAJOR:
+        main = ct.extract(total, (0, 0), (BN, BM))
+        mid = ct.extract(total, (1, 0), (BN, BM))
+        low = ct.extract(total, (2, 0), (BN, BM))
+        value = ct.transpose(
+            ct.add(
+                main,
+                ct.add(mid, low, rounding_mode=ct.RoundingMode.RN),
+                rounding_mode=ct.RoundingMode.RN,
+            )
+        )
+    else:
+        correction = ct.sum(
+            ct.reshape(residual, (BM, 2, BN)), axis=1, rounding_mode=ct.RoundingMode.RN
+        )
+        value = ct.add(total, correction, rounding_mode=ct.RoundingMode.RN)
+    ct.store(
+        Y,
+        (part, row, col),
+        ct.reshape(value, (1, BM, BN)),
+        allow_tma=not (BM == 256 and SPLIT_K > 1),
+    )
+
+
+@ct.kernel
+def _split_k_reduce_kernel(
+    P, Y, SPLIT_K: ct.Constant[int], PDL: ct.Constant[bool] = False
+):
+    if PDL:
+        ct.grid_dependency_control_wait()
+    if SPLIT_K == 2:
+        values = ct.load(
+            P,
+            (0, ct.bid(0)),
+            (SPLIT_K, 1024),
+            padding_mode=ct.PaddingMode.ZERO,
+            allow_tma=False,
+        )
+        total = ct.sum(values, axis=0, rounding_mode=ct.RoundingMode.RN)
+    else:
+        total = ct.full((1024,), 0, ct.float32)
+        for part in range(SPLIT_K):
+            value = ct.load(
+                P,
+                (part, ct.bid(0)),
+                (1, 1024),
+                padding_mode=ct.PaddingMode.ZERO,
+                allow_tma=False,
+            )
+            total = ct.add(
+                total, ct.reshape(value, (1024,)), rounding_mode=ct.RoundingMode.RN
+            )
+    ct.store(Y, (ct.bid(0),), total, allow_tma=False)
+
+
+_GEMM_KERNELS = {
+    1: _router_gemm_bf16x3_kernel,
+    2: _router_gemm_bf16x3_kernel.replace_hints(num_ctas=2),
+}
+# (max M, (BM, BN, BK, SPLIT_K, CTAS)); larger M uses the fallback below.
+_BLACKWELL_CONFIGS = (
+    (64, (32, 16, 256, 4, 1)),
+    (128, (64, 32, 256, 4, 1)),
+    (512, (128, 64, 256, 8, 1)),
+    (1024, (128, 64, 128, 4, 1)),
+    (2048, (128, 64, 128, 2, 1)),
+    (4096, (256, 128, 256, 2, 2)),
+)
+# Use smaller tiles on Hopper to limit register pressure.
+_HOPPER_CONFIGS = (
+    (256, (64, 64, 128, 8, 1)),
+    (512, (64, 64, 128, 4, 1)),
+    (1024, (64, 64, 128, 2, 1)),
+)
+
+
+@functools.cache
+def _is_hopper(device_index: int) -> bool:
+    return torch.cuda.get_device_capability(device_index)[0] == 9
+
+
+def _select_launch_config(m: int, device: torch.device) -> tuple:
+    if m > 128 and _is_hopper(device.index):
+        configs, fallback = _HOPPER_CONFIGS, (64, 64, 128, 1, 1)
+    else:
+        configs, fallback = _BLACKWELL_CONFIGS, (256, 128, 256, 1, 2)
+    return next((config for limit, config in configs if m <= limit), fallback)
+
+
+def _router_gemm_bf16x3_fake(x: torch.Tensor, parts: torch.Tensor) -> torch.Tensor:
+    return x.new_empty((x.shape[0], parts.shape[0] // 3), dtype=torch.float32)
+
+
+@register_custom_op(fake_impl=_router_gemm_bf16x3_fake)
+def router_gemm_bf16x3_cutile(x: torch.Tensor, parts: torch.Tensor) -> torch.Tensor:
+    """Multiply BF16 [M, K] by prepared BF16 [3N, K], producing FP32 [M, N]."""
+    m, k = x.shape
+    n = parts.shape[0] // 3
+    bm, bn, bk, splits, ctas = _select_launch_config(m, x.device)
+    pdl = bk == 128 and splits == 4
+    out = torch.empty((m, n), device=x.device, dtype=torch.float32)
+    if m == 0 or n == 0:
+        return out
+    partial = (
+        torch.empty((splits, m, n), device=x.device, dtype=torch.float32)
+        if splits > 1
+        else out.unsqueeze(0)
+    )
+    with torch.cuda.device(x.device):
+        stream = torch.cuda.current_stream()
+        ct.launch(
+            stream,
+            (ct.cdiv(n, bn), ct.cdiv(m, bm), splits),
+            _GEMM_KERNELS[ctas],
+            (
+                x,
+                parts.view(3, n, k),
+                parts[n:].view(2, n, k),
+                partial,
+                bm,
+                bn,
+                bk,
+                splits,
+                m <= 64,
+                pdl,
+            ),
+            programmatic_dependent_launch=pdl,
+        )
+        if splits > 1:
+            ct.launch(
+                stream,
+                (ct.cdiv(out.numel(), 1024),),
+                _split_k_reduce_kernel,
+                (partial.view(splits, -1), out.view(-1), splits, pdl),
+                programmatic_dependent_launch=pdl,
+            )
+    return out

python/sglang/kernels/ops/gemm/router_gemv.py

@@ -11,12 +11,13 @@
 
 from __future__ import annotations
 
-from typing import Dict, Optional, Tuple
+from typing import Dict, Tuple
 
 import torch
 import triton
 import triton.language as tl
 
+from sglang.kernels.ops.quantization.bf16x3_split import split_fp32_to_bf16x3
 from sglang.srt.utils import is_gfx95_supported, is_hip
 
 _is_hip = is_hip()
@@ -186,48 +187,6 @@ def router_gemv(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
     return out
 
 
-@triton.jit
-def _split_fp32_to_bf16x3_kernel(w_ptr, out_ptr, numel, BLOCK: tl.constexpr):
-    offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
-    mask = offs < numel
-    w = tl.load(w_ptr + offs, mask=mask)
-    hi = w.to(tl.bfloat16)
-    residual = w - hi.to(tl.float32)
-    mid = residual.to(tl.bfloat16)
-    lo = (residual - mid.to(tl.float32)).to(tl.bfloat16)
-    tl.store(out_ptr + offs, hi, mask=mask)
-    tl.store(out_ptr + numel + offs, mid, mask=mask)
-    tl.store(out_ptr + 2 * numel + offs, lo, mask=mask)
-
-
-def split_fp32_to_bf16x3(
-    w: torch.Tensor, out: Optional[torch.Tensor] = None
-) -> torch.Tensor:
-    """[N, K] fp32 -> [3N, K] bf16 as (hi, mid, lo) blocks with w == hi + mid + lo.
-
-    Exact for |w| >= 2**-110: each part takes the next 8 significant bits, and the
-    last residual then stays a normal number.
-    """
-    if out is None:
-        out = torch.empty(
-            3 * w.shape[0], w.shape[1], dtype=torch.bfloat16, device=w.device
-        )
-    if w.is_cuda and w.is_contiguous():
-        block = 1024
-        _split_fp32_to_bf16x3_kernel[(triton.cdiv(w.numel(), block),)](
-            w, out, w.numel(), BLOCK=block
-        )
-        return out
-    hi = w.to(torch.bfloat16)
-    residual = w - hi.float()
-    mid = residual.to(torch.bfloat16)
-    lo = (residual - mid.float()).to(torch.bfloat16)
-    return torch.cat([hi, mid, lo], out=out)
-
-
-_split_scratch: Dict[Tuple[torch.device, int, int], torch.Tensor] = {}
-
-
 def router_gemm_bf16x3_cublas(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
     """[M, K] bf16 x [N, K] fp32 router logits on BF16 tensor cores.
 
@@ -236,11 +195,6 @@ def router_gemm_bf16x3_cublas(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
     tensor-core accumulation is not IEEE fp32: the error is a few times that of
     an fp32 SGEMM.
     """
-    n, k = w.shape
-    key = (w.device, n, k)
-    if key not in _split_scratch:
-        _split_scratch[key] = torch.empty(
-            3 * n, k, dtype=torch.bfloat16, device=w.device
-        )
-    w3 = split_fp32_to_bf16x3(w, out=_split_scratch[key])
+    n = w.shape[0]
+    w3 = split_fp32_to_bf16x3(w)
     return torch.mm(x, w3.t(), out_dtype=torch.float32).view(-1, 3, n).sum(dim=1)

python/sglang/kernels/ops/quantization/__init__.py

@@ -142,3 +142,11 @@ def sgl_per_token_group_quant_8bit(
         target="sglang.kernels.ops.quantization.act_quant_triton:act_quant",
     )
 )
+register_kernel(
+    KernelSpec(
+        op="quantization.split_fp32_to_bf16x3",
+        backend=KernelBackend.TRITON,
+        target="sglang.kernels.ops.quantization.bf16x3_split:split_fp32_to_bf16x3",
+        capabilities=CUDA_ONLY,
+    )
+)