diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 78fc57bc415..756e28ddafe 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -106,6 +106,28 @@ for file in "${sharded_files[@]}"; do _ "$file" ) ;; + op_tests/multigpu_tests/test_mega_moe_gfx1250.py) + { + echo "Running gfx1250 MegaMoE fused-scatter accuracy on 8 GPUs when supported" + } | tee -a latest_test.log + test_cmd=( + timeout 60m + bash -c ' + set -euo pipefail + test_file=$1 + arch=$(python3 -c \ + "from aiter.jit.utils.chip_info import get_gfx; print(get_gfx())") + if [[ "$arch" != "gfx1250" ]]; then + echo "Skipping $test_file: requires gfx1250, got $arch" + exit 0 + fi + exec env MORI_SHMEM_HEAP_SIZE=40G \ + torchrun --standalone --nproc_per_node=8 "$test_file" \ + --combine scatter_fused --layers 2 --acc_verify 1 + ' + _ "$file" + ) + ;; op_tests/multigpu_tests/test_mega_moe_v2.py) { echo "Running MegaMoEV2 v4_pro fixed-slot and compact coverage on 8 GPUs" diff --git a/aiter/aot/flydsl/grouped_moe.py b/aiter/aot/flydsl/grouped_moe.py index dd338ef2928..f5c219915bb 100644 --- a/aiter/aot/flydsl/grouped_moe.py +++ b/aiter/aot/flydsl/grouped_moe.py @@ -386,7 +386,7 @@ def compile_one_config(**job): t0 = time.time() try: # TODO(aot): only the auxiliary (non-GEMM) kernels are precompiled here. - # The grouped GEMM itself moved to the felix TDM batched kernel + # The grouped GEMM itself moved to the TDM batched kernel # (aiter.ops.flydsl.batched_gemm_mxfp4), which has no AOT wiring yet, so # it JIT-compiles on first use. Add TDM GEMM jobs here to restore the # coverage the deleted moe_grouped_gemm_mxscale path used to provide. diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index 3a1bd702c98..f94621a3562 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -31,6 +31,7 @@ gfx_from_cu_num, ) from aiter.jit.utils.torch_guard import torch_compile_guard +from aiter.ops.flydsl.kernels.mega_moe_gfx1250.types import Stage2ScatterContext try: from aiter.ops.flydsl.moe_common import GateMode @@ -489,6 +490,7 @@ def fused_moe( shared_w1_scale: torch.Tensor | None = None, shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, + stage2_scatter: Stage2ScatterContext | None = None, ): if ( any( @@ -531,6 +533,8 @@ def fused_moe( ) if not block_size_M: block_size_M = -1 + enable_ep_scatter = stage2_scatter is not None + scatter_source_map = stage2_scatter.source_token_map if enable_ep_scatter else None return fused_moe_( hidden_states=hidden_states, w1=w1, @@ -557,6 +561,18 @@ def fused_moe( beta=beta, linear_beta=linear_beta, gate_mode=gate_mode, + ep_arena_handle=stage2_scatter.arena_handle if enable_ep_scatter else 0, + ep_combine_input_offset=( + stage2_scatter.combine_input_offset if enable_ep_scatter else 0 + ), + ep_slot_stride_bytes=( + stage2_scatter.slot_stride_bytes if enable_ep_scatter else 0 + ), + ep_max_tokens_per_rank=( + stage2_scatter.max_tokens_per_rank if enable_ep_scatter else 0 + ), + ep_world_size=stage2_scatter.world_size if enable_ep_scatter else 0, + ep_source_token_map=scatter_source_map, ) @@ -588,6 +604,12 @@ def fused_moe_fake( beta: float | None = None, linear_beta: float | None = None, gate_mode: str = GateMode.SEPARATED.value, + ep_arena_handle: int = 0, + ep_combine_input_offset: int = 0, + ep_slot_stride_bytes: int = 0, + ep_max_tokens_per_rank: int = 0, + ep_world_size: int = 0, + ep_source_token_map: torch.Tensor | None = None, ) -> torch.Tensor: device = topk_ids.device M, _topk = topk_ids.shape @@ -626,7 +648,23 @@ def fused_moe_( beta: float | None = None, linear_beta: float | None = None, gate_mode: str = GateMode.SEPARATED.value, + ep_arena_handle: int = 0, + ep_combine_input_offset: int = 0, + ep_slot_stride_bytes: int = 0, + ep_max_tokens_per_rank: int = 0, + ep_world_size: int = 0, + ep_source_token_map: torch.Tensor | None = None, ) -> torch.Tensor: + stage2_scatter = None + if ep_source_token_map is not None: + stage2_scatter = Stage2ScatterContext( + arena_handle=ep_arena_handle, + combine_input_offset=ep_combine_input_offset, + slot_stride_bytes=ep_slot_stride_bytes, + max_tokens_per_rank=ep_max_tokens_per_rank, + world_size=ep_world_size, + source_token_map=ep_source_token_map, + ) return _fused_moe_impl( hidden_states=hidden_states, w1=w1, @@ -653,6 +691,7 @@ def fused_moe_( beta=beta, linear_beta=linear_beta, gate_mode=gate_mode, + stage2_scatter=stage2_scatter, ) @@ -682,6 +721,7 @@ def _fused_moe_impl( beta: float | None = None, linear_beta: float | None = None, gate_mode: str = GateMode.SEPARATED.value, + stage2_scatter: Stage2ScatterContext | None = None, *, _q_dtype_a: torch.dtype | None = None, _metadata_transform: Callable | None = None, @@ -804,6 +844,7 @@ def _fused_moe_impl( num_local_tokens=num_local_tokens, situ_beta=1.0 if beta is None else float(beta), situ_linear_beta=1.0 if linear_beta is None else float(linear_beta), + stage2_scatter=stage2_scatter, ) if grouped_a8w4_out is not None: diff --git a/aiter/ops/flydsl/grouped_gemm_mxfp4.py b/aiter/ops/flydsl/grouped_gemm_mxfp4.py index 5018cbb638d..d310644e65b 100644 --- a/aiter/ops/flydsl/grouped_gemm_mxfp4.py +++ b/aiter/ops/flydsl/grouped_gemm_mxfp4.py @@ -9,6 +9,7 @@ import torch +from .kernels.mega_moe_gfx1250.types import Stage2ScatterContext from .kernels.tensor_shim import ptr_arg _SUPPORTED_CLUSTER_N = (4, 3, 2) @@ -90,6 +91,9 @@ def flydsl_grouped_gemm_a8w4_masked( cluster_n=-1, waves_per_tensor_tdm=-1, next_stage_prefetch=0, + stage2_scatter: Stage2ScatterContext | None = None, + ep_destination_stride=0, + ep_row_map=None, situ_beta=1.0, situ_linear_beta=1.0, ): @@ -115,6 +119,8 @@ def flydsl_grouped_gemm_a8w4_masked( f"[grouped-moe tdm] cluster_n={cluster_n} needs n_tiles={n_tiles} " f"(N={N}, tile_n={tile_n}) to be an exact multiple" ) + enable_ep_scatter = stage2_scatter is not None + ep_row_map_tensor = ep_row_map if ep_row_map is not None else out launch_gemm_a8w4_tdm( out, ptr_arg(a), @@ -145,7 +151,18 @@ def flydsl_grouped_gemm_a8w4_masked( cluster_n, _select_next_stage_prefetch(next_stage_prefetch), waves_per_tensor_tdm, - float(situ_beta), - float(situ_linear_beta), + enable_ep_scatter=int(enable_ep_scatter), + ep_arena_handle=(int(stage2_scatter.arena_handle) if enable_ep_scatter else 0), + ep_combine_input_offset=( + int(stage2_scatter.combine_input_offset) if enable_ep_scatter else 0 + ), + ep_slot_stride_bytes=( + int(stage2_scatter.slot_stride_bytes) if enable_ep_scatter else 0 + ), + ep_destination_stride=int(ep_destination_stride), + ep_world_size=int(stage2_scatter.world_size) if enable_ep_scatter else 0, + arg_ep_row_map=ep_row_map_tensor, + f32_situ_beta=float(situ_beta), + f32_situ_linear_beta=float(situ_linear_beta), ) return out diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index 75df6e167bc..98107ccedac 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -14,6 +14,7 @@ from aiter import ActivationType, QuantType, dtypes, logger from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.kernels.mega_moe_gfx1250.types import Stage2ScatterContext from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg # Opt-in switch for the gfx1250 FlyDSL grouped-GEMM path. @@ -442,6 +443,7 @@ def _grouped_a8w4_tdm_moe( data_format="a8w4", expert_mask=None, num_local_tokens=None, + stage2_scatter: Stage2ScatterContext | None = None, situ_beta=1.0, situ_linear_beta=1.0, ): @@ -457,6 +459,7 @@ def _grouped_a8w4_tdm_moe( device = hidden_states.device token_num, topk = topk_ids.shape + enable_ep_scatter = stage2_scatter is not None if tile_m2 is None: tile_m2 = tile_m if tile_n2 is None: @@ -481,9 +484,12 @@ def _grouped_a8w4_tdm_moe( # route kernel remaps them to local buckets via ``g2l_lut`` (sentinel E = # dropped/non-local route), casts the f32 route weights into ``_gather_w_buf`` # (kept -> weight_dtype, dropped -> 0), and skips the EP dead-tail (routes >= - # num_valid_routes / tokens >= num_valid_tokens). The felix TDM batched - # GEMMs themselves are EP-agnostic: they operate on the already-routed - # contiguous layout. + # num_valid_routes / tokens >= num_valid_tokens). Non-local routes claim no + # per-expert slot, so ``_masked_m`` / ``psum`` -- and the rows the GEMMs + # actually compute -- cover only the routes this rank owns. ``contiguous_m`` + # stays the static worst case to keep the launch grid CUDAGraph-safe; the + # surplus tiles fall past ``psum`` and exit early. The TDM batched GEMMs + # themselves are EP-agnostic: they operate on the routed contiguous layout. _is_ep = expert_mask is not None _g2l_lut = None _g2l_counter = None @@ -525,10 +531,45 @@ def _grouped_a8w4_tdm_moe( ) else: _masked_m, topids_to_rows = flydsl_moe_topids_to_rows(topk_ids, E, max_m) + # EP gemm2-fused scatter: build the ep_rowmap inside the remap pass, which + # already knows each route's final contiguous row, so the gemm2 TDM epilogue + # can P2P each weighted row into peers' comb_inp. + ep_scatter_params = None + ep_rowmap = None + if enable_ep_scatter: + ep_rowmap = torch.empty( + (int(contiguous_m) + 1, 2), dtype=torch.int32, device=device + ) + ep_scatter_params = { + "gather_w": _gather_w_buf, + "tis": stage2_scatter.source_token_map, + "ep_rowmap": ep_rowmap, + "topk": int(topk), + "max_tok": int(stage2_scatter.max_tokens_per_rank), + "slot_stride": int(stage2_scatter.max_tokens_per_rank) * int(topk), + } _starts, psum, _ = contiguous_psum_remap( - _masked_m, topids_to_rows, E, max_m, tile_m, num_valid_routes=_ep_nvr + _masked_m, + topids_to_rows, + E, + max_m, + tile_m, + num_valid_routes=_ep_nvr, + ep_scatter_params=ep_scatter_params, ) psum = psum.to(torch.int32).contiguous() + # Turns the TDM GEMM2 epilogue into the fused P2P scatter-combine. + _ep_gemm2_kwargs = ( + { + "stage2_scatter": stage2_scatter, + "ep_destination_stride": ( + int(stage2_scatter.max_tokens_per_rank) * int(topk) + ), + "ep_row_map": ep_rowmap, + } + if enable_ep_scatter + else {} + ) out_is_f16 = 1 if (dtype == torch.float16 or dtype == dtypes.fp16) else 0 two_inter = 2 * inter_dim @@ -694,6 +735,7 @@ def _grouped_a8w4_tdm_moe( cluster_n=cluster_n, waves_per_tensor_tdm=waves_per_tensor_tdm, next_stage_prefetch=next_stage_prefetch, + **_ep_gemm2_kwargs, ) if kernel_bench_callable is not None: @@ -800,10 +842,18 @@ def _grouped_a8w4_tdm_moe( ) ) + if enable_ep_scatter: + # GEMM2 already P2P-wrote every route-weighted result into the peers' + # combine buffers, so this only satisfies the custom-op's return contract; + # MegaMoE ignores its contents and consumes comb_inp instead. + return torch.empty((token_num, model_dim), dtype=dtype, device=device) + moe_out = torch.empty((token_num, model_dim), dtype=dtype, device=device) if _is_ep: # Route kernel already produced gather weights (dropped routes zeroed); - # the dead-tail (tokens >= total_recv) is skipped via num_valid_tokens. + # the dead-tail (tokens >= total_recv) is skipped via num_valid_tokens, and + # non-local routes read through a zero-sized descriptor (DROPPED_ROUTE_ROW), + # so a token whose every route is remote reduces to zeros. gather_w = _gather_w_buf else: gather_w = ( @@ -818,7 +868,6 @@ def _grouped_a8w4_tdm_moe( out=moe_out, num_valid_tokens=(_ep_nvt if _is_ep else None), ) - os.environ["AITER_LAST_FUSED_MOE_IMPL"] = "grouped_a8w4_tdm" return moe_out @@ -850,8 +899,9 @@ def grouped_gemm_gfx1250_a8w4( num_local_tokens: torch.Tensor | None = None, situ_beta: float = 1.0, situ_linear_beta: float = 1.0, + stage2_scatter: Stage2ScatterContext | None = None, ): - """Grouped a8w4/a4w4 MoE on the felix TDM batched GEMM (gfx1250). + """Grouped a8w4/a4w4 MoE on the TDM batched GEMM (gfx1250). ``w1`` MUST be GUGU: gate/up row-interleaved ([g0,u0,g1,u1,...]), i.e. what ``moe_shuffle_weight(..., is_guinterleave=True, gate_up=True)`` produces. @@ -920,7 +970,6 @@ def _fmt(v): if os.environ.get("AITER_DISABLE_GROUPED_A8W4", "0") == "1": _grouped_dbg("AITER_DISABLE_GROUPED_A8W4 enabled; skip grouped mode") return None - os.environ["AITER_LAST_FUSED_MOE_IMPL"] = "default" _is_ep = expert_mask is not None if _is_ep: _grouped_dbg(f"EP enabled: expert_mask numel={expert_mask.numel()}, E={E}") @@ -1117,12 +1166,13 @@ def _tdm_env(name): data_format=data_format, expert_mask=expert_mask, num_local_tokens=num_local_tokens, + stage2_scatter=stage2_scatter, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, **_tdm_kw, ) - # Only the felix TDM grouped path is kept; the previous non-TDM grouped + # Only the TDM grouped path is kept; the previous non-TDM grouped # GEMM (gemm_mxscale_gfx1250 / moe_grouped_gemm_mxscale_gfx1250) was # removed. Anything the TDM path cannot serve falls back to the caller's # generic MoE via None. @@ -1272,6 +1322,16 @@ def contiguous_psum(masked_m: torch.Tensor, experts: int, tile_m: int): return starts, psum, contiguous_m_t +@functools.cache +def _get_compiled_contiguous_psum_remap_ep(): + """psum + remap fused with the gemm2 EP ep_rowmap build.""" + from aiter.ops.flydsl.kernels.moe_contiguous_psum import ( + build_moe_contiguous_psum_remap_ep_module, + ) + + return build_moe_contiguous_psum_remap_ep_module() + + def contiguous_psum_remap( masked_m: torch.Tensor, topids_to_rows: torch.Tensor, @@ -1279,8 +1339,14 @@ def contiguous_psum_remap( route_max_m: int, tile_m: int, num_valid_routes: torch.Tensor | None = None, + ep_scatter_params: dict | None = None, ): - """Tile-aligned psum and in-place masked-row -> contiguous-row remap.""" + """Tile-aligned psum and in-place masked-row -> contiguous-row remap. + + With ``ep_scatter_params`` (gather_w/tis/ep_rowmap/topk/max_tok/slot_stride) + the same pass also scatters the gemm2-fused EP row map, reusing the final row + it just computed. + """ device = masked_m.device experts = int(experts) masked_m_i32 = masked_m[:experts].to(torch.int32) @@ -1300,6 +1366,30 @@ def contiguous_psum_remap( .to(device=device, dtype=torch.int32) .contiguous() ) + if ep_scatter_params is not None: + launch = _get_compiled_contiguous_psum_remap_ep() + # Init ep_rowmap to (-1, 0) with one int64 fill (low i32 = -1, high = 0); + # stream-ordered before the launch, whose scatter overwrites the kept rows. + ep_scatter_params["ep_rowmap"].view(torch.int64).fill_(0xFFFFFFFF) + launch( + ptr_arg(masked_m_i32), + ptr_arg(topids_flat), + ptr_arg(starts), + ptr_arg(psum), + ptr_arg(contiguous_m_t), + experts, + int(route_max_m), + int(tile_m), + ptr_arg(num_valid_routes_i32), + ptr_arg(ep_scatter_params["gather_w"].reshape(-1)), + ptr_arg(ep_scatter_params["tis"].reshape(-1)), + ptr_arg(ep_scatter_params["ep_rowmap"].reshape(-1)), + int(ep_scatter_params["topk"]), + int(ep_scatter_params["max_tok"]), + int(ep_scatter_params["slot_stride"]), + stream=torch.cuda.current_stream(), + ) + return starts, psum, contiguous_m_t launch = _get_compiled_contiguous_psum_remap() launch( ptr_arg(masked_m_i32), @@ -1355,7 +1445,9 @@ def flydsl_moe_gather_reduce( """One-pass gather-reduce: out[t] = sum_k w[t,k] * grouped[topids_to_rows[t,k]]. ``gather_w`` may be f32 (native route weights, no host-side cast) or match - ``grouped_out``'s bf16/f16; the kernel accumulates in f32 either way. + ``grouped_out``'s bf16/f16; the kernel accumulates in f32 either way. Slots + holding ``moe_route_maps.DROPPED_ROUTE_ROW`` (EP routes with no grouped row) + contribute 0 without touching ``grouped_out``. """ if grouped_out.dim() == 4: split_k, E, max_m, model_dim = grouped_out.shape diff --git a/aiter/ops/flydsl/kernels/communication_ops_utils.py b/aiter/ops/flydsl/kernels/communication_ops_utils.py index 6f7588ecf8d..de7446f3cf3 100644 --- a/aiter/ops/flydsl/kernels/communication_ops_utils.py +++ b/aiter/ops/flydsl/kernels/communication_ops_utils.py @@ -19,7 +19,10 @@ import flydsl.expr as fx from flydsl._mlir import ir from flydsl._mlir.dialects import llvm as _llvm_d +from flydsl._mlir.dialects import rocdl as _rocdl_d +from flydsl.compiler.ast_rewriter import ASTRewriter from flydsl.expr import arith +from flydsl.expr.typing import T __all__ = [ "GeometryTuningTable", @@ -32,9 +35,18 @@ "fence_release", "fence_system_acquire", "fence_system_release", + "load_i32_acquire", + "load_i32_nt", + "load_i64_acquire", "load_i64_global", + "load_v4i32_nt", + "spin_until_eq_i32", + "spin_until_ge_i64", + "spin_until_gt_i32", "store_i32_system", "store_i64_global_system", + "traced", + "waitcnt_all", ] @@ -45,6 +57,96 @@ def _to_ptr_global(v): ).result +def _ptr_plus(base_i64, offset, elem_bytes): + """Global pointer for base + offset*elem_bytes.""" + addr = ( + fx.Int64(arith.unwrap(base_i64)) + fx.Int64(arith.unwrap(offset)) * elem_bytes + ) + return _to_ptr_global(addr) + + +def traced(fn): + """Run FlyDSL's AST rewriting over a helper that kernel bodies call. + + ``@flyc.kernel`` and ``@flyc.jit`` apply this same transform to their own + source but do not recurse into callees, so a helper that wants ``if`` / + ``while`` over traced values (rather than a host-side truthiness test) has + to opt in. + """ + return ASTRewriter.transform(fn) + + +def waitcnt_all(): + """Drain outstanding gfx12 load/store counters before a grid barrier.""" + _rocdl_d.s_wait_storecnt(0) + _rocdl_d.s_wait_loadcnt(0) + + +def load_i32_acquire(addr_i64): + """Volatile monotonic i32 load suitable for a spin-wait.""" + return _llvm_d.LoadOp( + T.i32, + _to_ptr_global(addr_i64), + alignment=4, + volatile_=True, + ordering=_llvm_d.AtomicOrdering.monotonic, + syncscope="one-as", + ).res + + +def load_i64_acquire(addr_i64): + """Volatile monotonic i64 load suitable for a spin-wait.""" + return _llvm_d.LoadOp( + T.i64, + _to_ptr_global(addr_i64), + alignment=8, + volatile_=True, + ordering=_llvm_d.AtomicOrdering.monotonic, + syncscope="one-as", + ).res + + +def load_i32_nt(base_i64, offset): + """Non-temporal global i32 load at base + offset*4.""" + return _llvm_d.LoadOp( + T.i32, _ptr_plus(base_i64, offset, 4), alignment=4, nontemporal=True + ).res + + +def load_v4i32_nt(base_i64, offset): + """Non-temporal global vector<4xi32> load at base + offset*4.""" + return _llvm_d.LoadOp( + T.i32x4, _ptr_plus(base_i64, offset, 4), alignment=4, nontemporal=True + ).res + + +@traced +def spin_until_ge_i64(addr_i64, val): + """Spin until a monotonic cross-device i64 flag is at least ``val``.""" + cur = fx.Int64(load_i64_acquire(addr_i64)) + while cur < fx.Int64(val): + cur = fx.Int64(load_i64_acquire(addr_i64)) + return cur + + +@traced +def spin_until_eq_i32(addr_i64, val): + """Spin until an i32 flag equals ``val``.""" + cur = fx.Int32(load_i32_acquire(addr_i64)) + while cur != fx.Int32(val): + cur = fx.Int32(load_i32_acquire(addr_i64)) + return cur + + +@traced +def spin_until_gt_i32(addr_i64, val): + """Spin until an i32 flag exceeds ``val`` and return the observed value.""" + cur = fx.Int32(load_i32_acquire(addr_i64)) + while cur <= fx.Int32(val): + cur = fx.Int32(load_i32_acquire(addr_i64)) + return cur + + def store_i32_system(addr_i64, offset, val): """System-scope release i32 store at ``addr_i64 + offset*4``.""" base = arith.unwrap(addr_i64) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py new file mode 100644 index 00000000000..3526be5a7d8 --- /dev/null +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""gfx1250 MegaMoE with a single lazy public entry point.""" + +import importlib + +__all__ = ["MegaMoEGfx1250"] + + +def __getattr__(name): + if name != "MegaMoEGfx1250": + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(f"{__name__}.mega_moe"), name) + + +def __dir__(): + return sorted(list(globals()) + __all__) diff --git a/aiter/ops/flydsl/kernels/mega_moe_gfx1250/combine.py b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/combine.py new file mode 100644 index 00000000000..0e77a3ef45a --- /dev/null +++ b/aiter/ops/flydsl/kernels/mega_moe_gfx1250/combine.py @@ -0,0 +1,278 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Barrier-wait and top-k reduce kernel for gfx1250 Stage2-fused MegaMoE.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +import mori.cco.device.flydsl as cco +from flydsl.expr import arith, range_constexpr +from flydsl.expr.typing import Int32, Int64, T + +from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops +from aiter.ops.flydsl.kernels import vector +from aiter.ops.flydsl.kernels.buffer_ops import ( + buffer_load, + buffer_store, + create_buffer_resource_from_addr, +) + +from .config import ( + _LANE_MASK as LANE_MASK, +) +from .config import ( + _LOG2_WAVE_SIZE as LOG2_WAVE, +) +from .config import ( + _WAVE_SIZE as WAVE, +) + +# The cross-device xdb barrier is combine's own, not shared with dispatch's: it +# waits on monotonic per-rank phase slots, while dispatch gates on a grid-wide +# disp_bar count and then hands each peer its recv_num. Different state, so nothing +# to factor out. + + +def _V2BF16(): + return T.vec(2, T.bf16) + + +def _V2F32(): + return T.vec(2, T.f32) + + +def _V1I32(): + return T.vec(1, T.i32) + + +def _bf16_accum_funcs(): + def to_accum(i32_scalar): + return vector.bitcast( + _V2BF16(), vector.from_elements(_V1I32(), [i32_scalar]) + ).extf(_V2F32()) + + def from_accum(acc): + return vector.extract( + vector.bitcast(_V1I32(), acc.truncf(_V2BF16())), static_position=[0] + ) + + def zero_accum(): + return to_accum(arith.constant(0)) + + return to_accum, from_accum, zero_accum + + +def _make_combine_fused_sync( + *, + rank, + npes, + off_xdb_mem, +): + """Stage A: wait until every peer's gemm2 P2P writes into comb_inp land. + + Launch this before the reduce kernel; its retirement is stream-ordered, so + the reduce needs no in-kernel fence and its grid stays unconstrained. + + One thread per peer pushes the phase and polls that peer's local slot. The + block is rounded up to whole waves, so rack-scale domains larger than one + wave are covered without a cross-wave dependency or barrier -- the kernel's + cost is the peer wait, not thread count. + + Being a rendezvous, it also fences the next dispatch off the regions this + forward still reads, which is what lets ``Routing.source_token_map`` hand + gemm2 a live view of ``recv_to_src_token`` instead of a copy. + """ + + sync_block_size = ((npes + WAVE - 1) // WAVE) * WAVE + + @flyc.kernel(known_block_size=[sync_block_size, 1, 1]) + def ep_combine_fused_sync( + arena: Int64, + addr_xdb_flag: Int64, + my_lsa_rank: Int32, + ): + tid = fx.thread_idx.x + window = cco.Window(arena) + rsrc_xdb_flag = create_buffer_resource_from_addr(addr_xdb_flag) + phase = fx.Int64(buffer_load(rsrc_xdb_flag, 0, vec_width=1, dtype=T.i64)) + # push this call's phase to every peer's shared xdb slot [rank] + if tid < npes: + xdb_remote = fx.Int64(window.lsa_ptr(tid, off_xdb_mem)) + fx.Int64( + rank + ) * fx.Int64(8) + comm_ops.store_i64_global_system(xdb_remote, phase) + # advance the counter for the next call (single writer, no atomic) + if tid == 0: + buffer_store(phase + arith.constant(1, type=T.i64), rsrc_xdb_flag, 0) + # `>=` not `==`: a faster peer can lap us and overwrite its monotonic push + # with a higher call count before we read it. + if tid < npes: + xdb_peer_slot = fx.Int64( + window.lsa_ptr(my_lsa_rank, off_xdb_mem) + ) + fx.Int64(tid) * fx.Int64(8) + comm_ops.spin_until_ge_i64(xdb_peer_slot, phase) + + @flyc.jit + def run( + arena: Int64, + addr_xdb_flag: Int64, + my_lsa_rank: Int32, + stream=fx.Stream(None), # noqa: B008 + ): + ep_combine_fused_sync(arena, addr_xdb_flag, my_lsa_rank).launch( + grid=(1, 1, 1), + block=[sync_block_size, 1, 1], + stream=stream, + ) + + return run + + +def _make_combine_fused_reduce( + *, + experts_per_token, + hidden_dim, + block_num, + warp_num_per_block, + slot_stride_nbytes=None, +): + """Stage B of the GEMM2-fused scatter combine: the per-token topk sum. + + gemm2 has already P2P-written each token's WEIGHTED per-expert result into + this rank's comb_inp[origin_lid*topk + k] (one contiguous topk-block per + token), so this is an unweighted sum: out[t] = sum_{k