chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[DeepSeek V3.1/V3.2] Optimize fused moe configs for H20 & H20-3E based on swapab (#17133)

Browse Source
This commit is contained in:
Yongfei Xu
2026-01-17 00:10:52 +08:00
committed by GitHub
parent 6f10e17b4a
commit 82a1b645ba
6 changed files with 954 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py
+332 -210
View File
@@ -1,5 +1,6 @@
# Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py
import argparse
import dataclasses
import json
import os
import time
@@ -20,9 +21,7 @@ from common_utils import (
sort_config,
)
from ray.experimental.tqdm_ray import tqdm
from sgl_kernel import silu_and_mul
from sglang.srt.layers.moe.fused_moe_triton import override_config
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
get_config_dtype_str,
invoke_fused_moe_kernel,
@@ -38,6 +37,90 @@ from sglang.srt.utils import is_hip
_is_hip = is_hip()
@dataclasses.dataclass
class MoeInputs:
topk_ids: torch.Tensor
sorted_token_ids: torch.Tensor
expert_ids: torch.Tensor
num_tokens_post_padded: torch.Tensor
class KernelWrapper:
def __init__(self, moe_inputs, use_cuda_graph=True, inner_iter=10, **kwargs):
self.func = invoke_fused_moe_kernel
self.use_cuda_graph = use_cuda_graph
self.moe_inputs = moe_inputs
self.inner_iter = inner_iter
self.kwargs = kwargs
if use_cuda_graph:
self.graph = self.cuda_graph_wrapper()
else:
self.graph = None
def cuda_graph_wrapper(self):
moe_input = self.moe_inputs[0]
self.func(
**self.kwargs,
topk_ids=moe_input.topk_ids,
sorted_token_ids=moe_input.sorted_token_ids,
expert_ids=moe_input.expert_ids,
num_tokens_post_padded=moe_input.num_tokens_post_padded,
)
torch.cuda.synchronize()
# Capture 10 invocations with CUDA graph
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
for k in range(self.inner_iter):
moe_input = self.moe_inputs[k]
self.func(
**self.kwargs,
topk_ids=moe_input.topk_ids,
sorted_token_ids=moe_input.sorted_token_ids,
expert_ids=moe_input.expert_ids,
num_tokens_post_padded=moe_input.num_tokens_post_padded,
)
torch.cuda.synchronize()
# Warmup
for _ in range(5):
graph.replay()
torch.cuda.synchronize()
return graph
def forward_cost(self, try_cnt=2):
time_cost = float("inf")
for _ in range(try_cnt):
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
if self.use_cuda_graph:
self.graph.replay()
else:
for k in range(self.inner_iter):
moe_input = self.moe_inputs[k]
self.func(
**self.kwargs,
topk_ids=moe_input.topk_ids,
sorted_token_ids=moe_input.sorted_token_ids,
expert_ids=moe_input.expert_ids,
num_tokens_post_padded=moe_input.num_tokens_post_padded,
)
end_event.record()
torch.cuda.synchronize()
time_cost = min(time_cost, start_event.elapsed_time(end_event))
return time_cost
def load_topk_ids(topk_ids_dir, i: int):
num_layers = 61
dense_layers = 3
moe_layers = num_layers - dense_layers
return torch.load(
f"{topk_ids_dir}/topk_ids_layer{i % moe_layers + dense_layers}_idx{i // moe_layers}.pt"
)
def benchmark_config(
config: BenchmarkConfig,
num_tokens: int,
@@ -49,7 +132,7 @@ def benchmark_config(
use_fp8_w8a8: bool,
use_int8_w8a8: bool,
use_int8_w8a16: bool,
topk_ids_dir: str,
topk_ids_list,
block_shape: List[int] = None,
num_iters: int = 100,
) -> float:
@@ -86,7 +169,6 @@ def benchmark_config(
w2 = torch.randn(
num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype
)
gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32)
w1_scale = None
w2_scale = None
@@ -130,196 +212,187 @@ def benchmark_config(
top_k=topk,
renormalize=True,
)
topk_output = select_experts(hidden_states, input_gating, topk_config)
topk_output_ = select_experts(hidden_states, input_gating, topk_config)
sorted_token_ids_, expert_ids_, num_tokens_post_padded_ = moe_align_block_size(
topk_output_.topk_ids, config["BLOCK_SIZE_M"], num_experts
)
inner_iter = 10 if not ncu_enable else 1
moe_inputs = [
MoeInputs(
topk_output_.topk_ids.clone(),
sorted_token_ids_.clone(),
expert_ids_.clone(),
num_tokens_post_padded_.clone(),
)
for _ in range(inner_iter)
]
M = hidden_states.shape[0]
E, N, _ = w1.shape
def prepare(i: int):
input_gating = gating_output[i]
topk_ids = torch.load(f"{topk_ids_dir}/topk_ids_layer{i%58+3}_idx{i//58}.pt")
new_topk_output = select_experts(hidden_states, input_gating, topk_config)
topk_output.topk_weights.copy_(new_topk_output.topk_weights)
tokens, _topk = topk_output.topk_ids.shape
topk_output.topk_ids.copy_(topk_ids[:tokens, :_topk])
topk_output.router_logits.copy_(new_topk_output.router_logits)
padded_tokens = min(M * topk, E + 1) * (
config["BLOCK_SIZE_M"] - 1
) # if moe_use_tma else 0
total_tokens = M * topk + padded_tokens
cache = torch.empty(
total_tokens * max(N, w2.shape[1]),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
intermediate_cache1 = cache[: total_tokens * N].view(
(total_tokens, N),
)
intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
intermediate_cache3 = cache[: M * topk * w2.shape[1]].view(
(M, topk, w2.shape[1]),
)
moe_use_tma = False
def run():
moe_runner_config = MoeRunnerConfig(
inplace=True,
)
topk_weights, topk_ids, _ = topk_output
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
topk_ids, config["BLOCK_SIZE_M"], num_experts
)
M = hidden_states.shape[0]
E, N, _ = w1.shape
topk = topk_ids.shape[1]
padded_tokens = (
min(M * topk, E + 1) * (config["BLOCK_SIZE_M"] - 1) if moe_use_tma else 0
)
total_tokens = M * topk + padded_tokens
cache = torch.empty(
total_tokens * max(N, w2.shape[1]),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
intermediate_cache1 = cache[: total_tokens * N].view(
(total_tokens, N),
)
intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
intermediate_cache3 = cache[: M * topk * w2.shape[1]].view(
(M, topk, w2.shape[1]),
)
def prepare(i: int, inner_iter): # update inputs according to topk_ids
for k in range(inner_iter):
topk_ids = topk_ids_list[i * inner_iter + k]
tokens, _topk = moe_inputs[k].topk_ids.shape
moe_inputs[k].topk_ids.copy_(topk_ids[:tokens, :_topk])
sorted_token_ids_, expert_ids_, num_tokens_post_padded_ = (
moe_align_block_size(
moe_inputs[k].topk_ids, config["BLOCK_SIZE_M"], num_experts
)
)
moe_inputs[k].sorted_token_ids.copy_(sorted_token_ids_)
moe_inputs[k].expert_ids.copy_(expert_ids_)
moe_inputs[k].num_tokens_post_padded.copy_(num_tokens_post_padded_)
def get_kernel_wrapper(moe_use_tma, inner_iter, use_cuda_graph):
compute_type = (
tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
)
moe_runner_config = MoeRunnerConfig(
inplace=True,
)
apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input
kernel0 = KernelWrapper(
A=hidden_states,
B=w1,
bias=None,
C=intermediate_cache1,
A_scale=None,
B_scale=w1_scale,
B_zp=None,
topk_weights=topk_output_.topk_weights,
moe_inputs=moe_inputs,
mul_routed_weight=apply_router_weight_on_input,
top_k=topk,
config=config,
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
block_shape=block_shape,
b_use_tma=moe_use_tma,
c_sorted=moe_use_tma,
filter_expert=False,
use_cuda_graph=use_cuda_graph,
inner_iter=inner_iter,
)
kernel1 = KernelWrapper(
A=intermediate_cache2,
B=w2,
bias=None,
C=intermediate_cache3,
A_scale=a2_scale,
B_scale=w2_scale,
B_zp=None,
topk_weights=topk_output_.topk_weights,
moe_inputs=moe_inputs,
mul_routed_weight=not apply_router_weight_on_input,
top_k=1,
config=config,
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
block_shape=block_shape,
a_use_tma=moe_use_tma,
b_use_tma=moe_use_tma,
filter_expert=False,
use_cuda_graph=use_cuda_graph,
inner_iter=inner_iter,
)
return kernel0, kernel1
with override_config(config):
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
start_event.record()
for _ in range(10 if not ncu_enable else 1):
invoke_fused_moe_kernel(
hidden_states,
w1,
None,
intermediate_cache1,
None,
w1_scale,
None,
topk_weights,
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
apply_router_weight_on_input,
topk_ids.shape[1],
config,
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
block_shape=block_shape,
b_use_tma=moe_use_tma,
c_sorted=moe_use_tma,
filter_expert=False,
)
end_event.record()
end_event.synchronize()
time_cost0 = start_event.elapsed_time(end_event)
use_cuda_graph = True if not ncu_enable else False
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
start_event.record()
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
for _ in range(10 if not ncu_enable else 1):
invoke_fused_moe_kernel(
intermediate_cache2,
w2,
None,
intermediate_cache3,
a2_scale,
w2_scale,
None,
topk_weights,
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
not apply_router_weight_on_input,
1,
config,
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
block_shape=block_shape,
a_use_tma=moe_use_tma,
b_use_tma=moe_use_tma,
filter_expert=False,
)
end_event.record()
end_event.synchronize()
time_cost1 = start_event.elapsed_time(end_event)
return time_cost0, time_cost1
kernel0, kernel1 = get_kernel_wrapper(False, inner_iter, use_cuda_graph)
kernel_tma0, kernel_tma1 = get_kernel_wrapper(True, inner_iter, use_cuda_graph)
# JIT compilation & warmup
if not ncu_enable:
moe_use_tma = False
run()
moe_use_tma = True
run()
latencies: List[float] = []
latencies1: List[float] = []
latencies_tma: List[float] = []
latencies1_tma: List[float] = []
kernel0.forward_cost()
kernel1.forward_cost()
kernel_tma0.forward_cost()
kernel_tma1.forward_cost()
for i in range(num_iters):
prepare(i)
torch.cuda.synchronize()
moe_use_tma = False
t0, t1 = run()
torch.cuda.synchronize()
latencies.append(t0)
latencies1.append(t1)
ts0 = []
ts1 = []
ts_tma0 = []
ts_tma1 = []
moe_use_tma = True
t0, t1 = run()
torch.cuda.synchronize()
latencies_tma.append(t0)
latencies1_tma.append(t1)
for i in range(num_iters // inner_iter):
prepare(i, inner_iter)
ts0.append(kernel0.forward_cost())
ts1.append(kernel1.forward_cost())
ts_tma0.append(kernel_tma0.forward_cost())
ts_tma1.append(kernel_tma1.forward_cost())
torch.cuda.synchronize()
avg = sum(latencies) / (num_iters * 10) * 1000 # us
avg_tma = sum(latencies_tma) / (num_iters * 10) * 1000 # us
avg1 = sum(latencies1) / (num_iters * 10) * 1000 # us
avg1_tma = sum(latencies1_tma) / (num_iters * 10) * 1000 # us
avg = sum(ts0) / (num_iters) * 1000 # us
avg1 = sum(ts1) / (num_iters) * 1000 # us
avg_tma = sum(ts_tma0) / (num_iters) * 1000 # us
avg1_tma = sum(ts_tma1) / (num_iters) * 1000 # us
return avg, avg_tma, avg1, avg1_tma
class BestConfigTrace:
def __init__(self, name):
def __init__(self, name, down_moe=False):
self.name = name
self.config = None
self.time_cost = float("inf")
self.time_cost_all = None # kernel0 without tma,, kernel0 with tma, kernel1 without tma, kernel1 with tma
self.down_moe = down_moe
self.best_costs_m = {} # block_m: best_cost
def update(self, config, time_cost, time_cost_all):
if time_cost < self.time_cost:
print(
f"New best config for {self.name}: {config}, {time_cost=}, {time_cost_all=}, org: {self.config}, {self.time_cost_all}",
flush=True,
)
self.config = config
self.time_cost = time_cost
self.time_cost_all = time_cost_all
def update(self, config, time_cost_all):
block_m = config["BLOCK_SIZE_M"]
if not self.down_moe:
time_cost = time_cost_all[0]
else:
time_cost = min(time_cost_all[2], time_cost_all[3])
if (
block_m not in self.best_costs_m
or time_cost < self.best_costs_m[block_m][1]
):
self.best_costs_m[block_m] = config, time_cost, time_cost_all
@property
def total_time(self):
return self.time_cost_all[0] + min(self.time_cost_all[2], self.time_cost_all[3])
def time_cost(self, block_m):
if block_m not in self.best_costs_m:
return float("inf")
time_cost = self.best_costs_m[block_m][1]
return time_cost
def config_dict(self, down_moe=False):
if not down_moe:
return self.config
def config_dict(self, block_m):
if block_m not in self.best_costs_m:
return {}
config, _, time_cost_all = self.best_costs_m[block_m]
if not self.down_moe:
return config
else:
return {
**self.config,
"USE_TMA": self.time_cost_all[2] > self.time_cost_all[3],
**config,
"USE_TMA": time_cost_all[2] > time_cost_all[3],
}
@@ -349,13 +422,7 @@ class BenchmarkWorker:
topk_ids_dir: str,
) -> Tuple[Dict[str, int], float]:
torch.cuda.manual_seed_all(0)
dtype_str = get_config_dtype_str(
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
)
# NOTE(woosuk): The current naming convention uses w2.shape[2], which
# is the intermediate size after silu_and_mul.
block_n = block_shape[0] if block_shape else 0
block_k = block_shape[1] if block_shape else 0
topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)]
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
kernel_time = benchmark_config(
cfg,
@@ -368,7 +435,7 @@ class BenchmarkWorker:
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
topk_ids_dir,
topk_ids_list,
block_shape,
)
return cfg, kernel_time
@@ -388,9 +455,9 @@ class BenchmarkWorker:
search_space: List[Dict[str, int]],
topk_ids_dir: str,
) -> Dict[str, int]:
trace0 = BestConfigTrace("kernel0")
trace1 = BestConfigTrace("kernel1")
trace2 = BestConfigTrace("kernel all")
trace0 = BestConfigTrace("kernel0", down_moe=False)
trace1 = BestConfigTrace("kernel1", down_moe=True)
topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)]
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for config in tqdm(search_space):
@@ -406,56 +473,91 @@ class BenchmarkWorker:
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
topk_ids_dir,
topk_ids_list,
block_shape,
num_iters=10,
num_iters=100,
)
except triton.runtime.autotuner.OutOfResources:
# Some configurations may be invalid and fail to compile.
continue
kt0 = kt0_no_tma
kt1 = min(kt1_no_tma, kt1_tma)
trace0.update(
config,
kt0,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
trace1.update(
config,
kt1,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
trace2.update(
config,
kt0 + kt1,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
now = datetime.now()
print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}")
assert trace0.config is not None
assert trace1.config is not None
print(
f"{num_tokens=}, {trace0.config=}, {trace0.time_cost_all=}, {trace1.config=}, {trace1.time_cost_all=}"
)
if trace0.config["BLOCK_SIZE_M"] != trace1.config["BLOCK_SIZE_M"]:
best_trace = trace0 if trace0.total_time < trace1.total_time else trace1
best_trace = (
best_trace if best_trace.total_time < trace2.total_time else trace2
)
return (
best_trace.config_dict(),
best_trace.config_dict(True),
best_trace.time_cost_all,
best_trace.time_cost_all,
)
best_block_m = 16
for block_m in (32, 64, 128, 256):
if trace0.time_cost(block_m) + trace1.time_cost(block_m) < trace0.time_cost(
best_block_m
) + trace1.time_cost(best_block_m):
best_block_m = block_m
return (
trace0.config_dict(),
trace1.config_dict(True),
trace0.time_cost_all,
trace1.time_cost_all,
trace0.config_dict(best_block_m),
trace1.config_dict(best_block_m),
trace0.time_cost(best_block_m),
trace1.time_cost(best_block_m),
)
def cmp_configs(
self,
num_tokens: List[int],
num_experts: int,
shard_intermediate_size: int,
hidden_size: int,
topk: int,
dtype: torch.dtype,
use_fp8_w8a8: bool,
use_int8_w8a8: bool,
use_int8_w8a16: bool,
block_shape: List[int],
cmp_config_files: List[str],
topk_ids_dir: str,
):
# compare performance of different configs
cmp_configs = []
for file in cmp_config_files:
with open(file) as f:
cmp_configs.append({int(key): val for key, val in json.load(f).items()})
for i, file in enumerate(cmp_config_files):
print(f"config {i}: {file}")
topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)]
torch.cuda.manual_seed_all(0)
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for bs in num_tokens:
kernel_times = []
cfgs = []
for configs in cmp_configs:
cfg_org = configs[min(configs.keys(), key=lambda x: abs(x - bs))]
cfgs.append(cfg_org)
cfg = cfg_org.copy()
cfg.pop("USE_TMA", None)
kernel_time = benchmark_config(
cfg,
bs,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
topk_ids_list,
block_shape,
)
kernel_times.append(kernel_time)
print(f"batch_size={bs=}:")
for i, cfg in enumerate(cfgs):
print(f" config {i} {cfg}: {kernel_times[i]}")
def save_configs_sep(
configs: Dict[int, BenchmarkConfig],
@@ -521,6 +623,25 @@ def main(args: argparse.Namespace):
batch_sizes.reverse()
else:
batch_sizes = [args.batch_size]
if args.cmp_configs is not None:
worker = BenchmarkWorker(args.seed)
worker.cmp_configs(
batch_sizes,
E,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
block_shape,
args.cmp_configs,
topk_ids_dir,
)
return
if len(batch_sizes) == 1:
worker = BenchmarkWorker(args.seed)
if args.tune:
@@ -689,6 +810,7 @@ if __name__ == "__main__":
parser.add_argument("--disable-shared-experts-fusion", action="store_true")
parser.add_argument("--configs", type=int, nargs="+", required=False)
parser.add_argument("--topk-ids-dir", type=str, required=True)
parser.add_argument("--cmp-configs", type=str, nargs="+", required=False)
args = parser.parse_args()
main(args)
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_5_1/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json
+146
View File
@@ -0,0 +1,146 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 5
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 4
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"512": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 4
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"1536": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
}
}
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_5_1/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128]_down.json
+164
View File
@@ -0,0 +1,164 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"512": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"1536": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
}
}
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_5_1/E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json
+146
View File
@@ -0,0 +1,146 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 5
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 5
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 4
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 4
},
"512": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"1536": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
}
}
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_5_1/E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128]_down.json
+164
View File
@@ -0,0 +1,164 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"512": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"1536": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
}
}
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py
+2 -2
View File
@@ -744,8 +744,8 @@ def invoke_fused_moe_kernel(
B.stride(0),
B.stride(2),
B.stride(1),
C.stride(1),
C.stride(2),
C.stride(-2),
C.stride(-1),
B_scale.stride(0),
B_scale.stride(2),
B_scale.stride(1),