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" ) 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 720d1a13e2..ee7fec5a1a 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -854,6 +854,9 @@ def _fused_moe_impl( doweight_stage1=doweight_stage1, w1_scale=w1_scale, w2_scale=w2_scale, + # 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, intermediate_pad=intermediate_pad, diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index aee2fe0271..4b67bdc883 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 + # 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 = src_a1_scale is not None and hidden_states.dtype in ( + dtypes.fp8, + torch.uint8, + dtypes.fp4x2, + ) + 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}: " + "expected packed MX bytes (pre-quantized) or the model dtype (ignored)" + ) + # 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 + 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=src_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). @@ -988,19 +1020,17 @@ 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 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 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 @@ -1169,6 +1199,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/__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 8b02e53262..963eb44c9a 100644 --- a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py @@ -37,9 +37,58 @@ "dispOut": "disp_out", "outTok": "comb_inp", "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": "disp_out_scales", } +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 + 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, + 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 + 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 + + +_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(_DISPATCH_WIRE_SPECS) + + def _align_up(value: int, alignment: int) -> int: return (value + alignment - 1) // alignment * alignment @@ -108,19 +157,54 @@ 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 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 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_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_dispatch_wire and self.hidden_dim % 32: + raise ValueError( + "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( f"dispatch_backend must be one of {_DISPATCH_BACKENDS}, " @@ -164,13 +248,76 @@ def max_recv(self) -> int: return self.world_size * self.max_tokens_per_rank @property - def token_nbytes(self) -> int: + 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 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.dispatch_wire_spec.payload_bytes) + + @property + 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 + and would otherwise move two bytes per packed byte. + """ + return ( + self.dispatch_token_nbytes + if self.is_quant_dispatch_wire + else self.hidden_dim + ) + + @property + def combine_token_nbytes(self) -> int: + """Combine moves bf16 post-expert tokens, whatever the wire carried.""" return self.hidden_dim * 2 + @property + 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 + (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 dispatch wire without a repack. + """ + return self.hidden_dim // 32 if self.is_quant_dispatch_wire else 0 + + @property + 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_dispatch_wire: + return 0 + 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 + + return scale_stride_bytes(self.dispatch_scale_nbytes) + @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 +358,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. @@ -302,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, @@ -314,6 +462,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 read_dispatch_wire_env() + ), ), communicator, ) @@ -418,6 +571,14 @@ def forward( recv_x = recv_x[:bound] recv_weights = recv_weights[:bound] recv_ids = recv_ids[:bound] + if self._config.is_quant_dispatch_wire: + assert a1_scale is None, ( + "a1_scale is produced by the quantizing dispatch wire itself; a " + "caller-supplied one would be silently discarded" + ) + a1_scale = self._recv_dispatch_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 @@ -456,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()) @@ -470,8 +631,15 @@ 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), + *( + # Arrival stride, not the packed row: undersizing overruns + # the last slots. + [("disp_out_scales", max_recv * config.dispatch_scale_dst_nbytes)] + if config.dispatch_scale_dst_nbytes + else [] + ), ( "comb_inp", config.max_tokens_per_rank @@ -492,6 +660,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) + # Points at the quant op's own scale rows, set per dispatch on a quantizing + # wire; 0 (and unread) on bf16. + 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( @@ -561,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 @@ -588,17 +759,28 @@ 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 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 {} + ) plans = {} for spec in self._dispatch_specs: plan = EpDispatchPlan( world_size=config.world_size, - hidden_dim=config.hidden_dim, + # see dispatch_wire_elem_count; mori's plan_api: "the caller halves + # hiddenDim" + 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=torch.bfloat16, + dtype=config.dispatch_wire_spec.mori_dtype, use_weights=True, + **scale_kw, block_num=spec[0], warp_per_block=spec[1], arena=self._arena, @@ -636,6 +818,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._dispatch_sent_scales_ptr, ) return launch @@ -656,10 +844,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 + # 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 + ) return _from_gpu_ptr( self._arena.local_ptr("disp_out"), - (self._config.max_recv, self._config.hidden_dim), - torch.bfloat16, + (config.max_recv, width), + config.dispatch_wire_spec.recv_dtype, + ) + + def _recv_dispatch_scales(self) -> torch.Tensor | None: + """The forwarded e8m0 rows, or None on the bf16 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("disp_out_scales"), + (self._config.max_recv, self._config.dispatch_scale_dst_nbytes), + torch.uint8, ) def _recv_weights(self) -> torch.Tensor: @@ -685,9 +893,26 @@ 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_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 + # 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=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._dispatch_sent_scales_ptr = scale_rows.data_ptr() 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,54 @@ 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: 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): + 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 +1335,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 +1347,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 +1379,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 +1419,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 +1544,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 +1597,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 +1612,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 67111c2931..a183f41fa7 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,13 +127,40 @@ 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 (--combine fused) dispatch wire. +_DISPATCH_WIRE_FOR_QUANT = {"a8w4_mxfp4": "fp8", "a4w4_mxfp4": "fp4"} + + +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 + the grouped GEMM as its A operand, so it has to be the width that GEMM wants + (a8w4 -> fp8, a4w4 -> fp4), which is what ``auto`` resolves to. The other + pairing is a width error, not a slow path, so it is rejected here rather + than deep inside the gather. + """ + if wire == "auto": + return _DISPATCH_WIRE_FOR_QUANT.get(quant_key, "bf16") + if wire == "bf16": + return "bf16" + 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(_DISPATCH_WIRE_FOR_QUANT)}), got -q {quant_key}" + ) + if wire != want: + raise ValueError( + f"-q {quant_key} wants a {want} A operand, so --dispatch_wire={wire} " + "would hand the GEMM the wrong payload width" + ) + return wire + + # Weight quantization + shuffle (device path) / dequant (reference) def weight_per_128x128_quant(weight, quant_dtype): E, dim1, dim2 = weight.shape @@ -226,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, @@ -582,6 +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_DISPATCH_WIRE cannot change what is measured. + dispatch_wire=self.spec["dispatch_wire"], ) else: EpDispatchCombineConfig, EpDispatchCombineOp = _import_mori_v2() @@ -816,7 +829,11 @@ def main(): args = _parse_args() dist_ctx = Dist() dev = torch.device("cuda", dist_ctx.local_rank) - spec = resolve_spec(args.quant_type, args.dispatch_commu_dtype) + # 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) + 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: @@ -836,7 +853,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} 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, ) @@ -964,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", @@ -1011,11 +1033,14 @@ def _parse_args(): "can stall multi-rank graph-profile runs)", ) p.add_argument( - "--dispatch_commu_dtype", + "--dispatch_wire", type=str, - choices=["auto", "bf16", "fp8"], - default="auto", - help="dispatch transport (communication) dtype", + choices=["auto", "bf16", "fp8", "fp4"], + 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 " + "key's GEMM wants.", ) p.add_argument( "--combine",