[kernel] add triton moe TMA up support (#33559)

Co-authored-by: undefined <zhouchen.arrebol@jd.com>
Co-authored-by: xq25478 <xq25478@qq.com>
Co-authored-by: xieminghe.simon <xieminghe.simon@jd.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py

@@ -143,6 +143,9 @@ def benchmark_config(
     block_shape: List[int] = None,
     ep_size: int = 1,
     num_iters: int = 100,
+    enable_up_tma: bool = False,
+    tune_round: str = "both",
+    down_use_tma_map: dict = None,
 ) -> float:
     ncu_enable = os.getenv("NCU_ENABLE", "0") == "1"
     if ncu_enable:
@@ -312,14 +315,29 @@ def benchmark_config(
             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
-        compute_type = (
+    moe_runner_config = MoeRunnerConfig(
-            tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
+        inplace=True,
-        )
+    )
-        moe_runner_config = MoeRunnerConfig(
+    apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input
-            inplace=True,
+
-        )
+    use_cuda_graph = True if not ncu_enable else False
-        apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input
+
+    # Determine which kernels to build based on tune_round:
+    #   "both"  — up (c_sorted=False, no TMA) + down (no-tma + tma)  [default, enable_up_tma=False]
+    #   "down"  — down (no-tma + tma) only                            [round 1 of two-round tune]
+    #   "up"    — up (no-tma + tma) only, c_sorted from down_use_tma_map [round 2 of two-round tune]
+    build_up = tune_round in ("both", "up")
+    build_down = tune_round in ("both", "down")
+
+    # For "up" round, c_sorted must match what down TMA decided at runtime
+    if tune_round == "up":
+        c_sorted = down_use_tma_map.get(config["BLOCK_SIZE_M"], False)
+    else:
+        c_sorted = False
+
+    up_kernels = []
+    if build_up:
         kernel0 = KernelWrapper(
             A=hidden_states,
             B=w1,
@@ -340,12 +358,44 @@ def benchmark_config(
             use_int4_w4a16=use_int4_w4a16,
             per_channel_quant=False,
             block_shape=block_shape,
-            b_use_tma=moe_use_tma,
+            b_use_tma=False,
-            c_sorted=moe_use_tma,
+            c_sorted=c_sorted,
             filter_expert=False,
             use_cuda_graph=use_cuda_graph,
             inner_iter=inner_iter,
         )
+        up_kernels.append(kernel0)
+        if enable_up_tma or tune_round == "up":
+            kernel0_tma = KernelWrapper(
+                A=hidden_states,
+                B=w1,
+                bias=None,
+                C=intermediate_cache1,
+                A_scale=a1_scale,
+                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=use_int8_w8a8,
+                use_int8_w8a16=use_int8_w8a16,
+                use_int4_w4a16=use_int4_w4a16,
+                per_channel_quant=False,
+                block_shape=block_shape,
+                b_use_tma=True,
+                c_sorted=c_sorted,
+                filter_expert=False,
+                use_cuda_graph=use_cuda_graph,
+                inner_iter=inner_iter,
+            )
+            up_kernels.append(kernel0_tma)
+
+    down_kernels = []
+    if build_down:
         kernel1 = KernelWrapper(
             A=intermediate_cache2,
             B=w2,
@@ -366,25 +416,45 @@ def benchmark_config(
             use_int4_w4a16=use_int4_w4a16,
             per_channel_quant=False,
             block_shape=block_shape,
-            a_use_tma=moe_use_tma,
+            a_use_tma=False,
-            b_use_tma=moe_use_tma,
+            b_use_tma=False,
             filter_expert=False,
             use_cuda_graph=use_cuda_graph,
             inner_iter=inner_iter,
         )
-        return kernel0, kernel1
+        down_kernels.append(kernel1)
-
+        kernel1_tma = KernelWrapper(
-    use_cuda_graph = True if not ncu_enable else False
+            A=intermediate_cache2,
-
+            B=w2,
-    kernel0, kernel1 = get_kernel_wrapper(False, inner_iter, use_cuda_graph)
+            bias=None,
-    kernel_tma0, kernel_tma1 = get_kernel_wrapper(True, inner_iter, use_cuda_graph)
+            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=use_int8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            per_channel_quant=False,
+            block_shape=block_shape,
+            a_use_tma=True,
+            b_use_tma=True,
+            filter_expert=False,
+            use_cuda_graph=use_cuda_graph,
+            inner_iter=inner_iter,
+        )
+        down_kernels.append(kernel1_tma)
 
     # JIT compilation & warmup
     if not ncu_enable:
-        kernel0.forward_cost()
+        for k in up_kernels + down_kernels:
-        kernel1.forward_cost()
+            k.forward_cost()
-        kernel_tma0.forward_cost()
-        kernel_tma1.forward_cost()
 
     ts0 = []
     ts1 = []
@@ -393,31 +463,42 @@ def benchmark_config(
 
     for i in range(num_iters // inner_iter):
         prepare(i, inner_iter)
-        ts0.append(kernel0.forward_cost())
+        if build_up:
-        ts1.append(kernel1.forward_cost())
+            ts0.append(kernel0.forward_cost())  # up no-tma
-        ts_tma0.append(kernel_tma0.forward_cost())
+            if len(up_kernels) > 1:
-        ts_tma1.append(kernel_tma1.forward_cost())
+                ts_tma0.append(kernel0_tma.forward_cost())  # up tma
+        if build_down:
+            ts1.append(kernel1.forward_cost())  # down no-tma
+            ts_tma1.append(kernel1_tma.forward_cost())  # down tma
     torch.cuda.synchronize()
 
-    avg = sum(ts0) / (num_iters) * 1000  # us
+    avg = sum(ts0) / (num_iters) * 1000 if ts0 else float("inf")
-    avg1 = sum(ts1) / (num_iters) * 1000  # us
+    avg1 = sum(ts1) / (num_iters) * 1000 if ts1 else float("inf")
-    avg_tma = sum(ts_tma0) / (num_iters) * 1000  # us
+    avg_tma = sum(ts_tma0) / (num_iters) * 1000 if ts_tma0 else float("inf")
-    avg1_tma = sum(ts_tma1) / (num_iters) * 1000  # us
+    avg1_tma = sum(ts_tma1) / (num_iters) * 1000 if ts_tma1 else float("inf")
 
     return avg, avg_tma, avg1, avg1_tma
 
 
 class BestConfigTrace:
-    def __init__(self, name, down_moe=False):
+    def __init__(self, name, down_moe=False, enable_up_tma=False):
         self.name = name
         self.down_moe = down_moe
+        self.enable_up_tma = enable_up_tma
         self.best_costs_m = {}  # block_m: best_cost
 
     def update(self, config, time_cost_all):
         block_m = config["BLOCK_SIZE_M"]
         if not self.down_moe:
-            time_cost = time_cost_all[0]
+            # time_cost_all = (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma)
+            if self.enable_up_tma:
+                # For up_proj, pick the faster of TMA vs no-TMA.
+                time_cost = min(time_cost_all[0], time_cost_all[1])
+            else:
+                # Up TMA not enabled — always use no-TMA cost.
+                time_cost = time_cost_all[0]
         else:
+            # For down_proj, pick the faster of TMA vs no-TMA.
             time_cost = min(time_cost_all[2], time_cost_all[3])
         if (
             block_m not in self.best_costs_m
@@ -436,7 +517,16 @@ class BestConfigTrace:
             return {}
         config, _, time_cost_all = self.best_costs_m[block_m]
         if not self.down_moe:
-            return config
+            if self.enable_up_tma:
+                # up_proj: use TMA when the TMA variant is faster.
+                return {
+                    **config,
+                    "USE_TMA": time_cost_all[0] > time_cost_all[1],
+                }
+            else:
+                # Up TMA not enabled — do not add USE_TMA key so the runtime
+                # defaults to no-TMA for up-projection.
+                return config
         else:
             return {
                 **config,
@@ -471,26 +561,72 @@ class BenchmarkWorker:
         cfg: Dict[str, int],
         topk_ids_dir: str,
         ep_size: int = 1,
+        enable_up_tma: bool = False,
     ) -> Tuple[Dict[str, int], float]:
         torch.cuda.manual_seed_all(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(
+            if enable_up_tma:
-                cfg,
+                # Two-step: first measure down to determine c_sorted,
-                num_tokens,
+                # then measure up with the correct c_sorted.
-                num_experts,
+                _, _, kt1_no_tma, kt1_tma = benchmark_config(
-                shard_intermediate_size,
+                    cfg,
-                hidden_size,
+                    num_tokens,
-                topk,
+                    num_experts,
-                dtype,
+                    shard_intermediate_size,
-                use_fp8_w8a8,
+                    hidden_size,
-                use_int8_w8a8,
+                    topk,
-                use_int8_w8a16,
+                    dtype,
-                use_int4_w4a16,
+                    use_fp8_w8a8,
-                topk_ids_list,
+                    use_int8_w8a8,
-                block_shape,
+                    use_int8_w8a16,
-                ep_size=ep_size,
+                    use_int4_w4a16,
-            )
+                    topk_ids_list,
+                    block_shape,
+                    ep_size=ep_size,
+                    enable_up_tma=True,
+                    tune_round="down",
+                )
+                down_use_tma = kt1_no_tma > kt1_tma
+                kt0_no_tma, kt0_tma, _, _ = benchmark_config(
+                    cfg,
+                    num_tokens,
+                    num_experts,
+                    shard_intermediate_size,
+                    hidden_size,
+                    topk,
+                    dtype,
+                    use_fp8_w8a8,
+                    use_int8_w8a8,
+                    use_int8_w8a16,
+                    use_int4_w4a16,
+                    topk_ids_list,
+                    block_shape,
+                    ep_size=ep_size,
+                    enable_up_tma=True,
+                    tune_round="up",
+                    down_use_tma_map={cfg["BLOCK_SIZE_M"]: down_use_tma},
+                )
+                kernel_time = (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma)
+            else:
+                kernel_time = benchmark_config(
+                    cfg,
+                    num_tokens,
+                    num_experts,
+                    shard_intermediate_size,
+                    hidden_size,
+                    topk,
+                    dtype,
+                    use_fp8_w8a8,
+                    use_int8_w8a8,
+                    use_int8_w8a16,
+                    use_int4_w4a16,
+                    topk_ids_list,
+                    block_shape,
+                    ep_size=ep_size,
+                    enable_up_tma=False,
+                    tune_round="both",
+                )
         return cfg, kernel_time
 
     def tune(
@@ -509,42 +645,124 @@ class BenchmarkWorker:
         search_space: List[Dict[str, int]],
         topk_ids_dir: str,
         ep_size: int = 1,
+        enable_up_tma: bool = False,
     ) -> Dict[str, int]:
-        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():
+        if not enable_up_tma:
-            for config in tqdm(search_space):
+            # Default path: single round, up c_sorted=False, no up TMA.
-                try:
+            # Down TMA is still tuned.
-                    kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma = benchmark_config(
+            trace0 = BestConfigTrace("kernel0", down_moe=False, enable_up_tma=False)
+            trace1 = BestConfigTrace("kernel1", down_moe=True)
+
+            with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
+                for config in tqdm(search_space):
+                    try:
+                        kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma = benchmark_config(
+                            config,
+                            num_tokens,
+                            num_experts,
+                            shard_intermediate_size,
+                            hidden_size,
+                            topk,
+                            dtype,
+                            use_fp8_w8a8,
+                            use_int8_w8a8,
+                            use_int8_w8a16,
+                            use_int4_w4a16,
+                            topk_ids_list,
+                            block_shape,
+                            ep_size=ep_size,
+                            num_iters=100,
+                            enable_up_tma=False,
+                            tune_round="both",
+                        )
+                    except triton.runtime.autotuner.OutOfResources:
+                        continue
+                    trace0.update(
                         config,
-                        num_tokens,
+                        (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
-                        num_experts,
+                    )
-                        shard_intermediate_size,
+                    trace1.update(
-                        hidden_size,
+                        config,
-                        topk,
+                        (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
-                        dtype,
+                    )
-                        use_fp8_w8a8,
+        else:
-                        use_int8_w8a8,
+            # Two-round coupled tuning: first tune down to determine c_sorted,
-                        use_int8_w8a16,
+            # then tune up with the correct c_sorted that matches runtime.
-                        use_int4_w4a16,
+            trace0 = BestConfigTrace("kernel0", down_moe=False, enable_up_tma=True)
-                        topk_ids_list,
+            trace1 = BestConfigTrace("kernel1", down_moe=True)
-                        block_shape,
+
-                        ep_size=ep_size,
+            # === Round 1: Down-only ===
-                        num_iters=100,
+            with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
+                for config in tqdm(search_space, desc="Round 1 (down)"):
+                    try:
+                        _, _, kt1_no_tma, kt1_tma = benchmark_config(
+                            config,
+                            num_tokens,
+                            num_experts,
+                            shard_intermediate_size,
+                            hidden_size,
+                            topk,
+                            dtype,
+                            use_fp8_w8a8,
+                            use_int8_w8a8,
+                            use_int8_w8a16,
+                            use_int4_w4a16,
+                            topk_ids_list,
+                            block_shape,
+                            ep_size=ep_size,
+                            num_iters=100,
+                            enable_up_tma=True,
+                            tune_round="down",
+                        )
+                    except triton.runtime.autotuner.OutOfResources:
+                        continue
+                    trace1.update(
+                        config,
+                        (float("inf"), float("inf"), kt1_no_tma, kt1_tma),
+                    )
+
+            # Extract down TMA decision per BLOCK_SIZE_M from round 1 results
+            down_use_tma_map = {}
+            for block_m, (_, _, time_cost_all) in trace1.best_costs_m.items():
+                down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3]
+
+            print(
+                f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: "
+                f"{down_use_tma_map}"
+            )
+
+            # === Round 2: Up with c_sorted from round 1 ===
+            with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
+                for config in tqdm(search_space, desc="Round 2 (up)"):
+                    try:
+                        kt0_no_tma, kt0_tma, _, _ = benchmark_config(
+                            config,
+                            num_tokens,
+                            num_experts,
+                            shard_intermediate_size,
+                            hidden_size,
+                            topk,
+                            dtype,
+                            use_fp8_w8a8,
+                            use_int8_w8a8,
+                            use_int8_w8a16,
+                            use_int4_w4a16,
+                            topk_ids_list,
+                            block_shape,
+                            ep_size=ep_size,
+                            num_iters=100,
+                            enable_up_tma=True,
+                            tune_round="up",
+                            down_use_tma_map=down_use_tma_map,
+                        )
+                    except triton.runtime.autotuner.OutOfResources:
+                        continue
+                    trace0.update(
+                        config,
+                        (kt0_no_tma, kt0_tma, float("inf"), float("inf")),
                     )
-                except triton.runtime.autotuner.OutOfResources:
-                    # Some configurations may be invalid and fail to compile.
-                    continue
-                trace0.update(
-                    config,
-                    (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
-                )
-                trace1.update(
-                    config,
-                    (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
-                )
 
         now = datetime.now()
         print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}")
@@ -731,6 +949,7 @@ def main(args: argparse.Namespace):
                 search_space,
                 topk_ids_dir,
                 args.ep_size,
+                enable_up_tma=args.enable_tune_up_tma,
             )
         else:
             cfg = {
@@ -757,6 +976,7 @@ def main(args: argparse.Namespace):
                 cfg,
                 topk_ids_dir,
                 args.ep_size,
+                enable_up_tma=args.enable_tune_up_tma,
             )
             print(f"{t0=}, {t0_tma=}, {t1=}, {t1_tma=}")
         return
@@ -823,6 +1043,7 @@ def main(args: argparse.Namespace):
                 search_space,
                 topk_ids_dir,
                 args.ep_size,
+                args.enable_tune_up_tma,
             )
             for batch_size in batch_sizes
         ],
@@ -890,6 +1111,14 @@ if __name__ == "__main__":
     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)
+    parser.add_argument(
+        "--enable-tune-up-tma",
+        action="store_true",
+        help="Enable up-projection TMA tuning in addition to down-projection TMA. "
+        "When set, the up config file will contain a USE_TMA flag. "
+        "When not set (default), only down-projection TMA is tuned and the up "
+        "config will not contain a USE_TMA key.",
+    )
     args = parser.parse_args()
 
     main(args)

python/sglang/srt/layers/moe/moe_runner/triton.py

@@ -138,6 +138,7 @@ class TritonRunnerCore(MoeRunnerCore):
             running_state["config"],
             running_state.get("down_config"),
             running_state.get("down_moe_use_tma", False),
+            running_state.get("up_moe_use_tma", False),
             b1=quant_info.b13,
             b2=quant_info.b2,
             use_fp8_w8a8=quant_info.use_fp8_w8a8,
@@ -280,6 +281,7 @@ def pre_permute_standard_to_triton(
         config,
         down_config,
         down_moe_use_tma,
+        up_moe_use_tma,
         sorted_token_ids,
         expert_ids,
         num_tokens_post_padded,
@@ -299,6 +301,7 @@ def pre_permute_standard_to_triton(
     running_state["config"] = config
     running_state["down_config"] = down_config
     running_state["down_moe_use_tma"] = down_moe_use_tma
+    running_state["up_moe_use_tma"] = up_moe_use_tma
 
     return TritonRunnerInput(
         hidden_states=hidden_states,

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json

@@ -0,0 +1,164 @@
+{
+    "1": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": false
+    },
+    "2": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": false
+    },
+    "4": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": false
+    },
+    "8": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": false
+    },
+    "16": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": false
+    },
+    "24": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": true
+    },
+    "32": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": false
+    },
+    "48": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": true
+    },
+    "64": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 5,
+        "USE_TMA": true
+    },
+    "96": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "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": 3,
+        "USE_TMA": true
+    },
+    "256": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 4,
+        "USE_TMA": true
+    },
+    "512": {
+        "BLOCK_SIZE_M": 32,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "1024": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "1536": {
+        "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
+    },
+    "2048": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 4,
+        "USE_TMA": false
+    },
+    "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": false
+    },
+    "4096": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 4,
+        "USE_TMA": false
+    }
+}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,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": 128,
+        "GROUP_SIZE_M": 16,
+        "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": 16,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "4": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "8": {
+        "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
+    },
+    "16": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "24": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "32": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "48": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "64": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "96": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "128": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "256": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "512": {
+        "BLOCK_SIZE_M": 32,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "1024": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    },
+    "1536": {
+        "BLOCK_SIZE_M": 64,
+        "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": 2,
+        "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": 2,
+        "USE_TMA": true
+    },
+    "4096": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "USE_TMA": true
+    }
+}

