EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[chore] Remove unused ep_moe cuda kernels (#9956)

Browse Source
This commit is contained in:
hlu1
2025-09-06 01:35:50 -07:00
committed by GitHub
parent 039cef76aa
commit 5f1eb20484
13 changed files with 4 additions and 1110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
sgl-kernel/benchmark/bench_moe_ep_post_reorder.py
View File

@@ -1,6 +1,5 @@
import torch
import triton
from sgl_kernel import ep_moe_post_reorder
from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel
@@ -13,9 +12,9 @@ configs = [(bs,) for bs in batch_sizes]
x_names=["batch_size"],
x_vals=[list(_) for _ in configs],
line_arg="provider",
line_vals=["cuda", "triton"],
line_names=["CUDA Kernel", "Triton Kernel"],
styles=[("green", "-"), ("orange", "-")],
line_vals=["triton"],
line_names=["Triton Kernel"],
styles=[("orange", "-")],
ylabel="us",
plot_name="ep-moe-post-reorder-performance",
args={},
@@ -46,24 +45,7 @@ def benchmark(batch_size, provider):
quantiles = [0.5, 0.2, 0.8]
if provider == "cuda":
d_out, out, s2d, tk_ids, tk_weights = alloc_tensors()
def run_cuda():
ep_moe_post_reorder(
d_out,
out,
s2d,
tk_ids,
tk_weights,
start_expert_id,
end_expert_id,
topk,
)
ms, min_ms, max_ms = triton.testing.do_bench(run_cuda, quantiles=quantiles)
elif provider == "triton":
if provider == "triton":
d_out, out, s2d, tk_ids, tk_weights = alloc_tensors()
def run_triton():

103
sgl-kernel/benchmark/bench_moe_ep_pre_reorder.py
View File

@@ -1,103 +0,0 @@
import torch
import triton
from sgl_kernel import ep_moe_pre_reorder
from sglang.srt.layers.moe.ep_moe.kernels import pre_reorder_triton_kernel
batch_sizes = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]
configs = [(bs,) for bs in batch_sizes]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=[list(_) for _ in configs],
line_arg="provider",
line_vals=["cuda", "triton"],
line_names=["CUDA Kernel", "Triton Kernel"],
styles=[("green", "-"), ("orange", "-")],
ylabel="us",
plot_name="ep-moe-pre-reorder-performance",
args={},
)
)
def benchmark(batch_size, provider):
dtype = torch.bfloat16
device = torch.device("cuda")
hidden_size, topk, start_expert_id, end_expert_id, block_size = (
4096,
8,
0,
255,
512,
)
# Allocate fresh tensors for every run to match bench_moe_fused_gate style
def alloc_tensors():
input_ = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
gateup_input = torch.zeros(
batch_size * topk, hidden_size, dtype=dtype, device=device
)
src2dst = torch.randint(
0, batch_size * topk, (batch_size, topk), dtype=torch.int32, device=device
)
topk_ids = torch.randint(
start_expert_id,
end_expert_id + 1,
(batch_size, topk),
dtype=torch.int32,
device=device,
)
a1_scales = torch.rand(
end_expert_id - start_expert_id + 1, dtype=torch.float32, device=device
)
return input_, gateup_input, src2dst, topk_ids, a1_scales
quantiles = [0.5, 0.2, 0.8]
if provider == "cuda":
inp, gout, s2d, tk_ids, scales = alloc_tensors()
def run_cuda():
ep_moe_pre_reorder(
inp,
gout,
s2d,
tk_ids,
scales,
start_expert_id,
end_expert_id,
topk,
True,
)
ms, min_ms, max_ms = triton.testing.do_bench(run_cuda, quantiles=quantiles)
elif provider == "triton":
inp, gout, s2d, tk_ids, scales = alloc_tensors()
def run_triton():
pre_reorder_triton_kernel[(batch_size,)](
inp.view(-1),
gout.view(-1),
s2d.view(-1),
tk_ids.view(-1),
scales,
start_expert_id,
end_expert_id,
topk,
hidden_size,
block_size,
True,
)
ms, min_ms, max_ms = triton.testing.do_bench(run_triton, quantiles=quantiles)
else:
raise ValueError(f"Unknown provider: {provider}")
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
benchmark.run(print_data=True)

92
sgl-kernel/benchmark/bench_moe_silu_and_mul.py
View File

@@ -1,92 +0,0 @@
import itertools
import torch
import triton
from sgl_kernel import ep_moe_silu_and_mul
from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_triton_kernel
batch_size_range = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]
hidden_size_range = [1024, 2048, 4096, 8192]
block_size_range = [128, 256, 512]
configs = list(itertools.product(batch_size_range, hidden_size_range, block_size_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "hidden_size", "block_size"],
x_vals=[list(cfg) for cfg in configs],
line_arg="provider",
line_vals=["cuda", "triton"],
line_names=["CUDA Kernel", "Triton Kernel"],
styles=[("green", "-"), ("orange", "-")],
ylabel="us",
plot_name="ep-moe-silu-and-mul-performance",
args={},
)
)
def benchmark(batch_size, hidden_size, block_size, provider):
dtype = torch.bfloat16
device = torch.device("cuda")
half_hidden_size = hidden_size // 2
start_expert_id, end_expert_id = 0, 255
block_size = 512
quantiles = [0.5, 0.2, 0.8]
def alloc_tensors():
gateup_output = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
down_input = torch.empty(
batch_size, half_hidden_size, dtype=dtype, device=device
)
reorder_topk_ids = torch.randint(
start_expert_id,
end_expert_id + 1,
(batch_size,),
dtype=torch.int32,
device=device,
)
scales = torch.rand(
end_expert_id - start_expert_id + 1, dtype=torch.float32, device=device
)
return gateup_output, down_input, reorder_topk_ids, scales
if provider == "cuda":
gateup, down, ids, scales = alloc_tensors()
def run_cuda():
ep_moe_silu_and_mul(
gateup,
down,
ids,
scales,
start_expert_id,
end_expert_id,
)
ms, min_ms, max_ms = triton.testing.do_bench(run_cuda, quantiles=quantiles)
elif provider == "triton":
gateup, down, ids, scales = alloc_tensors()
def run_triton():
silu_and_mul_triton_kernel[(batch_size,)](
gateup.view(-1),
down.view(-1),
hidden_size,
ids,
scales,
start_expert_id,
end_expert_id,
block_size,
)
ms, min_ms, max_ms = triton.testing.do_bench(run_triton, quantiles=quantiles)
else:
raise ValueError(f"Unknown provider: {provider}")
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
benchmark.run(print_data=True)