test(ep): add configurable CLI knobs to dispatch/combine benchmark

tests/python/ops/bench_dispatch_combine.py

@@ -52,8 +52,10 @@ def __init__(
         force_scale_active=False,
         report_scale_stats=False,
         combine_scale_dim=None,
+        routing="random",
     ):
         super().__init__(config)
+        self.routing = routing
         self.combine_data_type = (
             combine_data_type if combine_data_type is not None else config.data_type
         )
@@ -83,6 +85,7 @@ def gen_test_data(self):
         self.config.hidden_dim = self.dispatch_hidden_dim
         result = super().gen_test_data(
             use_max_token_num=True,
+            routing=self.routing,
             input_dist=self.input_dist,
             input_scale=self.input_scale,
             input_shift=self.input_shift,
@@ -481,6 +484,7 @@ def run(
         graph_replay_iters=10,
         skip_e2e=False,
         call_local_expert_count=False,
+        verify=True,
     ):
         test_data = self.gen_test_data()
         for _ in range(warmup):
@@ -491,6 +495,7 @@ def run(
                 dispatch_warp_per_block,
                 combine_block_num,
                 combine_warp_per_block,
+                check=verify,
                 call_local_expert_count=call_local_expert_count,
             )
 
@@ -841,6 +846,29 @@ def _save_intranode_tuning_result(
 )
 
 
+def _optional_kwargs(**kwargs):
+    """Drop keys whose value is None, so the callee's own default applies."""
+    return {k: v for k, v in kwargs.items() if v is not None}
+
+
+KERNEL_TYPE_CHOICES = {
+    "IntraNode": mori.ops.EpDispatchCombineKernelType.IntraNode,
+    "IntraNodeLL": mori.ops.EpDispatchCombineKernelType.IntraNodeLL,
+    "InterNode": mori.ops.EpDispatchCombineKernelType.InterNode,
+    "InterNodeV1": mori.ops.EpDispatchCombineKernelType.InterNodeV1,
+    "InterNodeV1LL": mori.ops.EpDispatchCombineKernelType.InterNodeV1LL,
+    "AsyncLL": mori.ops.EpDispatchCombineKernelType.AsyncLL,
+}
+
+
+def _resolve_bench_kernel_type(name):
+    if name not in KERNEL_TYPE_CHOICES:
+        raise ValueError(
+            f"Unknown kernel_type={name!r}; choose from {sorted(KERNEL_TYPE_CHOICES)}"
+        )
+    return KERNEL_TYPE_CHOICES[name]
+
+
 def _get_default_launch_config(
     world_size,
     max_num_inp_token_per_rank,
@@ -898,6 +926,12 @@ def _bench_dispatch_combine(
     input_shift=0.0,
     force_scale_active=False,
     report_scale_stats=False,
+    warmup=None,
+    iters=None,
+    verify=True,
+    routing="random",
+    kernel_type="IntraNode",
+    graph_replay_iters=None,
 ):
     if combine_data_type is None:
         combine_data_type = data_type
@@ -940,6 +974,7 @@ def _bench_dispatch_combine(
         use_external_inp_buf=not zero_copy,  # zero-copy mode requires use_external_inp_buf=False
         gpu_per_node=world_size,
         quant_type=quant_type,
+        kernel_type=_resolve_bench_kernel_type(kernel_type),
     )
     with TorchDistContext(rank=rank, world_size=world_size, master_port=port):
         mori.shmem.shmem_torch_process_group_init("default")
@@ -963,6 +998,7 @@ def _bench_dispatch_combine(
             force_scale_active=force_scale_active,
             report_scale_stats=report_scale_stats,
             combine_scale_dim=bench_combine_scale_dim,
+            routing=routing,
         )
 
         (
@@ -1000,6 +1036,10 @@ def _bench_dispatch_combine(
                 combine_block_num=combine_block_num,
                 combine_warp_per_block=combine_warp_per_block,
                 call_local_expert_count=call_local_expert_count,
+                verify=verify,
+                **_optional_kwargs(
+                    warmup=warmup, iters=iters, graph_replay_iters=graph_replay_iters
+                ),
             )
 
         elif cmd == "stress":
@@ -1033,6 +1073,7 @@ def _bench_dispatch_combine(
                 combine_block_num=combine_block_num,
                 combine_warp_per_block=combine_warp_per_block,
                 call_local_expert_count=call_local_expert_count,
+                **_optional_kwargs(warmup=warmup, capture_iters=iters),
             )
 
         elif cmd == "tuning":
@@ -1236,6 +1277,12 @@ def bench_dispatch_combine(
     input_shift=0.0,
     force_scale_active=False,
     report_scale_stats=False,
+    warmup=None,
+    iters=None,
+    verify=True,
+    routing="random",
+    kernel_type="IntraNode",
+    graph_replay_iters=None,
 ):
     if combine_data_type is None:
         combine_data_type = dtype
@@ -1268,6 +1315,12 @@ def bench_dispatch_combine(
             input_shift,
             force_scale_active,
             report_scale_stats,
+            warmup,
+            iters,
+            verify,
+            routing,
+            kernel_type,
+            graph_replay_iters,
         ),
         nprocs=world_size,
         join=True,
@@ -1468,6 +1521,66 @@ def bench_dispatch_combine(
             "p50/p90/p99/max) for the generated input."
         ),
     )
+    parser.add_argument(
+        "--warmup",
+        type=int,
+        default=None,
+        help=(
+            "Number of warmup iterations before timing. Applies to --cmd bench "
+            "(default: 1) and --cmd profile (default: 5);"
+        ),
+    )
+    parser.add_argument(
+        "--iters",
+        type=int,
+        default=None,
+        help=(
+            "Number of timed/captured iterations. Applies to --cmd bench "
+            "(default: 10) and --cmd profile (default: 3) "
+        ),
+    )
+    parser.add_argument(
+        "--graph-replay-iters",
+        type=int,
+        default=None,
+        help=(
+            "Number of times each captured CUDA graph is replayed per "
+            "timed --iters sample (default: 10). --cmd bench only"
+        ),
+    )
+    parser.add_argument(
+        "--routing",
+        type=str,
+        default="random",
+        choices=[
+            "random",
+            "round_robin",
+            "remote_round_robin",
+            "spread",
+            "all_to_one",
+        ],
+        help=(
+            "Token-to-expert routing pattern used to generate test data "
+            "(default: random)"
+        ),
+    )
+    parser.add_argument(
+        "--verify",
+        type=int,
+        default=1,
+        choices=[0, 1],
+        help=(
+            "When 1 (default), verify dispatch/combine correctness during "
+            "warmup. --cmd bench only"
+        ),
+    )
+    parser.add_argument(
+        "--kernel-type",
+        type=str,
+        default="IntraNode",
+        choices=sorted(KERNEL_TYPE_CHOICES),
+        help="Dispatch/combine kernel implementation to benchmark (default: IntraNode)",
+    )
     args = parser.parse_args()
 
     if args.num_experts_per_rank is None:
@@ -1510,7 +1623,9 @@ def bench_dispatch_combine(
         f"dispatch_block_num: {args.dispatch_block_num}, "
         f"dispatch_warp_per_block: {args.dispatch_warp_per_block}, "
         f"combine_block_num: {args.combine_block_num}, "
-        f"combine_warp_per_block: {args.combine_warp_per_block}"
+        f"combine_warp_per_block: {args.combine_warp_per_block}, "
+        f"kernel_type: {args.kernel_type}, "
+        f"graph_replay_iters: {args.graph_replay_iters}"
     )
     print("-" * 60)
     bench_dispatch_combine(
@@ -1537,4 +1652,10 @@ def bench_dispatch_combine(
         input_shift=args.input_shift,
         force_scale_active=bool(args.force_scale_active),
         report_scale_stats=bool(args.report_scale_stats),
+        warmup=args.warmup,
+        iters=args.iters,
+        verify=bool(args.verify),
+        routing=args.routing,
+        kernel_type=args.kernel_type,
+        graph_replay_iters=args.graph_replay_iters,
     )

tests/python/ops/dispatch_combine_test_utils.py

@@ -411,6 +411,24 @@ def gen_test_data(
                     ) * self.config.num_experts_per_token
                     for j in range(self.config.num_experts_per_token):
                         indices[i, j] = (base + j) % total_experts
+            elif routing == "remote_round_robin":
+                # Balanced, no-skew round-robin over every destination rank
+                # except the local one, then round-robins over that rank's
+                # own experts to pick which one receives the token.
+                assert (
+                    self.config.world_size > 1
+                ), "remote_round_robin routing requires world_size > 1"
+                indices = torch.empty(
+                    n, self.config.num_experts_per_token, dtype=torch.int64
+                )
+                ws = self.config.world_size
+                nepr = self.config.num_experts_per_rank
+                for i in range(n):
+                    for j in range(self.config.num_experts_per_token):
+                        rec = i * self.config.num_experts_per_token + j
+                        dst = (r + 1 + (rec % (ws - 1))) % ws
+                        local_e = (rec // (ws - 1)) % nepr
+                        indices[i, j] = dst * nepr + local_e
             elif routing == "spread":
                 # Sends exactly one expert to every rank (requires num_experts_per_token ==
                 # world_size). After per-rank deduplication each rank receives every source