perf(EPv2/gfx1250): retune EP4 topk-9 fp4 combine; un-shadow the topk-9 dispatch row

python/mori/ops/dispatch_combine_v2/README.md

@@ -51,7 +51,7 @@ lazily, only when selected, so the package imports without FlyDSL installed.
 | `flydsl_prims.py` | FlyDSL device primitives: system atomics / ordered stores / fences / volatile-spin waits |
 | `intranode_kernels.py` | FlyDSL kernel factories: `make_dispatch` (+scales/replay), `make_combine` (gather) / `make_combine_scatter` (`_nop2p`, bf16/f32/fp8/fp4), `make_convert_dispatch_output` / `make_convert_combine_input` (StdMoE), `make_local_expert_count` |
 | `tuning_configs.py` | **flydsl** kernel geometry: per-(world,hidden,topk) block/warp lookup |
-| `hip_tuning_configs.py` | **hip** kernel geometry, separate table (never borrows flydsl's); same `lookup` contract. Independent dispatch/combine tables, keyed by device, shape, topk and (dispatch only) dtype; an unswept shape gets a single-shot default |
+| `hip_tuning_configs.py` | **hip** kernel geometry, separate table (never borrows flydsl's); same `lookup` contract. Independent dispatch/combine tables, keyed by device, shape and topk, plus the dispatch dtype (`dispatch_dtype_str`: bf16/fp8/fp4) for dispatch and `quant_type` for combine; an unswept shape gets a single-shot default |
 | `internode_tuning_configs.py` | **hip** internode kernel geometry, a third table: token-count buckets keyed by device, shape, topk and dispatch dtype, carrying `(block_num, rdma_block_num, warp_num)` **per phase** — dispatch and combine are tuned to different values over one shared arena. The internode plans are compiled per geometry, so the backend walks the whole table at build time and `lookup` only picks a prebuilt bucket; `block_num` is clamped to the CU count |
 
 ## gfx125x dispatch: slot reservation

python/mori/ops/dispatch_combine_v2/dispatch_combine_op.py

@@ -378,7 +378,7 @@ def _resolve_geometry(self):
                     self.world_size,
                     self.hidden_dim,
                     self.num_experts_per_token,
-                    dtype=self.dtype_str,
+                    dtype=self.dispatch_dtype_str,
                     experts_per_rank=self.num_experts_per_rank,
                     quant_type=self.quant_type,
                 )
@@ -552,9 +552,22 @@ def combine_token_nbytes(self):
     def is_asymmetric_dtype(self):
         return self.dispatch_dtype != self.combine_dtype
 
+    @property
+    def dispatch_dtype_str(self):
+        """What dispatch transports, "fp4" / "fp8" / "bf16": hip_tuning_configs' key.
+
+        HIP tunes dispatch and combine apart, and its combine half is picked by
+        quant_type, so the combine dtype has no place in this key.
+        """
+        if self.is_fp4:
+            return "fp4"
+        if self.is_fp8:
+            return "fp8"
+        return "bf16"
+
     @property
     def dtype_str(self):
-        """Token/dispatch dtype key for tuning_configs.lookup (fp4/fp8/default)."""
+        """FlyDSL's tuning key (tuning_configs.lookup), which names both halves."""
         if self.is_fp4:
             # fp4 dispatch + non-fp4 combine (asymmetric) moves 2 B/elem on the
             # combine side, so it needs the bf16 combine geometry, not the

python/mori/ops/dispatch_combine_v2/hip_tuning_configs.py

@@ -36,9 +36,12 @@
 An unswept shape returns schedule=None and the single-shot default below. Add one by
 sweeping with ``bench_ep.py``. fp32 combine is untuned and takes the bf16 buckets.
 
-dtype keys are whatever ``EpDispatchCombineConfig.dtype_str`` produces, hence
-"fp4_disp_bf16_comb": hip rejects an fp4 combine outright, so an fp4 dispatch here is
-always paired with bf16 -- which is the configuration the sweep measured.
+Dispatch dtype keys name only what dispatch transports: "bf16", "fp8", "fp4"
+(``EpDispatchCombineConfig.dispatch_dtype_str``). The combine side has no dtype here --
+it always reduces bf16 rows, and whether those go out as mxfp4 is ``quant_type``'s
+table. "fp4_disp_bf16_comb" is FlyDSL's key for the same fp4 dispatch (its table does
+name both halves); ``lookup`` still reads it as "fp4" for callers that pass
+``dtype_str``.
 """
 
 from __future__ import annotations
@@ -97,11 +100,11 @@ def _hip_default() -> dict:
     "mi355x": {
         (8, 7168, 8, None): {
             None: ((None, 64, 8),),
-            "fp4_disp_bf16_comb": ((None, 128, 8),),
+            "fp4": ((None, 128, 8),),
         },
         (8, 7168, 6, None): {
             None: ((None, 64, 8),),
-            "fp4_disp_bf16_comb": ((None, 128, 8),),
+            "fp4": ((None, 128, 8),),
         },
     },
     # 4x gfx1250 at EP4, hidden 7168, 2026-08-11. topk moves the edges (it sets _tpi),
@@ -122,7 +125,14 @@ def _hip_default() -> dict:
         # 256-CU GPU at the ~3456-recv operating point. Block 256 saturates it.
         (4, 7168, 9, None): {
             None: ((512, 64, 8), (4096, 256, 8), (None, 256, 16)),
-            "fp4_disp_bf16_comb": ((512, 64, 8), (1024, 64, 16), (None, 128, 16)),
+            "fp4": (
+                (512, 64, 8),
+                (1024, 128, 8),
+                (1536, 192, 8),
+                (2048, 256, 8),
+                (8192, 192, 24),
+                (None, 192, 32),
+            ),
         },
         # topk 6 (384 experts at EP4). The edges move in: 64x8 stops paying at 512.
         #   ct     64x8   64x16  128x16 256x16      (bf16 / fp8 / fp4)
@@ -132,7 +142,7 @@ def _hip_default() -> dict:
         #   16384 551.7   518.5  478.5  470.7  | 423.8 303. 264.9 266.| 414.9 233.7 172.6 166.4
         (4, 7168, 6, None): {
             None: ((512, 64, 8), (2048, 64, 16), (None, 128, 16)),
-            "fp4_disp_bf16_comb": ((512, 64, 8), (1024, 64, 16), (None, 128, 16)),
+            "fp4": ((512, 64, 8), (1024, 64, 16), (None, 128, 16)),
         },
         # EP8 on two 4-GPU hosts, topk 6, fp4 dispatch, 2026-10-03; the grid stays at 64 blocks
         # (the rest of the CUs belong to the co-resident GEMM). What wins is one token a warp:
@@ -155,7 +165,7 @@ def _hip_default() -> dict:
         #   5120   149.4         161.0
         (8, 7168, 6, None): {
             None: ((2048, 64, 16), (None, 128, 16)),
-            "fp4_disp_bf16_comb": (
+            "fp4": (
                 (1024, 64, 16),
                 (1536, 64, 24),
                 (2048, 64, 32),
@@ -169,7 +179,7 @@ def _hip_default() -> dict:
         # -80.9..-83.3. 2048 is the only point 64x16 wins; the edge between them is unswept.
         (16, 7168, 6, None): {
             None: ((None, 64, 16),),
-            "fp4_disp_bf16_comb": ((4095, 64, 16), (None, 128, 16)),
+            "fp4": ((4095, 64, 16), (None, 128, 16)),
         },
     },
 }
@@ -188,11 +198,15 @@ def _hip_default() -> dict:
     },
     "gfx1250": {
         (4, 7168, 8, None): ((None, 64, 8),),
-        # topk 9 = 8 routed + 1 shared (ATOM's actual dispatch key). Needed so the
-        # dispatch schedule has a combine half -- lookup() returns the block-64
-        # single-shot default unless BOTH halves exist. Combine geometry mirrors
-        # topk 8 (unchanged); only the dispatch half is retuned to block 256.
-        (4, 7168, 9, None): ((None, 64, 8),),
+        # topk 9 = 8 routed + 1 shared. Also the combine half of the topk-9 dispatch
+        # schedule: lookup() returns the single-shot default unless both halves exist.
+        (4, 7168, 9, None): (
+            (64, 64, 4),
+            (128, 64, 8),
+            (256, 128, 8),
+            (1024, 256, 8),
+            (None, 192, 8),
+        ),
         (4, 7168, 6, None): ((None, 64, 8),),
         # The single-shot default, listed so the EP8 and EP16 dispatch schedules have a combine half.
         (8, 7168, 6, None): ((None, 64, 8),),
@@ -214,11 +228,12 @@ def _hip_default() -> dict:
                 (None, 64, 16),
             ),
             (4, 7168, 9, None): (
-                (256, 64, 8),
-                (2047, 64, 16),
-                (2048, 64, 24),
-                (16384, 128, 24),
-                (None, 64, 16),
+                (128, 64, 8),
+                (512, 192, 8),
+                (1024, 256, 8),
+                (1536, 192, 16),
+                (2048, 256, 16),
+                (None, 192, 24),
             ),
             # 64x24 runs the overlapped path with eight producers and twelve reducing waves:
             # 329.6 -> 314.8 us at 16384 tokens a rank. It tied with 64x16 at 4096 until its store
@@ -311,6 +326,10 @@ def pick(buckets, edge):
     return tuple((edge,) + pick(disp, edge) + pick(comb, edge) for edge in edges)
 
 
+# FlyDSL's dtype_str spelling, still accepted from callers that pass it (aiter's MegaMoE).
+_DISPATCH_DTYPE_ALIAS = {"fp4_disp_bf16_comb": "fp4"}
+
+
 def lookup(
     world_size,
     hidden_dim,
@@ -321,9 +340,11 @@ def lookup(
 ) -> dict:
     """HIP geometry for this device/shape/dtype, composed from HIP's own two tables.
 
