Qwen3.8-27B Model Support (#34859)

Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai>
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Zijie Xia <zijie.xia@radixark.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>

test/registered/attention/test_gdn_fused_split_head_ratios.py

@@ -0,0 +1,91 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
+    fused_qkvzba_split_reshape_cat_contiguous,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=3, stage="base-b", runner_config="1-gpu-large")
+
+
+def _reference_split(mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_qk, head_v):
+    """Plain-slicing reference for the contiguous [Q|K|V|Z] / [B|A] layouts."""
+    batch = mixed_qkvz.shape[0]
+    total_q = num_heads_qk * head_qk
+    total_v = num_heads_v * head_v
+    q = mixed_qkvz[:, :total_q]
+    k = mixed_qkvz[:, total_q : 2 * total_q]
+    v = mixed_qkvz[:, 2 * total_q : 2 * total_q + total_v]
+    z = mixed_qkvz[:, 2 * total_q + total_v :]
+    mixed_qkv = torch.cat((q, k, v), dim=-1).contiguous()
+    b = mixed_ba[:, :num_heads_v].contiguous()
+    a = mixed_ba[:, num_heads_v:].contiguous()
+    return (
+        mixed_qkv,
+        z.reshape(batch, num_heads_v, head_v).contiguous(),
+        b,
+        a,
+    )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestGdnFusedSplitHeadRatios(unittest.TestCase):
+    """The fused contiguous split must be exact for every supported v/k head
+    ratio, including the non-power-of-two ratio 3 of the dense 27B hybrids
+    (served by the per-head walk instead of one wide vector access)."""
+
+    HEAD_QK = 128
+    HEAD_V = 128
+    NUM_HEADS_QK = 16
+
+    def _run_ratio(self, ratio: int, batch: int = 33) -> None:
+        torch.manual_seed(ratio)
+        num_heads_v = self.NUM_HEADS_QK * ratio
+        total_qkvz = (
+            2 * self.NUM_HEADS_QK * self.HEAD_QK + 2 * num_heads_v * self.HEAD_V
+        )
+        mixed_qkvz = torch.randn(batch, total_qkvz, dtype=torch.bfloat16, device="cuda")
+        mixed_ba = torch.randn(
+            batch, 2 * num_heads_v, dtype=torch.bfloat16, device="cuda"
+        )
+
+        got_qkv, got_z, got_b, got_a = fused_qkvzba_split_reshape_cat_contiguous(
+            mixed_qkvz,
+            mixed_ba,
+            self.NUM_HEADS_QK,
+            num_heads_v,
+            self.HEAD_QK,
+            self.HEAD_V,
+        )
+        ref_qkv, ref_z, ref_b, ref_a = _reference_split(
+            mixed_qkvz,
+            mixed_ba,
+            self.NUM_HEADS_QK,
+            num_heads_v,
+            self.HEAD_QK,
+            self.HEAD_V,
+        )
+
+        # A pure data-movement kernel must be bitwise exact.
+        torch.testing.assert_close(got_qkv.view(-1), ref_qkv.view(-1), rtol=0, atol=0)
+        torch.testing.assert_close(got_z.reshape(-1), ref_z.reshape(-1), rtol=0, atol=0)
+        torch.testing.assert_close(got_b.view(-1), ref_b.view(-1), rtol=0, atol=0)
+        torch.testing.assert_close(got_a.view(-1), ref_a.view(-1), rtol=0, atol=0)
+
+    def test_ratio_1(self):
+        self._run_ratio(1)
+
+    def test_ratio_2(self):
+        self._run_ratio(2)
+
+    def test_ratio_3(self):
+        self._run_ratio(3)
+
+    def test_ratio_4(self):
+        self._run_ratio(4)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/gemm/test_hopper_bf16_gemv.py

@@ -0,0 +1,68 @@
+"""
+Tests the Hopper single-token bf16 GEMV JIT kernel against torch (cuBLAS +
+fp32 reference) on the dispatch domains where the backend enables it.
+"""
+
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=6, stage="base-b", runner_config="1-gpu-large")
+
+
+def _is_sm90() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9
+
+
+@unittest.skipIf(not _is_sm90(), "Hopper bf16 GEMV requires an SM90 GPU")
+class TestHopperBf16Gemv(unittest.TestCase):
+    def _run_case(self, n, k, seed=0):
+        from sglang.kernels.ops.gemm.hopper_bf16_gemv import hopper_bf16_gemv
+
+        torch.manual_seed(seed)
+        x = torch.randn(1, k, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.05
+        out = hopper_bf16_gemv(x, w)
+        ref = x.float() @ w.float().t()
+        cub = (x @ w.t()).float()
+        err = (out.float() - ref).abs().max().item()
+        err_cub = (cub - ref).abs().max().item()
+        # fp32 accumulation + single warp-tree reduction: at least as tight as
+        # cuBLAS (which split-K reduces) against the fp32 reference.
+        self.assertLessEqual(err, max(err_cub * 2.0, 1e-2), (n, k, err, err_cub))
+        self.assertFalse(torch.isnan(out).any().item(), (n, k))
+
+    def test_dispatch_domain_shapes(self):
+        # Representative dense-decode shapes (Qwen3.6-27B): out_proj/attn_o,
+        # attn_qkv, mlp_down, mlp_gate_up, in_proj_ba.
+        for n, k in [
+            (5120, 6144),
+            (8192, 5120),
+            (5120, 17408),
+            (34816, 5120),
+            (96, 5120),
+        ]:
+            self._run_case(n, k)
+
+    def test_tail_rows(self):
+        # N not divisible by rows_per_block exercises the guarded tail path.
+        for n in [104, 5128, 8200]:
+            self._run_case(n, 5120)
+
+    def test_predicate(self):
+        from sglang.kernels.ops.gemm.hopper_bf16_gemv import use_hopper_bf16_gemv
+
+        self.assertTrue(use_hopper_bf16_gemv(1, 5120, 6144))
+        self.assertTrue(use_hopper_bf16_gemv(1, 34816, 5120))
+        # batched decode, odd K, huge N (lm_head), and the near-optimal-cuBLAS
+        # mid-N band must all fall back.
+        self.assertFalse(use_hopper_bf16_gemv(2, 5120, 6144))
+        self.assertFalse(use_hopper_bf16_gemv(1, 5120, 6000))
+        self.assertFalse(use_hopper_bf16_gemv(1, 248320, 5120))
+        self.assertFalse(use_hopper_bf16_gemv(1, 16384, 5120))
+
+
+if __name__ == "__main__":
+    unittest.main()