[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>
This commit is contained in:
meinie
co-authored by
Guohao Shao
Brayden Zhong
Guohao Shao
Brayden Zhong
parent
c98d31143d
commit
714011a40f
@@ -264,9 +264,6 @@ set(SOURCES
|
||||
"csrc/gemm/awq_kernel.cu"
|
||||
"csrc/gemm/bmm_fp8.cu"
|
||||
"csrc/gemm/dsv3_fused_a_gemm.cu"
|
||||
"csrc/gemm/dsv3_router_gemm_bf16_out.cu"
|
||||
"csrc/gemm/dsv3_router_gemm_entry.cu"
|
||||
"csrc/gemm/dsv3_router_gemm_float_out.cu"
|
||||
"csrc/gemm/fp8_blockwise_gemm_kernel.cu"
|
||||
"csrc/gemm/fp8_gemm_kernel.cu"
|
||||
"csrc/gemm/int8_gemm_kernel.cu"
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.testing
|
||||
from sgl_kernel import dsv3_router_gemm
|
||||
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
# CI environment uses simplified parameters
|
||||
if IS_CI:
|
||||
num_tokens_vals = [1] # Only test 1 value in CI
|
||||
line_vals = ["sgl-kernel-256"] # Only test one implementation in CI
|
||||
else:
|
||||
num_tokens_vals = [i + 1 for i in range(16)] # Test 1-16 in full mode
|
||||
line_vals = ["torch-256", "sgl-kernel-256", "torch-384", "sgl-kernel-384"]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_tokens"],
|
||||
x_vals=num_tokens_vals,
|
||||
x_log=False,
|
||||
line_arg="impl",
|
||||
line_vals=line_vals,
|
||||
line_names=(
|
||||
[
|
||||
"torch-256",
|
||||
"dsv3_router_gemm-256",
|
||||
"torch-384",
|
||||
"dsv3_router_gemm-384",
|
||||
]
|
||||
if not IS_CI
|
||||
else ["dsv3_router_gemm-256"]
|
||||
),
|
||||
styles=(
|
||||
[("blue", "-"), ("orange", "-"), ("green", "-"), ("red", "-")]
|
||||
if not IS_CI
|
||||
else [("orange", "-")]
|
||||
),
|
||||
ylabel="TFLOPs",
|
||||
plot_name="input-bf16-output-bf16 dsv3 router gemm throughput",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_bf16_output(num_tokens, impl):
|
||||
# M: num_tokens, K: hidden_dim, N: num_experts
|
||||
M, K = num_tokens, 7168
|
||||
|
||||
if impl == "torch-256" or impl == "sgl-kernel-256":
|
||||
N = 256
|
||||
elif impl == "torch-384" or impl == "sgl-kernel-384":
|
||||
N = 384
|
||||
else:
|
||||
raise ValueError(f"Unknown impl: {impl}")
|
||||
|
||||
mat_a = torch.randn((M, K), dtype=torch.bfloat16, device="cuda").contiguous()
|
||||
mat_b = torch.randn((N, K), dtype=torch.bfloat16, device="cuda").contiguous()
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
if impl == "torch-256" or impl == "torch-384":
|
||||
|
||||
def runner():
|
||||
F.linear(mat_a, mat_b)
|
||||
|
||||
elif impl == "sgl-kernel-256" or impl == "sgl-kernel-384":
|
||||
|
||||
def runner():
|
||||
dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.bfloat16)
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(runner, quantiles=quantiles)
|
||||
|
||||
def tflops(t_ms):
|
||||
flops = 2 * M * K * N
|
||||
return flops / (t_ms * 1e-3) / 1e12
|
||||
|
||||
return tflops(ms), tflops(max_ms), tflops(min_ms)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_tokens"],
|
||||
x_vals=num_tokens_vals,
|
||||
x_log=False,
|
||||
line_arg="impl",
|
||||
line_vals=line_vals,
|
||||
line_names=(
|
||||
[
|
||||
"torch-256",
|
||||
"dsv3_router_gemm-256",
|
||||
"torch-384",
|
||||
"dsv3_router_gemm-384",
|
||||
]
|
||||
if not IS_CI
|
||||
else ["dsv3_router_gemm-256"]
|
||||
),
|
||||
styles=(
|
||||
[("blue", "-"), ("orange", "-"), ("green", "-"), ("red", "-")]
|
||||
if not IS_CI
|
||||
else [("orange", "-")]
|
||||
),
|
||||
ylabel="TFLOPs",
|
||||
plot_name="input-bf16-output-fp32 dsv3 router gemm throughput",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_float_output(num_tokens, impl):
|
||||
# M: num_tokens, K: hidden_dim, N: num_experts
|
||||
M, K = num_tokens, 7168
|
||||
|
||||
if impl == "torch-256" or impl == "sgl-kernel-256":
|
||||
N = 256
|
||||
elif impl == "torch-384" or impl == "sgl-kernel-384":
|
||||
N = 384
|
||||
else:
|
||||
raise ValueError(f"Unknown impl: {impl}")
|
||||
|
||||
mat_a = torch.randn((M, K), dtype=torch.bfloat16, device="cuda").contiguous()
|
||||
mat_b = torch.randn((N, K), dtype=torch.bfloat16, device="cuda").contiguous()
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
if impl == "torch-256" or impl == "torch-384":
|
||||
|
||||
def runner():
|
||||
F.linear(mat_a, mat_b).to(torch.float32)
|
||||
|
||||
elif impl == "sgl-kernel-256" or impl == "sgl-kernel-384":
|
||||
|
||||
def runner():
|
||||
dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.float32)
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(runner, quantiles=quantiles)
|
||||
|
||||
def tflops(t_ms):
|
||||
flops = 2 * M * K * N
|
||||
return flops / (t_ms * 1e-3) / 1e12
|
||||
|
||||
return tflops(ms), tflops(max_ms), tflops(min_ms)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
args = parser.parse_args()
|
||||
|
||||
benchmark_bf16_output.run(print_data=True)
|
||||
benchmark_float_output.run(print_data=True)
|
||||
@@ -135,9 +135,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
|
||||
m.def("dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
m.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
|
||||
|
||||
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
m.impl("dsv3_router_gemm", torch::kCUDA, &dsv3_router_gemm);
|
||||
|
||||
/*
|
||||
* From csrc/gemm/gptq
|
||||
*/
|
||||
|
||||
@@ -251,7 +251,6 @@ void bmm_fp8(
|
||||
at::Tensor B_scale,
|
||||
at::Tensor workspace_buffer,
|
||||
int64_t cublas_handle);
|
||||
void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a, const torch::Tensor& mat_b);
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b);
|
||||
|
||||
torch::Tensor gptq_gemm(
|
||||
|
||||
@@ -57,7 +57,6 @@ else:
|
||||
awq_dequantize,
|
||||
bmm_fp8,
|
||||
dsv3_fused_a_gemm,
|
||||
dsv3_router_gemm,
|
||||
fp8_blockwise_scaled_mm,
|
||||
fp8_scaled_mm,
|
||||
gptq_gemm,
|
||||
|
||||
@@ -191,25 +191,6 @@ def qserve_w4a8_per_group_gemm(
|
||||
return out_feats
|
||||
|
||||
|
||||
def dsv3_router_gemm(
|
||||
hidden_states: torch.Tensor,
|
||||
router_weights: torch.Tensor,
|
||||
out_dtype: torch.dtype = torch.bfloat16,
|
||||
) -> torch.Tensor:
|
||||
output = torch.empty(
|
||||
hidden_states.shape[0],
|
||||
router_weights.shape[0],
|
||||
device=hidden_states.device,
|
||||
dtype=out_dtype,
|
||||
)
|
||||
torch.ops.sgl_kernel.dsv3_router_gemm(
|
||||
output,
|
||||
hidden_states,
|
||||
router_weights,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape):
|
||||
output_tensor = torch.empty(
|
||||
output_tensor_shape,
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sgl_kernel import dsv3_router_gemm
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [i + 1 for i in range(16)])
|
||||
@pytest.mark.parametrize("num_experts", [256, 384])
|
||||
def test_dsv3_router_gemm(num_tokens, num_experts):
|
||||
hidden_dim = 7168
|
||||
|
||||
mat_a = torch.randn(
|
||||
(num_tokens, hidden_dim), dtype=torch.bfloat16, device="cuda"
|
||||
).contiguous()
|
||||
mat_b = torch.randn(
|
||||
(num_experts, hidden_dim), dtype=torch.bfloat16, device="cuda"
|
||||
).contiguous()
|
||||
|
||||
bf16_ref = F.linear(mat_a, mat_b)
|
||||
float_ref = bf16_ref.to(torch.float32)
|
||||
|
||||
bf16_output = dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.bfloat16)
|
||||
float_output = dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.float32)
|
||||
|
||||
assert torch.allclose(
|
||||
bf16_output, bf16_ref, rtol=1e-2, atol=1e-3
|
||||
), "Router GEMM output in bf16 dtype mismatch with torch.nn.functional.linear reference"
|
||||
|
||||
assert torch.allclose(
|
||||
float_output, float_ref, rtol=1e-2, atol=1e-3
|
||||
), "Router GEMM output in float32 dtype mismatch with torch.nn.functional.linear reference"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user