118 lines
3.1 KiB
Python
118 lines
3.1 KiB
Python
import itertools
|
|
import math
|
|
import os
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
from sgl_kernel import kimi_k2_moe_fused_gate
|
|
|
|
from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl
|
|
|
|
# CI environment detection
|
|
IS_CI = (
|
|
os.getenv("CI", "false").lower() == "true"
|
|
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
|
)
|
|
|
|
|
|
def kimi_k2_biased_topk_torch_compile(scores, bias, topk, routed_scaling_factor):
|
|
"""Original torch.compile-based implementation"""
|
|
return kimi_k2_biased_topk_impl(
|
|
scores,
|
|
scores,
|
|
bias,
|
|
topk=topk,
|
|
renormalize=True,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
)
|
|
|
|
|
|
def kimi_k2_biased_topk_fused_kernel(scores, bias, topk, routed_scaling_factor):
|
|
"""Our fused CUDA kernel implementation"""
|
|
return kimi_k2_moe_fused_gate(
|
|
scores,
|
|
bias,
|
|
topk=topk,
|
|
renormalize=True,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
)
|
|
|
|
|
|
# CI environment uses simplified parameters
|
|
if IS_CI:
|
|
seq_length_range = [5000] # Only test one sequence length in CI
|
|
else:
|
|
seq_length_range = [
|
|
1,
|
|
8,
|
|
16,
|
|
32,
|
|
64,
|
|
128,
|
|
256,
|
|
512,
|
|
1024,
|
|
2048,
|
|
4096,
|
|
10000,
|
|
15000,
|
|
20000,
|
|
25000,
|
|
30000,
|
|
35000,
|
|
40000,
|
|
]
|
|
|
|
configs = [(sq,) for sq in seq_length_range]
|
|
|
|
|
|
@triton.testing.perf_report(
|
|
triton.testing.Benchmark(
|
|
x_names=["seq_length"],
|
|
x_vals=[list(_) for _ in configs],
|
|
line_arg="provider",
|
|
line_vals=["torch_compile", "fused_kernel"],
|
|
line_names=["Torch Compile", "Fused Kernel"],
|
|
styles=[("blue", "-"), ("red", "-")],
|
|
ylabel="us",
|
|
plot_name="kimi-k2-moe-fused-gate-performance",
|
|
args={},
|
|
)
|
|
)
|
|
def benchmark(seq_length, provider):
|
|
dtype = torch.float32
|
|
device = torch.device("cuda")
|
|
num_experts, topk = 384, 6 # Kimi K2 configuration
|
|
routed_scaling_factor = 2.872 # Kimi K2's routed scaling factor
|
|
|
|
scores = torch.randn((seq_length, num_experts), device=device, dtype=dtype)
|
|
bias = torch.rand(num_experts, device=device, dtype=dtype)
|
|
|
|
quantiles = [0.5, 0.2, 0.8]
|
|
|
|
if provider == "torch_compile":
|
|
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
|
lambda: kimi_k2_biased_topk_torch_compile(
|
|
scores.clone(), bias.clone(), topk, routed_scaling_factor
|
|
),
|
|
quantiles=quantiles,
|
|
)
|
|
elif provider == "fused_kernel":
|
|
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
|
lambda: kimi_k2_biased_topk_fused_kernel(
|
|
scores.clone(), bias.clone(), topk, routed_scaling_factor
|
|
),
|
|
quantiles=quantiles,
|
|
)
|
|
|
|
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 80)
|
|
print("Benchmarking Kimi K2 MoE Fused Gate Performance")
|
|
print("=" * 80)
|
|
print("\nPerformance vs Sequence Length (384 experts, topk=6)")
|
|
benchmark.run(print_data=True, save_path=".")
|