diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py index afee79940..d0c922a4c 100644 --- a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py +++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py @@ -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) diff --git a/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 b/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 new file mode 100644 index 000000000..317aa18cd --- /dev/null +++ b/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 @@ -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 + } +} diff --git a/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 b/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 new file mode 100644 index 000000000..8cfd871fe --- /dev/null +++ b/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 @@ -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 + } +} diff --git a/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 b/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 new file mode 100644 index 000000000..a240dd111 --- /dev/null +++ b/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 @@ -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 + } +} diff --git a/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 b/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 new file mode 100644 index 000000000..a2f6db88b --- /dev/null +++ b/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 @@ -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 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py index 4f0f6bb76..230b64057 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -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),