diff --git a/python/mori/ops/dispatch_combine_v2/README.md b/python/mori/ops/dispatch_combine_v2/README.md index 962b90382..77fda10e3 100644 --- a/python/mori/ops/dispatch_combine_v2/README.md +++ b/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 diff --git a/python/mori/ops/dispatch_combine_v2/dispatch_combine_op.py b/python/mori/ops/dispatch_combine_v2/dispatch_combine_op.py index 8c9eb8817..48b693542 100644 --- a/python/mori/ops/dispatch_combine_v2/dispatch_combine_op.py +++ b/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 diff --git a/python/mori/ops/dispatch_combine_v2/hip_tuning_configs.py b/python/mori/ops/dispatch_combine_v2/hip_tuning_configs.py index 62a896742..e415a09ca 100644 --- a/python/mori/ops/dispatch_combine_v2/hip_tuning_configs.py +++ b/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) diff --git a/tests/python/ops/dispatch_combine_v2/bench_ep.py b/tests/python/ops/dispatch_combine_v2/bench_ep.py index a9d96981a..7c66b7e7b 100644 --- a/tests/python/ops/dispatch_combine_v2/bench_ep.py +++ b/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