[Fix] Triton TP MoE Dpsk V3/Qwen3 Coder with SwapAB (#17965)

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py

@@ -134,6 +134,7 @@ def benchmark_config(
     use_int8_w8a16: bool,
     topk_ids_list,
     block_shape: List[int] = None,
+    ep_size: int = 1,
     num_iters: int = 100,
 ) -> float:
     ncu_enable = os.getenv("NCU_ENABLE", "0") == "1"
@@ -253,6 +254,12 @@ def benchmark_config(
     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]
+            # With EP, saved topk_ids are global expert indices; remap to local.
+            if ep_size > 1:
+                topk_ids = (topk_ids // ep_size).to(
+                    device=moe_inputs[k].topk_ids.device,
+                    dtype=moe_inputs[k].topk_ids.dtype,
+                )
             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_ = (
@@ -420,6 +427,7 @@ class BenchmarkWorker:
         block_shape: List[int],
         cfg: Dict[str, int],
         topk_ids_dir: str,
+        ep_size: int = 1,
     ) -> 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)]
@@ -437,6 +445,7 @@ class BenchmarkWorker:
                 use_int8_w8a16,
                 topk_ids_list,
                 block_shape,
+                ep_size=ep_size,
             )
         return cfg, kernel_time
 
@@ -454,6 +463,7 @@ class BenchmarkWorker:
         block_shape: List[int],
         search_space: List[Dict[str, int]],
         topk_ids_dir: str,
+        ep_size: int = 1,
     ) -> Dict[str, int]:
         trace0 = BestConfigTrace("kernel0", down_moe=False)
         trace1 = BestConfigTrace("kernel1", down_moe=True)
@@ -475,6 +485,7 @@ class BenchmarkWorker:
                         use_int8_w8a16,
                         topk_ids_list,
                         block_shape,
+                        ep_size=ep_size,
                         num_iters=100,
                     )
                 except triton.runtime.autotuner.OutOfResources:
@@ -519,6 +530,7 @@ class BenchmarkWorker:
         block_shape: List[int],
         cmp_config_files: List[str],
         topk_ids_dir: str,
+        ep_size: int = 1,
     ):
         # compare performance of different configs
         cmp_configs = []
@@ -552,6 +564,7 @@ class BenchmarkWorker:
                         use_int8_w8a16,
                         topk_ids_list,
                         block_shape,
+                        ep_size=ep_size,
                     )
                     kernel_times.append(kernel_time)
                 print(f"batch_size={bs=}:")
@@ -639,6 +652,7 @@ def main(args: argparse.Namespace):
             block_shape,
             args.cmp_configs,
             topk_ids_dir,
+            args.ep_size,
         )
         return
 
@@ -659,6 +673,7 @@ def main(args: argparse.Namespace):
                 block_shape,
                 search_space,
                 topk_ids_dir,
+                args.ep_size,
             )
         else:
             cfg = {
@@ -683,6 +698,7 @@ def main(args: argparse.Namespace):
                 block_shape,
                 cfg,
                 topk_ids_dir,
+                args.ep_size,
             )
             print(f"{t0=}, {t0_tma=}, {t1=}, {t1_tma=}")
         return
@@ -746,6 +762,7 @@ def main(args: argparse.Namespace):
                 block_shape,
                 search_space,
                 topk_ids_dir,
+                args.ep_size,
             )
             for batch_size in batch_sizes
         ],