Support JIT fused A GEMM (MLA down projection) and support GLM-5 hidden size, SM120 (#27397)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py

@@ -0,0 +1,94 @@
+"""Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs
+sgl_kernel AOT vs torch.
+
+Run on SM90+ (Hopper or later):
+    python test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py
+"""
+
+import torch
+import torch.nn.functional as F
+import triton.testing
+from sgl_kernel import dsv3_fused_a_gemm as sgl_kernel_dsv3_fused_a_gemm
+
+from sglang.jit_kernel.benchmark import marker
+from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import (
+    dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm,
+)
+from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
+from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.srt.utils.common import is_sm120_supported
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.utils import is_in_ci
+
+register_cuda_ci(est_time=12, suite="base-b-kernel-benchmark-1-gpu-large")
+
+IS_CI = is_in_ci()
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+HD_OUT = 2112
+HD_IN_LIST = [6144, 7168]
+
+AOT_HD_IN = 7168
+HAS_AOT = not is_sm120_supported()
+
+NUM_TOKENS_LIST = [1, 8, 16] if IS_CI else list(range(1, 17))
+
+LINE_VALS = ["cutedsl", "jit", "sgl_kernel", "torch"]
+LINE_NAMES = ["CuTe DSL", "CUDA JIT", "sgl_kernel AOT", "torch F.linear"]
+STYLES = [("blue", "-"), ("orange", "--"), ("red", ":"), ("green", "-.")]
+
+
+def _median_us(fn, *args) -> float:
+    result = marker.do_bench(
+        fn,
+        input_args=args,
+        use_cuda_graph=True,
+        metrics=(0.5,),
+        disable_log_bandwidth=True,
+    )
+    return result.times[0] * 1e6
+
+
+def _bench(num_tokens, provider, hd_in):
+    if provider == "sgl_kernel" and not (HAS_AOT and hd_in == AOT_HD_IN):
+        return float("nan")
+
+    mat_a = torch.randn((num_tokens, hd_in), dtype=DTYPE, device=DEVICE)
+    mat_b = torch.randn((HD_OUT, hd_in), dtype=DTYPE, device=DEVICE).transpose(0, 1)
+    fn_map = {
+        "cutedsl": cutedsl_dsv3_fused_a_gemm,
+        "jit": dsv3_fused_a_gemm,
+        "sgl_kernel": sgl_kernel_dsv3_fused_a_gemm,
+        "torch": lambda a, b: F.linear(a, b.T),
+    }
+    return _median_us(fn_map[provider], mat_a, mat_b)
+
+
+@triton.testing.perf_report(
+    [
+        triton.testing.Benchmark(
+            x_names=["num_tokens"],
+            x_vals=NUM_TOKENS_LIST,
+            line_arg="provider",
+            line_vals=LINE_VALS,
+            line_names=LINE_NAMES,
+            styles=STYLES,
+            ylabel="us",
+            plot_name=f"dsv3-fused-a-gemm-bf16-K{hd_in}-N{HD_OUT}",
+            args={"hd_in": hd_in},
+        )
+        for hd_in in HD_IN_LIST
+    ]
+)
+def benchmark(num_tokens, provider, hd_in):
+    return _bench(num_tokens, provider, hd_in)
+
+
+if __name__ == "__main__":
+    if is_hip_runtime() or get_jit_cuda_arch().major < 9:
+        print(
+            "dsv3_fused_a_gemm JIT kernel requires SM90+ (Hopper). Skipping benchmark."
+        )
+    else:
+        benchmark.run(print_data=True)

test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py

@@ -0,0 +1,42 @@
+"""Tests for the CuTe DSL DeepSeek-V3 fused-A GEMM kernel."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import dsv3_fused_a_gemm
+from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+
+# hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes.
+HD_INS = [6144, 7168]
+# hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants.
+HD_OUTS = [2112, 2624]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("hd_out", HD_OUTS)
+@pytest.mark.parametrize("hd_in", HD_INS)
+@pytest.mark.parametrize("num_tokens", list(range(1, 17)))
+def test_dsv3_fused_a_gemm(num_tokens, hd_in, hd_out):
+    if is_hip_runtime() or get_jit_cuda_arch().major < 9:
+        pytest.skip("SM90+ required")
+
+    torch.manual_seed(num_tokens)
+    weight = torch.randn(hd_out, hd_in, dtype=torch.bfloat16, device="cuda")
+    mat_a = torch.randn(num_tokens, hd_in, dtype=torch.bfloat16, device="cuda")
+    mat_b = weight.t()
+
+    out = dsv3_fused_a_gemm(mat_a, mat_b)
+    assert out.shape == (num_tokens, hd_out)
+    assert out.dtype == torch.bfloat16
+
+    ref = (mat_a.float() @ weight.float().T).bfloat16()
+    torch.testing.assert_close(out, ref, rtol=2e-2, atol=2.5)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/jit/test_dsv3_fused_a_gemm.py

@@ -0,0 +1,43 @@
+"""Tests for JIT dsv3_fused_a_gemm kernel."""
+
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
+from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+
+# hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes.
+HD_INS = [6144, 7168]
+# hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants.
+HD_OUTS = [2112, 2624]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("hd_out", HD_OUTS)
+@pytest.mark.parametrize("hd_in", HD_INS)
+@pytest.mark.parametrize("num_tokens", list(range(1, 17)))
+def test_dsv3_fused_a_gemm(num_tokens, hd_in, hd_out):
+    if is_hip_runtime() or get_jit_cuda_arch().major < 9:
+        pytest.skip("SM90+ required")
+
+    mat_a = torch.randn(num_tokens, hd_in, dtype=torch.bfloat16, device="cuda")
+    mat_b = torch.randn(hd_out, hd_in, dtype=torch.bfloat16, device="cuda").transpose(
+        0, 1
+    )
+
+    ref = F.linear(mat_a, mat_b.T)
+    out = dsv3_fused_a_gemm(mat_a, mat_b)
+
+    assert out.shape == (num_tokens, hd_out)
+    assert out.dtype == torch.bfloat16
+    torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-3)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))