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

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):
+    padded_tokens = min(M * topk, E + 1) * (
-        input_gating = gating_output[i]
+        config["BLOCK_SIZE_M"] - 1
-        topk_ids = torch.load(f"{topk_ids_dir}/topk_ids_layer{i%58+3}_idx{i//58}.pt")
+    )  # if moe_use_tma else 0
-        new_topk_output = select_experts(hidden_states, input_gating, topk_config)
+    total_tokens = M * topk + padded_tokens
-        topk_output.topk_weights.copy_(new_topk_output.topk_weights)
+    cache = torch.empty(
-        tokens, _topk = topk_output.topk_ids.shape
+        total_tokens * max(N, w2.shape[1]),
-        topk_output.topk_ids.copy_(topk_ids[:tokens, :_topk])
+        device=hidden_states.device,
-        topk_output.router_logits.copy_(new_topk_output.router_logits)
+        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 prepare(i: int, inner_iter):  # update inputs according to topk_ids
-
+        for k in range(inner_iter):
-    def run():
+            topk_ids = topk_ids_list[i * inner_iter + k]
-        moe_runner_config = MoeRunnerConfig(
+            tokens, _topk = moe_inputs[k].topk_ids.shape
-            inplace=True,
+            moe_inputs[k].topk_ids.copy_(topk_ids[:tokens, :_topk])
-        )
+            sorted_token_ids_, expert_ids_, num_tokens_post_padded_ = (
-        topk_weights, topk_ids, _ = topk_output
+                moe_align_block_size(
-
+                    moe_inputs[k].topk_ids, config["BLOCK_SIZE_M"], num_experts
-        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                )
-            topk_ids, config["BLOCK_SIZE_M"], num_experts
+            )
-        )
+            moe_inputs[k].sorted_token_ids.copy_(sorted_token_ids_)
-        M = hidden_states.shape[0]
+            moe_inputs[k].expert_ids.copy_(expert_ids_)
-        E, N, _ = w1.shape
+            moe_inputs[k].num_tokens_post_padded.copy_(num_tokens_post_padded_)
-
-        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 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):
+    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()
-            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)
 
-            start_event = torch.cuda.Event(enable_timing=True)
+    kernel0, kernel1 = get_kernel_wrapper(False, inner_iter, use_cuda_graph)
-            end_event = torch.cuda.Event(enable_timing=True)
+    kernel_tma0, kernel_tma1 = get_kernel_wrapper(True, inner_iter, use_cuda_graph)
-            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
 
     # JIT compilation & warmup
     if not ncu_enable:
-        moe_use_tma = False
+        kernel0.forward_cost()
-        run()
+        kernel1.forward_cost()
-        moe_use_tma = True
+        kernel_tma0.forward_cost()
-        run()
+        kernel_tma1.forward_cost()
-    latencies: List[float] = []
-    latencies1: List[float] = []
-    latencies_tma: List[float] = []
-    latencies1_tma: List[float] = []
 
-    for i in range(num_iters):
+    ts0 = []
-        prepare(i)
+    ts1 = []
-        torch.cuda.synchronize()
+    ts_tma0 = []
-        moe_use_tma = False
+    ts_tma1 = []
-        t0, t1 = run()
-        torch.cuda.synchronize()
-        latencies.append(t0)
-        latencies1.append(t1)
 
-        moe_use_tma = True
+    for i in range(num_iters // inner_iter):
-        t0, t1 = run()
+        prepare(i, inner_iter)
-        torch.cuda.synchronize()
+        ts0.append(kernel0.forward_cost())
-        latencies_tma.append(t0)
+        ts1.append(kernel1.forward_cost())
-        latencies1_tma.append(t1)
+        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 = sum(ts0) / (num_iters) * 1000  # us
-    avg_tma = sum(latencies_tma) / (num_iters * 10) * 1000  # us
+    avg1 = sum(ts1) / (num_iters) * 1000  # us
-    avg1 = sum(latencies1) / (num_iters * 10) * 1000  # us
+    avg_tma = sum(ts_tma0) / (num_iters) * 1000  # us
-    avg1_tma = sum(latencies1_tma) / (num_iters * 10) * 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.down_moe = down_moe
-        self.time_cost = float("inf")
+        self.best_costs_m = {}  # block_m: best_cost
-        self.time_cost_all = None  # kernel0 without tma,, kernel0 with tma, kernel1 without tma, kernel1 with tma
 
-    def update(self, config, time_cost, time_cost_all):
+    def update(self, config, time_cost_all):
-        if time_cost < self.time_cost:
+        block_m = config["BLOCK_SIZE_M"]
-            print(
+        if not self.down_moe:
-                f"New best config for {self.name}: {config}, {time_cost=}, {time_cost_all=}, org: {self.config}, {self.time_cost_all}",
+            time_cost = time_cost_all[0]
-                flush=True,
+        else:
-            )
+            time_cost = min(time_cost_all[2], time_cost_all[3])
-            self.config = config
+        if (
-            self.time_cost = time_cost
+            block_m not in self.best_costs_m
-            self.time_cost_all = time_cost_all
+            or time_cost < self.best_costs_m[block_m][1]
+        ):
+            self.best_costs_m[block_m] = config, time_cost, time_cost_all
 
-    @property
+    def time_cost(self, block_m):
-    def total_time(self):
+        if block_m not in self.best_costs_m:
-        return self.time_cost_all[0] + min(self.time_cost_all[2], self.time_cost_all[3])
+            return float("inf")
+        time_cost = self.best_costs_m[block_m][1]
+        return time_cost
 
-    def config_dict(self, down_moe=False):
+    def config_dict(self, block_m):
-        if not down_moe:
+        if block_m not in self.best_costs_m:
-            return self.config
+            return {}
+        config, _, time_cost_all = self.best_costs_m[block_m]
+        if not self.down_moe:
+            return config
         else:
             return {
-                **self.config,
+                **config,
-                "USE_TMA": self.time_cost_all[2] > self.time_cost_all[3],
+                "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(
+        topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)]
-            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
         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")
+        trace0 = BestConfigTrace("kernel0", down_moe=False)
-        trace1 = BestConfigTrace("kernel1")
+        trace1 = BestConfigTrace("kernel1", down_moe=True)
-        trace2 = BestConfigTrace("kernel all")
+        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
+        best_block_m = 16
-        assert trace1.config is not None
+        for block_m in (32, 64, 128, 256):
-        print(
+            if trace0.time_cost(block_m) + trace1.time_cost(block_m) < trace0.time_cost(
-            f"{num_tokens=}, {trace0.config=}, {trace0.time_cost_all=}, {trace1.config=}, {trace1.time_cost_all=}"
+                best_block_m
-        )
+            ) + trace1.time_cost(best_block_m):
-        if trace0.config["BLOCK_SIZE_M"] != trace1.config["BLOCK_SIZE_M"]:
+                best_block_m = block_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,
-            )
         return (
-            trace0.config_dict(),
+            trace0.config_dict(best_block_m),
-            trace1.config_dict(True),
+            trace1.config_dict(best_block_m),
-            trace0.time_cost_all,
+            trace0.time_cost(best_block_m),
-            trace1.time_cost_all,
+            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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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),