From 58c18e0fce61252f8c7be1ea965d9dd046f95018 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=A0=C3=A7kawhil-amd?= Date: Tue, 27 Oct 2026 11:17:32 +0000 Subject: [PATCH 1/9] example(cco): MoE GEMM2 fused with an SDMA combine (gfx950) GEMM2 writes its rows into a local send buffer grouped by origin rank (stable, expert-major inside a run), and cco SDMA ships each rank's run as contiguous puts, so the transfer does not occupy CUs. Three kernels on one stream (the P2P-scatter baseline has two): 1. rank_sort_plan: persistent grid, count -> scan -> assign with grid barriers; a stable send-buffer index per row. 2. gemm2: writes the send buffer; extra comm waves (one per peer) wait for finished chunks of sort tiles and put them on their own SDMA queue while the remaining tiles still compute. 3. combine: persistent grid; waits for the peers' flags, builds the slot -> row table, reduces the top-k rows, and tells the peers their data has been consumed. Reuse of the landing buffers across calls is guarded by per-peer "consumed" flags, so a fast rank cannot overwrite a slow one's data. main.py checks the result against aiter's MegaMoEV2 stage2+combine on the same stage1 state; bench.py compares latency against it per token count. Compile-checked only; not yet run on GPUs. Co-Authored-By: Claude Code --- examples/cco/README.md | 14 + .../python/11_flydsl_mega_moe_sdma/bench.py | 301 ++++++++ .../python/11_flydsl_mega_moe_sdma/main.py | 263 +++++++ .../moe_gemm2_localsort.py | 549 ++++++++++++++ .../11_flydsl_mega_moe_sdma/sorted_combine.py | 667 ++++++++++++++++++ 5 files changed, 1794 insertions(+) create mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/bench.py create mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/main.py create mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py create mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py diff --git a/examples/cco/README.md b/examples/cco/README.md index 09bc0100d4..1cfaf60c86 100644 --- a/examples/cco/README.md +++ b/examples/cco/README.md @@ -22,6 +22,7 @@ examples/cco/ | `python/08_torch_symm_import` | py | import a `torch.symm_mem` tensor into the flat LSA space, no copy | | `python/09_torch_symm_lsa_all2all` | py | mori's own torch symm-mem backend + Triton `lsa_ptr` all-to-all | | `python/10_flydsl_sdma_allgather_async` | py | FlyDSL SDMA all-gather split into `start_async` / `wait_async`, overlapped with a GEMM | +| `python/11_flydsl_mega_moe_sdma` | py | MoE GEMM2 (gfx950) writing a rank-sorted send buffer, shipped by SDMA `put`s from comm waves (one per peer) inside the GEMM kernel (needs aiter) | | `ir/test_triton_cco.py` | py | Triton DevComm queries and LSA/SDMA/GDA device APIs | | `cpp/01_lsa_put.cpp` | c++ | intra-node LSA put (includes only `cco.hpp`) | | `cpp/02_gda_put.cpp` | c++ | GPU-initiated RDMA put + signal/wait (includes only `cco_scale_out.hpp`) | @@ -90,6 +91,19 @@ MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ examples/cco/python/10_flydsl_sdma_allgather_async/main.py ``` +Example 11 also uses `torchrun` and SDMA, and additionally needs aiter (it reuses +`MegaMoEV2` for the reference and the GEMM2 mainloop) on a gfx950 node with 8 GPUs: + +```bash +MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ + examples/cco/python/11_flydsl_mega_moe_sdma/main.py --network glm52 --bs 128 + +# latency vs aiter's MegaMoEV2 (mori.shmem P2P scatter), sweeping token counts +MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ + examples/cco/python/11_flydsl_mega_moe_sdma/bench.py --network glm52 \ + --bs-list 16,128,1024,4096 --chunk-tiles 2,4,8,16 +``` + `MORI_CCO_GDA_CONN=full` is required for GDA (03, 06) when both ranks share one node; LSA examples (02, 04, 05) ignore it. Each example prints `SUCCESS` on pass. diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py new file mode 100644 index 0000000000..c4c8d3bb1c --- /dev/null +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +# 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. +"""Latency: aiter MegaMoEV2 stage2+combine vs the rank-sorted SDMA path (gfx950) + +Baseline: aiter's ``MegaMoEV2`` (mori.shmem symmetric heap). Its fused GEMM2 scatters +every output row into the origin rank's ``comb_inp`` with P2P stores from the compute +units, then a combine kernel runs a cross-rank barrier and the top-k reduce. No cco, +no LSA, no copy engine. + +New path (``sorted_combine.SortedCombine``), three kernels: a plan kernel, GEMM2 writing +a rank-sorted local send buffer with comm waves inside it that put finished chunks with +cco SDMA while the tiles are still computing, and a combine kernel that waits for the +peers and does an indexed top-k reduce. + +Per token count it reports, as mean/max over ranks: + + stage2+combine baseline vs the new path (the part that is replaced) + end to end stage1 + that (dispatch and stage1 are shared, so e2e shows + how much of the layer the replaced part is) + breakdown per-stage time of one eager call (rank 0), to see where the + new path spends its time + +CUDA Graph timing by default (what the production path uses); ``--no-graph`` times +eager launches. + + MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 bench.py \\ + --network glm52 --bs-list 16,128,1024,4096 +""" + +import argparse + +from main import NETWORKS, _make_inputs, _next_pow2, _quantize_weights, _setup + +import torch +import torch.distributed as dist + +import mori.shmem as ms + +from aiter.ops.flydsl.kernels.mega_moe import MegaMoEV2 + +from sorted_combine import SortedCombine + + +def _over_ranks(local_ms): + """(mean, max) of a per-rank time over all ranks.""" + t = torch.tensor([local_ms], dtype=torch.float64) + mx = t.clone() + dist.all_reduce(t, op=dist.ReduceOp.SUM) + dist.all_reduce(mx, op=dist.ReduceOp.MAX) + return float(t.item()) / dist.get_world_size(), float(mx.item()) + + +def _sync(): + torch.cuda.synchronize() + dist.barrier() + + +def _time(fn, iters, graph): + """Mean/max-over-ranks ms per call. Falls back to eager if the graph cannot be captured.""" + _sync() + fn() + _sync() + replay = fn + if graph: + try: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch.cuda.Stream()): + fn() + replay = g.replay + except Exception as exc: # noqa: BLE001 + if dist.get_rank() == 0: + print( + f" (graph capture failed, timing eagerly: {type(exc).__name__}: {exc})", + flush=True, + ) + torch.cuda.synchronize() + for _ in range(5): + replay() + _sync() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + start.record() + for _ in range(iters): + replay() + end.record() + torch.cuda.synchronize() + return _over_ranks(start.elapsed_time(end) / iters) + + +def _breakdown(sc, bs, config): + """Per-stage ms of one eager call, averaged over ranks.""" + sc.run(bs, config) # warm + _sync() + marks = [] + sc.run(bs, config, profile=marks) + torch.cuda.synchronize() + names = [n for n, _ in marks[1:]] + ms_ = [marks[i][1].elapsed_time(marks[i + 1][1]) for i in range(len(marks) - 1)] + t = torch.tensor(ms_, dtype=torch.float64) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + t /= dist.get_world_size() + return dict(zip(names, t.tolist())) + + +def _rel_l2(out, ref): + diff = out.float() - ref.float() + local = torch.tensor( + [float(diff.norm() ** 2), float(ref.float().norm() ** 2)], dtype=torch.float64 + ) + dist.all_reduce(local) + return float((local[0] / local[1].clamp_min(1e-30)).sqrt()) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--network", choices=NETWORKS, default="glm52") + parser.add_argument("--bs-list", default="16,128,1024,4096", help="tokens per rank") + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--rtol", type=float, default=0.05) + parser.add_argument( + "--chunk-tiles", default="4", help="fused chunk sizes to sweep, e.g. 2,4,8,16" + ) + parser.add_argument("--queues", type=int, default=8) + parser.add_argument("--no-graph", action="store_true") + args = parser.parse_args() + + batch_sizes = [int(v) for v in args.bs_list.split(",")] + chunk_list = [int(v) for v in args.chunk_tiles.split(",")] + graph = not args.no_graph + + rank, world, device = _setup() + rows = [] # (bs, label, s2_mean, s2_max, e2e_mean, e2e_max, relL2) + try: + net = NETWORKS[args.network] + w1, w1s, w2, w2s = _quantize_weights( + net["model_dim"], + net["inter_dim"], + net["experts"] // world, + rank, + args.seed, + device, + ) + max_bs = max(batch_sizes) + x_all, wts_all, ids_all = _make_inputs( + max_bs, + net["model_dim"], + net["experts"], + net["topk"], + rank, + args.seed, + device, + ) + # One operator and one SortedCombine for every token count: building more would + # re-create symmetric heaps and SDMA queues. + moe = MegaMoEV2( + rank=rank, + world_size=world, + quant="a8w4", + w1=w1, + w1_scale=w1s, + w2=w2, + w2_scale=w2s, + max_tok_per_rank=max(16, _next_pow2(max_bs)), + fanout_masks=(), + **net, + ) + sc = SortedCombine(moe, queues=args.queues) + + for bs in batch_sizes: + x, wts, ids = ( + x_all[:bs].contiguous(), + wts_all[:bs].contiguous(), + ids_all[:bs].contiguous(), + ) + ref = moe(x, wts, ids)[:bs].clone() + torch.cuda.synchronize() + config = moe._active_config + if rank == 0: + print( + f"\n=== bs={bs} tokens/rank p2p_quant={config.p2p_quant} " + f"BM={moe._g2_active_block_m} SBM={moe._s1_active_tile_m} " + f"aligned_pair={config.stage2.aligned_pair} ===", + flush=True, + ) + if config.stage2.aligned_pair: + if rank == 0: + print( + " aligned-pair stage2 is not covered by SortedCombine; skipped", + flush=True, + ) + continue + + x_q, scale = moe.quantize(x) + + def stage1(): + moe._run_fused_stage1(x_q, wts, scale, ids) + + def base_s2(): + moe._run_stage2(bs, None, True, config) + + def base_e2e(): + moe(x, wts, ids) + + s1 = _time(stage1, args.iters, graph) + stage1() + _sync() + b_s2 = _time(base_s2, args.iters, graph) + b_e2e = _time(base_e2e, args.iters, graph) + rows.append((bs, "aiter MegaMoEV2 (P2P scatter)", *b_s2, *b_e2e, None)) + if rank == 0: + print( + f" stage1 alone {s1[0]:9.4f} / {s1[1]:9.4f} ms (mean/max)", + flush=True, + ) + print( + f" baseline stage2+combine {b_s2[0]:9.4f} / {b_s2[1]:9.4f} ms", + flush=True, + ) + print( + f" baseline end to end {b_e2e[0]:9.4f} / {b_e2e[1]:9.4f} ms", + flush=True, + ) + + stage1() + _sync() + for ct in chunk_list: + sc.set_chunk_tiles(ct) + label = f"SDMA 3-kernel CT={ct}" + out = sc.run(bs, config).clone() + torch.cuda.synchronize() + rel = _rel_l2(out, ref) + n_s2 = _time(lambda: sc.run(bs, config), args.iters, graph) + + def e2e(): + stage1() + sc.run(bs, config) + + n_e2e = _time(e2e, args.iters, graph) + rows.append((bs, label, *n_s2, *n_e2e, rel)) + br = _breakdown(sc, bs, config) + if rank == 0: + parts = " ".join(f"{k}={v:.3f}" for k, v in br.items()) + print( + f" {label:<22} stage2+combine {n_s2[0]:9.4f} / {n_s2[1]:9.4f} ms " + f"e2e {n_e2e[0]:9.4f} / {n_e2e[1]:9.4f} ms relL2={rel:.2e}", + flush=True, + ) + print(f" breakdown (eager, ms): {parts}", flush=True) + if rel >= args.rtol: + raise AssertionError( + f"bs={bs} {label}: relL2={rel:.3e} exceeds {args.rtol}" + ) + + if rank == 0: + print( + "\n| bs | path | stage2+combine mean ms | max ms | vs baseline | " + "e2e mean ms | vs baseline | relL2 |" + ) + print("|---:|---|---:|---:|---:|---:|---:|---:|") + base = {} + for bs, label, s2m, s2x, e2m, e2x, rel in rows: + if rel is None: + base[bs] = (s2m, e2m) + b_s2, b_e2e = base.get(bs, (s2m, e2m)) + print( + f"| {bs} | {label} | {s2m:.4f} | {s2x:.4f} | {b_s2 / s2m:.2f}x | " + f"{e2m:.4f} | {b_e2e / e2m:.2f}x | {'-' if rel is None else f'{rel:.1e}'} |" + ) + print("\nSUCCESS", flush=True) + sc.close() + finally: + try: + ms.shmem_finalize() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py new file mode 100644 index 0000000000..5458e1eacd --- /dev/null +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +# 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. +"""CCO Example 11 — MoE GEMM2 fused with an SDMA combine (gfx950) + +MegaMoE's GEMM2 normally scatters every output row straight into the origin +rank's ``comb_inp`` with P2P stores from the compute units. Here GEMM2 writes its +rows into a *local* send buffer grouped by origin rank, and cco's SDMA copy engine +ships each rank's run -- one contiguous ``put`` per peer -- so the transfer does +not occupy CUs: + + S = [ origin rank 0: rows | origin rank 1: rows | ... ] (expert-major inside a run) + +Three kernels on one stream (see ``sorted_combine.py``; the P2P baseline has two): + + 1. rank_sort_plan count -> scan -> assign: a stable send-buffer index per row + 2. gemm2 writes S[out_row] (``moe_gemm2_localsort.py``); extra comm waves, + one per peer, wait for a chunk of sort tiles to be written (a + per-chunk counter the GEMM workgroups bump after their stores) and + put that chunk for their peer on their own SDMA queue while the + remaining tiles are still computing + 3. combine wait for the peers' flags, build slot -> row, top-k reduce + +The example runs MegaMoEV2 (aiter) once to get the reference output and the +stage1/dispatch state, then redoes stage2 + transfer + reduce on the same state and +compares. + +Needs aiter (MegaMoEV2 and the GEMM2 mainloop), FlyDSL, and mori built with SDMA: + + MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 main.py +""" + +import argparse +import os + +os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "20G") +os.environ.setdefault("MORI_ENABLE_SDMA", "1") + +import torch # must be imported before anything that pulls in mori.cco +import torch.distributed as dist + +import mori.shmem as ms + +import aiter +from aiter import dtypes +from aiter.ops.flydsl.kernels.mega_moe import MegaMoEV2 +from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4 + +from sorted_combine import SortedCombine + +NETWORKS = { + "glm52": dict( + model_dim=6144, inter_dim=2048, experts=256, topk=8, swiglu_limit=0.0 + ), + "v4_pro": dict( + model_dim=7168, inter_dim=3072, experts=384, topk=6, swiglu_limit=10.0 + ), +} + + +def _next_pow2(v): + return 1 << (int(v) - 1).bit_length() + + +def _make_inputs(tokens, model_dim, experts, topk, rank, seed, device): + g = torch.Generator(device=device).manual_seed(seed + rank) + x = torch.randn( + (tokens, model_dim), dtype=torch.bfloat16, device=device, generator=g + ) + scores = torch.randn( + (tokens, experts), dtype=torch.float32, device=device, generator=g + ) + values, ids = torch.topk(scores, topk, dim=-1) + return ( + x.contiguous(), + values.softmax(dim=-1).contiguous(), + ids.to(torch.int32).contiguous(), + ) + + +def _quantize_weights(model_dim, inter_dim, local_experts, rank, seed, device): + g = torch.Generator(device=device).manual_seed(seed + 1000 + rank) + quantize = aiter.get_torch_quant(aiter.QuantType.per_1x32) + + w1 = torch.randn( + (local_experts, 2 * inter_dim, model_dim), + dtype=torch.bfloat16, + device=device, + generator=g, + ).mul_(model_dim**-0.25) + w1_q, w1_scale = quantize(w1, quant_dtype=dtypes.fp4x2) + w1_q = w1_q.view(local_experts, 2 * inter_dim, model_dim // 2) + del w1 + w2 = torch.randn( + (local_experts, model_dim, inter_dim), + dtype=torch.bfloat16, + device=device, + generator=g, + ).mul_(inter_dim**-0.25) + w2_q, w2_scale = quantize(w2, quant_dtype=dtypes.fp4x2) + w2_q = w2_q.view(local_experts, model_dim, inter_dim // 2) + del w2 + torch.cuda.empty_cache() + return ( + shuffle_weight_a16w4(w1_q, 16, True).contiguous(), + shuffle_scale_a16w4(w1_scale, local_experts, True).contiguous(), + shuffle_weight_a16w4(w2_q, 16, False).contiguous(), + shuffle_scale_a16w4(w2_scale, local_experts, False).contiguous(), + ) + + +def _setup(): + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + local = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local) + # gloo only: no NCCL, so this also runs where /dev/shm is too small for it. + dist.init_process_group("gloo") + import torch._C._distributed_c10d as c10d + + c10d._register_process_group("default", dist.group.WORLD) + ms.shmem_torch_process_group_init("default") + return rank, world, torch.device("cuda", local) + + +def _timeit(fn, iters): + torch.cuda.synchronize() + dist.barrier() + fn() + torch.cuda.synchronize() + dist.barrier() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--network", choices=NETWORKS, default="glm52") + parser.add_argument("--bs", type=int, default=128, help="tokens per rank") + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--rtol", type=float, default=0.05) + parser.add_argument( + "--chunk-tiles", + type=int, + default=4, + help="sort tiles per shipped chunk", + ) + parser.add_argument("--queues", type=int, default=8) + args = parser.parse_args() + + rank, world, device = _setup() + try: + net = NETWORKS[args.network] + w1, w1s, w2, w2s = _quantize_weights( + net["model_dim"], + net["inter_dim"], + net["experts"] // world, + rank, + args.seed, + device, + ) + x, wts, ids = _make_inputs( + args.bs, + net["model_dim"], + net["experts"], + net["topk"], + rank, + args.seed, + device, + ) + moe = MegaMoEV2( + rank=rank, + world_size=world, + quant="a8w4", + w1=w1, + w1_scale=w1s, + w2=w2, + w2_scale=w2s, + max_tok_per_rank=max(16, _next_pow2(args.bs)), + fanout_masks=(), + **net, + ) + ref = moe(x, wts, ids)[: args.bs].clone() + torch.cuda.synchronize() + config = moe._active_config + if config.stage2.aligned_pair: + raise SystemExit( + "this config selects aligned-pair stage2, which SortedCombine does not cover" + ) + if rank == 0: + print( + f"reference: p2p_quant={config.p2p_quant} BM={moe._g2_active_block_m} " + f"SBM={moe._s1_active_tile_m}", + flush=True, + ) + + ref_ms = _timeit( + lambda: moe._run_stage2(args.bs, None, True, config), args.iters + ) + failures = [] + sc = SortedCombine(moe, queues=args.queues, chunk_tiles=args.chunk_tiles) + worst = 0.0 + # Repeated calls with no host sync between them also exercise the epoch counter, + # the flags and the "consumed" handshake that stops a fast rank overrunning a slow one. + for _ in range(4): + out = sc.run(args.bs, config).clone() + torch.cuda.synchronize() + diff = out.float() - ref.float() + worst = max(worst, float(diff.norm() / ref.float().norm().clamp_min(1e-9))) + t = torch.tensor([worst], dtype=torch.float64) + dist.all_reduce(t, op=dist.ReduceOp.MAX) + worst = float(t.item()) + ms_per = _timeit(lambda: sc.run(args.bs, config), args.iters) + ok = worst < args.rtol + if not ok: + failures.append("sdma") + if rank == 0: + print( + f"relL2={worst:.3e} {'PASS' if ok else 'FAIL'} " + f"time={ms_per:.4f}ms (reference stage2+combine {ref_ms:.4f}ms)", + flush=True, + ) + sc.close() + if failures: + raise AssertionError(f"relL2 above {args.rtol} for {failures}") + if rank == 0: + print("SUCCESS", flush=True) + finally: + try: + ms.shmem_finalize() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py new file mode 100644 index 0000000000..5f219fba88 --- /dev/null +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py @@ -0,0 +1,549 @@ +# 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. +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +# ruff: noqa: B023, I001 +"""GEMM2 whose epilogue writes weighted bf16 rows into a local, rank-sorted send buffer. + +Same GEMM2 mainloop and input layout as ``mega_moe_stage2``; only the output side +differs. Instead of scattering each row into the origin rank's ``comb_inp`` it +writes the row to ``send[out_row[row]]``, where ``out_row`` was assigned once per +row beforehand (``sorted_combine.build_rank_sort_plan``) so every N +tile of a row agrees on it. A negative ``out_row`` drops the row. + +With ``comm`` the kernel also carries ``comm.comm_wgs`` extra leading workgroups that +ship finished chunks of the send buffer to the peers with cco SDMA while the GEMM +workgroups are still computing. A GEMM unit signals ``cdone[chunk]`` after its stores. +Each comm WAVE owns its own peers (peer ``p`` belongs to wave ``p % n_comm_waves``): it +spins on the chunk counter itself, then puts that chunk's rows for its peers on its own +SDMA queue, so peers neither wait for each other's spin nor share a wave's issue slot. +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr.typing import Int8, T +from flydsl.runtime.device import get_rocm_arch + +from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops +from aiter.ops.flydsl.kernels.buffer_ops import ( + buffer_load, + create_buffer_resource_from_addr, +) +from aiter.ops.flydsl.kernels.mxfp4_gemm_common import ( + global_typed_ptr, + lds_typed_ptr, + lds_vec_load, +) +from aiter.ops.flydsl.kernels.tensor_shim import ( + _run_compiled, + buf_copy_store, + ptr_buf_tensor, +) + +from aiter.ops.flydsl.kernels.mega_moe.gemm2 import ( + _resolve_g2_knobs, + _spart_output_tile_index, + gemm2_compute_v2, + issue_a_load_lds_dt, + kStages, +) +from sorted_combine import ( + ArenaLayout, + sdma_finish_peer, + sdma_put_run, + sdma_wait_consumed, +) +from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage2 import ( + _BUFFER_OFFSET_ABI_BYTES, + _stage2_lds_bytes, +) + + +@dataclass(frozen=True) +class CommSpec: + """Compile-time description of the fused SDMA comm workgroups. + + ``comm_wgs`` workgroups of 4 waves each; one wave per peer needs ``comm_wgs * 4 >= npes``. + """ + + layout: ArenaLayout + rank: int + queues: int + chunk_tiles: int + comm_wgs: int = 1 + + +# fmt: off +def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, N_OUT, BM, BN, + send_nbytes, lds_outrow_off, lds_weight_off, g2_bf16_lds=False, scatter_vec=8): +# fmt: on + """CShuffle one GEMM2 tile into weighted BF16 rows and store them at their send-buffer index.""" + kMChunks = BM // 16 + numAccN = (BN // 4) // 16 + wave_n = BN // 4 + lane_div_16 = lane // 16 + lane_mod_16 = lane % 16 + lds_base_fptr = lds_typed_ptr(lds_acc_base, T.f32) + lds_base_bf16 = ( + lds_typed_ptr(lds_acc_base, T.bf16, align=2) + if const_expr(g2_bf16_lds) + else None + ) + row_nbytes = N_OUT * 2 + + fx.barrier() + + for i in range_constexpr(kMChunks): + row_base = fx.Int32(i * 16) + lane_div_16 * 4 + if const_expr(g2_bf16_lds): + w_row = [ + fx.ptr_load( + lds_typed_ptr( + fx.Int32(lds_weight_off) + (row_base + v) * fx.Int32(4), + T.f32, + align=4, + ) + ) + for v in range_constexpr(4) + ] + for J in range_constexpr(numAccN): + col = wave * fx.Int32(wave_n) + J * 16 + lane_mod_16 + vec = fx.Vector(accm[i][J]) + for v in range_constexpr(4): + idx = (row_base + v) * BN + col + if const_expr(g2_bf16_lds): + lds_base_bf16[idx] = fx.BFloat16(fx.Float32(vec[v]) * fx.Float32(w_row[v])) + else: + lds_base_fptr[idx] = fx.Float32(vec[v]) + + fx.barrier() + + output_buf = ptr_buf_tensor( + send_base, + fx.BFloat16, + unit_elems=scatter_vec, + num_records_bytes=send_nbytes, + ) + # Inactive lanes read safe LDS and issue a bounded OOB store. + active = lane < fx.Int32(BN // scatter_vec) + col = active.select(lane * fx.Int32(scatter_vec), fx.Int32(0)) + + for row_iter in range_constexpr(BM // 4): + row = wave + fx.Int32(row_iter * 4) + row_byte_off = row * fx.Int32(4) + out_row = fx.ptr_load( + lds_typed_ptr(fx.Int32(lds_outrow_off) + row_byte_off, T.i32, align=4) + ) + out_row = rocdl.readfirstlane(T.i32, out_row.ir_value()) + idx0 = row * fx.Int32(BN) + col + if const_expr(g2_bf16_lds): + pk = fx.Vector( + lds_vec_load( + lds_acc_base, + idx0 * fx.Int32(2), + fx.Vector.make_type(scatter_vec, fx.BFloat16), + fx.BFloat16, + align=16, + ) + ) + else: + weight = fx.ptr_load( + lds_typed_ptr(fx.Int32(lds_weight_off) + row_byte_off, T.f32, align=4) + ) + weight = rocdl.readfirstlane(T.f32, weight.ir_value()) + v8 = fx.Vector( + lds_vec_load( + lds_acc_base, + idx0 * fx.Int32(4), + fx.Vector.make_type(scatter_vec, fx.Float32), + fx.Float32, + align=16, + ) + ) + pk = fx.Vector.from_elements( + [v8[i] * weight for i in range_constexpr(scatter_vec)], + fx.Float32, + ).to(fx.BFloat16) + row_off = out_row * fx.Int32(row_nbytes) + n_block_idx * fx.Int32(BN * 2) + off = ((out_row >= fx.Int32(0)) & active).select( + row_off + col * fx.Int32(2), fx.Int32(send_nbytes) + ) + buf_copy_store( + output_buf, + off // fx.Int32(scatter_vec * 2), + pk, + fx.BFloat16, + unit_elems=scatter_vec, + cache_modifier=0, + ) + + +# fmt: off +def compile_moe_gemm2_localsort(*, model_dim: int, inter_dim: int, experts: int, rank: int, + send_nbytes: int, BM: int = 32, BN: int = 256, BK: int = 256, use_nt: bool = True, + HIDDEN_MAX: int = 8192, INTER_MAX: int = 8192, a_dtype: str = "fp8", SBM: int | None = None, + persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, + g2_ascale_pf=None, g2_spart=None, persist_strided: bool = False, g2_bf16_lds: bool = False, + skew_cu: int = 0, scatter_vec: int = 8, comm=None): +# fmt: on + """Compile GEMM2 with the rank-sorted local-output epilogue (optionally fused with SDMA).""" + arch = str(get_rocm_arch() or "") + if not arch.startswith("gfx95"): + raise RuntimeError(f"MegaMoE v2 stage2 requires CDNA4 (gfx95x), got {arch or 'unknown'}") + assert model_dim % BN == 0 and HIDDEN_MAX % BN == 0 + assert INTER_MAX % BK == 0, f"INTER_MAX must be a multiple of {BK}" + if BM not in (16, 32, 64, 128): + raise AssertionError(f"BM must be in {{16,32,64,128}}, got {BM}") + SBM = BM if SBM is None else int(SBM) + if SBM % BM != 0: + raise AssertionError(f"SBM ({SBM}) must be a multiple of BM ({BM})") + if scatter_vec not in (8, 16): + raise ValueError("Stage2 scatter_vec must be 8 or 16") + if BN % scatter_vec: + raise ValueError("Stage2 scatter_vec must divide BN") + if a_dtype not in ("fp4", "fp8"): + raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") + if persist and cu_num <= 0: + raise AssertionError(f"persist=True requires cu_num>0, got {cu_num}") + if skew_cu and (not persist or not 0 < skew_cu <= cu_num): + raise AssertionError(f"skew_cu={skew_cu} requires persist=True and 0= 1 and chunk_tiles >= 1") + N_OUT = model_dim + g2_bhoist, g2_ascale_pf, g2_spart, g2_group_num, g2_m01, _g2_bf16_lds = _resolve_g2_knobs( + g2_bhoist, g2_ascale_pf, g2_spart, False, False + ) + is_f8 = a_dtype == "fp8" + aStages = kStages + 1 + KH_TILE_A = BK // (1 if is_f8 else 2) + compute_lds_bytes = _stage2_lds_bytes(BM, BN, BK, a_dtype, aStages, g2_bf16_lds) + lds_outrow_off = compute_lds_bytes + lds_weight_off = lds_outrow_off + BM * 4 + lds_bytes = lds_weight_off + BM * 4 + if lds_bytes > 160 * 1024: + raise ValueError(f"Stage2 LDS use {lds_bytes} exceeds 160 KiB") + _expert_offset = rank * experts + + @fx.struct + class SharedStorage: + buf: fx.Array[Int8, lds_bytes, 16] + + kernel_name = ( + f"megamoe_stage2_localsort_t{BM}x{BN}x{BK}" + f"_sbm{SBM}_{a_dtype}_nt{int(use_nt)}" + f"_p{int(persist)}cu{cu_num}s{int(persist_strided)}_pad{int(has_pad)}" + f"_sk{skew_cu}" + f"_bh{int(g2_bhoist)}apf{int(g2_ascale_pf)}sp{g2_group_num}x{g2_m01}" + f"_bf16lds{int(g2_bf16_lds)}_sv{scatter_vec}" + + (f"_fuse{comm.chunk_tiles}q{comm.queues}r{comm.rank}" if fuse else "") + ) + + # fmt: off + @flyc.kernel(name=kernel_name, known_block_size=[256, 1, 1]) + def kernel_localsort(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, + arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, + arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, + arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, + arg_cursor: fx.Int64, i32_max_m_blocks: fx.Int32, i32_inter: fx.Int32, i32_hidden: fx.Int32, + i32_kpad: fx.Int32, i32_npad: fx.Int32): + # fmt: on + tx_i32 = fx.thread_idx.x + bid = fx.block_idx.x + lane = tx_i32 % fx.Int32(64) + wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_base_i32 = fx.Int32(fx.ptrtoint(lds.buf.ptr)) + + num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) + k_bytes = fx.Int32(i32_inter) // fx.Int32(1 if is_f8 else 2) + trb_buf = ptr_buf_tensor(arg_trb, fx.Int32) + sweights_buf = ptr_buf_tensor(arg_sweights, fx.Float32) + outrow_buf = ptr_buf_tensor(arg_out_row, fx.Int32) + + def issue_all_a_loads(m_row0): + for slot in range_constexpr(kStages): + issue_a_load_lds_dt(arg_aq, lds_base_i32, slot, slot, m_row0, wave, lane, + is_f8, KH_TILE_A, k_bytes, BM=BM) + + def run_unit(unit_bx, m_block_idx): + # Map each Stage2 BM sub-tile to its Stage1 SBM metadata row. + m_row = m_block_idx * fx.Int32(BM) + sort_block_idx = m_row // fx.Int32(SBM) + row_in_sort_block = m_row - sort_block_idx * fx.Int32(SBM) + srcmap_row_base = trb_buf[sort_block_idx] + row_in_sort_block + if tx_i32 < fx.Int32(BM): + sorted_pos = srcmap_row_base + tx_i32 + fx.ptr_store( + outrow_buf[m_row + tx_i32], + lds_typed_ptr( + fx.Int32(lds_outrow_off) + tx_i32 * fx.Int32(4), + T.i32, + align=4, + ), + ) + fx.ptr_store( + sweights_buf[sorted_pos], + lds_typed_ptr( + fx.Int32(lds_weight_off) + tx_i32 * fx.Int32(4), + T.f32, + align=4, + ), + ) + # fmt: off + accm_vecs, m_row, n_block_idx, _n_out_rt = gemm2_compute_v2(lds_base_i32, arg_ascale, arg_bq, + arg_bscale, arg_eids, arg_aq, i32_max_m_blocks, unit_bx, lane, wave, i32_inter, i32_hidden, + i32_kpad, i32_npad, BM=BM, BN=BN, BK=BK, use_nt=use_nt, INTER_MAX=INTER_MAX, aStages=aStages, + a_dtype=a_dtype, has_pad=has_pad, SBM=SBM, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, + expert_offset=_expert_offset) + send_base = rocdl.readfirstlane(T.i64, fx.Int64(arg_send).ir_value()) + local_sort_epilog(lds_base_i32, accm_vecs, n_block_idx, wave, lane, send_base, + N_OUT=N_OUT, BM=BM, BN=BN, send_nbytes=send_nbytes, lds_outrow_off=lds_outrow_off, + lds_weight_off=lds_weight_off, g2_bf16_lds=g2_bf16_lds, scatter_vec=scatter_vec) + # fmt: on + + def run_one(unit_bx, m_block_idx): + issue_all_a_loads(m_block_idx * fx.Int32(BM)) + rocdl.sched_barrier(0) + run_unit(unit_bx, m_block_idx) + if const_expr(fuse): + # Every thread's stores must be complete before one thread publishes + # the unit; the copy engine reads them from memory, so release at + # system scope. + rocdl.s_waitcnt(0) + fx.barrier() + if tx_i32 == fx.Int32(0): + comm_ops.fence_system_release() + chunk = (m_block_idx * fx.Int32(BM)) // fx.Int32(SBM * comm.chunk_tiles) + comm_ops.atomic_add_agent( + fx.Int64(arg_cdone) + fx.Int64(chunk) * fx.Int64(4), fx.Int32(1) + ) + + cumsum0 = global_typed_ptr(arg_cumsum, T.i32)[0] + total_m_blocks = (cumsum0 + fx.Int32(BM - 1)) // fx.Int32(BM) + + def gemm_role(bx_i32): + + if const_expr(not persist and g2_spart <= 0): + bound = total_m_blocks * fx.Int32(num_n_blocks) + if fx.Int32(bx_i32) < bound: + run_one(bx_i32, bx_i32 // num_n_blocks) + elif const_expr(not persist): + bound = total_m_blocks * fx.Int32(num_n_blocks) + if fx.Int32(bx_i32) < bound: + m_block_idx, n_block_idx = _spart_output_tile_index( + bx_i32, total_m_blocks, num_n_blocks, g2_group_num, g2_m01 + ) + unit_bx = m_block_idx * fx.Int32(num_n_blocks) + n_block_idx + run_one(unit_bx, m_block_idx) + elif const_expr(skew_cu > 0): + total_stage1_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) + max_expert_tiles = global_typed_ptr(arg_max_expert_tiles, T.i32)[0] + skewed = max_expert_tiles * fx.Int32(4) > total_stage1_tiles + skew_n_block = bx_i32 // fx.Int32(cu_num) + skew_m_slot = bx_i32 - skew_n_block * fx.Int32(cu_num) + normal_m_slot = bx_i32 // fx.Int32(num_n_blocks) + normal_n_block = bx_i32 - normal_m_slot * fx.Int32(num_n_blocks) + n_block = skewed.select(skew_n_block, normal_n_block) + m_slot = skewed.select(skew_m_slot, normal_m_slot) + active_cu = skewed.select(fx.Int32(skew_cu), fx.Int32(cu_num)) + strided_diff = total_m_blocks - m_slot + strided_rem = (strided_diff > fx.Int32(0)).select(strided_diff, fx.Int32(0)) + strided_iters = (strided_rem + active_cu - fx.Int32(1)) // active_cu + tiles_per_slot = ( + total_m_blocks + active_cu - fx.Int32(1) + ) // active_cu + m_tile0 = m_slot * tiles_per_slot + contiguous_diff = total_m_blocks - m_tile0 + contiguous_rem = (contiguous_diff > fx.Int32(0)).select( + contiguous_diff, fx.Int32(0) + ) + contiguous_iters = (contiguous_rem < tiles_per_slot).select( + contiguous_rem, tiles_per_slot + ) + n_iters = skewed.select(strided_iters, contiguous_iters) + active = m_slot < active_cu + for _it in range(fx.Int32(0), n_iters, fx.Int32(1)): + strided_m = m_slot + fx.Int32(_it) * active_cu + contiguous_m = m_tile0 + fx.Int32(_it) + m_block = skewed.select(strided_m, contiguous_m) + if active: + unit_bx = m_block * fx.Int32(num_n_blocks) + n_block + fx.barrier() + if m_block < total_m_blocks: + run_one(unit_bx, m_block) + else: + m_slot = bx_i32 // fx.Int32(num_n_blocks) + n_block = bx_i32 - m_slot * fx.Int32(num_n_blocks) + if const_expr(persist_strided): + diff = total_m_blocks - m_slot + rem = (diff > fx.Int32(0)).select(diff, fx.Int32(0)) + n_iters = (rem + fx.Int32(cu_num - 1)) // fx.Int32(cu_num) + else: + tiles_per_slot = ( + total_m_blocks + fx.Int32(cu_num - 1) + ) // fx.Int32(cu_num) + m_tile0 = m_slot * tiles_per_slot + diff = total_m_blocks - m_tile0 + rem = (diff > fx.Int32(0)).select(diff, fx.Int32(0)) + n_iters = (rem < tiles_per_slot).select(rem, tiles_per_slot) + for _it in range(fx.Int32(0), n_iters, fx.Int32(1)): + if const_expr(persist_strided): + m_block = m_slot + fx.Int32(_it) * fx.Int32(cu_num) + else: + m_block = m_tile0 + fx.Int32(_it) + unit_bx = m_block * fx.Int32(num_n_blocks) + n_block + fx.barrier() # separate prev-iter epilog LDS reads from this iter's A-load into the LDS union + if fx.Int32(m_block) < total_m_blocks: + run_one(unit_bx, m_block) + + def comm_role(cw): + import mori.cco.device.flydsl as cco + + lay, rank, npes, CT = comm.layout, comm.rank, comm.layout.npes, comm.chunk_tiles + n_waves = NC * 4 + n_mine = (npes + n_waves - 1) // n_waves + mb_per_tile = SBM // BM + # This wave owns peers owner, owner + n_waves, ...; a wave past the last + # peer has nothing to do. + owner = cw * fx.Int32(4) + wave + is_lane0 = lane == fx.Int32(0) + if owner < fx.Int32(npes): + w = cco.Window(arg_win) + sdma = cco.DevComm(arg_dev_comm).sdma() + rsrc_counter = create_buffer_resource_from_addr( + fx.Int64(w.lsa_ptr(rank, lay.off_counter)) + ) + # Read-only here: the combine kernel advances the counter after the call, + # so every comm wave sees the same phase whenever it starts. + phase = fx.Int32(buffer_load(rsrc_counter, 0, vec_width=1, dtype=T.i32)) + fx.Int32(1) + cursor_buf = ptr_buf_tensor(arg_cursor, fx.Int32) + tile_base_buf = ptr_buf_tensor(arg_tile_base, fx.Int32) + # Do not overwrite a peer's landing buffer until it has read the last call's. + for j in range_constexpr(n_mine): + peer = owner + fx.Int32(j * n_waves) + if peer < fx.Int32(npes): + if peer != fx.Int32(rank): + if is_lane0: + sdma_wait_consumed(w, peer, phase, layout=lay, rank=rank) + total_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) + n_chunks = (total_tiles + fx.Int32(CT - 1)) // fx.Int32(CT) + for c in range(fx.Int32(0), n_chunks, fx.Int32(1)): + ci = fx.Int32(c) + m_lo = ci * fx.Int32(CT * mb_per_tile) + m_hi = m_lo + fx.Int32(CT * mb_per_tile) + m_hi = (m_hi < total_m_blocks).select(m_hi, total_m_blocks) + if is_lane0: + comm_ops.spin_until_ge_i32_agent( + fx.Int64(arg_cdone) + fx.Int64(ci) * fx.Int64(4), + (m_hi - m_lo) * num_n_blocks, + ) + comm_ops.fence_agent_acquire() + t0 = ci * fx.Int32(CT) + t1 = t0 + fx.Int32(CT) + t1 = (t1 < total_tiles).select(t1, total_tiles) + for j in range_constexpr(n_mine): + peer = owner + fx.Int32(j * n_waves) + if peer < fx.Int32(npes): + if peer != fx.Int32(rank): + lo = tile_base_buf[t0 * fx.Int32(npes) + peer] + hi = tile_base_buf[t1 * fx.Int32(npes) + peer] + if hi > lo: + if is_lane0: + sdma_put_run(sdma, arg_win, peer, peer % fx.Int32(comm.queues), + layout=lay, rank=rank, lo=lo, hi=hi, cco=cco) + for j in range_constexpr(n_mine): + peer = owner + fx.Int32(j * n_waves) + if peer < fx.Int32(npes): + if peer != fx.Int32(rank): + if is_lane0: + sdma_finish_peer(sdma, w, arg_win, peer, peer % fx.Int32(comm.queues), + cursor_buf[peer], phase, layout=lay, rank=rank, cco=cco) + + if const_expr(fuse): + if bid < fx.Int32(NC): + comm_role(bid) + else: + gemm_role(bid - fx.Int32(NC)) + else: + gemm_role(bid) + + # fmt: off + @flyc.jit + def launch(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, + arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, + arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, + arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, + arg_cursor: fx.Int64, i32_max_m_blocks: fx.Int32, i32_grid_blocks: fx.Int32, i32_inter: fx.Int32, + i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32, stream: fx.Stream): + # fmt: on + num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) + grid_x = i32_grid_blocks * num_n_blocks + fx.Int32(NC) + kernel_localsort( + arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, + arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, arg_cdone, + arg_tile_base, arg_cursor, i32_max_m_blocks, i32_inter, i32_hidden, i32_kpad, i32_npad, + ).launch(grid=(grid_x, 1, 1), block=(256, 1, 1), stream=stream) + + return launch + + +_LAUNCH_CACHE = {} + + +# fmt: off +def run_moe_gemm2_localsort(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, + arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, + arg_cdone, arg_tile_base, arg_cursor, row_capacity, i32_inter, i32_hidden, stream, *, + model_dim, inter_dim, experts, rank, send_nbytes, BM, SBM, HIDDEN_MAX, INTER_MAX, cu_num, + BN=256, BK=256, use_nt=True, g2_bhoist=True, g2_ascale_pf=True, g2_spart=402, persist=False, + persist_cu=0, persist_strided=False, g2_bf16_lds=False, skew_cu=0, scatter_vec=8, comm=None): + # fmt: on + """Compile or reuse one configuration and launch it.""" + launch_cu_num = min(cu_num, persist_cu) if persist and persist_cu > 0 else cu_num + compile_kw = dict( + model_dim=model_dim, inter_dim=inter_dim, experts=experts, rank=rank, + send_nbytes=send_nbytes, BM=BM, BN=BN, BK=BK, use_nt=use_nt, HIDDEN_MAX=HIDDEN_MAX, + INTER_MAX=INTER_MAX, SBM=SBM, persist=persist, cu_num=launch_cu_num, + g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, + persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, skew_cu=skew_cu, + scatter_vec=scatter_vec, comm=comm, + ) + key = tuple(sorted(compile_kw.items())) + launch = _LAUNCH_CACHE.get(key) + if launch is None: + launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_localsort(**compile_kw) + max_m_blocks = (row_capacity + BM - 1) // BM + grid_blocks = launch_cu_num if persist else max_m_blocks + _run_compiled( + launch, arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, + arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, + arg_cdone, arg_tile_base, arg_cursor, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), fx.Int32(i32_inter), fx.Int32(i32_hidden), + fx.Int32(0), fx.Int32(0), stream, + ) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py new file mode 100644 index 0000000000..ac9f297c2f --- /dev/null +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py @@ -0,0 +1,667 @@ +#!/usr/bin/env python3 +# 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. +"""Rank-sorted send buffer for gfx950 MegaMoE stage2 + combine, shipped by cco SDMA. + +Stage2 normally scatters every GEMM2 row straight into the origin rank's +``comb_inp``. Here the rows go to a local send buffer ``S`` grouped by origin +rank, so one contiguous run per peer can be moved by a copy engine:: + + S = [ origin rank 0: rows... | origin rank 1: rows... | ... ] (cap rows each) + +Within a rank's run rows are in sorted-row order, i.e. expert-major, so a prefix +of every run is finished as soon as the first experts are. That is what lets the +GEMM2 kernel ship chunks while it is still computing. + +Three kernels on one stream (the P2P-scatter baseline has two): + + 1. rank_sort_plan persistent grid; count -> scan -> assign with grid barriers: + out_row per sorted row, idx_send, tile_base, cursor; zeroes cdone + 2. gemm2_localsort writes S[out_row]; extra comm waves (one per peer) put finished + chunks with SDMA while the GEMM tiles are still computing + 3. combine persistent grid; wait for the peers' flags -> build slot->row + table -> grid barrier -> top-k reduce -> publish "consumed" + +Everything lives in ONE cco window, so ``put`` names the same window for source +and destination and only the byte offsets differ:: + + [ S | R | idx_send | idx_land | flags | done | counter ] + +``flags[src]`` says "src's data for call n has landed" (and how many rows). +``done[dst]`` says "dst has finished reading call n": a sender waits for it before +overwriting that peer's ``R``/``idx_land`` with call n+1, which keeps a fast rank +from running into a slow one. ``counter`` is the call number; only the combine +kernel advances it, after every reader of it has run. + +A row's *landing row* is its row index counted from the start of ``S``: ``S`` rows +are ``[0, npes*cap)`` and ``R`` rows ``[npes*cap, 2*npes*cap)``. My own run never +travels; it is read from ``S`` in place. + +Both persistent kernels use grid-wide barriers, so their grid must be fully +resident (one block per CU is). +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, range_constexpr, rocdl +from flydsl.expr.typing import T + +from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops +from aiter.ops.flydsl.kernels.buffer_ops import ( + buffer_load, + buffer_store, + create_buffer_resource_from_addr, +) +from aiter.ops.flydsl.kernels.mxfp4_gemm_common import global_typed_ptr +from aiter.ops.flydsl.kernels.tensor_shim import ( + buf_copy_load, + buf_copy_store, + ptr_buf_tensor, +) + +BLOCK = 256 +MIN_SBM = 16 +_ALIGN = 128 + + +def _align(n): + return (n + _ALIGN - 1) // _ALIGN * _ALIGN + + +@dataclass(frozen=True) +class ArenaLayout: + """Byte offsets inside the cco window; identical on every rank.""" + + npes: int + cap: int + row_bytes: int + + @property + def n_rows(self): + return self.npes * self.cap + + @property + def off_s(self): + return 0 + + @property + def off_r(self): + return self.off_s + self.n_rows * self.row_bytes + + @property + def off_idx_send(self): + return self.off_r + self.n_rows * self.row_bytes + + @property + def off_idx_land(self): + return self.off_idx_send + _align(self.n_rows * 4) + + @property + def off_flags(self): + return self.off_idx_land + _align(self.n_rows * 4) + + @property + def off_done(self): + return self.off_flags + _align(self.npes * 8) + + @property + def off_counter(self): + return self.off_done + _align(self.npes * 4) + + @property + def total_bytes(self): + return self.off_counter + _ALIGN + + def check(self): + # The reduce addresses S and R through one 32-bit buffer-resource offset. + if 2 * self.n_rows * self.row_bytes >= 1 << 32: + raise ValueError("S+R exceed the 4 GiB buffer-resource range") + + +def _uniform_i64(x): + """A wave-uniform i64 (SGPR) copy of ``x``, as a buffer-resource base needs.""" + return fx.Int64(rocdl.readfirstlane(T.i64, fx.Int64(x).ir_value())) + + +def _decode(p, *, log2_max_tok, recv_cap, topk, npes): + """Packed srcmap word -> (valid, dest_pe, dest_lid, slot_k), as the stage2 epilogue does.""" + t = p & fx.Int32(0x00FFFFFF) + s = p.shrui(fx.Int32(24)) & fx.Int32(0xFF) + dest_pe = t >> fx.Int32(log2_max_tok) + dest_lid = t & fx.Int32((1 << log2_max_tok) - 1) + valid = (t < fx.Int32(recv_cap)) & (s < fx.Int32(topk)) & (dest_pe < fx.Int32(npes)) + return valid, dest_pe, dest_lid, s + + +@comm_ops.traced +def grid_barrier(bar_addr, nblocks, tid): + """All blocks of a resident grid rendezvous. ``bar_addr`` is a never-reset i32 counter. + + Barrier k of a launch sees arrivals ``(k-1)*nblocks .. k*nblocks-1``, so the target + ``(old // nblocks + 1) * nblocks`` works for any number of barriers without a reset. + """ + rocdl.s_waitcnt(0) + fx.barrier() + if tid == fx.Int32(0): + comm_ops.fence_agent_release() + old = fx.Int32(comm_ops.atomic_add_agent(bar_addr, fx.Int32(1))) + target = (old // nblocks + fx.Int32(1)) * nblocks + comm_ops.spin_until_ge_i32_agent(bar_addr, target) + comm_ops.fence_agent_acquire() + fx.barrier() + + +def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap): + """Kernel 1: give every valid sorted row a stable index in its rank's run. + + * count: ``tile_cnt[tile][pe]`` = rows of the sort tile bound for ``pe`` + * scan: ``tile_base[tile][pe]`` = rows for ``pe`` in earlier tiles (so the + rows of a tile range are contiguous in the run), ``cursor[pe]`` = total + * assign: ``out_row = pe*cap + tile_base + (earlier rows of this tile for pe)``, + and ``idx_send[out_row] = origin_lid*topk + k`` + + The phases are separated by grid barriers inside one persistent launch. Row + loops read the tile's SBM rows from global memory instead of using LDS scans; + SBM is small and the rows are L2-resident. + """ + dec = dict(log2_max_tok=log2_max_tok, recv_cap=recv_cap, topk=topk, npes=npes) + + # fmt: off + @flyc.kernel(name=f"mega_rank_sort_plan_sbm{SBM}_k{topk}_p{npes}", known_block_size=[BLOCK, 1, 1]) + def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, + arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, + arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, + i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32): + # fmt: on + tid = fx.thread_idx.x + bid = fx.block_idx.x + gid = bid * fx.Int32(BLOCK) + tid + stride = i32_nblocks * fx.Int32(BLOCK) + num_valid = global_typed_ptr(arg_num_valid, T.i32)[0] + tiles = (num_valid + fx.Int32(SBM - 1)) // fx.Int32(SBM) + trb = ptr_buf_tensor(arg_trb, fx.Int32) + stids = ptr_buf_tensor(arg_stids, fx.Int32) + tile_cnt = ptr_buf_tensor(arg_tile_cnt, fx.Int32) + tile_base = ptr_buf_tensor(arg_tile_base, fx.Int32) + out_buf = ptr_buf_tensor(arg_out_row, fx.Int32) + cdone = ptr_buf_tensor(arg_cdone, fx.Int32) + + # The GEMM2 kernel's per-chunk completion counters start from zero every call. + for g in range(gid, i32_cdone_n, stride): + cdone[g] = fx.Int32(0) + + for g in range(gid, tiles * fx.Int32(npes), stride): + tile = g // fx.Int32(npes) + pe = g - tile * fx.Int32(npes) + base = trb[tile] + cnt = fx.Int32(0) + for r in range_constexpr(SBM): + valid, dest_pe, _, _ = _decode(stids[base + fx.Int32(r)], **dec) + cnt = cnt + (valid & (dest_pe == pe)).select(fx.Int32(1), fx.Int32(0)) + tile_cnt[g] = cnt + grid_barrier(arg_bar, i32_nblocks, tid) + + if bid == fx.Int32(0): + if tid < fx.Int32(npes): + run = arith.constant(0) + for t in range(fx.Int32(0), tiles, fx.Int32(1)): + tile_base[t * fx.Int32(npes) + tid] = run + run = run + tile_cnt[t * fx.Int32(npes) + tid] + tile_base[tiles * fx.Int32(npes) + tid] = run + ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = run + grid_barrier(arg_bar, i32_nblocks, tid) + + idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) + for g in range(gid, tiles * fx.Int32(SBM), stride): + tile = g // fx.Int32(SBM) + row = g - tile * fx.Int32(SBM) + base = trb[tile] + valid, dest_pe, dest_lid, s = _decode(stids[base + row], **dec) + out_buf[g] = fx.Int32(-1) + if valid: + within = fx.Int32(0) + for rp in range_constexpr(SBM): + v2, pe2, _, _ = _decode(stids[base + fx.Int32(rp)], **dec) + earlier = (fx.Int32(rp) < row) & v2 & (pe2 == dest_pe) + within = within + earlier.select(fx.Int32(1), fx.Int32(0)) + out_row = dest_pe * fx.Int32(cap) + tile_base[tile * fx.Int32(npes) + dest_pe] + within + out_buf[g] = out_row + idx_buf[out_row] = dest_lid * fx.Int32(topk) + s + + # fmt: off + @flyc.jit + def launch(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, + arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, + arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, + i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32, stream: fx.Stream): + # fmt: on + plan(arg_stids, arg_trb, arg_num_valid, arg_tile_cnt, arg_tile_base, arg_cursor, + arg_out_row, arg_idx, arg_cdone, arg_bar, i32_nblocks, i32_cdone_n).launch( + grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + return launch + + +def sdma_wait_consumed(w, peer, phase, *, layout, rank): + """Block until ``peer`` has finished reading the previous call's data from my sends.""" + comm_ops.spin_until_ge_i32_system( + fx.Int64(w.lsa_ptr(rank, layout.off_done)) + fx.Int64(peer) * fx.Int64(4), + phase - fx.Int32(1), + acquire=True, + ) + + +def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): + """Ship rows ``[lo, hi)`` of my run for ``peer`` (data only; idx goes once at the end).""" + rb, cap = layout.row_bytes, layout.cap + sdma.put( + peer, + win, + fx.Int64(layout.off_r + rank * cap * rb) + fx.Int64(lo) * fx.Int64(rb), + win, + fx.Int64(layout.off_s) + + (fx.Int64(peer) * fx.Int64(cap) + fx.Int64(lo)) * fx.Int64(rb), + fx.Int64(hi - lo) * fx.Int64(rb), + queue, + coop=cco.CoopScope.THREAD, + ) + + +@comm_ops.traced +def sdma_finish_peer(sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco): + """After the last data put: ship idx, drain the queue, publish ``(cnt << 32) | phase``.""" + if cnt > fx.Int32(0): + sdma.put( + peer, + win, + fx.Int64(layout.off_idx_land + rank * layout.cap * 4), + win, + fx.Int64(layout.off_idx_send) + fx.Int64(peer) * fx.Int64(layout.cap * 4), + fx.Int64(cnt) * fx.Int64(4), + queue, + coop=cco.CoopScope.THREAD, + ) + sdma.quiet_queue(peer, queue) + # The bytes were moved by the copy engine, not by this CU; the fence only keeps + # the flag store from being hoisted above the drain, which is the ordering the + # receiver relies on. + comm_ops.fence_system_release() + comm_ops.store_i64_global_system( + fx.Int64(w.lsa_ptr(peer, layout.off_flags)) + fx.Int64(rank * 8), + (fx.Int64(cnt) << fx.Int64(32)) | fx.Int64(phase), + ) + + +def build_combine(*, layout, rank, hidden, topk): + """Kernel 3: finish the transfer and reduce. + + One persistent launch (grid resident), in phases: + + A. wait for every peer's flag for this call (acquire), which also says how many + rows each sent; + B. ``inv[idx[src][j]] = landing row`` for every landed row, spread over the grid; + -- grid barrier -- + C. ``out[t] = sum_k row[inv[t*topk+k]]`` in f32, one block per token in turn. A slot's + ``inv`` entry is reset to -1 once read, so it starts clean next call and nothing + has to refill it; slots no route filled stay -1 and contribute zero; + D. the last block to finish publishes ``done`` to every peer and advances ``counter``. + """ + import mori.cco.device.flydsl as cco + + if hidden % 8: + raise ValueError(f"hidden must be a multiple of 8, got {hidden}") + npes, cap = layout.npes, layout.cap + units = hidden // 8 + chunks = (units + BLOCK - 1) // BLOCK + + # fmt: off + @flyc.kernel(name=f"mega_combine_p{npes}_c{cap}_h{hidden}_k{topk}_r{rank}", known_block_size=[BLOCK, 1, 1]) + def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, + arg_bar: fx.Int64, arg_fin: fx.Int64, i32_nblocks: fx.Int32, i32_tokens: fx.Int32): + # fmt: on + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.Window(win) + rsrc_counter = create_buffer_resource_from_addr( + _uniform_i64(w.lsa_ptr(rank, layout.off_counter)) + ) + # Advanced only by the last block below, after every block has read it. + phase = fx.Int32(buffer_load(rsrc_counter, 0, vec_width=1, dtype=T.i32)) + fx.Int32(1) + flags_base = fx.Int64(w.lsa_ptr(rank, layout.off_flags)) + + # A. the peers' data has landed + if tid < fx.Int32(npes): + if tid != fx.Int32(rank): + comm_ops.spin_until_ge_i32_system( + flags_base + fx.Int64(tid) * fx.Int64(8), phase, acquire=True + ) + fx.barrier() + + # B. slot -> landing row + idx_send = fx.Int64(w.lsa_ptr(rank, layout.off_idx_send)) + idx_land = fx.Int64(w.lsa_ptr(rank, layout.off_idx_land)) + inv_buf = ptr_buf_tensor(arg_inv, fx.Int32) + cursor = ptr_buf_tensor(arg_cursor, fx.Int32) + gid = bid * fx.Int32(BLOCK) + tid + stride = i32_nblocks * fx.Int32(BLOCK) + for g in range(gid, fx.Int32(npes * cap), stride): + src = g // fx.Int32(cap) + is_self = src == fx.Int32(rank) + theirs = fx.Int32( + comm_ops.load_i32_global_system(flags_base + fx.Int64(src) * fx.Int64(8) + fx.Int64(4)) + ) + cnt = is_self.select(cursor[fx.Int32(rank)], theirs) + if (g - src * fx.Int32(cap)) < cnt: + base = is_self.select(idx_send, idx_land) + slot = global_typed_ptr(base + fx.Int64(g) * fx.Int64(4), T.i32)[0] + inv_buf[slot] = is_self.select(g, g + fx.Int32(npes * cap)) + grid_barrier(arg_bar, i32_nblocks, tid) + + # C. top-k reduce + rows = ptr_buf_tensor(_uniform_i64(w.lsa_ptr(rank, layout.off_s)), fx.BFloat16, unit_elems=8) + out = ptr_buf_tensor(arg_out, fx.BFloat16, unit_elems=8) + for tok in range(bid, i32_tokens, i32_nblocks): + slots = [] + for k in range_constexpr(topk): + slots.append(inv_buf[tok * fx.Int32(topk) + fx.Int32(k)]) + for chunk in range_constexpr(chunks): + u = fx.Int32(chunk * BLOCK) + tid + if u < fx.Int32(units): + acc = fx.Vector.filled(8, 0.0, fx.Float32) + for k in range_constexpr(topk): + ok = slots[k] >= fx.Int32(0) + safe = ok.select(slots[k], fx.Int32(0)) + raw = buf_copy_load(rows, safe * fx.Int32(units) + u, fx.BFloat16, unit_elems=8) + keep = ok.select(fx.Float32(1.0), fx.Float32(0.0)) + acc = acc + fx.Vector(raw).to(fx.Float32) * keep + buf_copy_store( + out, tok * fx.Int32(units) + u, acc.to(fx.BFloat16), fx.BFloat16, unit_elems=8 + ) + # Every thread has the slots in registers; only then may they be reset. + fx.barrier() + if tid < fx.Int32(topk): + inv_buf[tok * fx.Int32(topk) + tid] = fx.Int32(-1) + + # D. the last block out tells the peers this call's data has been consumed + rocdl.s_waitcnt(0) + fx.barrier() + if tid == fx.Int32(0): + comm_ops.fence_agent_release() + old = fx.Int32(comm_ops.atomic_add_agent(arg_fin, fx.Int32(1))) + if (old + fx.Int32(1)) % i32_nblocks == fx.Int32(0): + comm_ops.fence_agent_acquire() + buffer_store(phase, rsrc_counter, 0) + comm_ops.fence_system_release() + for peer in range_constexpr(npes): + if peer != rank: + comm_ops.store_i32_global_system_release( + fx.Int64(w.lsa_ptr(peer, layout.off_done)) + fx.Int64(rank * 4), phase + ) + + # fmt: off + @flyc.jit + def launch(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, + arg_bar: fx.Int64, arg_fin: fx.Int64, i32_nblocks: fx.Int32, i32_tokens: fx.Int32, + stream: fx.Stream): + # fmt: on + combine(win, arg_cursor, arg_inv, arg_out, arg_bar, arg_fin, i32_nblocks, i32_tokens).launch( + grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + return launch + + +def _torch_view(ptr, nbytes, dtype): + import torch + + class _View: + def __init__(self): + self.__cuda_array_interface__ = { + "data": (int(ptr), False), + "shape": (int(nbytes),), + "strides": None, + "typestr": "|u1", + "version": 3, + } + + return torch.as_tensor(_View(), device=f"cuda:{torch.cuda.current_device()}").view( + dtype + ) + + +class SortedCombine: + """GEMM2 with a rank-sorted local output, shipped by SDMA from inside the GEMM2 kernel. + + Reuses a ``MegaMoEV2``'s stage1/dispatch state, so it can run right after + ``moe(...)`` and be compared with that call's own combine output. + + ``chunk_tiles`` is how many sort tiles make one shipped chunk; ``comm_wgs`` the number + of comm workgroups (4 waves each, one wave per peer by default). + + Needs ``MORI_ENABLE_SDMA=1`` and a mori built with SDMA. + """ + + def __init__(self, moe, queues=8, chunk_tiles=4, comm_wgs=None): + import torch + import torch.distributed as dist + from mori.cco import ( + GDA_CONNECTION_NONE, + CCODevCommRequirements, + Communicator, + ) + + from moe_gemm2_localsort import CommSpec + + cfg = moe.comb_cfg + self.moe = moe + self.rank, self.npes = int(cfg.rank), int(cfg.world_size) + self.hidden = int(cfg.hidden_dim) + self.topk = int(cfg.num_experts_per_token) + self.max_tok = int(cfg.max_num_inp_token_per_rank) + self.cap = self.max_tok * self.topk + self.layout = ArenaLayout(self.npes, self.cap, self.hidden * 2) + self.layout.check() + dev = torch.device("cuda", self.rank) + lay = self.layout + + uid = Communicator.get_unique_id() if self.rank == 0 else None + holder = [uid] + dist.broadcast_object_list(holder, src=0) + self.comm = Communicator.init(self.npes, self.rank, holder[0]) + self.mem = self.comm.alloc_mem(lay.total_bytes) + self.win = self.comm.register_window(self.mem.ptr, self.mem.size) + _torch_view(self.mem.ptr, lay.total_bytes, torch.uint8).zero_() + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + reqs.sdma_queue_count = queues + self.dev_comm = self.comm.create_dev_comm(reqs) + + # The sort block size can change with the token count, so size the per-tile arrays + # for the smallest one and build the plan kernel on first use of each size. + nvm = int(moe._s1_nvm) + max_tiles = nvm // MIN_SBM + 2 + self.max_tiles = max_tiles + i32 = dict(dtype=torch.int32, device=dev) + self.out_row = torch.full((nvm + 4 * 128,), -1, **i32) + self.tile_cnt = torch.zeros(max_tiles * self.npes, **i32) + self.tile_base = torch.zeros((max_tiles + 1) * self.npes, **i32) + self.cursor = torch.zeros(self.npes, **i32) + self.cdone = torch.zeros(max_tiles + 1, **i32) + # Never-reset counters for the two grid barriers and the combine's last-block count. + self.plan_bar = torch.zeros(1, **i32) + self.combine_bar = torch.zeros(1, **i32) + self.combine_fin = torch.zeros(1, **i32) + # -1 means "no row for this slot"; the combine kernel puts it back after reading. + self.inv = torch.full((self.cap,), -1, **i32) + self.out = torch.empty( + self.max_tok * self.hidden, dtype=torch.bfloat16, device=dev + ) + # Both persistent kernels barrier across the grid, so keep it resident: one block per CU. + self.nblocks = torch.cuda.get_device_properties(dev).multi_processor_count + torch.cuda.synchronize() + self.comm.barrier() + + # One comm wave per peer (4 waves per workgroup). + comm_wgs = max(1, (self.npes + 3) // 4) if comm_wgs is None else comm_wgs + self.comm_spec = CommSpec( + layout=lay, + rank=self.rank, + queues=queues, + chunk_tiles=chunk_tiles, + comm_wgs=comm_wgs, + ) + self._plans = {} + self._combine = build_combine( + layout=lay, rank=self.rank, hidden=self.hidden, topk=self.topk + ) + + def _plan_for(self, sbm): + if sbm < MIN_SBM: + raise ValueError( + f"sort_block_m={sbm} is below the supported minimum {MIN_SBM}" + ) + if sbm not in self._plans: + self._plans[sbm] = build_rank_sort_plan( + SBM=sbm, + topk=self.topk, + log2_max_tok=self.max_tok.bit_length() - 1, + recv_cap=int(self.moe.max_recv), + npes=self.npes, + cap=self.cap, + ) + return self._plans[sbm] + + def set_chunk_tiles(self, chunk_tiles): + """Change the shipped chunk size; only the gemm2 kernel is (re)compiled.""" + from dataclasses import replace + + self.comm_spec = replace(self.comm_spec, chunk_tiles=chunk_tiles) + + def _launch_gemm2(self, config, s_fx): + import flydsl.expr as fx + from moe_gemm2_localsort import run_moe_gemm2_localsort + + moe, op, lay = self.moe, self.moe._s1_op, self.layout + ws, stage2 = moe._s1_dispatch_workspace, config.stage2 + inv = moe._g2_invariants_by_quant["none"] + run_moe_gemm2_localsort( + fx.Int64(moe._s1_out.view(-1).data_ptr()), + fx.Int64(moe._s1_osd.data_ptr()), + fx.Int64(moe.w2.data_ptr()), + fx.Int64(moe.w2_scale.data_ptr()), + fx.Int64(op.sorted_expert_ids.data_ptr()), + fx.Int64(op.num_valid.data_ptr()), + fx.Int64(ws["max_expert_tiles"].data_ptr()), + fx.Int64(op.wts_em.data_ptr()), + fx.Int64(op.tile_row_base.data_ptr()), + fx.Int64(self.mem.ptr + lay.off_s), + fx.Int64(self.out_row.data_ptr()), + fx.Int64(self.dev_comm.ptr), + fx.Int64(self.win.handle), + fx.Int64(self.cdone.data_ptr()), + fx.Int64(self.tile_base.data_ptr()), + fx.Int64(self.cursor.data_ptr()), + moe._s1_nvm, + moe._g2v2_inter, + moe._g2v2_hidden, + s_fx, + model_dim=inv["model_dim"], + inter_dim=inv["inter_dim"], + experts=inv["experts"], + rank=inv["rank"], + send_nbytes=int(lay.n_rows * lay.row_bytes), + BM=stage2.block_m, + SBM=config.stage1.sort_block_m, + HIDDEN_MAX=inv["HIDDEN_MAX"], + INTER_MAX=inv["INTER_MAX"], + cu_num=inv["cu_num"], + BN=stage2.block_n, + BK=stage2.block_k, + use_nt=stage2.use_nt, + g2_bhoist=stage2.b_hoist, + g2_ascale_pf=stage2.ascale_prefetch, + g2_spart=stage2.spatial_partition, + persist=stage2.persist, + persist_cu=stage2.persist_cu, + persist_strided=stage2.persist_strided, + skew_cu=stage2.skew_cu, + g2_bf16_lds=stage2.bf16_lds, + comm=self.comm_spec, + ) + + def run(self, run_tokens, config=None, profile=None): + """Three kernel launches. ``profile``: a list that receives ``(stage, cuda_event)`` marks.""" + import flydsl.expr as fx + import torch + from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled + + moe, lay = self.moe, self.layout + config = moe._active_config if config is None else config + s_fx = fx.Stream(torch.cuda.current_stream().cuda_stream) + op = moe._s1_op + + def mark(stage): + if profile is not None: + ev = torch.cuda.Event(enable_timing=True) + ev.record() + profile.append((stage, ev)) + + mark("start") + _run_compiled( + self._plan_for(int(config.stage1.sort_block_m)), + fx.Int64(op.srcmap_em.data_ptr()), + fx.Int64(op.tile_row_base.data_ptr()), + fx.Int64(op.num_valid.data_ptr()), + fx.Int64(self.tile_cnt.data_ptr()), + fx.Int64(self.tile_base.data_ptr()), + fx.Int64(self.cursor.data_ptr()), + fx.Int64(self.out_row.data_ptr()), + fx.Int64(self.mem.ptr + lay.off_idx_send), + fx.Int64(self.cdone.data_ptr()), + fx.Int64(self.plan_bar.data_ptr()), + fx.Int32(self.nblocks), + fx.Int32(self.cdone.numel()), + s_fx, + ) + mark("plan") + self._launch_gemm2(config, s_fx) + mark("gemm2+sdma") + _run_compiled( + self._combine, + fx.Int64(self.win.handle), + fx.Int64(self.cursor.data_ptr()), + fx.Int64(self.inv.data_ptr()), + fx.Int64(self.out.data_ptr()), + fx.Int64(self.combine_bar.data_ptr()), + fx.Int64(self.combine_fin.data_ptr()), + fx.Int32(self.nblocks), + fx.Int32(run_tokens), + s_fx, + ) + mark("combine") + return self.out[: run_tokens * self.hidden].view(run_tokens, self.hidden) + + def close(self): + self.comm.close_all_resources() From c7a8acfdc2b798ab35ea5ce9c740f6981a23765b Mon Sep 17 00:00:00 2001 From: kawhil-amd Date: Tue, 27 Oct 2026 15:31:13 +0000 Subject: [PATCH 2/9] example(cco): run the SDMA MoE combine on GPUs, and make it win above 4k tokens The example was compile-checked only. Running it needed two fixes (aiter here does not export comm_ops.traced, and mori's SDMA queues are per (src,dst) pair so --queues 8 asks for 64 per rank and fails), and then it was 0.48x aiter's stage2+combine at 2048 tokens/rank. It is now 1.09x at 4096 and 1.18x at 8192, matching aiter's output bit for bit wherever aiter also picks fp8. Four changes, in order of what they were worth: fp8 transport. bf16 moved twice the bytes and the transfer stopped fitting under the GEMM; a truncated-shipping probe showed the comm cost tracked bytes, not sync. Rows now carry aiter's fp8_blockwise_1x32 in aiter's own layout -- payload then the e8m0 scales, inline -- so one contiguous put still ships a whole row. Write-through epilogue stores. The per-unit publish released at system scope, which on gfx950 is buffer_wbl2: a write-back of the *whole* L2, charged once per (m_block, n_block) unit, 12288 times a launch. Storing sc0|sc1 instead makes retiring the stores the release the copy engine needs. 2.4x on the GEMM2 kernel. Same recipe as ops/gemm_ar/kernels_sdma.py. The per-chunk counters are also a cache line apart now, which is worth 8% on its own. A per-pair row reservation. `cap` was the whole token-slot space reserved for every (src,dst) pair, 16x what a pair can use; that put bs=8192 over the 4 GiB buffer-resource range. It is now the even share times --row-slack, with the plan kernel refusing to place a row past it, which keeps the handshake-free write into the peer's landing buffer. A cheaper plan kernel and completion signal. The plan had two grid barriers around a prefix sum walked by 8 threads; it is one barrier and a Hillis-Steele scan every block runs in its own LDS, 0.066 -> 0.030ms. And the comm wave hangs the copy engine's own trailing ATOMIC on its last put instead of draining the queue, so it no longer pins the tail of the GEMM2 kernel. bench.py times eager by default -- prefill does not capture CUDA graphs -- with --graph for the two-graph slope bench_ep.py uses; the two agree within ~1%. --no-aligned-pair lets bs=8192 compare, since aiter picks its common-pair stage2 there and this path has no equivalent (measured 1.5% slower than without it on uniform routing anyway). sdmabw.py is new: one cco SDMA channel with no GEMM in the way, which says t_put = 5.8us + bytes / 60.9 GB/s and sizes the above against the hardware. stage2+combine vs aiter, glm52 on 8x MI355, eager, --chunk-tiles 16: tokens/rank 128 512 2048 4096 8192 vs aiter 0.64x 0.72x 0.97x 1.09x 1.18x whole layer 0.91x 0.90x 1.00x 1.05x 1.08x Below ~2048 the plan launch and the combine's index inversion outweigh what the copy engine saves, and aiter's P2P scatter wins. Co-Authored-By: Claude Opus 5 (1M context) --- examples/cco/README.md | 22 +- .../python/11_flydsl_mega_moe_sdma/bench.py | 109 ++++- .../python/11_flydsl_mega_moe_sdma/main.py | 35 +- .../moe_gemm2_localsort.py | 203 +++++++-- .../python/11_flydsl_mega_moe_sdma/sdmabw.py | 189 +++++++++ .../11_flydsl_mega_moe_sdma/sorted_combine.py | 386 +++++++++++++++--- 6 files changed, 820 insertions(+), 124 deletions(-) create mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py diff --git a/examples/cco/README.md b/examples/cco/README.md index 1cfaf60c86..67e2af95ad 100644 --- a/examples/cco/README.md +++ b/examples/cco/README.md @@ -23,6 +23,7 @@ examples/cco/ | `python/09_torch_symm_lsa_all2all` | py | mori's own torch symm-mem backend + Triton `lsa_ptr` all-to-all | | `python/10_flydsl_sdma_allgather_async` | py | FlyDSL SDMA all-gather split into `start_async` / `wait_async`, overlapped with a GEMM | | `python/11_flydsl_mega_moe_sdma` | py | MoE GEMM2 (gfx950) writing a rank-sorted send buffer, shipped by SDMA `put`s from comm waves (one per peer) inside the GEMM kernel (needs aiter) | +| `python/11_flydsl_mega_moe_sdma/sdmabw.py` | py | what one cco SDMA channel moves, with no GEMM in the way | | `ir/test_triton_cco.py` | py | Triton DevComm queries and LSA/SDMA/GDA device APIs | | `cpp/01_lsa_put.cpp` | c++ | intra-node LSA put (includes only `cco.hpp`) | | `cpp/02_gda_put.cpp` | c++ | GPU-initiated RDMA put + signal/wait (includes only `cco_scale_out.hpp`) | @@ -96,14 +97,31 @@ Example 11 also uses `torchrun` and SDMA, and additionally needs aiter (it reuse ```bash MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ - examples/cco/python/11_flydsl_mega_moe_sdma/main.py --network glm52 --bs 128 + examples/cco/python/11_flydsl_mega_moe_sdma/main.py --network glm52 --bs 4096 # latency vs aiter's MegaMoEV2 (mori.shmem P2P scatter), sweeping token counts MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ examples/cco/python/11_flydsl_mega_moe_sdma/bench.py --network glm52 \ - --bs-list 16,128,1024,4096 --chunk-tiles 2,4,8,16 + --bs-list 128,512,2048,4096,8192 --no-aligned-pair + +# what one SDMA channel actually moves, to size the above against the hardware +MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ + examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py ``` +It is worth it only at large per-rank batches. `stage2+combine` against aiter, +eager, fp8 transport, `--chunk-tiles 16`, glm52 on 8x MI355, matching aiter's +output exactly (relL2 0) at every size: + +| tokens/rank | 128 | 512 | 2048 | 4096 | 8192 | +|---|---|---|---|---|---| +| vs aiter | 0.63x | 0.69x | 0.97x | **1.08x** | **1.18x** | +| whole MoE layer | 0.86x | 0.84x | 1.00x | 1.05x | **1.08x** | + +Below ~2048 the three kernels' fixed costs -- the plan launch and its grid +barrier, and the combine's extra index inversion -- outweigh what the copy +engine saves, and aiter's P2P scatter wins. + `MORI_CCO_GDA_CONN=full` is required for GDA (03, 06) when both ranks share one node; LSA examples (02, 04, 05) ignore it. Each example prints `SUCCESS` on pass. diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py index c4c8d3bb1c..6156788868 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py @@ -40,14 +40,15 @@ breakdown per-stage time of one eager call (rank 0), to see where the new path spends its time -CUDA Graph timing by default (what the production path uses); ``--no-graph`` times -eager launches. +Eager launches by default, which is what prefill does; ``--graph`` switches to +the two-graph slope used by ``bench_ep.py``. MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 bench.py \\ --network glm52 --bs-list 16,128,1024,4096 """ import argparse +from dataclasses import replace from main import NETWORKS, _make_inputs, _next_pow2, _quantize_weights, _setup @@ -75,18 +76,55 @@ def _sync(): dist.barrier() -def _time(fn, iters, graph): - """Mean/max-over-ranks ms per call. Falls back to eager if the graph cannot be captured.""" +#: Calls captured per graph for the slope measurement, as ``bench_ep.py`` does. +GRAPH_REPS = 10 + + +def _replay_ms(g, n): + """Mean ms of one ``g.replay()`` over ``n`` replays.""" + _sync() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + start.record() + for _ in range(n): + g.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / n + + +def _graph_of(fn, reps): + """Capture ``reps`` back-to-back calls, warm it up, lockstep the ranks.""" + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch.cuda.Stream()): + for _ in range(reps): + fn() + _sync() + for _ in range(5): + g.replay() + _sync() + return g + + +def _time(fn, iters, graph, reps=GRAPH_REPS): + """Mean/max-over-ranks ms per call. + + With ``graph``, this is the *marginal* cost of one call: two graphs holding + ``reps`` and ``2*reps`` calls are replayed and differenced, so the per-replay + launch overhead -- which at bs=16 is the same order as the work -- cancels + instead of being charged to every call. Same two-graph slope as + ``mori/tests/python/ops/dispatch_combine_v2/bench_ep.py``. Falls back to + eager timing if the capture fails. + """ _sync() fn() _sync() - replay = fn if graph: try: - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g, stream=torch.cuda.Stream()): - fn() - replay = g.replay + g1, g2 = _graph_of(fn, reps), _graph_of(fn, 2 * reps) + t1, t2 = _replay_ms(g1, iters), _replay_ms(g2, iters) + return _over_ranks((t2 - t1) / reps) except Exception as exc: # noqa: BLE001 if dist.get_rank() == 0: print( @@ -95,14 +133,14 @@ def _time(fn, iters, graph): ) torch.cuda.synchronize() for _ in range(5): - replay() + fn() _sync() start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( enable_timing=True ) start.record() for _ in range(iters): - replay() + fn() end.record() torch.cuda.synchronize() return _over_ranks(start.elapsed_time(end) / iters) @@ -135,20 +173,48 @@ def _rel_l2(out, ref): def main(): parser = argparse.ArgumentParser() parser.add_argument("--network", choices=NETWORKS, default="glm52") - parser.add_argument("--bs-list", default="16,128,1024,4096", help="tokens per rank") + parser.add_argument( + "--bs-list", default="128,512,2048,4096,8192", help="tokens per rank" + ) parser.add_argument("--iters", type=int, default=30) parser.add_argument("--seed", type=int, default=123) parser.add_argument("--rtol", type=float, default=0.05) parser.add_argument( - "--chunk-tiles", default="4", help="fused chunk sizes to sweep, e.g. 2,4,8,16" + "--chunk-tiles", default="16", help="fused chunk sizes to sweep, e.g. 8,16,32" + ) + parser.add_argument( + "--queues", type=int, default=1, + help="SDMA channels per (src, dst) pair. More than one measured " + "slower: KFD recommends a single engine per peer link here, so " + "extra channels share it and only add descriptors.", + ) + parser.add_argument("--row-slack", type=float, default=1.5) + parser.add_argument( + "--signal-mode", choices=("quiet", "remote"), default="remote", + help="how a sender tells a peer its rows landed", + ) + parser.add_argument( + "--no-aligned-pair", action="store_true", + help="force aiter off its common-pair stage2 so both paths run the " + "same config. It is picked only at bucket 8192, where SortedCombine " + "has no equivalent -- and measured 1.5%% slower than without it on " + "uniform routing anyway.", + ) + parser.add_argument( + "--quant", choices=("none", "fp8_blockwise_1x32"), default="fp8_blockwise_1x32", + help="quantise the send buffer; halves the bytes SDMA moves", + ) + parser.add_argument( + "--graph", + action="store_true", + help="time CUDA Graph replays instead of eager launches. Off by " + "default: prefill does not capture, so eager is the shape that matters.", ) - parser.add_argument("--queues", type=int, default=8) - parser.add_argument("--no-graph", action="store_true") args = parser.parse_args() batch_sizes = [int(v) for v in args.bs_list.split(",")] chunk_list = [int(v) for v in args.chunk_tiles.split(",")] - graph = not args.no_graph + graph = args.graph rank, world, device = _setup() rows = [] # (bs, label, s2_mean, s2_max, e2e_mean, e2e_max, relL2) @@ -186,7 +252,9 @@ def main(): fanout_masks=(), **net, ) - sc = SortedCombine(moe, queues=args.queues) + sc = SortedCombine(moe, queues=args.queues, quant=args.quant, + row_slack=args.row_slack, + signal_mode=args.signal_mode) for bs in batch_sizes: x, wts, ids = ( @@ -197,6 +265,13 @@ def main(): ref = moe(x, wts, ids)[:bs].clone() torch.cuda.synchronize() config = moe._active_config + if args.no_aligned_pair and config.stage2.aligned_pair: + config = replace( + config, + stage2=replace( + config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96 + ), + ) if rank == 0: print( f"\n=== bs={bs} tokens/rank p2p_quant={config.p2p_quant} " diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py index 5458e1eacd..75c2609a9a 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py @@ -50,6 +50,7 @@ """ import argparse +from dataclasses import replace import os os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "20G") @@ -169,10 +170,31 @@ def main(): parser.add_argument( "--chunk-tiles", type=int, - default=4, + default=16, help="sort tiles per shipped chunk", ) - parser.add_argument("--queues", type=int, default=8) + parser.add_argument( + "--queues", type=int, default=1, + help="SDMA channels per (src, dst) pair. More than one measured " + "slower: KFD recommends a single engine per peer link here, so " + "extra channels share it and only add descriptors.", + ) + parser.add_argument("--row-slack", type=float, default=1.5) + parser.add_argument( + "--signal-mode", choices=("quiet", "remote"), default="remote", + help="how a sender tells a peer its rows landed", + ) + parser.add_argument( + "--no-aligned-pair", action="store_true", + help="force aiter off its common-pair stage2 so both paths run the " + "same config. It is picked only at bucket 8192, where SortedCombine " + "has no equivalent -- and measured 1.5%% slower than without it on " + "uniform routing anyway.", + ) + parser.add_argument( + "--quant", choices=("none", "fp8_blockwise_1x32"), default="fp8_blockwise_1x32", + help="quantise the send buffer; halves the bytes SDMA moves", + ) args = parser.parse_args() rank, world, device = _setup() @@ -210,6 +232,11 @@ def main(): ref = moe(x, wts, ids)[: args.bs].clone() torch.cuda.synchronize() config = moe._active_config + if args.no_aligned_pair and config.stage2.aligned_pair: + config = replace( + config, + stage2=replace(config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96), + ) if config.stage2.aligned_pair: raise SystemExit( "this config selects aligned-pair stage2, which SortedCombine does not cover" @@ -225,7 +252,9 @@ def main(): lambda: moe._run_stage2(args.bs, None, True, config), args.iters ) failures = [] - sc = SortedCombine(moe, queues=args.queues, chunk_tiles=args.chunk_tiles) + sc = SortedCombine(moe, queues=args.queues, chunk_tiles=args.chunk_tiles, quant=args.quant, + row_slack=args.row_slack, + signal_mode=args.signal_mode) worst = 0.0 # Repeated calls with no host sync between them also exercise the epoch counter, # the flags and the "consumed" handshake that stops a fast rank overrunning a slow one. diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py index 5f219fba88..f8aecebe33 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py @@ -36,10 +36,23 @@ Each comm WAVE owns its own peers (peer ``p`` belongs to wave ``p % n_comm_waves``): it spins on the chunk counter itself, then puts that chunk's rows for its peers on its own SDMA queue, so peers neither wait for each other's spin nor share a wave's issue slot. + +Two things about the publish are load-bearing and were each worth more than a +factor on this kernel. The epilogue stores the rows write-through (sc0|sc1) so +that retiring them is the release a copy engine needs; releasing at system scope +instead is a whole-L2 write-back per unit. And the per-chunk counters sit a +cache line apart, because every unit of a chunk bumps one. """ from dataclasses import dataclass +#: i32 slots between consecutive chunk-completion counters. Every unit of a +#: chunk bumps its counter, so packing them adjacent puts all of a launch's +#: atomics on one or two cache lines: 128 B apart instead is worth ~8% of the +#: GEMM2 kernel at bs=4096. Spreading a chunk's count over several counters on +#: top of this buys nothing -- the atomic itself is not the cost. +CDONE_STRIDE = 32 + import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import const_expr, range_constexpr, rocdl @@ -49,6 +62,7 @@ from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops from aiter.ops.flydsl.kernels.buffer_ops import ( buffer_load, + buffer_store, create_buffer_resource_from_addr, ) from aiter.ops.flydsl.kernels.mxfp4_gemm_common import ( @@ -77,15 +91,22 @@ ) from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage2 import ( _BUFFER_OFFSET_ABI_BYTES, + _fp8_scale_for_leader, _stage2_lds_bytes, ) +from aiter.ops.flydsl.kernels.mxfp4_gemm_common import _fabs_f32 as fabs_f32 @dataclass(frozen=True) class CommSpec: """Compile-time description of the fused SDMA comm workgroups. - ``comm_wgs`` workgroups of 4 waves each; one wave per peer needs ``comm_wgs * 4 >= npes``. + ``comm_wgs`` workgroups of 4 waves each; one wave per peer needs + ``comm_wgs * 4 >= npes``. ``queues`` is the SDMA channels reserved per + (src, dst) pair; peer ``p`` uses channel ``p % queues`` for everything it + sends, so a chunk is never split across channels -- splitting measured + slower, and on this topology KFD recommends one engine per peer link + anyway, which makes extra channels pure overhead. """ layout: ArenaLayout @@ -93,11 +114,16 @@ class CommSpec: queues: int chunk_tiles: int comm_wgs: int = 1 + #: How a sender tells a peer its rows have landed; see + #: ``sorted_combine.sdma_finish_peer``. "remote" needs the signal to ride + #: the channel the data went out on, which ``queues`` already guarantees. + signal_mode: str = "remote" # fmt: off def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, N_OUT, BM, BN, - send_nbytes, lds_outrow_off, lds_weight_off, g2_bf16_lds=False, scatter_vec=8): + send_nbytes, lds_outrow_off, lds_weight_off, g2_bf16_lds=False, scatter_vec=8, store_cm=0, + quant="none"): # fmt: on """CShuffle one GEMM2 tile into weighted BF16 rows and store them at their send-buffer index.""" kMChunks = BM // 16 @@ -111,7 +137,15 @@ def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, if const_expr(g2_bf16_lds) else None ) - row_nbytes = N_OUT * 2 + quant_fp8 = quant == "fp8_blockwise_1x32" + if quant not in ("none", "fp8_blockwise_1x32"): + raise ValueError(f"unsupported quant={quant!r}") + if quant_fp8 and g2_bf16_lds: + raise ValueError("fp8_blockwise_1x32 needs the f32 CShuffle (g2_bf16_lds=False)") + out_elem_bytes = 1 if quant_fp8 else 2 + # Row layout matches aiter's P2P scatter: N_OUT payload bytes followed by the + # N_OUT/32 e8m0 block scales, so one contiguous SDMA put ships both. + row_nbytes = N_OUT + N_OUT // 32 if quant_fp8 else N_OUT * 2 fx.barrier() @@ -140,12 +174,21 @@ def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, fx.barrier() - output_buf = ptr_buf_tensor( - send_base, - fx.BFloat16, - unit_elems=scatter_vec, - num_records_bytes=send_nbytes, - ) + # ``buf_copy_store``'s cache_modifier is a BufferCopy enum (0=cached, 2=nt), + # so sc0|sc1 cannot be reached through it -- flydsl's own type stub says so, + # and so does mori/ops/gemm_ar/_compat.py. ``buffer_store`` passes the value + # through as the instruction's aux field, so the write-through path uses it. + if const_expr(store_cm or quant_fp8): + output_buf = create_buffer_resource_from_addr( + send_base, num_records_bytes=send_nbytes + ) + else: + output_buf = ptr_buf_tensor( + send_base, + fx.BFloat16, + unit_elems=scatter_vec, + num_records_bytes=send_nbytes, + ) # Inactive lanes read safe LDS and issue a bounded OOB store. active = lane < fx.Int32(BN // scatter_vec) col = active.select(lane * fx.Int32(scatter_vec), fx.Int32(0)) @@ -182,22 +225,84 @@ def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, align=16, ) ) - pk = fx.Vector.from_elements( + weighted = fx.Vector.from_elements( [v8[i] * weight for i in range_constexpr(scatter_vec)], fx.Float32, - ).to(fx.BFloat16) - row_off = out_row * fx.Int32(row_nbytes) + n_block_idx * fx.Int32(BN * 2) - off = ((out_row >= fx.Int32(0)) & active).select( - row_off + col * fx.Int32(2), fx.Int32(send_nbytes) - ) - buf_copy_store( - output_buf, - off // fx.Int32(scatter_vec * 2), - pk, - fx.BFloat16, - unit_elems=scatter_vec, - cache_modifier=0, - ) + ) + if const_expr(not quant_fp8): + pk = weighted.to(fx.BFloat16) + row_base = out_row * fx.Int32(row_nbytes) + valid = (out_row >= fx.Int32(0)) & active + if const_expr(quant_fp8): + # One e8m0 per 32 contiguous values, i.e. per 32//scatter_vec lanes. + vals = [fx.Float32(weighted[i]) for i in range_constexpr(scatter_vec)] + local_max = fabs_f32(vals[0]) + for q in range_constexpr(1, scatter_vec): + local_max = local_max.maximumf(fabs_f32(vals[q])) + max_bits = local_max.bitcast(fx.Int32) + for xor_lane in (1, 2): + if xor_lane < 32 // scatter_vec: + remote = rocdl.ds_bpermute( + T.i32, (lane ^ fx.Int32(xor_lane)) * fx.Int32(4), max_bits + ) + local_max = local_max.maximumf(fx.Int32(remote).bitcast(fx.Float32)) + max_bits = local_max.bitcast(fx.Int32) + group_lanes = 32 // scatter_vec + leader_lane = lane & fx.Int32(~(group_lanes - 1)) + is_leader = (lane & fx.Int32(group_lanes - 1)) == fx.Int32(0) + e8m0 = fx.Int32( + rocdl.ds_bpermute( + T.i32, + leader_lane * fx.Int32(4), + _fp8_scale_for_leader(is_leader, local_max), + ) + ) + block_scale = (e8m0 << fx.Int32(23)).bitcast(fx.Float32) + pk_ty = T.vec(2, T.i16) + words = [] + for word in range_constexpr(scatter_vec // 4): + packed = fx.Vector.filled(2, 0, fx.Int16).ir_value() + for pair in range_constexpr(2): + e = word * 4 + pair * 2 + packed = rocdl.cvt_scalef32_pk_fp8_f32( + pk_ty, packed, vals[e].ir_value(), vals[e + 1].ir_value(), + block_scale.ir_value(), pair, + ) + words.append(fx.Vector(packed).bitcast(fx.Int32)[0]) + payload = fx.Vector.from_elements(words, fx.Int32) + off = valid.select( + row_base + n_block_idx * fx.Int32(BN) + col, fx.Int32(send_nbytes) + ) + buffer_store( + payload, output_buf, off, + cache_modifier=store_cm, offset_is_bytes=True, + ) + # Only the lane that owns a 32-value block writes its scale byte. + scale_off = (valid & is_leader).select( + row_base + fx.Int32(N_OUT) + n_block_idx * fx.Int32(BN // 32) + + lane // fx.Int32(group_lanes), + fx.Int32(send_nbytes), + ) + buffer_store( + e8m0.to(fx.Int8), output_buf, scale_off, + cache_modifier=store_cm, offset_is_bytes=True, + ) + else: + off = valid.select( + row_base + n_block_idx * fx.Int32(BN * 2) + col * fx.Int32(2), + fx.Int32(send_nbytes), + ) + if const_expr(store_cm): + buffer_store(pk, output_buf, off, cache_modifier=store_cm, offset_is_bytes=True) + else: + buf_copy_store( + output_buf, + off // fx.Int32(scatter_vec * 2), + pk, + fx.BFloat16, + unit_elems=scatter_vec, + cache_modifier=0, + ) # fmt: off @@ -206,7 +311,7 @@ def compile_moe_gemm2_localsort(*, model_dim: int, inter_dim: int, experts: int, HIDDEN_MAX: int = 8192, INTER_MAX: int = 8192, a_dtype: str = "fp8", SBM: int | None = None, persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, g2_ascale_pf=None, g2_spart=None, persist_strided: bool = False, g2_bf16_lds: bool = False, - skew_cu: int = 0, scatter_vec: int = 8, comm=None): + skew_cu: int = 0, scatter_vec: int = 8, comm=None, quant: str = "none"): # fmt: on """Compile GEMM2 with the rank-sorted local-output epilogue (optionally fused with SDMA).""" arch = str(get_rocm_arch() or "") @@ -233,8 +338,17 @@ def compile_moe_gemm2_localsort(*, model_dim: int, inter_dim: int, experts: int, raise ValueError("MegaMoE v2 stage2 send buffer exceeds the 32-bit buffer-resource ABI") fuse = comm is not None NC = comm.comm_wgs if fuse else 0 + # comm_wgs == 0 is a diagnostic build: the GEMM units still publish each chunk + # (s_waitcnt + barrier + atomic), but no comm wave is traced and nothing ships. if fuse and (comm.comm_wgs < 1 or comm.chunk_tiles < 1): raise ValueError("fused comm needs comm_wgs >= 1 and chunk_tiles >= 1") + # The rows are stored sc0|sc1 -- past L1 (GLC) and past L2 (SCC) -- so they + # are in memory when the store retires and the ``s_waitcnt vmcnt(0)`` in + # front of the publish is the release the copy engine needs. Leaving them in + # L2 and issuing a system-scope release instead costs a whole-L2 write-back + # per unit (``buffer_wbl2`` on gfx950): measured 2.4x on the GEMM2 kernel. + # See mori/ops/gemm_ar/_compat.py for the encoding, and why sc1 is 16 not 4. + STORE_CM = 1 | 16 if fuse else 0 N_OUT = model_dim g2_bhoist, g2_ascale_pf, g2_spart, g2_group_num, g2_m01, _g2_bf16_lds = _resolve_g2_knobs( g2_bhoist, g2_ascale_pf, g2_spart, False, False @@ -260,8 +374,13 @@ class SharedStorage: f"_p{int(persist)}cu{cu_num}s{int(persist_strided)}_pad{int(has_pad)}" f"_sk{skew_cu}" f"_bh{int(g2_bhoist)}apf{int(g2_ascale_pf)}sp{g2_group_num}x{g2_m01}" - f"_bf16lds{int(g2_bf16_lds)}_sv{scatter_vec}" - + (f"_fuse{comm.chunk_tiles}q{comm.queues}r{comm.rank}" if fuse else "") + f"_bf16lds{int(g2_bf16_lds)}_sv{scatter_vec}_{quant}" + + ( + f"_fuse{comm.chunk_tiles}q{comm.queues}w{comm.comm_wgs}" + f"sm{comm.signal_mode[0]}r{comm.rank}" + if fuse + else "" + ) ) # fmt: off @@ -325,7 +444,8 @@ def run_unit(unit_bx, m_block_idx): send_base = rocdl.readfirstlane(T.i64, fx.Int64(arg_send).ir_value()) local_sort_epilog(lds_base_i32, accm_vecs, n_block_idx, wave, lane, send_base, N_OUT=N_OUT, BM=BM, BN=BN, send_nbytes=send_nbytes, lds_outrow_off=lds_outrow_off, - lds_weight_off=lds_weight_off, g2_bf16_lds=g2_bf16_lds, scatter_vec=scatter_vec) + lds_weight_off=lds_weight_off, g2_bf16_lds=g2_bf16_lds, scatter_vec=scatter_vec, + store_cm=STORE_CM, quant=quant) # fmt: on def run_one(unit_bx, m_block_idx): @@ -333,16 +453,16 @@ def run_one(unit_bx, m_block_idx): rocdl.sched_barrier(0) run_unit(unit_bx, m_block_idx) if const_expr(fuse): - # Every thread's stores must be complete before one thread publishes - # the unit; the copy engine reads them from memory, so release at - # system scope. - rocdl.s_waitcnt(0) + # The rows went past L2 on the way out, so retiring the stores + # is all the release the copy engine needs; the barrier makes it + # the whole workgroup's stores, not just this wave's. + rocdl.s_waitcnt(vmcnt=0) fx.barrier() if tx_i32 == fx.Int32(0): - comm_ops.fence_system_release() chunk = (m_block_idx * fx.Int32(BM)) // fx.Int32(SBM * comm.chunk_tiles) comm_ops.atomic_add_agent( - fx.Int64(arg_cdone) + fx.Int64(chunk) * fx.Int64(4), fx.Int32(1) + fx.Int64(arg_cdone) + fx.Int64(chunk) * fx.Int64(CDONE_STRIDE * 4), + fx.Int32(1), ) cumsum0 = global_typed_ptr(arg_cumsum, T.i32)[0] @@ -459,9 +579,10 @@ def comm_role(cw): m_lo = ci * fx.Int32(CT * mb_per_tile) m_hi = m_lo + fx.Int32(CT * mb_per_tile) m_hi = (m_hi < total_m_blocks).select(m_hi, total_m_blocks) + # Every (m_block, n_block) unit of this chunk has stored. if is_lane0: comm_ops.spin_until_ge_i32_agent( - fx.Int64(arg_cdone) + fx.Int64(ci) * fx.Int64(4), + fx.Int64(arg_cdone) + fx.Int64(ci) * fx.Int64(CDONE_STRIDE * 4), (m_hi - m_lo) * num_n_blocks, ) comm_ops.fence_agent_acquire() @@ -476,15 +597,18 @@ def comm_role(cw): hi = tile_base_buf[t1 * fx.Int32(npes) + peer] if hi > lo: if is_lane0: - sdma_put_run(sdma, arg_win, peer, peer % fx.Int32(comm.queues), + sdma_put_run(sdma, arg_win, peer, + peer % fx.Int32(comm.queues), layout=lay, rank=rank, lo=lo, hi=hi, cco=cco) for j in range_constexpr(n_mine): peer = owner + fx.Int32(j * n_waves) if peer < fx.Int32(npes): if peer != fx.Int32(rank): if is_lane0: - sdma_finish_peer(sdma, w, arg_win, peer, peer % fx.Int32(comm.queues), - cursor_buf[peer], phase, layout=lay, rank=rank, cco=cco) + sdma_finish_peer(sdma, w, arg_win, peer, + peer % fx.Int32(comm.queues), + cursor_buf[peer], phase, layout=lay, rank=rank, + cco=cco, signal_mode=comm.signal_mode) if const_expr(fuse): if bid < fx.Int32(NC): @@ -523,7 +647,8 @@ def run_moe_gemm2_localsort(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar arg_cdone, arg_tile_base, arg_cursor, row_capacity, i32_inter, i32_hidden, stream, *, model_dim, inter_dim, experts, rank, send_nbytes, BM, SBM, HIDDEN_MAX, INTER_MAX, cu_num, BN=256, BK=256, use_nt=True, g2_bhoist=True, g2_ascale_pf=True, g2_spart=402, persist=False, - persist_cu=0, persist_strided=False, g2_bf16_lds=False, skew_cu=0, scatter_vec=8, comm=None): + persist_cu=0, persist_strided=False, g2_bf16_lds=False, skew_cu=0, scatter_vec=8, comm=None, + quant="none"): # fmt: on """Compile or reuse one configuration and launch it.""" launch_cu_num = min(cu_num, persist_cu) if persist and persist_cu > 0 else cu_num @@ -533,7 +658,7 @@ def run_moe_gemm2_localsort(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar INTER_MAX=INTER_MAX, SBM=SBM, persist=persist, cu_num=launch_cu_num, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, skew_cu=skew_cu, - scatter_vec=scatter_vec, comm=comm, + scatter_vec=scatter_vec, comm=comm, quant=quant, ) key = tuple(sorted(compile_kw.items())) launch = _LAUNCH_CACHE.get(key) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py new file mode 100644 index 0000000000..e97959ce14 --- /dev/null +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +"""What one cco SDMA channel actually moves, with no GEMM in the way. + +Same shape as the fused combine's transport -- one comm wave per peer, each +putting a contiguous run from its own window into the peer's, then draining -- +but nothing else in the kernel. Reports per-link and per-rank bandwidth for a +sweep of transfer sizes and channel counts, so the fused kernel's exposed +transfer time can be checked against what the hardware will actually do. + + MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 sdmabw.py \\ + --mb 0.36,0.72,1.44,2.89,5.78 --queues 1,2 +""" + +import argparse +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import range_constexpr, rocdl +from flydsl.expr.typing import T + +import torch +import torch.distributed as dist + +import mori.shmem as ms +from mori.cco import GDA_CONNECTION_NONE, CCODevCommRequirements, Communicator + +WAVE = 64 + + +def build_xfer(*, npes, rank, comm_wgs, queues, split, puts): + """One wave per peer: ``puts`` back-to-back runs, then drain every channel.""" + import mori.cco.device.flydsl as cco + + n_waves = comm_wgs * 4 + + @flyc.kernel( + name=f"sdmabw_p{npes}_w{comm_wgs}q{queues}s{split}n{puts}_r{rank}", + known_block_size=[WAVE * 4, 1, 1], + ) + def xfer(dev_comm: fx.Int64, win: fx.Int64, i64_bytes: fx.Int64, i64_stride: fx.Int64): + tid = fx.thread_idx.x + lane = tid % fx.Int32(WAVE) + wave = rocdl.readfirstlane(T.i32, tid // fx.Int32(WAVE)) + owner = fx.block_idx.x * fx.Int32(4) + wave + if owner < fx.Int32(npes): + sdma = cco.DevComm(dev_comm).sdma() + if owner != fx.Int32(rank): + if lane == fx.Int32(0): + for p in range_constexpr(puts): + # Distinct source/destination offsets per put so the copy + # engine cannot coalesce the repeats into one transfer. + off = fx.Int64(p) * i64_stride + for q in range_constexpr(split): + part = i64_bytes // fx.Int64(split) + sdma.put( + owner, + win, + off + fx.Int64(q) * part, + win, + off + fx.Int64(q) * part, + part, + (p + q) % queues if split > 1 else p % queues, + coop=cco.CoopScope.THREAD, + ) + sdma.quiet(owner) + + @flyc.jit + def launch(dev_comm: fx.Int64, win: fx.Int64, i64_bytes: fx.Int64, + i64_stride: fx.Int64, stream: fx.Stream): + xfer(dev_comm, win, i64_bytes, i64_stride).launch( + grid=(comm_wgs, 1, 1), block=(WAVE * 4, 1, 1), stream=stream + ) + + return launch + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--mb", default="0.36,0.72,1.44,2.89,5.78", help="MB per put") + parser.add_argument("--queues", default="1,2", help="channels per peer") + parser.add_argument("--split", default="1", help="split one put across channels") + parser.add_argument("--puts", type=int, default=8, help="back-to-back puts per peer") + parser.add_argument("--iters", type=int, default=50) + args = parser.parse_args() + + sizes = [float(v) for v in args.mb.split(",")] + queue_list = [int(v) for v in args.queues.split(",")] + splits = [int(v) for v in args.split.split(",")] + + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + dist.init_process_group("gloo") + import torch._C._distributed_c10d as c10d + + c10d._register_process_group("default", dist.group.WORLD) + ms.shmem_torch_process_group_init("default") + + try: + max_bytes = int(max(sizes) * (1 << 20)) + stride = (max_bytes + 4095) // 4096 * 4096 + total = stride * args.puts * 2 # src half and dst half never overlap + rows = [] + for queues in queue_list: + uid = Communicator.get_unique_id() if rank == 0 else None + holder = [uid] + dist.broadcast_object_list(holder, src=0) + comm = Communicator.init(world, rank, holder[0]) + mem = comm.alloc_mem(total) + win = comm.register_window(mem.ptr, mem.size) + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + reqs.sdma_queue_count = queues + dev_comm = comm.create_dev_comm(reqs) + torch.cuda.synchronize() + comm.barrier() + + comm_wgs = max(1, (world + 3) // 4) + for split in splits: + if split > queues: + continue + fn = build_xfer( + npes=world, rank=rank, comm_wgs=comm_wgs, + queues=queues, split=split, puts=args.puts, + ) + for mb in sizes: + nbytes = int(mb * (1 << 20)) + s_fx = fx.Stream(torch.cuda.current_stream().cuda_stream) + + def once(): + from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled + + _run_compiled( + fn, fx.Int64(dev_comm.ptr), fx.Int64(win.handle), + fx.Int64(nbytes), fx.Int64(stride), s_fx, + ) + + once() + torch.cuda.synchronize() + dist.barrier() + for _ in range(3): + once() + torch.cuda.synchronize() + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(args.iters): + once() + end.record() + torch.cuda.synchronize() + ms_per = start.elapsed_time(end) / args.iters + # Each rank sends to world-1 peers, puts runs of nbytes each. + per_link = nbytes * args.puts / (ms_per * 1e-3) / 1e9 + per_rank = per_link * (world - 1) + t = torch.tensor([per_link, per_rank], dtype=torch.float64) + dist.all_reduce(t) + t /= world + rows.append((queues, split, mb, ms_per, float(t[0]), float(t[1]))) + if rank == 0: + print( + f" queues={queues} split={split} put={mb:6.2f} MB " + f"x{args.puts} {ms_per:8.4f} ms " + f"per-link {t[0]:7.1f} GB/s per-rank {t[1]:7.1f} GB/s", + flush=True, + ) + comm.finalize() if hasattr(comm, "finalize") else None + + if rank == 0: + print("\n| queues | split | MB/put | ms | per-link GB/s | per-rank GB/s |") + print("|---:|---:|---:|---:|---:|---:|") + for q, sp, mb, t, pl, pr in rows: + print(f"| {q} | {sp} | {mb:.2f} | {t:.4f} | {pl:.1f} | {pr:.1f} |") + print("\nSUCCESS", flush=True) + finally: + try: + ms.shmem_finalize() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py index ac9f297c2f..e24b08a818 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py @@ -56,6 +56,19 @@ are ``[0, npes*cap)`` and ``R`` rows ``[npes*cap, 2*npes*cap)``. My own run never travels; it is read from ``S`` in place. +``cap`` is the rows reserved for ONE (src, dst) pair, not the whole token-slot +space: a sender writes into the peer's ``R`` at a statically known offset with +no handshake, so the reservation has to be per pair. Sizing it at the worst case +(every token of a peer routing all of top-k here) costs ``npes`` times the +window and put bs=8192 over the 4 GiB buffer-resource range; the real counts sit +at ``slot_cap/npes``, so ``row_slack`` reserves that share with headroom and the +plan kernel refuses to place a row past it. + +Rows are transported as fp8 with one e8m0 scale per 32 values, laid out as +aiter's P2P scatter lays them out -- payload then scales, inline -- so one +contiguous put still carries a whole row. bf16 doubles the bytes and the +transfer stops fitting under the GEMM. + Both persistent kernels use grid-wide barriers, so their grid must be fully resident (one block per CU is). """ @@ -64,16 +77,30 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl.expr import arith, range_constexpr, rocdl -from flydsl.expr.typing import T +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import Int8, T +from flydsl.compiler.ast_rewriter import ASTRewriter from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops + +# A device-side helper defined outside a kernel body is not rewritten by +# @flyc.kernel, so a data-dependent ``if`` inside it would stay a plain Python +# branch on a traced value. flyc.jit turns such a body into scf ops with +# ASTRewriter.transform; apply the same thing by hand on aiter builds that do +# not export comm_ops.traced yet. +if not hasattr(comm_ops, "traced"): + comm_ops.traced = ASTRewriter.transform + from aiter.ops.flydsl.kernels.buffer_ops import ( buffer_load, buffer_store, create_buffer_resource_from_addr, ) -from aiter.ops.flydsl.kernels.mxfp4_gemm_common import global_typed_ptr +from aiter.ops.flydsl.kernels.mxfp4_gemm_common import ( + global_typed_ptr, + lds_typed_ptr, +) from aiter.ops.flydsl.kernels.tensor_shim import ( buf_copy_load, buf_copy_store, @@ -172,7 +199,8 @@ def grid_barrier(bar_addr, nblocks, tid): fx.barrier() -def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap): +def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap, slot_cap, + local_experts): """Kernel 1: give every valid sorted row a stable index in its rank's run. * count: ``tile_cnt[tile][pe]`` = rows of the sort tile bound for ``pe`` @@ -181,18 +209,43 @@ def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap): * assign: ``out_row = pe*cap + tile_base + (earlier rows of this tile for pe)``, and ``idx_send[out_row] = origin_lid*topk + k`` - The phases are separated by grid barriers inside one persistent launch. Row - loops read the tile's SBM rows from global memory instead of using LDS scans; - SBM is small and the rows are L2-resident. + Only ONE grid barrier, between count and scan. The scan has to see every + block's counts, so that one is unavoidable; a second barrier after it is not, + because the thing being scanned is tiny -- at most ``cap/SBM`` tiles times + ``npes``, i.e. 2048 int32 -- so every block redundantly scans the whole array + into its own LDS instead of one block scanning it and the rest waiting to read + the result. That also replaces the serial prefix sum (8 threads walking + ``tiles`` dependent global loads, the single most expensive thing in this + kernel) with a Hillis-Steele scan in LDS. + + Row loops read the tile's SBM rows from global memory instead of using LDS + scans; SBM is small and the rows are L2-resident. """ dec = dict(log2_max_tok=log2_max_tok, recv_cap=recv_cap, topk=topk, npes=npes) + # Bound on the tile count, so the scan buffers can live in LDS. Tiles cover + # the sorted rows, so this is sized from the token-slot space, NOT from + # ``cap`` -- ``cap`` is the per-peer row reservation and is npes times + # smaller. Stage1 also pads each local expert's rows up to an SBM boundary, + # so num_valid runs past slot_cap by up to local_experts * (SBM - 1). + TILES_MAX = (slot_cap + SBM - 1) // SBM + local_experts + 1 + SCAN_N = (TILES_MAX + 1) * npes + STEPS = max(1, (TILES_MAX - 1).bit_length()) + lds_bytes = 2 * SCAN_N * 4 + if lds_bytes > 160 * 1024: + raise ValueError(f"plan scan needs {lds_bytes} B of LDS, over the 160 KiB budget") + + @fx.struct + class ScanShared: + buf: fx.Array[Int8, lds_bytes, 16] + # fmt: off - @flyc.kernel(name=f"mega_rank_sort_plan_sbm{SBM}_k{topk}_p{npes}", known_block_size=[BLOCK, 1, 1]) + @flyc.kernel(name=f"mega_rank_sort_plan_sbm{SBM}_k{topk}_p{npes}_c{cap}", + known_block_size=[BLOCK, 1, 1]) def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, - i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32): + arg_over: fx.Int64, i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32): # fmt: on tid = fx.thread_idx.x bid = fx.block_idx.x @@ -207,6 +260,15 @@ def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, out_buf = ptr_buf_tensor(arg_out_row, fx.Int32) cdone = ptr_buf_tensor(arg_cdone, fx.Int32) + lds = fx.SharedAllocator().allocate(ScanShared).peek() + lds_i32 = fx.Int32(fx.ptrtoint(lds.buf.ptr)) + + def sc(half, i): + """LDS i32 slot ``i`` of ping(0)/pong(1).""" + return lds_typed_ptr( + lds_i32 + (fx.Int32(half * SCAN_N) + i) * fx.Int32(4), T.i32, align=4 + ) + # The GEMM2 kernel's per-chunk completion counters start from zero every call. for g in range(gid, i32_cdone_n, stride): cdone[g] = fx.Int32(0) @@ -222,17 +284,38 @@ def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, tile_cnt[g] = cnt grid_barrier(arg_bar, i32_nblocks, tid) + # Inclusive Hillis-Steele over the tile axis, all npes lanes at once, + # run identically by every block. + n = tiles * fx.Int32(npes) + for i in range(tid, n, fx.Int32(BLOCK)): + fx.ptr_store(tile_cnt[i], sc(0, i)) + fx.barrier() + for step in range_constexpr(STEPS): + src, dst = step % 2, (step + 1) % 2 + back = fx.Int32((1 << step) * npes) + for i in range(tid, n, fx.Int32(BLOCK)): + v = fx.Int32(fx.ptr_load(sc(src, i))) + v = (i >= back).select( + v + fx.Int32(fx.ptr_load(sc(src, i - back))), v + ) + fx.ptr_store(v, sc(dst, i)) + fx.barrier() + final = STEPS % 2 + # Exclusive = inclusive - own count; the extra row holds the totals the + # comm waves read as a chunk's end. if bid == fx.Int32(0): + for i in range(tid, n, fx.Int32(BLOCK)): + tile_base[i] = fx.Int32(fx.ptr_load(sc(final, i))) - tile_cnt[i] if tid < fx.Int32(npes): - run = arith.constant(0) - for t in range(fx.Int32(0), tiles, fx.Int32(1)): - tile_base[t * fx.Int32(npes) + tid] = run - run = run + tile_cnt[t * fx.Int32(npes) + tid] - tile_base[tiles * fx.Int32(npes) + tid] = run - ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = run - grid_barrier(arg_bar, i32_nblocks, tid) + total = fx.Int32( + fx.ptr_load(sc(final, (tiles - fx.Int32(1)) * fx.Int32(npes) + tid)) + ) + capped = (total < fx.Int32(cap)).select(total, fx.Int32(cap)) + tile_base[tiles * fx.Int32(npes) + tid] = capped + ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = capped idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) + final = STEPS % 2 for g in range(gid, tiles * fx.Int32(SBM), stride): tile = g // fx.Int32(SBM) row = g - tile * fx.Int32(SBM) @@ -245,19 +328,30 @@ def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, v2, pe2, _, _ = _decode(stids[base + fx.Int32(rp)], **dec) earlier = (fx.Int32(rp) < row) & v2 & (pe2 == dest_pe) within = within + earlier.select(fx.Int32(1), fx.Int32(0)) - out_row = dest_pe * fx.Int32(cap) + tile_base[tile * fx.Int32(npes) + dest_pe] + within - out_buf[g] = out_row - idx_buf[out_row] = dest_lid * fx.Int32(topk) + s + slot = tile * fx.Int32(npes) + dest_pe + # This block scanned the whole array itself, so the prefix is + # in its own LDS -- no barrier between the scan and this. + excl = fx.Int32(fx.ptr_load(sc(final, slot))) - tile_cnt[slot] + place = excl + within + # Past the per-pair reservation there is nowhere to put this + # row. Drop it rather than write into the next peer's run, + # and leave a mark the host can assert on. + if place >= fx.Int32(cap): + comm_ops.atomic_add_agent(arg_over, fx.Int32(1)) + else: + out_row = dest_pe * fx.Int32(cap) + place + out_buf[g] = out_row + idx_buf[out_row] = dest_lid * fx.Int32(topk) + s # fmt: off @flyc.jit def launch(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, - i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32, stream: fx.Stream): + arg_over: fx.Int64, i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32, stream: fx.Stream): # fmt: on plan(arg_stids, arg_trb, arg_num_valid, arg_tile_cnt, arg_tile_base, arg_cursor, - arg_out_row, arg_idx, arg_cdone, arg_bar, i32_nblocks, i32_cdone_n).launch( + arg_out_row, arg_idx, arg_cdone, arg_bar, arg_over, i32_nblocks, i32_cdone_n).launch( grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) return launch @@ -273,7 +367,16 @@ def sdma_wait_consumed(w, peer, phase, *, layout, rank): def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): - """Ship rows ``[lo, hi)`` of my run for ``peer`` (data only; idx goes once at the end).""" + """Ship rows ``[lo, hi)`` of my run for ``peer`` (data only; idx goes at the end). + + One put for the whole run. Puts are asynchronous -- each reserves a slot, + writes the copy descriptor and rings the doorbell -- but a queue is a ring + its engine drains in order, so splitting a run across channels only buys + parallelism if those channels sit on different engines. KFD recommends a + single SDMA engine per peer link on this topology (one bit per link in + ``recommended_sdma_engine_id_mask``), so they do not: splitting a 1.44 MB run + in two measured 3% slower, rotating whole runs across two channels 2%. + """ rb, cap = layout.row_bytes, layout.cap sdma.put( peer, @@ -289,31 +392,60 @@ def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): @comm_ops.traced -def sdma_finish_peer(sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco): - """After the last data put: ship idx, drain the queue, publish ``(cnt << 32) | phase``.""" - if cnt > fx.Int32(0): +def sdma_finish_peer(sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco, + signal_mode="remote"): + """Ship idx, tell ``peer`` its rows have landed, and publish ``cnt``. + + Two ways to say "landed": + + ``quiet`` drain our queue here, then store the flag. The drain waits for + the last chunk to reach the peer, and that wait sits at the tail + of the GEMM2 kernel, after the last tile -- 42us of a 772us + kernel at bs=4096, CT=16. + ``remote`` hang the copy engine's own trailing ATOMIC on this put, so the + engine bumps a counter on the peer once the bytes are there. + Nothing is drained: the comm wave issues and leaves, and the + peer's combine -- which was going to wait for us anyway -- polls + the counter. Worth 6-12us end to end. + + ``remote`` needs this put to be the LAST op on the channel the data went out + on. It is: everything for ``peer`` uses channel ``peer % queues``, and a + queue is a FIFO its engine drains in order, so the trailing ATOMIC fires + only once every chunk ahead of it has landed. + + The flag carries ``cnt`` either way. Under ``remote`` it no longer says + anything about the data, so it needs no release -- the signal does that. + """ + if signal_mode not in ("quiet", "remote"): + raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") + remote = signal_mode == "remote" + # Under `remote` this put carries the signal, so it is issued even when + # there is no idx to ship -- 4 bytes the peer never reads. + if const_expr(remote) or cnt > fx.Int32(0): sdma.put( peer, win, fx.Int64(layout.off_idx_land + rank * layout.cap * 4), win, fx.Int64(layout.off_idx_send) + fx.Int64(peer) * fx.Int64(layout.cap * 4), - fx.Int64(cnt) * fx.Int64(4), + fx.Int64((cnt > fx.Int32(0)).select(cnt, fx.Int32(1))) * fx.Int64(4), queue, coop=cco.CoopScope.THREAD, + remote_signal=True if remote else None, ) - sdma.quiet_queue(peer, queue) - # The bytes were moved by the copy engine, not by this CU; the fence only keeps - # the flag store from being hoisted above the drain, which is the ordering the - # receiver relies on. - comm_ops.fence_system_release() + if const_expr(not remote): + sdma.quiet_queue(peer, queue) + # The bytes moved by the copy engine, not by this CU, so nothing of ours + # has to be written back -- this only keeps the flag store from being + # hoisted above the drain. + rocdl.s_waitcnt(vmcnt=0) comm_ops.store_i64_global_system( fx.Int64(w.lsa_ptr(peer, layout.off_flags)) + fx.Int64(rank * 8), (fx.Int64(cnt) << fx.Int64(32)) | fx.Int64(phase), ) -def build_combine(*, layout, rank, hidden, topk): +def build_combine(*, layout, rank, hidden, topk, quant="none", signal_mode="remote"): """Kernel 3: finish the transfer and reduce. One persistent launch (grid resident), in phases: @@ -331,14 +463,22 @@ def build_combine(*, layout, rank, hidden, topk): if hidden % 8: raise ValueError(f"hidden must be a multiple of 8, got {hidden}") + quant_fp8 = quant == "fp8_blockwise_1x32" + if quant not in ("none", "fp8_blockwise_1x32"): + raise ValueError(f"unsupported quant={quant!r}") + if quant_fp8 and hidden % 32: + raise ValueError(f"fp8_blockwise_1x32 needs hidden % 32 == 0, got {hidden}") npes, cap = layout.npes, layout.cap units = hidden // 8 chunks = (units + BLOCK - 1) // BLOCK + row_bytes = layout.row_bytes # fmt: off - @flyc.kernel(name=f"mega_combine_p{npes}_c{cap}_h{hidden}_k{topk}_r{rank}", known_block_size=[BLOCK, 1, 1]) + @flyc.kernel(name=f"mega_combine_p{npes}_c{cap}_h{hidden}_k{topk}_{quant}_{signal_mode}" + f"_r{rank}", known_block_size=[BLOCK, 1, 1]) def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, - arg_bar: fx.Int64, arg_fin: fx.Int64, i32_nblocks: fx.Int32, i32_tokens: fx.Int32): + arg_bar: fx.Int64, arg_fin: fx.Int64, arg_dev_comm: fx.Int64, + i32_nblocks: fx.Int32, i32_tokens: fx.Int32): # fmt: on tid = fx.thread_idx.x bid = fx.block_idx.x @@ -353,9 +493,19 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. # A. the peers' data has landed if tid < fx.Int32(npes): if tid != fx.Int32(rank): + # The flag carries cnt (and, under "quiet", data-readiness too). comm_ops.spin_until_ge_i32_system( flags_base + fx.Int64(tid) * fx.Int64(8), phase, acquire=True ) + if const_expr(signal_mode == "remote"): + # The sender did not drain its queue; its copy engine bumps + # this counter once per call when the last put lands. One + # signal per (src, queue) per call, so the running total is + # the call number. + cco.DevComm(arg_dev_comm).sdma().wait_signal( + tid, 0, fx.Int64(phase) + ) + comm_ops.fence_system_acquire() fx.barrier() # B. slot -> landing row @@ -379,7 +529,15 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. grid_barrier(arg_bar, i32_nblocks, tid) # C. top-k reduce - rows = ptr_buf_tensor(_uniform_i64(w.lsa_ptr(rank, layout.off_s)), fx.BFloat16, unit_elems=8) + s_base = _uniform_i64(w.lsa_ptr(rank, layout.off_s)) + if const_expr(quant_fp8): + # Payload is 8 fp8 bytes (2 i32) per thread; the 32-value block they + # fall in is u // 4, so each thread reads its own scale byte rather + # than electing a leader and broadcasting. + rows = ptr_buf_tensor(s_base, fx.Int32, unit_elems=2) + scales = ptr_buf_tensor(s_base, fx.Int8) + else: + rows = ptr_buf_tensor(s_base, fx.BFloat16, unit_elems=8) out = ptr_buf_tensor(arg_out, fx.BFloat16, unit_elems=8) for tok in range(bid, i32_tokens, i32_nblocks): slots = [] @@ -392,9 +550,35 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. for k in range_constexpr(topk): ok = slots[k] >= fx.Int32(0) safe = ok.select(slots[k], fx.Int32(0)) - raw = buf_copy_load(rows, safe * fx.Int32(units) + u, fx.BFloat16, unit_elems=8) keep = ok.select(fx.Float32(1.0), fx.Float32(0.0)) - acc = acc + fx.Vector(raw).to(fx.Float32) * keep + if const_expr(quant_fp8): + row0 = safe * fx.Int32(row_bytes) + raw = fx.Vector( + buf_copy_load( + rows, (row0 + u * fx.Int32(8)) // fx.Int32(8), + fx.Int32, unit_elems=2, + ) + ) + sc = scales[row0 + fx.Int32(hidden) + u // fx.Int32(4)] + scale = ( + fx.Uint8(sc).to(fx.Uint32).bitcast(fx.Int32) + << fx.Int32(23) + ).bitcast(fx.Float32) + v2f32 = T.vec(2, T.f32) + parts = [] + for word in range_constexpr(2): + lo = cvt_pk_f32_fp8(res=v2f32, src=raw[word], word_sel=False) + hi = cvt_pk_f32_fp8(res=v2f32, src=raw[word], word_sel=True) + parts.append(lo.shuffle(hi, [0, 1, 2, 3])) + vals = fx.Vector(parts[0]).shuffle( + fx.Vector(parts[1]), [0, 1, 2, 3, 4, 5, 6, 7] + ) + acc = acc + fx.Vector(vals) * (scale * keep) + else: + raw = buf_copy_load( + rows, safe * fx.Int32(units) + u, fx.BFloat16, unit_elems=8 + ) + acc = acc + fx.Vector(raw).to(fx.Float32) * keep buf_copy_store( out, tok * fx.Int32(units) + u, acc.to(fx.BFloat16), fx.BFloat16, unit_elems=8 ) @@ -422,10 +606,11 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. # fmt: off @flyc.jit def launch(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, - arg_bar: fx.Int64, arg_fin: fx.Int64, i32_nblocks: fx.Int32, i32_tokens: fx.Int32, - stream: fx.Stream): + arg_bar: fx.Int64, arg_fin: fx.Int64, arg_dev_comm: fx.Int64, + i32_nblocks: fx.Int32, i32_tokens: fx.Int32, stream: fx.Stream): # fmt: on - combine(win, arg_cursor, arg_inv, arg_out, arg_bar, arg_fin, i32_nblocks, i32_tokens).launch( + combine(win, arg_cursor, arg_inv, arg_out, arg_bar, arg_fin, arg_dev_comm, + i32_nblocks, i32_tokens).launch( grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) return launch @@ -461,7 +646,16 @@ class SortedCombine: Needs ``MORI_ENABLE_SDMA=1`` and a mori built with SDMA. """ - def __init__(self, moe, queues=8, chunk_tiles=4, comm_wgs=None): + def __init__( + self, + moe, + queues=8, + chunk_tiles=4, + comm_wgs=None, + row_slack=1.5, + signal_mode="quiet", + quant="none", + ): import torch import torch.distributed as dist from mori.cco import ( @@ -478,8 +672,40 @@ def __init__(self, moe, queues=8, chunk_tiles=4, comm_wgs=None): self.hidden = int(cfg.hidden_dim) self.topk = int(cfg.num_experts_per_token) self.max_tok = int(cfg.max_num_inp_token_per_rank) - self.cap = self.max_tok * self.topk - self.layout = ArenaLayout(self.npes, self.cap, self.hidden * 2) + # Token-slot space: one entry per (token, k). This is what idx values and + # ``inv`` are indexed by. + self.slot_cap = self.max_tok * self.topk + # Row space: how many rows one rank reserves for ONE peer, in both the + # send buffer and the peer's landing buffer. A sender writes into the + # peer's landing buffer at a statically known offset, with no handshake, + # so this has to be reserved per (src, dst) pair rather than packed. The + # worst case is slot_cap (every token of the peer routing all of top-k + # here), but the real counts sit at slot_cap/npes -- 4012..4221 against + # an even share of 4096, measured at bs=4096 -- so reserve that share + # times ``row_slack`` and have the plan kernel refuse to write past it. + # Reserving the worst case instead costs npes times the window and is + # what put bs=8192 over the 4 GiB buffer-resource range. + self.row_slack = row_slack + if signal_mode not in ("quiet", "remote"): + raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") + if signal_mode == "remote" and queues != 1: + raise ValueError( + "signal_mode='remote' hangs the completion signal on the last put " + f"of one channel; queues={queues} spreads the data over channels " + "it would not cover" + ) + self.signal_mode = signal_mode + share = -(-self.slot_cap // self.npes) + self.cap = -(-int(share * row_slack) // 128) * 128 + self.quant = quant + # fp8 rows carry their e8m0 block scales inline, so one contiguous SDMA + # put still ships a whole row. + row_bytes = ( + self.hidden + self.hidden // 32 + if quant == "fp8_blockwise_1x32" + else self.hidden * 2 + ) + self.layout = ArenaLayout(self.npes, self.cap, row_bytes) self.layout.check() dev = torch.device("cuda", self.rank) lay = self.layout @@ -508,13 +734,20 @@ def __init__(self, moe, queues=8, chunk_tiles=4, comm_wgs=None): self.tile_cnt = torch.zeros(max_tiles * self.npes, **i32) self.tile_base = torch.zeros((max_tiles + 1) * self.npes, **i32) self.cursor = torch.zeros(self.npes, **i32) - self.cdone = torch.zeros(max_tiles + 1, **i32) + # One completion counter per chunk, each on its own cache line; the plan + # kernel zeroes the prefix a launch can touch. + from moe_gemm2_localsort import CDONE_STRIDE + + self.cdone_stride = CDONE_STRIDE + self.cdone = torch.zeros((max_tiles + 1) * CDONE_STRIDE, **i32) # Never-reset counters for the two grid barriers and the combine's last-block count. self.plan_bar = torch.zeros(1, **i32) + # Rows the per-pair reservation had no room for; must stay 0. + self.overflow = torch.zeros(1, **i32) self.combine_bar = torch.zeros(1, **i32) self.combine_fin = torch.zeros(1, **i32) # -1 means "no row for this slot"; the combine kernel puts it back after reading. - self.inv = torch.full((self.cap,), -1, **i32) + self.inv = torch.full((self.slot_cap,), -1, **i32) self.out = torch.empty( self.max_tok * self.hidden, dtype=torch.bfloat16, device=dev ) @@ -531,10 +764,12 @@ def __init__(self, moe, queues=8, chunk_tiles=4, comm_wgs=None): queues=queues, chunk_tiles=chunk_tiles, comm_wgs=comm_wgs, + signal_mode=signal_mode, ) self._plans = {} self._combine = build_combine( - layout=lay, rank=self.rank, hidden=self.hidden, topk=self.topk + layout=lay, rank=self.rank, hidden=self.hidden, topk=self.topk, quant=quant, + signal_mode=signal_mode, ) def _plan_for(self, sbm): @@ -550,6 +785,8 @@ def _plan_for(self, sbm): recv_cap=int(self.moe.max_recv), npes=self.npes, cap=self.cap, + slot_cap=self.slot_cap, + local_experts=int(self.moe.w2.shape[0]), ) return self._plans[sbm] @@ -559,6 +796,36 @@ def set_chunk_tiles(self, chunk_tiles): self.comm_spec = replace(self.comm_spec, chunk_tiles=chunk_tiles) + def _cdone_used(self, config): + """Prefix of cdone a launch at this config can touch; the rest stays zero.""" + spec = self.comm_spec + ct = 1 if spec is None else spec.chunk_tiles + n_chunks = (self.max_tiles + ct - 1) // ct + 1 + return min(self.cdone.numel(), n_chunks * self.cdone_stride) + + def _launch_plan(self, config, s_fx): + import flydsl.expr as fx + from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled + + op, lay = self.moe._s1_op, self.layout + _run_compiled( + self._plan_for(int(config.stage1.sort_block_m)), + fx.Int64(op.srcmap_em.data_ptr()), + fx.Int64(op.tile_row_base.data_ptr()), + fx.Int64(op.num_valid.data_ptr()), + fx.Int64(self.tile_cnt.data_ptr()), + fx.Int64(self.tile_base.data_ptr()), + fx.Int64(self.cursor.data_ptr()), + fx.Int64(self.out_row.data_ptr()), + fx.Int64(self.mem.ptr + lay.off_idx_send), + fx.Int64(self.cdone.data_ptr()), + fx.Int64(self.plan_bar.data_ptr()), + fx.Int64(self.overflow.data_ptr()), + fx.Int32(self.nblocks), + fx.Int32(self._cdone_used(config)), + s_fx, + ) + def _launch_gemm2(self, config, s_fx): import flydsl.expr as fx from moe_gemm2_localsort import run_moe_gemm2_localsort @@ -604,11 +871,18 @@ def _launch_gemm2(self, config, s_fx): g2_ascale_pf=stage2.ascale_prefetch, g2_spart=stage2.spatial_partition, persist=stage2.persist, + # Pipeline depth: each workgroup runs ceil(total_m_blocks/persist_cu) + # iterations, and that -- not chunk_tiles -- is how many batches the + # GEMM hands the comm waves. Strided order also decides *which* + # m_blocks finish first: strided completes them roughly in order, so + # a low chunk is whole early; contiguous spreads each slot over its + # own range, so no chunk is whole until near the end. persist_cu=stage2.persist_cu, persist_strided=stage2.persist_strided, skew_cu=stage2.skew_cu, - g2_bf16_lds=stage2.bf16_lds, + g2_bf16_lds=False if self.quant != "none" else stage2.bf16_lds, comm=self.comm_spec, + quant=self.quant, ) def run(self, run_tokens, config=None, profile=None): @@ -629,22 +903,7 @@ def mark(stage): profile.append((stage, ev)) mark("start") - _run_compiled( - self._plan_for(int(config.stage1.sort_block_m)), - fx.Int64(op.srcmap_em.data_ptr()), - fx.Int64(op.tile_row_base.data_ptr()), - fx.Int64(op.num_valid.data_ptr()), - fx.Int64(self.tile_cnt.data_ptr()), - fx.Int64(self.tile_base.data_ptr()), - fx.Int64(self.cursor.data_ptr()), - fx.Int64(self.out_row.data_ptr()), - fx.Int64(self.mem.ptr + lay.off_idx_send), - fx.Int64(self.cdone.data_ptr()), - fx.Int64(self.plan_bar.data_ptr()), - fx.Int32(self.nblocks), - fx.Int32(self.cdone.numel()), - s_fx, - ) + self._launch_plan(config, s_fx) mark("plan") self._launch_gemm2(config, s_fx) mark("gemm2+sdma") @@ -656,6 +915,7 @@ def mark(stage): fx.Int64(self.out.data_ptr()), fx.Int64(self.combine_bar.data_ptr()), fx.Int64(self.combine_fin.data_ptr()), + fx.Int64(self.dev_comm.ptr), fx.Int32(self.nblocks), fx.Int32(run_tokens), s_fx, From ac0a306d7d035ec8ab3585937b0b8415730d78e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=A0=C3=A7kawhil-amd?= Date: Wed, 28 Oct 2026 03:49:40 +0000 Subject: [PATCH 3/9] perf(cco/example): cut the SDMA MoE plan kernel from 43 to 21 us at 8192 tokens/rank Three serial costs in the rank-sort plan kernel, found by a GPU-side phase ladder (rocprof) rather than host timing: - assign re-decoded the whole sort tile for every row (O(SBM) per row, ~12 us). Count and assign now run one block per tile with one thread per row; a wave ballot per pe plus a popcount gives a row's rank among earlier rows for its pe, and per-wave counts cross waves through a few LDS words. - the grid barrier had every block atomic-add one address, which serialises in L2 (~7 us). Each block now bumps its own cache-line slot and polls all slots. - the Hillis-Steele scan ran redundantly in every block (10 barriers, ~15 us). Each block now sums the earlier tiles' counts for the tiles it owns with one parallel pass and a short LDS reduction, and writes tile_base/cursor itself. Outputs (out_row, idx, tile_cnt, tile_base, cursor) are bit-identical to the previous kernel on the same stage1 output at 8192/512/128 tokens per rank, and the bench relL2 stays 0 from 128 to 16384. stage2+combine vs aiter on 8x MI355: 1.20x at 8192 (was 1.17x), 1.31x at 16384, 0.72x at 128 (was 0.65x). plan_bar is now one 128 B slot per block, so a plan kernel and an older one must not share it. Co-Authored-By: Claude Code --- .../11_flydsl_mega_moe_sdma/sorted_combine.py | 242 +++++++++++------- 1 file changed, 154 insertions(+), 88 deletions(-) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py index e24b08a818..e7357e99f9 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py @@ -108,6 +108,8 @@ ) BLOCK = 256 +WAVE = 64 +BAR_STRIDE = 32 # ints per flag_barrier slot: one 128 B cache line MIN_SBM = 16 _ALIGN = 128 @@ -199,44 +201,82 @@ def grid_barrier(bar_addr, nblocks, tid): fx.barrier() -def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap, slot_cap, - local_experts): +@comm_ops.traced +def flag_barrier(bar_addr, nblocks, tid, bid): + """All blocks of a resident grid rendezvous without a contended atomic. + + Block ``i`` bumps its own cache-line-sized slot ``bar[i]`` and every block polls + all slots, so arrival costs one uncontended atomic instead of ``nblocks`` of them + queueing on one address. Slots are never reset: a block's target is its own slot's + value at entry plus one, and every block bumps once per launch. + """ + target = fx.Int32( + comm_ops.load_i32_global_agent( + bar_addr + fx.Int64(bid) * fx.Int64(BAR_STRIDE * 4) + ) + ) + fx.Int32(1) + rocdl.s_waitcnt(0) + fx.barrier() + if tid == fx.Int32(0): + comm_ops.fence_agent_release() + comm_ops.atomic_add_agent( + bar_addr + fx.Int64(bid) * fx.Int64(BAR_STRIDE * 4), fx.Int32(1) + ) + for i in range(tid, nblocks, fx.Int32(BLOCK)): + comm_ops.spin_until_ge_i32_agent( + bar_addr + fx.Int64(i) * fx.Int64(BAR_STRIDE * 4), target + ) + comm_ops.fence_agent_acquire() + fx.barrier() + + +def build_rank_sort_plan( + *, SBM, topk, log2_max_tok, recv_cap, npes, cap, slot_cap, local_experts +): """Kernel 1: give every valid sorted row a stable index in its rank's run. * count: ``tile_cnt[tile][pe]`` = rows of the sort tile bound for ``pe`` - * scan: ``tile_base[tile][pe]`` = rows for ``pe`` in earlier tiles (so the + * prefix: ``tile_base[tile][pe]`` = rows for ``pe`` in earlier tiles (so the rows of a tile range are contiguous in the run), ``cursor[pe]`` = total * assign: ``out_row = pe*cap + tile_base + (earlier rows of this tile for pe)``, and ``idx_send[out_row] = origin_lid*topk + k`` - Only ONE grid barrier, between count and scan. The scan has to see every - block's counts, so that one is unavoidable; a second barrier after it is not, - because the thing being scanned is tiny -- at most ``cap/SBM`` tiles times - ``npes``, i.e. 2048 int32 -- so every block redundantly scans the whole array - into its own LDS instead of one block scanning it and the rest waiting to read - the result. That also replaces the serial prefix sum (8 threads walking - ``tiles`` dependent global loads, the single most expensive thing in this - kernel) with a Hillis-Steele scan in LDS. - - Row loops read the tile's SBM rows from global memory instead of using LDS - scans; SBM is small and the rows are L2-resident. + One rendezvous, between count and prefix; the prefix needs every earlier tile's + count. It is a ``flag_barrier`` (no single contended atomic). After it, a block + does not scan the whole ``tile_cnt`` array: for each tile it owns it sums the + counts of the earlier tiles with one pass of the block's threads (thread ``t`` + handles pe ``t % npes``) and a short LDS reduction -- a parallel sum, not a + log-depth scan with a barrier per step. + + Count and assign each run one block per tile with one thread per row: a row + is decoded once, a wave ballot per pe gives the rows of the wave bound for it, + and the per-wave counts go through a few LDS words. A row's rank among the + earlier rows for its pe is then a popcount, not a walk over the tile. """ dec = dict(log2_max_tok=log2_max_tok, recv_cap=recv_cap, topk=topk, npes=npes) - # Bound on the tile count, so the scan buffers can live in LDS. Tiles cover - # the sorted rows, so this is sized from the token-slot space, NOT from + if SBM > BLOCK: + raise ValueError( + f"plan needs one thread per row of a sort tile: SBM={SBM} > {BLOCK}" + ) + if BLOCK % npes: + raise ValueError(f"plan needs npes to divide {BLOCK}, got {npes}") + # Bound on the tile count, so the prefix loop has a static trip count. Tiles + # cover the sorted rows, so this is sized from the token-slot space, NOT from # ``cap`` -- ``cap`` is the per-peer row reservation and is npes times # smaller. Stage1 also pads each local expert's rows up to an SBM boundary, # so num_valid runs past slot_cap by up to local_experts * (SBM - 1). TILES_MAX = (slot_cap + SBM - 1) // SBM + local_experts + 1 - SCAN_N = (TILES_MAX + 1) * npes - STEPS = max(1, (TILES_MAX - 1).bit_length()) - lds_bytes = 2 * SCAN_N * 4 - if lds_bytes > 160 * 1024: - raise ValueError(f"plan scan needs {lds_bytes} B of LDS, over the 160 KiB budget") + PREFIX_J = -(-TILES_MAX * npes // BLOCK) + GROUPS = BLOCK // npes + WAVES = BLOCK // WAVE + # LDS words: per-wave per-pe counts | per-thread partial sums | the tile's prefix. + OFF_PART = WAVES * npes + OFF_PREF = OFF_PART + BLOCK + lds_bytes = (OFF_PREF + npes) * 4 @fx.struct - class ScanShared: + class PlanShared: buf: fx.Array[Int8, lds_bytes, 16] # fmt: off @@ -260,78 +300,94 @@ def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, out_buf = ptr_buf_tensor(arg_out_row, fx.Int32) cdone = ptr_buf_tensor(arg_cdone, fx.Int32) - lds = fx.SharedAllocator().allocate(ScanShared).peek() + lds = fx.SharedAllocator().allocate(PlanShared).peek() lds_i32 = fx.Int32(fx.ptrtoint(lds.buf.ptr)) - def sc(half, i): - """LDS i32 slot ``i`` of ping(0)/pong(1).""" - return lds_typed_ptr( - lds_i32 + (fx.Int32(half * SCAN_N) + i) * fx.Int32(4), T.i32, align=4 - ) + def lw(i): + """LDS i32 word ``i`` of the block's scratch.""" + return lds_typed_ptr(lds_i32 + i * fx.Int32(4), T.i32, align=4) # The GEMM2 kernel's per-chunk completion counters start from zero every call. for g in range(gid, i32_cdone_n, stride): cdone[g] = fx.Int32(0) - for g in range(gid, tiles * fx.Int32(npes), stride): - tile = g // fx.Int32(npes) - pe = g - tile * fx.Int32(npes) - base = trb[tile] - cnt = fx.Int32(0) - for r in range_constexpr(SBM): - valid, dest_pe, _, _ = _decode(stids[base + fx.Int32(r)], **dec) - cnt = cnt + (valid & (dest_pe == pe)).select(fx.Int32(1), fx.Int32(0)) - tile_cnt[g] = cnt - grid_barrier(arg_bar, i32_nblocks, tid) + lane = tid % fx.Int32(WAVE) + wave = tid // fx.Int32(WAVE) + row_ok = tid < fx.Int32(SBM) - # Inclusive Hillis-Steele over the tile axis, all npes lanes at once, - # run identically by every block. - n = tiles * fx.Int32(npes) - for i in range(tid, n, fx.Int32(BLOCK)): - fx.ptr_store(tile_cnt[i], sc(0, i)) - fx.barrier() - for step in range_constexpr(STEPS): - src, dst = step % 2, (step + 1) % 2 - back = fx.Int32((1 << step) * npes) - for i in range(tid, n, fx.Int32(BLOCK)): - v = fx.Int32(fx.ptr_load(sc(src, i))) - v = (i >= back).select( - v + fx.Int32(fx.ptr_load(sc(src, i - back))), v - ) - fx.ptr_store(v, sc(dst, i)) + def tile_masks(tile): + """One thread per row of ``tile``: decode it once, then a ballot per pe + gives the wave's rows bound for that pe. Returns the row's decode and + those masks; the per-wave counts are published to LDS and made + visible (barrier) before returning.""" + base = trb[tile] + valid, dest_pe, dest_lid, s = _decode(stids[base + row_ok.select(tid, fx.Int32(0))], **dec) + valid = valid & row_ok + masks = [ + fx.Int64(rocdl.ballot(T.i64, valid & (dest_pe == fx.Int32(p)))) + for p in range_constexpr(npes) + ] + if lane == fx.Int32(0): + for p in range_constexpr(npes): + fx.ptr_store(fx.Int32(fx.math.ctpop(masks[p])), lw(wave * fx.Int32(npes) + fx.Int32(p))) fx.barrier() - final = STEPS % 2 - # Exclusive = inclusive - own count; the extra row holds the totals the - # comm waves read as a chunk's end. - if bid == fx.Int32(0): - for i in range(tid, n, fx.Int32(BLOCK)): - tile_base[i] = fx.Int32(fx.ptr_load(sc(final, i))) - tile_cnt[i] + return valid, dest_pe, dest_lid, s, masks + + # One block per tile; no thread walks a tile's rows. + for tile in range(bid, tiles, i32_nblocks): + tile_masks(tile) if tid < fx.Int32(npes): - total = fx.Int32( - fx.ptr_load(sc(final, (tiles - fx.Int32(1)) * fx.Int32(npes) + tid)) - ) - capped = (total < fx.Int32(cap)).select(total, fx.Int32(cap)) - tile_base[tiles * fx.Int32(npes) + tid] = capped - ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = capped + total = fx.Int32(0) + for w in range_constexpr(WAVES): + total = total + fx.Int32(fx.ptr_load(lw(fx.Int32(w * npes) + tid))) + tile_cnt[tile * fx.Int32(npes) + tid] = total + fx.barrier() + flag_barrier(arg_bar, i32_nblocks, tid, bid) idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) - final = STEPS % 2 - for g in range(gid, tiles * fx.Int32(SBM), stride): - tile = g // fx.Int32(SBM) - row = g - tile * fx.Int32(SBM) - base = trb[tile] - valid, dest_pe, dest_lid, s = _decode(stids[base + row], **dec) - out_buf[g] = fx.Int32(-1) + lane_lt = (fx.Int64(1) << fx.Int64(lane)) - fx.Int64(1) + for tile in range(bid, tiles, i32_nblocks): + # Rows for pe in the tiles before this one. Element e of tile_cnt is + # (tile e // npes, pe e % npes), and BLOCK % npes == 0, so each thread + # always sums one pe. + lim = tile * fx.Int32(npes) + part = fx.Int32(0) + for j in range_constexpr(PREFIX_J): + e = tid + fx.Int32(j * BLOCK) + before = e < lim + part = part + before.select(tile_cnt[before.select(e, fx.Int32(0))], fx.Int32(0)) + fx.ptr_store(part, lw(fx.Int32(OFF_PART) + tid)) + fx.barrier() + if tid < fx.Int32(npes): + pre = fx.Int32(0) + for k in range_constexpr(GROUPS): + pre = pre + fx.Int32(fx.ptr_load(lw(fx.Int32(OFF_PART + k * npes) + tid))) + fx.ptr_store(pre, lw(fx.Int32(OFF_PREF) + tid)) + tile_base[tile * fx.Int32(npes) + tid] = pre + # The extra row holds the totals the comm waves read as a chunk's end. + if tile == tiles - fx.Int32(1): + total = pre + tile_cnt[tile * fx.Int32(npes) + tid] + capped = (total < fx.Int32(cap)).select(total, fx.Int32(cap)) + tile_base[tiles * fx.Int32(npes) + tid] = capped + ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = capped + fx.barrier() + + valid, dest_pe, dest_lid, s, masks = tile_masks(tile) + g = tile * fx.Int32(SBM) + tid + if row_ok: + out_buf[g] = fx.Int32(-1) if valid: - within = fx.Int32(0) - for rp in range_constexpr(SBM): - v2, pe2, _, _ = _decode(stids[base + fx.Int32(rp)], **dec) - earlier = (fx.Int32(rp) < row) & v2 & (pe2 == dest_pe) - within = within + earlier.select(fx.Int32(1), fx.Int32(0)) - slot = tile * fx.Int32(npes) + dest_pe - # This block scanned the whole array itself, so the prefix is - # in its own LDS -- no barrier between the scan and this. - excl = fx.Int32(fx.ptr_load(sc(final, slot))) - tile_cnt[slot] + # Rank among this wave's rows for my pe, plus the rows of the + # earlier waves of the tile for it. + mine = masks[0] + for p in range_constexpr(1, npes): + mine = (dest_pe == fx.Int32(p)).select(masks[p], mine) + within = fx.Int32(fx.math.ctpop(mine & lane_lt)) + for w in range_constexpr(WAVES): + within = within + (fx.Int32(w) < wave).select( + fx.Int32(fx.ptr_load(lw(fx.Int32(w * npes) + dest_pe))), fx.Int32(0) + ) + excl = fx.Int32(fx.ptr_load(lw(fx.Int32(OFF_PREF) + dest_pe))) place = excl + within # Past the per-pair reservation there is nowhere to put this # row. Drop it rather than write into the next peer's run, @@ -342,6 +398,7 @@ def sc(half, i): out_row = dest_pe * fx.Int32(cap) + place out_buf[g] = out_row idx_buf[out_row] = dest_lid * fx.Int32(topk) + s + fx.barrier() # fmt: off @flyc.jit @@ -392,8 +449,9 @@ def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): @comm_ops.traced -def sdma_finish_peer(sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco, - signal_mode="remote"): +def sdma_finish_peer( + sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco, signal_mode="remote" +): """Ship idx, tell ``peer`` its rows have landed, and publish ``cnt``. Two ways to say "landed": @@ -687,7 +745,9 @@ def __init__( # what put bs=8192 over the 4 GiB buffer-resource range. self.row_slack = row_slack if signal_mode not in ("quiet", "remote"): - raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") + raise ValueError( + f"signal_mode must be quiet or remote, got {signal_mode!r}" + ) if signal_mode == "remote" and queues != 1: raise ValueError( "signal_mode='remote' hangs the completion signal on the last put " @@ -741,7 +801,10 @@ def __init__( self.cdone_stride = CDONE_STRIDE self.cdone = torch.zeros((max_tiles + 1) * CDONE_STRIDE, **i32) # Never-reset counters for the two grid barriers and the combine's last-block count. - self.plan_bar = torch.zeros(1, **i32) + self.plan_bar = torch.zeros( + torch.cuda.get_device_properties(dev).multi_processor_count * BAR_STRIDE, + **i32, + ) # Rows the per-pair reservation had no room for; must stay 0. self.overflow = torch.zeros(1, **i32) self.combine_bar = torch.zeros(1, **i32) @@ -768,7 +831,11 @@ def __init__( ) self._plans = {} self._combine = build_combine( - layout=lay, rank=self.rank, hidden=self.hidden, topk=self.topk, quant=quant, + layout=lay, + rank=self.rank, + hidden=self.hidden, + topk=self.topk, + quant=quant, signal_mode=signal_mode, ) @@ -891,10 +958,9 @@ def run(self, run_tokens, config=None, profile=None): import torch from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled - moe, lay = self.moe, self.layout + moe = self.moe config = moe._active_config if config is None else config s_fx = fx.Stream(torch.cuda.current_stream().cuda_stream) - op = moe._s1_op def mark(stage): if profile is not None: From bea7f03747463d271f4aa6d3e0b02f86070633c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=A0=C3=A7kawhil-amd?= Date: Wed, 28 Oct 2026 03:57:00 +0000 Subject: [PATCH 4/9] fix(cco/example): plan kernel correctness under skewed routing and overflow Review findings on #859: - The plan's prefix was bounded by a tile count sized from the send side (slot_cap), but a rank receives rows from every peer, so under skewed routing the bound was exceeded and the prefix was silently truncated: wrong output on every rank with overflow 0. The prefix now sums up to the tile's own index, so no bound is needed. main.py --hot-bias reproduces it (bias 0.1 failed before, is bit-exact now). - tile_base is clamped to cap, so a chunk that crosses the per-pair reservation is cut at it instead of reading the next peer's run and writing past this rank's slot in the peer's landing buffer. - Nothing read the overflow counter. SortedCombine.check_overflow() raises, and main.py / bench.py call it. main.py also requires zero mismatched rows when both sides move fp8, since relL2 alone let a few wrong rows through. - tiles == 0 left cursor and tile_base stale; they are now zeroed. - The plan waits for every peer's done flag before rewriting idx_send, and close() waits for the same flags, so this rank's own outbound puts have finished reading S before it is reused or freed. Co-Authored-By: Claude Code --- .../python/11_flydsl_mega_moe_sdma/bench.py | 26 +++-- .../python/11_flydsl_mega_moe_sdma/main.py | 66 +++++++++-- .../11_flydsl_mega_moe_sdma/sorted_combine.py | 105 ++++++++++++++---- 3 files changed, 156 insertions(+), 41 deletions(-) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py index 6156788868..3e2a41cfb6 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py @@ -183,25 +183,32 @@ def main(): "--chunk-tiles", default="16", help="fused chunk sizes to sweep, e.g. 8,16,32" ) parser.add_argument( - "--queues", type=int, default=1, + "--queues", + type=int, + default=1, help="SDMA channels per (src, dst) pair. More than one measured " "slower: KFD recommends a single engine per peer link here, so " "extra channels share it and only add descriptors.", ) parser.add_argument("--row-slack", type=float, default=1.5) parser.add_argument( - "--signal-mode", choices=("quiet", "remote"), default="remote", + "--signal-mode", + choices=("quiet", "remote"), + default="remote", help="how a sender tells a peer its rows landed", ) parser.add_argument( - "--no-aligned-pair", action="store_true", + "--no-aligned-pair", + action="store_true", help="force aiter off its common-pair stage2 so both paths run the " "same config. It is picked only at bucket 8192, where SortedCombine " "has no equivalent -- and measured 1.5%% slower than without it on " "uniform routing anyway.", ) parser.add_argument( - "--quant", choices=("none", "fp8_blockwise_1x32"), default="fp8_blockwise_1x32", + "--quant", + choices=("none", "fp8_blockwise_1x32"), + default="fp8_blockwise_1x32", help="quantise the send buffer; halves the bytes SDMA moves", ) parser.add_argument( @@ -252,9 +259,13 @@ def main(): fanout_masks=(), **net, ) - sc = SortedCombine(moe, queues=args.queues, quant=args.quant, - row_slack=args.row_slack, - signal_mode=args.signal_mode) + sc = SortedCombine( + moe, + queues=args.queues, + quant=args.quant, + row_slack=args.row_slack, + signal_mode=args.signal_mode, + ) for bs in batch_sizes: x, wts, ids = ( @@ -347,6 +358,7 @@ def e2e(): raise AssertionError( f"bs={bs} {label}: relL2={rel:.3e} exceeds {args.rtol}" ) + sc.check_overflow() if rank == 0: print( diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py index 75c2609a9a..0221ed540f 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py @@ -82,7 +82,11 @@ def _next_pow2(v): return 1 << (int(v) - 1).bit_length() -def _make_inputs(tokens, model_dim, experts, topk, rank, seed, device): +def _make_inputs( + tokens, model_dim, experts, topk, rank, seed, device, hot_bias=0.0, hot_experts=0 +): + """``hot_bias`` raises the router score of the first ``hot_experts`` experts, i.e. + skews routing toward rank 0; 0 is the uniform routing everything else uses.""" g = torch.Generator(device=device).manual_seed(seed + rank) x = torch.randn( (tokens, model_dim), dtype=torch.bfloat16, device=device, generator=g @@ -90,6 +94,8 @@ def _make_inputs(tokens, model_dim, experts, topk, rank, seed, device): scores = torch.randn( (tokens, experts), dtype=torch.float32, device=device, generator=g ) + if hot_bias: + scores[:, :hot_experts] += hot_bias values, ids = torch.topk(scores, topk, dim=-1) return ( x.contiguous(), @@ -174,25 +180,39 @@ def main(): help="sort tiles per shipped chunk", ) parser.add_argument( - "--queues", type=int, default=1, + "--queues", + type=int, + default=1, help="SDMA channels per (src, dst) pair. More than one measured " "slower: KFD recommends a single engine per peer link here, so " "extra channels share it and only add descriptors.", ) parser.add_argument("--row-slack", type=float, default=1.5) parser.add_argument( - "--signal-mode", choices=("quiet", "remote"), default="remote", + "--hot-bias", + type=float, + default=0.0, + help="add this to the router scores of rank 0's experts: skewed routing, " + "where one rank receives more than an even share of the rows", + ) + parser.add_argument( + "--signal-mode", + choices=("quiet", "remote"), + default="remote", help="how a sender tells a peer its rows landed", ) parser.add_argument( - "--no-aligned-pair", action="store_true", + "--no-aligned-pair", + action="store_true", help="force aiter off its common-pair stage2 so both paths run the " "same config. It is picked only at bucket 8192, where SortedCombine " "has no equivalent -- and measured 1.5%% slower than without it on " "uniform routing anyway.", ) parser.add_argument( - "--quant", choices=("none", "fp8_blockwise_1x32"), default="fp8_blockwise_1x32", + "--quant", + choices=("none", "fp8_blockwise_1x32"), + default="fp8_blockwise_1x32", help="quantise the send buffer; halves the bytes SDMA moves", ) args = parser.parse_args() @@ -216,6 +236,8 @@ def main(): rank, args.seed, device, + hot_bias=args.hot_bias, + hot_experts=net["experts"] // world, ) moe = MegaMoEV2( rank=rank, @@ -235,7 +257,9 @@ def main(): if args.no_aligned_pair and config.stage2.aligned_pair: config = replace( config, - stage2=replace(config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96), + stage2=replace( + config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96 + ), ) if config.stage2.aligned_pair: raise SystemExit( @@ -252,10 +276,16 @@ def main(): lambda: moe._run_stage2(args.bs, None, True, config), args.iters ) failures = [] - sc = SortedCombine(moe, queues=args.queues, chunk_tiles=args.chunk_tiles, quant=args.quant, + sc = SortedCombine( + moe, + queues=args.queues, + chunk_tiles=args.chunk_tiles, + quant=args.quant, row_slack=args.row_slack, - signal_mode=args.signal_mode) + signal_mode=args.signal_mode, + ) worst = 0.0 + mismatched = 0 # Repeated calls with no host sync between them also exercise the epoch counter, # the flags and the "consumed" handshake that stops a fast rank overrunning a slow one. for _ in range(4): @@ -263,22 +293,34 @@ def main(): torch.cuda.synchronize() diff = out.float() - ref.float() worst = max(worst, float(diff.norm() / ref.float().norm().clamp_min(1e-9))) - t = torch.tensor([worst], dtype=torch.float64) + mismatched = max(mismatched, int((out != ref).any(dim=1).sum())) + t = torch.tensor([worst, mismatched], dtype=torch.float64) dist.all_reduce(t, op=dist.ReduceOp.MAX) - worst = float(t.item()) + worst, mismatched = float(t[0].item()), int(t[1].item()) ms_per = _timeit(lambda: sc.run(args.bs, config), args.iters) ok = worst < args.rtol + # Both sides move fp8 rows in aiter's layout and reduce in the same order, so + # they must agree exactly; relL2 alone would let a few wrong rows through. + exact = config.p2p_quant == args.quant == "fp8_blockwise_1x32" + if exact and mismatched: + ok = False if not ok: failures.append("sdma") + try: + sc.check_overflow() + except RuntimeError as e: + failures.append(f"overflow ({e})") + ok = False if rank == 0: print( - f"relL2={worst:.3e} {'PASS' if ok else 'FAIL'} " + f"relL2={worst:.3e} mismatched_rows={mismatched} " + f"{'PASS' if ok else 'FAIL'} " f"time={ms_per:.4f}ms (reference stage2+combine {ref_ms:.4f}ms)", flush=True, ) sc.close() if failures: - raise AssertionError(f"relL2 above {args.rtol} for {failures}") + raise AssertionError(f"sdma path failed: {failures}") if rank == 0: print("SUCCESS", flush=True) finally: diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py index e7357e99f9..9204c68cfb 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py @@ -230,9 +230,7 @@ def flag_barrier(bar_addr, nblocks, tid, bid): fx.barrier() -def build_rank_sort_plan( - *, SBM, topk, log2_max_tok, recv_cap, npes, cap, slot_cap, local_experts -): +def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap, rank): """Kernel 1: give every valid sorted row a stable index in its rank's run. * count: ``tile_cnt[tile][pe]`` = rows of the sort tile bound for ``pe`` @@ -246,7 +244,15 @@ def build_rank_sort_plan( does not scan the whole ``tile_cnt`` array: for each tile it owns it sums the counts of the earlier tiles with one pass of the block's threads (thread ``t`` handles pe ``t % npes``) and a short LDS reduction -- a parallel sum, not a - log-depth scan with a barrier per step. + log-depth scan with a barrier per step. The sum runs to the tile's own index, + so nothing here is sized by a bound on the tile count: a rank that receives + far more rows than an even share (skewed routing) is handled the same way. + + Before it writes anything the kernel waits until every peer has consumed the + previous call: that is what guarantees this rank's own outbound puts have + finished reading ``S`` and ``idx_send``, which this call is about to rewrite. + ``tile_base`` is clamped to ``cap``, so a chunk that crosses the per-pair + reservation is cut at it instead of shipping the next peer's rows. Count and assign each run one block per tile with one thread per row: a row is decoded once, a wave ballot per pe gives the rows of the wave bound for it, @@ -261,13 +267,6 @@ def build_rank_sort_plan( ) if BLOCK % npes: raise ValueError(f"plan needs npes to divide {BLOCK}, got {npes}") - # Bound on the tile count, so the prefix loop has a static trip count. Tiles - # cover the sorted rows, so this is sized from the token-slot space, NOT from - # ``cap`` -- ``cap`` is the per-peer row reservation and is npes times - # smaller. Stage1 also pads each local expert's rows up to an SBM boundary, - # so num_valid runs past slot_cap by up to local_experts * (SBM - 1). - TILES_MAX = (slot_cap + SBM - 1) // SBM + local_experts + 1 - PREFIX_J = -(-TILES_MAX * npes // BLOCK) GROUPS = BLOCK // npes WAVES = BLOCK // WAVE # LDS words: per-wave per-pe counts | per-thread partial sums | the tile's prefix. @@ -285,7 +284,8 @@ class PlanShared: def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, - arg_over: fx.Int64, i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32): + arg_over: fx.Int64, arg_done: fx.Int64, arg_counter: fx.Int64, + i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32): # fmt: on tid = fx.thread_idx.x bid = fx.block_idx.x @@ -307,6 +307,19 @@ def lw(i): """LDS i32 word ``i`` of the block's scratch.""" return lds_typed_ptr(lds_i32 + i * fx.Int32(4), T.i32, align=4) + # S and idx_send are rewritten from here on. A peer's done flag says it has + # read everything I sent it last call, and with the copy engine's FIFO that + # means my puts have finished reading them. + consumed = global_typed_ptr(arg_counter, T.i32)[0] + if tid < fx.Int32(npes): + if tid != fx.Int32(rank): + comm_ops.spin_until_ge_i32_system( + fx.Int64(arg_done) + fx.Int64(tid) * fx.Int64(4), + consumed, + acquire=True, + ) + fx.barrier() + # The GEMM2 kernel's per-chunk completion counters start from zero every call. for g in range(gid, i32_cdone_n, stride): cdone[g] = fx.Int32(0) @@ -344,18 +357,20 @@ def tile_masks(tile): fx.barrier() flag_barrier(arg_bar, i32_nblocks, tid, bid) + # No rows at all: the loop below never runs, so say so explicitly. + if (tiles == fx.Int32(0)) & (bid == fx.Int32(0)) & (tid < fx.Int32(npes)): + tile_base[tid] = fx.Int32(0) + ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = fx.Int32(0) + idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) lane_lt = (fx.Int64(1) << fx.Int64(lane)) - fx.Int64(1) for tile in range(bid, tiles, i32_nblocks): # Rows for pe in the tiles before this one. Element e of tile_cnt is # (tile e // npes, pe e % npes), and BLOCK % npes == 0, so each thread # always sums one pe. - lim = tile * fx.Int32(npes) part = fx.Int32(0) - for j in range_constexpr(PREFIX_J): - e = tid + fx.Int32(j * BLOCK) - before = e < lim - part = part + before.select(tile_cnt[before.select(e, fx.Int32(0))], fx.Int32(0)) + for e in range(tid, tile * fx.Int32(npes), fx.Int32(BLOCK)): + part = part + tile_cnt[e] fx.ptr_store(part, lw(fx.Int32(OFF_PART) + tid)) fx.barrier() if tid < fx.Int32(npes): @@ -363,7 +378,9 @@ def tile_masks(tile): for k in range_constexpr(GROUPS): pre = pre + fx.Int32(fx.ptr_load(lw(fx.Int32(OFF_PART + k * npes) + tid))) fx.ptr_store(pre, lw(fx.Int32(OFF_PREF) + tid)) - tile_base[tile * fx.Int32(npes) + tid] = pre + tile_base[tile * fx.Int32(npes) + tid] = (pre < fx.Int32(cap)).select( + pre, fx.Int32(cap) + ) # The extra row holds the totals the comm waves read as a chunk's end. if tile == tiles - fx.Int32(1): total = pre + tile_cnt[tile * fx.Int32(npes) + tid] @@ -405,10 +422,12 @@ def tile_masks(tile): def launch(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, - arg_over: fx.Int64, i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32, stream: fx.Stream): + arg_over: fx.Int64, arg_done: fx.Int64, arg_counter: fx.Int64, + i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32, stream: fx.Stream): # fmt: on plan(arg_stids, arg_trb, arg_num_valid, arg_tile_cnt, arg_tile_base, arg_cursor, - arg_out_row, arg_idx, arg_cdone, arg_bar, arg_over, i32_nblocks, i32_cdone_n).launch( + arg_out_row, arg_idx, arg_cdone, arg_bar, arg_over, arg_done, arg_counter, + i32_nblocks, i32_cdone_n).launch( grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) return launch @@ -852,8 +871,7 @@ def _plan_for(self, sbm): recv_cap=int(self.moe.max_recv), npes=self.npes, cap=self.cap, - slot_cap=self.slot_cap, - local_experts=int(self.moe.w2.shape[0]), + rank=self.rank, ) return self._plans[sbm] @@ -888,6 +906,8 @@ def _launch_plan(self, config, s_fx): fx.Int64(self.cdone.data_ptr()), fx.Int64(self.plan_bar.data_ptr()), fx.Int64(self.overflow.data_ptr()), + fx.Int64(self.mem.ptr + lay.off_done), + fx.Int64(self.mem.ptr + lay.off_counter), fx.Int32(self.nblocks), fx.Int32(self._cdone_used(config)), s_fx, @@ -989,5 +1009,46 @@ def mark(stage): mark("combine") return self.out[: run_tokens * self.hidden].view(run_tokens, self.hidden) + def check_overflow(self): + """Raise if the plan dropped rows since the last check (syncs the device). + + A row that does not fit its peer's per-pair reservation is dropped, so + the output is silently wrong unless the host looks at this counter. + """ + import torch + + torch.cuda.synchronize() + dropped = int(self.overflow.item()) + self.overflow.zero_() + if dropped: + raise RuntimeError( + f"rank {self.rank}: {dropped} rows did not fit the per-pair reservation " + f"(cap={self.cap} rows, row_slack={self.row_slack}); raise row_slack" + ) + + def _wait_peers_consumed(self, timeout_s=30.0): + """Wait until every peer has read what this rank last sent it. + + Under ``remote`` signalling a rank's last call can return while its own + puts are still in flight; freeing the window then would pull it out from + under the copy engine. + """ + import time + + import torch + + torch.cuda.synchronize() + lay = self.layout + counter = _torch_view(self.mem.ptr + lay.off_counter, 4, torch.int32) + done = _torch_view(self.mem.ptr + lay.off_done, self.npes * 4, torch.int32) + calls = int(counter.item()) + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + flags = done.tolist() + if all(f >= calls for r, f in enumerate(flags) if r != self.rank): + return + time.sleep(0.001) + def close(self): + self._wait_peers_consumed() self.comm.close_all_resources() From 186b0031c23f942b3ad9ebb63cb0f017e0e0e1c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=A0=C3=A7kawhil-amd?= Date: Wed, 28 Oct 2026 04:46:45 +0000 Subject: [PATCH 5/9] style(cco/example): make the PR pass pre-commit CI runs the hooks on every file the PR touches, and these two had never been through them: sdmabw.py lacked the license header and was not black-formatted, moe_gemm2_localsort.py defined a constant between imports (E402), and both had an unused local (F841). No behaviour change. Co-Authored-By: Claude Code --- .../moe_gemm2_localsort.py | 16 +++--- .../python/11_flydsl_mega_moe_sdma/sdmabw.py | 55 +++++++++++++++---- 2 files changed, 52 insertions(+), 19 deletions(-) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py index f8aecebe33..dba8fb3a65 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py @@ -46,13 +46,6 @@ from dataclasses import dataclass -#: i32 slots between consecutive chunk-completion counters. Every unit of a -#: chunk bumps its counter, so packing them adjacent puts all of a launch's -#: atomics on one or two cache lines: 128 B apart instead is worth ~8% of the -#: GEMM2 kernel at bs=4096. Spreading a chunk's count over several counters on -#: top of this buys nothing -- the atomic itself is not the cost. -CDONE_STRIDE = 32 - import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import const_expr, range_constexpr, rocdl @@ -97,6 +90,14 @@ from aiter.ops.flydsl.kernels.mxfp4_gemm_common import _fabs_f32 as fabs_f32 +#: i32 slots between consecutive chunk-completion counters. Every unit of a +#: chunk bumps its counter, so packing them adjacent puts all of a launch's +#: atomics on one or two cache lines: 128 B apart instead is worth ~8% of the +#: GEMM2 kernel at bs=4096. Spreading a chunk's count over several counters on +#: top of this buys nothing -- the atomic itself is not the cost. +CDONE_STRIDE = 32 + + @dataclass(frozen=True) class CommSpec: """Compile-time description of the fused SDMA comm workgroups. @@ -142,7 +143,6 @@ def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, raise ValueError(f"unsupported quant={quant!r}") if quant_fp8 and g2_bf16_lds: raise ValueError("fp8_blockwise_1x32 needs the f32 CShuffle (g2_bf16_lds=False)") - out_elem_bytes = 1 if quant_fp8 else 2 # Row layout matches aiter's P2P scatter: N_OUT payload bytes followed by the # N_OUT/32 e8m0 block scales, so one contiguous SDMA put ships both. row_nbytes = N_OUT + N_OUT // 32 if quant_fp8 else N_OUT * 2 diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py index e97959ce14..e5082f06a1 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py @@ -1,7 +1,25 @@ #!/usr/bin/env python3 # Copyright © Advanced Micro Devices, Inc. All rights reserved. # -# SPDX-License-Identifier: MIT +# 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. """What one cco SDMA channel actually moves, with no GEMM in the way. Same shape as the fused combine's transport -- one comm wave per peer, each @@ -35,13 +53,13 @@ def build_xfer(*, npes, rank, comm_wgs, queues, split, puts): """One wave per peer: ``puts`` back-to-back runs, then drain every channel.""" import mori.cco.device.flydsl as cco - n_waves = comm_wgs * 4 - @flyc.kernel( name=f"sdmabw_p{npes}_w{comm_wgs}q{queues}s{split}n{puts}_r{rank}", known_block_size=[WAVE * 4, 1, 1], ) - def xfer(dev_comm: fx.Int64, win: fx.Int64, i64_bytes: fx.Int64, i64_stride: fx.Int64): + def xfer( + dev_comm: fx.Int64, win: fx.Int64, i64_bytes: fx.Int64, i64_stride: fx.Int64 + ): tid = fx.thread_idx.x lane = tid % fx.Int32(WAVE) wave = rocdl.readfirstlane(T.i32, tid // fx.Int32(WAVE)) @@ -69,8 +87,13 @@ def xfer(dev_comm: fx.Int64, win: fx.Int64, i64_bytes: fx.Int64, i64_stride: fx. sdma.quiet(owner) @flyc.jit - def launch(dev_comm: fx.Int64, win: fx.Int64, i64_bytes: fx.Int64, - i64_stride: fx.Int64, stream: fx.Stream): + def launch( + dev_comm: fx.Int64, + win: fx.Int64, + i64_bytes: fx.Int64, + i64_stride: fx.Int64, + stream: fx.Stream, + ): xfer(dev_comm, win, i64_bytes, i64_stride).launch( grid=(comm_wgs, 1, 1), block=(WAVE * 4, 1, 1), stream=stream ) @@ -83,7 +106,9 @@ def main(): parser.add_argument("--mb", default="0.36,0.72,1.44,2.89,5.78", help="MB per put") parser.add_argument("--queues", default="1,2", help="channels per peer") parser.add_argument("--split", default="1", help="split one put across channels") - parser.add_argument("--puts", type=int, default=8, help="back-to-back puts per peer") + parser.add_argument( + "--puts", type=int, default=8, help="back-to-back puts per peer" + ) parser.add_argument("--iters", type=int, default=50) args = parser.parse_args() @@ -126,8 +151,12 @@ def main(): if split > queues: continue fn = build_xfer( - npes=world, rank=rank, comm_wgs=comm_wgs, - queues=queues, split=split, puts=args.puts, + npes=world, + rank=rank, + comm_wgs=comm_wgs, + queues=queues, + split=split, + puts=args.puts, ) for mb in sizes: nbytes = int(mb * (1 << 20)) @@ -137,8 +166,12 @@ def once(): from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled _run_compiled( - fn, fx.Int64(dev_comm.ptr), fx.Int64(win.handle), - fx.Int64(nbytes), fx.Int64(stride), s_fx, + fn, + fx.Int64(dev_comm.ptr), + fx.Int64(win.handle), + fx.Int64(nbytes), + fx.Int64(stride), + s_fx, ) once() From 919fba7f0bbd2290b563d70c3f822011ba1844bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=A0=C3=A7kawhil-amd?= Date: Wed, 28 Oct 2026 07:55:09 +0000 Subject: [PATCH 6/9] perf(cco/example): sender-side inv, tree barrier and 2x combine blocks (+1-3% stage2+combine) The combine used to build the slot -> landing-row table itself and rendezvous the whole grid before reducing; the plan kernel's grid barrier and its remote inv stores were also a large part of its time. Found with a GPU-side phase ladder (rocprof / CUDA-graph replay) rather than host timing. - inv is built by the SENDER and lives in the cco window: the plan kernel fills it for rows that stay on this rank, and the GEMM2 kernel's comm waves push the peers' entries with direct stores while the GEMM computes its first chunk. No idx travels with the data, and the combine loses its build phase and its grid barrier (a 4-byte signal put replaces the idx put). - plan kernel: a two-level tree barrier (no address takes more than 32 atomics, one poll per block) instead of every thread of every block polling every block's flag; plan 43 -> ~17 us at 8192 tokens/rank. - combine runs two blocks per CU now that it has no grid-wide rendezvous: more row loads in flight for the reduce. - the GEMM2 file is renamed moe_gemm2_localsort.py -> moe_gemm2_with_sdma.py. stage2+combine on 8x MI355, same harness, A/B against the previous commit: 2048 0.507 -> 0.492 ms, 4096 0.882 -> 0.855 ms, 8192 1.537 -> 1.490 ms, 16384 2.922 -> 2.840 ms (vs aiter: 1.00x -> 1.03x, 1.10x -> 1.14x, 1.19x -> 1.23x, 1.31x -> 1.35x). Output is bit-identical to aiter at every size and to the previous commit on the same inputs. Co-Authored-By: Claude Code --- .../python/11_flydsl_mega_moe_sdma/main.py | 2 +- ...m2_localsort.py => moe_gemm2_with_sdma.py} | 59 ++++- .../11_flydsl_mega_moe_sdma/sorted_combine.py | 223 ++++++++++-------- 3 files changed, 172 insertions(+), 112 deletions(-) rename examples/cco/python/11_flydsl_mega_moe_sdma/{moe_gemm2_localsort.py => moe_gemm2_with_sdma.py} (91%) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py index 0221ed540f..7888b68f84 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py @@ -33,7 +33,7 @@ Three kernels on one stream (see ``sorted_combine.py``; the P2P baseline has two): 1. rank_sort_plan count -> scan -> assign: a stable send-buffer index per row - 2. gemm2 writes S[out_row] (``moe_gemm2_localsort.py``); extra comm waves, + 2. gemm2 writes S[out_row] (``moe_gemm2_with_sdma.py``); extra comm waves, one per peer, wait for a chunk of sort tiles to be written (a per-chunk counter the GEMM workgroups bump after their stores) and put that chunk for their peer on their own SDMA queue while the diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py similarity index 91% rename from examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py rename to examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py index dba8fb3a65..1bc955a5c1 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_localsort.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py @@ -78,6 +78,7 @@ ) from sorted_combine import ( ArenaLayout, + _store_i32_relaxed_system, sdma_finish_peer, sdma_put_run, sdma_wait_consumed, @@ -97,6 +98,10 @@ #: top of this buys nothing -- the atomic itself is not the cost. CDONE_STRIDE = 32 +#: Rows a comm wave pushes into a peer's inv per lane per loop trip, so that many +#: independent idx loads are in flight before their stores. +PUSH_U = 8 + @dataclass(frozen=True) class CommSpec: @@ -306,7 +311,7 @@ def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, # fmt: off -def compile_moe_gemm2_localsort(*, model_dim: int, inter_dim: int, experts: int, rank: int, +def compile_moe_gemm2_with_sdma(*, model_dim: int, inter_dim: int, experts: int, rank: int, send_nbytes: int, BM: int = 32, BN: int = 256, BK: int = 256, use_nt: bool = True, HIDDEN_MAX: int = 8192, INTER_MAX: int = 8192, a_dtype: str = "fp8", SBM: int | None = None, persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, @@ -369,7 +374,7 @@ class SharedStorage: buf: fx.Array[Int8, lds_bytes, 16] kernel_name = ( - f"megamoe_stage2_localsort_t{BM}x{BN}x{BK}" + f"megamoe_stage2_with_sdma_t{BM}x{BN}x{BK}" f"_sbm{SBM}_{a_dtype}_nt{int(use_nt)}" f"_p{int(persist)}cu{cu_num}s{int(persist_strided)}_pad{int(has_pad)}" f"_sk{skew_cu}" @@ -385,12 +390,12 @@ class SharedStorage: # fmt: off @flyc.kernel(name=kernel_name, known_block_size=[256, 1, 1]) - def kernel_localsort(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, + def kernel_with_sdma(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, - arg_cursor: fx.Int64, i32_max_m_blocks: fx.Int32, i32_inter: fx.Int32, i32_hidden: fx.Int32, - i32_kpad: fx.Int32, i32_npad: fx.Int32): + arg_cursor: fx.Int64, arg_idx: fx.Int64, i32_max_m_blocks: fx.Int32, i32_inter: fx.Int32, + i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32): # fmt: on tx_i32 = fx.thread_idx.x bid = fx.block_idx.x @@ -572,6 +577,33 @@ def comm_role(cw): if peer != fx.Int32(rank): if is_lane0: sdma_wait_consumed(w, peer, phase, layout=lay, rank=rank) + # Fill each peer's inv for the rows I send it: slot -> where the row lands + # there (its R rows follow its S rows, my run starts at rank * cap). This + # runs while the GEMM computes its first chunk, and has to be performed + # before any put or signal of mine, hence the wait on vmcnt. + idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) + land0 = fx.Int32(npes * lay.cap + rank * lay.cap) + for j in range_constexpr(n_mine): + peer = owner + fx.Int32(j * n_waves) + if peer < fx.Int32(npes): + if peer != fx.Int32(rank): + cnt = cursor_buf[peer] + inv_peer = fx.Int64(w.lsa_ptr(peer, lay.off_inv)) + for base in range(fx.Int32(0), cnt, fx.Int32(64 * PUSH_U)): + ok, js, sl = [], [], [] + for u in range_constexpr(PUSH_U): + jj = base + fx.Int32(u * 64) + lane + o = jj < cnt + ok.append(o) + js.append(jj) + sl.append(idx_buf[peer * fx.Int32(lay.cap) + o.select(jj, fx.Int32(0))]) + for u in range_constexpr(PUSH_U): + if ok[u]: + _store_i32_relaxed_system( + inv_peer + fx.Int64(sl[u]) * fx.Int64(4), + land0 + js[u], + ) + rocdl.s_waitcnt(vmcnt=0) total_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) n_chunks = (total_tiles + fx.Int32(CT - 1)) // fx.Int32(CT) for c in range(fx.Int32(0), n_chunks, fx.Int32(1)): @@ -624,15 +656,16 @@ def launch(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, - arg_cursor: fx.Int64, i32_max_m_blocks: fx.Int32, i32_grid_blocks: fx.Int32, i32_inter: fx.Int32, - i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32, stream: fx.Stream): + arg_cursor: fx.Int64, arg_idx: fx.Int64, i32_max_m_blocks: fx.Int32, i32_grid_blocks: fx.Int32, + i32_inter: fx.Int32, i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32, + stream: fx.Stream): # fmt: on num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) grid_x = i32_grid_blocks * num_n_blocks + fx.Int32(NC) - kernel_localsort( + kernel_with_sdma( arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, arg_cdone, - arg_tile_base, arg_cursor, i32_max_m_blocks, i32_inter, i32_hidden, i32_kpad, i32_npad, + arg_tile_base, arg_cursor, arg_idx, i32_max_m_blocks, i32_inter, i32_hidden, i32_kpad, i32_npad, ).launch(grid=(grid_x, 1, 1), block=(256, 1, 1), stream=stream) return launch @@ -642,9 +675,9 @@ def launch(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: # fmt: off -def run_moe_gemm2_localsort(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, +def run_moe_gemm2_with_sdma(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, - arg_cdone, arg_tile_base, arg_cursor, row_capacity, i32_inter, i32_hidden, stream, *, + arg_cdone, arg_tile_base, arg_cursor, arg_idx, row_capacity, i32_inter, i32_hidden, stream, *, model_dim, inter_dim, experts, rank, send_nbytes, BM, SBM, HIDDEN_MAX, INTER_MAX, cu_num, BN=256, BK=256, use_nt=True, g2_bhoist=True, g2_ascale_pf=True, g2_spart=402, persist=False, persist_cu=0, persist_strided=False, g2_bf16_lds=False, skew_cu=0, scatter_vec=8, comm=None, @@ -663,12 +696,12 @@ def run_moe_gemm2_localsort(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar key = tuple(sorted(compile_kw.items())) launch = _LAUNCH_CACHE.get(key) if launch is None: - launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_localsort(**compile_kw) + launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_with_sdma(**compile_kw) max_m_blocks = (row_capacity + BM - 1) // BM grid_blocks = launch_cu_num if persist else max_m_blocks _run_compiled( launch, arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, - arg_cdone, arg_tile_base, arg_cursor, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), fx.Int32(i32_inter), fx.Int32(i32_hidden), + arg_cdone, arg_tile_base, arg_cursor, arg_idx, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), fx.Int32(i32_inter), fx.Int32(i32_hidden), fx.Int32(0), fx.Int32(0), stream, ) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py index 9204c68cfb..17c52a424f 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py @@ -34,21 +34,28 @@ Three kernels on one stream (the P2P-scatter baseline has two): - 1. rank_sort_plan persistent grid; count -> scan -> assign with grid barriers: - out_row per sorted row, idx_send, tile_base, cursor; zeroes cdone - 2. gemm2_localsort writes S[out_row]; extra comm waves (one per peer) put finished + 1. rank_sort_plan persistent grid; count -> prefix -> assign with one rendezvous: + out_row per sorted row, tile_base, cursor; zeroes cdone; and writes + every row's slot -> landing-row entry into its DESTINATION's ``inv`` + 2. gemm2_with_sdma writes S[out_row]; extra comm waves (one per peer) put finished chunks with SDMA while the GEMM tiles are still computing - 3. combine persistent grid; wait for the peers' flags -> build slot->row - table -> grid barrier -> top-k reduce -> publish "consumed" + 3. combine persistent grid; wait for the peers' flags -> top-k reduce + -> publish "consumed" Everything lives in ONE cco window, so ``put`` names the same window for source and destination and only the byte offsets differ:: - [ S | R | idx_send | idx_land | flags | done | counter ] + [ S | R | inv | sig | flags | done | counter ] + +``inv[slot]`` is the landing row of the row for token slot ``slot`` (``origin_lid*topk ++ k``), or -1. The SENDER writes it, in the plan kernel, where a row's destination, +slot and position are all known at once -- so no index travels with the data and the +receiver builds no table. ``sig`` is scratch for the 4-byte put that carries the +completion signal. ``flags[src]`` says "src's data for call n has landed" (and how many rows). ``done[dst]`` says "dst has finished reading call n": a sender waits for it before -overwriting that peer's ``R``/``idx_land`` with call n+1, which keeps a fast rank +overwriting that peer's ``R`` and ``inv`` with call n+1, which keeps a fast rank from running into a slow one. ``counter`` is the call number; only the combine kernel advances it, after every reader of it has run. @@ -69,8 +76,8 @@ contiguous put still carries a whole row. bf16 doubles the bytes and the transfer stops fitting under the GEMM. -Both persistent kernels use grid-wide barriers, so their grid must be fully -resident (one block per CU is). +The plan kernel rendezvouses across its grid, so that grid must be fully resident +(one block per CU is). The combine has no grid-wide rendezvous and is not resident. """ from dataclasses import dataclass @@ -109,7 +116,8 @@ BLOCK = 256 WAVE = 64 -BAR_STRIDE = 32 # ints per flag_barrier slot: one 128 B cache line +BAR_STRIDE = 32 # ints per tree_barrier counter: one 128 B cache line +BAR_GROUP = 32 # blocks per tree_barrier group MIN_SBM = 16 _ALIGN = 128 @@ -125,6 +133,7 @@ class ArenaLayout: npes: int cap: int row_bytes: int + slot_cap: int @property def n_rows(self): @@ -139,16 +148,16 @@ def off_r(self): return self.off_s + self.n_rows * self.row_bytes @property - def off_idx_send(self): + def off_inv(self): return self.off_r + self.n_rows * self.row_bytes @property - def off_idx_land(self): - return self.off_idx_send + _align(self.n_rows * 4) + def off_sig(self): + return self.off_inv + _align(self.slot_cap * 4) @property def off_flags(self): - return self.off_idx_land + _align(self.n_rows * 4) + return self.off_sig + _align(self.npes * 4) @property def off_done(self): @@ -173,6 +182,21 @@ def _uniform_i64(x): return fx.Int64(rocdl.readfirstlane(T.i64, fx.Int64(x).ir_value())) +def _store_i32_relaxed_system(addr_i64, val): + """Relaxed i32 store that is visible system-wide once it retires. + + Unlike ``comm_ops.store_i32_global_system_release`` it carries no release, so + it does not write back the whole L2; it is the store for a row's ``inv`` entry + in a peer's window, one per row. + """ + comm_ops._llvm_d.StoreOp( + val.ir_value(), + comm_ops._to_ptr_global(addr_i64), + alignment=4, + ordering=comm_ops._llvm_d.AtomicOrdering.monotonic, + ) + + def _decode(p, *, log2_max_tok, recv_cap, topk, npes): """Packed srcmap word -> (valid, dest_pe, dest_lid, slot_k), as the stage2 epilogue does.""" t = p & fx.Int32(0x00FFFFFF) @@ -202,31 +226,40 @@ def grid_barrier(bar_addr, nblocks, tid): @comm_ops.traced -def flag_barrier(bar_addr, nblocks, tid, bid): - """All blocks of a resident grid rendezvous without a contended atomic. - - Block ``i`` bumps its own cache-line-sized slot ``bar[i]`` and every block polls - all slots, so arrival costs one uncontended atomic instead of ``nblocks`` of them - queueing on one address. Slots are never reset: a block's target is its own slot's - value at entry plus one, and every block bumps once per launch. +def tree_barrier(bar_addr, nblocks, tid, bid): + """All blocks of a resident grid rendezvous with two levels of small atomics. + + Blocks are grouped ``BAR_GROUP`` to a counter; the last block of a group bumps a + top counter; everyone polls only the top one. No address sees more than + ``BAR_GROUP`` atomics, and a round of polling is one load per block, not + ``nblocks`` of them (which, spun on by every thread of every block, slowed the + very stores it was waiting for). Counters are never reset: a group's counter + advances by its size per launch, so the launch number is ``old // size``. + Each counter sits on its own cache line. """ - target = fx.Int32( - comm_ops.load_i32_global_agent( - bar_addr + fx.Int64(bid) * fx.Int64(BAR_STRIDE * 4) - ) - ) + fx.Int32(1) + group = bid // fx.Int32(BAR_GROUP) + n_groups = (nblocks + fx.Int32(BAR_GROUP - 1)) // fx.Int32(BAR_GROUP) + size = nblocks - group * fx.Int32(BAR_GROUP) + size = (size < fx.Int32(BAR_GROUP)).select(size, fx.Int32(BAR_GROUP)) rocdl.s_waitcnt(0) fx.barrier() if tid == fx.Int32(0): comm_ops.fence_agent_release() - comm_ops.atomic_add_agent( - bar_addr + fx.Int64(bid) * fx.Int64(BAR_STRIDE * 4), fx.Int32(1) + old = fx.Int32( + comm_ops.atomic_add_agent( + bar_addr + fx.Int64(group) * fx.Int64(BAR_STRIDE * 4), fx.Int32(1) + ) ) - for i in range(tid, nblocks, fx.Int32(BLOCK)): + launch = old // size + if old - launch * size == size - fx.Int32(1): + comm_ops.atomic_add_agent( + bar_addr + fx.Int64(n_groups) * fx.Int64(BAR_STRIDE * 4), fx.Int32(1) + ) comm_ops.spin_until_ge_i32_agent( - bar_addr + fx.Int64(i) * fx.Int64(BAR_STRIDE * 4), target + bar_addr + fx.Int64(n_groups) * fx.Int64(BAR_STRIDE * 4), + (launch + fx.Int32(1)) * n_groups, ) - comm_ops.fence_agent_acquire() + comm_ops.fence_agent_acquire() fx.barrier() @@ -237,10 +270,12 @@ def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap, rank): * prefix: ``tile_base[tile][pe]`` = rows for ``pe`` in earlier tiles (so the rows of a tile range are contiguous in the run), ``cursor[pe]`` = total * assign: ``out_row = pe*cap + tile_base + (earlier rows of this tile for pe)``, - and ``idx_send[out_row] = origin_lid*topk + k`` + and ``idx_send[out_row] = origin_lid*topk + k``. Rows bound for this rank + also get their ``inv`` entry here; the other ranks' entries are pushed by + the GEMM2 kernel's comm waves, off this kernel's critical path One rendezvous, between count and prefix; the prefix needs every earlier tile's - count. It is a ``flag_barrier`` (no single contended atomic). After it, a block + count. It is a ``tree_barrier`` (no address takes more than a group's worth of atomics). After it, a block does not scan the whole ``tile_cnt`` array: for each tile it owns it sums the counts of the earlier tiles with one pass of the block's threads (thread ``t`` handles pe ``t % npes``) and a short LDS reduction -- a parallel sum, not a @@ -250,7 +285,7 @@ def build_rank_sort_plan(*, SBM, topk, log2_max_tok, recv_cap, npes, cap, rank): Before it writes anything the kernel waits until every peer has consumed the previous call: that is what guarantees this rank's own outbound puts have - finished reading ``S`` and ``idx_send``, which this call is about to rewrite. + finished reading ``S`` (rewritten by the GEMM2 that follows). ``tile_base`` is clamped to ``cap``, so a chunk that crosses the per-pair reservation is cut at it instead of shipping the next peer's rows. @@ -283,8 +318,8 @@ class PlanShared: known_block_size=[BLOCK, 1, 1]) def plan(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, - arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, - arg_over: fx.Int64, arg_done: fx.Int64, arg_counter: fx.Int64, + arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_inv: fx.Int64, arg_cdone: fx.Int64, + arg_bar: fx.Int64, arg_over: fx.Int64, arg_done: fx.Int64, arg_counter: fx.Int64, i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32): # fmt: on tid = fx.thread_idx.x @@ -307,9 +342,9 @@ def lw(i): """LDS i32 word ``i`` of the block's scratch.""" return lds_typed_ptr(lds_i32 + i * fx.Int32(4), T.i32, align=4) - # S and idx_send are rewritten from here on. A peer's done flag says it has - # read everything I sent it last call, and with the copy engine's FIFO that - # means my puts have finished reading them. + # S and the peers' inv are rewritten from here on. A peer's done flag says it + # has read everything I sent it last call, and with the copy engine's FIFO that + # means my puts have finished reading S and it has reset its inv. consumed = global_typed_ptr(arg_counter, T.i32)[0] if tid < fx.Int32(npes): if tid != fx.Int32(rank): @@ -355,7 +390,7 @@ def tile_masks(tile): total = total + fx.Int32(fx.ptr_load(lw(fx.Int32(w * npes) + tid))) tile_cnt[tile * fx.Int32(npes) + tid] = total fx.barrier() - flag_barrier(arg_bar, i32_nblocks, tid, bid) + tree_barrier(arg_bar, i32_nblocks, tid, bid) # No rows at all: the loop below never runs, so say so explicitly. if (tiles == fx.Int32(0)) & (bid == fx.Int32(0)) & (tid < fx.Int32(npes)): @@ -363,6 +398,7 @@ def tile_masks(tile): ptr_buf_tensor(arg_cursor, fx.Int32)[tid] = fx.Int32(0) idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) + inv_buf = ptr_buf_tensor(arg_inv, fx.Int32) lane_lt = (fx.Int64(1) << fx.Int64(lane)) - fx.Int64(1) for tile in range(bid, tiles, i32_nblocks): # Rows for pe in the tiles before this one. Element e of tile_cnt is @@ -414,19 +450,22 @@ def tile_masks(tile): else: out_row = dest_pe * fx.Int32(cap) + place out_buf[g] = out_row - idx_buf[out_row] = dest_lid * fx.Int32(topk) + s + slot = dest_lid * fx.Int32(topk) + s + idx_buf[out_row] = slot + if dest_pe == fx.Int32(rank): + inv_buf[slot] = out_row fx.barrier() # fmt: off @flyc.jit def launch(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, arg_tile_cnt: fx.Int64, arg_tile_base: fx.Int64, arg_cursor: fx.Int64, - arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_cdone: fx.Int64, arg_bar: fx.Int64, - arg_over: fx.Int64, arg_done: fx.Int64, arg_counter: fx.Int64, + arg_out_row: fx.Int64, arg_idx: fx.Int64, arg_inv: fx.Int64, arg_cdone: fx.Int64, + arg_bar: fx.Int64, arg_over: fx.Int64, arg_done: fx.Int64, arg_counter: fx.Int64, i32_nblocks: fx.Int32, i32_cdone_n: fx.Int32, stream: fx.Stream): # fmt: on plan(arg_stids, arg_trb, arg_num_valid, arg_tile_cnt, arg_tile_base, arg_cursor, - arg_out_row, arg_idx, arg_cdone, arg_bar, arg_over, arg_done, arg_counter, + arg_out_row, arg_idx, arg_inv, arg_cdone, arg_bar, arg_over, arg_done, arg_counter, i32_nblocks, i32_cdone_n).launch( grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) @@ -471,7 +510,7 @@ def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): def sdma_finish_peer( sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco, signal_mode="remote" ): - """Ship idx, tell ``peer`` its rows have landed, and publish ``cnt``. + """Tell ``peer`` its rows have landed, and publish ``cnt``. Two ways to say "landed": @@ -496,19 +535,18 @@ def sdma_finish_peer( if signal_mode not in ("quiet", "remote"): raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") remote = signal_mode == "remote" - # Under `remote` this put carries the signal, so it is issued even when - # there is no idx to ship -- 4 bytes the peer never reads. - if const_expr(remote) or cnt > fx.Int32(0): + # Under `remote` this put carries the signal: 4 bytes the peer never reads. + if const_expr(remote): sdma.put( peer, win, - fx.Int64(layout.off_idx_land + rank * layout.cap * 4), + fx.Int64(layout.off_sig + rank * 4), win, - fx.Int64(layout.off_idx_send) + fx.Int64(peer) * fx.Int64(layout.cap * 4), - fx.Int64((cnt > fx.Int32(0)).select(cnt, fx.Int32(1))) * fx.Int64(4), + fx.Int64(layout.off_sig + peer * 4), + fx.Int64(4), queue, coop=cco.CoopScope.THREAD, - remote_signal=True if remote else None, + remote_signal=True, ) if const_expr(not remote): sdma.quiet_queue(peer, queue) @@ -525,16 +563,13 @@ def sdma_finish_peer( def build_combine(*, layout, rank, hidden, topk, quant="none", signal_mode="remote"): """Kernel 3: finish the transfer and reduce. - One persistent launch (grid resident), in phases: + One launch, in phases: - A. wait for every peer's flag for this call (acquire), which also says how many - rows each sent; - B. ``inv[idx[src][j]] = landing row`` for every landed row, spread over the grid; - -- grid barrier -- - C. ``out[t] = sum_k row[inv[t*topk+k]]`` in f32, one block per token in turn. A slot's + A. wait for every peer's flag for this call (acquire); + B. ``out[t] = sum_k row[inv[t*topk+k]]`` in f32, one block per token in turn. A slot's ``inv`` entry is reset to -1 once read, so it starts clean next call and nothing has to refill it; slots no route filled stay -1 and contribute zero; - D. the last block to finish publishes ``done`` to every peer and advances ``counter``. + C. the last block to finish publishes ``done`` to every peer and advances ``counter``. """ import mori.cco.device.flydsl as cco @@ -553,8 +588,8 @@ def build_combine(*, layout, rank, hidden, topk, quant="none", signal_mode="remo # fmt: off @flyc.kernel(name=f"mega_combine_p{npes}_c{cap}_h{hidden}_k{topk}_{quant}_{signal_mode}" f"_r{rank}", known_block_size=[BLOCK, 1, 1]) - def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, - arg_bar: fx.Int64, arg_fin: fx.Int64, arg_dev_comm: fx.Int64, + def combine(win: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, + arg_fin: fx.Int64, arg_dev_comm: fx.Int64, i32_nblocks: fx.Int32, i32_tokens: fx.Int32): # fmt: on tid = fx.thread_idx.x @@ -585,27 +620,9 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. comm_ops.fence_system_acquire() fx.barrier() - # B. slot -> landing row - idx_send = fx.Int64(w.lsa_ptr(rank, layout.off_idx_send)) - idx_land = fx.Int64(w.lsa_ptr(rank, layout.off_idx_land)) inv_buf = ptr_buf_tensor(arg_inv, fx.Int32) - cursor = ptr_buf_tensor(arg_cursor, fx.Int32) - gid = bid * fx.Int32(BLOCK) + tid - stride = i32_nblocks * fx.Int32(BLOCK) - for g in range(gid, fx.Int32(npes * cap), stride): - src = g // fx.Int32(cap) - is_self = src == fx.Int32(rank) - theirs = fx.Int32( - comm_ops.load_i32_global_system(flags_base + fx.Int64(src) * fx.Int64(8) + fx.Int64(4)) - ) - cnt = is_self.select(cursor[fx.Int32(rank)], theirs) - if (g - src * fx.Int32(cap)) < cnt: - base = is_self.select(idx_send, idx_land) - slot = global_typed_ptr(base + fx.Int64(g) * fx.Int64(4), T.i32)[0] - inv_buf[slot] = is_self.select(g, g + fx.Int32(npes * cap)) - grid_barrier(arg_bar, i32_nblocks, tid) - - # C. top-k reduce + + # B. top-k reduce s_base = _uniform_i64(w.lsa_ptr(rank, layout.off_s)) if const_expr(quant_fp8): # Payload is 8 fp8 bytes (2 i32) per thread; the 32-value block they @@ -664,7 +681,7 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. if tid < fx.Int32(topk): inv_buf[tok * fx.Int32(topk) + tid] = fx.Int32(-1) - # D. the last block out tells the peers this call's data has been consumed + # C. the last block out tells the peers this call's data has been consumed rocdl.s_waitcnt(0) fx.barrier() if tid == fx.Int32(0): @@ -682,11 +699,11 @@ def combine(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx. # fmt: off @flyc.jit - def launch(win: fx.Int64, arg_cursor: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, - arg_bar: fx.Int64, arg_fin: fx.Int64, arg_dev_comm: fx.Int64, + def launch(win: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, + arg_fin: fx.Int64, arg_dev_comm: fx.Int64, i32_nblocks: fx.Int32, i32_tokens: fx.Int32, stream: fx.Stream): # fmt: on - combine(win, arg_cursor, arg_inv, arg_out, arg_bar, arg_fin, arg_dev_comm, + combine(win, arg_inv, arg_out, arg_fin, arg_dev_comm, i32_nblocks, i32_tokens).launch( grid=(i32_nblocks, 1, 1), block=(BLOCK, 1, 1), stream=stream) @@ -741,7 +758,7 @@ def __init__( Communicator, ) - from moe_gemm2_localsort import CommSpec + from moe_gemm2_with_sdma import CommSpec cfg = moe.comb_cfg self.moe = moe @@ -784,7 +801,7 @@ def __init__( if quant == "fp8_blockwise_1x32" else self.hidden * 2 ) - self.layout = ArenaLayout(self.npes, self.cap, row_bytes) + self.layout = ArenaLayout(self.npes, self.cap, row_bytes, self.slot_cap) self.layout.check() dev = torch.device("cuda", self.rank) lay = self.layout @@ -815,7 +832,7 @@ def __init__( self.cursor = torch.zeros(self.npes, **i32) # One completion counter per chunk, each on its own cache line; the plan # kernel zeroes the prefix a launch can touch. - from moe_gemm2_localsort import CDONE_STRIDE + from moe_gemm2_with_sdma import CDONE_STRIDE self.cdone_stride = CDONE_STRIDE self.cdone = torch.zeros((max_tiles + 1) * CDONE_STRIDE, **i32) @@ -826,15 +843,25 @@ def __init__( ) # Rows the per-pair reservation had no room for; must stay 0. self.overflow = torch.zeros(1, **i32) - self.combine_bar = torch.zeros(1, **i32) self.combine_fin = torch.zeros(1, **i32) - # -1 means "no row for this slot"; the combine kernel puts it back after reading. - self.inv = torch.full((self.slot_cap,), -1, **i32) + # In the window, so a sender can fill it: -1 means "no row for this slot"; the + # combine kernel puts it back after reading. Set before the barrier below, so + # no peer's first plan can write into it first. + self.inv = _torch_view( + self.mem.ptr + lay.off_inv, self.slot_cap * 4, torch.int32 + ) + self.inv.fill_(-1) + # The plan's per-row slot, read back by the comm waves that push the peers' inv. + self.idx_send = torch.zeros(self.npes * self.cap, **i32) self.out = torch.empty( self.max_tok * self.hidden, dtype=torch.bfloat16, device=dev ) - # Both persistent kernels barrier across the grid, so keep it resident: one block per CU. + # The plan kernel barriers across its grid, so keep it resident: one block per CU. self.nblocks = torch.cuda.get_device_properties(dev).multi_processor_count + # The combine has no grid-wide rendezvous, so it need not be resident. Two blocks + # per CU keep more row loads in flight for the reduce: 1.5-2x measured 1-2% faster + # stage2+combine at 2048-8192 tokens per rank, 4x and up was already slower. + self.comb_blocks = 2 * self.nblocks torch.cuda.synchronize() self.comm.barrier() @@ -902,7 +929,8 @@ def _launch_plan(self, config, s_fx): fx.Int64(self.tile_base.data_ptr()), fx.Int64(self.cursor.data_ptr()), fx.Int64(self.out_row.data_ptr()), - fx.Int64(self.mem.ptr + lay.off_idx_send), + fx.Int64(self.idx_send.data_ptr()), + fx.Int64(self.inv.data_ptr()), fx.Int64(self.cdone.data_ptr()), fx.Int64(self.plan_bar.data_ptr()), fx.Int64(self.overflow.data_ptr()), @@ -915,12 +943,12 @@ def _launch_plan(self, config, s_fx): def _launch_gemm2(self, config, s_fx): import flydsl.expr as fx - from moe_gemm2_localsort import run_moe_gemm2_localsort + from moe_gemm2_with_sdma import run_moe_gemm2_with_sdma moe, op, lay = self.moe, self.moe._s1_op, self.layout ws, stage2 = moe._s1_dispatch_workspace, config.stage2 inv = moe._g2_invariants_by_quant["none"] - run_moe_gemm2_localsort( + run_moe_gemm2_with_sdma( fx.Int64(moe._s1_out.view(-1).data_ptr()), fx.Int64(moe._s1_osd.data_ptr()), fx.Int64(moe.w2.data_ptr()), @@ -937,6 +965,7 @@ def _launch_gemm2(self, config, s_fx): fx.Int64(self.cdone.data_ptr()), fx.Int64(self.tile_base.data_ptr()), fx.Int64(self.cursor.data_ptr()), + fx.Int64(self.idx_send.data_ptr()), moe._s1_nvm, moe._g2v2_inter, moe._g2v2_hidden, @@ -996,13 +1025,11 @@ def mark(stage): _run_compiled( self._combine, fx.Int64(self.win.handle), - fx.Int64(self.cursor.data_ptr()), fx.Int64(self.inv.data_ptr()), fx.Int64(self.out.data_ptr()), - fx.Int64(self.combine_bar.data_ptr()), fx.Int64(self.combine_fin.data_ptr()), fx.Int64(self.dev_comm.ptr), - fx.Int32(self.nblocks), + fx.Int32(self.comb_blocks), fx.Int32(run_tokens), s_fx, ) From 0ad1c7b7136aeb4b381c8b793a80e253c4ba7349 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=A0=C3=A7kawhil-amd?= Date: Wed, 28 Oct 2026 07:55:24 +0000 Subject: [PATCH 7/9] refactor(cco/example): one file for the SDMA mega-moe GEMM2, one for its test and benchmark Layout of examples/cco/python/11_flydsl_mega_moe_sdma/ before -> after: sorted_combine.py + moe_gemm2_with_sdma.py -> mega_moe_gemm2_with_sdma.py the host op MegaMoESdma (was SortedCombine) and its three kernels (plan, GEMM2 with comm waves, combine) with the window layout and the sender side of the SDMA protocol they share. The two files imported each other; they no longer do. main.py + bench.py -> main.py one harness around aiter's MegaMoEV2 with `test` and `bench` subcommands; input generation, weight quantisation, the aiter operator, the reference output and the timing helpers exist once. sdmabw.py -> sdma_bandwidth.py Cleanup of what the previous commit left behind or made stale: - remove grid_barrier (unused), the row count that flags[src] still carried and _cdone_used's unused argument; flags are one i32 per peer now. - comm_role: the `for peer ... if peer < npes: if peer != rank:` boilerplate that was repeated four times is one `sends_to`, and the inv push is `push_inv`. - use aiter's comm_ops.store_i32_global_system_monotonic for the per-row inv stores instead of a private helper that reached into comm_ops._llvm_d, which aiter main no longer has. The example now also runs on aiter main (9b7bb36ad, FlyDSL 0.4.0) with the same performance and bit-exact output. - docstrings and comments that still described the idx transfer, two grid barriers and a combine build phase. `bench` now compares stage2+combine only: no stage1 timing and no end-to-end rows, which were not a like-for-like comparison (aiter's included its input quantisation). README commands, file table and results are updated. No behaviour change: the SDMA output is bit-identical to the previous commit's on 8 configurations x 8 ranks (bs 128/2048/4096/8192, two skewed routings, quiet signalling, bf16 rows), and stage2+combine is unchanged within noise. Co-Authored-By: Claude Code --- examples/cco/README.md | 26 +- .../python/11_flydsl_mega_moe_sdma/bench.py | 388 -------- .../python/11_flydsl_mega_moe_sdma/main.py | 538 ++++++++--- ...combine.py => mega_moe_gemm2_with_sdma.py} | 907 +++++++++++++++--- .../moe_gemm2_with_sdma.py | 707 -------------- .../{sdmabw.py => sdma_bandwidth.py} | 4 +- 6 files changed, 1175 insertions(+), 1395 deletions(-) delete mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/bench.py rename examples/cco/python/11_flydsl_mega_moe_sdma/{sorted_combine.py => mega_moe_gemm2_with_sdma.py} (55%) delete mode 100644 examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py rename examples/cco/python/11_flydsl_mega_moe_sdma/{sdmabw.py => sdma_bandwidth.py} (98%) diff --git a/examples/cco/README.md b/examples/cco/README.md index 67e2af95ad..d7b3ab0b96 100644 --- a/examples/cco/README.md +++ b/examples/cco/README.md @@ -22,8 +22,8 @@ examples/cco/ | `python/08_torch_symm_import` | py | import a `torch.symm_mem` tensor into the flat LSA space, no copy | | `python/09_torch_symm_lsa_all2all` | py | mori's own torch symm-mem backend + Triton `lsa_ptr` all-to-all | | `python/10_flydsl_sdma_allgather_async` | py | FlyDSL SDMA all-gather split into `start_async` / `wait_async`, overlapped with a GEMM | -| `python/11_flydsl_mega_moe_sdma` | py | MoE GEMM2 (gfx950) writing a rank-sorted send buffer, shipped by SDMA `put`s from comm waves (one per peer) inside the GEMM kernel (needs aiter) | -| `python/11_flydsl_mega_moe_sdma/sdmabw.py` | py | what one cco SDMA channel moves, with no GEMM in the way | +| `python/11_flydsl_mega_moe_sdma` | py | MegaMoE stage2 + combine (gfx950) with the transfer done by SDMA: GEMM2 writes a rank-sorted send buffer and comm waves (one per peer) inside the GEMM kernel `put` it; `mega_moe_gemm2_with_sdma.py` is the op and its kernels, `main.py` its test and benchmark against aiter (needs aiter) | +| `python/11_flydsl_mega_moe_sdma/sdma_bandwidth.py` | py | what one cco SDMA channel moves, with no GEMM in the way | | `ir/test_triton_cco.py` | py | Triton DevComm queries and LSA/SDMA/GDA device APIs | | `cpp/01_lsa_put.cpp` | c++ | intra-node LSA put (includes only `cco.hpp`) | | `cpp/02_gda_put.cpp` | c++ | GPU-initiated RDMA put + signal/wait (includes only `cco_scale_out.hpp`) | @@ -97,30 +97,28 @@ Example 11 also uses `torchrun` and SDMA, and additionally needs aiter (it reuse ```bash MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ - examples/cco/python/11_flydsl_mega_moe_sdma/main.py --network glm52 --bs 4096 + examples/cco/python/11_flydsl_mega_moe_sdma/main.py test --network glm52 --bs 4096 # latency vs aiter's MegaMoEV2 (mori.shmem P2P scatter), sweeping token counts MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ - examples/cco/python/11_flydsl_mega_moe_sdma/bench.py --network glm52 \ - --bs-list 128,512,2048,4096,8192 --no-aligned-pair + examples/cco/python/11_flydsl_mega_moe_sdma/main.py bench --network glm52 \ + --bs-list 128,512,2048,4096,8192,16384 --no-aligned-pair # what one SDMA channel actually moves, to size the above against the hardware MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 \ - examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py + examples/cco/python/11_flydsl_mega_moe_sdma/sdma_bandwidth.py ``` It is worth it only at large per-rank batches. `stage2+combine` against aiter, eager, fp8 transport, `--chunk-tiles 16`, glm52 on 8x MI355, matching aiter's -output exactly (relL2 0) at every size: +output exactly (relL2 0) at every size (`--no-aligned-pair` for 8192 and up): -| tokens/rank | 128 | 512 | 2048 | 4096 | 8192 | -|---|---|---|---|---|---| -| vs aiter | 0.63x | 0.69x | 0.97x | **1.08x** | **1.18x** | -| whole MoE layer | 0.86x | 0.84x | 1.00x | 1.05x | **1.08x** | +| tokens/rank | 128 | 512 | 2048 | 4096 | 8192 | 16384 | +|---|---|---|---|---|---|---| +| vs aiter | 0.71x | 0.78x | 1.01x | **1.12x** | **1.23x** | **1.35x** | -Below ~2048 the three kernels' fixed costs -- the plan launch and its grid -barrier, and the combine's extra index inversion -- outweigh what the copy -engine saves, and aiter's P2P scatter wins. +Below ~2048 the fixed costs of the extra plan kernel and of the combine outweigh +what the copy engine saves, and aiter's P2P scatter wins. `MORI_CCO_GDA_CONN=full` is required for GDA (03, 06) when both ranks share one node; LSA examples (02, 04, 05) ignore it. Each example prints `SUCCESS` on pass. diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py b/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py deleted file mode 100644 index 3e2a41cfb6..0000000000 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/bench.py +++ /dev/null @@ -1,388 +0,0 @@ -#!/usr/bin/env python3 -# 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. -"""Latency: aiter MegaMoEV2 stage2+combine vs the rank-sorted SDMA path (gfx950) - -Baseline: aiter's ``MegaMoEV2`` (mori.shmem symmetric heap). Its fused GEMM2 scatters -every output row into the origin rank's ``comb_inp`` with P2P stores from the compute -units, then a combine kernel runs a cross-rank barrier and the top-k reduce. No cco, -no LSA, no copy engine. - -New path (``sorted_combine.SortedCombine``), three kernels: a plan kernel, GEMM2 writing -a rank-sorted local send buffer with comm waves inside it that put finished chunks with -cco SDMA while the tiles are still computing, and a combine kernel that waits for the -peers and does an indexed top-k reduce. - -Per token count it reports, as mean/max over ranks: - - stage2+combine baseline vs the new path (the part that is replaced) - end to end stage1 + that (dispatch and stage1 are shared, so e2e shows - how much of the layer the replaced part is) - breakdown per-stage time of one eager call (rank 0), to see where the - new path spends its time - -Eager launches by default, which is what prefill does; ``--graph`` switches to -the two-graph slope used by ``bench_ep.py``. - - MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 bench.py \\ - --network glm52 --bs-list 16,128,1024,4096 -""" - -import argparse -from dataclasses import replace - -from main import NETWORKS, _make_inputs, _next_pow2, _quantize_weights, _setup - -import torch -import torch.distributed as dist - -import mori.shmem as ms - -from aiter.ops.flydsl.kernels.mega_moe import MegaMoEV2 - -from sorted_combine import SortedCombine - - -def _over_ranks(local_ms): - """(mean, max) of a per-rank time over all ranks.""" - t = torch.tensor([local_ms], dtype=torch.float64) - mx = t.clone() - dist.all_reduce(t, op=dist.ReduceOp.SUM) - dist.all_reduce(mx, op=dist.ReduceOp.MAX) - return float(t.item()) / dist.get_world_size(), float(mx.item()) - - -def _sync(): - torch.cuda.synchronize() - dist.barrier() - - -#: Calls captured per graph for the slope measurement, as ``bench_ep.py`` does. -GRAPH_REPS = 10 - - -def _replay_ms(g, n): - """Mean ms of one ``g.replay()`` over ``n`` replays.""" - _sync() - start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( - enable_timing=True - ) - start.record() - for _ in range(n): - g.replay() - end.record() - torch.cuda.synchronize() - return start.elapsed_time(end) / n - - -def _graph_of(fn, reps): - """Capture ``reps`` back-to-back calls, warm it up, lockstep the ranks.""" - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g, stream=torch.cuda.Stream()): - for _ in range(reps): - fn() - _sync() - for _ in range(5): - g.replay() - _sync() - return g - - -def _time(fn, iters, graph, reps=GRAPH_REPS): - """Mean/max-over-ranks ms per call. - - With ``graph``, this is the *marginal* cost of one call: two graphs holding - ``reps`` and ``2*reps`` calls are replayed and differenced, so the per-replay - launch overhead -- which at bs=16 is the same order as the work -- cancels - instead of being charged to every call. Same two-graph slope as - ``mori/tests/python/ops/dispatch_combine_v2/bench_ep.py``. Falls back to - eager timing if the capture fails. - """ - _sync() - fn() - _sync() - if graph: - try: - g1, g2 = _graph_of(fn, reps), _graph_of(fn, 2 * reps) - t1, t2 = _replay_ms(g1, iters), _replay_ms(g2, iters) - return _over_ranks((t2 - t1) / reps) - except Exception as exc: # noqa: BLE001 - if dist.get_rank() == 0: - print( - f" (graph capture failed, timing eagerly: {type(exc).__name__}: {exc})", - flush=True, - ) - torch.cuda.synchronize() - for _ in range(5): - fn() - _sync() - start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( - enable_timing=True - ) - start.record() - for _ in range(iters): - fn() - end.record() - torch.cuda.synchronize() - return _over_ranks(start.elapsed_time(end) / iters) - - -def _breakdown(sc, bs, config): - """Per-stage ms of one eager call, averaged over ranks.""" - sc.run(bs, config) # warm - _sync() - marks = [] - sc.run(bs, config, profile=marks) - torch.cuda.synchronize() - names = [n for n, _ in marks[1:]] - ms_ = [marks[i][1].elapsed_time(marks[i + 1][1]) for i in range(len(marks) - 1)] - t = torch.tensor(ms_, dtype=torch.float64) - dist.all_reduce(t, op=dist.ReduceOp.SUM) - t /= dist.get_world_size() - return dict(zip(names, t.tolist())) - - -def _rel_l2(out, ref): - diff = out.float() - ref.float() - local = torch.tensor( - [float(diff.norm() ** 2), float(ref.float().norm() ** 2)], dtype=torch.float64 - ) - dist.all_reduce(local) - return float((local[0] / local[1].clamp_min(1e-30)).sqrt()) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--network", choices=NETWORKS, default="glm52") - parser.add_argument( - "--bs-list", default="128,512,2048,4096,8192", help="tokens per rank" - ) - parser.add_argument("--iters", type=int, default=30) - parser.add_argument("--seed", type=int, default=123) - parser.add_argument("--rtol", type=float, default=0.05) - parser.add_argument( - "--chunk-tiles", default="16", help="fused chunk sizes to sweep, e.g. 8,16,32" - ) - parser.add_argument( - "--queues", - type=int, - default=1, - help="SDMA channels per (src, dst) pair. More than one measured " - "slower: KFD recommends a single engine per peer link here, so " - "extra channels share it and only add descriptors.", - ) - parser.add_argument("--row-slack", type=float, default=1.5) - parser.add_argument( - "--signal-mode", - choices=("quiet", "remote"), - default="remote", - help="how a sender tells a peer its rows landed", - ) - parser.add_argument( - "--no-aligned-pair", - action="store_true", - help="force aiter off its common-pair stage2 so both paths run the " - "same config. It is picked only at bucket 8192, where SortedCombine " - "has no equivalent -- and measured 1.5%% slower than without it on " - "uniform routing anyway.", - ) - parser.add_argument( - "--quant", - choices=("none", "fp8_blockwise_1x32"), - default="fp8_blockwise_1x32", - help="quantise the send buffer; halves the bytes SDMA moves", - ) - parser.add_argument( - "--graph", - action="store_true", - help="time CUDA Graph replays instead of eager launches. Off by " - "default: prefill does not capture, so eager is the shape that matters.", - ) - args = parser.parse_args() - - batch_sizes = [int(v) for v in args.bs_list.split(",")] - chunk_list = [int(v) for v in args.chunk_tiles.split(",")] - graph = args.graph - - rank, world, device = _setup() - rows = [] # (bs, label, s2_mean, s2_max, e2e_mean, e2e_max, relL2) - try: - net = NETWORKS[args.network] - w1, w1s, w2, w2s = _quantize_weights( - net["model_dim"], - net["inter_dim"], - net["experts"] // world, - rank, - args.seed, - device, - ) - max_bs = max(batch_sizes) - x_all, wts_all, ids_all = _make_inputs( - max_bs, - net["model_dim"], - net["experts"], - net["topk"], - rank, - args.seed, - device, - ) - # One operator and one SortedCombine for every token count: building more would - # re-create symmetric heaps and SDMA queues. - moe = MegaMoEV2( - rank=rank, - world_size=world, - quant="a8w4", - w1=w1, - w1_scale=w1s, - w2=w2, - w2_scale=w2s, - max_tok_per_rank=max(16, _next_pow2(max_bs)), - fanout_masks=(), - **net, - ) - sc = SortedCombine( - moe, - queues=args.queues, - quant=args.quant, - row_slack=args.row_slack, - signal_mode=args.signal_mode, - ) - - for bs in batch_sizes: - x, wts, ids = ( - x_all[:bs].contiguous(), - wts_all[:bs].contiguous(), - ids_all[:bs].contiguous(), - ) - ref = moe(x, wts, ids)[:bs].clone() - torch.cuda.synchronize() - config = moe._active_config - if args.no_aligned_pair and config.stage2.aligned_pair: - config = replace( - config, - stage2=replace( - config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96 - ), - ) - if rank == 0: - print( - f"\n=== bs={bs} tokens/rank p2p_quant={config.p2p_quant} " - f"BM={moe._g2_active_block_m} SBM={moe._s1_active_tile_m} " - f"aligned_pair={config.stage2.aligned_pair} ===", - flush=True, - ) - if config.stage2.aligned_pair: - if rank == 0: - print( - " aligned-pair stage2 is not covered by SortedCombine; skipped", - flush=True, - ) - continue - - x_q, scale = moe.quantize(x) - - def stage1(): - moe._run_fused_stage1(x_q, wts, scale, ids) - - def base_s2(): - moe._run_stage2(bs, None, True, config) - - def base_e2e(): - moe(x, wts, ids) - - s1 = _time(stage1, args.iters, graph) - stage1() - _sync() - b_s2 = _time(base_s2, args.iters, graph) - b_e2e = _time(base_e2e, args.iters, graph) - rows.append((bs, "aiter MegaMoEV2 (P2P scatter)", *b_s2, *b_e2e, None)) - if rank == 0: - print( - f" stage1 alone {s1[0]:9.4f} / {s1[1]:9.4f} ms (mean/max)", - flush=True, - ) - print( - f" baseline stage2+combine {b_s2[0]:9.4f} / {b_s2[1]:9.4f} ms", - flush=True, - ) - print( - f" baseline end to end {b_e2e[0]:9.4f} / {b_e2e[1]:9.4f} ms", - flush=True, - ) - - stage1() - _sync() - for ct in chunk_list: - sc.set_chunk_tiles(ct) - label = f"SDMA 3-kernel CT={ct}" - out = sc.run(bs, config).clone() - torch.cuda.synchronize() - rel = _rel_l2(out, ref) - n_s2 = _time(lambda: sc.run(bs, config), args.iters, graph) - - def e2e(): - stage1() - sc.run(bs, config) - - n_e2e = _time(e2e, args.iters, graph) - rows.append((bs, label, *n_s2, *n_e2e, rel)) - br = _breakdown(sc, bs, config) - if rank == 0: - parts = " ".join(f"{k}={v:.3f}" for k, v in br.items()) - print( - f" {label:<22} stage2+combine {n_s2[0]:9.4f} / {n_s2[1]:9.4f} ms " - f"e2e {n_e2e[0]:9.4f} / {n_e2e[1]:9.4f} ms relL2={rel:.2e}", - flush=True, - ) - print(f" breakdown (eager, ms): {parts}", flush=True) - if rel >= args.rtol: - raise AssertionError( - f"bs={bs} {label}: relL2={rel:.3e} exceeds {args.rtol}" - ) - sc.check_overflow() - - if rank == 0: - print( - "\n| bs | path | stage2+combine mean ms | max ms | vs baseline | " - "e2e mean ms | vs baseline | relL2 |" - ) - print("|---:|---|---:|---:|---:|---:|---:|---:|") - base = {} - for bs, label, s2m, s2x, e2m, e2x, rel in rows: - if rel is None: - base[bs] = (s2m, e2m) - b_s2, b_e2e = base.get(bs, (s2m, e2m)) - print( - f"| {bs} | {label} | {s2m:.4f} | {s2x:.4f} | {b_s2 / s2m:.2f}x | " - f"{e2m:.4f} | {b_e2e / e2m:.2f}x | {'-' if rel is None else f'{rel:.1e}'} |" - ) - print("\nSUCCESS", flush=True) - sc.close() - finally: - try: - ms.shmem_finalize() - finally: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py index 7888b68f84..355f4f1287 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py @@ -20,7 +20,7 @@ # 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. -"""CCO Example 11 — MoE GEMM2 fused with an SDMA combine (gfx950) +"""CCO Example 11 -- MegaMoE stage2 + combine with the transfer done by SDMA (gfx950) MegaMoE's GEMM2 normally scatters every output row straight into the origin rank's ``comb_inp`` with P2P stores from the compute units. Here GEMM2 writes its @@ -30,23 +30,24 @@ S = [ origin rank 0: rows | origin rank 1: rows | ... ] (expert-major inside a run) -Three kernels on one stream (see ``sorted_combine.py``; the P2P baseline has two): +The SDMA path is ``mega_moe_gemm2_with_sdma.MegaMoESdma`` (the op and its three kernels -- a plan +kernel, GEMM2 with comm waves, and a combine kernel -- all in +``mega_moe_gemm2_with_sdma.py``). This file is the +harness around aiter's ``MegaMoEV2``: it runs aiter once for stage1/dispatch and the +reference output, then redoes stage2 + transfer + reduce on the same state. - 1. rank_sort_plan count -> scan -> assign: a stable send-buffer index per row - 2. gemm2 writes S[out_row] (``moe_gemm2_with_sdma.py``); extra comm waves, - one per peer, wait for a chunk of sort tiles to be written (a - per-chunk counter the GEMM workgroups bump after their stores) and - put that chunk for their peer on their own SDMA queue while the - remaining tiles are still computing - 3. combine wait for the peers' flags, build slot -> row, top-k reduce - -The example runs MegaMoEV2 (aiter) once to get the reference output and the -stage1/dispatch state, then redoes stage2 + transfer + reduce on the same state and -compares. + test the SDMA output against aiter's, on one token count; bit-exact when + both sides move fp8 rows + bench latency of aiter's stage2+combine vs the SDMA path, per token count Needs aiter (MegaMoEV2 and the GEMM2 mainloop), FlyDSL, and mori built with SDMA: - MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 main.py + MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 main.py test --bs 4096 + MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 main.py bench \\ + --bs-list 128,1024,4096,8192 + +``bench`` times eager launches, which is what prefill does; ``--graph`` times CUDA +Graph replays with the two-graph slope used by ``bench_ep.py`` instead. """ import argparse @@ -66,7 +67,7 @@ from aiter.ops.flydsl.kernels.mega_moe import MegaMoEV2 from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4 -from sorted_combine import SortedCombine +from mega_moe_gemm2_with_sdma import MegaMoESdma NETWORKS = { "glm52": dict( @@ -78,6 +79,9 @@ } +# --- Inputs, weights and the aiter operator ------------------------------------------- + + def _next_pow2(v): return 1 << (int(v) - 1).bit_length() @@ -149,12 +153,132 @@ def _setup(): return rank, world, torch.device("cuda", local) -def _timeit(fn, iters): +def _build_moe(args, max_tokens, rank, world, device): + """aiter's operator, sized for ``max_tokens`` per rank, with its quantised weights.""" + net = NETWORKS[args.network] + w1, w1s, w2, w2s = _quantize_weights( + net["model_dim"], + net["inter_dim"], + net["experts"] // world, + rank, + args.seed, + device, + ) + moe = MegaMoEV2( + rank=rank, + world_size=world, + quant="a8w4", + w1=w1, + w1_scale=w1s, + w2=w2, + w2_scale=w2s, + max_tok_per_rank=max(16, _next_pow2(max_tokens)), + fanout_masks=(), + **net, + ) + return net, moe + + +def _reference(moe, x, wts, ids, no_aligned_pair): + """aiter's own output for these inputs, and the stage2 config both paths then run.""" + ref = moe(x, wts, ids)[: x.shape[0]].clone() torch.cuda.synchronize() - dist.barrier() - fn() + config = moe._active_config + if no_aligned_pair and config.stage2.aligned_pair: + config = replace( + config, + stage2=replace(config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96), + ) + return ref, config + + +def _make_sdma(moe, args, chunk_tiles): + return MegaMoESdma( + moe, + queues=args.queues, + chunk_tiles=chunk_tiles, + quant=args.quant, + row_slack=args.row_slack, + signal_mode=args.signal_mode, + ) + + +# --- Timing --------------------------------------------------------------------------- + + +def _sync(): torch.cuda.synchronize() dist.barrier() + + +def _over_ranks(local_ms): + """(mean, max) of a per-rank time over all ranks.""" + t = torch.tensor([local_ms], dtype=torch.float64) + mx = t.clone() + dist.all_reduce(t, op=dist.ReduceOp.SUM) + dist.all_reduce(mx, op=dist.ReduceOp.MAX) + return float(t.item()) / dist.get_world_size(), float(mx.item()) + + +#: Calls captured per graph for the slope measurement, as ``bench_ep.py`` does. +GRAPH_REPS = 10 + + +def _replay_ms(g, n): + """Mean ms of one ``g.replay()`` over ``n`` replays.""" + _sync() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + start.record() + for _ in range(n): + g.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / n + + +def _graph_of(fn, reps): + """Capture ``reps`` back-to-back calls, warm it up, lockstep the ranks.""" + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch.cuda.Stream()): + for _ in range(reps): + fn() + _sync() + for _ in range(5): + g.replay() + _sync() + return g + + +def _time(fn, iters, graph=False, reps=GRAPH_REPS): + """Mean/max-over-ranks ms per call. + + With ``graph``, this is the *marginal* cost of one call: two graphs holding + ``reps`` and ``2*reps`` calls are replayed and differenced, so the per-replay + launch overhead -- which at bs=16 is the same order as the work -- cancels + instead of being charged to every call. Same two-graph slope as + ``mori/tests/python/ops/dispatch_combine_v2/bench_ep.py``. Falls back to + eager timing if the capture fails. + """ + _sync() + fn() + _sync() + if graph: + try: + g1, g2 = _graph_of(fn, reps), _graph_of(fn, 2 * reps) + t1, t2 = _replay_ms(g1, iters), _replay_ms(g2, iters) + return _over_ranks((t2 - t1) / reps) + except Exception as exc: # noqa: BLE001 + if dist.get_rank() == 0: + print( + f" (graph capture failed, timing eagerly: {type(exc).__name__}: {exc})", + flush=True, + ) + torch.cuda.synchronize() + for _ in range(5): + fn() + _sync() start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( enable_timing=True ) @@ -163,23 +287,210 @@ def _timeit(fn, iters): fn() end.record() torch.cuda.synchronize() - return start.elapsed_time(end) / iters + return _over_ranks(start.elapsed_time(end) / iters) -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--network", choices=NETWORKS, default="glm52") - parser.add_argument("--bs", type=int, default=128, help="tokens per rank") - parser.add_argument("--seed", type=int, default=123) - parser.add_argument("--iters", type=int, default=20) - parser.add_argument("--rtol", type=float, default=0.05) - parser.add_argument( - "--chunk-tiles", - type=int, - default=16, - help="sort tiles per shipped chunk", +def _breakdown(sc, bs, config): + """Per-stage ms of one eager call, averaged over ranks.""" + sc.run(bs, config) # warm + _sync() + marks = [] + sc.run(bs, config, profile=marks) + torch.cuda.synchronize() + names = [n for n, _ in marks[1:]] + ms_ = [marks[i][1].elapsed_time(marks[i + 1][1]) for i in range(len(marks) - 1)] + t = torch.tensor(ms_, dtype=torch.float64) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + t /= dist.get_world_size() + return dict(zip(names, t.tolist())) + + +def _rel_l2(out, ref): + diff = out.float() - ref.float() + local = torch.tensor( + [float(diff.norm() ** 2), float(ref.float().norm() ** 2)], dtype=torch.float64 + ) + dist.all_reduce(local) + return float((local[0] / local[1].clamp_min(1e-30)).sqrt()) + + +# --- test ----------------------------------------------------------------------------- + + +def run_test(args, rank, world, device): + net, moe = _build_moe(args, args.bs, rank, world, device) + x, wts, ids = _make_inputs( + args.bs, + net["model_dim"], + net["experts"], + net["topk"], + rank, + args.seed, + device, + hot_bias=args.hot_bias, + hot_experts=net["experts"] // world, + ) + ref, config = _reference(moe, x, wts, ids, args.no_aligned_pair) + if config.stage2.aligned_pair: + raise SystemExit( + "this config selects aligned-pair stage2, which MegaMoESdma does not cover" + ) + if rank == 0: + print( + f"reference: p2p_quant={config.p2p_quant} BM={moe._g2_active_block_m} " + f"SBM={moe._s1_active_tile_m}", + flush=True, + ) + ref_ms = _time(lambda: moe._run_stage2(args.bs, None, True, config), args.iters)[0] + + sc = _make_sdma(moe, args, args.chunk_tiles) + worst = 0.0 + mismatched = 0 + # Repeated calls with no host sync between them also exercise the epoch counter, + # the flags and the "consumed" handshake that stops a fast rank overrunning a slow one. + for _ in range(4): + out = sc.run(args.bs, config).clone() + torch.cuda.synchronize() + diff = out.float() - ref.float() + worst = max(worst, float(diff.norm() / ref.float().norm().clamp_min(1e-9))) + mismatched = max(mismatched, int((out != ref).any(dim=1).sum())) + t = torch.tensor([worst, mismatched], dtype=torch.float64) + dist.all_reduce(t, op=dist.ReduceOp.MAX) + worst, mismatched = float(t[0].item()), int(t[1].item()) + ms_per = _time(lambda: sc.run(args.bs, config), args.iters)[0] + + failures = [] + ok = worst < args.rtol + # Both sides move fp8 rows in aiter's layout and reduce in the same order, so + # they must agree exactly; relL2 alone would let a few wrong rows through. + exact = config.p2p_quant == args.quant == "fp8_blockwise_1x32" + if exact and mismatched: + ok = False + if not ok: + failures.append("sdma") + try: + sc.check_overflow() + except RuntimeError as e: + failures.append(f"overflow ({e})") + ok = False + if rank == 0: + print( + f"relL2={worst:.3e} mismatched_rows={mismatched} " + f"{'PASS' if ok else 'FAIL'} " + f"time={ms_per:.4f}ms (reference stage2+combine {ref_ms:.4f}ms)", + flush=True, + ) + sc.close() + if failures: + raise AssertionError(f"sdma path failed: {failures}") + if rank == 0: + print("SUCCESS", flush=True) + + +# --- bench ---------------------------------------------------------------------------- + + +def run_bench(args, rank, world, device): + batch_sizes = [int(v) for v in args.bs_list.split(",")] + chunk_list = [int(v) for v in args.chunk_tiles.split(",")] + graph = args.graph + rows = [] # (bs, label, s2_mean, s2_max, relL2) + + max_bs = max(batch_sizes) + net, moe = _build_moe(args, max_bs, rank, world, device) + x_all, wts_all, ids_all = _make_inputs( + max_bs, + net["model_dim"], + net["experts"], + net["topk"], + rank, + args.seed, + device, ) - parser.add_argument( + # One operator and one MegaMoESdma for every token count: building more would + # re-create symmetric heaps and SDMA queues. + sc = _make_sdma(moe, args, chunk_list[0]) + + for bs in batch_sizes: + x, wts, ids = ( + x_all[:bs].contiguous(), + wts_all[:bs].contiguous(), + ids_all[:bs].contiguous(), + ) + ref, config = _reference(moe, x, wts, ids, args.no_aligned_pair) + if rank == 0: + print( + f"\n=== bs={bs} tokens/rank p2p_quant={config.p2p_quant} " + f"BM={moe._g2_active_block_m} SBM={moe._s1_active_tile_m} " + f"aligned_pair={config.stage2.aligned_pair} ===", + flush=True, + ) + if config.stage2.aligned_pair: + if rank == 0: + print( + " aligned-pair stage2 is not covered by MegaMoESdma; skipped", + flush=True, + ) + continue + + def base_s2(): + moe._run_stage2(bs, None, True, config) + + b_s2 = _time(base_s2, args.iters, graph) + rows.append((bs, "aiter MegaMoEV2 (P2P scatter)", *b_s2, None)) + if rank == 0: + print( + f" baseline stage2+combine {b_s2[0]:9.4f} / {b_s2[1]:9.4f} ms (mean/max)", + flush=True, + ) + + for ct in chunk_list: + sc.set_chunk_tiles(ct) + label = f"SDMA 3-kernel CT={ct}" + out = sc.run(bs, config).clone() + torch.cuda.synchronize() + rel = _rel_l2(out, ref) + n_s2 = _time(lambda: sc.run(bs, config), args.iters, graph) + rows.append((bs, label, *n_s2, rel)) + br = _breakdown(sc, bs, config) + if rank == 0: + parts = " ".join(f"{k}={v:.3f}" for k, v in br.items()) + print( + f" {label:<22} stage2+combine {n_s2[0]:9.4f} / {n_s2[1]:9.4f} ms " + f"relL2={rel:.2e}", + flush=True, + ) + print(f" breakdown (eager, ms): {parts}", flush=True) + if rel >= args.rtol: + raise AssertionError( + f"bs={bs} {label}: relL2={rel:.3e} exceeds {args.rtol}" + ) + sc.check_overflow() + + if rank == 0: + print("\n| bs | path | stage2+combine mean ms | max ms | vs baseline | relL2 |") + print("|---:|---|---:|---:|---:|---:|") + base = {} + for bs, label, s2m, s2x, rel in rows: + if rel is None: + base[bs] = s2m + print( + f"| {bs} | {label} | {s2m:.4f} | {s2x:.4f} | {base.get(bs, s2m) / s2m:.2f}x | " + f"{'-' if rel is None else f'{rel:.1e}'} |" + ) + print("\nSUCCESS", flush=True) + sc.close() + + +# --- command line --------------------------------------------------------------------- + + +def _parse_args(): + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--network", choices=NETWORKS, default="glm52") + common.add_argument("--seed", type=int, default=123) + common.add_argument("--rtol", type=float, default=0.05) + common.add_argument( "--queues", type=int, default=1, @@ -187,142 +498,71 @@ def main(): "slower: KFD recommends a single engine per peer link here, so " "extra channels share it and only add descriptors.", ) - parser.add_argument("--row-slack", type=float, default=1.5) - parser.add_argument( - "--hot-bias", - type=float, - default=0.0, - help="add this to the router scores of rank 0's experts: skewed routing, " - "where one rank receives more than an even share of the rows", - ) - parser.add_argument( + common.add_argument("--row-slack", type=float, default=1.5) + common.add_argument( "--signal-mode", choices=("quiet", "remote"), default="remote", help="how a sender tells a peer its rows landed", ) - parser.add_argument( + common.add_argument( "--no-aligned-pair", action="store_true", help="force aiter off its common-pair stage2 so both paths run the " - "same config. It is picked only at bucket 8192, where SortedCombine " + "same config. It is picked only at bucket 8192, where MegaMoESdma " "has no equivalent -- and measured 1.5%% slower than without it on " "uniform routing anyway.", ) - parser.add_argument( + common.add_argument( "--quant", choices=("none", "fp8_blockwise_1x32"), default="fp8_blockwise_1x32", help="quantise the send buffer; halves the bytes SDMA moves", ) - args = parser.parse_args() + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + sub = parser.add_subparsers(dest="cmd", required=True) + + test = sub.add_parser( + "test", parents=[common], help="SDMA output vs aiter's, one token count" + ) + test.add_argument("--bs", type=int, default=128, help="tokens per rank") + test.add_argument("--iters", type=int, default=20) + test.add_argument( + "--chunk-tiles", type=int, default=16, help="sort tiles per shipped chunk" + ) + test.add_argument( + "--hot-bias", + type=float, + default=0.0, + help="add this to the router scores of rank 0's experts: skewed routing, " + "where one rank receives more than an even share of the rows", + ) + + bench = sub.add_parser( + "bench", parents=[common], help="latency vs aiter, per token count" + ) + bench.add_argument( + "--bs-list", default="128,512,2048,4096,8192", help="tokens per rank" + ) + bench.add_argument("--iters", type=int, default=30) + bench.add_argument( + "--chunk-tiles", default="16", help="fused chunk sizes to sweep, e.g. 8,16,32" + ) + bench.add_argument( + "--graph", + action="store_true", + help="time CUDA Graph replays instead of eager launches. Off by " + "default: prefill does not capture, so eager is the shape that matters.", + ) + return parser.parse_args() + + +def main(): + args = _parse_args() rank, world, device = _setup() try: - net = NETWORKS[args.network] - w1, w1s, w2, w2s = _quantize_weights( - net["model_dim"], - net["inter_dim"], - net["experts"] // world, - rank, - args.seed, - device, - ) - x, wts, ids = _make_inputs( - args.bs, - net["model_dim"], - net["experts"], - net["topk"], - rank, - args.seed, - device, - hot_bias=args.hot_bias, - hot_experts=net["experts"] // world, - ) - moe = MegaMoEV2( - rank=rank, - world_size=world, - quant="a8w4", - w1=w1, - w1_scale=w1s, - w2=w2, - w2_scale=w2s, - max_tok_per_rank=max(16, _next_pow2(args.bs)), - fanout_masks=(), - **net, - ) - ref = moe(x, wts, ids)[: args.bs].clone() - torch.cuda.synchronize() - config = moe._active_config - if args.no_aligned_pair and config.stage2.aligned_pair: - config = replace( - config, - stage2=replace( - config.stage2, aligned_pair=False, pair_cu=0, skew_cu=96 - ), - ) - if config.stage2.aligned_pair: - raise SystemExit( - "this config selects aligned-pair stage2, which SortedCombine does not cover" - ) - if rank == 0: - print( - f"reference: p2p_quant={config.p2p_quant} BM={moe._g2_active_block_m} " - f"SBM={moe._s1_active_tile_m}", - flush=True, - ) - - ref_ms = _timeit( - lambda: moe._run_stage2(args.bs, None, True, config), args.iters - ) - failures = [] - sc = SortedCombine( - moe, - queues=args.queues, - chunk_tiles=args.chunk_tiles, - quant=args.quant, - row_slack=args.row_slack, - signal_mode=args.signal_mode, - ) - worst = 0.0 - mismatched = 0 - # Repeated calls with no host sync between them also exercise the epoch counter, - # the flags and the "consumed" handshake that stops a fast rank overrunning a slow one. - for _ in range(4): - out = sc.run(args.bs, config).clone() - torch.cuda.synchronize() - diff = out.float() - ref.float() - worst = max(worst, float(diff.norm() / ref.float().norm().clamp_min(1e-9))) - mismatched = max(mismatched, int((out != ref).any(dim=1).sum())) - t = torch.tensor([worst, mismatched], dtype=torch.float64) - dist.all_reduce(t, op=dist.ReduceOp.MAX) - worst, mismatched = float(t[0].item()), int(t[1].item()) - ms_per = _timeit(lambda: sc.run(args.bs, config), args.iters) - ok = worst < args.rtol - # Both sides move fp8 rows in aiter's layout and reduce in the same order, so - # they must agree exactly; relL2 alone would let a few wrong rows through. - exact = config.p2p_quant == args.quant == "fp8_blockwise_1x32" - if exact and mismatched: - ok = False - if not ok: - failures.append("sdma") - try: - sc.check_overflow() - except RuntimeError as e: - failures.append(f"overflow ({e})") - ok = False - if rank == 0: - print( - f"relL2={worst:.3e} mismatched_rows={mismatched} " - f"{'PASS' if ok else 'FAIL'} " - f"time={ms_per:.4f}ms (reference stage2+combine {ref_ms:.4f}ms)", - flush=True, - ) - sc.close() - if failures: - raise AssertionError(f"sdma path failed: {failures}") - if rank == 0: - print("SUCCESS", flush=True) + {"test": run_test, "bench": run_bench}[args.cmd](args, rank, world, device) finally: try: ms.shmem_finalize() diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py b/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py similarity index 55% rename from examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py rename to examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py index 17c52a424f..a961a7f9da 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sorted_combine.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 # Copyright © Advanced Micro Devices, Inc. All rights reserved. # # MIT License @@ -20,7 +19,11 @@ # 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. -"""Rank-sorted send buffer for gfx950 MegaMoE stage2 + combine, shipped by cco SDMA. +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +# ruff: noqa: B023, I001 +"""MegaMoE stage2 + combine for gfx950 with the transfer done by cco SDMA: the host op ``MegaMoESdma`` and +its three kernels (the plan kernel, GEMM2 with comm waves, and the combine kernel). Stage2 normally scatters every GEMM2 row straight into the origin rank's ``comb_inp``. Here the rows go to a local send buffer ``S`` grouped by origin @@ -35,12 +38,12 @@ Three kernels on one stream (the P2P-scatter baseline has two): 1. rank_sort_plan persistent grid; count -> prefix -> assign with one rendezvous: - out_row per sorted row, tile_base, cursor; zeroes cdone; and writes - every row's slot -> landing-row entry into its DESTINATION's ``inv`` - 2. gemm2_with_sdma writes S[out_row]; extra comm waves (one per peer) put finished - chunks with SDMA while the GEMM tiles are still computing - 3. combine persistent grid; wait for the peers' flags -> top-k reduce - -> publish "consumed" + out_row per sorted row, idx_send, tile_base, cursor, and ``inv`` + for the rows that stay on this rank; zeroes cdone + 2. gemm2_with_sdma writes S[out_row]; extra comm waves (one per peer) fill that peer's + ``inv``, then put finished chunks with SDMA while the GEMM tiles + are still computing + 3. combine wait for the peers' flags -> top-k reduce -> publish "consumed" Everything lives in ONE cco window, so ``put`` names the same window for source and destination and only the byte offsets differ:: @@ -48,12 +51,12 @@ [ S | R | inv | sig | flags | done | counter ] ``inv[slot]`` is the landing row of the row for token slot ``slot`` (``origin_lid*topk -+ k``), or -1. The SENDER writes it, in the plan kernel, where a row's destination, -slot and position are all known at once -- so no index travels with the data and the -receiver builds no table. ``sig`` is scratch for the 4-byte put that carries the -completion signal. ++ k``), or -1. The SENDER writes it -- the plan kernel for rows that stay here, the +sender's comm waves (direct stores into the peer's window) for the rest -- so no index +travels with the data and the receiver builds no table. ``sig`` is scratch for the +4-byte put that carries the completion signal. -``flags[src]`` says "src's data for call n has landed" (and how many rows). +``flags[src]`` says "src's data for call n has landed". ``done[dst]`` says "dst has finished reading call n": a sender waits for it before overwriting that peer's ``R`` and ``inv`` with call n+1, which keeps a fast rank from running into a slow one. ``counter`` is the call number; only the combine @@ -79,7 +82,6 @@ The plan kernel rendezvouses across its grid, so that grid must be fully resident (one block per CU is). The combine has no grid-wide rendezvous and is not resident. """ - from dataclasses import dataclass import flydsl.compiler as flyc @@ -88,6 +90,7 @@ from flydsl.expr.rocdl import cvt_pk_f32_fp8 from flydsl.expr.typing import Int8, T from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.runtime.device import get_rocm_arch from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops @@ -98,27 +101,58 @@ # not export comm_ops.traced yet. if not hasattr(comm_ops, "traced"): comm_ops.traced = ASTRewriter.transform - from aiter.ops.flydsl.kernels.buffer_ops import ( buffer_load, buffer_store, create_buffer_resource_from_addr, ) from aiter.ops.flydsl.kernels.mxfp4_gemm_common import ( + _fabs_f32 as fabs_f32, global_typed_ptr, lds_typed_ptr, + lds_vec_load, ) from aiter.ops.flydsl.kernels.tensor_shim import ( + _run_compiled, buf_copy_load, buf_copy_store, ptr_buf_tensor, ) +from aiter.ops.flydsl.kernels.mega_moe.gemm2 import ( + _resolve_g2_knobs, + _spart_output_tile_index, + gemm2_compute_v2, + issue_a_load_lds_dt, + kStages, +) +from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage2 import ( + _BUFFER_OFFSET_ABI_BYTES, + _fp8_scale_for_leader, + _stage2_lds_bytes, +) + + +#: i32 slots between consecutive chunk-completion counters. Every unit of a +#: chunk bumps its counter, so packing them adjacent puts all of a launch's +#: atomics on one or two cache lines: 128 B apart instead is worth ~8% of the +#: GEMM2 kernel at bs=4096. Spreading a chunk's count over several counters on +#: top of this buys nothing -- the atomic itself is not the cost. +CDONE_STRIDE = 32 + +#: Rows per lane per loop trip when a comm wave pushes inv entries to a peer. +PUSH_U = 8 + + BLOCK = 256 WAVE = 64 BAR_STRIDE = 32 # ints per tree_barrier counter: one 128 B cache line BAR_GROUP = 32 # blocks per tree_barrier group MIN_SBM = 16 + + +# --- The cco window and the sender side of the SDMA protocol --------------------------- + _ALIGN = 128 @@ -161,7 +195,7 @@ def off_flags(self): @property def off_done(self): - return self.off_flags + _align(self.npes * 8) + return self.off_flags + _align(self.npes * 4) @property def off_counter(self): @@ -177,26 +211,101 @@ def check(self): raise ValueError("S+R exceed the 4 GiB buffer-resource range") -def _uniform_i64(x): - """A wave-uniform i64 (SGPR) copy of ``x``, as a buffer-resource base needs.""" - return fx.Int64(rocdl.readfirstlane(T.i64, fx.Int64(x).ir_value())) +def sdma_wait_consumed(w, peer, phase, *, layout, rank): + """Block until ``peer`` has finished reading the previous call's data from my sends.""" + comm_ops.spin_until_ge_i32_system( + fx.Int64(w.lsa_ptr(rank, layout.off_done)) + fx.Int64(peer) * fx.Int64(4), + phase - fx.Int32(1), + acquire=True, + ) -def _store_i32_relaxed_system(addr_i64, val): - """Relaxed i32 store that is visible system-wide once it retires. +def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): + """Ship rows ``[lo, hi)`` of my run for ``peer``. + + One put for the whole run. Puts are asynchronous -- each reserves a slot, + writes the copy descriptor and rings the doorbell -- but a queue is a ring + its engine drains in order, so splitting a run across channels only buys + parallelism if those channels sit on different engines. KFD recommends a + single SDMA engine per peer link on this topology (one bit per link in + ``recommended_sdma_engine_id_mask``), so they do not: splitting a 1.44 MB run + in two measured 3% slower, rotating whole runs across two channels 2%. + """ + rb, cap = layout.row_bytes, layout.cap + sdma.put( + peer, + win, + fx.Int64(layout.off_r + rank * cap * rb) + fx.Int64(lo) * fx.Int64(rb), + win, + fx.Int64(layout.off_s) + + (fx.Int64(peer) * fx.Int64(cap) + fx.Int64(lo)) * fx.Int64(rb), + fx.Int64(hi - lo) * fx.Int64(rb), + queue, + coop=cco.CoopScope.THREAD, + ) + + +@comm_ops.traced +def sdma_finish_peer( + sdma, w, win, peer, queue, phase, *, layout, rank, cco, signal_mode="remote" +): + """Tell ``peer`` its rows have landed. + + Two ways to say "landed": + + ``quiet`` drain our queue here, then store the flag. The drain waits for + the last chunk to reach the peer, and that wait sits at the tail + of the GEMM2 kernel, after the last tile -- 42us of a 772us + kernel at bs=4096, CT=16. + ``remote`` hang the copy engine's own trailing ATOMIC on this put, so the + engine bumps a counter on the peer once the bytes are there. + Nothing is drained: the comm wave issues and leaves, and the + peer's combine -- which was going to wait for us anyway -- polls + the counter. Worth 6-12us end to end. + + ``remote`` needs this put to be the LAST op on the channel the data went out + on. It is: everything for ``peer`` uses channel ``peer % queues``, and a + queue is a FIFO its engine drains in order, so the trailing ATOMIC fires + only once every chunk ahead of it has landed. - Unlike ``comm_ops.store_i32_global_system_release`` it carries no release, so - it does not write back the whole L2; it is the store for a row's ``inv`` entry - in a peer's window, one per row. + Either way the flag carries the call number. Under ``remote`` it says nothing + about the data -- the signal does that. """ - comm_ops._llvm_d.StoreOp( - val.ir_value(), - comm_ops._to_ptr_global(addr_i64), - alignment=4, - ordering=comm_ops._llvm_d.AtomicOrdering.monotonic, + if signal_mode not in ("quiet", "remote"): + raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") + remote = signal_mode == "remote" + # Under `remote` this put carries the signal: 4 bytes the peer never reads. + if const_expr(remote): + sdma.put( + peer, + win, + fx.Int64(layout.off_sig + rank * 4), + win, + fx.Int64(layout.off_sig + peer * 4), + fx.Int64(4), + queue, + coop=cco.CoopScope.THREAD, + remote_signal=True, + ) + if const_expr(not remote): + sdma.quiet_queue(peer, queue) + # The bytes moved by the copy engine, not by this CU, so nothing of ours + # has to be written back -- this only keeps the flag store from being + # hoisted above the drain. + rocdl.s_waitcnt(vmcnt=0) + comm_ops.store_i32_global_system_release( + fx.Int64(w.lsa_ptr(peer, layout.off_flags)) + fx.Int64(rank * 4), phase ) +# --- Kernel 1: rank sort plan -------------------------------------------------- + + +def _uniform_i64(x): + """A wave-uniform i64 (SGPR) copy of ``x``, as a buffer-resource base needs.""" + return fx.Int64(rocdl.readfirstlane(T.i64, fx.Int64(x).ir_value())) + + def _decode(p, *, log2_max_tok, recv_cap, topk, npes): """Packed srcmap word -> (valid, dest_pe, dest_lid, slot_k), as the stage2 epilogue does.""" t = p & fx.Int32(0x00FFFFFF) @@ -207,24 +316,6 @@ def _decode(p, *, log2_max_tok, recv_cap, topk, npes): return valid, dest_pe, dest_lid, s -@comm_ops.traced -def grid_barrier(bar_addr, nblocks, tid): - """All blocks of a resident grid rendezvous. ``bar_addr`` is a never-reset i32 counter. - - Barrier k of a launch sees arrivals ``(k-1)*nblocks .. k*nblocks-1``, so the target - ``(old // nblocks + 1) * nblocks`` works for any number of barriers without a reset. - """ - rocdl.s_waitcnt(0) - fx.barrier() - if tid == fx.Int32(0): - comm_ops.fence_agent_release() - old = fx.Int32(comm_ops.atomic_add_agent(bar_addr, fx.Int32(1))) - target = (old // nblocks + fx.Int32(1)) * nblocks - comm_ops.spin_until_ge_i32_agent(bar_addr, target) - comm_ops.fence_agent_acquire() - fx.barrier() - - @comm_ops.traced def tree_barrier(bar_addr, nblocks, tid, bid): """All blocks of a resident grid rendezvous with two levels of small atomics. @@ -472,93 +563,641 @@ def launch(arg_stids: fx.Int64, arg_trb: fx.Int64, arg_num_valid: fx.Int64, return launch -def sdma_wait_consumed(w, peer, phase, *, layout, rank): - """Block until ``peer`` has finished reading the previous call's data from my sends.""" - comm_ops.spin_until_ge_i32_system( - fx.Int64(w.lsa_ptr(rank, layout.off_done)) + fx.Int64(peer) * fx.Int64(4), - phase - fx.Int32(1), - acquire=True, - ) - +# --- Kernel 2: GEMM2 with comm waves ---------------------------------------- +# +# GEMM2 whose epilogue writes weighted bf16 rows into a local, rank-sorted send buffer. +# +# Same GEMM2 mainloop and input layout as ``mega_moe_stage2``; only the output side +# differs. Instead of scattering each row into the origin rank's ``comb_inp`` it +# writes the row to ``send[out_row[row]]``, where ``out_row`` was assigned once per +# row beforehand (``build_rank_sort_plan``) so every N +# tile of a row agrees on it. A negative ``out_row`` drops the row. +# +# With ``comm`` the kernel also carries ``comm.comm_wgs`` extra leading workgroups that +# ship finished chunks of the send buffer to the peers with cco SDMA while the GEMM +# workgroups are still computing. A GEMM unit signals ``cdone[chunk]`` after its stores. +# Each comm WAVE owns its own peers (peer ``p`` belongs to wave ``p % n_comm_waves``): it +# spins on the chunk counter itself, then puts that chunk's rows for its peers on its own +# SDMA queue, so peers neither wait for each other's spin nor share a wave's issue slot. +# +# Two things about the publish are load-bearing and were each worth more than a +# factor on this kernel. The epilogue stores the rows write-through (sc0|sc1) so +# that retiring them is the release a copy engine needs; releasing at system scope +# instead is a whole-L2 write-back per unit. And the per-chunk counters sit a +# cache line apart, because every unit of a chunk bumps one. -def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): - """Ship rows ``[lo, hi)`` of my run for ``peer`` (data only; idx goes at the end). - One put for the whole run. Puts are asynchronous -- each reserves a slot, - writes the copy descriptor and rings the doorbell -- but a queue is a ring - its engine drains in order, so splitting a run across channels only buys - parallelism if those channels sit on different engines. KFD recommends a - single SDMA engine per peer link on this topology (one bit per link in - ``recommended_sdma_engine_id_mask``), so they do not: splitting a 1.44 MB run - in two measured 3% slower, rotating whole runs across two channels 2%. +@dataclass(frozen=True) +class CommSpec: + """Compile-time description of the fused SDMA comm workgroups. + + ``comm_wgs`` workgroups of 4 waves each; one wave per peer needs + ``comm_wgs * 4 >= npes``. ``queues`` is the SDMA channels reserved per + (src, dst) pair; peer ``p`` uses channel ``p % queues`` for everything it + sends, so a chunk is never split across channels -- splitting measured + slower, and on this topology KFD recommends one engine per peer link + anyway, which makes extra channels pure overhead. """ - rb, cap = layout.row_bytes, layout.cap - sdma.put( - peer, - win, - fx.Int64(layout.off_r + rank * cap * rb) + fx.Int64(lo) * fx.Int64(rb), - win, - fx.Int64(layout.off_s) - + (fx.Int64(peer) * fx.Int64(cap) + fx.Int64(lo)) * fx.Int64(rb), - fx.Int64(hi - lo) * fx.Int64(rb), - queue, - coop=cco.CoopScope.THREAD, + + layout: ArenaLayout + rank: int + queues: int + chunk_tiles: int + comm_wgs: int = 1 + #: How a sender tells a peer its rows have landed; see + #: ``sdma_finish_peer``. "remote" needs the signal to ride + #: the channel the data went out on, which ``queues`` already guarantees. + signal_mode: str = "remote" + + +# fmt: off +def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, N_OUT, BM, BN, + send_nbytes, lds_outrow_off, lds_weight_off, g2_bf16_lds=False, scatter_vec=8, store_cm=0, + quant="none"): +# fmt: on + """CShuffle one GEMM2 tile into weighted BF16 rows and store them at their send-buffer index.""" + kMChunks = BM // 16 + numAccN = (BN // 4) // 16 + wave_n = BN // 4 + lane_div_16 = lane // 16 + lane_mod_16 = lane % 16 + lds_base_fptr = lds_typed_ptr(lds_acc_base, T.f32) + lds_base_bf16 = ( + lds_typed_ptr(lds_acc_base, T.bf16, align=2) + if const_expr(g2_bf16_lds) + else None ) + quant_fp8 = quant == "fp8_blockwise_1x32" + if quant not in ("none", "fp8_blockwise_1x32"): + raise ValueError(f"unsupported quant={quant!r}") + if quant_fp8 and g2_bf16_lds: + raise ValueError("fp8_blockwise_1x32 needs the f32 CShuffle (g2_bf16_lds=False)") + # Row layout matches aiter's P2P scatter: N_OUT payload bytes followed by the + # N_OUT/32 e8m0 block scales, so one contiguous SDMA put ships both. + row_nbytes = N_OUT + N_OUT // 32 if quant_fp8 else N_OUT * 2 + fx.barrier() -@comm_ops.traced -def sdma_finish_peer( - sdma, w, win, peer, queue, cnt, phase, *, layout, rank, cco, signal_mode="remote" -): - """Tell ``peer`` its rows have landed, and publish ``cnt``. + for i in range_constexpr(kMChunks): + row_base = fx.Int32(i * 16) + lane_div_16 * 4 + if const_expr(g2_bf16_lds): + w_row = [ + fx.ptr_load( + lds_typed_ptr( + fx.Int32(lds_weight_off) + (row_base + v) * fx.Int32(4), + T.f32, + align=4, + ) + ) + for v in range_constexpr(4) + ] + for J in range_constexpr(numAccN): + col = wave * fx.Int32(wave_n) + J * 16 + lane_mod_16 + vec = fx.Vector(accm[i][J]) + for v in range_constexpr(4): + idx = (row_base + v) * BN + col + if const_expr(g2_bf16_lds): + lds_base_bf16[idx] = fx.BFloat16(fx.Float32(vec[v]) * fx.Float32(w_row[v])) + else: + lds_base_fptr[idx] = fx.Float32(vec[v]) - Two ways to say "landed": + fx.barrier() - ``quiet`` drain our queue here, then store the flag. The drain waits for - the last chunk to reach the peer, and that wait sits at the tail - of the GEMM2 kernel, after the last tile -- 42us of a 772us - kernel at bs=4096, CT=16. - ``remote`` hang the copy engine's own trailing ATOMIC on this put, so the - engine bumps a counter on the peer once the bytes are there. - Nothing is drained: the comm wave issues and leaves, and the - peer's combine -- which was going to wait for us anyway -- polls - the counter. Worth 6-12us end to end. + # ``buf_copy_store``'s cache_modifier is a BufferCopy enum (0=cached, 2=nt), + # so sc0|sc1 cannot be reached through it -- flydsl's own type stub says so, + # and so does mori/ops/gemm_ar/_compat.py. ``buffer_store`` passes the value + # through as the instruction's aux field, so the write-through path uses it. + if const_expr(store_cm or quant_fp8): + output_buf = create_buffer_resource_from_addr( + send_base, num_records_bytes=send_nbytes + ) + else: + output_buf = ptr_buf_tensor( + send_base, + fx.BFloat16, + unit_elems=scatter_vec, + num_records_bytes=send_nbytes, + ) + # Inactive lanes read safe LDS and issue a bounded OOB store. + active = lane < fx.Int32(BN // scatter_vec) + col = active.select(lane * fx.Int32(scatter_vec), fx.Int32(0)) + + for row_iter in range_constexpr(BM // 4): + row = wave + fx.Int32(row_iter * 4) + row_byte_off = row * fx.Int32(4) + out_row = fx.ptr_load( + lds_typed_ptr(fx.Int32(lds_outrow_off) + row_byte_off, T.i32, align=4) + ) + out_row = rocdl.readfirstlane(T.i32, out_row.ir_value()) + idx0 = row * fx.Int32(BN) + col + if const_expr(g2_bf16_lds): + pk = fx.Vector( + lds_vec_load( + lds_acc_base, + idx0 * fx.Int32(2), + fx.Vector.make_type(scatter_vec, fx.BFloat16), + fx.BFloat16, + align=16, + ) + ) + else: + weight = fx.ptr_load( + lds_typed_ptr(fx.Int32(lds_weight_off) + row_byte_off, T.f32, align=4) + ) + weight = rocdl.readfirstlane(T.f32, weight.ir_value()) + v8 = fx.Vector( + lds_vec_load( + lds_acc_base, + idx0 * fx.Int32(4), + fx.Vector.make_type(scatter_vec, fx.Float32), + fx.Float32, + align=16, + ) + ) + weighted = fx.Vector.from_elements( + [v8[i] * weight for i in range_constexpr(scatter_vec)], + fx.Float32, + ) + if const_expr(not quant_fp8): + pk = weighted.to(fx.BFloat16) + row_base = out_row * fx.Int32(row_nbytes) + valid = (out_row >= fx.Int32(0)) & active + if const_expr(quant_fp8): + # One e8m0 per 32 contiguous values, i.e. per 32//scatter_vec lanes. + vals = [fx.Float32(weighted[i]) for i in range_constexpr(scatter_vec)] + local_max = fabs_f32(vals[0]) + for q in range_constexpr(1, scatter_vec): + local_max = local_max.maximumf(fabs_f32(vals[q])) + max_bits = local_max.bitcast(fx.Int32) + for xor_lane in (1, 2): + if xor_lane < 32 // scatter_vec: + remote = rocdl.ds_bpermute( + T.i32, (lane ^ fx.Int32(xor_lane)) * fx.Int32(4), max_bits + ) + local_max = local_max.maximumf(fx.Int32(remote).bitcast(fx.Float32)) + max_bits = local_max.bitcast(fx.Int32) + group_lanes = 32 // scatter_vec + leader_lane = lane & fx.Int32(~(group_lanes - 1)) + is_leader = (lane & fx.Int32(group_lanes - 1)) == fx.Int32(0) + e8m0 = fx.Int32( + rocdl.ds_bpermute( + T.i32, + leader_lane * fx.Int32(4), + _fp8_scale_for_leader(is_leader, local_max), + ) + ) + block_scale = (e8m0 << fx.Int32(23)).bitcast(fx.Float32) + pk_ty = T.vec(2, T.i16) + words = [] + for word in range_constexpr(scatter_vec // 4): + packed = fx.Vector.filled(2, 0, fx.Int16).ir_value() + for pair in range_constexpr(2): + e = word * 4 + pair * 2 + packed = rocdl.cvt_scalef32_pk_fp8_f32( + pk_ty, packed, vals[e].ir_value(), vals[e + 1].ir_value(), + block_scale.ir_value(), pair, + ) + words.append(fx.Vector(packed).bitcast(fx.Int32)[0]) + payload = fx.Vector.from_elements(words, fx.Int32) + off = valid.select( + row_base + n_block_idx * fx.Int32(BN) + col, fx.Int32(send_nbytes) + ) + buffer_store( + payload, output_buf, off, + cache_modifier=store_cm, offset_is_bytes=True, + ) + # Only the lane that owns a 32-value block writes its scale byte. + scale_off = (valid & is_leader).select( + row_base + fx.Int32(N_OUT) + n_block_idx * fx.Int32(BN // 32) + + lane // fx.Int32(group_lanes), + fx.Int32(send_nbytes), + ) + buffer_store( + e8m0.to(fx.Int8), output_buf, scale_off, + cache_modifier=store_cm, offset_is_bytes=True, + ) + else: + off = valid.select( + row_base + n_block_idx * fx.Int32(BN * 2) + col * fx.Int32(2), + fx.Int32(send_nbytes), + ) + if const_expr(store_cm): + buffer_store(pk, output_buf, off, cache_modifier=store_cm, offset_is_bytes=True) + else: + buf_copy_store( + output_buf, + off // fx.Int32(scatter_vec * 2), + pk, + fx.BFloat16, + unit_elems=scatter_vec, + cache_modifier=0, + ) - ``remote`` needs this put to be the LAST op on the channel the data went out - on. It is: everything for ``peer`` uses channel ``peer % queues``, and a - queue is a FIFO its engine drains in order, so the trailing ATOMIC fires - only once every chunk ahead of it has landed. - The flag carries ``cnt`` either way. Under ``remote`` it no longer says - anything about the data, so it needs no release -- the signal does that. - """ - if signal_mode not in ("quiet", "remote"): - raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") - remote = signal_mode == "remote" - # Under `remote` this put carries the signal: 4 bytes the peer never reads. - if const_expr(remote): - sdma.put( - peer, - win, - fx.Int64(layout.off_sig + rank * 4), - win, - fx.Int64(layout.off_sig + peer * 4), - fx.Int64(4), - queue, - coop=cco.CoopScope.THREAD, - remote_signal=True, +# fmt: off +def compile_moe_gemm2_with_sdma(*, model_dim: int, inter_dim: int, experts: int, rank: int, + send_nbytes: int, BM: int = 32, BN: int = 256, BK: int = 256, use_nt: bool = True, + HIDDEN_MAX: int = 8192, INTER_MAX: int = 8192, a_dtype: str = "fp8", SBM: int | None = None, + persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, + g2_ascale_pf=None, g2_spart=None, persist_strided: bool = False, g2_bf16_lds: bool = False, + skew_cu: int = 0, scatter_vec: int = 8, comm=None, quant: str = "none"): +# fmt: on + """Compile GEMM2 with the rank-sorted local-output epilogue (optionally fused with SDMA).""" + arch = str(get_rocm_arch() or "") + if not arch.startswith("gfx95"): + raise RuntimeError(f"MegaMoE v2 stage2 requires CDNA4 (gfx95x), got {arch or 'unknown'}") + assert model_dim % BN == 0 and HIDDEN_MAX % BN == 0 + assert INTER_MAX % BK == 0, f"INTER_MAX must be a multiple of {BK}" + if BM not in (16, 32, 64, 128): + raise AssertionError(f"BM must be in {{16,32,64,128}}, got {BM}") + SBM = BM if SBM is None else int(SBM) + if SBM % BM != 0: + raise AssertionError(f"SBM ({SBM}) must be a multiple of BM ({BM})") + if scatter_vec not in (8, 16): + raise ValueError("Stage2 scatter_vec must be 8 or 16") + if BN % scatter_vec: + raise ValueError("Stage2 scatter_vec must divide BN") + if a_dtype not in ("fp4", "fp8"): + raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") + if persist and cu_num <= 0: + raise AssertionError(f"persist=True requires cu_num>0, got {cu_num}") + if skew_cu and (not persist or not 0 < skew_cu <= cu_num): + raise AssertionError(f"skew_cu={skew_cu} requires persist=True and 0= 1 and chunk_tiles >= 1") + # The rows are stored sc0|sc1 -- past L1 (GLC) and past L2 (SCC) -- so they + # are in memory when the store retires and the ``s_waitcnt vmcnt(0)`` in + # front of the publish is the release the copy engine needs. Leaving them in + # L2 and issuing a system-scope release instead costs a whole-L2 write-back + # per unit (``buffer_wbl2`` on gfx950): measured 2.4x on the GEMM2 kernel. + # See mori/ops/gemm_ar/_compat.py for the encoding, and why sc1 is 16 not 4. + STORE_CM = 1 | 16 if fuse else 0 + N_OUT = model_dim + g2_bhoist, g2_ascale_pf, g2_spart, g2_group_num, g2_m01, _g2_bf16_lds = _resolve_g2_knobs( + g2_bhoist, g2_ascale_pf, g2_spart, False, False + ) + is_f8 = a_dtype == "fp8" + aStages = kStages + 1 + KH_TILE_A = BK // (1 if is_f8 else 2) + compute_lds_bytes = _stage2_lds_bytes(BM, BN, BK, a_dtype, aStages, g2_bf16_lds) + lds_outrow_off = compute_lds_bytes + lds_weight_off = lds_outrow_off + BM * 4 + lds_bytes = lds_weight_off + BM * 4 + if lds_bytes > 160 * 1024: + raise ValueError(f"Stage2 LDS use {lds_bytes} exceeds 160 KiB") + _expert_offset = rank * experts + + @fx.struct + class SharedStorage: + buf: fx.Array[Int8, lds_bytes, 16] + + kernel_name = ( + f"megamoe_stage2_with_sdma_t{BM}x{BN}x{BK}" + f"_sbm{SBM}_{a_dtype}_nt{int(use_nt)}" + f"_p{int(persist)}cu{cu_num}s{int(persist_strided)}_pad{int(has_pad)}" + f"_sk{skew_cu}" + f"_bh{int(g2_bhoist)}apf{int(g2_ascale_pf)}sp{g2_group_num}x{g2_m01}" + f"_bf16lds{int(g2_bf16_lds)}_sv{scatter_vec}_{quant}" + + ( + f"_fuse{comm.chunk_tiles}q{comm.queues}w{comm.comm_wgs}" + f"sm{comm.signal_mode[0]}r{comm.rank}" + if fuse + else "" ) - if const_expr(not remote): - sdma.quiet_queue(peer, queue) - # The bytes moved by the copy engine, not by this CU, so nothing of ours - # has to be written back -- this only keeps the flag store from being - # hoisted above the drain. - rocdl.s_waitcnt(vmcnt=0) - comm_ops.store_i64_global_system( - fx.Int64(w.lsa_ptr(peer, layout.off_flags)) + fx.Int64(rank * 8), - (fx.Int64(cnt) << fx.Int64(32)) | fx.Int64(phase), ) + # fmt: off + @flyc.kernel(name=kernel_name, known_block_size=[256, 1, 1]) + def kernel_with_sdma(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, + arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, + arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, + arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, + arg_cursor: fx.Int64, arg_idx: fx.Int64, i32_max_m_blocks: fx.Int32, i32_inter: fx.Int32, + i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32): + # fmt: on + tx_i32 = fx.thread_idx.x + bid = fx.block_idx.x + lane = tx_i32 % fx.Int32(64) + wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_base_i32 = fx.Int32(fx.ptrtoint(lds.buf.ptr)) + + num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) + k_bytes = fx.Int32(i32_inter) // fx.Int32(1 if is_f8 else 2) + trb_buf = ptr_buf_tensor(arg_trb, fx.Int32) + sweights_buf = ptr_buf_tensor(arg_sweights, fx.Float32) + outrow_buf = ptr_buf_tensor(arg_out_row, fx.Int32) + + def issue_all_a_loads(m_row0): + for slot in range_constexpr(kStages): + issue_a_load_lds_dt(arg_aq, lds_base_i32, slot, slot, m_row0, wave, lane, + is_f8, KH_TILE_A, k_bytes, BM=BM) + + def run_unit(unit_bx, m_block_idx): + # Map each Stage2 BM sub-tile to its Stage1 SBM metadata row. + m_row = m_block_idx * fx.Int32(BM) + sort_block_idx = m_row // fx.Int32(SBM) + row_in_sort_block = m_row - sort_block_idx * fx.Int32(SBM) + srcmap_row_base = trb_buf[sort_block_idx] + row_in_sort_block + if tx_i32 < fx.Int32(BM): + sorted_pos = srcmap_row_base + tx_i32 + fx.ptr_store( + outrow_buf[m_row + tx_i32], + lds_typed_ptr( + fx.Int32(lds_outrow_off) + tx_i32 * fx.Int32(4), + T.i32, + align=4, + ), + ) + fx.ptr_store( + sweights_buf[sorted_pos], + lds_typed_ptr( + fx.Int32(lds_weight_off) + tx_i32 * fx.Int32(4), + T.f32, + align=4, + ), + ) + # fmt: off + accm_vecs, m_row, n_block_idx, _n_out_rt = gemm2_compute_v2(lds_base_i32, arg_ascale, arg_bq, + arg_bscale, arg_eids, arg_aq, i32_max_m_blocks, unit_bx, lane, wave, i32_inter, i32_hidden, + i32_kpad, i32_npad, BM=BM, BN=BN, BK=BK, use_nt=use_nt, INTER_MAX=INTER_MAX, aStages=aStages, + a_dtype=a_dtype, has_pad=has_pad, SBM=SBM, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, + expert_offset=_expert_offset) + send_base = rocdl.readfirstlane(T.i64, fx.Int64(arg_send).ir_value()) + local_sort_epilog(lds_base_i32, accm_vecs, n_block_idx, wave, lane, send_base, + N_OUT=N_OUT, BM=BM, BN=BN, send_nbytes=send_nbytes, lds_outrow_off=lds_outrow_off, + lds_weight_off=lds_weight_off, g2_bf16_lds=g2_bf16_lds, scatter_vec=scatter_vec, + store_cm=STORE_CM, quant=quant) + # fmt: on + + def run_one(unit_bx, m_block_idx): + issue_all_a_loads(m_block_idx * fx.Int32(BM)) + rocdl.sched_barrier(0) + run_unit(unit_bx, m_block_idx) + if const_expr(fuse): + # The rows went past L2 on the way out, so retiring the stores + # is all the release the copy engine needs; the barrier makes it + # the whole workgroup's stores, not just this wave's. + rocdl.s_waitcnt(vmcnt=0) + fx.barrier() + if tx_i32 == fx.Int32(0): + chunk = (m_block_idx * fx.Int32(BM)) // fx.Int32(SBM * comm.chunk_tiles) + comm_ops.atomic_add_agent( + fx.Int64(arg_cdone) + fx.Int64(chunk) * fx.Int64(CDONE_STRIDE * 4), + fx.Int32(1), + ) + + cumsum0 = global_typed_ptr(arg_cumsum, T.i32)[0] + total_m_blocks = (cumsum0 + fx.Int32(BM - 1)) // fx.Int32(BM) + + def gemm_role(bx_i32): + + if const_expr(not persist and g2_spart <= 0): + bound = total_m_blocks * fx.Int32(num_n_blocks) + if fx.Int32(bx_i32) < bound: + run_one(bx_i32, bx_i32 // num_n_blocks) + elif const_expr(not persist): + bound = total_m_blocks * fx.Int32(num_n_blocks) + if fx.Int32(bx_i32) < bound: + m_block_idx, n_block_idx = _spart_output_tile_index( + bx_i32, total_m_blocks, num_n_blocks, g2_group_num, g2_m01 + ) + unit_bx = m_block_idx * fx.Int32(num_n_blocks) + n_block_idx + run_one(unit_bx, m_block_idx) + elif const_expr(skew_cu > 0): + total_stage1_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) + max_expert_tiles = global_typed_ptr(arg_max_expert_tiles, T.i32)[0] + skewed = max_expert_tiles * fx.Int32(4) > total_stage1_tiles + skew_n_block = bx_i32 // fx.Int32(cu_num) + skew_m_slot = bx_i32 - skew_n_block * fx.Int32(cu_num) + normal_m_slot = bx_i32 // fx.Int32(num_n_blocks) + normal_n_block = bx_i32 - normal_m_slot * fx.Int32(num_n_blocks) + n_block = skewed.select(skew_n_block, normal_n_block) + m_slot = skewed.select(skew_m_slot, normal_m_slot) + active_cu = skewed.select(fx.Int32(skew_cu), fx.Int32(cu_num)) + strided_diff = total_m_blocks - m_slot + strided_rem = (strided_diff > fx.Int32(0)).select(strided_diff, fx.Int32(0)) + strided_iters = (strided_rem + active_cu - fx.Int32(1)) // active_cu + tiles_per_slot = ( + total_m_blocks + active_cu - fx.Int32(1) + ) // active_cu + m_tile0 = m_slot * tiles_per_slot + contiguous_diff = total_m_blocks - m_tile0 + contiguous_rem = (contiguous_diff > fx.Int32(0)).select( + contiguous_diff, fx.Int32(0) + ) + contiguous_iters = (contiguous_rem < tiles_per_slot).select( + contiguous_rem, tiles_per_slot + ) + n_iters = skewed.select(strided_iters, contiguous_iters) + active = m_slot < active_cu + for _it in range(fx.Int32(0), n_iters, fx.Int32(1)): + strided_m = m_slot + fx.Int32(_it) * active_cu + contiguous_m = m_tile0 + fx.Int32(_it) + m_block = skewed.select(strided_m, contiguous_m) + if active: + unit_bx = m_block * fx.Int32(num_n_blocks) + n_block + fx.barrier() + if m_block < total_m_blocks: + run_one(unit_bx, m_block) + else: + m_slot = bx_i32 // fx.Int32(num_n_blocks) + n_block = bx_i32 - m_slot * fx.Int32(num_n_blocks) + if const_expr(persist_strided): + diff = total_m_blocks - m_slot + rem = (diff > fx.Int32(0)).select(diff, fx.Int32(0)) + n_iters = (rem + fx.Int32(cu_num - 1)) // fx.Int32(cu_num) + else: + tiles_per_slot = ( + total_m_blocks + fx.Int32(cu_num - 1) + ) // fx.Int32(cu_num) + m_tile0 = m_slot * tiles_per_slot + diff = total_m_blocks - m_tile0 + rem = (diff > fx.Int32(0)).select(diff, fx.Int32(0)) + n_iters = (rem < tiles_per_slot).select(rem, tiles_per_slot) + for _it in range(fx.Int32(0), n_iters, fx.Int32(1)): + if const_expr(persist_strided): + m_block = m_slot + fx.Int32(_it) * fx.Int32(cu_num) + else: + m_block = m_tile0 + fx.Int32(_it) + unit_bx = m_block * fx.Int32(num_n_blocks) + n_block + fx.barrier() # separate prev-iter epilog LDS reads from this iter's A-load into the LDS union + if fx.Int32(m_block) < total_m_blocks: + run_one(unit_bx, m_block) + + def comm_role(cw): + import mori.cco.device.flydsl as cco + + lay, rank, npes, CT = comm.layout, comm.rank, comm.layout.npes, comm.chunk_tiles + n_waves = NC * 4 + n_mine = (npes + n_waves - 1) // n_waves + mb_per_tile = SBM // BM + # This wave owns peers owner, owner + n_waves, ...; a wave past the last + # peer has nothing to do. + owner = cw * fx.Int32(4) + wave + is_lane0 = lane == fx.Int32(0) + my_peers = [owner + fx.Int32(j * n_waves) for j in range_constexpr(n_mine)] + + def sends_to(peer): + return (peer < fx.Int32(npes)) & (peer != fx.Int32(rank)) + + def queue_of(peer): + return peer % fx.Int32(comm.queues) + + if owner < fx.Int32(npes): + w = cco.Window(arg_win) + sdma = cco.DevComm(arg_dev_comm).sdma() + rsrc_counter = create_buffer_resource_from_addr( + fx.Int64(w.lsa_ptr(rank, lay.off_counter)) + ) + # Read-only here: the combine kernel advances the counter after the call, + # so every comm wave sees the same phase whenever it starts. + phase = fx.Int32(buffer_load(rsrc_counter, 0, vec_width=1, dtype=T.i32)) + fx.Int32(1) + cursor_buf = ptr_buf_tensor(arg_cursor, fx.Int32) + tile_base_buf = ptr_buf_tensor(arg_tile_base, fx.Int32) + idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) + + def push_inv(peer): + """Fill ``peer``'s inv for the rows I send it: slot -> where the row lands + there (its R rows follow its S rows, my run starts at rank * cap). All 64 + lanes work, PUSH_U rows each per trip, so that many idx loads are in + flight before their stores.""" + cnt = cursor_buf[peer] + inv_peer = fx.Int64(w.lsa_ptr(peer, lay.off_inv)) + land0 = fx.Int32(npes * lay.cap + rank * lay.cap) + for base in range(fx.Int32(0), cnt, fx.Int32(64 * PUSH_U)): + ok, js, slots = [], [], [] + for u in range_constexpr(PUSH_U): + j = base + fx.Int32(u * 64) + lane + in_run = j < cnt + ok.append(in_run) + js.append(j) + slots.append( + idx_buf[peer * fx.Int32(lay.cap) + in_run.select(j, fx.Int32(0))] + ) + for u in range_constexpr(PUSH_U): + if ok[u]: + comm_ops.store_i32_global_system_monotonic( + inv_peer + fx.Int64(slots[u]) * fx.Int64(4), land0 + js[u] + ) + + # Do not overwrite a peer's landing buffer or inv until it has read the last call's. + for peer in my_peers: + if sends_to(peer): + if is_lane0: + sdma_wait_consumed(w, peer, phase, layout=lay, rank=rank) + # The inv pushes run while the GEMM computes its first chunk. They have to be + # performed before any put or signal of mine, hence the wait on vmcnt. + for peer in my_peers: + if sends_to(peer): + push_inv(peer) + rocdl.s_waitcnt(vmcnt=0) + + total_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) + n_chunks = (total_tiles + fx.Int32(CT - 1)) // fx.Int32(CT) + for c in range(fx.Int32(0), n_chunks, fx.Int32(1)): + ci = fx.Int32(c) + m_lo = ci * fx.Int32(CT * mb_per_tile) + m_hi = m_lo + fx.Int32(CT * mb_per_tile) + m_hi = (m_hi < total_m_blocks).select(m_hi, total_m_blocks) + # Every (m_block, n_block) unit of this chunk has stored. + if is_lane0: + comm_ops.spin_until_ge_i32_agent( + fx.Int64(arg_cdone) + fx.Int64(ci) * fx.Int64(CDONE_STRIDE * 4), + (m_hi - m_lo) * num_n_blocks, + ) + comm_ops.fence_agent_acquire() + t0 = ci * fx.Int32(CT) + t1 = t0 + fx.Int32(CT) + t1 = (t1 < total_tiles).select(t1, total_tiles) + for peer in my_peers: + if sends_to(peer): + lo = tile_base_buf[t0 * fx.Int32(npes) + peer] + hi = tile_base_buf[t1 * fx.Int32(npes) + peer] + if hi > lo: + if is_lane0: + sdma_put_run(sdma, arg_win, peer, queue_of(peer), + layout=lay, rank=rank, lo=lo, hi=hi, cco=cco) + for peer in my_peers: + if sends_to(peer): + if is_lane0: + sdma_finish_peer(sdma, w, arg_win, peer, queue_of(peer), phase, + layout=lay, rank=rank, cco=cco, + signal_mode=comm.signal_mode) + + if const_expr(fuse): + if bid < fx.Int32(NC): + comm_role(bid) + else: + gemm_role(bid - fx.Int32(NC)) + else: + gemm_role(bid) + + # fmt: off + @flyc.jit + def launch(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, + arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, + arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, + arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, + arg_cursor: fx.Int64, arg_idx: fx.Int64, i32_max_m_blocks: fx.Int32, i32_grid_blocks: fx.Int32, + i32_inter: fx.Int32, i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32, + stream: fx.Stream): + # fmt: on + num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) + grid_x = i32_grid_blocks * num_n_blocks + fx.Int32(NC) + kernel_with_sdma( + arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, + arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, arg_cdone, + arg_tile_base, arg_cursor, arg_idx, i32_max_m_blocks, i32_inter, i32_hidden, i32_kpad, i32_npad, + ).launch(grid=(grid_x, 1, 1), block=(256, 1, 1), stream=stream) + + return launch + + +_LAUNCH_CACHE = {} + + +# fmt: off +def run_moe_gemm2_with_sdma(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, + arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, + arg_cdone, arg_tile_base, arg_cursor, arg_idx, row_capacity, i32_inter, i32_hidden, stream, *, + model_dim, inter_dim, experts, rank, send_nbytes, BM, SBM, HIDDEN_MAX, INTER_MAX, cu_num, + BN=256, BK=256, use_nt=True, g2_bhoist=True, g2_ascale_pf=True, g2_spart=402, persist=False, + persist_cu=0, persist_strided=False, g2_bf16_lds=False, skew_cu=0, scatter_vec=8, comm=None, + quant="none"): + # fmt: on + """Compile or reuse one configuration and launch it.""" + launch_cu_num = min(cu_num, persist_cu) if persist and persist_cu > 0 else cu_num + compile_kw = dict( + model_dim=model_dim, inter_dim=inter_dim, experts=experts, rank=rank, + send_nbytes=send_nbytes, BM=BM, BN=BN, BK=BK, use_nt=use_nt, HIDDEN_MAX=HIDDEN_MAX, + INTER_MAX=INTER_MAX, SBM=SBM, persist=persist, cu_num=launch_cu_num, + g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, + persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, skew_cu=skew_cu, + scatter_vec=scatter_vec, comm=comm, quant=quant, + ) + key = tuple(sorted(compile_kw.items())) + launch = _LAUNCH_CACHE.get(key) + if launch is None: + launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_with_sdma(**compile_kw) + max_m_blocks = (row_capacity + BM - 1) // BM + grid_blocks = launch_cu_num if persist else max_m_blocks + _run_compiled( + launch, arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, + arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, + arg_cdone, arg_tile_base, arg_cursor, arg_idx, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), fx.Int32(i32_inter), fx.Int32(i32_hidden), + fx.Int32(0), fx.Int32(0), stream, + ) + + +# --- Kernel 3: combine -------------------------------------------------------- def build_combine(*, layout, rank, hidden, topk, quant="none", signal_mode="remote"): """Kernel 3: finish the transfer and reduce. @@ -605,9 +1244,9 @@ def combine(win: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, # A. the peers' data has landed if tid < fx.Int32(npes): if tid != fx.Int32(rank): - # The flag carries cnt (and, under "quiet", data-readiness too). + # The flag is the call number (and, under "quiet", data-readiness too). comm_ops.spin_until_ge_i32_system( - flags_base + fx.Int64(tid) * fx.Int64(8), phase, acquire=True + flags_base + fx.Int64(tid) * fx.Int64(4), phase, acquire=True ) if const_expr(signal_mode == "remote"): # The sender did not drain its queue; its copy engine bumps @@ -710,6 +1349,8 @@ def launch(win: fx.Int64, arg_inv: fx.Int64, arg_out: fx.Int64, return launch +# --- The host op -------------------------------------------------------------- + def _torch_view(ptr, nbytes, dtype): import torch @@ -728,11 +1369,12 @@ def __init__(self): ) -class SortedCombine: - """GEMM2 with a rank-sorted local output, shipped by SDMA from inside the GEMM2 kernel. +class MegaMoESdma: + """MegaMoE stage2 + combine with the transfer done by SDMA from inside the GEMM2 kernel. - Reuses a ``MegaMoEV2``'s stage1/dispatch state, so it can run right after - ``moe(...)`` and be compared with that call's own combine output. + Reuses a ``MegaMoEV2``'s stage1/dispatch state (and reads its private attributes), + so it can run right after ``moe(...)`` and be compared with that call's own + combine output. ``chunk_tiles`` is how many sort tiles make one shipped chunk; ``comm_wgs`` the number of comm workgroups (4 waves each, one wave per peer by default). @@ -758,8 +1400,6 @@ def __init__( Communicator, ) - from moe_gemm2_with_sdma import CommSpec - cfg = moe.comb_cfg self.moe = moe self.rank, self.npes = int(cfg.rank), int(cfg.world_size) @@ -832,11 +1472,9 @@ def __init__( self.cursor = torch.zeros(self.npes, **i32) # One completion counter per chunk, each on its own cache line; the plan # kernel zeroes the prefix a launch can touch. - from moe_gemm2_with_sdma import CDONE_STRIDE - self.cdone_stride = CDONE_STRIDE self.cdone = torch.zeros((max_tiles + 1) * CDONE_STRIDE, **i32) - # Never-reset counters for the two grid barriers and the combine's last-block count. + # Never-reset counters: the plan kernel's barrier, and the combine's last-block count. self.plan_bar = torch.zeros( torch.cuda.get_device_properties(dev).multi_processor_count * BAR_STRIDE, **i32, @@ -908,7 +1546,7 @@ def set_chunk_tiles(self, chunk_tiles): self.comm_spec = replace(self.comm_spec, chunk_tiles=chunk_tiles) - def _cdone_used(self, config): + def _cdone_used(self): """Prefix of cdone a launch at this config can touch; the rest stays zero.""" spec = self.comm_spec ct = 1 if spec is None else spec.chunk_tiles @@ -937,13 +1575,12 @@ def _launch_plan(self, config, s_fx): fx.Int64(self.mem.ptr + lay.off_done), fx.Int64(self.mem.ptr + lay.off_counter), fx.Int32(self.nblocks), - fx.Int32(self._cdone_used(config)), + fx.Int32(self._cdone_used()), s_fx, ) def _launch_gemm2(self, config, s_fx): import flydsl.expr as fx - from moe_gemm2_with_sdma import run_moe_gemm2_with_sdma moe, op, lay = self.moe, self.moe._s1_op, self.layout ws, stage2 = moe._s1_dispatch_workspace, config.stage2 diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py b/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py deleted file mode 100644 index 1bc955a5c1..0000000000 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/moe_gemm2_with_sdma.py +++ /dev/null @@ -1,707 +0,0 @@ -# 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. -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors -# ruff: noqa: B023, I001 -"""GEMM2 whose epilogue writes weighted bf16 rows into a local, rank-sorted send buffer. - -Same GEMM2 mainloop and input layout as ``mega_moe_stage2``; only the output side -differs. Instead of scattering each row into the origin rank's ``comb_inp`` it -writes the row to ``send[out_row[row]]``, where ``out_row`` was assigned once per -row beforehand (``sorted_combine.build_rank_sort_plan``) so every N -tile of a row agrees on it. A negative ``out_row`` drops the row. - -With ``comm`` the kernel also carries ``comm.comm_wgs`` extra leading workgroups that -ship finished chunks of the send buffer to the peers with cco SDMA while the GEMM -workgroups are still computing. A GEMM unit signals ``cdone[chunk]`` after its stores. -Each comm WAVE owns its own peers (peer ``p`` belongs to wave ``p % n_comm_waves``): it -spins on the chunk counter itself, then puts that chunk's rows for its peers on its own -SDMA queue, so peers neither wait for each other's spin nor share a wave's issue slot. - -Two things about the publish are load-bearing and were each worth more than a -factor on this kernel. The epilogue stores the rows write-through (sc0|sc1) so -that retiring them is the release a copy engine needs; releasing at system scope -instead is a whole-L2 write-back per unit. And the per-chunk counters sit a -cache line apart, because every unit of a chunk bumps one. -""" - -from dataclasses import dataclass - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl.expr import const_expr, range_constexpr, rocdl -from flydsl.expr.typing import Int8, T -from flydsl.runtime.device import get_rocm_arch - -from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops -from aiter.ops.flydsl.kernels.buffer_ops import ( - buffer_load, - buffer_store, - create_buffer_resource_from_addr, -) -from aiter.ops.flydsl.kernels.mxfp4_gemm_common import ( - global_typed_ptr, - lds_typed_ptr, - lds_vec_load, -) -from aiter.ops.flydsl.kernels.tensor_shim import ( - _run_compiled, - buf_copy_store, - ptr_buf_tensor, -) - -from aiter.ops.flydsl.kernels.mega_moe.gemm2 import ( - _resolve_g2_knobs, - _spart_output_tile_index, - gemm2_compute_v2, - issue_a_load_lds_dt, - kStages, -) -from sorted_combine import ( - ArenaLayout, - _store_i32_relaxed_system, - sdma_finish_peer, - sdma_put_run, - sdma_wait_consumed, -) -from aiter.ops.flydsl.kernels.mega_moe.mega_moe_stage2 import ( - _BUFFER_OFFSET_ABI_BYTES, - _fp8_scale_for_leader, - _stage2_lds_bytes, -) -from aiter.ops.flydsl.kernels.mxfp4_gemm_common import _fabs_f32 as fabs_f32 - - -#: i32 slots between consecutive chunk-completion counters. Every unit of a -#: chunk bumps its counter, so packing them adjacent puts all of a launch's -#: atomics on one or two cache lines: 128 B apart instead is worth ~8% of the -#: GEMM2 kernel at bs=4096. Spreading a chunk's count over several counters on -#: top of this buys nothing -- the atomic itself is not the cost. -CDONE_STRIDE = 32 - -#: Rows a comm wave pushes into a peer's inv per lane per loop trip, so that many -#: independent idx loads are in flight before their stores. -PUSH_U = 8 - - -@dataclass(frozen=True) -class CommSpec: - """Compile-time description of the fused SDMA comm workgroups. - - ``comm_wgs`` workgroups of 4 waves each; one wave per peer needs - ``comm_wgs * 4 >= npes``. ``queues`` is the SDMA channels reserved per - (src, dst) pair; peer ``p`` uses channel ``p % queues`` for everything it - sends, so a chunk is never split across channels -- splitting measured - slower, and on this topology KFD recommends one engine per peer link - anyway, which makes extra channels pure overhead. - """ - - layout: ArenaLayout - rank: int - queues: int - chunk_tiles: int - comm_wgs: int = 1 - #: How a sender tells a peer its rows have landed; see - #: ``sorted_combine.sdma_finish_peer``. "remote" needs the signal to ride - #: the channel the data went out on, which ``queues`` already guarantees. - signal_mode: str = "remote" - - -# fmt: off -def local_sort_epilog(lds_acc_base, accm, n_block_idx, wave, lane, send_base, *, N_OUT, BM, BN, - send_nbytes, lds_outrow_off, lds_weight_off, g2_bf16_lds=False, scatter_vec=8, store_cm=0, - quant="none"): -# fmt: on - """CShuffle one GEMM2 tile into weighted BF16 rows and store them at their send-buffer index.""" - kMChunks = BM // 16 - numAccN = (BN // 4) // 16 - wave_n = BN // 4 - lane_div_16 = lane // 16 - lane_mod_16 = lane % 16 - lds_base_fptr = lds_typed_ptr(lds_acc_base, T.f32) - lds_base_bf16 = ( - lds_typed_ptr(lds_acc_base, T.bf16, align=2) - if const_expr(g2_bf16_lds) - else None - ) - quant_fp8 = quant == "fp8_blockwise_1x32" - if quant not in ("none", "fp8_blockwise_1x32"): - raise ValueError(f"unsupported quant={quant!r}") - if quant_fp8 and g2_bf16_lds: - raise ValueError("fp8_blockwise_1x32 needs the f32 CShuffle (g2_bf16_lds=False)") - # Row layout matches aiter's P2P scatter: N_OUT payload bytes followed by the - # N_OUT/32 e8m0 block scales, so one contiguous SDMA put ships both. - row_nbytes = N_OUT + N_OUT // 32 if quant_fp8 else N_OUT * 2 - - fx.barrier() - - for i in range_constexpr(kMChunks): - row_base = fx.Int32(i * 16) + lane_div_16 * 4 - if const_expr(g2_bf16_lds): - w_row = [ - fx.ptr_load( - lds_typed_ptr( - fx.Int32(lds_weight_off) + (row_base + v) * fx.Int32(4), - T.f32, - align=4, - ) - ) - for v in range_constexpr(4) - ] - for J in range_constexpr(numAccN): - col = wave * fx.Int32(wave_n) + J * 16 + lane_mod_16 - vec = fx.Vector(accm[i][J]) - for v in range_constexpr(4): - idx = (row_base + v) * BN + col - if const_expr(g2_bf16_lds): - lds_base_bf16[idx] = fx.BFloat16(fx.Float32(vec[v]) * fx.Float32(w_row[v])) - else: - lds_base_fptr[idx] = fx.Float32(vec[v]) - - fx.barrier() - - # ``buf_copy_store``'s cache_modifier is a BufferCopy enum (0=cached, 2=nt), - # so sc0|sc1 cannot be reached through it -- flydsl's own type stub says so, - # and so does mori/ops/gemm_ar/_compat.py. ``buffer_store`` passes the value - # through as the instruction's aux field, so the write-through path uses it. - if const_expr(store_cm or quant_fp8): - output_buf = create_buffer_resource_from_addr( - send_base, num_records_bytes=send_nbytes - ) - else: - output_buf = ptr_buf_tensor( - send_base, - fx.BFloat16, - unit_elems=scatter_vec, - num_records_bytes=send_nbytes, - ) - # Inactive lanes read safe LDS and issue a bounded OOB store. - active = lane < fx.Int32(BN // scatter_vec) - col = active.select(lane * fx.Int32(scatter_vec), fx.Int32(0)) - - for row_iter in range_constexpr(BM // 4): - row = wave + fx.Int32(row_iter * 4) - row_byte_off = row * fx.Int32(4) - out_row = fx.ptr_load( - lds_typed_ptr(fx.Int32(lds_outrow_off) + row_byte_off, T.i32, align=4) - ) - out_row = rocdl.readfirstlane(T.i32, out_row.ir_value()) - idx0 = row * fx.Int32(BN) + col - if const_expr(g2_bf16_lds): - pk = fx.Vector( - lds_vec_load( - lds_acc_base, - idx0 * fx.Int32(2), - fx.Vector.make_type(scatter_vec, fx.BFloat16), - fx.BFloat16, - align=16, - ) - ) - else: - weight = fx.ptr_load( - lds_typed_ptr(fx.Int32(lds_weight_off) + row_byte_off, T.f32, align=4) - ) - weight = rocdl.readfirstlane(T.f32, weight.ir_value()) - v8 = fx.Vector( - lds_vec_load( - lds_acc_base, - idx0 * fx.Int32(4), - fx.Vector.make_type(scatter_vec, fx.Float32), - fx.Float32, - align=16, - ) - ) - weighted = fx.Vector.from_elements( - [v8[i] * weight for i in range_constexpr(scatter_vec)], - fx.Float32, - ) - if const_expr(not quant_fp8): - pk = weighted.to(fx.BFloat16) - row_base = out_row * fx.Int32(row_nbytes) - valid = (out_row >= fx.Int32(0)) & active - if const_expr(quant_fp8): - # One e8m0 per 32 contiguous values, i.e. per 32//scatter_vec lanes. - vals = [fx.Float32(weighted[i]) for i in range_constexpr(scatter_vec)] - local_max = fabs_f32(vals[0]) - for q in range_constexpr(1, scatter_vec): - local_max = local_max.maximumf(fabs_f32(vals[q])) - max_bits = local_max.bitcast(fx.Int32) - for xor_lane in (1, 2): - if xor_lane < 32 // scatter_vec: - remote = rocdl.ds_bpermute( - T.i32, (lane ^ fx.Int32(xor_lane)) * fx.Int32(4), max_bits - ) - local_max = local_max.maximumf(fx.Int32(remote).bitcast(fx.Float32)) - max_bits = local_max.bitcast(fx.Int32) - group_lanes = 32 // scatter_vec - leader_lane = lane & fx.Int32(~(group_lanes - 1)) - is_leader = (lane & fx.Int32(group_lanes - 1)) == fx.Int32(0) - e8m0 = fx.Int32( - rocdl.ds_bpermute( - T.i32, - leader_lane * fx.Int32(4), - _fp8_scale_for_leader(is_leader, local_max), - ) - ) - block_scale = (e8m0 << fx.Int32(23)).bitcast(fx.Float32) - pk_ty = T.vec(2, T.i16) - words = [] - for word in range_constexpr(scatter_vec // 4): - packed = fx.Vector.filled(2, 0, fx.Int16).ir_value() - for pair in range_constexpr(2): - e = word * 4 + pair * 2 - packed = rocdl.cvt_scalef32_pk_fp8_f32( - pk_ty, packed, vals[e].ir_value(), vals[e + 1].ir_value(), - block_scale.ir_value(), pair, - ) - words.append(fx.Vector(packed).bitcast(fx.Int32)[0]) - payload = fx.Vector.from_elements(words, fx.Int32) - off = valid.select( - row_base + n_block_idx * fx.Int32(BN) + col, fx.Int32(send_nbytes) - ) - buffer_store( - payload, output_buf, off, - cache_modifier=store_cm, offset_is_bytes=True, - ) - # Only the lane that owns a 32-value block writes its scale byte. - scale_off = (valid & is_leader).select( - row_base + fx.Int32(N_OUT) + n_block_idx * fx.Int32(BN // 32) - + lane // fx.Int32(group_lanes), - fx.Int32(send_nbytes), - ) - buffer_store( - e8m0.to(fx.Int8), output_buf, scale_off, - cache_modifier=store_cm, offset_is_bytes=True, - ) - else: - off = valid.select( - row_base + n_block_idx * fx.Int32(BN * 2) + col * fx.Int32(2), - fx.Int32(send_nbytes), - ) - if const_expr(store_cm): - buffer_store(pk, output_buf, off, cache_modifier=store_cm, offset_is_bytes=True) - else: - buf_copy_store( - output_buf, - off // fx.Int32(scatter_vec * 2), - pk, - fx.BFloat16, - unit_elems=scatter_vec, - cache_modifier=0, - ) - - -# fmt: off -def compile_moe_gemm2_with_sdma(*, model_dim: int, inter_dim: int, experts: int, rank: int, - send_nbytes: int, BM: int = 32, BN: int = 256, BK: int = 256, use_nt: bool = True, - HIDDEN_MAX: int = 8192, INTER_MAX: int = 8192, a_dtype: str = "fp8", SBM: int | None = None, - persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, - g2_ascale_pf=None, g2_spart=None, persist_strided: bool = False, g2_bf16_lds: bool = False, - skew_cu: int = 0, scatter_vec: int = 8, comm=None, quant: str = "none"): -# fmt: on - """Compile GEMM2 with the rank-sorted local-output epilogue (optionally fused with SDMA).""" - arch = str(get_rocm_arch() or "") - if not arch.startswith("gfx95"): - raise RuntimeError(f"MegaMoE v2 stage2 requires CDNA4 (gfx95x), got {arch or 'unknown'}") - assert model_dim % BN == 0 and HIDDEN_MAX % BN == 0 - assert INTER_MAX % BK == 0, f"INTER_MAX must be a multiple of {BK}" - if BM not in (16, 32, 64, 128): - raise AssertionError(f"BM must be in {{16,32,64,128}}, got {BM}") - SBM = BM if SBM is None else int(SBM) - if SBM % BM != 0: - raise AssertionError(f"SBM ({SBM}) must be a multiple of BM ({BM})") - if scatter_vec not in (8, 16): - raise ValueError("Stage2 scatter_vec must be 8 or 16") - if BN % scatter_vec: - raise ValueError("Stage2 scatter_vec must divide BN") - if a_dtype not in ("fp4", "fp8"): - raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") - if persist and cu_num <= 0: - raise AssertionError(f"persist=True requires cu_num>0, got {cu_num}") - if skew_cu and (not persist or not 0 < skew_cu <= cu_num): - raise AssertionError(f"skew_cu={skew_cu} requires persist=True and 0= 1 and chunk_tiles >= 1") - # The rows are stored sc0|sc1 -- past L1 (GLC) and past L2 (SCC) -- so they - # are in memory when the store retires and the ``s_waitcnt vmcnt(0)`` in - # front of the publish is the release the copy engine needs. Leaving them in - # L2 and issuing a system-scope release instead costs a whole-L2 write-back - # per unit (``buffer_wbl2`` on gfx950): measured 2.4x on the GEMM2 kernel. - # See mori/ops/gemm_ar/_compat.py for the encoding, and why sc1 is 16 not 4. - STORE_CM = 1 | 16 if fuse else 0 - N_OUT = model_dim - g2_bhoist, g2_ascale_pf, g2_spart, g2_group_num, g2_m01, _g2_bf16_lds = _resolve_g2_knobs( - g2_bhoist, g2_ascale_pf, g2_spart, False, False - ) - is_f8 = a_dtype == "fp8" - aStages = kStages + 1 - KH_TILE_A = BK // (1 if is_f8 else 2) - compute_lds_bytes = _stage2_lds_bytes(BM, BN, BK, a_dtype, aStages, g2_bf16_lds) - lds_outrow_off = compute_lds_bytes - lds_weight_off = lds_outrow_off + BM * 4 - lds_bytes = lds_weight_off + BM * 4 - if lds_bytes > 160 * 1024: - raise ValueError(f"Stage2 LDS use {lds_bytes} exceeds 160 KiB") - _expert_offset = rank * experts - - @fx.struct - class SharedStorage: - buf: fx.Array[Int8, lds_bytes, 16] - - kernel_name = ( - f"megamoe_stage2_with_sdma_t{BM}x{BN}x{BK}" - f"_sbm{SBM}_{a_dtype}_nt{int(use_nt)}" - f"_p{int(persist)}cu{cu_num}s{int(persist_strided)}_pad{int(has_pad)}" - f"_sk{skew_cu}" - f"_bh{int(g2_bhoist)}apf{int(g2_ascale_pf)}sp{g2_group_num}x{g2_m01}" - f"_bf16lds{int(g2_bf16_lds)}_sv{scatter_vec}_{quant}" - + ( - f"_fuse{comm.chunk_tiles}q{comm.queues}w{comm.comm_wgs}" - f"sm{comm.signal_mode[0]}r{comm.rank}" - if fuse - else "" - ) - ) - - # fmt: off - @flyc.kernel(name=kernel_name, known_block_size=[256, 1, 1]) - def kernel_with_sdma(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, - arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, - arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, - arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, - arg_cursor: fx.Int64, arg_idx: fx.Int64, i32_max_m_blocks: fx.Int32, i32_inter: fx.Int32, - i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32): - # fmt: on - tx_i32 = fx.thread_idx.x - bid = fx.block_idx.x - lane = tx_i32 % fx.Int32(64) - wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) - - lds = fx.SharedAllocator().allocate(SharedStorage).peek() - lds_base_i32 = fx.Int32(fx.ptrtoint(lds.buf.ptr)) - - num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) - k_bytes = fx.Int32(i32_inter) // fx.Int32(1 if is_f8 else 2) - trb_buf = ptr_buf_tensor(arg_trb, fx.Int32) - sweights_buf = ptr_buf_tensor(arg_sweights, fx.Float32) - outrow_buf = ptr_buf_tensor(arg_out_row, fx.Int32) - - def issue_all_a_loads(m_row0): - for slot in range_constexpr(kStages): - issue_a_load_lds_dt(arg_aq, lds_base_i32, slot, slot, m_row0, wave, lane, - is_f8, KH_TILE_A, k_bytes, BM=BM) - - def run_unit(unit_bx, m_block_idx): - # Map each Stage2 BM sub-tile to its Stage1 SBM metadata row. - m_row = m_block_idx * fx.Int32(BM) - sort_block_idx = m_row // fx.Int32(SBM) - row_in_sort_block = m_row - sort_block_idx * fx.Int32(SBM) - srcmap_row_base = trb_buf[sort_block_idx] + row_in_sort_block - if tx_i32 < fx.Int32(BM): - sorted_pos = srcmap_row_base + tx_i32 - fx.ptr_store( - outrow_buf[m_row + tx_i32], - lds_typed_ptr( - fx.Int32(lds_outrow_off) + tx_i32 * fx.Int32(4), - T.i32, - align=4, - ), - ) - fx.ptr_store( - sweights_buf[sorted_pos], - lds_typed_ptr( - fx.Int32(lds_weight_off) + tx_i32 * fx.Int32(4), - T.f32, - align=4, - ), - ) - # fmt: off - accm_vecs, m_row, n_block_idx, _n_out_rt = gemm2_compute_v2(lds_base_i32, arg_ascale, arg_bq, - arg_bscale, arg_eids, arg_aq, i32_max_m_blocks, unit_bx, lane, wave, i32_inter, i32_hidden, - i32_kpad, i32_npad, BM=BM, BN=BN, BK=BK, use_nt=use_nt, INTER_MAX=INTER_MAX, aStages=aStages, - a_dtype=a_dtype, has_pad=has_pad, SBM=SBM, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, - expert_offset=_expert_offset) - send_base = rocdl.readfirstlane(T.i64, fx.Int64(arg_send).ir_value()) - local_sort_epilog(lds_base_i32, accm_vecs, n_block_idx, wave, lane, send_base, - N_OUT=N_OUT, BM=BM, BN=BN, send_nbytes=send_nbytes, lds_outrow_off=lds_outrow_off, - lds_weight_off=lds_weight_off, g2_bf16_lds=g2_bf16_lds, scatter_vec=scatter_vec, - store_cm=STORE_CM, quant=quant) - # fmt: on - - def run_one(unit_bx, m_block_idx): - issue_all_a_loads(m_block_idx * fx.Int32(BM)) - rocdl.sched_barrier(0) - run_unit(unit_bx, m_block_idx) - if const_expr(fuse): - # The rows went past L2 on the way out, so retiring the stores - # is all the release the copy engine needs; the barrier makes it - # the whole workgroup's stores, not just this wave's. - rocdl.s_waitcnt(vmcnt=0) - fx.barrier() - if tx_i32 == fx.Int32(0): - chunk = (m_block_idx * fx.Int32(BM)) // fx.Int32(SBM * comm.chunk_tiles) - comm_ops.atomic_add_agent( - fx.Int64(arg_cdone) + fx.Int64(chunk) * fx.Int64(CDONE_STRIDE * 4), - fx.Int32(1), - ) - - cumsum0 = global_typed_ptr(arg_cumsum, T.i32)[0] - total_m_blocks = (cumsum0 + fx.Int32(BM - 1)) // fx.Int32(BM) - - def gemm_role(bx_i32): - - if const_expr(not persist and g2_spart <= 0): - bound = total_m_blocks * fx.Int32(num_n_blocks) - if fx.Int32(bx_i32) < bound: - run_one(bx_i32, bx_i32 // num_n_blocks) - elif const_expr(not persist): - bound = total_m_blocks * fx.Int32(num_n_blocks) - if fx.Int32(bx_i32) < bound: - m_block_idx, n_block_idx = _spart_output_tile_index( - bx_i32, total_m_blocks, num_n_blocks, g2_group_num, g2_m01 - ) - unit_bx = m_block_idx * fx.Int32(num_n_blocks) + n_block_idx - run_one(unit_bx, m_block_idx) - elif const_expr(skew_cu > 0): - total_stage1_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) - max_expert_tiles = global_typed_ptr(arg_max_expert_tiles, T.i32)[0] - skewed = max_expert_tiles * fx.Int32(4) > total_stage1_tiles - skew_n_block = bx_i32 // fx.Int32(cu_num) - skew_m_slot = bx_i32 - skew_n_block * fx.Int32(cu_num) - normal_m_slot = bx_i32 // fx.Int32(num_n_blocks) - normal_n_block = bx_i32 - normal_m_slot * fx.Int32(num_n_blocks) - n_block = skewed.select(skew_n_block, normal_n_block) - m_slot = skewed.select(skew_m_slot, normal_m_slot) - active_cu = skewed.select(fx.Int32(skew_cu), fx.Int32(cu_num)) - strided_diff = total_m_blocks - m_slot - strided_rem = (strided_diff > fx.Int32(0)).select(strided_diff, fx.Int32(0)) - strided_iters = (strided_rem + active_cu - fx.Int32(1)) // active_cu - tiles_per_slot = ( - total_m_blocks + active_cu - fx.Int32(1) - ) // active_cu - m_tile0 = m_slot * tiles_per_slot - contiguous_diff = total_m_blocks - m_tile0 - contiguous_rem = (contiguous_diff > fx.Int32(0)).select( - contiguous_diff, fx.Int32(0) - ) - contiguous_iters = (contiguous_rem < tiles_per_slot).select( - contiguous_rem, tiles_per_slot - ) - n_iters = skewed.select(strided_iters, contiguous_iters) - active = m_slot < active_cu - for _it in range(fx.Int32(0), n_iters, fx.Int32(1)): - strided_m = m_slot + fx.Int32(_it) * active_cu - contiguous_m = m_tile0 + fx.Int32(_it) - m_block = skewed.select(strided_m, contiguous_m) - if active: - unit_bx = m_block * fx.Int32(num_n_blocks) + n_block - fx.barrier() - if m_block < total_m_blocks: - run_one(unit_bx, m_block) - else: - m_slot = bx_i32 // fx.Int32(num_n_blocks) - n_block = bx_i32 - m_slot * fx.Int32(num_n_blocks) - if const_expr(persist_strided): - diff = total_m_blocks - m_slot - rem = (diff > fx.Int32(0)).select(diff, fx.Int32(0)) - n_iters = (rem + fx.Int32(cu_num - 1)) // fx.Int32(cu_num) - else: - tiles_per_slot = ( - total_m_blocks + fx.Int32(cu_num - 1) - ) // fx.Int32(cu_num) - m_tile0 = m_slot * tiles_per_slot - diff = total_m_blocks - m_tile0 - rem = (diff > fx.Int32(0)).select(diff, fx.Int32(0)) - n_iters = (rem < tiles_per_slot).select(rem, tiles_per_slot) - for _it in range(fx.Int32(0), n_iters, fx.Int32(1)): - if const_expr(persist_strided): - m_block = m_slot + fx.Int32(_it) * fx.Int32(cu_num) - else: - m_block = m_tile0 + fx.Int32(_it) - unit_bx = m_block * fx.Int32(num_n_blocks) + n_block - fx.barrier() # separate prev-iter epilog LDS reads from this iter's A-load into the LDS union - if fx.Int32(m_block) < total_m_blocks: - run_one(unit_bx, m_block) - - def comm_role(cw): - import mori.cco.device.flydsl as cco - - lay, rank, npes, CT = comm.layout, comm.rank, comm.layout.npes, comm.chunk_tiles - n_waves = NC * 4 - n_mine = (npes + n_waves - 1) // n_waves - mb_per_tile = SBM // BM - # This wave owns peers owner, owner + n_waves, ...; a wave past the last - # peer has nothing to do. - owner = cw * fx.Int32(4) + wave - is_lane0 = lane == fx.Int32(0) - if owner < fx.Int32(npes): - w = cco.Window(arg_win) - sdma = cco.DevComm(arg_dev_comm).sdma() - rsrc_counter = create_buffer_resource_from_addr( - fx.Int64(w.lsa_ptr(rank, lay.off_counter)) - ) - # Read-only here: the combine kernel advances the counter after the call, - # so every comm wave sees the same phase whenever it starts. - phase = fx.Int32(buffer_load(rsrc_counter, 0, vec_width=1, dtype=T.i32)) + fx.Int32(1) - cursor_buf = ptr_buf_tensor(arg_cursor, fx.Int32) - tile_base_buf = ptr_buf_tensor(arg_tile_base, fx.Int32) - # Do not overwrite a peer's landing buffer until it has read the last call's. - for j in range_constexpr(n_mine): - peer = owner + fx.Int32(j * n_waves) - if peer < fx.Int32(npes): - if peer != fx.Int32(rank): - if is_lane0: - sdma_wait_consumed(w, peer, phase, layout=lay, rank=rank) - # Fill each peer's inv for the rows I send it: slot -> where the row lands - # there (its R rows follow its S rows, my run starts at rank * cap). This - # runs while the GEMM computes its first chunk, and has to be performed - # before any put or signal of mine, hence the wait on vmcnt. - idx_buf = ptr_buf_tensor(arg_idx, fx.Int32) - land0 = fx.Int32(npes * lay.cap + rank * lay.cap) - for j in range_constexpr(n_mine): - peer = owner + fx.Int32(j * n_waves) - if peer < fx.Int32(npes): - if peer != fx.Int32(rank): - cnt = cursor_buf[peer] - inv_peer = fx.Int64(w.lsa_ptr(peer, lay.off_inv)) - for base in range(fx.Int32(0), cnt, fx.Int32(64 * PUSH_U)): - ok, js, sl = [], [], [] - for u in range_constexpr(PUSH_U): - jj = base + fx.Int32(u * 64) + lane - o = jj < cnt - ok.append(o) - js.append(jj) - sl.append(idx_buf[peer * fx.Int32(lay.cap) + o.select(jj, fx.Int32(0))]) - for u in range_constexpr(PUSH_U): - if ok[u]: - _store_i32_relaxed_system( - inv_peer + fx.Int64(sl[u]) * fx.Int64(4), - land0 + js[u], - ) - rocdl.s_waitcnt(vmcnt=0) - total_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) - n_chunks = (total_tiles + fx.Int32(CT - 1)) // fx.Int32(CT) - for c in range(fx.Int32(0), n_chunks, fx.Int32(1)): - ci = fx.Int32(c) - m_lo = ci * fx.Int32(CT * mb_per_tile) - m_hi = m_lo + fx.Int32(CT * mb_per_tile) - m_hi = (m_hi < total_m_blocks).select(m_hi, total_m_blocks) - # Every (m_block, n_block) unit of this chunk has stored. - if is_lane0: - comm_ops.spin_until_ge_i32_agent( - fx.Int64(arg_cdone) + fx.Int64(ci) * fx.Int64(CDONE_STRIDE * 4), - (m_hi - m_lo) * num_n_blocks, - ) - comm_ops.fence_agent_acquire() - t0 = ci * fx.Int32(CT) - t1 = t0 + fx.Int32(CT) - t1 = (t1 < total_tiles).select(t1, total_tiles) - for j in range_constexpr(n_mine): - peer = owner + fx.Int32(j * n_waves) - if peer < fx.Int32(npes): - if peer != fx.Int32(rank): - lo = tile_base_buf[t0 * fx.Int32(npes) + peer] - hi = tile_base_buf[t1 * fx.Int32(npes) + peer] - if hi > lo: - if is_lane0: - sdma_put_run(sdma, arg_win, peer, - peer % fx.Int32(comm.queues), - layout=lay, rank=rank, lo=lo, hi=hi, cco=cco) - for j in range_constexpr(n_mine): - peer = owner + fx.Int32(j * n_waves) - if peer < fx.Int32(npes): - if peer != fx.Int32(rank): - if is_lane0: - sdma_finish_peer(sdma, w, arg_win, peer, - peer % fx.Int32(comm.queues), - cursor_buf[peer], phase, layout=lay, rank=rank, - cco=cco, signal_mode=comm.signal_mode) - - if const_expr(fuse): - if bid < fx.Int32(NC): - comm_role(bid) - else: - gemm_role(bid - fx.Int32(NC)) - else: - gemm_role(bid) - - # fmt: off - @flyc.jit - def launch(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, arg_bscale: fx.Int64, - arg_eids: fx.Int64, arg_cumsum: fx.Int64, arg_max_expert_tiles: fx.Int64, - arg_sweights: fx.Int64, arg_trb: fx.Int64, arg_send: fx.Int64, arg_out_row: fx.Int64, - arg_dev_comm: fx.Int64, arg_win: fx.Int64, arg_cdone: fx.Int64, arg_tile_base: fx.Int64, - arg_cursor: fx.Int64, arg_idx: fx.Int64, i32_max_m_blocks: fx.Int32, i32_grid_blocks: fx.Int32, - i32_inter: fx.Int32, i32_hidden: fx.Int32, i32_kpad: fx.Int32, i32_npad: fx.Int32, - stream: fx.Stream): - # fmt: on - num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) - grid_x = i32_grid_blocks * num_n_blocks + fx.Int32(NC) - kernel_with_sdma( - arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, - arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, arg_cdone, - arg_tile_base, arg_cursor, arg_idx, i32_max_m_blocks, i32_inter, i32_hidden, i32_kpad, i32_npad, - ).launch(grid=(grid_x, 1, 1), block=(256, 1, 1), stream=stream) - - return launch - - -_LAUNCH_CACHE = {} - - -# fmt: off -def run_moe_gemm2_with_sdma(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, - arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, - arg_cdone, arg_tile_base, arg_cursor, arg_idx, row_capacity, i32_inter, i32_hidden, stream, *, - model_dim, inter_dim, experts, rank, send_nbytes, BM, SBM, HIDDEN_MAX, INTER_MAX, cu_num, - BN=256, BK=256, use_nt=True, g2_bhoist=True, g2_ascale_pf=True, g2_spart=402, persist=False, - persist_cu=0, persist_strided=False, g2_bf16_lds=False, skew_cu=0, scatter_vec=8, comm=None, - quant="none"): - # fmt: on - """Compile or reuse one configuration and launch it.""" - launch_cu_num = min(cu_num, persist_cu) if persist and persist_cu > 0 else cu_num - compile_kw = dict( - model_dim=model_dim, inter_dim=inter_dim, experts=experts, rank=rank, - send_nbytes=send_nbytes, BM=BM, BN=BN, BK=BK, use_nt=use_nt, HIDDEN_MAX=HIDDEN_MAX, - INTER_MAX=INTER_MAX, SBM=SBM, persist=persist, cu_num=launch_cu_num, - g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, - persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, skew_cu=skew_cu, - scatter_vec=scatter_vec, comm=comm, quant=quant, - ) - key = tuple(sorted(compile_kw.items())) - launch = _LAUNCH_CACHE.get(key) - if launch is None: - launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_with_sdma(**compile_kw) - max_m_blocks = (row_capacity + BM - 1) // BM - grid_blocks = launch_cu_num if persist else max_m_blocks - _run_compiled( - launch, arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, - arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, - arg_cdone, arg_tile_base, arg_cursor, arg_idx, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), fx.Int32(i32_inter), fx.Int32(i32_hidden), - fx.Int32(0), fx.Int32(0), stream, - ) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py b/examples/cco/python/11_flydsl_mega_moe_sdma/sdma_bandwidth.py similarity index 98% rename from examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py rename to examples/cco/python/11_flydsl_mega_moe_sdma/sdma_bandwidth.py index e5082f06a1..febb442bc5 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/sdmabw.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/sdma_bandwidth.py @@ -28,7 +28,7 @@ sweep of transfer sizes and channel counts, so the fused kernel's exposed transfer time can be checked against what the hardware will actually do. - MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 sdmabw.py \\ + MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=8 sdma_bandwidth.py \\ --mb 0.36,0.72,1.44,2.89,5.78 --queues 1,2 """ @@ -54,7 +54,7 @@ def build_xfer(*, npes, rank, comm_wgs, queues, split, puts): import mori.cco.device.flydsl as cco @flyc.kernel( - name=f"sdmabw_p{npes}_w{comm_wgs}q{queues}s{split}n{puts}_r{rank}", + name=f"sdma_bandwidth_p{npes}_w{comm_wgs}q{queues}s{split}n{puts}_r{rank}", known_block_size=[WAVE * 4, 1, 1], ) def xfer( From b6e4404c2d24edb34e5873c4e0418cbde2b3dbaf Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 29 Oct 2026 02:54:35 +0000 Subject: [PATCH 8/9] perf(cco/example): drop the unused local signal, and signal on the last put Two packets the copy engines did not need: * Every data put trailed a local signal. A put that names no signal keeps cco's old default -- an ATOMIC into this rank's own signal pool -- and nothing reads that pool: completion is the peer's remote signal, or quiet_queue. Puts now pass local_signal=False. * Under signal_mode="remote" the completion signal rode a separate 4-byte put after the last chunk. It now rides the last chunk's own put; only a peer whose last chunk is empty (or a call with no chunks) still gets the separate put. stage2+combine on glm52, 8x MI355X, fp8, CT=16, median of interleaved rounds, bit-exact against aiter at every size: tokens/rank 2048 4096 8192 16384 no local signal -3.6% -0.6% -0.5% -0.55% both -1.1% -0.5% Co-Authored-By: Claude --- .../mega_moe_gemm2_with_sdma.py | 71 ++++++++++++++++--- 1 file changed, 63 insertions(+), 8 deletions(-) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py b/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py index a961a7f9da..960aade4d1 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py @@ -220,9 +220,18 @@ def sdma_wait_consumed(w, peer, phase, *, layout, rank): ) -def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): +def sdma_put_run( + sdma, win, peer, queue, *, layout, rank, lo, hi, cco, remote_signal=False +): """Ship rows ``[lo, hi)`` of my run for ``peer``. + No local signal: a put that names none keeps cco's old default, a trailing + ATOMIC into our own signal pool, and nothing here waits on it -- completion + is the peer's remote signal or ``quiet_queue``. Dropping it is one packet + less per put (-3.6% stage2+combine at 2048 tokens/rank, -0.5% at 8192). + ``remote_signal`` hangs the completion signal on this put instead (see + ``sdma_finish_peer``). + One put for the whole run. Puts are asynchronous -- each reserves a slot, writes the copy descriptor and rings the doorbell -- but a queue is a ring its engine drains in order, so splitting a run across channels only buys @@ -242,12 +251,25 @@ def sdma_put_run(sdma, win, peer, queue, *, layout, rank, lo, hi, cco): fx.Int64(hi - lo) * fx.Int64(rb), queue, coop=cco.CoopScope.THREAD, + local_signal=False, + remote_signal=remote_signal, ) @comm_ops.traced def sdma_finish_peer( - sdma, w, win, peer, queue, phase, *, layout, rank, cco, signal_mode="remote" + sdma, + w, + win, + peer, + queue, + phase, + *, + layout, + rank, + cco, + signal_mode="remote", + signalled=False, ): """Tell ``peer`` its rows have landed. @@ -270,12 +292,15 @@ def sdma_finish_peer( Either way the flag carries the call number. Under ``remote`` it says nothing about the data -- the signal does that. + + ``signalled``: the last data put for ``peer`` already carried the signal, so + the 4-byte put is skipped. Only a peer whose last chunk was empty needs it. """ if signal_mode not in ("quiet", "remote"): raise ValueError(f"signal_mode must be quiet or remote, got {signal_mode!r}") remote = signal_mode == "remote" # Under `remote` this put carries the signal: 4 bytes the peer never reads. - if const_expr(remote): + if const_expr(remote and not signalled): sdma.put( peer, win, @@ -1102,6 +1127,9 @@ def push_inv(peer): total_tiles = (cumsum0 + fx.Int32(SBM - 1)) // fx.Int32(SBM) n_chunks = (total_tiles + fx.Int32(CT - 1)) // fx.Int32(CT) + last = n_chunks - fx.Int32(1) + # Under "remote" the last chunk's put carries the completion signal. + sig_on_put = comm.signal_mode == "remote" for c in range(fx.Int32(0), n_chunks, fx.Int32(1)): ci = fx.Int32(c) m_lo = ci * fx.Int32(CT * mb_per_tile) @@ -1123,14 +1151,41 @@ def push_inv(peer): hi = tile_base_buf[t1 * fx.Int32(npes) + peer] if hi > lo: if is_lane0: - sdma_put_run(sdma, arg_win, peer, queue_of(peer), - layout=lay, rank=rank, lo=lo, hi=hi, cco=cco) + if const_expr(sig_on_put): + if ci == last: + sdma_put_run(sdma, arg_win, peer, queue_of(peer), + layout=lay, rank=rank, lo=lo, hi=hi, + cco=cco, remote_signal=True) + else: + sdma_put_run(sdma, arg_win, peer, queue_of(peer), + layout=lay, rank=rank, lo=lo, hi=hi, + cco=cco) + else: + sdma_put_run(sdma, arg_win, peer, queue_of(peer), + layout=lay, rank=rank, lo=lo, hi=hi, cco=cco) for peer in my_peers: if sends_to(peer): if is_lane0: - sdma_finish_peer(sdma, w, arg_win, peer, queue_of(peer), phase, - layout=lay, rank=rank, cco=cco, - signal_mode=comm.signal_mode) + if const_expr(sig_on_put): + # Did the last chunk's put carry the signal? Not if that + # chunk was empty for this peer, or there were no chunks. + t_last = last * fx.Int32(CT) + t_last = (t_last > fx.Int32(0)).select(t_last, fx.Int32(0)) + lo_l = tile_base_buf[t_last * fx.Int32(npes) + peer] + hi_l = tile_base_buf[total_tiles * fx.Int32(npes) + peer] + if (n_chunks > fx.Int32(0)) & (hi_l > lo_l): + sdma_finish_peer(sdma, w, arg_win, peer, queue_of(peer), + phase, layout=lay, rank=rank, cco=cco, + signal_mode=comm.signal_mode, + signalled=True) + else: + sdma_finish_peer(sdma, w, arg_win, peer, queue_of(peer), + phase, layout=lay, rank=rank, cco=cco, + signal_mode=comm.signal_mode) + else: + sdma_finish_peer(sdma, w, arg_win, peer, queue_of(peer), phase, + layout=lay, rank=rank, cco=cco, + signal_mode=comm.signal_mode) if const_expr(fuse): if bid < fx.Int32(NC): From f5771d11729106bb121f88c4e2c6ae67eef3bf4d Mon Sep 17 00:00:00 2001 From: kaiclian Date: Thu, 29 Oct 2026 14:03:23 +0000 Subject: [PATCH 9/9] perf(cco/example): unbreak the GEMM2 main loop schedule and give the fused path every CU Two independent costs the fused stage2 was paying, both on top of a GEMM that was already running at full speed. Measured EP8 on MI355X, fp8, uniform routing, row_slack 1.5, --no-aligned-pair, chunk_tiles 16; every number below is bit-exact against aiter. 1. The AMDGPU UnclusteredHighRPReschedule stage. Fusing the comm role into this kernel cost 8.5% of GEMM2 at 16384 tokens per rank even when the comm role never ran -- a diagnostic build that compiled it in behind a runtime-false guard paid the same 0.175 ms. It is not register pressure in the main loop and not instruction cache: the K loop holds the same opcode multiset either way (one s_nop and one s_waitcnt apart), the 13 SGPR spills all sit in the prologue filling s_nop slots, and duplicating the comm role four times to take .text from 23.2 KB to 45.8 KB moved the cost 2.3%. What changes is the order. The comm role's live ranges lift the function's register pressure over the threshold that turns the stage on, and it then breaks the main loop's clustering: the MFMAs move 23 instructions earlier and the mean load to s_waitcnt distance falls from 20.9 to 13.6, so the loop hides less latency on every trip. Compiling with the stage off restores the original schedule. 2. persist_cu on the fused path. A comm workgroup holds one of its CU's two workgroup slots for the whole kernel. ktrace at 16384 tokens per rank: those two CUs run 8 GEMM workgroups while the other 254 run 15 to 18, every workgroup keeps its normal duration (p50 25396 to 25092 ticks), and mean per-CU finish does not move -- only the last CU finishes 3.3% late. The displaced work has to be redistributed and, at the token bucket's tuned persist_cu of 192, one workgroup is ~13% of the kernel span, so the tail quantises badly. Handing the fused path every CU makes the workgroups shorter and more numerous, and the tail absorbs the displacement: the fused kernel's cost over plain GEMM2 drops a third (V4.1-Flash 0.200 to 0.133 ms, V4-Pro 0.274 to 0.184 ms). Dropping persist_cu instead, or moving the comm workgroups elsewhere in the grid, does nothing -- as do poll backoff, comm workgroup count, and XCD alignment. stage2+combine against aiter MegaMoEV2, tokens per rank: 4096 8192 16384 GLM-5.2 1.15 -> 1.27 1.26 -> 1.38 1.39 -> 1.52 V4.1-Flash 1.03 -> 1.15 1.15 -> 1.27 1.15 -> 1.33 V4-Pro 0.96 -> 1.10 1.07 -> 1.20 1.08 -> 1.21 Also adds the V4.1-Flash shape to the harness networks so those rows reproduce. Note aiter has no tuned profile for it at 48 experts per rank -- its v41_flash_ep2 profile wants 192 -- so all three networks here run aiter's generic stage2 config. Co-Authored-By: Claude Opus 5 (1M context) --- .../python/11_flydsl_mega_moe_sdma/main.py | 3 ++ .../mega_moe_gemm2_with_sdma.py | 39 ++++++++++++++++--- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py index 355f4f1287..2138361395 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/main.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/main.py @@ -76,6 +76,9 @@ "v4_pro": dict( model_dim=7168, inter_dim=3072, experts=384, topk=6, swiglu_limit=10.0 ), + "flash": dict( + model_dim=5120, inter_dim=2304, experts=384, topk=6, swiglu_limit=10.0 + ), } diff --git a/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py b/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py index 960aade4d1..66e4803326 100644 --- a/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py +++ b/examples/cco/python/11_flydsl_mega_moe_sdma/mega_moe_gemm2_with_sdma.py @@ -90,6 +90,7 @@ from flydsl.expr.rocdl import cvt_pk_f32_fp8 from flydsl.expr.typing import Int8, T from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.compiler.llvm_options import llvm_options from flydsl.runtime.device import get_rocm_arch from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops @@ -1229,6 +1230,16 @@ def run_moe_gemm2_with_sdma(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar quant="none"): # fmt: on """Compile or reuse one configuration and launch it.""" + if comm is not None and persist and 0 < persist_cu < cu_num: + # A comm workgroup holds one of its CU's two workgroup slots for the whole + # kernel: ktrace shows those CUs run 8 GEMM workgroups where the rest run + # 15-18. The displaced work has to go somewhere, and at the bucket's tuned + # persist_cu=192 one workgroup is ~13% of the kernel span, so the tail + # quantises badly -- the last CU finishes 3.3% late while every workgroup + # runs at its normal speed. More, shorter workgroups absorb it: the fused + # kernel's cost over plain GEMM2 drops a third at 16384 tokens per rank + # (V4.1-Flash 0.200 -> 0.133 ms, V4-Pro 0.274 -> 0.184 ms). + persist_cu = cu_num launch_cu_num = min(cu_num, persist_cu) if persist and persist_cu > 0 else cu_num compile_kw = dict( model_dim=model_dim, inter_dim=inter_dim, experts=experts, rank=rank, @@ -1240,16 +1251,32 @@ def run_moe_gemm2_with_sdma(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar ) key = tuple(sorted(compile_kw.items())) launch = _LAUNCH_CACHE.get(key) - if launch is None: - launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_with_sdma(**compile_kw) max_m_blocks = (row_capacity + BM - 1) // BM grid_blocks = launch_cu_num if persist else max_m_blocks - _run_compiled( - launch, arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, + run_args = ( + arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cumsum, arg_max_expert_tiles, arg_sweights, arg_trb, arg_send, arg_out_row, arg_dev_comm, arg_win, - arg_cdone, arg_tile_base, arg_cursor, arg_idx, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), fx.Int32(i32_inter), fx.Int32(i32_hidden), - fx.Int32(0), fx.Int32(0), stream, + arg_cdone, arg_tile_base, arg_cursor, arg_idx, fx.Int32(max_m_blocks), fx.Int32(grid_blocks), + fx.Int32(i32_inter), fx.Int32(i32_hidden), fx.Int32(0), fx.Int32(0), stream, ) + if launch is None: + # The comm role's live ranges lift this function's register pressure over the + # threshold that turns on the AMDGPU UnclusteredHighRPReschedule stage, and + # that stage then breaks the GEMM main loop's clustering: the MFMAs move 23 + # instructions earlier and the mean load -> s_waitcnt distance falls from + # 20.9 to 13.6, so the loop hides less latency. Same instruction multiset, + # worse order, and the loop pays it on every trip -- 8.5% of GEMM2 at 16384 + # tokens per rank, present even when the comm role never runs. The register + # allocation is fine either way: 13 SGPR spills, all in the prologue, all + # filling s_nop slots. + # + # FlyDSL emits code on the first launch, not in the call below, so the option + # has to stay set across that launch as well. + with llvm_options({"amdgpu-disable-unclustered-high-rp-reschedule": True}): + launch = _LAUNCH_CACHE[key] = compile_moe_gemm2_with_sdma(**compile_kw) + _run_compiled(launch, *run_args) + return + _run_compiled(launch, *run_args) # --- Kernel 3: combine --------------------------------------------------------