python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import functools
+import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import torch
@@ -93,6 +94,8 @@ if not _is_cuda and not _is_hip and not _is_xpu:
 
 padding_size = get_moe_padding_size(_use_aiter)
 
+logger = logging.getLogger(__name__)
+
 
 def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool:
     return num_tokens <= 32 and not is_batch_invariant_mode_enabled()
@@ -363,7 +366,7 @@ def swiglu_no_interleaved_with_alpha_and_limit(x, gemm1_alpha, gemm1_limit):
 
 
 @functools.lru_cache()
-def _down_moe_use_tma():
+def _moe_support_tma():
     return support_tensor_descriptor()
 
 
@@ -409,11 +412,26 @@ def _prepare_fused_moe_run(
         per_channel_quant=per_channel_quant,
         return_down_config=True,
     )
-    down_moe_use_tma = (
+    # Copy config to avoid mutating the lru_cached dict returned by
-        _down_moe_use_tma()
+    # get_moe_configs; we pop USE_TMA below.
-        and down_config is not None
+    config = dict(config)
-        and down_config.pop("USE_TMA", False)
+    # Up-projection TMA is opt-in: only enabled when the up config file
-    )
+    # explicitly carries "USE_TMA": true (produced by tuning). By default the
+    # existing up config files do not contain this key, so existing users are
+    # unaffected unless they re-tune with the updated script.
+    up_tma_requested = config.pop("USE_TMA", False)
+    up_moe_use_tma = _moe_support_tma() and up_tma_requested
+    if up_moe_use_tma:
+        logger.warning_once(
+            "Up MoE TMA is enabled (USE_TMA=true in the up-projection config). "
+            "This requires a config produced by the updated tuning script. "
+        )
+    down_tma_requested = down_config is not None and down_config.pop("USE_TMA", False)
+    down_moe_use_tma = _moe_support_tma() and down_tma_requested
+    if down_moe_use_tma:
+        logger.warning_once(
+            "Down MoE TMA is enabled (USE_TMA=true in the down-projection config)."
+        )
 
     sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
         topk_ids, config["BLOCK_SIZE_M"], E
