[Kernel] Replace dsv3_router_gemm with the unified tiny GEMM (#34693)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
co-authored by
Claude Opus 5
Claude Opus 5
parent
ee462b5899
commit
cb6dd58fbe
@@ -1,53 +0,0 @@
|
||||
"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs torch).
|
||||
|
||||
Run on a Hopper (SM90+) GPU:
|
||||
python -m sglang.kernels.jit.benchmark.bench_dsv3_router_gemm
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.kernels.jit.benchmark import marker
|
||||
from sglang.kernels.jit.benchmark.utils import create_random
|
||||
from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
|
||||
from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
register_amd_ci(est_time=5, stage="jit-kernel-benchmark", runner_config="amd")
|
||||
|
||||
|
||||
def _torch(mat_a, mat_b, out_dtype):
|
||||
return F.linear(mat_a, mat_b).to(out_dtype)
|
||||
|
||||
|
||||
FN_MAP = {
|
||||
"jit": 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", "torch"])
|
||||
def benchmark(num_experts, hidden_dim, num_tokens, out_dtype, provider):
|
||||
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()
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Benchmark for the tiny GEMM (JIT kernel vs torch).
|
||||
|
||||
Run on a Hopper (SM90+) GPU:
|
||||
python -m sglang.kernels.jit.benchmark.bench_tiny_gemm
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.kernels.jit.benchmark import marker
|
||||
from sglang.kernels.jit.benchmark.utils import create_random
|
||||
from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
|
||||
def _torch(x, w, out_dtype):
|
||||
return F.linear(x, w).to(out_dtype)
|
||||
|
||||
|
||||
def _jit(x, w, out_dtype):
|
||||
return tiny_gemm_bf16(x, w, out_dtype=out_dtype, max_m=16)
|
||||
|
||||
|
||||
FN_MAP = {"jit": _jit, "torch": _torch}
|
||||
|
||||
SHAPES = [(256, 7168), (384, 7168), (256, 4096), (896, 7168), (144, 7168), (1536, 128)]
|
||||
|
||||
|
||||
@marker.parametrize("out_dtype", [torch.bfloat16, torch.float32])
|
||||
@marker.parametrize("shape", SHAPES, [(384, 7168)])
|
||||
@marker.parametrize("num_tokens", list(range(1, 17)), [1, 8, 16])
|
||||
@marker.benchmark("provider", ["jit", "torch"])
|
||||
def benchmark(shape, num_tokens, out_dtype, provider):
|
||||
n, k = shape
|
||||
x = create_random(num_tokens, k)
|
||||
w = create_random(n, k)
|
||||
return marker.do_bench(
|
||||
FN_MAP[provider],
|
||||
input_args=(x, w),
|
||||
input_kwargs={"out_dtype": out_dtype},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -1,68 +0,0 @@
|
||||
"""Tests for JIT dsv3_router_gemm kernel."""
|
||||
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
get_ci_test_range,
|
||||
get_jit_cuda_arch,
|
||||
is_hip_runtime,
|
||||
)
|
||||
from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=37, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
|
||||
register_cuda_ci(est_time=110, stage="nightly", runner_config="1-gpu-large")
|
||||
|
||||
HIDDEN_DIMS = [1024, 4096, 5120, 6144, 7168]
|
||||
ROUTER_GEMM_CASES = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[256, 384],
|
||||
HIDDEN_DIMS,
|
||||
list(range(1, 17)),
|
||||
[torch.bfloat16, torch.float32],
|
||||
)
|
||||
),
|
||||
[
|
||||
(256, 1024, 1, torch.bfloat16),
|
||||
(256, 7168, 6, torch.bfloat16),
|
||||
(256, 6144, 4, torch.float32),
|
||||
(384, 7168, 8, torch.bfloat16),
|
||||
(256, 7168, 16, torch.float32),
|
||||
(384, 5120, 16, torch.float32),
|
||||
],
|
||||
)
|
||||
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,hidden_dim,num_tokens,out_dtype", ROUTER_GEMM_CASES
|
||||
)
|
||||
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"]))
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Tests for the JIT tiny_gemm kernels."""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
get_ci_test_range,
|
||||
get_jit_cuda_arch,
|
||||
is_hip_runtime,
|
||||
)
|
||||
from sglang.kernels.ops.gemm.tiny_gemm import can_use_tiny_gemm, tiny_gemm_bf16
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
register_cuda_ci(est_time=300, stage="nightly", runner_config="1-gpu-large")
|
||||
|
||||
# One kernel is built per m in [1, MAX_M], so hold max_m fixed across the sweep:
|
||||
# every num_tokens of a shape then shares one JIT module.
|
||||
MAX_M = 16
|
||||
SHAPES = [(256, 7168), (384, 7168), (256, 4096), (896, 7168), (144, 7168), (1536, 128)]
|
||||
|
||||
TINY_GEMM_CASES = get_ci_test_range(
|
||||
[
|
||||
(n, k, num_tokens, dtype)
|
||||
for n, k in SHAPES
|
||||
for num_tokens in range(1, MAX_M + 1)
|
||||
for dtype in (torch.bfloat16, torch.float32)
|
||||
],
|
||||
[
|
||||
(384, 7168, 1, torch.float32),
|
||||
(384, 7168, 4, torch.float32),
|
||||
(896, 7168, 8, torch.float32),
|
||||
(1536, 128, 16, torch.bfloat16),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize("n,k,num_tokens,out_dtype", TINY_GEMM_CASES)
|
||||
def test_tiny_gemm(n, k, num_tokens, out_dtype):
|
||||
if is_hip_runtime() or get_jit_cuda_arch().major < 9:
|
||||
pytest.skip("SM90+ required")
|
||||
|
||||
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
|
||||
w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
assert can_use_tiny_gemm(n, k, MAX_M)
|
||||
out = tiny_gemm_bf16(x, w, out_dtype=out_dtype, max_m=MAX_M)
|
||||
ref = torch.nn.functional.linear(x, w)
|
||||
torch.testing.assert_close(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -24,10 +24,7 @@ from sglang.kernels.ops.attention.vision_rope import (
|
||||
prepare_fused_qk_complex_rope_inplace,
|
||||
)
|
||||
from sglang.kernels.ops.elementwise import add3
|
||||
from sglang.kernels.ops.gemm.tiny_gemm import (
|
||||
tiny_k_gemm_bf16,
|
||||
tiny_n_gemm_bf16,
|
||||
)
|
||||
from sglang.kernels.ops.gemm.tiny_gemm import tiny_gemm_bf16
|
||||
from sglang.kernels.ops.kvcache.set_mla_kv_buffer import set_mla_kv_buffer
|
||||
from sglang.kernels.ops.mm.process.image import (
|
||||
_normalize_and_patchify_torch,
|
||||
@@ -379,17 +376,19 @@ class TestKimiK3PrerequisiteOps(CustomTestCase):
|
||||
)
|
||||
|
||||
def test_tiny_gemm_variants(self):
|
||||
"""Both K3 gate-projection shapes, one per kernel variant: N=144 is the
|
||||
tiny dimension for the first, K=128 for the second."""
|
||||
torch.manual_seed(2)
|
||||
x = torch.randn(2, 7168, device="cuda", dtype=torch.bfloat16) / 8
|
||||
weight = torch.randn(144, 7168, device="cuda", dtype=torch.bfloat16) / 8
|
||||
actual = tiny_n_gemm_bf16(x, weight, out_dtype=torch.float32)
|
||||
actual = tiny_gemm_bf16(x, weight, out_dtype=torch.float32)
|
||||
torch.testing.assert_close(
|
||||
actual.double(), x.double() @ weight.double().t(), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
x = torch.randn(7, 128, device="cuda", dtype=torch.bfloat16) / 4
|
||||
weight = torch.randn(1536, 128, device="cuda", dtype=torch.bfloat16) / 4
|
||||
actual = tiny_k_gemm_bf16(x, weight)
|
||||
actual = tiny_gemm_bf16(x, weight)
|
||||
torch.testing.assert_close(
|
||||
actual.double(), x.double() @ weight.double().t(), rtol=2e-2, atol=2e-2
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user