-    An unswept shape gets the HIP single-shot default (schedule=None). A swept one
-    gets a per-token-count schedule built from the dispatch and combine tables
-    independently, so either half can be re-tuned without touching the other.
+    ``dtype`` is the dispatch wire ("bf16" / "fp8" / "fp4") and picks the dispatch
+    half; ``quant_type`` picks the combine half. An unswept shape gets the HIP
+    single-shot default (schedule=None). A swept one gets a per-token-count schedule
+    built from the dispatch and combine tables independently, so either half can be
+    re-tuned without touching the other.
     """
     base = _hip_default()
     dev = _device_key()
@@ -346,7 +367,7 @@ def lookup(
     if disp is None or comb is None:
         return base  # half a schedule is not a schedule
     # None is the "every dtype measured the same" key; an exact dtype overrides it.
-    disp = disp.get(dtype) or disp.get(None)
+    disp = disp.get(_DISPATCH_DTYPE_ALIAS.get(dtype, dtype)) or disp.get(None)
     if disp is None:
         return base
     base["schedule"] = _merge(disp, comb)

tests/python/ops/dispatch_combine_v2/bench_ep.py

@@ -119,10 +119,13 @@
 CHECK_REPEAT = int(os.environ.get("CHECK_REPEAT", 0))
 PUSH_QCHECK_DEV = int(os.environ.get("PUSH_QCHECK_DEV", 0))
 # ROUTE=rand (default): TOPK distinct experts per token, drawn at random.
-# ROUTE=ring: ATOM's --fake-eplb placement, where position p = token * TOPK + j lands
-# on rank p % world. With TOPK >= world every token reaches every rank, so the
-# largest SWEEP point fills each receive buffer to exactly its capacity -- the case
-# that overwrites anything parked at the end of the landing zone.
+# ROUTE=ring: position p = token * TOPK + j lands on rank p % world. With TOPK >= world
+# every token reaches every rank, so the largest SWEEP point fills each receive buffer
+# to exactly its capacity -- the case that overwrites anything parked at the end of the
+# landing zone.
+# ROUTE=grouped: group-limited routing (EXPERT_GROUPS=8, TOPK_GROUPS=4 by default), plus
+# SHARED_EXPERTS=1 trailing columns pinned to the token's own rank. At EP4 a token then
+# reaches 2-3 ranks, not every one, and the ranks receive unevenly.
 # ROUTE=noself: rand without the experts of the token's own rank, so no rank
 # receives from itself.
 ROUTE = os.environ.get("ROUTE", "rand")
@@ -213,8 +216,39 @@ def main():
             [torch.randperm(n_experts - EPR, generator=gr)[:TOPK] for _ in range(M)]
         )
         idx = (idx + (idx >= rank * EPR).long() * EPR).to(torch.int32).to(dev)
+    elif ROUTE == "grouped":
+        # A token picks TOPK_GROUPS of EXPERT_GROUPS groups on a ring, then the same
+        # number of experts in each. Its SHARED_EXPERTS columns stay on its own rank.
+        expert_groups = int(os.environ.get("EXPERT_GROUPS", 8))
+        topk_groups = int(os.environ.get("TOPK_GROUPS", 4))
+        shared_experts = int(os.environ.get("SHARED_EXPERTS", 1))
+        routed_topk = TOPK - shared_experts
+        routed_per_rank = EPR - shared_experts
+        experts_per_group = world * routed_per_rank // expert_groups
+        ranks_per_group = max(1, experts_per_group // routed_per_rank)
+        experts_per_group_rank = experts_per_group // ranks_per_group
+        picks_per_group = routed_topk // topk_groups
+        # Tokens interleave over the ranks: rank r holds tokens r, r + world, ...
+        token = rank + torch.arange(M) * world
+        group_slot = torch.arange(topk_groups).view(1, -1, 1)
+        pick = torch.arange(picks_per_group).view(1, 1, -1)
+        # Position on the group ring; its group is group_pos % expert_groups.
+        group_pos = token.view(-1, 1, 1) * topk_groups + group_slot
+        # How many times that group was picked before, so its experts rotate.
+        group_visit = (group_pos // expert_groups) * picks_per_group + pick
+        routed = (
+            (group_pos % expert_groups) * experts_per_group
+            + (group_visit % ranks_per_group) * experts_per_group_rank
+            + (group_visit // ranks_per_group) % experts_per_group_rank
+        ).reshape(M, routed_topk)
+        # Routed expert ids -> dispatch ids, whose ranks hold EPR slots each.
+        routed_ids = (routed // routed_per_rank) * EPR + routed % routed_per_rank
+        shared_ids = (
+            rank * EPR + routed_per_rank + torch.arange(shared_experts)
+        ).expand(M, shared_experts)
+        idx = torch.cat([routed_ids, shared_ids], 1).to(torch.int32).to(dev)
     else:
-        raise ValueError(f"ROUTE={ROUTE!r}: expected rand, ring or noself")
+        raise ValueError(f"ROUTE={ROUTE!r}: expected rand, ring, grouped or noself")
     # Per-token scale rows, transported when SCALE_DIM>0 (fp8/fp4 by default). Shaped
     # exactly as repro_epv2_topk9.py: sc_n_i32 int32 lanes viewed as bytes, trimmed to
     # SCALE_DIM so scale_type_size=1 * scale_dim holds. Deterministic (arange + rank