@@ -423,6 +441,7 @@ def _prepare_fused_moe_run(
         config,
         down_config,
         down_moe_use_tma,
+        up_moe_use_tma,
         sorted_token_ids,
         expert_ids,
         num_tokens_post_padded,
@@ -441,6 +460,7 @@ def _fused_moe_kernel_sequence(
     config: Dict[str, Any],
     down_config: Optional[Dict[str, Any]],
     down_moe_use_tma: bool,
+    up_moe_use_tma: bool,
     *,
     b1: Optional[torch.Tensor],
     b2: Optional[torch.Tensor],
@@ -568,6 +588,7 @@ def _fused_moe_kernel_sequence(
         per_channel_quant=per_channel_quant,
         block_shape=block_shape,
         c_sorted=down_moe_use_tma,
+        b_use_tma=up_moe_use_tma,
         filter_expert=filter_expert,
     )
 
@@ -927,6 +948,7 @@ def fused_experts_impl(
         config,
         down_config,
         down_moe_use_tma,
+        up_moe_use_tma,
         sorted_token_ids,
         expert_ids,
         num_tokens_post_padded,
@@ -955,6 +977,7 @@ def fused_experts_impl(
         config,
         down_config,
         down_moe_use_tma,
+        up_moe_use_tma,
         b1=b1,
         b2=b2,
         use_fp8_w8a8=use_fp8_w8a8,