[JIT Kernel] Migrate dsv3_router_gemm from AOT sgl-kernel to JIT kernel (#21531)

Co-authored-by: Guohao Shao <shao.gh.98@gmail.com>
Co-authored-by: Brayden Zhong <brayden@radixark.ai>

test/registered/jit/benchmark/bench_dsv3_router_gemm.py

@@ -0,0 +1,64 @@
+"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs sgl_kernel AOT vs torch).
+
+Run on a Hopper (SM90+) GPU:
+    python -m sglang.jit_kernel.benchmark.bench_dsv3_router_gemm
+"""
+
+import torch
+import torch.nn.functional as F
+
+try:
+    from sgl_kernel import dsv3_router_gemm as sgl_kernel_dsv3_router_gemm
+except ImportError:
+    sgl_kernel_dsv3_router_gemm = None
+
+from sglang.jit_kernel.benchmark import marker
+from sglang.jit_kernel.benchmark.utils import create_random
+from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_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=5, suite="base-b-kernel-benchmark-1-gpu-large")
+
+# sgl_kernel AOT kernel is specialized for hidden_dim=7168 only.
+SGL_KERNEL_HIDDEN_DIM = 7168
+
+
+def _torch(mat_a, mat_b, out_dtype):
+    return F.linear(mat_a, mat_b).to(out_dtype)
+
+
+FN_MAP = {
+    "jit": dsv3_router_gemm,
+    "sgl_kernel": sgl_kernel_dsv3_router_gemm,
+    "torch": _torch,
+}
+
+
+@marker.parametrize("num_experts", [256, 384], [256])
+@marker.parametrize("hidden_dim", [6144, 7168], [7168])
+@marker.parametrize("num_tokens", list(range(1, 17)), [1, 8, 16])
+@marker.parametrize("out_dtype", [torch.bfloat16, torch.float32])
+@marker.benchmark("provider", ["jit", "sgl_kernel", "torch"])
+def benchmark(num_experts, hidden_dim, num_tokens, out_dtype, provider):
+    if provider == "sgl_kernel":
+        if sgl_kernel_dsv3_router_gemm is None:
+            marker.skip("sgl_kernel dsv3_router_gemm not available in this build")
+        if hidden_dim != SGL_KERNEL_HIDDEN_DIM:
+            marker.skip("sgl_kernel AOT only supports hidden_dim=7168")
+    mat_a = create_random(num_tokens, hidden_dim)
+    mat_b = create_random(num_experts, hidden_dim)
+    return marker.do_bench(
+        FN_MAP[provider],
+        input_args=(mat_a, mat_b),
+        input_kwargs={"out_dtype": out_dtype},
+    )
+
+
+if __name__ == "__main__":
+    if is_hip_runtime() or get_jit_cuda_arch().major < 9:
+        print(
+            "dsv3_router_gemm JIT kernel requires SM90+ (Hopper). Skipping benchmark."
+        )
+    else:
+        benchmark.run()

test/registered/jit/test_dsv3_router_gemm.py

@@ -0,0 +1,45 @@
+"""Tests for JIT dsv3_router_gemm kernel."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_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=37, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=148, suite="nightly-kernel-1-gpu", nightly=True)
+
+HIDDEN_DIMS = [1024, 4096, 5120, 6144, 7168]
+ATOL = 1e-2
+RTOL = 1e-2
+
+
+def _ref(hidden_states, router_weights, out_dtype):
+    return (hidden_states.float() @ router_weights.float().T).to(out_dtype)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("num_experts", [256, 384])
+@pytest.mark.parametrize("hidden_dim", HIDDEN_DIMS)
+@pytest.mark.parametrize("num_tokens", list(range(1, 17)))
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32])
+def test_dsv3_router_gemm(num_experts, hidden_dim, num_tokens, out_dtype):
+    if is_hip_runtime() or get_jit_cuda_arch().major < 9:
+        pytest.skip("SM90+ required")
+
+    mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda")
+    mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.bfloat16, device="cuda")
+
+    ref = _ref(mat_a, mat_b, out_dtype)
+    out = dsv3_router_gemm(mat_a, mat_b, out_dtype=out_dtype)
+
+    assert out.shape == (num_tokens, num_experts)
+    assert out.dtype == out_dtype
+    torch.testing.assert_close(out.float(), ref.float(), atol=ATOL, rtol=RTOL)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))