From f73152410957d191fe407b542c1f9e708466fd1a Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 24 Aug 2026 09:07:59 +0000 Subject: [PATCH 01/13] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Dispatch sends bf16 and every receiver quantizes each copy it gets, so a token routed to topk peers is quantized topk times on the same values. Quantizing on the sender does it once per LOCAL token and halves (fp8) or quarters (fp4) what crosses the fabric; only the WMMA scale preshuffle has to stay on the receiver, because its destination is the grouped row that rank assigns. The wire must match what the expert GEMM wants for its A operand -- a8w4 -> fp8, a4w4 -> fp4 -- so it is checked, not inferred: a mismatch is a row-width error that would read into the next token's bytes rather than a slow path. MegaMoEStage2Config.dispatch_wire: bf16 | fp8 | fp4 (mori backend only, which is the one with a channel for the scale row). combine_token_nbytes stays bf16 and separate, so the combine slot stride does not follow the wire down. The e8m0 row is padded to 128 B. Dword is all mori's validator asks for, but TdmWholeOrSplit128 only yields a body for the part of a run that starts 128 B aligned, and at the natural 224 B stride only every 4th token does. mori is told hidden_dim in ELEMENTS at its own element size: fp8 and fp4 both transport as byte8, so an fp4 wire halves the count itself. Receiver side, the route gather learns a prequantized producer: it loads the payload dwords and the e8m0 byte the quant pass would have computed, and the store pass -- the only copy of the preshuffled-scale addressing -- is shared verbatim. No software pipelining needed here: this tree's quant/store split already lets the loads cluster. test_mega_moe_gfx1250.py gains --mega_wire {auto,bf16,fp8,fp4}, allows scatter_fused for a4w4_mxfp4, and drives AITER_FORCE_A8W4 off the quant key -- otherwise -q a4w4_mxfp4 silently measured a8w4. Compile-checked only (no run): 8/8 route-gather variants, 6/6 mori dispatch plans (bf16/fp8/fp4 x 8/16 warps) emit .hsaco, and the four host paths produce the payload/scale widths the GEMM expects. --- aiter/aot/flydsl/grouped_moe.py | 5 + aiter/fused_moe.py | 4 + aiter/ops/flydsl/grouped_moe_gfx1250.py | 35 ++- .../kernels/mega_moe_gfx1250/mega_moe.py | 207 +++++++++++++++++- .../flydsl/kernels/mega_moe_gfx1250/types.py | 6 + .../kernels/moe_fused_route_quant_scatter.py | 87 +++++++- aiter/ops/flydsl/moe_kernels.py | 69 +++++- .../multigpu_tests/test_mega_moe_gfx1250.py | 58 ++++- 8 files changed, 453 insertions(+), 18 deletions(-) diff --git a/aiter/aot/flydsl/grouped_moe.py b/aiter/aot/flydsl/grouped_moe.py index f5c219915b..4e3e22c43f 100644 --- a/aiter/aot/flydsl/grouped_moe.py +++ b/aiter/aot/flydsl/grouped_moe.py @@ -209,6 +209,11 @@ def _route_ksplit(feat_dim, source_topk, out_e, out_m): 1, numel, ptr_arg(torch.empty(0, dtype=i32, device=dev)), + # src_scale: read only by the prequantized build, which this job + # does not emit; the pointer still has to be passed. Missing it + # would not fail loudly -- compile_one_config swallows the + # TypeError and the routeks kernels just stop being precompiled. + ptr_arg(torch.empty(0, dtype=u8, device=dev)), grid, stream=0, ) diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index e7bf02b088..a80e8b477a 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -847,6 +847,10 @@ def _fused_moe_impl( doweight_stage1=doweight_stage1, w1_scale=w1_scale, w2_scale=w2_scale, + # Forwarded, not dropped: with an fp8 hidden_states this is the + # e8m0 row of an activation the caller already quantized (fp8 EP + # dispatch), and the gfx1250 path needs it to skip its own quant. + a1_scale=a1_scale, expert_mask=expert_mask, hidden_pad=hidden_pad, intermediate_pad=intermediate_pad, diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index aee2fe0271..e26ef171bc 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -446,6 +446,7 @@ def _grouped_a8w4_tdm_moe( stage2_scatter: Stage2ScatterContext | None = None, situ_beta=1.0, situ_linear_beta=1.0, + a1_scale=None, ): import functools @@ -603,8 +604,37 @@ def _grouped_a8w4_tdm_moe( _quant_mode = "fp4" if _is_fp4 else "fp8" _a_is_fp4 = 1 if _is_fp4 else 0 + # Pre-quantized activation: an MX payload plus its e8m0 row is what a + # quantizing EP dispatch delivers, and it is also aiter's standing meaning + # for this pair. The fused pass then keeps its route gather and its scale + # preshuffle and drops only the quant -- the preshuffle cannot move to the + # sender anyway, because its destination is a function of the grouped row + # THIS rank assigns. + _prequantized = a1_scale is not None and hidden_states.dtype in ( + dtypes.fp8, + torch.uint8, + dtypes.fp4x2, + ) + if a1_scale is not None and not _prequantized: + # Loud rather than silently re-quantizing something already quantized. + assert hidden_states.dtype == dtype, ( + f"a1_scale given with hidden_states dtype {hidden_states.dtype}: " + "expected packed MX bytes (pre-quantized) or the model dtype (ignored)" + ) + # A payload row is model_dim bytes on an fp8 wire and model_dim//2 on an fp4 + # one. Checked rather than inferred: a wire that disagrees with the GEMM's + # data_format would otherwise read into the next token's bytes and produce + # plausible garbage. + _src_width = hidden_states.shape[-1] + if _prequantized: + _want_width = model_dim // 2 if _is_fp4 else model_dim + assert _src_width == _want_width, ( + f"prequantized {_quant_mode} payload should be {_want_width} B per " + f"row at model_dim {model_dim}, got {_src_width}" + ) + a1_payload, a1_scale = flydsl_moe_fused_quant_preshuffle( - hidden_states.reshape(1, token_num, model_dim), + hidden_states.reshape(1, token_num, _src_width), 1, contiguous_m, wmma_rep=wmma_rep, @@ -613,6 +643,7 @@ def _grouped_a8w4_tdm_moe( topids_to_rows=topids_to_rows, source_topk=topk, num_valid_routes=_ep_nvr, + prequantized_scale=a1_scale if _prequantized else None, ) # Fuse gemm1 activation + MX quantization + scale preshuffle into the @@ -900,6 +931,7 @@ def grouped_gemm_gfx1250_a8w4( situ_beta: float = 1.0, situ_linear_beta: float = 1.0, stage2_scatter: Stage2ScatterContext | None = None, + a1_scale: torch.Tensor | None = None, ): """Grouped a8w4/a4w4 MoE on the TDM batched GEMM (gfx1250). @@ -1169,6 +1201,7 @@ def _tdm_env(name): stage2_scatter=stage2_scatter, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, + a1_scale=a1_scale, **_tdm_kw, ) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index 8b02e53262..6493000c08 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -37,8 +37,29 @@ "dispOut": "disp_out", "outTok": "comb_inp", "xdb": "cross_device_barrier", + # Only laid out on a quantizing wire. plan_api binds a region the arena does + # not carry to offset 0, and the kernel's own `if constexpr` is what keeps a 0 + # from being read, so the bf16 arena needs no placeholder. + "outScales": "out_scales", } +# What dispatch puts on the wire, and what each mode costs per token at +# hidden 7168: bf16 14336 B, fp8 7168 + 256, fp4 3584 + 256. +_DISPATCH_WIRES = ("bf16", "fp8", "fp4") +# Payload bytes per feature. fp4 packs two features per byte, which is why the +# wire's element COUNT and the model's feature count part ways below. +_WIRE_PAYLOAD_BYTES = {"bf16": 2, "fp8": 1, "fp4": 0.5} +# What mori's Cfg is told the payload is. fp8 and fp4 both land on its byte8 +# transport, which only copies -- fp4x2 is declared for what it is rather than +# borrowed from fp8, so a reader of the plan sees the wire, not a stand-in. +_MORI_WIRE_DTYPE = { + "bf16": torch.bfloat16, + "fp8": dtypes.fp8, + "fp4": dtypes.fp4x2, +} +# What the sender quantizes to. +_QUANT_WIRE_DTYPE = {"fp8": dtypes.fp8, "fp4": dtypes.fp4x2} + def _align_up(value: int, alignment: int) -> int: return (value + alignment - 1) // alignment * alignment @@ -119,8 +140,38 @@ class MegaMoEStage2Config: dispatch_warp_num_per_block: int | None = None schedule: tuple | None = None dispatch_backend: str = "flydsl" + # What dispatch puts on the wire. fp8 halves the payload and fp4 quarters it, + # each sending a per-token e8m0 row along; the receiver then skips its own + # quant. Combine is unaffected -- it moves post-expert tokens, which are bf16 + # whatever the wire carried. + # + # The wire must MATCH what the expert GEMM wants for its A operand (a8w4 -> + # fp8, a4w4 -> fp4). It is not a free choice: the receiver hands the payload + # to the grouped GEMM as-is, so a mismatch is a width error, not a slow path. + dispatch_wire: str = "bf16" def __post_init__(self): + if self.dispatch_wire not in _DISPATCH_WIRES: + raise ValueError( + f"dispatch_wire must be one of {_DISPATCH_WIRES}, " + f"got {self.dispatch_wire!r}" + ) + if self.is_quant_wire and self.dispatch_backend != "mori": + # Only mori's kernel carries the scale row; this package's own + # dispatch has no channel for it. + raise ValueError( + f"dispatch_wire={self.dispatch_wire!r} requires " + f"dispatch_backend='mori' (got {self.dispatch_backend!r})" + ) + if self.is_quant_wire and self.hidden_dim % 32: + raise ValueError( + "one e8m0 scale covers 32 features, so a quantizing wire needs " + f"hidden_dim % 32 == 0, got {self.hidden_dim}" + ) + if self.dispatch_wire == "fp4" and self.hidden_dim % 2: + raise ValueError( + f"fp4 packs two features per byte, got hidden_dim={self.hidden_dim}" + ) if self.dispatch_backend not in _DISPATCH_BACKENDS: raise ValueError( f"dispatch_backend must be one of {_DISPATCH_BACKENDS}, " @@ -163,14 +214,78 @@ def __post_init__(self): def max_recv(self) -> int: return self.world_size * self.max_tokens_per_rank + @property + def is_quant_wire(self) -> bool: + """The wire carries an MX payload plus its e8m0 row, not bf16.""" + return self.dispatch_wire in ("fp8", "fp4") + + @property + def is_fp8_wire(self) -> bool: + return self.dispatch_wire == "fp8" + + @property + def is_fp4_wire(self) -> bool: + return self.dispatch_wire == "fp4" + + @property + def quant_mode(self) -> str | None: + """aiter's name for the MX format on the wire, or None on bf16.""" + return self.dispatch_wire if self.is_quant_wire else None + @property def token_nbytes(self) -> int: + """Payload bytes per dispatched token -- the DISPATCH direction only.""" + return int(self.hidden_dim * _WIRE_PAYLOAD_BYTES[self.dispatch_wire]) + + @property + def wire_elem_count(self) -> int: + """What mori's Cfg calls hidden_dim: ELEMENTS, at its own element size. + + fp8 and fp4 both transport as one byte per element, so an fp4 wire has to + halve the count itself -- mori sizes the token as hidden_dim * elem_size + and would otherwise move two bytes per packed byte. + """ + return self.token_nbytes if self.is_quant_wire else self.hidden_dim + + @property + def combine_token_nbytes(self) -> int: + """Combine always moves bf16 post-expert tokens, whatever the wire carried. + + Kept separate from token_nbytes on purpose: deriving the combine slot + stride from a now-wire-dependent token size silently halves it on fp8 + (quarters it on fp4) while the rows staged into comb_inp are still bf16. + """ return self.hidden_dim * 2 + @property + def scale_nbytes(self) -> int: + """Per-token e8m0 row, padded to 128 B. Same width for fp8 and fp4. + + Dword padding is all mori's validator demands, but it is not what makes + the transfer fast. mori's gfx1250 dispatch batches these rows through TDM, + and TdmWholeOrSplit128 only yields a `body` for the part of a run that + starts on a 128 B boundary. At the natural 224 B stride only every 4th + token starts aligned, so most of the run fell out to the scalar + global->global fallback -- measured as dispatch 87.8us (bf16) -> 157.7us + (fp8) at 512 tokens/rank. At a 128 B multiple every token is aligned. + + The padding is a bigger share of an fp4 wire (256 B against 3584 B of + payload, versus 7168) -- but it buys the same alignment, and the fp4 + payload row is 128 B-aligned for the same reason. + """ + if not self.is_quant_wire: + return 0 + return (self.hidden_dim // 32 + 127) // 128 * 128 + + @property + def scale_used_nbytes(self) -> int: + """The e8m0 bytes that carry information; the rest of the row is padding.""" + return self.hidden_dim // 32 if self.is_quant_wire else 0 + @property def combine_slot_stride_bytes(self) -> int: stride = 1 - while stride < self.token_nbytes: + while stride < self.combine_token_nbytes: stride <<= 1 return stride @@ -211,6 +326,7 @@ def __init__( situ_beta: torch.Tensor | None = None, situ_linear_beta: torch.Tensor | None = None, dispatch_backend: str | None = None, + dispatch_wire: str | None = None, ): """Everything here is fixed for the whole model; forward() takes the rest. @@ -314,6 +430,11 @@ def __init__( if dispatch_backend is not None else os.environ.get("MEGA_DISPATCH", "flydsl") ), + dispatch_wire=( + dispatch_wire + if dispatch_wire is not None + else os.environ.get("MEGA_WIRE", "bf16") + ), ), communicator, ) @@ -418,6 +539,14 @@ def forward( recv_x = recv_x[:bound] recv_weights = recv_weights[:bound] recv_ids = recv_ids[:bound] + if self._config.is_quant_wire: + assert a1_scale is None, ( + "a1_scale is produced by the quantizing wire itself; a " + "caller-supplied one would be silently discarded" + ) + a1_scale = self._recv_scales() + if recv_token_bound is not None: + a1_scale = a1_scale[: int(recv_token_bound)] extra = {} if self.activation == ActivationType.Situv2: extra["beta"] = self.situ_beta @@ -472,6 +601,11 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): ("out_wts", max_recv * config.topk * 4), ("disp_out", max_recv * config.token_nbytes), ("cross_device_barrier", config.world_size * 8), + *( + [("out_scales", max_recv * config.scale_nbytes)] + if config.scale_nbytes + else [] + ), ( "comb_inp", config.max_tokens_per_rank @@ -492,6 +626,9 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): config.world_size, dtype=torch.int32, device=device ) self._dispatch_barrier = torch.zeros(1, dtype=torch.int32, device=device) + # Set per dispatch on the fp8 wire; 0 (and unread) on bf16. + self._sent_scales = None + self._sent_scales_ptr = 0 self._total_recv = torch.zeros(1, dtype=torch.int32, device=device) self._cross_device_flag = torch.ones(1, dtype=torch.int64, device=device) self._combine_output = torch.zeros( @@ -592,13 +729,18 @@ def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: for spec in self._dispatch_specs: plan = EpDispatchPlan( world_size=config.world_size, - hidden_dim=config.hidden_dim, + # ELEMENTS at mori's element size, which is not the feature count + # on an fp4 wire: fp8 and fp4 both transport as byte8, so mori + # sizes a token as hidden_dim * 1 and fp4 has to halve the count + # itself (plan_api says as much: "the caller halves hiddenDim"). + hidden_dim=config.wire_elem_count, max_tok_per_rank=config.max_tokens_per_rank, num_expert_per_rank=config.experts_per_rank, num_expert_per_token=config.topk, max_recv=config.max_recv, - dtype=torch.bfloat16, + dtype=_MORI_WIRE_DTYPE[config.dispatch_wire], use_weights=True, + scale_bytes=config.scale_nbytes, block_num=spec[0], warp_per_block=spec[1], arena=self._arena, @@ -636,6 +778,12 @@ def launch( total_recv_token_num=addr_total_recv, grid_barrier=addr_disp_bar, num_tokens=inp_cur_tok, + # Read off self rather than through the variant's argument + # list: the list is shared with the FlyDSL dispatch, whose + # launcher is a traced @flyc.jit signature, and widening it + # would put a dead kernarg on a path that can never carry + # scales (fp8 requires dispatch_backend='mori'). + scales_buf=self._sent_scales_ptr, ) return launch @@ -656,10 +804,30 @@ def _select_dispatch(self, token_count: int) -> tuple[int, int]: return self._dispatch_specs[-1] def _recv_tokens(self) -> torch.Tensor: + config = self._config + if config.is_fp4_wire: + # uint8, not fp4x2: this is what the grouped GEMM's prequantized + # gather reads, and it addresses the row in BYTES. A packed dtype + # would make shape[-1] the feature count and hide the packing. + return _from_gpu_ptr( + self._arena.local_ptr("disp_out"), + (config.max_recv, config.token_nbytes), + torch.uint8, + ) return _from_gpu_ptr( self._arena.local_ptr("disp_out"), - (self._config.max_recv, self._config.hidden_dim), - torch.bfloat16, + (config.max_recv, config.hidden_dim), + dtypes.fp8 if config.is_fp8_wire else torch.bfloat16, + ) + + def _recv_scales(self) -> torch.Tensor | None: + """The forwarded e8m0 rows, or None on the bf16 wire.""" + if not self._config.is_quant_wire: + return None + return _from_gpu_ptr( + self._arena.local_ptr("out_scales"), + (self._config.max_recv, self._config.scale_nbytes), + torch.uint8, ) def _recv_weights(self) -> torch.Tensor: @@ -685,9 +853,36 @@ def _dispatch( token_count = hidden_states.shape[0] spec = self._select_dispatch(token_count) stream = fx.Stream(torch.cuda.current_stream()) + payload = hidden_states + if self._config.is_quant_wire: + # Quantize ONCE PER LOCAL TOKEN here, instead of once per received + # copy on the far side. Destination-independent, so the bytes are the + # same either way; the preshuffle cannot move with it, because its + # destination is the grouped row the receiver assigns. + from aiter.ops.quant import per_1x32_mx_quant_hip + + payload, scale_rows = per_1x32_mx_quant_hip( + hidden_states, + quant_dtype=_QUANT_WIRE_DTYPE[self._config.dispatch_wire], + scale_type=dtypes.fp8_e8m0, + shuffle=False, + ) + # Copied into a 128 B-strided buffer rather than sent as-is: see + # scale_nbytes. Allocated once at capacity and reused, so the copy is + # the only per-call cost and the pointer is stable under graph capture. + used = self._config.scale_used_nbytes + if self._sent_scales is None: + self._sent_scales = torch.zeros( + (self._config.max_tokens_per_rank, self._config.scale_nbytes), + dtype=torch.uint8, + device=hidden_states.device, + ) + self._sent_scales_ptr = self._sent_scales.data_ptr() + rows_u8 = scale_rows.view(torch.uint8).reshape(token_count, -1) + self._sent_scales[:token_count, :used] = rows_u8[:, :used] self._dispatch_variants[spec]( self._arena.handle, - hidden_states.data_ptr(), + payload.data_ptr(), topk_ids.data_ptr(), topk_weights.data_ptr(), self._token_destination_map.data_ptr(), diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/types.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/types.py index 799afaa111..4974ad860a 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/types.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/types.py @@ -14,6 +14,12 @@ torch.int32: (" None: experiment with multi-destination scattering without changing the quant math. Shared verbatim by both stage1 and stage2; only the preamble that computes ``c.dests`` differs. + + With ``c.prequantized`` the source row is already an MX payload plus a + separate e8m0 row (``c.src_scale_base`` / ``c.src_scale_bytes_per_row``, and + ``c.feat_bytes_per_row`` sized for the payload, not for bf16): the first pass + loads what it would otherwise have computed, and the store pass is unchanged. """ i32 = c.i32 f32 = c.f32 @@ -290,13 +295,56 @@ def _emit_quant_block_loop(c: SimpleNamespace) -> None: hidden_row_addr, num_records_bytes=c.feat_bytes_per_row ) feat_elem_base = arith.constant(0, type=i32) + + # Pre-quantized source: the sender already produced the MX payload and its + # e8m0 row, so this pass loads what the quant pass would have computed. The + # store pass below is then shared verbatim -- which is the point. The + # preshuffled scale destination has already moved once (to WMMA-contiguous), + # and a second copy of that arithmetic would fail by writing to the wrong + # offset, silently, rather than as a merge conflict. + prequantized = getattr(c, "prequantized", False) + src_scale_rsrc = None + if const_expr(prequantized): + c2_i32 = arith.constant(2, type=i32) + # Its own stride: a sender pads the e8m0 row (mori pads to 128 B so every + # token's TDM run starts aligned), so it is a build constant rather than + # feat_dim // 32. + src_scale_rsrc = buffer_ops.create_buffer_resource_from_addr( + c.src_scale_base + fx.Uint64(c.feat_row_i32) * c.src_scale_bytes_per_row, + num_records_bytes=c.src_scale_bytes_per_row, + ) + quant_results = [] for it in range_constexpr(c.block_iters): # MX block (along K) this lane works on this iteration. mx_block = (mx_group_base + arith.constant(it, type=i32)) * arith.constant( c.mx_blocks_per_wave_iter, type=i32 ) + c.block_in_wave - if const_expr(c.use_pk8): + if const_expr(prequantized): + # This lane's payload bytes sit at exactly the offset the store pass + # writes them to, so both sides share the expression and cannot drift. + # fp8: 8 B/lane = 2 dwords; fp4: 4 B/lane = 1 dword -- the same types + # the pk8 converts produce, so the store needs no special case. + byte_off = ( + mx_block * c.c_payload_bytes_per_block + + c.lane_in_block * c.c_payload_bytes_per_lane + ) + payload_val = buffer_ops.buffer_load( + hidden_rsrc, + byte_off >> c2_i32, + vec_width=c.payload_dwords_per_lane, + dtype=i32, + ) + # Every lane of an MX block loads the same e8m0 byte (one cache line) + # and only the lead lane stores it. Unconditional on purpose: a value + # defined inside an scf.if would not dominate the store pass below. + e8m0_byte = buffer_ops.buffer_load( + src_scale_rsrc, mx_block, vec_width=1, dtype=T.i8 + ) + # Widen so the store pass's trunci sees the same i32 it does on the + # quant path, where e8m0 comes out of emit_mx_e8m0_scale as i32. + e8m0_scale = arith.extui(i32, ArithValue(e8m0_byte)) + elif const_expr(c.use_pk8): # gfx1250 native pk8: 8 contiguous bf16 cols this lane. # col_base = mx_block*32 + lane_in_block*8. col_base = ( @@ -1289,6 +1337,8 @@ def build_moe_fused_quant_preshuffle_route_ksplit_module( source_topk: int = 0, remap_rows: bool = False, ksplit: bool = True, + prequantized: bool = False, + src_scale_bytes_per_row: int = 0, ): """Route-indexed grouped quant+preshuffle. @@ -1299,6 +1349,13 @@ def build_moe_fused_quant_preshuffle_route_ksplit_module( across ``grid.y = block_iters`` so each workgroup handles one K-group. When ``ksplit=False`` (large token counts where grid.x already saturates), ``grid.y = 1`` and each warp loops over all K-groups internally. + + ``prequantized`` says ``grouped_in`` is an MX payload the sender already + produced (fp8 or fp4 EP dispatch) and ``src_scale`` its row-major e8m0 rows, + ``src_scale_bytes_per_row`` apart. The kernel then keeps the route gather and + the scale preshuffle and drops only the quant -- the preshuffle cannot move + to the sender, because its destination is a function of the grouped row THIS + rank assigns, which no sender knows. """ L = _quant_layout(feat_dim, quant_mode, wmma_rep) if not L.use_pk8: @@ -1324,15 +1381,34 @@ def build_moe_fused_quant_preshuffle_route_ksplit_module( block_iters = L.block_iters amax_shuffle_dists = L.amax_shuffle_dists + if prequantized: + assert src_scale_bytes_per_row >= L.scale_bytes_per_row, ( + f"src_scale_bytes_per_row {src_scale_bytes_per_row} cannot hold " + f"{L.scale_bytes_per_row} e8m0 bytes for feat_dim {feat_dim}" + ) + # The payload row IS the source row here, so the loop's bounds and its + # per-lane offsets both come off the payload geometry. + src_bytes_per_row = payload_bytes_per_row if prequantized else feat_dim * 2 + payload_dwords_per_lane = payload_bytes_per_lane // 4 + if prequantized: + assert payload_bytes_per_lane % 4 == 0, ( + f"prequantized load is dword-wide; {quant_mode} gives " + f"{payload_bytes_per_lane} B/lane" + ) + source_tag = f"srctk{source_topk}" if source_topk > 0 else "srcrow" remap_tag = "_remap" if remap_rows else "" ksplit_tag = "" if ksplit else "_noKS" + # In the name because it changes what the kernel READS, not just how fast: + # two builds with the same feat_dim/quant_mode are not interchangeable. + prequant_tag = f"_pq{src_scale_bytes_per_row}" if prequantized else "" source_topk_is_pow2 = source_topk > 0 and (source_topk & (source_topk - 1)) == 0 source_topk_shift = source_topk.bit_length() - 1 if source_topk_is_pow2 else 0 module_name = ( f"moe_fused_quant_preshuffle_routeks_fd{feat_dim}_r{wmma_rep}" f"_{quant_mode}_{L.native_tag}_{source_tag}{remap_tag}{ksplit_tag}" + f"{prequant_tag}" ) @flyc.kernel(name=module_name, known_block_size=[BLOCK_THREADS, 1, 1]) @@ -1345,6 +1421,7 @@ def fused_kernel( route_max_m: Int32, # masked route stride, read iff remap_rows numel: Int32, num_valid_routes: fx.Pointer, # (1,) int32: routes >= this are dead-tail padding (EP dynamic token count); skip + src_scale: fx.Pointer, # (tokens, src_scale_bytes_per_row) e8m0, read iff prequantized ): """Write masked or contiguous ``(Mtile, K//128, wmma_rep, 16, 4)`` scales.""" i32 = T.i32 @@ -1469,8 +1546,12 @@ def fused_kernel( payload_base=payload_base, payload_bytes_per_row=payload_bytes_per_row, hidden_base=hidden_base, - feat_bytes_per_row=feat_dim * 2, + feat_bytes_per_row=src_bytes_per_row, feat_row_i32=feat_row_i32, + prequantized=prequantized, + payload_dwords_per_lane=payload_dwords_per_lane, + src_scale_base=fx.Int64(ptrtoint(src_scale)), + src_scale_bytes_per_row=src_scale_bytes_per_row, mx_blocks_per_wave_iter=mx_blocks_per_wave_iter, mx_blocks_per_row=mx_blocks_per_row, amax_shuffle_dists=amax_shuffle_dists, @@ -1518,6 +1599,7 @@ def launch_fused( route_max_m: fx.Int32, numel: fx.Int32, num_valid_routes: fx.Pointer, + src_scale: fx.Pointer, grid_route_blocks: fx.Int32, stream: fx.Stream = fx.Stream(None), # noqa: B008 ): @@ -1532,6 +1614,7 @@ def launch_fused( route_max_m, numel, num_valid_routes, + src_scale, ).launch( grid=(grid_x, grid_y, 1), block=(BLOCK_THREADS, 1, 1), diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 8fcacc8c39..17cdb68457 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -2714,6 +2714,8 @@ def flydsl_moe_fused_route_quant_scatter( source_topk=topk, ksplit=use_ksplit_s1, ) + _null_i32 = torch.empty(0, dtype=torch.int32, device=device) + assert _null_i32.data_ptr() == 0, "expected a null data_ptr" launch_routeks( ptr_arg(hidden_flat), ptr_arg(grouped_a1.view(-1)), @@ -2722,6 +2724,12 @@ def flydsl_moe_fused_route_quant_scatter( ptr_arg(counter), # dummy row_starts; unused because remap_rows=False 1, numel, + # Pre-existing omission, not fallout of the prequantized change: this + # branch never passed num_valid_routes. A 0-element tensor has a null + # data_ptr, which the kernel tests for before dereferencing. + ptr_arg(_null_i32), + # src_scale: read only by the prequantized build, which this is not. + ptr_arg(grouped_a1_scale.view(-1)), grid_blocks, stream=torch.cuda.current_stream(), ) @@ -2958,6 +2966,8 @@ def _get_compiled_fused_quant_preshuffle_route_ksplit( source_topk: int = 0, remap_rows: bool = False, ksplit: bool = True, + prequantized: bool = False, + src_scale_bytes_per_row: int = 0, ): from aiter.ops.flydsl.kernels.moe_fused_route_quant_scatter import ( build_moe_fused_quant_preshuffle_route_ksplit_module, @@ -2970,6 +2980,8 @@ def _get_compiled_fused_quant_preshuffle_route_ksplit( source_topk=source_topk, remap_rows=remap_rows, ksplit=ksplit, + prequantized=prequantized, + src_scale_bytes_per_row=src_scale_bytes_per_row, ) @@ -2987,9 +2999,12 @@ def flydsl_moe_fused_quant_preshuffle( route_max_m: int = 0, out_payload: torch.Tensor | None = None, # (E, max_m, Pb) uint8 out_scale: torch.Tensor | None = None, # (E, max_m//wmma_rep, Ws*wmma_rep) - num_valid_routes: ( - torch.Tensor | None - ) = None, # (1,) int32; route-branch only: skip routes >= this (EP dead-tail) + # (1,) int32; route-branch only: skip routes >= this (EP dead-tail) + num_valid_routes: torch.Tensor | None = None, + # (tokens, Ws) uint8 e8m0. When given, grouped_in IS the MX payload for + # ``quant_mode``: the sender already quantized, so the kernel only scatters + # + preshuffles. + prequantized_scale: torch.Tensor | None = None, ): """Fused grouped quant + e8m0 scale-preshuffle in one kernel pass. @@ -3000,11 +3015,44 @@ def flydsl_moe_fused_quant_preshuffle( f"flydsl_moe_fused_quant_preshuffle: quant_mode={quant_mode!r} " "unsupported (expected 'fp4' or 'fp8')." ) - assert ( - grouped_in.dtype == torch.bfloat16 - ), f"fused grouped quant+preshuffle requires bf16 input (got {grouped_in.dtype})" + # A quantizing EP dispatch (fp8 or fp4) already put the payload and its e8m0 + # row on the wire: nothing left to convert, only scatter + preshuffle. + prequantized = prequantized_scale is not None + if prequantized: + # torch dtypes, not aiter.dtypes: this module deliberately imports only + # torch and the tensor shim. + _packed = tuple( + d + for d in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.uint8, + getattr(torch, "float4_e2m1fn_x2", None), + ) + if d is not None + ) + assert grouped_in.dtype in _packed, ( + "prequantized payload must be packed MX bytes " f"(got {grouped_in.dtype})" + ) + assert ( + topids_to_rows is not None + ), "prequantized mode exists only on the route-indexed branch" + assert ( + prequantized_scale.dtype == torch.uint8 + and prequantized_scale.is_contiguous() + ), "prequantized scale must be a contiguous uint8 (tokens, Ws) tensor" + else: + assert grouped_in.dtype == torch.bfloat16, ( + "fused grouped quant+preshuffle requires bf16 input " + f"(got {grouped_in.dtype})" + ) device = grouped_in.device + # feat_dim is the FEATURE count, and a prequantized fp4 row carries two + # features per byte -- taking shape[-1] there would halve every derived + # geometry (Pb, Ws, the module name) without tripping a single assert. feat_dim = grouped_in.shape[-1] + if prequantized and quant_mode == "fp4": + feat_dim *= 2 rows_per_tile = wmma_rep * 16 assert ( max_m % rows_per_tile == 0 @@ -3057,6 +3105,10 @@ def flydsl_moe_fused_quant_preshuffle( source_topk=source_topk, remap_rows=remap_rows, ksplit=use_ksplit, + prequantized=prequantized, + src_scale_bytes_per_row=( + int(prequantized_scale.shape[-1]) if prequantized else 0 + ), ) # Dead-tail skip (EP dynamic token count): routes >= num_valid_routes are # padding rows of the dispatch buffer and are not gathered/quantized. When @@ -3077,6 +3129,11 @@ def flydsl_moe_fused_quant_preshuffle( route_max_m_arg, numel, ptr_arg(num_valid_routes_i32), + # Read only when prequantized; the quant path must still pass a valid + # pointer, so hand it the output scale, which the kernel never loads. + ptr_arg( + prequantized_scale.view(-1) if prequantized else out_scale.view(-1) + ), grid_blocks, stream=torch.cuda.current_stream(), ) diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index ec9d6fdc8e..1ed5baf15e 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -130,6 +130,38 @@ def resolve_spec(quant_key, transport): } +# The MegaMoE (scatter_fused) wire. Unrelated to `transport` above, which belongs +# to the mori-v1 dispatch the other combine modes use. +_MEGA_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} + + +def resolve_mega_wire(mega_wire, quant_key): + """What MegaMoE's dispatch puts on the wire: bf16 | fp8 | fp4. + + A quantizing wire is not a free choice -- the receiver hands the payload to + the grouped GEMM as its A operand, so it has to be the format that GEMM wants + (a8w4 -> fp8, a4w4 -> fp4), which is what ``auto`` resolves to. Picking the + other one is a width error, not a slow path, so it is rejected here rather + than deep inside the gather. + """ + if mega_wire == "auto": + return _MEGA_WIRE_FOR_QUANT.get(quant_key, "bf16") + if mega_wire == "bf16": + return "bf16" + want = _MEGA_WIRE_FOR_QUANT.get(quant_key) + if want is None: + raise ValueError( + f"--mega_wire={mega_wire} needs an MX quant key " + f"({'/'.join(_MEGA_WIRE_FOR_QUANT)}), got -q {quant_key}" + ) + if mega_wire != want: + raise ValueError( + f"-q {quant_key} wants a {want} A operand, so --mega_wire={mega_wire} " + "would hand the GEMM the wrong payload width" + ) + return mega_wire + + # Weight quantization + shuffle (device path) / dequant (reference) def weight_per_128x128_quant(weight, quant_dtype): E, dim1, dim2 = weight.shape @@ -522,9 +554,9 @@ def setup(self, x0): uid = self.dist_ctx.bcast_uid(uid) self.comm = Communicator.init(self.dist_ctx.world, r, uid) if self.combine_mode == "scatter_fused": - if self.spec["key"] != "a8w4_mxfp4": + if self.spec["key"] not in _MXFP4_KEYS: raise NotImplementedError( - "scatter_fused is available only for a8w4_mxfp4" + f"scatter_fused is available only for {'/'.join(_MXFP4_KEYS)}" ) from aiter.ops.flydsl.kernels.mega_moe_gfx1250 import MegaMoEGfx1250 @@ -542,6 +574,9 @@ def setup(self, x0): activation=self.spec["activation"], gate_mode=self.spec["gate_mode"].value, quant_type=self.spec["aiter_qtype"], + # Explicit, not left to $MEGA_WIRE: the harness owns this now, and + # a stale env would otherwise silently change what is measured. + dispatch_wire=self.spec.get("mega_wire", "bf16"), ) else: EpDispatchCombineConfig, EpDispatchCombineOp = _import_mori_v2() @@ -768,9 +803,15 @@ def _device_shared_ffn(tokens, sw1, sw2): # Driver def main(): args = _parse_args() + # The import-time setdefault above pins a8w4; a4w4 is the other half of the + # same switch (aiter/fused_moe.py reads it per call on gfx1250, defaulting to + # fp4x2 unless this is 1), so the quant key has to drive it or -q a4w4_mxfp4 + # silently measures a8w4. + os.environ["AITER_FORCE_A8W4"] = "0" if args.quant_type == "a4w4_mxfp4" else "1" dist_ctx = Dist() dev = torch.device("cuda", dist_ctx.local_rank) spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype) + spec["mega_wire"] = resolve_mega_wire(args.mega_wire, args.quant_type) if spec["is_mxfp4"] and get_gfx() not in ("gfx950", "gfx1250"): if dist_ctx.rank == 0: @@ -790,7 +831,8 @@ def main(): print( f"[cfg] world={dist_ctx.world} layers={n_layers} tokens/rank={ct} hidden={hdim} " f"inter={idim} E={E} topk={topk} EPR={E // dist_ctx.world} quant={args.quant_type} " - f"combine={args.combine} " + f"combine={args.combine} mega_wire={spec['mega_wire']} " + f"force_a8w4={os.environ['AITER_FORCE_A8W4']} " f"gate={spec['gate_mode'].name} shared_E={args.shared_experts} gfx={get_gfx()}", flush=True, ) @@ -960,6 +1002,16 @@ def _parse_args(): default="auto", help="dispatch transport (communication) dtype", ) + p.add_argument( + "--mega_wire", + type=str, + choices=["auto", "bf16", "fp8", "fp4"], + default=os.environ.get("MEGA_WIRE", "bf16"), + help="MegaMoE (scatter_fused) dispatch wire: bf16 sends activations and " + "the receiver quantizes each copy; fp8/fp4 quantize once on the sender " + "and forward the e8m0 row. 'auto' picks what the quant key's GEMM wants. " + "Needs dispatch_backend=mori (MEGA_DISPATCH=mori).", + ) p.add_argument( "--combine", type=str, From 0a61dad8216573bd2aca3fd67dfef3dcacf9a71e Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 24 Aug 2026 09:51:01 +0000 Subject: [PATCH 02/13] fix(gfx1250): let a4w4 reach the grouped GEMM, and MegaMoE reach a4w4 Three things kept an fp4 wire from ever being measurable, none of them in the wire itself. grouped_moe_gfx1250: mxfp4 weights arrive either as fp4x2 or as the uint8 view of the same bytes (ATOM's loader keeps them uint8; MegaMoE accepts both). Only the a8w4 arm said so, so a4w4-with-uint8-weights failed the eligibility test and fell through to the 2-stage mxfp4 kernels -- a silent detour to a different kernel family, which on this shape has no tuned config and aborts. The next statement already normalized both spellings for the CSV key, so the asymmetry was an oversight, not a rule. test harness: MegaMoE rejects anything but g1u1 interleave, while resolve_spec gave a4w4 the SEPARATED default, so -q a4w4_mxfp4 --combine scatter_fused could not start. Forcing INTERLEAVE alone then produced uncorrelated output, because shuffle_group still laid the weights out for the 2-stage family (e8m0_shuffle) rather than the grouped one (n32k4 + gguu->gugu rows) -- the MX keys differ only in ACTIVATION dtype, so under MegaMoE both take the grouped prep. run_matrix.sh: never check accuracy in a timed run. The fp32 reference is a per-expert torch loop; at 16384 tokens/rank it pins all four ranks for tens of minutes and reads as a hang. Correctness now runs once per wire at 256 tokens. --- aiter/ops/flydsl/grouped_moe_gfx1250.py | 20 ++++++++-------- .../multigpu_tests/test_mega_moe_gfx1250.py | 24 +++++++++++++++---- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index e26ef171bc..06b945a8d7 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -1020,19 +1020,19 @@ def _fmt(v): ): _grouped_dbg("unsupported activation") return None - is_grouped_a4w4 = q_dtype_a == dtypes.fp4x2 and q_dtype_w == dtypes.fp4x2 - is_grouped_a8w4 = q_dtype_a == dtypes.fp8 and ( - q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8 - ) + # mxfp4 weights reach here either as fp4x2 or as the uint8 VIEW of the same + # bytes -- ATOM's loader keeps them uint8, and MegaMoE accepts both. Both + # arms have to say so: requiring the packed dtype on the a4w4 arm alone sent + # a4w4-with-uint8-weights to the 2-stage fallback instead, which is a silent + # detour to a different kernel family, not an error. The very next statement + # already normalizes the two spellings for the CSV key. + w_is_mxfp4 = q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8 + is_grouped_a4w4 = q_dtype_a == dtypes.fp4x2 and w_is_mxfp4 + is_grouped_a8w4 = q_dtype_a == dtypes.fp8 and w_is_mxfp4 if not (is_grouped_a4w4 or is_grouped_a8w4): return None data_format = "fp4" if is_grouped_a4w4 else "a8w4" - # Normalize uint8-viewed fp4 weights back to fp4x2 for CSV key matching. - q_dtype_w_key = ( - dtypes.fp4x2 - if (q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8) - else q_dtype_w - ) + q_dtype_w_key = dtypes.fp4x2 if w_is_mxfp4 else q_dtype_w _grouped_dbg(f"eligible data_format={data_format}") if w1_scale is None or w2_scale is None: return None diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index 1ed5baf15e..d64f6adbe1 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -95,7 +95,7 @@ def _import_mori_v2(): # Config / quant-path spec -def resolve_spec(quant_key, transport): +def resolve_spec(quant_key, transport, combine_mode="gather"): """How to prepare weights / quantize activations / call fused_moe for a quant key, plus the dispatch transport dtype. transport: auto|bf16|fp8.""" is_mxfp4 = quant_key in _MXFP4_KEYS @@ -116,11 +116,26 @@ def resolve_spec(quant_key, transport): aiter_qtype = QuantType.per_1x32 gate_mode = GateMode.INTERLEAVE if quant_key == "a8w4_mxfp4" else GateMode.SEPARATED + if is_mxfp4 and combine_mode == "scatter_fused": + # MegaMoE rejects anything but g1u1 interleave outright (its gemm2-fused + # scatter is built on that layout), so the a4w4 key has to follow a8w4 + # here rather than keep the SEPARATED default it uses elsewhere. Weight + # prep and the fp32 reference both read gate_mode off this spec, so + # overriding it in one place keeps all three consistent. + gate_mode = GateMode.INTERLEAVE return { "key": quant_key, "aiter_qtype": aiter_qtype, "gate_mode": gate_mode, + # Which family of expert kernels the weights are laid out for. The MX + # keys differ only in the ACTIVATION dtype -- the weights are the same + # mxfp4 either way -- but a4w4 has historically meant "the 2-stage mxfp4 + # kernels" here (e8m0_shuffle, is_shuffled) while a8w4 meant "the grouped + # n32k4 ones". MegaMoE is grouped-only, so under it both keys take the + # grouped prep; getting this wrong is silent, and costs a full run to + # find (the output is uncorrelated, not merely imprecise). + "grouped_weights": is_mxfp4 and combine_mode == "scatter_fused", "activation": ActivationType.Silu, "is_mxfp4": is_mxfp4, "is_fp8": is_fp8, @@ -224,7 +239,7 @@ def shuffle_group(w1_qt, w1_s, w2_qt, w2_s, spec, n_experts): key = spec["key"] if key in ("No", "per_Token", "per_128x128"): return shuffle_weight(w1_qt), shuffle_weight(w2_qt), w1_s, w2_s - if key == "a8w4_mxfp4": + if key == "a8w4_mxfp4" or spec.get("grouped_weights"): if spec["gate_mode"] == GateMode.INTERLEAVE: w1_phys = _gguu_to_gugu_rows(w1_qt.view(torch.uint8)) w1_a = shuffle_weight(w1_phys, layout=(16, 16)) @@ -240,7 +255,8 @@ def shuffle_group(w1_qt, w1_s, w2_qt, w2_s, spec, n_experts): w2_a = shuffle_weight(w2_qt.view(torch.uint8), layout=(16, 16)) w2_ss = moe_shuffle_scale(w2_s.contiguous(), experts_cnt=n_experts) return w1_a, w2_a, w1_ss, w2_ss - # a4w4_mxfp4 + # a4w4_mxfp4 on the 2-stage kernels (a different B layout from the grouped + # branch above -- e8m0_shuffle, not the n32k4 fold). w1_a = shuffle_weight(w1_qt, layout=(16, 16)) w2_a = shuffle_weight(w2_qt, layout=(16, 16)) w1_ss = fp4_utils.e8m0_shuffle(w1_s) @@ -810,7 +826,7 @@ def main(): os.environ["AITER_FORCE_A8W4"] = "0" if args.quant_type == "a4w4_mxfp4" else "1" dist_ctx = Dist() dev = torch.device("cuda", dist_ctx.local_rank) - spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype) + spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype, args.combine) spec["mega_wire"] = resolve_mega_wire(args.mega_wire, args.quant_type) if spec["is_mxfp4"] and get_gfx() not in ("gfx950", "gfx1250"): From d265dc6ec96a7497da5d4feb8ce7c029af290a51 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 24 Aug 2026 10:54:43 +0000 Subject: [PATCH 03/13] perf(mega_moe): hand mori the quantizer's scale rows as they come out The wire's e8m0 rows were being copied into a 128 B-strided buffer before dispatch, because mori laid a row down at exactly the width it was given and the alignment is what makes the transfer fast. mori now derives that stride itself, so the copy -- one extra kernel per dispatch, over every token's scales -- goes away and the quant op's output goes straight onto the wire. scale_nbytes is now what we SEND (hidden/32, packed); scale_dst_nbytes is what ARRIVES (mori's stride), and it is asked of mori rather than recomputed, so a change to the alignment cannot leave the two out of step. --- aiter/fused_moe.py | 5 +- aiter/ops/flydsl/grouped_moe_gfx1250.py | 21 +-- .../kernels/mega_moe_gfx1250/mega_moe.py | 175 ++++++++---------- .../kernels/moe_fused_route_quant_scatter.py | 10 +- .../multigpu_tests/test_mega_moe_gfx1250.py | 12 +- 5 files changed, 95 insertions(+), 128 deletions(-) diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index a80e8b477a..b644fdd402 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -847,9 +847,8 @@ def _fused_moe_impl( doweight_stage1=doweight_stage1, w1_scale=w1_scale, w2_scale=w2_scale, - # Forwarded, not dropped: with an fp8 hidden_states this is the - # e8m0 row of an activation the caller already quantized (fp8 EP - # dispatch), and the gfx1250 path needs it to skip its own quant. + # A quantizing EP dispatch puts the caller's e8m0 row here; + # the gfx1250 path uses it to skip a quant it would redo. a1_scale=a1_scale, expert_mask=expert_mask, hidden_pad=hidden_pad, diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index 06b945a8d7..b4c4b7571f 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -606,10 +606,7 @@ def _grouped_a8w4_tdm_moe( # Pre-quantized activation: an MX payload plus its e8m0 row is what a # quantizing EP dispatch delivers, and it is also aiter's standing meaning - # for this pair. The fused pass then keeps its route gather and its scale - # preshuffle and drops only the quant -- the preshuffle cannot move to the - # sender anyway, because its destination is a function of the grouped row - # THIS rank assigns. + # for this pair. _prequantized = a1_scale is not None and hidden_states.dtype in ( dtypes.fp8, torch.uint8, @@ -621,10 +618,8 @@ def _grouped_a8w4_tdm_moe( f"a1_scale given with hidden_states dtype {hidden_states.dtype}: " "expected packed MX bytes (pre-quantized) or the model dtype (ignored)" ) - # A payload row is model_dim bytes on an fp8 wire and model_dim//2 on an fp4 - # one. Checked rather than inferred: a wire that disagrees with the GEMM's - # data_format would otherwise read into the next token's bytes and produce - # plausible garbage. + # Checked, not inferred: a wire disagreeing with the GEMM's data_format would + # read into the next token's bytes and produce plausible garbage. _src_width = hidden_states.shape[-1] if _prequantized: _want_width = model_dim // 2 if _is_fp4 else model_dim @@ -1020,12 +1015,10 @@ def _fmt(v): ): _grouped_dbg("unsupported activation") return None - # mxfp4 weights reach here either as fp4x2 or as the uint8 VIEW of the same - # bytes -- ATOM's loader keeps them uint8, and MegaMoE accepts both. Both - # arms have to say so: requiring the packed dtype on the a4w4 arm alone sent - # a4w4-with-uint8-weights to the 2-stage fallback instead, which is a silent - # detour to a different kernel family, not an error. The very next statement - # already normalizes the two spellings for the CSV key. + # mxfp4 weights arrive as fp4x2 or as the uint8 view of the same bytes -- + # ATOM's loader keeps them uint8, and MegaMoE accepts both. Requiring the + # packed dtype on the a4w4 arm alone silently routed a4w4-with-uint8-weights + # to the 2-stage fallback. w_is_mxfp4 = q_dtype_w == dtypes.fp4x2 or w1.dtype == torch.uint8 is_grouped_a4w4 = q_dtype_a == dtypes.fp4x2 and w_is_mxfp4 is_grouped_a8w4 = q_dtype_a == dtypes.fp8 and w_is_mxfp4 diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index 6493000c08..80fa0d623d 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -37,28 +37,32 @@ "dispOut": "disp_out", "outTok": "comb_inp", "xdb": "cross_device_barrier", - # Only laid out on a quantizing wire. plan_api binds a region the arena does - # not carry to offset 0, and the kernel's own `if constexpr` is what keeps a 0 - # from being read, so the bf16 arena needs no placeholder. + # Only laid out on a quantizing wire; plan_api binds a missing region to 0 and + # the kernel's `if constexpr` keeps that 0 from being read. "outScales": "out_scales", } -# What dispatch puts on the wire, and what each mode costs per token at -# hidden 7168: bf16 14336 B, fp8 7168 + 256, fp4 3584 + 256. -_DISPATCH_WIRES = ("bf16", "fp8", "fp4") -# Payload bytes per feature. fp4 packs two features per byte, which is why the -# wire's element COUNT and the model's feature count part ways below. -_WIRE_PAYLOAD_BYTES = {"bf16": 2, "fp8": 1, "fp4": 0.5} -# What mori's Cfg is told the payload is. fp8 and fp4 both land on its byte8 -# transport, which only copies -- fp4x2 is declared for what it is rather than -# borrowed from fp8, so a reader of the plan sees the wire, not a stand-in. -_MORI_WIRE_DTYPE = { - "bf16": torch.bfloat16, - "fp8": dtypes.fp8, - "fp4": dtypes.fp4x2, + +@dataclass(frozen=True) +class _Wire: + """One dispatch wire. Per token at hidden 7168: bf16 14336 B, fp8 7168 + 256, + fp4 3584 + 256. + """ + + payload_bytes: float # PER FEATURE; fp4 packs two features into a byte + mori_dtype: torch.dtype + quant_dtype: torch.dtype | None # None: nothing for the sender to quantize to + # fp4 is viewed as raw bytes: the gather addresses a row in BYTES, and a + # packed dtype would make shape[-1] read as a feature count. + recv_dtype: torch.dtype + + +_WIRES = { + "bf16": _Wire(2, torch.bfloat16, None, torch.bfloat16), + "fp8": _Wire(1, dtypes.fp8, dtypes.fp8, dtypes.fp8), + "fp4": _Wire(0.5, dtypes.fp4x2, dtypes.fp4x2, torch.uint8), } -# What the sender quantizes to. -_QUANT_WIRE_DTYPE = {"fp8": dtypes.fp8, "fp4": dtypes.fp4x2} +_DISPATCH_WIRES = tuple(_WIRES) def _align_up(value: int, alignment: int) -> int: @@ -168,10 +172,6 @@ def __post_init__(self): "one e8m0 scale covers 32 features, so a quantizing wire needs " f"hidden_dim % 32 == 0, got {self.hidden_dim}" ) - if self.dispatch_wire == "fp4" and self.hidden_dim % 2: - raise ValueError( - f"fp4 packs two features per byte, got hidden_dim={self.hidden_dim}" - ) if self.dispatch_backend not in _DISPATCH_BACKENDS: raise ValueError( f"dispatch_backend must be one of {_DISPATCH_BACKENDS}, " @@ -220,22 +220,12 @@ def is_quant_wire(self) -> bool: return self.dispatch_wire in ("fp8", "fp4") @property - def is_fp8_wire(self) -> bool: - return self.dispatch_wire == "fp8" + def wire(self) -> "_Wire": + return _WIRES[self.dispatch_wire] @property - def is_fp4_wire(self) -> bool: - return self.dispatch_wire == "fp4" - - @property - def quant_mode(self) -> str | None: - """aiter's name for the MX format on the wire, or None on bf16.""" - return self.dispatch_wire if self.is_quant_wire else None - - @property - def token_nbytes(self) -> int: - """Payload bytes per dispatched token -- the DISPATCH direction only.""" - return int(self.hidden_dim * _WIRE_PAYLOAD_BYTES[self.dispatch_wire]) + def dispatch_token_nbytes(self) -> int: + return int(self.hidden_dim * self.wire.payload_bytes) @property def wire_elem_count(self) -> int: @@ -245,42 +235,46 @@ def wire_elem_count(self) -> int: halve the count itself -- mori sizes the token as hidden_dim * elem_size and would otherwise move two bytes per packed byte. """ - return self.token_nbytes if self.is_quant_wire else self.hidden_dim + return self.dispatch_token_nbytes if self.is_quant_wire else self.hidden_dim @property def combine_token_nbytes(self) -> int: - """Combine always moves bf16 post-expert tokens, whatever the wire carried. - - Kept separate from token_nbytes on purpose: deriving the combine slot - stride from a now-wire-dependent token size silently halves it on fp8 - (quarters it on fp4) while the rows staged into comb_inp are still bf16. - """ + """Combine moves bf16 post-expert tokens, whatever the wire carried.""" return self.hidden_dim * 2 @property def scale_nbytes(self) -> int: - """Per-token e8m0 row, padded to 128 B. Same width for fp8 and fp4. - - Dword padding is all mori's validator demands, but it is not what makes - the transfer fast. mori's gfx1250 dispatch batches these rows through TDM, - and TdmWholeOrSplit128 only yields a `body` for the part of a run that - starts on a 128 B boundary. At the natural 224 B stride only every 4th - token starts aligned, so most of the run fell out to the scalar - global->global fallback -- measured as dispatch 87.8us (bf16) -> 157.7us - (fp8) at 512 tokens/rank. At a 128 B multiple every token is aligned. - - The padding is a bigger share of an fp4 wire (256 B against 3584 B of - payload, versus 7168) -- but it buys the same alignment, and the fp4 - payload row is 128 B-aligned for the same reason. + """Per-token e8m0 row as WE produce it: one byte per 32 features, packed. + + Handed to mori as-is. mori lays it down at its own, 128 B-aligned stride + (scale_dst_nbytes) because that is what keeps a TDM run's start aligned; + that padding is mori's business, and the quant op's output can go straight + onto the wire without a repack. + """ + return self.hidden_dim // 32 if self.is_quant_wire else 0 + + @property + def scale_dst_nbytes(self) -> int: + """The stride the rows ARRIVE at, which the receiving gather addresses by. + + Asked of mori rather than recomputed: it is the transport's layout + decision, and a local copy of the rule would drift the first time the + alignment changes. """ if not self.is_quant_wire: return 0 - return (self.hidden_dim // 32 + 127) // 128 * 128 + try: + from mori.ops.dispatch_combine_v2.hip_backend import scale_stride_bytes + except ImportError as e: + # Imported here, not at module scope: a bf16 wire needs none of this, + # so an older mori keeps working until someone asks for fp8/fp4. + raise RuntimeError( + f"dispatch_wire={self.dispatch_wire!r} needs a mori whose EP " + "dispatch carries a per-token scale row (ROCm/mori#593 or later); " + "the installed one has no scale_stride_bytes" + ) from e - @property - def scale_used_nbytes(self) -> int: - """The e8m0 bytes that carry information; the rest of the row is padding.""" - return self.hidden_dim // 32 if self.is_quant_wire else 0 + return scale_stride_bytes(self.scale_nbytes) @property def combine_slot_stride_bytes(self) -> int: @@ -599,11 +593,13 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): ("recv_to_src_token", max_recv * 4), ("out_idx", max_recv * config.topk * 4), ("out_wts", max_recv * config.topk * 4), - ("disp_out", max_recv * config.token_nbytes), + ("disp_out", max_recv * config.dispatch_token_nbytes), ("cross_device_barrier", config.world_size * 8), *( - [("out_scales", max_recv * config.scale_nbytes)] - if config.scale_nbytes + # Arrival stride, not the packed row: undersizing overruns + # the last slots. + [("out_scales", max_recv * config.scale_dst_nbytes)] + if config.scale_dst_nbytes else [] ), ( @@ -626,8 +622,8 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): config.world_size, dtype=torch.int32, device=device ) self._dispatch_barrier = torch.zeros(1, dtype=torch.int32, device=device) - # Set per dispatch on the fp8 wire; 0 (and unread) on bf16. - self._sent_scales = None + # Points at the quant op's own scale rows, set per dispatch on a quantizing + # wire; 0 (and unread) on bf16. self._sent_scales_ptr = 0 self._total_recv = torch.zeros(1, dtype=torch.int32, device=device) self._cross_device_flag = torch.ones(1, dtype=torch.int64, device=device) @@ -729,16 +725,14 @@ def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: for spec in self._dispatch_specs: plan = EpDispatchPlan( world_size=config.world_size, - # ELEMENTS at mori's element size, which is not the feature count - # on an fp4 wire: fp8 and fp4 both transport as byte8, so mori - # sizes a token as hidden_dim * 1 and fp4 has to halve the count - # itself (plan_api says as much: "the caller halves hiddenDim"). + # see wire_elem_count; mori's plan_api: "the caller halves + # hiddenDim" hidden_dim=config.wire_elem_count, max_tok_per_rank=config.max_tokens_per_rank, num_expert_per_rank=config.experts_per_rank, num_expert_per_token=config.topk, max_recv=config.max_recv, - dtype=_MORI_WIRE_DTYPE[config.dispatch_wire], + dtype=config.wire.mori_dtype, use_weights=True, scale_bytes=config.scale_nbytes, block_num=spec[0], @@ -805,28 +799,25 @@ def _select_dispatch(self, token_count: int) -> tuple[int, int]: def _recv_tokens(self) -> torch.Tensor: config = self._config - if config.is_fp4_wire: - # uint8, not fp4x2: this is what the grouped GEMM's prequantized - # gather reads, and it addresses the row in BYTES. A packed dtype - # would make shape[-1] the feature count and hide the packing. - return _from_gpu_ptr( - self._arena.local_ptr("disp_out"), - (config.max_recv, config.token_nbytes), - torch.uint8, - ) + # Width in whatever recv_dtype counts: features for bf16/fp8, bytes for + # fp4 -- see _Wire.recv_dtype. + width = config.dispatch_token_nbytes // config.wire.recv_dtype.itemsize return _from_gpu_ptr( self._arena.local_ptr("disp_out"), - (config.max_recv, config.hidden_dim), - dtypes.fp8 if config.is_fp8_wire else torch.bfloat16, + (config.max_recv, width), + config.wire.recv_dtype, ) def _recv_scales(self) -> torch.Tensor | None: """The forwarded e8m0 rows, or None on the bf16 wire.""" if not self._config.is_quant_wire: return None + # Full padded rows, not a trimmed view: this goes to the gather kernel as a + # base pointer plus a build-constant pitch, and that pitch is the arrival + # stride. The kernel reads only the meaningful bytes of each row. return _from_gpu_ptr( self._arena.local_ptr("out_scales"), - (self._config.max_recv, self._config.scale_nbytes), + (self._config.max_recv, self._config.scale_dst_nbytes), torch.uint8, ) @@ -863,23 +854,13 @@ def _dispatch( payload, scale_rows = per_1x32_mx_quant_hip( hidden_states, - quant_dtype=_QUANT_WIRE_DTYPE[self._config.dispatch_wire], + quant_dtype=self._config.wire.quant_dtype, scale_type=dtypes.fp8_e8m0, shuffle=False, ) - # Copied into a 128 B-strided buffer rather than sent as-is: see - # scale_nbytes. Allocated once at capacity and reused, so the copy is - # the only per-call cost and the pointer is stable under graph capture. - used = self._config.scale_used_nbytes - if self._sent_scales is None: - self._sent_scales = torch.zeros( - (self._config.max_tokens_per_rank, self._config.scale_nbytes), - dtype=torch.uint8, - device=hidden_states.device, - ) - self._sent_scales_ptr = self._sent_scales.data_ptr() - rows_u8 = scale_rows.view(torch.uint8).reshape(token_count, -1) - self._sent_scales[:token_count, :used] = rows_u8[:, :used] + # Straight onto the wire; mori restrides these packed rows while it + # stages them, so there is no repack here. + self._sent_scales_ptr = scale_rows.data_ptr() self._dispatch_variants[spec]( self._arena.handle, payload.data_ptr(), diff --git a/aiter/ops/flydsl/kernels/moe_fused_route_quant_scatter.py b/aiter/ops/flydsl/kernels/moe_fused_route_quant_scatter.py index 796ca0744b..e8f73f18fd 100644 --- a/aiter/ops/flydsl/kernels/moe_fused_route_quant_scatter.py +++ b/aiter/ops/flydsl/kernels/moe_fused_route_quant_scatter.py @@ -296,12 +296,10 @@ def _emit_quant_block_loop(c: SimpleNamespace) -> None: ) feat_elem_base = arith.constant(0, type=i32) - # Pre-quantized source: the sender already produced the MX payload and its - # e8m0 row, so this pass loads what the quant pass would have computed. The - # store pass below is then shared verbatim -- which is the point. The - # preshuffled scale destination has already moved once (to WMMA-contiguous), - # and a second copy of that arithmetic would fail by writing to the wrong - # offset, silently, rather than as a merge conflict. + # Pre-quantized source: load what the quant pass would have computed, so the + # store pass below stays shared. That destination arithmetic has already + # moved once (to WMMA-contiguous); a second copy of it would fail silently, + # by writing to the wrong offset. prequantized = getattr(c, "prequantized", False) src_scale_rsrc = None if const_expr(prequantized): diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index d64f6adbe1..5de82fdb78 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -117,11 +117,8 @@ def resolve_spec(quant_key, transport, combine_mode="gather"): gate_mode = GateMode.INTERLEAVE if quant_key == "a8w4_mxfp4" else GateMode.SEPARATED if is_mxfp4 and combine_mode == "scatter_fused": - # MegaMoE rejects anything but g1u1 interleave outright (its gemm2-fused - # scatter is built on that layout), so the a4w4 key has to follow a8w4 - # here rather than keep the SEPARATED default it uses elsewhere. Weight - # prep and the fp32 reference both read gate_mode off this spec, so - # overriding it in one place keeps all three consistent. + # MegaMoE's gemm2-fused scatter is g1u1-interleave only, so a4w4 cannot + # keep the SEPARATED default it uses elsewhere. gate_mode = GateMode.INTERLEAVE return { @@ -590,9 +587,8 @@ def setup(self, x0): activation=self.spec["activation"], gate_mode=self.spec["gate_mode"].value, quant_type=self.spec["aiter_qtype"], - # Explicit, not left to $MEGA_WIRE: the harness owns this now, and - # a stale env would otherwise silently change what is measured. - dispatch_wire=self.spec.get("mega_wire", "bf16"), + # Explicit so a stale $MEGA_WIRE cannot change what is measured. + dispatch_wire=self.spec["mega_wire"], ) else: EpDispatchCombineConfig, EpDispatchCombineOp = _import_mori_v2() From 36dac5979a8e92ca5cc5abd941c91a1270ee93cc Mon Sep 17 00:00:00 2001 From: yanboshao Date: Wed, 26 Aug 2026 11:21:59 +0000 Subject: [PATCH 04/13] support a4w4 test --- .../kernels/mega_moe_gfx1250/mega_moe.py | 10 +- .../multigpu_tests/test_mega_moe_gfx1250.py | 384 ++++++++---------- 2 files changed, 176 insertions(+), 218 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index 80fa0d623d..9a8e63fee7 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -143,7 +143,7 @@ class MegaMoEStage2Config: dispatch_block_num: int | None = None dispatch_warp_num_per_block: int | None = None schedule: tuple | None = None - dispatch_backend: str = "flydsl" + dispatch_backend: str = "mori" # What dispatch puts on the wire. fp8 halves the payload and fp4 quarters it, # each sending a per-token e8m0 row along; the receiver then skips its own # quant. Combine is unaffected -- it moves post-expert tokens, which are bf16 @@ -422,7 +422,7 @@ def __init__( dispatch_backend=( dispatch_backend if dispatch_backend is not None - else os.environ.get("MEGA_DISPATCH", "flydsl") + else os.environ.get("MEGA_DISPATCH", "mori") ), dispatch_wire=( dispatch_wire @@ -721,6 +721,10 @@ def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: "built by mori's CMake and is not shipped by every install" ) from error + # Passed only on a quantizing wire, matching scale_dst_nbytes: mori grew + # scale_bytes in #593 and rejects UNKNOWN kwargs outright, so sending the + # bf16 wire's harmless 0 would make an older mori refuse the whole plan. + scale_kw = {"scale_bytes": config.scale_nbytes} if config.is_quant_wire else {} plans = {} for spec in self._dispatch_specs: plan = EpDispatchPlan( @@ -734,7 +738,7 @@ def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: max_recv=config.max_recv, dtype=config.wire.mori_dtype, use_weights=True, - scale_bytes=config.scale_nbytes, + **scale_kw, block_num=spec[0], warp_per_block=spec[1], arena=self._arena, diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index 5de82fdb78..442995678a 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -2,8 +2,8 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. """Multi-layer EP MoE end-to-end perf + accuracy on the mori v2 cco/FlyDSL op-layer. -N (default 61, DeepSeek-V4-Pro) MoE layers are chained. Gather/scatter use Mori -v2 dispatch -> AITER fused_moe -> Mori v2 combine. Scatter-fused calls only +N (default 61, DeepSeek-V4-Pro) MoE layers are chained. The ``base`` mode uses +Mori v2 dispatch -> AITER fused_moe -> Mori v2 combine. ``fused`` calls only ``MegaMoEGfx1250``, which owns AITER's dispatch -> fused_moe -> fused-combine pipeline. The combined output plus residual feeds the next layer. @@ -20,14 +20,15 @@ Launcher: torchrun (one process per rank / GPU), mirroring test_moe_layer_ep.py. -Launch (4x gfx1250, must build CK-free on gfx1250 -> ENABLE_CK=0): +Launch (4x gfx1250; every env knob below is already the script's default): cd # avoid the /app/triton namespace shadow - ENABLE_CK=0 AITER_FORCE_A8W4=1 AITER_USE_GROUPED_GEMM=1 AITER_BF16_FP8_MOE_BOUND=0 \ torchrun --standalone --nproc_per_node=4 test_mega_moe_gfx1250.py \ - -q a8w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 + -q a4w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 --combine fused \ + --dispatch_commu_dtype fp4 # Set MORI_CCO_BC to a prebuilt libmori_cco_device.bc to skip CCO JIT. -Env / CLI: --layers --logits_tol --acc_verify --dispatch_commu_dtype -tpr -hd -id -e -k --shared_E -q +Env / CLI: --layers --logits_tol --acc_verify --dispatch_commu_dtype --combine + -tpr -hd -id -e -k --shared_E -q """ import argparse @@ -37,15 +38,12 @@ import torch.distributed as dist import torch.profiler as tprof -import aiter from aiter import ( ActivationType, QuantType, dtypes, get_gfx, - get_hip_quant, get_torch_quant, - pertoken_quant, ) from aiter.fused_moe import fused_moe from aiter.ops.flydsl.moe_common import GateMode @@ -57,19 +55,25 @@ except Exception: # noqa: BLE001 # pragma: no cover get_trace_perf = None -# a8w4 (fp8 activation + mxfp4 weight) grouped kernel knobs. Force the real -# fp8/mxfp4 grouped path regardless of token count (mirrors test_moe_ep.py). +# gfx1250 grouped mxfp4 kernel knobs. Force the real grouped path regardless of +# token count (mirrors test_moe_ep.py). AITER_FORCE_A8W4 is deliberately absent: +# it selects the grouped kernel's ACTIVATION dtype (0 -> fp4, 1 -> fp8) and main() +# derives it from -q, so a stale value in the environment cannot make -q a4w4 +# silently measure a8w4. os.environ.setdefault("ENABLE_CK", "0") -os.environ.setdefault("AITER_FORCE_A8W4", "1") os.environ.setdefault("AITER_USE_GROUPED_GEMM", "1") os.environ.setdefault("AITER_BF16_FP8_MOE_BOUND", "0") +# The `base` path's mori EpDispatchCombineOp and the dispatch inside +# MegaMoEGfx1250 both go through mori's HIP/JIT kernels, so the two modes' +# kernel tables differ only in the combine. MegaMoE's own default is already +# mori (see MegaMoEStage2Config.dispatch_backend). +os.environ.setdefault("MORI_V2_KERNEL_BACKEND", "hip") os.environ.setdefault("FLYDSL_GPU_ARCH", get_gfx()) -_FP8_DTYPE = dtypes.fp8 -QUANT_KEYS = ["No", "per_Token", "per_128x128", "a8w4_mxfp4", "a4w4_mxfp4"] -_MXFP4_KEYS = ("a8w4_mxfp4", "a4w4_mxfp4") -_FP8_KEYS = ("per_Token", "per_128x128") +# Both keys carry mxfp4 WEIGHTS and differ only in the activation dtype the +# grouped GEMM wants for its A operand: a8w4 -> fp8, a4w4 -> fp4. +QUANT_KEYS = ["a8w4_mxfp4", "a4w4_mxfp4"] def _import_mori_comm(): @@ -82,7 +86,7 @@ def _import_mori_comm(): def _import_mori_v2(): """Import Mori's non-fused dispatch/combine v2 path. - Deferred: scatter_fused runs entirely on aiter's own mega_moe kernels and + Deferred: the fused mode runs entirely on aiter's own mega_moe kernels and only needs the communicator, so importing this eagerly would make the fused path fail whenever Mori's copy lags the installed flydsl API. """ @@ -95,96 +99,56 @@ def _import_mori_v2(): # Config / quant-path spec -def resolve_spec(quant_key, transport, combine_mode="gather"): - """How to prepare weights / quantize activations / call fused_moe for a quant - key, plus the dispatch transport dtype. transport: auto|bf16|fp8.""" - is_mxfp4 = quant_key in _MXFP4_KEYS - is_fp8 = quant_key in _FP8_KEYS - - if transport == "auto": - transport = "fp8" if is_fp8 else "bf16" - if transport == "fp8" and not is_fp8: - transport = "bf16" - - if quant_key == "No": - aiter_qtype = QuantType.No - elif quant_key == "per_Token": - aiter_qtype = QuantType.per_Token - elif quant_key == "per_128x128": - aiter_qtype = QuantType.per_128x128 - else: # a8w4_mxfp4 / a4w4_mxfp4 - aiter_qtype = QuantType.per_1x32 - - gate_mode = GateMode.INTERLEAVE if quant_key == "a8w4_mxfp4" else GateMode.SEPARATED - if is_mxfp4 and combine_mode == "scatter_fused": - # MegaMoE's gemm2-fused scatter is g1u1-interleave only, so a4w4 cannot - # keep the SEPARATED default it uses elsewhere. - gate_mode = GateMode.INTERLEAVE - +def resolve_spec(quant_key, wire): + """How to prepare weights / call fused_moe for a quant key, plus the resolved + dispatch wire. Both quant keys use the grouped gfx1250 kernels.""" return { "key": quant_key, - "aiter_qtype": aiter_qtype, - "gate_mode": gate_mode, - # Which family of expert kernels the weights are laid out for. The MX - # keys differ only in the ACTIVATION dtype -- the weights are the same - # mxfp4 either way -- but a4w4 has historically meant "the 2-stage mxfp4 - # kernels" here (e8m0_shuffle, is_shuffled) while a8w4 meant "the grouped - # n32k4 ones". MegaMoE is grouped-only, so under it both keys take the - # grouped prep; getting this wrong is silent, and costs a full run to - # find (the output is uncorrelated, not merely imprecise). - "grouped_weights": is_mxfp4 and combine_mode == "scatter_fused", + "aiter_qtype": QuantType.per_1x32, + # The gfx1250 grouped MoE GEMM reads GUGU (gate/up row-interleaved) w1 + # only, so both mxfp4 keys have to ask for INTERLEAVE; a SEPARATED layout + # silently falls through to the generic 2-stage MoE. + "gate_mode": GateMode.INTERLEAVE, "activation": ActivationType.Silu, - "is_mxfp4": is_mxfp4, - "is_fp8": is_fp8, - "transport": transport, - "prequant": transport == "fp8", - "fp8_dtype": _FP8_DTYPE, + "wire": resolve_wire(wire, quant_key), } -# The MegaMoE (scatter_fused) wire. Unrelated to `transport` above, which belongs -# to the mori-v1 dispatch the other combine modes use. -_MEGA_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} +# What the dispatch puts on the wire. A quantizing wire moves the activation +# quant to the SENDER: the payload shrinks (fp8 halves it, fp4 quarters it), a +# per-token e8m0 row rides along, and the receiver skips its own quant. Combine +# is unaffected -- it carries post-expert tokens, which are bf16 regardless. +_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} -def resolve_mega_wire(mega_wire, quant_key): - """What MegaMoE's dispatch puts on the wire: bf16 | fp8 | fp4. +def resolve_wire(wire, quant_key): + """Validate the dispatch wire against the quant key: bf16 | fp8 | fp4. A quantizing wire is not a free choice -- the receiver hands the payload to - the grouped GEMM as its A operand, so it has to be the format that GEMM wants - (a8w4 -> fp8, a4w4 -> fp4), which is what ``auto`` resolves to. Picking the - other one is a width error, not a slow path, so it is rejected here rather - than deep inside the gather. + the grouped GEMM as its A operand, so it has to be the width that GEMM wants + (a8w4 -> fp8, a4w4 -> fp4). Picking the other one is a width error, not a + slow path, so it is rejected here rather than deep inside the gather. This + is why the knob is not simply a bool: naming the format keeps the intended + pairing visible at the call site. + + `None` means "unset", and resolves to the quant key's own width so that the + quantizing wire is what you get by default on either key. """ - if mega_wire == "auto": - return _MEGA_WIRE_FOR_QUANT.get(quant_key, "bf16") - if mega_wire == "bf16": + if wire is None: + return _WIRE_FOR_QUANT[quant_key] + if wire == "bf16": return "bf16" - want = _MEGA_WIRE_FOR_QUANT.get(quant_key) - if want is None: - raise ValueError( - f"--mega_wire={mega_wire} needs an MX quant key " - f"({'/'.join(_MEGA_WIRE_FOR_QUANT)}), got -q {quant_key}" - ) - if mega_wire != want: + want = _WIRE_FOR_QUANT[quant_key] + if wire != want: raise ValueError( - f"-q {quant_key} wants a {want} A operand, so --mega_wire={mega_wire} " - "would hand the GEMM the wrong payload width" + f"-q {quant_key} wants a {want} A operand, so " + f"--dispatch_commu_dtype={wire} would hand the GEMM the wrong " + "payload width" ) - return mega_wire + return wire # Weight quantization + shuffle (device path) / dequant (reference) -def weight_per_128x128_quant(weight, quant_dtype): - E, dim1, dim2 = weight.shape - wb = weight.view(E, dim1 // 128, 128, dim2 // 128, 128) - wb = wb.permute(0, 1, 3, 2, 4).contiguous().view(E, -1, 128 * 128) - w_qt, w_s = aiter.pertoken_quant(wb, quant_dtype=quant_dtype) - w_qt = w_qt.view(E, dim1 // 128, dim2 // 128, 128, 128) - w_qt = w_qt.permute(0, 1, 3, 2, 4).contiguous().view(E, dim1, dim2) - return w_qt, w_s.view(E, dim1 // 128, dim2 // 128) - - def _mxfp4_quant(w): """per_1x32 mxfp4 quant: packed fp4x2 weight [E, d1, d2//2] + e8m0 scale.""" tq = get_torch_quant(QuantType.per_1x32) @@ -212,64 +176,34 @@ def _gguu_to_gugu_rows(t): def raw_quant_weights(w1, w2, spec): """Quantize (unshuffled) a group of routed-expert weights.""" - key = spec["key"] - if key == "No": - tq = get_torch_quant(QuantType.No) - w1_qt, _ = tq(w1, quant_dtype=None) - w2_qt, _ = tq(w2, quant_dtype=None) - return w1_qt.view(w1.shape), None, w2_qt.view(w2.shape), None - if key == "per_Token": - w1_qt, w1_s = pertoken_quant(w1, quant_dtype=_FP8_DTYPE) - w2_qt, w2_s = pertoken_quant(w2, quant_dtype=_FP8_DTYPE) - return w1_qt, w1_s, w2_qt, w2_s - if key == "per_128x128": - w1_qt, w1_s = weight_per_128x128_quant(w1, quant_dtype=_FP8_DTYPE) - w2_qt, w2_s = weight_per_128x128_quant(w2, quant_dtype=_FP8_DTYPE) - return w1_qt, w1_s, w2_qt, w2_s w1_qt, w1_s = _mxfp4_quant(w1) w2_qt, w2_s = _mxfp4_quant(w2) return w1_qt, w1_s, w2_qt, w2_s def shuffle_group(w1_qt, w1_s, w2_qt, w2_s, spec, n_experts): - """Layout-shuffle a group of `n_experts` quantized experts for the kernel.""" - key = spec["key"] - if key in ("No", "per_Token", "per_128x128"): - return shuffle_weight(w1_qt), shuffle_weight(w2_qt), w1_s, w2_s - if key == "a8w4_mxfp4" or spec.get("grouped_weights"): - if spec["gate_mode"] == GateMode.INTERLEAVE: - w1_phys = _gguu_to_gugu_rows(w1_qt.view(torch.uint8)) - w1_a = shuffle_weight(w1_phys, layout=(16, 16)) - w1_ss = moe_shuffle_scale( - w1_s.contiguous(), - experts_cnt=n_experts, - is_guinterleave=True, - gate_up=True, - ) - else: - w1_a = shuffle_weight(w1_qt.view(torch.uint8), layout=(16, 16)) - w1_ss = moe_shuffle_scale(w1_s.contiguous(), experts_cnt=n_experts) - w2_a = shuffle_weight(w2_qt.view(torch.uint8), layout=(16, 16)) - w2_ss = moe_shuffle_scale(w2_s.contiguous(), experts_cnt=n_experts) - return w1_a, w2_a, w1_ss, w2_ss - # a4w4_mxfp4 on the 2-stage kernels (a different B layout from the grouped - # branch above -- e8m0_shuffle, not the n32k4 fold). - w1_a = shuffle_weight(w1_qt, layout=(16, 16)) - w2_a = shuffle_weight(w2_qt, layout=(16, 16)) - w1_ss = fp4_utils.e8m0_shuffle(w1_s) - w2_ss = fp4_utils.e8m0_shuffle(w2_s) - w1_a.is_shuffled = True - w2_a.is_shuffled = True - return w1_a, w2_a, w1_ss, w2_ss - + """Layout-shuffle a group of `n_experts` quantized experts for the kernel. -def quant_tokens_fp8(tokens, spec): - """Per-token / per-block fp8 quant of the activations (fp8 pre-quant transport).""" - qt = spec["aiter_qtype"] - quant_func = get_hip_quant( - qt if qt != QuantType.per_128x128 else QuantType.per_1x128 + Both mxfp4 keys share the grouped gfx1250 layout (GUGU-interleaved w1 + the + n32k4 e8m0 B-scale). The only difference is the weight DTYPE handed to + ``fused_moe``: it keys off ``w1.dtype`` to pick the activation dtype, so + uint8 selects the fp8-activation (a8w4) kernel and fp4x2 the fp4-activation + (a4w4) one. See ``grouped_moe_gfx1250._grouped_a8w4_tdm_moe``. + """ + w1_phys = _gguu_to_gugu_rows(w1_qt.view(torch.uint8)) + w1_a = shuffle_weight(w1_phys, layout=(16, 16)) + w2_a = shuffle_weight(w2_qt.view(torch.uint8), layout=(16, 16)) + w1_ss = moe_shuffle_scale( + w1_s.contiguous(), + experts_cnt=n_experts, + is_guinterleave=True, + gate_up=True, ) - return quant_func(tokens, quant_dtype=spec["fp8_dtype"]) + w2_ss = moe_shuffle_scale(w2_s.contiguous(), experts_cnt=n_experts) + if spec["key"] == "a4w4_mxfp4": + w1_a = w1_a.view(dtypes.fp4x2) + w2_a = w2_a.view(dtypes.fp4x2) + return w1_a, w2_a, w1_ss, w2_ss def moe_forward( @@ -282,7 +216,6 @@ def moe_forward( topk_ids, expert_mask, spec, - a1_scale=None, num_local_tokens=None, ): """Single fused_moe call (device path). ``num_local_tokens`` (device int32 @@ -293,35 +226,20 @@ def moe_forward( num_local_tokens = torch.tensor( [hidden.shape[0]], dtype=dtypes.i32, device=hidden.device ) - if spec["is_mxfp4"]: - return fused_moe( - hidden, - w1_a, - w2_a, - topk_weights, - topk_ids, - expert_mask=expert_mask, - activation=spec["activation"], - gate_mode=spec["gate_mode"].value, - quant_type=spec["aiter_qtype"], - w1_scale=w1_s, - w2_scale=w2_s, - dtype=dtypes.bf16, - num_local_tokens=num_local_tokens, - ) return fused_moe( hidden, w1_a, w2_a, topk_weights, topk_ids, - expert_mask, - num_local_tokens=num_local_tokens, + expert_mask=expert_mask, + activation=spec["activation"], + gate_mode=spec["gate_mode"].value, + quant_type=spec["aiter_qtype"], w1_scale=w1_s, w2_scale=w2_s, - quant_type=spec["aiter_qtype"], - a1_scale=a1_scale, dtype=dtypes.bf16, + num_local_tokens=num_local_tokens, ) @@ -379,8 +297,9 @@ def make_routings(n_layers, ct, E, topk, dev, seed): def _rmsnorm(x, eps=1e-6): """RMSNorm (no learnable gain) on the last dim. Applied to each layer's MoE input so activations stay unit-scale across the 61-layer residual chain -- - without it the a8w4 fp8 activation quant (max ~448) overflows to NaN after a - few layers. Both device and reference use the SAME normalization.""" + without it the narrow activation quant (fp4/fp8) saturates and the chain + diverges to NaN after a few layers. Both device and reference use the SAME + normalization.""" xf = x.float() n = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) return n.to(x.dtype) @@ -395,6 +314,36 @@ def _calc_diff(x, y): return float(1 - 2 * (x * y).sum() / denom) +# Accuracy budget, measured on gfx1250 (2 ranks, 1024 tok/rank, 7168x3072, E=384, +# topk=6). +# +# _calc_diff is ||x-y||^2 / (||x||^2 + ||y||^2), a SQUARED error, and the per-layer +# errors accumulate as a random walk: r ~ sqrt(L) makes r^2 ~ L, so the metric grows +# about linearly in the layer count, then saturates as it approaches the bound: +# +# L=1 L=2 L=4 L=8 +# a4w4 0.021877 0.042683 0.080897 0.144881 +# a8w4 0.001433 0.002871 0.005742 0.011369 +# +# slope * L / (1 + sat * L) reproduces both rows within 1%, so scaling that curve +# keeps the SAME headroom at every layer count. A flat tol cannot: 0.1 rejects a +# healthy 8-layer a4w4 run (0.145) yet passes anything at all on a8w4. The rows +# were taken with an mxfp8 combine wire in the loop, which is a touch pessimistic +# here (it cost +32% on a8w4 and +1.6% on a4w4), so these budgets are slightly +# loose rather than tight. +_ACC_TOL = { # quant key -> (per-layer slope, saturation) + "a4w4_mxfp4": (0.0225, 0.031), + "a8w4_mxfp4": (0.00143, 0.0012), +} +_ACC_TOL_SAFETY = 1.5 + + +def default_logits_tol(quant_key, n_layers): + # Per-quant tol for an n_layers chain; see _ACC_TOL for the calibration. + slope, sat = _ACC_TOL[quant_key] + return _ACC_TOL_SAFETY * slope * n_layers / (1.0 + sat * n_layers) + + # torchrun rendezvous helper class Dist: def __init__(self): @@ -452,15 +401,10 @@ def _expert(self, g): if wd is None: w1_g = self.w1_bf[g : g + 1] w2_g = self.w2_bf[g : g + 1] - if self.spec["is_mxfp4"]: - w1_qt, w1_s = _mxfp4_quant(w1_g) - w2_qt, w2_s = _mxfp4_quant(w2_g) - w1d = _mxfp4_dequant(w1_qt, w1_s, (1, *w1_g.shape[1:]))[0] - w2d = _mxfp4_dequant(w2_qt, w2_s, (1, *w2_g.shape[1:]))[0] - else: - # No / fp8 paths: use the bf16 weights directly (approximate ref). - w1d = w1_g[0].float() - w2d = w2_g[0].float() + w1_qt, w1_s = _mxfp4_quant(w1_g) + w2_qt, w2_s = _mxfp4_quant(w2_g) + w1d = _mxfp4_dequant(w1_qt, w1_s, (1, *w1_g.shape[1:]))[0] + w2d = _mxfp4_dequant(w2_qt, w2_s, (1, *w2_g.shape[1:]))[0] wd = self._cache[g] = (w1d, w2d) return wd @@ -521,7 +465,7 @@ def __init__( sw2, routings, ct, - combine_mode="gather", + combine_mode="base", ): self.dist_ctx = dist_ctx self.E, self.hdim, self.idim, self.topk = E, hdim, idim, topk @@ -566,11 +510,7 @@ def setup(self, x0): uid = Communicator.get_unique_id() if r == 0 else None uid = self.dist_ctx.bcast_uid(uid) self.comm = Communicator.init(self.dist_ctx.world, r, uid) - if self.combine_mode == "scatter_fused": - if self.spec["key"] not in _MXFP4_KEYS: - raise NotImplementedError( - f"scatter_fused is available only for {'/'.join(_MXFP4_KEYS)}" - ) + if self.combine_mode == "fused": from aiter.ops.flydsl.kernels.mega_moe_gfx1250 import MegaMoEGfx1250 # Geometry + the expert-GEMM recipe are per-model, so they are fixed @@ -587,8 +527,9 @@ def setup(self, x0): activation=self.spec["activation"], gate_mode=self.spec["gate_mode"].value, quant_type=self.spec["aiter_qtype"], - # Explicit so a stale $MEGA_WIRE cannot change what is measured. - dispatch_wire=self.spec["mega_wire"], + # Passed explicitly: MegaMoE otherwise falls back to $MEGA_WIRE, + # and a stale one there would change what is measured. + dispatch_wire=self.spec["wire"], ) else: EpDispatchCombineConfig, EpDispatchCombineOp = _import_mori_v2() @@ -600,7 +541,7 @@ def setup(self, x0): num_experts_per_rank=self.EPR, num_experts_per_token=self.topk, data_type=self.transport_dtype, - combine_mode=self.combine_mode, + combine_mode="gather", # mori's name for the `base` combine ) self.op = EpDispatchCombineOp(cfg, self.comm) self.comm.barrier() @@ -608,7 +549,7 @@ def setup(self, x0): # ---- one graph-capturable layer + full chain (calls grouped together) ---- # def _layer_step(self, x, layer_idx): ids, wts = self.routings[layer_idx] - xn = _rmsnorm(x) # keep a8w4 fp8 activations in range across 61 layers + xn = _rmsnorm(x) # keep the quantized activations in range across 61 layers if self.mega is not None: y = self.mega( xn, @@ -779,12 +720,16 @@ def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): rows.append((avg_self, name, per_call, pc_avg, avg_count)) rows.sort(key=lambda r: (-r[0], r[1])) dev_per_layer = total_self / per_layer_denom if per_layer_denom else 0.0 + # Wide enough for a full TDM GEMM name, whose tile/warp/buffer recipe and its + # `_epscatter` / `_prefetch` suffix are the whole point of reading this table + # (e.g. a8w4_tdm_fp4_t256x256x256_w2x2_b3_K3072_e96_cn4_prefetch_epscatter). + name_w = 72 lines = [ ( f"# per-call self device time (us) by rank, {world} ranks " f"(rows sorted by total self time):" ), - f"{'Name':<52}" + f"{'Name':<{name_w}}" + "".join(f"{f'rank{r}':>11}" for r in range(world)) + f"{'avg':>11}{'calls':>8}", ] @@ -792,7 +737,9 @@ def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): cells = "".join( f"{v:>11.3f}" if v is not None else f"{'-':>11}" for v in per_call ) - lines.append(f"{name[:52]:<52}{cells}{pc_avg:>11.3f}{avg_count:>8.1f}") + lines.append( + f"{name[:name_w]:<{name_w}}{cells}{pc_avg:>11.3f}{avg_count:>8.1f}" + ) lines.append( f"# TOTAL self device time over ALL {len(rows)} kernels = {total_self:.1f} us " f"-> {dev_per_layer:.1f} us/layer (device-busy; compare to per_layer wall)" @@ -815,17 +762,18 @@ def _device_shared_ffn(tokens, sw1, sw2): # Driver def main(): args = _parse_args() - # The import-time setdefault above pins a8w4; a4w4 is the other half of the - # same switch (aiter/fused_moe.py reads it per call on gfx1250, defaulting to - # fp4x2 unless this is 1), so the quant key has to drive it or -q a4w4_mxfp4 - # silently measures a8w4. + # Set, not setdefault: this is the grouped kernel's activation-dtype switch + # (aiter/fused_moe.py reads it per call on gfx1250, defaulting to fp4x2 unless + # it is 1), so the quant key has to drive it or a stale environment value makes + # -q a4w4_mxfp4 silently measure a8w4. os.environ["AITER_FORCE_A8W4"] = "0" if args.quant_type == "a4w4_mxfp4" else "1" + # Ahead of Dist(): a bad -q/--dispatch_commu_dtype pairing is a width error, + # so report it as one traceback per rank rather than after a rendezvous. + spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype) dist_ctx = Dist() dev = torch.device("cuda", dist_ctx.local_rank) - spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype, args.combine) - spec["mega_wire"] = resolve_mega_wire(args.mega_wire, args.quant_type) - if spec["is_mxfp4"] and get_gfx() not in ("gfx950", "gfx1250"): + if get_gfx() not in ("gfx950", "gfx1250"): if dist_ctx.rank == 0: print( f"skip {args.quant_type}: mxfp4 requires gfx950/gfx1250, got {get_gfx()}" @@ -843,7 +791,7 @@ def main(): print( f"[cfg] world={dist_ctx.world} layers={n_layers} tokens/rank={ct} hidden={hdim} " f"inter={idim} E={E} topk={topk} EPR={E // dist_ctx.world} quant={args.quant_type} " - f"combine={args.combine} mega_wire={spec['mega_wire']} " + f"combine={args.combine} wire={spec['wire']} " f"force_a8w4={os.environ['AITER_FORCE_A8W4']} " f"gate={spec['gate_mode'].name} shared_E={args.shared_experts} gfx={get_gfx()}", flush=True, @@ -939,23 +887,30 @@ def main(): # ---- accuracy (isolated CPU/fp32 reference): end-to-end accumulated compare. accuracy_failure = None if args.acc_verify: + auto_tol = args.logits_tol is None + tol = ( + default_logits_tol(args.quant_type, n_layers) + if auto_tol + else args.logits_tol + ) + tol_desc = f"{tol:.6f}{' auto' if auto_tol else ''}" out_dev = pipe.final_output().float() ref = RefModel(w1_bf, w2_bf, sw1, sw2, spec, dev) ref_out = ref.run(x0, routings).float() logits_diff = _calc_diff(ref_out, out_dev) - errs = dist_ctx.allreduce_sum(0 if logits_diff < args.logits_tol else 1) + errs = dist_ctx.allreduce_sum(0 if logits_diff < tol else 1) avg_diff = dist_ctx.allreduce_avg_float(logits_diff) if dist_ctx.rank == 0: print( f"# MEGA-CHECK layers={n_layers}: {'PASS' if errs == 0 else 'FAIL'} " f"(avg logits_diff={avg_diff:.6f} over {dist_ctx.world} ranks, " - f"tol={args.logits_tol})", + f"tol={tol_desc})", flush=True, ) if errs != 0: accuracy_failure = ( f"MegaMoE accuracy check failed on {errs}/{dist_ctx.world} ranks: " - f"average logits_diff={avg_diff:.6f}, tolerance={args.logits_tol}" + f"average logits_diff={avg_diff:.6f}, tolerance={tol_desc}" ) pipe.teardown() @@ -971,8 +926,9 @@ def _parse_args(): "--quant_type", type=str, choices=QUANT_KEYS, - default="a8w4_mxfp4", - help="quantization type", + default="a4w4_mxfp4", + help="quantization type: mxfp4 weights either way, the prefix picks the " + "grouped GEMM's activation dtype (a8w4 -> fp8, a4w4 -> fp4)", ) p.add_argument( "-tpr", "--token_per_rank", type=int, default=128, help="tokens per rank" @@ -992,7 +948,11 @@ def _parse_args(): help="base RNG seed for weights/tokens/routing (optional; default 0)", ) p.add_argument( - "--logits_tol", type=float, default=0.1, help="end-to-end 1-cosine tol" + "--logits_tol", + type=float, + default=None, + help="end-to-end accuracy tol; default: the per-quant budget for --layers " + "(see _ACC_TOL)", ) p.add_argument( "--acc_verify", type=int, default=1, help="run fp32 reference accuracy check" @@ -1010,27 +970,21 @@ def _parse_args(): p.add_argument( "--dispatch_commu_dtype", type=str, - choices=["auto", "bf16", "fp8"], - default="auto", - help="dispatch transport (communication) dtype", - ) - p.add_argument( - "--mega_wire", - type=str, - choices=["auto", "bf16", "fp8", "fp4"], - default=os.environ.get("MEGA_WIRE", "bf16"), - help="MegaMoE (scatter_fused) dispatch wire: bf16 sends activations and " - "the receiver quantizes each copy; fp8/fp4 quantize once on the sender " - "and forward the e8m0 row. 'auto' picks what the quant key's GEMM wants. " - "Needs dispatch_backend=mori (MEGA_DISPATCH=mori).", + choices=["bf16", "fp8", "fp4"], + default=os.environ.get("DISPATCH_COMMU_DTYPE"), + help="dispatch wire (communication) dtype: bf16 sends activations and the " + "receiver quantizes each copy; fp8/fp4 quantize once on the sender and " + "forward the e8m0 row. A quantizing wire must match what -q's GEMM wants " + "for its A operand (a8w4 -> fp8, a4w4 -> fp4), and the other pairing is " + "rejected. Falls back to $DISPATCH_COMMU_DTYPE, then to -q's own width.", ) p.add_argument( "--combine", type=str, - choices=["gather", "scatter", "scatter_fused"], - default=os.environ.get("COMBINE", "gather"), - help="EP combine mode: gather | scatter | scatter_fused " - "(gemm2-fused P2P scatter; a8w4 only). Falls back to $COMBINE.", + choices=["base", "fused"], + default=os.environ.get("COMBINE", "fused"), + help="EP combine mode: base (mori v2 dispatch/combine around fused_moe) " + "| fused (gemm2-fused P2P scatter). Falls back to $COMBINE.", ) return p.parse_args() From e7560feb2612887f17c09557e0dac6f8b77d49d3 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 02:09:15 +0000 Subject: [PATCH 05/13] revert(test): drop the duplicated a4w4 test rewrite yanbo's 36dac5979 was cherry-picked here and then landed upstream separately, in a revised form, as #5052. Keeping both makes the two versions collide; the mega_moe.py half of that commit stays, since it is not on main. --- .../multigpu_tests/test_mega_moe_gfx1250.py | 384 ++++++++++-------- 1 file changed, 215 insertions(+), 169 deletions(-) diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index 442995678a..5de82fdb78 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -2,8 +2,8 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. """Multi-layer EP MoE end-to-end perf + accuracy on the mori v2 cco/FlyDSL op-layer. -N (default 61, DeepSeek-V4-Pro) MoE layers are chained. The ``base`` mode uses -Mori v2 dispatch -> AITER fused_moe -> Mori v2 combine. ``fused`` calls only +N (default 61, DeepSeek-V4-Pro) MoE layers are chained. Gather/scatter use Mori +v2 dispatch -> AITER fused_moe -> Mori v2 combine. Scatter-fused calls only ``MegaMoEGfx1250``, which owns AITER's dispatch -> fused_moe -> fused-combine pipeline. The combined output plus residual feeds the next layer. @@ -20,15 +20,14 @@ Launcher: torchrun (one process per rank / GPU), mirroring test_moe_layer_ep.py. -Launch (4x gfx1250; every env knob below is already the script's default): +Launch (4x gfx1250, must build CK-free on gfx1250 -> ENABLE_CK=0): cd # avoid the /app/triton namespace shadow + ENABLE_CK=0 AITER_FORCE_A8W4=1 AITER_USE_GROUPED_GEMM=1 AITER_BF16_FP8_MOE_BOUND=0 \ torchrun --standalone --nproc_per_node=4 test_mega_moe_gfx1250.py \ - -q a4w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 --combine fused \ - --dispatch_commu_dtype fp4 + -q a8w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 # Set MORI_CCO_BC to a prebuilt libmori_cco_device.bc to skip CCO JIT. -Env / CLI: --layers --logits_tol --acc_verify --dispatch_commu_dtype --combine - -tpr -hd -id -e -k --shared_E -q +Env / CLI: --layers --logits_tol --acc_verify --dispatch_commu_dtype -tpr -hd -id -e -k --shared_E -q """ import argparse @@ -38,12 +37,15 @@ import torch.distributed as dist import torch.profiler as tprof +import aiter from aiter import ( ActivationType, QuantType, dtypes, get_gfx, + get_hip_quant, get_torch_quant, + pertoken_quant, ) from aiter.fused_moe import fused_moe from aiter.ops.flydsl.moe_common import GateMode @@ -55,25 +57,19 @@ except Exception: # noqa: BLE001 # pragma: no cover get_trace_perf = None -# gfx1250 grouped mxfp4 kernel knobs. Force the real grouped path regardless of -# token count (mirrors test_moe_ep.py). AITER_FORCE_A8W4 is deliberately absent: -# it selects the grouped kernel's ACTIVATION dtype (0 -> fp4, 1 -> fp8) and main() -# derives it from -q, so a stale value in the environment cannot make -q a4w4 -# silently measure a8w4. +# a8w4 (fp8 activation + mxfp4 weight) grouped kernel knobs. Force the real +# fp8/mxfp4 grouped path regardless of token count (mirrors test_moe_ep.py). os.environ.setdefault("ENABLE_CK", "0") +os.environ.setdefault("AITER_FORCE_A8W4", "1") os.environ.setdefault("AITER_USE_GROUPED_GEMM", "1") os.environ.setdefault("AITER_BF16_FP8_MOE_BOUND", "0") -# The `base` path's mori EpDispatchCombineOp and the dispatch inside -# MegaMoEGfx1250 both go through mori's HIP/JIT kernels, so the two modes' -# kernel tables differ only in the combine. MegaMoE's own default is already -# mori (see MegaMoEStage2Config.dispatch_backend). -os.environ.setdefault("MORI_V2_KERNEL_BACKEND", "hip") os.environ.setdefault("FLYDSL_GPU_ARCH", get_gfx()) -# Both keys carry mxfp4 WEIGHTS and differ only in the activation dtype the -# grouped GEMM wants for its A operand: a8w4 -> fp8, a4w4 -> fp4. -QUANT_KEYS = ["a8w4_mxfp4", "a4w4_mxfp4"] +_FP8_DTYPE = dtypes.fp8 +QUANT_KEYS = ["No", "per_Token", "per_128x128", "a8w4_mxfp4", "a4w4_mxfp4"] +_MXFP4_KEYS = ("a8w4_mxfp4", "a4w4_mxfp4") +_FP8_KEYS = ("per_Token", "per_128x128") def _import_mori_comm(): @@ -86,7 +82,7 @@ def _import_mori_comm(): def _import_mori_v2(): """Import Mori's non-fused dispatch/combine v2 path. - Deferred: the fused mode runs entirely on aiter's own mega_moe kernels and + Deferred: scatter_fused runs entirely on aiter's own mega_moe kernels and only needs the communicator, so importing this eagerly would make the fused path fail whenever Mori's copy lags the installed flydsl API. """ @@ -99,56 +95,96 @@ def _import_mori_v2(): # Config / quant-path spec -def resolve_spec(quant_key, wire): - """How to prepare weights / call fused_moe for a quant key, plus the resolved - dispatch wire. Both quant keys use the grouped gfx1250 kernels.""" +def resolve_spec(quant_key, transport, combine_mode="gather"): + """How to prepare weights / quantize activations / call fused_moe for a quant + key, plus the dispatch transport dtype. transport: auto|bf16|fp8.""" + is_mxfp4 = quant_key in _MXFP4_KEYS + is_fp8 = quant_key in _FP8_KEYS + + if transport == "auto": + transport = "fp8" if is_fp8 else "bf16" + if transport == "fp8" and not is_fp8: + transport = "bf16" + + if quant_key == "No": + aiter_qtype = QuantType.No + elif quant_key == "per_Token": + aiter_qtype = QuantType.per_Token + elif quant_key == "per_128x128": + aiter_qtype = QuantType.per_128x128 + else: # a8w4_mxfp4 / a4w4_mxfp4 + aiter_qtype = QuantType.per_1x32 + + gate_mode = GateMode.INTERLEAVE if quant_key == "a8w4_mxfp4" else GateMode.SEPARATED + if is_mxfp4 and combine_mode == "scatter_fused": + # MegaMoE's gemm2-fused scatter is g1u1-interleave only, so a4w4 cannot + # keep the SEPARATED default it uses elsewhere. + gate_mode = GateMode.INTERLEAVE + return { "key": quant_key, - "aiter_qtype": QuantType.per_1x32, - # The gfx1250 grouped MoE GEMM reads GUGU (gate/up row-interleaved) w1 - # only, so both mxfp4 keys have to ask for INTERLEAVE; a SEPARATED layout - # silently falls through to the generic 2-stage MoE. - "gate_mode": GateMode.INTERLEAVE, + "aiter_qtype": aiter_qtype, + "gate_mode": gate_mode, + # Which family of expert kernels the weights are laid out for. The MX + # keys differ only in the ACTIVATION dtype -- the weights are the same + # mxfp4 either way -- but a4w4 has historically meant "the 2-stage mxfp4 + # kernels" here (e8m0_shuffle, is_shuffled) while a8w4 meant "the grouped + # n32k4 ones". MegaMoE is grouped-only, so under it both keys take the + # grouped prep; getting this wrong is silent, and costs a full run to + # find (the output is uncorrelated, not merely imprecise). + "grouped_weights": is_mxfp4 and combine_mode == "scatter_fused", "activation": ActivationType.Silu, - "wire": resolve_wire(wire, quant_key), + "is_mxfp4": is_mxfp4, + "is_fp8": is_fp8, + "transport": transport, + "prequant": transport == "fp8", + "fp8_dtype": _FP8_DTYPE, } -# What the dispatch puts on the wire. A quantizing wire moves the activation -# quant to the SENDER: the payload shrinks (fp8 halves it, fp4 quarters it), a -# per-token e8m0 row rides along, and the receiver skips its own quant. Combine -# is unaffected -- it carries post-expert tokens, which are bf16 regardless. -_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} +# The MegaMoE (scatter_fused) wire. Unrelated to `transport` above, which belongs +# to the mori-v1 dispatch the other combine modes use. +_MEGA_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} -def resolve_wire(wire, quant_key): - """Validate the dispatch wire against the quant key: bf16 | fp8 | fp4. +def resolve_mega_wire(mega_wire, quant_key): + """What MegaMoE's dispatch puts on the wire: bf16 | fp8 | fp4. A quantizing wire is not a free choice -- the receiver hands the payload to - the grouped GEMM as its A operand, so it has to be the width that GEMM wants - (a8w4 -> fp8, a4w4 -> fp4). Picking the other one is a width error, not a - slow path, so it is rejected here rather than deep inside the gather. This - is why the knob is not simply a bool: naming the format keeps the intended - pairing visible at the call site. - - `None` means "unset", and resolves to the quant key's own width so that the - quantizing wire is what you get by default on either key. + the grouped GEMM as its A operand, so it has to be the format that GEMM wants + (a8w4 -> fp8, a4w4 -> fp4), which is what ``auto`` resolves to. Picking the + other one is a width error, not a slow path, so it is rejected here rather + than deep inside the gather. """ - if wire is None: - return _WIRE_FOR_QUANT[quant_key] - if wire == "bf16": + if mega_wire == "auto": + return _MEGA_WIRE_FOR_QUANT.get(quant_key, "bf16") + if mega_wire == "bf16": return "bf16" - want = _WIRE_FOR_QUANT[quant_key] - if wire != want: + want = _MEGA_WIRE_FOR_QUANT.get(quant_key) + if want is None: + raise ValueError( + f"--mega_wire={mega_wire} needs an MX quant key " + f"({'/'.join(_MEGA_WIRE_FOR_QUANT)}), got -q {quant_key}" + ) + if mega_wire != want: raise ValueError( - f"-q {quant_key} wants a {want} A operand, so " - f"--dispatch_commu_dtype={wire} would hand the GEMM the wrong " - "payload width" + f"-q {quant_key} wants a {want} A operand, so --mega_wire={mega_wire} " + "would hand the GEMM the wrong payload width" ) - return wire + return mega_wire # Weight quantization + shuffle (device path) / dequant (reference) +def weight_per_128x128_quant(weight, quant_dtype): + E, dim1, dim2 = weight.shape + wb = weight.view(E, dim1 // 128, 128, dim2 // 128, 128) + wb = wb.permute(0, 1, 3, 2, 4).contiguous().view(E, -1, 128 * 128) + w_qt, w_s = aiter.pertoken_quant(wb, quant_dtype=quant_dtype) + w_qt = w_qt.view(E, dim1 // 128, dim2 // 128, 128, 128) + w_qt = w_qt.permute(0, 1, 3, 2, 4).contiguous().view(E, dim1, dim2) + return w_qt, w_s.view(E, dim1 // 128, dim2 // 128) + + def _mxfp4_quant(w): """per_1x32 mxfp4 quant: packed fp4x2 weight [E, d1, d2//2] + e8m0 scale.""" tq = get_torch_quant(QuantType.per_1x32) @@ -176,34 +212,64 @@ def _gguu_to_gugu_rows(t): def raw_quant_weights(w1, w2, spec): """Quantize (unshuffled) a group of routed-expert weights.""" + key = spec["key"] + if key == "No": + tq = get_torch_quant(QuantType.No) + w1_qt, _ = tq(w1, quant_dtype=None) + w2_qt, _ = tq(w2, quant_dtype=None) + return w1_qt.view(w1.shape), None, w2_qt.view(w2.shape), None + if key == "per_Token": + w1_qt, w1_s = pertoken_quant(w1, quant_dtype=_FP8_DTYPE) + w2_qt, w2_s = pertoken_quant(w2, quant_dtype=_FP8_DTYPE) + return w1_qt, w1_s, w2_qt, w2_s + if key == "per_128x128": + w1_qt, w1_s = weight_per_128x128_quant(w1, quant_dtype=_FP8_DTYPE) + w2_qt, w2_s = weight_per_128x128_quant(w2, quant_dtype=_FP8_DTYPE) + return w1_qt, w1_s, w2_qt, w2_s w1_qt, w1_s = _mxfp4_quant(w1) w2_qt, w2_s = _mxfp4_quant(w2) return w1_qt, w1_s, w2_qt, w2_s def shuffle_group(w1_qt, w1_s, w2_qt, w2_s, spec, n_experts): - """Layout-shuffle a group of `n_experts` quantized experts for the kernel. + """Layout-shuffle a group of `n_experts` quantized experts for the kernel.""" + key = spec["key"] + if key in ("No", "per_Token", "per_128x128"): + return shuffle_weight(w1_qt), shuffle_weight(w2_qt), w1_s, w2_s + if key == "a8w4_mxfp4" or spec.get("grouped_weights"): + if spec["gate_mode"] == GateMode.INTERLEAVE: + w1_phys = _gguu_to_gugu_rows(w1_qt.view(torch.uint8)) + w1_a = shuffle_weight(w1_phys, layout=(16, 16)) + w1_ss = moe_shuffle_scale( + w1_s.contiguous(), + experts_cnt=n_experts, + is_guinterleave=True, + gate_up=True, + ) + else: + w1_a = shuffle_weight(w1_qt.view(torch.uint8), layout=(16, 16)) + w1_ss = moe_shuffle_scale(w1_s.contiguous(), experts_cnt=n_experts) + w2_a = shuffle_weight(w2_qt.view(torch.uint8), layout=(16, 16)) + w2_ss = moe_shuffle_scale(w2_s.contiguous(), experts_cnt=n_experts) + return w1_a, w2_a, w1_ss, w2_ss + # a4w4_mxfp4 on the 2-stage kernels (a different B layout from the grouped + # branch above -- e8m0_shuffle, not the n32k4 fold). + w1_a = shuffle_weight(w1_qt, layout=(16, 16)) + w2_a = shuffle_weight(w2_qt, layout=(16, 16)) + w1_ss = fp4_utils.e8m0_shuffle(w1_s) + w2_ss = fp4_utils.e8m0_shuffle(w2_s) + w1_a.is_shuffled = True + w2_a.is_shuffled = True + return w1_a, w2_a, w1_ss, w2_ss - Both mxfp4 keys share the grouped gfx1250 layout (GUGU-interleaved w1 + the - n32k4 e8m0 B-scale). The only difference is the weight DTYPE handed to - ``fused_moe``: it keys off ``w1.dtype`` to pick the activation dtype, so - uint8 selects the fp8-activation (a8w4) kernel and fp4x2 the fp4-activation - (a4w4) one. See ``grouped_moe_gfx1250._grouped_a8w4_tdm_moe``. - """ - w1_phys = _gguu_to_gugu_rows(w1_qt.view(torch.uint8)) - w1_a = shuffle_weight(w1_phys, layout=(16, 16)) - w2_a = shuffle_weight(w2_qt.view(torch.uint8), layout=(16, 16)) - w1_ss = moe_shuffle_scale( - w1_s.contiguous(), - experts_cnt=n_experts, - is_guinterleave=True, - gate_up=True, + +def quant_tokens_fp8(tokens, spec): + """Per-token / per-block fp8 quant of the activations (fp8 pre-quant transport).""" + qt = spec["aiter_qtype"] + quant_func = get_hip_quant( + qt if qt != QuantType.per_128x128 else QuantType.per_1x128 ) - w2_ss = moe_shuffle_scale(w2_s.contiguous(), experts_cnt=n_experts) - if spec["key"] == "a4w4_mxfp4": - w1_a = w1_a.view(dtypes.fp4x2) - w2_a = w2_a.view(dtypes.fp4x2) - return w1_a, w2_a, w1_ss, w2_ss + return quant_func(tokens, quant_dtype=spec["fp8_dtype"]) def moe_forward( @@ -216,6 +282,7 @@ def moe_forward( topk_ids, expert_mask, spec, + a1_scale=None, num_local_tokens=None, ): """Single fused_moe call (device path). ``num_local_tokens`` (device int32 @@ -226,20 +293,35 @@ def moe_forward( num_local_tokens = torch.tensor( [hidden.shape[0]], dtype=dtypes.i32, device=hidden.device ) + if spec["is_mxfp4"]: + return fused_moe( + hidden, + w1_a, + w2_a, + topk_weights, + topk_ids, + expert_mask=expert_mask, + activation=spec["activation"], + gate_mode=spec["gate_mode"].value, + quant_type=spec["aiter_qtype"], + w1_scale=w1_s, + w2_scale=w2_s, + dtype=dtypes.bf16, + num_local_tokens=num_local_tokens, + ) return fused_moe( hidden, w1_a, w2_a, topk_weights, topk_ids, - expert_mask=expert_mask, - activation=spec["activation"], - gate_mode=spec["gate_mode"].value, - quant_type=spec["aiter_qtype"], + expert_mask, + num_local_tokens=num_local_tokens, w1_scale=w1_s, w2_scale=w2_s, + quant_type=spec["aiter_qtype"], + a1_scale=a1_scale, dtype=dtypes.bf16, - num_local_tokens=num_local_tokens, ) @@ -297,9 +379,8 @@ def make_routings(n_layers, ct, E, topk, dev, seed): def _rmsnorm(x, eps=1e-6): """RMSNorm (no learnable gain) on the last dim. Applied to each layer's MoE input so activations stay unit-scale across the 61-layer residual chain -- - without it the narrow activation quant (fp4/fp8) saturates and the chain - diverges to NaN after a few layers. Both device and reference use the SAME - normalization.""" + without it the a8w4 fp8 activation quant (max ~448) overflows to NaN after a + few layers. Both device and reference use the SAME normalization.""" xf = x.float() n = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) return n.to(x.dtype) @@ -314,36 +395,6 @@ def _calc_diff(x, y): return float(1 - 2 * (x * y).sum() / denom) -# Accuracy budget, measured on gfx1250 (2 ranks, 1024 tok/rank, 7168x3072, E=384, -# topk=6). -# -# _calc_diff is ||x-y||^2 / (||x||^2 + ||y||^2), a SQUARED error, and the per-layer -# errors accumulate as a random walk: r ~ sqrt(L) makes r^2 ~ L, so the metric grows -# about linearly in the layer count, then saturates as it approaches the bound: -# -# L=1 L=2 L=4 L=8 -# a4w4 0.021877 0.042683 0.080897 0.144881 -# a8w4 0.001433 0.002871 0.005742 0.011369 -# -# slope * L / (1 + sat * L) reproduces both rows within 1%, so scaling that curve -# keeps the SAME headroom at every layer count. A flat tol cannot: 0.1 rejects a -# healthy 8-layer a4w4 run (0.145) yet passes anything at all on a8w4. The rows -# were taken with an mxfp8 combine wire in the loop, which is a touch pessimistic -# here (it cost +32% on a8w4 and +1.6% on a4w4), so these budgets are slightly -# loose rather than tight. -_ACC_TOL = { # quant key -> (per-layer slope, saturation) - "a4w4_mxfp4": (0.0225, 0.031), - "a8w4_mxfp4": (0.00143, 0.0012), -} -_ACC_TOL_SAFETY = 1.5 - - -def default_logits_tol(quant_key, n_layers): - # Per-quant tol for an n_layers chain; see _ACC_TOL for the calibration. - slope, sat = _ACC_TOL[quant_key] - return _ACC_TOL_SAFETY * slope * n_layers / (1.0 + sat * n_layers) - - # torchrun rendezvous helper class Dist: def __init__(self): @@ -401,10 +452,15 @@ def _expert(self, g): if wd is None: w1_g = self.w1_bf[g : g + 1] w2_g = self.w2_bf[g : g + 1] - w1_qt, w1_s = _mxfp4_quant(w1_g) - w2_qt, w2_s = _mxfp4_quant(w2_g) - w1d = _mxfp4_dequant(w1_qt, w1_s, (1, *w1_g.shape[1:]))[0] - w2d = _mxfp4_dequant(w2_qt, w2_s, (1, *w2_g.shape[1:]))[0] + if self.spec["is_mxfp4"]: + w1_qt, w1_s = _mxfp4_quant(w1_g) + w2_qt, w2_s = _mxfp4_quant(w2_g) + w1d = _mxfp4_dequant(w1_qt, w1_s, (1, *w1_g.shape[1:]))[0] + w2d = _mxfp4_dequant(w2_qt, w2_s, (1, *w2_g.shape[1:]))[0] + else: + # No / fp8 paths: use the bf16 weights directly (approximate ref). + w1d = w1_g[0].float() + w2d = w2_g[0].float() wd = self._cache[g] = (w1d, w2d) return wd @@ -465,7 +521,7 @@ def __init__( sw2, routings, ct, - combine_mode="base", + combine_mode="gather", ): self.dist_ctx = dist_ctx self.E, self.hdim, self.idim, self.topk = E, hdim, idim, topk @@ -510,7 +566,11 @@ def setup(self, x0): uid = Communicator.get_unique_id() if r == 0 else None uid = self.dist_ctx.bcast_uid(uid) self.comm = Communicator.init(self.dist_ctx.world, r, uid) - if self.combine_mode == "fused": + if self.combine_mode == "scatter_fused": + if self.spec["key"] not in _MXFP4_KEYS: + raise NotImplementedError( + f"scatter_fused is available only for {'/'.join(_MXFP4_KEYS)}" + ) from aiter.ops.flydsl.kernels.mega_moe_gfx1250 import MegaMoEGfx1250 # Geometry + the expert-GEMM recipe are per-model, so they are fixed @@ -527,9 +587,8 @@ def setup(self, x0): activation=self.spec["activation"], gate_mode=self.spec["gate_mode"].value, quant_type=self.spec["aiter_qtype"], - # Passed explicitly: MegaMoE otherwise falls back to $MEGA_WIRE, - # and a stale one there would change what is measured. - dispatch_wire=self.spec["wire"], + # Explicit so a stale $MEGA_WIRE cannot change what is measured. + dispatch_wire=self.spec["mega_wire"], ) else: EpDispatchCombineConfig, EpDispatchCombineOp = _import_mori_v2() @@ -541,7 +600,7 @@ def setup(self, x0): num_experts_per_rank=self.EPR, num_experts_per_token=self.topk, data_type=self.transport_dtype, - combine_mode="gather", # mori's name for the `base` combine + combine_mode=self.combine_mode, ) self.op = EpDispatchCombineOp(cfg, self.comm) self.comm.barrier() @@ -549,7 +608,7 @@ def setup(self, x0): # ---- one graph-capturable layer + full chain (calls grouped together) ---- # def _layer_step(self, x, layer_idx): ids, wts = self.routings[layer_idx] - xn = _rmsnorm(x) # keep the quantized activations in range across 61 layers + xn = _rmsnorm(x) # keep a8w4 fp8 activations in range across 61 layers if self.mega is not None: y = self.mega( xn, @@ -720,16 +779,12 @@ def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): rows.append((avg_self, name, per_call, pc_avg, avg_count)) rows.sort(key=lambda r: (-r[0], r[1])) dev_per_layer = total_self / per_layer_denom if per_layer_denom else 0.0 - # Wide enough for a full TDM GEMM name, whose tile/warp/buffer recipe and its - # `_epscatter` / `_prefetch` suffix are the whole point of reading this table - # (e.g. a8w4_tdm_fp4_t256x256x256_w2x2_b3_K3072_e96_cn4_prefetch_epscatter). - name_w = 72 lines = [ ( f"# per-call self device time (us) by rank, {world} ranks " f"(rows sorted by total self time):" ), - f"{'Name':<{name_w}}" + f"{'Name':<52}" + "".join(f"{f'rank{r}':>11}" for r in range(world)) + f"{'avg':>11}{'calls':>8}", ] @@ -737,9 +792,7 @@ def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): cells = "".join( f"{v:>11.3f}" if v is not None else f"{'-':>11}" for v in per_call ) - lines.append( - f"{name[:name_w]:<{name_w}}{cells}{pc_avg:>11.3f}{avg_count:>8.1f}" - ) + lines.append(f"{name[:52]:<52}{cells}{pc_avg:>11.3f}{avg_count:>8.1f}") lines.append( f"# TOTAL self device time over ALL {len(rows)} kernels = {total_self:.1f} us " f"-> {dev_per_layer:.1f} us/layer (device-busy; compare to per_layer wall)" @@ -762,18 +815,17 @@ def _device_shared_ffn(tokens, sw1, sw2): # Driver def main(): args = _parse_args() - # Set, not setdefault: this is the grouped kernel's activation-dtype switch - # (aiter/fused_moe.py reads it per call on gfx1250, defaulting to fp4x2 unless - # it is 1), so the quant key has to drive it or a stale environment value makes - # -q a4w4_mxfp4 silently measure a8w4. + # The import-time setdefault above pins a8w4; a4w4 is the other half of the + # same switch (aiter/fused_moe.py reads it per call on gfx1250, defaulting to + # fp4x2 unless this is 1), so the quant key has to drive it or -q a4w4_mxfp4 + # silently measures a8w4. os.environ["AITER_FORCE_A8W4"] = "0" if args.quant_type == "a4w4_mxfp4" else "1" - # Ahead of Dist(): a bad -q/--dispatch_commu_dtype pairing is a width error, - # so report it as one traceback per rank rather than after a rendezvous. - spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype) dist_ctx = Dist() dev = torch.device("cuda", dist_ctx.local_rank) + spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype, args.combine) + spec["mega_wire"] = resolve_mega_wire(args.mega_wire, args.quant_type) - if get_gfx() not in ("gfx950", "gfx1250"): + if spec["is_mxfp4"] and get_gfx() not in ("gfx950", "gfx1250"): if dist_ctx.rank == 0: print( f"skip {args.quant_type}: mxfp4 requires gfx950/gfx1250, got {get_gfx()}" @@ -791,7 +843,7 @@ def main(): print( f"[cfg] world={dist_ctx.world} layers={n_layers} tokens/rank={ct} hidden={hdim} " f"inter={idim} E={E} topk={topk} EPR={E // dist_ctx.world} quant={args.quant_type} " - f"combine={args.combine} wire={spec['wire']} " + f"combine={args.combine} mega_wire={spec['mega_wire']} " f"force_a8w4={os.environ['AITER_FORCE_A8W4']} " f"gate={spec['gate_mode'].name} shared_E={args.shared_experts} gfx={get_gfx()}", flush=True, @@ -887,30 +939,23 @@ def main(): # ---- accuracy (isolated CPU/fp32 reference): end-to-end accumulated compare. accuracy_failure = None if args.acc_verify: - auto_tol = args.logits_tol is None - tol = ( - default_logits_tol(args.quant_type, n_layers) - if auto_tol - else args.logits_tol - ) - tol_desc = f"{tol:.6f}{' auto' if auto_tol else ''}" out_dev = pipe.final_output().float() ref = RefModel(w1_bf, w2_bf, sw1, sw2, spec, dev) ref_out = ref.run(x0, routings).float() logits_diff = _calc_diff(ref_out, out_dev) - errs = dist_ctx.allreduce_sum(0 if logits_diff < tol else 1) + errs = dist_ctx.allreduce_sum(0 if logits_diff < args.logits_tol else 1) avg_diff = dist_ctx.allreduce_avg_float(logits_diff) if dist_ctx.rank == 0: print( f"# MEGA-CHECK layers={n_layers}: {'PASS' if errs == 0 else 'FAIL'} " f"(avg logits_diff={avg_diff:.6f} over {dist_ctx.world} ranks, " - f"tol={tol_desc})", + f"tol={args.logits_tol})", flush=True, ) if errs != 0: accuracy_failure = ( f"MegaMoE accuracy check failed on {errs}/{dist_ctx.world} ranks: " - f"average logits_diff={avg_diff:.6f}, tolerance={tol_desc}" + f"average logits_diff={avg_diff:.6f}, tolerance={args.logits_tol}" ) pipe.teardown() @@ -926,9 +971,8 @@ def _parse_args(): "--quant_type", type=str, choices=QUANT_KEYS, - default="a4w4_mxfp4", - help="quantization type: mxfp4 weights either way, the prefix picks the " - "grouped GEMM's activation dtype (a8w4 -> fp8, a4w4 -> fp4)", + default="a8w4_mxfp4", + help="quantization type", ) p.add_argument( "-tpr", "--token_per_rank", type=int, default=128, help="tokens per rank" @@ -948,11 +992,7 @@ def _parse_args(): help="base RNG seed for weights/tokens/routing (optional; default 0)", ) p.add_argument( - "--logits_tol", - type=float, - default=None, - help="end-to-end accuracy tol; default: the per-quant budget for --layers " - "(see _ACC_TOL)", + "--logits_tol", type=float, default=0.1, help="end-to-end 1-cosine tol" ) p.add_argument( "--acc_verify", type=int, default=1, help="run fp32 reference accuracy check" @@ -970,21 +1010,27 @@ def _parse_args(): p.add_argument( "--dispatch_commu_dtype", type=str, - choices=["bf16", "fp8", "fp4"], - default=os.environ.get("DISPATCH_COMMU_DTYPE"), - help="dispatch wire (communication) dtype: bf16 sends activations and the " - "receiver quantizes each copy; fp8/fp4 quantize once on the sender and " - "forward the e8m0 row. A quantizing wire must match what -q's GEMM wants " - "for its A operand (a8w4 -> fp8, a4w4 -> fp4), and the other pairing is " - "rejected. Falls back to $DISPATCH_COMMU_DTYPE, then to -q's own width.", + choices=["auto", "bf16", "fp8"], + default="auto", + help="dispatch transport (communication) dtype", + ) + p.add_argument( + "--mega_wire", + type=str, + choices=["auto", "bf16", "fp8", "fp4"], + default=os.environ.get("MEGA_WIRE", "bf16"), + help="MegaMoE (scatter_fused) dispatch wire: bf16 sends activations and " + "the receiver quantizes each copy; fp8/fp4 quantize once on the sender " + "and forward the e8m0 row. 'auto' picks what the quant key's GEMM wants. " + "Needs dispatch_backend=mori (MEGA_DISPATCH=mori).", ) p.add_argument( "--combine", type=str, - choices=["base", "fused"], - default=os.environ.get("COMBINE", "fused"), - help="EP combine mode: base (mori v2 dispatch/combine around fused_moe) " - "| fused (gemm2-fused P2P scatter). Falls back to $COMBINE.", + choices=["gather", "scatter", "scatter_fused"], + default=os.environ.get("COMBINE", "gather"), + help="EP combine mode: gather | scatter | scatter_fused " + "(gemm2-fused P2P scatter; a8w4 only). Falls back to $COMBINE.", ) return p.parse_args() From edd47d69c4bc5034bdd6fb8a11ecb8099dcce7c3 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 14:45:20 +0000 Subject: [PATCH 06/13] fix(ci): the gfx1250 MegaMoE job passes a --combine value that no longer exists #5052 narrowed --combine to base|fused; the CI line still says scatter_fused, so the job has been exiting on argparse rather than running the test. --- .github/scripts/aiter_test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 2a26c59315..5126140600 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -101,7 +101,7 @@ for file in "${sharded_files[@]}"; do fi exec env MORI_SHMEM_HEAP_SIZE=40G \ torchrun --standalone --nproc_per_node=8 "$test_file" \ - --combine scatter_fused --layers 2 --acc_verify 1 + --combine fused --layers 2 --acc_verify 1 ' _ "$file" ) From d2a7dc5b19ff7e58111872765bcade4fc10bf287 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 14:46:09 +0000 Subject: [PATCH 07/13] refactor(mega_moe): leave the dispatch backend default alone Flipping it to mori was an out-of-scope behaviour change: a quantizing wire already requires the mori backend, and __post_init__ says so with a message naming the kwarg, so the default never had to move. The dataclass field default is unreachable anyway -- MegaMoEGfx1250.__init__ always supplies a value -- but the $MEGA_DISPATCH fallback beside it is live, and it decided the backend for every caller that sets neither. --- aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index 9a8e63fee7..32ad4653ec 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -143,7 +143,7 @@ class MegaMoEStage2Config: dispatch_block_num: int | None = None dispatch_warp_num_per_block: int | None = None schedule: tuple | None = None - dispatch_backend: str = "mori" + dispatch_backend: str = "flydsl" # What dispatch puts on the wire. fp8 halves the payload and fp4 quarters it, # each sending a per-token e8m0 row along; the receiver then skips its own # quant. Combine is unaffected -- it moves post-expert tokens, which are bf16 @@ -422,7 +422,7 @@ def __init__( dispatch_backend=( dispatch_backend if dispatch_backend is not None - else os.environ.get("MEGA_DISPATCH", "mori") + else os.environ.get("MEGA_DISPATCH", "flydsl") ), dispatch_wire=( dispatch_wire From 31fea2513126c92c816967ba3ee5c68079167066 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 14:47:33 +0000 Subject: [PATCH 08/13] refactor(test): one dispatch-wire flag, and drop the one that never worked --dispatch_commu_dtype has had no effect since the file was created in #4785: resolve_spec wrote transport/prequant/is_fp8/fp8_dtype into the spec and nothing ever read them, quant_tokens_fp8() had no callers, and the base combine hardcodes transport_dtype = bf16. So this is not two overlapping knobs -- it is one live knob and one corpse. --mega_wire takes over the name (--dispatch_wire, matching the kwarg it feeds) and the dead half goes with the flag it belonged to. _FP8_DTYPE stays: the per_Token / per_128x128 weight quant still uses it. --- .../multigpu_tests/test_mega_moe_gfx1250.py | 81 +++++++------------ 1 file changed, 27 insertions(+), 54 deletions(-) diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index dd9118b4c1..c1798fa50b 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -26,7 +26,8 @@ -q a4w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 --combine base # Set MORI_CCO_BC to a prebuilt libmori_cco_device.bc to skip CCO JIT. -Env / CLI: --layers --logits_tol --acc_verify --dispatch_commu_dtype -tpr -hd -id -e -k --shared_E -q +Env / CLI: --layers --logits_tol --acc_verify --dispatch_wire --combine + -tpr -hd -id -e -k --shared_E -q """ import argparse @@ -42,7 +43,6 @@ QuantType, dtypes, get_gfx, - get_hip_quant, get_torch_quant, pertoken_quant, ) @@ -77,7 +77,6 @@ _FP8_DTYPE = dtypes.fp8 QUANT_KEYS = ["No", "per_Token", "per_128x128", "a8w4_mxfp4", "a4w4_mxfp4"] _MXFP4_KEYS = ("a8w4_mxfp4", "a4w4_mxfp4") -_FP8_KEYS = ("per_Token", "per_128x128") def _import_mori_comm(): @@ -103,16 +102,10 @@ def _import_mori_v2(): # Config / quant-path spec -def resolve_spec(quant_key, transport): +def resolve_spec(quant_key): """How to prepare weights / quantize activations / call fused_moe for a quant - key, plus the dispatch transport dtype. transport: auto|bf16|fp8.""" + key.""" is_mxfp4 = quant_key in _MXFP4_KEYS - is_fp8 = quant_key in _FP8_KEYS - - if transport == "auto": - transport = "fp8" if is_fp8 else "bf16" - if transport == "fp8" and not is_fp8: - transport = "bf16" if quant_key == "No": aiter_qtype = QuantType.No @@ -134,19 +127,14 @@ def resolve_spec(quant_key, transport): "gate_mode": gate_mode, "activation": ActivationType.Silu, "is_mxfp4": is_mxfp4, - "is_fp8": is_fp8, - "transport": transport, - "prequant": transport == "fp8", - "fp8_dtype": _FP8_DTYPE, } -# The MegaMoE (fused) dispatch wire. Unrelated to `transport` above, which -# belongs to the mori-v1 dispatch the base combine uses. -_MEGA_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} +# The MegaMoE (--combine fused) dispatch wire. +_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} -def resolve_mega_wire(mega_wire, quant_key): +def resolve_dispatch_wire(wire, quant_key): """What MegaMoE's dispatch puts on the wire: bf16 | fp8 | fp4. A quantizing wire is not a free choice -- the receiver hands the payload to @@ -155,22 +143,22 @@ def resolve_mega_wire(mega_wire, quant_key): pairing is a width error, not a slow path, so it is rejected here rather than deep inside the gather. """ - if mega_wire == "auto": - return _MEGA_WIRE_FOR_QUANT.get(quant_key, "bf16") - if mega_wire == "bf16": + if wire == "auto": + return _WIRE_FOR_QUANT.get(quant_key, "bf16") + if wire == "bf16": return "bf16" - want = _MEGA_WIRE_FOR_QUANT.get(quant_key) + want = _WIRE_FOR_QUANT.get(quant_key) if want is None: raise ValueError( - f"--mega_wire={mega_wire} needs an MX quant key " - f"({'/'.join(_MEGA_WIRE_FOR_QUANT)}), got -q {quant_key}" + f"--dispatch_wire={wire} needs an MX quant key " + f"({'/'.join(_WIRE_FOR_QUANT)}), got -q {quant_key}" ) - if mega_wire != want: + if wire != want: raise ValueError( - f"-q {quant_key} wants a {want} A operand, so --mega_wire={mega_wire} " + f"-q {quant_key} wants a {want} A operand, so --dispatch_wire={wire} " "would hand the GEMM the wrong payload width" ) - return mega_wire + return wire # Weight quantization + shuffle (device path) / dequant (reference) @@ -258,15 +246,6 @@ def shuffle_group(w1_qt, w1_s, w2_qt, w2_s, spec, n_experts): return w1_a, w2_a, w1_ss, w2_ss -def quant_tokens_fp8(tokens, spec): - """Per-token / per-block fp8 quant of the activations (fp8 pre-quant transport).""" - qt = spec["aiter_qtype"] - quant_func = get_hip_quant( - qt if qt != QuantType.per_128x128 else QuantType.per_1x128 - ) - return quant_func(tokens, quant_dtype=spec["fp8_dtype"]) - - def moe_forward( hidden, w1_a, @@ -614,8 +593,8 @@ def setup(self, x0): activation=self.spec["activation"], gate_mode=self.spec["gate_mode"].value, quant_type=self.spec["aiter_qtype"], - # Explicit so a stale $MEGA_WIRE cannot change what is measured. - dispatch_wire=self.spec["mega_wire"], + # Explicit so a stale $MEGA_DISPATCH_WIRE cannot change what is measured. + dispatch_wire=self.spec["dispatch_wire"], ) else: EpDispatchCombineConfig, EpDispatchCombineOp = _import_mori_v2() @@ -853,8 +832,8 @@ def main(): # Set, not setdefault: the wire has to match this, so a stale environment # value would otherwise make -q a4w4_mxfp4 silently measure a8w4. os.environ["AITER_FORCE_A8W4"] = "0" if args.quant_type == "a4w4_mxfp4" else "1" - spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype) - spec["mega_wire"] = resolve_mega_wire(args.mega_wire, args.quant_type) + spec = resolve_spec(args.quant_type) + spec["dispatch_wire"] = resolve_dispatch_wire(args.dispatch_wire, args.quant_type) if spec["is_mxfp4"] and get_gfx() not in ("gfx950", "gfx1250"): if dist_ctx.rank == 0: @@ -874,7 +853,7 @@ def main(): print( f"[cfg] world={dist_ctx.world} layers={n_layers} tokens/rank={ct} hidden={hdim} " f"inter={idim} E={E} topk={topk} EPR={E // dist_ctx.world} quant={args.quant_type} " - f"combine={args.combine} mega_wire={spec['mega_wire']} " + f"combine={args.combine} dispatch_wire={spec['dispatch_wire']} " f"force_a8w4={os.environ['AITER_FORCE_A8W4']} " f"gate={spec['gate_mode'].name} shared_E={args.shared_experts} gfx={get_gfx()}", flush=True, @@ -1050,20 +1029,14 @@ def _parse_args(): "can stall multi-rank graph-profile runs)", ) p.add_argument( - "--dispatch_commu_dtype", - type=str, - choices=["auto", "bf16", "fp8"], - default="auto", - help="dispatch transport (communication) dtype", - ) - p.add_argument( - "--mega_wire", + "--dispatch_wire", type=str, choices=["auto", "bf16", "fp8", "fp4"], - default=os.environ.get("MEGA_WIRE", "bf16"), - help="MegaMoE (--combine fused) dispatch wire: bf16 sends activations and " - "the receiver quantizes each copy; fp8/fp4 quantize once on the sender " - "and forward the e8m0 row. 'auto' picks what the quant key's GEMM wants.", + default=os.environ.get("MEGA_DISPATCH_WIRE", "bf16"), + help="what dispatch puts on the wire (--combine fused only): bf16 sends " + "activations and the receiver quantizes each copy; fp8/fp4 quantize once " + "on the sender and forward the e8m0 row. 'auto' picks what the quant " + "key's GEMM wants.", ) p.add_argument( "--combine", From 67b547b086fe43dd9eab89725cb2c664454b8580 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 14:49:04 +0000 Subject: [PATCH 09/13] refactor(mega_moe): say dispatch out loud on everything that means the dispatch wire Half the names this PR added already carried the direction (dispatch_token_nbytes, combine_token_nbytes); the other half did not, and a combine wire is coming. is_quant_wire is the one that mattered: it is a pure dispatch predicate with nine call sites, so 'if config.is_quant_wire' written inside _combine() would take the wrong branch AND RUN -- a bf16 reduce over fp8 bytes, wrong numbers, no error. _DispatchWire is deliberately not generalised: three of its four fields are structurally dispatch-only, so sharing the table would hand combine three dead ones. --- .../kernels/mega_moe_gfx1250/mega_moe.py | 105 ++++++++++-------- 1 file changed, 58 insertions(+), 47 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index 32ad4653ec..8a8d638fa2 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -39,14 +39,20 @@ "xdb": "cross_device_barrier", # Only laid out on a quantizing wire; plan_api binds a missing region to 0 and # the kernel's `if constexpr` keeps that 0 from being read. - "outScales": "out_scales", + "outScales": "disp_out_scales", } @dataclass(frozen=True) -class _Wire: - """One dispatch wire. Per token at hidden 7168: bf16 14336 B, fp8 7168 + 256, +class _DispatchWire: + """One DISPATCH wire. Per token at hidden 7168: bf16 14336 B, fp8 7168 + 256, fp4 3584 + 256. + + Dispatch-only on purpose: a quantized combine cannot reuse this. mori + carries no scales on a combine, the quant would have to happen inside the + gemm2 epilogue on an LDS tile rather than host-side on a whole tensor, and + recv_dtype exists only to build a torch view combine has no equivalent of. + Only payload_bytes would carry over. """ payload_bytes: float # PER FEATURE; fp4 packs two features into a byte @@ -57,12 +63,12 @@ class _Wire: recv_dtype: torch.dtype -_WIRES = { - "bf16": _Wire(2, torch.bfloat16, None, torch.bfloat16), - "fp8": _Wire(1, dtypes.fp8, dtypes.fp8, dtypes.fp8), - "fp4": _Wire(0.5, dtypes.fp4x2, dtypes.fp4x2, torch.uint8), +_DISPATCH_WIRE_SPECS = { + "bf16": _DispatchWire(2, torch.bfloat16, None, torch.bfloat16), + "fp8": _DispatchWire(1, dtypes.fp8, dtypes.fp8, dtypes.fp8), + "fp4": _DispatchWire(0.5, dtypes.fp4x2, dtypes.fp4x2, torch.uint8), } -_DISPATCH_WIRES = tuple(_WIRES) +_DISPATCH_WIRES = tuple(_DISPATCH_WIRE_SPECS) def _align_up(value: int, alignment: int) -> int: @@ -140,6 +146,9 @@ class MegaMoEStage2Config: max_tokens_per_rank: int experts_per_rank: int topk: int + # Dispatch (stage1) knobs. They live here because this package has no + # stage1 config yet -- dispatch and gemm1 are not fused. When that fusion + # lands they move together into whatever config it brings. dispatch_block_num: int | None = None dispatch_warp_num_per_block: int | None = None schedule: tuple | None = None @@ -160,16 +169,17 @@ def __post_init__(self): f"dispatch_wire must be one of {_DISPATCH_WIRES}, " f"got {self.dispatch_wire!r}" ) - if self.is_quant_wire and self.dispatch_backend != "mori": + if self.is_quant_dispatch_wire and self.dispatch_backend != "mori": # Only mori's kernel carries the scale row; this package's own # dispatch has no channel for it. raise ValueError( f"dispatch_wire={self.dispatch_wire!r} requires " f"dispatch_backend='mori' (got {self.dispatch_backend!r})" ) - if self.is_quant_wire and self.hidden_dim % 32: + if self.is_quant_dispatch_wire and self.hidden_dim % 32: raise ValueError( - "one e8m0 scale covers 32 features, so a quantizing wire needs " + "one e8m0 scale covers 32 features, so a quantizing dispatch wire " + "needs " f"hidden_dim % 32 == 0, got {self.hidden_dim}" ) if self.dispatch_backend not in _DISPATCH_BACKENDS: @@ -215,27 +225,28 @@ def max_recv(self) -> int: return self.world_size * self.max_tokens_per_rank @property - def is_quant_wire(self) -> bool: + def is_quant_dispatch_wire(self) -> bool: """The wire carries an MX payload plus its e8m0 row, not bf16.""" return self.dispatch_wire in ("fp8", "fp4") @property - def wire(self) -> "_Wire": - return _WIRES[self.dispatch_wire] + def dispatch_wire_spec(self) -> "_DispatchWire": + return _DISPATCH_WIRE_SPECS[self.dispatch_wire] @property def dispatch_token_nbytes(self) -> int: - return int(self.hidden_dim * self.wire.payload_bytes) + return int(self.hidden_dim * self.dispatch_wire_spec.payload_bytes) @property - def wire_elem_count(self) -> int: + def dispatch_wire_elem_count(self) -> int: """What mori's Cfg calls hidden_dim: ELEMENTS, at its own element size. - fp8 and fp4 both transport as one byte per element, so an fp4 wire has to + fp8 and fp4 both transport as one byte per element, so an fp4 dispatch + wire has to halve the count itself -- mori sizes the token as hidden_dim * elem_size and would otherwise move two bytes per packed byte. """ - return self.dispatch_token_nbytes if self.is_quant_wire else self.hidden_dim + return self.dispatch_token_nbytes if self.is_quant_dispatch_wire else self.hidden_dim @property def combine_token_nbytes(self) -> int: @@ -243,25 +254,25 @@ def combine_token_nbytes(self) -> int: return self.hidden_dim * 2 @property - def scale_nbytes(self) -> int: + def dispatch_scale_nbytes(self) -> int: """Per-token e8m0 row as WE produce it: one byte per 32 features, packed. Handed to mori as-is. mori lays it down at its own, 128 B-aligned stride - (scale_dst_nbytes) because that is what keeps a TDM run's start aligned; + (dispatch_scale_dst_nbytes) because that is what keeps a TDM run's start aligned; that padding is mori's business, and the quant op's output can go straight - onto the wire without a repack. + onto the dispatch wire without a repack. """ - return self.hidden_dim // 32 if self.is_quant_wire else 0 + return self.hidden_dim // 32 if self.is_quant_dispatch_wire else 0 @property - def scale_dst_nbytes(self) -> int: + def dispatch_scale_dst_nbytes(self) -> int: """The stride the rows ARRIVE at, which the receiving gather addresses by. Asked of mori rather than recomputed: it is the transport's layout decision, and a local copy of the rule would drift the first time the alignment changes. """ - if not self.is_quant_wire: + if not self.is_quant_dispatch_wire: return 0 try: from mori.ops.dispatch_combine_v2.hip_backend import scale_stride_bytes @@ -274,7 +285,7 @@ def scale_dst_nbytes(self) -> int: "the installed one has no scale_stride_bytes" ) from e - return scale_stride_bytes(self.scale_nbytes) + return scale_stride_bytes(self.dispatch_scale_nbytes) @property def combine_slot_stride_bytes(self) -> int: @@ -533,12 +544,12 @@ def forward( recv_x = recv_x[:bound] recv_weights = recv_weights[:bound] recv_ids = recv_ids[:bound] - if self._config.is_quant_wire: + if self._config.is_quant_dispatch_wire: assert a1_scale is None, ( - "a1_scale is produced by the quantizing wire itself; a " + "a1_scale is produced by the quantizing dispatch wire itself; a " "caller-supplied one would be silently discarded" ) - a1_scale = self._recv_scales() + a1_scale = self._recv_dispatch_scales() if recv_token_bound is not None: a1_scale = a1_scale[: int(recv_token_bound)] extra = {} @@ -598,8 +609,8 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): *( # Arrival stride, not the packed row: undersizing overruns # the last slots. - [("out_scales", max_recv * config.scale_dst_nbytes)] - if config.scale_dst_nbytes + [("disp_out_scales", max_recv * config.dispatch_scale_dst_nbytes)] + if config.dispatch_scale_dst_nbytes else [] ), ( @@ -624,7 +635,7 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): self._dispatch_barrier = torch.zeros(1, dtype=torch.int32, device=device) # Points at the quant op's own scale rows, set per dispatch on a quantizing # wire; 0 (and unread) on bf16. - self._sent_scales_ptr = 0 + self._dispatch_sent_scales_ptr = 0 self._total_recv = torch.zeros(1, dtype=torch.int32, device=device) self._cross_device_flag = torch.ones(1, dtype=torch.int64, device=device) self._combine_output = torch.zeros( @@ -721,22 +732,22 @@ def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: "built by mori's CMake and is not shipped by every install" ) from error - # Passed only on a quantizing wire, matching scale_dst_nbytes: mori grew + # Passed only on a quantizing wire, matching dispatch_scale_dst_nbytes: mori grew # scale_bytes in #593 and rejects UNKNOWN kwargs outright, so sending the # bf16 wire's harmless 0 would make an older mori refuse the whole plan. - scale_kw = {"scale_bytes": config.scale_nbytes} if config.is_quant_wire else {} + scale_kw = {"scale_bytes": config.dispatch_scale_nbytes} if config.is_quant_dispatch_wire else {} plans = {} for spec in self._dispatch_specs: plan = EpDispatchPlan( world_size=config.world_size, - # see wire_elem_count; mori's plan_api: "the caller halves + # see dispatch_wire_elem_count; mori's plan_api: "the caller halves # hiddenDim" - hidden_dim=config.wire_elem_count, + hidden_dim=config.dispatch_wire_elem_count, max_tok_per_rank=config.max_tokens_per_rank, num_expert_per_rank=config.experts_per_rank, num_expert_per_token=config.topk, max_recv=config.max_recv, - dtype=config.wire.mori_dtype, + dtype=config.dispatch_wire_spec.mori_dtype, use_weights=True, **scale_kw, block_num=spec[0], @@ -781,7 +792,7 @@ def launch( # launcher is a traced @flyc.jit signature, and widening it # would put a dead kernarg on a path that can never carry # scales (fp8 requires dispatch_backend='mori'). - scales_buf=self._sent_scales_ptr, + scales_buf=self._dispatch_sent_scales_ptr, ) return launch @@ -804,24 +815,24 @@ def _select_dispatch(self, token_count: int) -> tuple[int, int]: def _recv_tokens(self) -> torch.Tensor: config = self._config # Width in whatever recv_dtype counts: features for bf16/fp8, bytes for - # fp4 -- see _Wire.recv_dtype. - width = config.dispatch_token_nbytes // config.wire.recv_dtype.itemsize + # fp4 -- see _DispatchWire.recv_dtype. + width = config.dispatch_token_nbytes // config.dispatch_wire_spec.recv_dtype.itemsize return _from_gpu_ptr( self._arena.local_ptr("disp_out"), (config.max_recv, width), - config.wire.recv_dtype, + config.dispatch_wire_spec.recv_dtype, ) - def _recv_scales(self) -> torch.Tensor | None: + def _recv_dispatch_scales(self) -> torch.Tensor | None: """The forwarded e8m0 rows, or None on the bf16 wire.""" - if not self._config.is_quant_wire: + if not self._config.is_quant_dispatch_wire: return None # Full padded rows, not a trimmed view: this goes to the gather kernel as a # base pointer plus a build-constant pitch, and that pitch is the arrival # stride. The kernel reads only the meaningful bytes of each row. return _from_gpu_ptr( - self._arena.local_ptr("out_scales"), - (self._config.max_recv, self._config.scale_dst_nbytes), + self._arena.local_ptr("disp_out_scales"), + (self._config.max_recv, self._config.dispatch_scale_dst_nbytes), torch.uint8, ) @@ -849,7 +860,7 @@ def _dispatch( spec = self._select_dispatch(token_count) stream = fx.Stream(torch.cuda.current_stream()) payload = hidden_states - if self._config.is_quant_wire: + if self._config.is_quant_dispatch_wire: # Quantize ONCE PER LOCAL TOKEN here, instead of once per received # copy on the far side. Destination-independent, so the bytes are the # same either way; the preshuffle cannot move with it, because its @@ -858,13 +869,13 @@ def _dispatch( payload, scale_rows = per_1x32_mx_quant_hip( hidden_states, - quant_dtype=self._config.wire.quant_dtype, + quant_dtype=self._config.dispatch_wire_spec.quant_dtype, scale_type=dtypes.fp8_e8m0, shuffle=False, ) # Straight onto the wire; mori restrides these packed rows while it # stages them, so there is no repack here. - self._sent_scales_ptr = scale_rows.data_ptr() + self._dispatch_sent_scales_ptr = scale_rows.data_ptr() self._dispatch_variants[spec]( self._arena.handle, payload.data_ptr(), From 3af755aa4c7a12a42ec9ceb0894a3c620530c9f4 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 14:49:51 +0000 Subject: [PATCH 10/13] refactor(mega_moe): MEGA_WIRE becomes MEGA_DISPATCH_WIRE, and the old name is fatal Combine will want a wire of its own, so the unqualified name had to go while it is still unreleased. The stale name raises rather than falling back: an env var that is silently ignored sends a run that asked for fp4 down the bf16 path and reports nothing, which is exactly the failure a wire benchmark cannot survive. Setting both names to the same value is allowed, so a fleet can be rolled over one script at a time. --- .../kernels/mega_moe_gfx1250/mega_moe.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index 8a8d638fa2..c1ea49cd36 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -43,6 +43,24 @@ } +def _dispatch_wire_from_env() -> str: + """$MEGA_DISPATCH_WIRE, and a loud death for the name it replaced. + + Not a fallback: an env var that is silently ignored sends a run that asked + for fp4 down the bf16 path and reports nothing, which is the one failure + mode a wire benchmark cannot survive. + """ + stale, current = os.environ.get("MEGA_WIRE"), os.environ.get("MEGA_DISPATCH_WIRE") + if stale is not None and current != stale: + raise RuntimeError( + "MEGA_WIRE was renamed to MEGA_DISPATCH_WIRE (combine gets its own " + f"wire); found MEGA_WIRE={stale!r} with MEGA_DISPATCH_WIRE=" + f"{current!r}. Update the launch script rather than relying on the " + "old name -- it is no longer read." + ) + return current or "bf16" + + @dataclass(frozen=True) class _DispatchWire: """One DISPATCH wire. Per token at hidden 7168: bf16 14336 B, fp8 7168 + 256, @@ -438,7 +456,7 @@ def __init__( dispatch_wire=( dispatch_wire if dispatch_wire is not None - else os.environ.get("MEGA_WIRE", "bf16") + else _dispatch_wire_from_env() ), ), communicator, From e033ed96c37eb465e5db5fd9213f4ef64f2e84f3 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 15:04:07 +0000 Subject: [PATCH 11/13] fix(gfx1250): stop reusing a1_scale for two different tensors The parameter is the caller's per-token e8m0 rows; forty lines later the quant pass rebinds the same name to the preshuffled grouped scale. Both are uint8 with a plausible shape, so a wrong read is silent. The incoming one now has its own name and the rebinding introduces the grouped meaning exactly once. --- aiter/ops/flydsl/grouped_moe_gfx1250.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index b4c4b7571f..4b67bdc883 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -604,15 +604,20 @@ def _grouped_a8w4_tdm_moe( _quant_mode = "fp4" if _is_fp4 else "fp8" _a_is_fp4 = 1 if _is_fp4 else 0 + # Bound once, because the quant pass below rebinds a1_scale to the + # PRESHUFFLED GROUPED scale. Both are uint8 and both have a plausible + # shape, so nothing downstream could tell which one it was handed. + src_a1_scale = a1_scale + # Pre-quantized activation: an MX payload plus its e8m0 row is what a # quantizing EP dispatch delivers, and it is also aiter's standing meaning # for this pair. - _prequantized = a1_scale is not None and hidden_states.dtype in ( + _prequantized = src_a1_scale is not None and hidden_states.dtype in ( dtypes.fp8, torch.uint8, dtypes.fp4x2, ) - if a1_scale is not None and not _prequantized: + if src_a1_scale is not None and not _prequantized: # Loud rather than silently re-quantizing something already quantized. assert hidden_states.dtype == dtype, ( f"a1_scale given with hidden_states dtype {hidden_states.dtype}: " @@ -638,7 +643,7 @@ def _grouped_a8w4_tdm_moe( topids_to_rows=topids_to_rows, source_topk=topk, num_valid_routes=_ep_nvr, - prequantized_scale=a1_scale if _prequantized else None, + prequantized_scale=src_a1_scale if _prequantized else None, ) # Fuse gemm1 activation + MX quantization + scale preshuffle into the From ebc95068b7abff7ca6a6ba0bd10b990eab01ce60 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 15:04:24 +0000 Subject: [PATCH 12/13] refactor(mega_moe): one copy of the stale-env check, and satisfy black Three renamed lines went past 88 columns, which the pre-checks job enforces. The env guard was also true of one caller only: the library evaluates it just when dispatch_wire is None, and the test always passes the kwarg, so the harness -- the launch path most likely to carry a stale MEGA_WIRE -- would have measured bf16 in silence. The helper is now public and the test's argparse default calls it, imported lazily so the mega package is not pulled in before FLYDSL_GPU_ARCH is set. _WIRE_FOR_QUANT picks up the dispatch prefix too; it was the same ambiguity the commit before this one exists to remove. --- .../kernels/mega_moe_gfx1250/__init__.py | 4 +-- .../kernels/mega_moe_gfx1250/mega_moe.py | 30 ++++++++++++------- .../multigpu_tests/test_mega_moe_gfx1250.py | 14 +++++---- 3 files changed, 31 insertions(+), 17 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py index 3526be5a7d..d005eece54 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py @@ -5,11 +5,11 @@ import importlib -__all__ = ["MegaMoEGfx1250"] +__all__ = ["MegaMoEGfx1250", "read_dispatch_wire_env"] def __getattr__(name): - if name != "MegaMoEGfx1250": + if name not in __all__: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") return getattr(importlib.import_module(f"{__name__}.mega_moe"), name) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index c1ea49cd36..b511babfcb 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -43,7 +43,7 @@ } -def _dispatch_wire_from_env() -> str: +def read_dispatch_wire_env() -> str: """$MEGA_DISPATCH_WIRE, and a loud death for the name it replaced. Not a fallback: an env var that is silently ignored sends a run that asked @@ -196,9 +196,8 @@ def __post_init__(self): ) if self.is_quant_dispatch_wire and self.hidden_dim % 32: raise ValueError( - "one e8m0 scale covers 32 features, so a quantizing dispatch wire " - "needs " - f"hidden_dim % 32 == 0, got {self.hidden_dim}" + "one e8m0 scale covers 32 features, so a quantizing dispatch " + f"wire needs hidden_dim % 32 == 0, got {self.hidden_dim}" ) if self.dispatch_backend not in _DISPATCH_BACKENDS: raise ValueError( @@ -260,11 +259,15 @@ def dispatch_wire_elem_count(self) -> int: """What mori's Cfg calls hidden_dim: ELEMENTS, at its own element size. fp8 and fp4 both transport as one byte per element, so an fp4 dispatch - wire has to - halve the count itself -- mori sizes the token as hidden_dim * elem_size + wire has to halve the count itself -- mori sizes the token as + hidden_dim * elem_size and would otherwise move two bytes per packed byte. """ - return self.dispatch_token_nbytes if self.is_quant_dispatch_wire else self.hidden_dim + return ( + self.dispatch_token_nbytes + if self.is_quant_dispatch_wire + else self.hidden_dim + ) @property def combine_token_nbytes(self) -> int: @@ -456,7 +459,7 @@ def __init__( dispatch_wire=( dispatch_wire if dispatch_wire is not None - else _dispatch_wire_from_env() + else read_dispatch_wire_env() ), ), communicator, @@ -753,7 +756,11 @@ def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: # Passed only on a quantizing wire, matching dispatch_scale_dst_nbytes: mori grew # scale_bytes in #593 and rejects UNKNOWN kwargs outright, so sending the # bf16 wire's harmless 0 would make an older mori refuse the whole plan. - scale_kw = {"scale_bytes": config.dispatch_scale_nbytes} if config.is_quant_dispatch_wire else {} + scale_kw = ( + {"scale_bytes": config.dispatch_scale_nbytes} + if config.is_quant_dispatch_wire + else {} + ) plans = {} for spec in self._dispatch_specs: plan = EpDispatchPlan( @@ -834,7 +841,10 @@ def _recv_tokens(self) -> torch.Tensor: config = self._config # Width in whatever recv_dtype counts: features for bf16/fp8, bytes for # fp4 -- see _DispatchWire.recv_dtype. - width = config.dispatch_token_nbytes // config.dispatch_wire_spec.recv_dtype.itemsize + width = ( + config.dispatch_token_nbytes + // config.dispatch_wire_spec.recv_dtype.itemsize + ) return _from_gpu_ptr( self._arena.local_ptr("disp_out"), (config.max_recv, width), diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index c1798fa50b..a183f41fa7 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -131,7 +131,7 @@ def resolve_spec(quant_key): # The MegaMoE (--combine fused) dispatch wire. -_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} +_DISPATCH_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} def resolve_dispatch_wire(wire, quant_key): @@ -144,14 +144,14 @@ def resolve_dispatch_wire(wire, quant_key): than deep inside the gather. """ if wire == "auto": - return _WIRE_FOR_QUANT.get(quant_key, "bf16") + return _DISPATCH_WIRE_FOR_QUANT.get(quant_key, "bf16") if wire == "bf16": return "bf16" - want = _WIRE_FOR_QUANT.get(quant_key) + want = _DISPATCH_WIRE_FOR_QUANT.get(quant_key) if want is None: raise ValueError( f"--dispatch_wire={wire} needs an MX quant key " - f"({'/'.join(_WIRE_FOR_QUANT)}), got -q {quant_key}" + f"({'/'.join(_DISPATCH_WIRE_FOR_QUANT)}), got -q {quant_key}" ) if wire != want: raise ValueError( @@ -982,6 +982,10 @@ def main(): def _parse_args(): + # Imported here, not at module scope: pulling in the mega package before + # FLYDSL_GPU_ARCH is set below would hand flydsl the wrong arch. + from aiter.ops.flydsl.kernels.mega_moe_gfx1250 import read_dispatch_wire_env + p = argparse.ArgumentParser(description="multi-layer EP MoE perf + accuracy") p.add_argument( "-q", @@ -1032,7 +1036,7 @@ def _parse_args(): "--dispatch_wire", type=str, choices=["auto", "bf16", "fp8", "fp4"], - default=os.environ.get("MEGA_DISPATCH_WIRE", "bf16"), + default=read_dispatch_wire_env(), help="what dispatch puts on the wire (--combine fused only): bf16 sends " "activations and the receiver quantizes each copy; fp8/fp4 quantize once " "on the sender and forward the e8m0 row. 'auto' picks what the quant " From b21cdcf5c0c542a73434955b9e15ca325661c1a4 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Tue, 1 Sep 2026 02:30:40 +0000 Subject: [PATCH 13/13] refactor(mega_moe): the op-level config is not a stage2 config It carries geometry and the dispatch knobs and not one stage2 parameter. Stage2 in this package is the gemm2 epilogue fused into combine, which Stage2ScatterContext already names correctly and which takes nothing from here -- so the config was the only place the word was wrong. Name only: the fields, the flat kwargs and the public surface are unchanged, so no caller moves. --- .../flydsl/kernels/mega_moe_gfx1250/mega_moe.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py index b511babfcb..963eb44c9a 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -157,7 +157,13 @@ def close(self): @dataclass -class MegaMoEStage2Config: +class MegaMoEConfig: + """Op-level config: geometry, plus the dispatch knobs. + + Nothing here tunes stage2 -- that is the gemm2 epilogue fused into combine + (Stage2ScatterContext), which takes no parameter from this side. + """ + rank: int world_size: int hidden_dim: int @@ -444,7 +450,7 @@ def __init__( self.expert_mask[first_expert : first_expert + self.experts_per_rank] = 1 self._initialize_pipeline( - MegaMoEStage2Config( + MegaMoEConfig( rank=int(rank), world_size=int(world_size), hidden_dim=self.model_dim, @@ -611,7 +617,7 @@ def __enter__(self): def __exit__(self, *exc): self.close() - def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): + def _initialize_pipeline(self, config: MegaMoEConfig, communicator): self._config = config self._closed = False device = torch.device("cuda", torch.cuda.current_device()) @@ -726,7 +732,7 @@ def _initialize_pipeline(self, config: MegaMoEStage2Config, communicator): off_xdb_mem=self._arena.offset("cross_device_barrier"), ) - def _build_mori_dispatch(self, config: MegaMoEStage2Config) -> dict: + def _build_mori_dispatch(self, config: MegaMoEConfig) -> dict: """mori's HIP/JIT dispatch, wearing `_make_dispatch`'s calling convention. Only the kernel changes: mori leaves the same arena state this package's