diff --git a/.github/workflows/ci_cco.yml b/.github/workflows/ci_cco.yml index 626c8a86d..e78f9d184 100644 --- a/.github/workflows/ci_cco.yml +++ b/.github/workflows/ci_cco.yml @@ -33,7 +33,7 @@ jobs: MORI_RDMA_DEVICES: ${{ matrix.rdma_devices }} MORI_RDMA_SL: ${{ matrix.rdma_sl }} MORI_RDMA_TC: ${{ matrix.rdma_tc }} - timeout-minutes: 30 + timeout-minutes: 60 steps: - name: Checkout uses: actions/checkout@v4 @@ -179,6 +179,34 @@ jobs: cd $GITHUB_WORKSPACE && \ timeout 300 pytest tests/python/cco/test_sdma_api.py -v" + - name: Run CCO GEMM+AR numerics (up to 8 ranks) + run: | + # The op ships in mori.ops.gemm_ar, so run pytest from /tmp: from a + # repo cwd `python/mori/` shadows the install and a packaging miss + # (a subpackage find_packages never picked up) would read as green. + # + # test_gemm_ar_op.py is the public wrapper's own tests. The other two + # go through the benchmark, which calls the low-level builders, so + # they said nothing about that layer -- and three correctness bugs + # lived there behind a green suite. + $CT exec -e MORI_ENABLE_SDMA=1 -e MORI_SOCKET_IFNAME=lo $CONTAINER bash -c "\ + cd /tmp && python -c 'import mori.ops.gemm_ar as g; print(g.GemmAllReduceOp)' && \ + timeout 1800 python -m pytest -v \ + $GITHUB_WORKSPACE/tests/python/cco/test_gemm_ar.py \ + $GITHUB_WORKSPACE/tests/python/cco/test_gemm_ar_op.py \ + $GITHUB_WORKSPACE/tests/python/cco/test_flydsl_ar.py" + + - name: Run CCO GEMM+AR benchmark (8 ranks, accuracy-gated) + run: | + # Accuracy is a gate (validated + rel_l2); absolute timings are only + # recorded -- this runner is shared, so a microsecond threshold would + # fail on a busy box and pass a real regression on an idle one. The + # one perf assertion is relative (fused beats split, measured back to + # back). -s so the numbers reach the log for a bisect to read. + $CT exec -e MORI_ENABLE_SDMA=1 -e MORI_SOCKET_IFNAME=lo $CONTAINER bash -c "\ + cd $GITHUB_WORKSPACE && \ + timeout 1800 pytest -v -s tests/python/cco/test_gemm_ar_bench.py" + - name: Validate CCO ROCm runtime import order run: | $CT exec $CONTAINER bash -c "\ diff --git a/benchmark/cco/flydsl/ar/_probe_lsa.py b/benchmark/cco/flydsl/ar/_probe_lsa.py new file mode 100644 index 000000000..790f18a5a --- /dev/null +++ b/benchmark/cco/flydsl/ar/_probe_lsa.py @@ -0,0 +1,203 @@ +#!/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. +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +"""Toolchain probe: example 05's LSA all-reduce, ported onto ``_compat``. + +Validates the pieces the real kernels depend on, on whichever flydsl is +installed: ``Window.lsa_ptr``, the uncached signal barrier, vectorized peer +reads, and the cco device-bitcode JIT. + + MORI_SOCKET_IFNAME=lo mpirun --allow-run-as-root -np 2 python _probe_lsa.py +""" + +import os +import sys + +from mpi4py import MPI + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf +from flydsl.expr import gpu as fgpu +from flydsl.expr import range_constexpr +from flydsl.expr.typing import Int32, Int64 + +import mori.cco.device.flydsl as cco +from mori.cco import CCODevCommRequirements, Communicator, GDA_CONNECTION_NONE + +from mori.ops.gemm_ar._compat import ( + CM_CACHED, + FLYDSL_VERSION, + buffer_load, + buffer_store, + create_buffer_resource_from_addr, + i32_type, + signal_load_u32, + signal_ptr, + signal_store_u32, +) + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "..") +) +sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "..", + "..", + "examples", + "cco", + "python", + ), +) +from cco_example_common import F32, fill, read, set_device, sync, zero # noqa: E402 + +WS = 2 +THREADS = 256 +NUM_ELEMS = 256 * 1024 +ELEMS_PER_PACK = 4 +NUM_PACKS = NUM_ELEMS // ELEMS_PER_PACK + +SIG_OFF = 0 +SIG_BYTES = 256 +IN_OFF = SIG_BYTES +OUT_OFF = IN_OFF + NUM_ELEMS * 4 +WIN_BYTES = OUT_OFF + NUM_ELEMS * 4 + + +@flyc.kernel(known_block_size=[THREADS, 1, 1]) +def probe_ar_kernel(dev_comm: Int64, win: Int64, flag: Int32): + tid = fx.thread_idx.x + dc = cco.DevComm(dev_comm) + w = cco.Window(win) + rank = dc.lsa_rank + + ins = [fx.Int64(w.lsa_ptr(p, IN_OFF)) for p in range(WS)] + self_sig = fx.Int64(w.lsa_ptr(rank, SIG_OFF)) + out = fx.Int64(w.lsa_ptr(rank, OUT_OFF)) + + # start-sync: publish my arrival to every peer, then wait for all of them. + if tid < WS: + peer_base = fx.Int64(w.lsa_ptr(tid, SIG_OFF)) + signal_store_u32(signal_ptr(peer_base + fx.Int64(rank) * fx.Int64(4)), flag) + + wait_addr = self_sig + fx.Int64(tid) * fx.Int64(4) + wait_rsrc = signal_ptr(wait_addr) + i32 = i32_type() + first = signal_load_u32(wait_rsrc) + first_v = first.ir_value() if hasattr(first, "ir_value") else first + loop = scf.WhileOp([i32], [first_v]) + cond = ir.Block.create_at_start(loop.before, [i32]) + body = ir.Block.create_at_start(loop.after, [i32]) + with ir.InsertionPoint(cond): + cur = fx.Int32(cond.arguments[0]) + should_wait = fx.Uint32(cur) < fx.Uint32(flag) + scf.ConditionOp(should_wait.ir_value(), [cond.arguments[0]]) + with ir.InsertionPoint(body): + nxt = signal_load_u32(wait_rsrc) + scf.YieldOp([nxt.ir_value() if hasattr(nxt, "ir_value") else nxt]) + fgpu.barrier() + + out_rsrc = create_buffer_resource_from_addr(out) + in_rsrc = [create_buffer_resource_from_addr(ins[p]) for p in range(WS)] + for pk in range(tid, NUM_PACKS, THREADS): + elem_off = pk * ELEMS_PER_PACK + acc = None + for p in range_constexpr(WS): + raw = fx.Vector( + buffer_load(in_rsrc[p], elem_off, vec_width=4, dtype=i32_type()) + ) + vf = raw.bitcast(fx.Float32) + acc = vf if acc is None else acc + vf + buffer_store( + acc.bitcast(fx.Int32), out_rsrc, elem_off, cache_modifier=CM_CACHED + ) + + +@flyc.jit +def run_probe(dev_comm: Int64, win: Int64, flag: Int32, stream=fx.Stream(None)): + probe_ar_kernel(dev_comm, win, flag).launch( + grid=(1, 1, 1), block=[THREADS, 1, 1], stream=stream + ) + + +def main() -> int: + mpi = MPI.COMM_WORLD + rank, nranks = mpi.Get_rank(), mpi.Get_size() + if nranks != WS: + if rank == 0: + print(f"probe is built for world_size={WS} (mpirun -n {WS})") + return 1 + set_device(rank) + uid = Communicator.get_unique_id() if rank == 0 else None + uid = mpi.bcast(uid, root=0) + + errors = 0 + with Communicator.init(nranks, rank, uid, per_rank_vmm=256 * 1024 * 1024) as comm: + mem = comm.alloc_mem(WIN_BYTES) + win = comm.register_window(mem.ptr, mem.size) + zero(win.local_ptr, WIN_BYTES) + fill( + win.local_ptr + IN_OFF, + [(rank + 1) * (i + 1) for i in range(NUM_ELEMS)], + F32, + ) + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + dc = comm.create_dev_comm(reqs) + + comm.barrier() + run_probe(dc.ptr, win.handle, 1) + sync() + comm.barrier() + + S = WS * (WS + 1) // 2 + host = read(win.local_ptr + OUT_OFF, NUM_ELEMS, F32) + for i in (0, 1, NUM_ELEMS // 2, NUM_ELEMS - 1): + exp = float(S * (i + 1)) + if abs(host[i] - exp) > 1e-3 * max(1.0, exp): + print(f"[rank {rank}] MISMATCH [{i}]: {host[i]} != {exp}", flush=True) + errors += 1 + print( + f"[rank {rank}] flydsl={FLYDSL_VERSION} " + f"{'PROBE OK' if errors == 0 else 'PROBE FAILED'} " + f"out[0,1,-1]={[host[0], host[1], host[NUM_ELEMS - 1]]}", + flush=True, + ) + + all_err = mpi.allreduce(errors, op=MPI.SUM) + if rank == 0: + print("SUCCESS" if all_err == 0 else "FAILED", flush=True) + return 0 if all_err == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py new file mode 100644 index 000000000..f3a1cc6f6 --- /dev/null +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -0,0 +1,411 @@ +#!/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. +"""Benchmark: FlyDSL cco all-reduce (LSA / SDMA) vs aiter's custom all-reduce. + +Reports one ``RESULT_JSON`` line per run, using the same critical-rank +decomposition as ``benchmark/cco/triton/gemm_a2a/bench_gemm_a2a.py`` so the two +benchmarks can be read side by side. + + MORI_SOCKET_IFNAME=lo MORI_ENABLE_SDMA=1 \\ + torchrun --standalone --nproc_per_node=8 bench_ar.py \\ + --backend lsa -m 4096 -n 7168 --warmup 5 --iters 51 + +Measured on 8x MI355X (gfx950), ROCm 7.2, ``[M, 7168]`` bf16, CUDA-graph replay, +best of 3 runs of 101 iterations each. aiter baseline is +``tensor_model_parallel_all_reduce`` with ``set_custom_all_reduce(True)``: + + M aiter LSA (/aiter) SDMA (/LSA) + 64 25.6us 30.7us 1.20 45.6us 1.49 + 512 53.6 66.5 1.24 73.6 1.11 + 2048 146.1 160.3 1.10 167.6 1.05 + 4096 271.1 285.9 1.05 294.5 1.03 + 8192 558.8 531.7 0.95 544.6 1.02 + +Read those with three caveats. First, all three have a payload-independent floor +-- LSA's is ~28us and SDMA's ~40us, neither moving between M=8 and M=64 -- so the +small-M ratios are two constants divided, not a throughput result. The 12us +between the two floors is the copy engine's ~6us dispatch cost, paid once per +transfer phase. Second, single runs vary by up to 1.7x (aiter's M=64 spanned +25.6-44.0us across three runs), which is why this uses best-of-3; do not compare +single runs. Third, SDMA's remaining deficit is almost entirely its ~40us floor: +past M=2048 it is within 5% of LSA and at M=8192 it beats aiter. It got there by +sizing the reduce kernel's grid independently -- see ``SDMA_REDUCE_BLOCK_CAP`` in +layout.py, worth 48us at M=4096 on its own. + +This box is affected by mori known-issue 3 (ROCm 7.2 routes uncached VMM +allocations to the coarse-grained pool): ``CCO_UNCACHED_WINDOW=0`` changes nothing, +which is the documented tell. aiter is on IPC handles and does not pay it. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys + +import flydsl.expr as fx +import torch +import torch.distributed as dist + +from mori.cco import ( + CCODevCommRequirements, + Communicator, + GDA_CONNECTION_NONE, + UniqueId, +) +from mori.tensor_utils import from_gpu_ptr + +from mori.ops.gemm_ar import ArConfig + +#: Accepted relL2 for the fp8 wire. +#: +#: e4m3 carries 3 mantissa bits, so one rounding at per-row scale costs about +#: **2.1e-2** on a standard-normal payload -- measured, on this benchmark, at +#: [4096, 7168] over 8 ranks. That is an order of magnitude more than the +#: bf16 path (which is exact) and about 5x the 3.7-4.0e-3 the fused wo_b +#: currently shows end to end in the model, so it is a real cost, not a +#: rounding detail. +#: +#: 4e-2 is a small multiple of that: it passes the quantisation the option asks +#: for and still catches a structurally wrong kernel, which lands near 0.9. +FP8_WIRE_TOL = 4e-2 + +VMM_SLACK = 256 * 1024 * 1024 + + +def _setup_distributed(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="cpu:gloo") + rank, world_size = dist.get_rank(), dist.get_world_size() + payload = [bytes(Communicator.get_unique_id()) if rank == 0 else None] + dist.broadcast_object_list(payload, src=0) + return local_rank, rank, world_size, UniqueId.from_bytes(payload[0]) + + +def make_input(rank: int, cfg: ArConfig) -> torch.Tensor: + """Deterministic, rank-dependent, and cheap to reproduce as a reference. + + Values stay small so the fp32 accumulation of ``world_size`` terms is exact + in bf16 -- the test asserts bit-exactness, so the data must not be what + introduces rounding. + """ + idx = torch.arange(cfg.num_elems, device="cuda", dtype=torch.float32) + vals = ((idx % 7.0) - 3.0) * 0.25 + float(rank) + return vals.view(cfg.m, cfg.n).to(torch.bfloat16) + + +def reference(world_size: int, cfg: ArConfig) -> torch.Tensor: + acc = torch.zeros(cfg.num_elems, device="cuda", dtype=torch.float32) + for r in range(world_size): + acc += make_input(r, cfg).view(-1).float() + return acc.view(cfg.m, cfg.n).to(torch.bfloat16) + + +def _median_us(fn, warmup: int, iters: int, *, graph: bool = True) -> float: + """Median wall time of one all-reduce. + + Graph replay is the default because the eager path pays flydsl's Python + launch wrapper on every call, which at decode sizes is larger than the + collective itself. Capture is safe here: the barrier flags are monotonic and + live in device memory, so a replay advances them exactly like a fresh launch + (the same property that lets aiter's custom AR be captured). + """ + if graph: + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(warmup): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + torch.cuda.synchronize() + return _median_us(g.replay, warmup, iters, graph=False) + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples = [] + for _ in range(iters): + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0) + return statistics.median(samples) + + +def run(args) -> int: + local_rank, rank, world_size, uid = _setup_distributed() + cfg = ArConfig( + world_size=world_size, + m=args.m, + n=args.n, + max_blocks=max(args.max_blocks, args.blocks or 0), + force_blocks=args.blocks, + force_reduce_blocks=args.reduce_blocks, + # LSA reduces straight out of the peers' input regions; SDMA has to be + # given somewhere for the copy engine to land each peer's slice. + recv_slots=world_size if args.backend == "sdma" else 0, + scatter_dtype=args.scatter_dtype or args.comm_dtype, + gather_dtype=args.gather_dtype or args.comm_dtype, + ) + cfg.validate() + + vmm = max(4 * cfg.window_bytes + VMM_SLACK, VMM_SLACK) + result = None + + with Communicator.init(world_size, rank, uid, per_rank_vmm=vmm) as comm: + mem = comm.alloc_mem(cfg.window_bytes) + win = comm.register_window(mem.ptr, mem.size) + + whole = from_gpu_ptr(mem.ptr, (cfg.window_bytes,), torch.uint8) + whole.zero_() + inp = from_gpu_ptr(mem.ptr + cfg.input_off, (cfg.m, cfg.n), torch.bfloat16) + out = from_gpu_ptr(mem.ptr + cfg.output_off, (cfg.m, cfg.n), torch.bfloat16) + inp.copy_(make_input(rank, cfg)) + torch.cuda.synchronize() + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + if args.backend == "sdma": + reqs.sdma_queue_count = args.sdma_queues + dc = comm.create_dev_comm(reqs) + + if args.backend == "lsa": + from mori.ops.gemm_ar import build_lsa_ar + + launch, stage = build_lsa_ar( + cfg, rank, force_stage=args.force_stage, gather=args.gather + ) + elif args.backend == "sdma": + from mori.ops.gemm_ar import build_sdma_phases + + parts = build_sdma_phases( + cfg, + rank, + queues=args.sdma_queues, + signal=args.sdma_signal, + reduce_blocks=args.reduce_blocks, + ) + + # parts["order"] rather than a literal list: the fp8 wire inserts + # a quantise and a dequantise around the gather, and hardcoding the + # three-phase order here would skip them and reduce into zeros. + def launch(dc_ptr, win_h, stream=None, _seq=None): + for k in parts["order"]: + parts[k](dc_ptr, win_h, stream=stream) + + stage = 2 + else: + raise ValueError(f"unknown backend {args.backend!r}") + + def once(): + # The stream must be passed explicitly. flydsl's default + # ``fx.Stream(None)`` resolves to the legacy default stream, which is + # not torch's capture stream, so under ``torch.cuda.graph`` the work + # lands outside the graph and replay measures an empty graph (a + # constant ~4.5us whatever the size). + launch( + dc.ptr, + win.handle, + stream=fx.Stream(torch.cuda.current_stream()), + ) + + # warm up + correctness on the same path the timing loop uses + comm.barrier() + once() + torch.cuda.synchronize() + comm.barrier() + + validated = True + rel_l2 = 0.0 + if not args.skip_validation: + ref = reference(world_size, cfg) + diff = (out.float() - ref.float()).norm().item() + denom = ref.float().norm().item() + rel_l2 = diff / denom if denom else diff + # The bf16 wire moves the bytes untouched and reduces in a fixed peer + # order, so it is *exactly* the reference and anything else is a bug. + # The fp8 wire rounds every element that crosses a link, so it gets a + # tolerance. + tol = FP8_WIRE_TOL if cfg.fp8_wire else 0.0 + validated = rel_l2 <= tol + if not validated: + # Which owner's slice is wrong localises the bug immediately: + # only-mine-right => the all-gather or its barrier; mine-wrong + # too => the reduce-scatter or the index math. + packs = cfg.packs_per_rank + per_pack = out.view(-1).numel() // cfg.num_packs + bad = [] + for r in range(world_size): + lo = r * packs * per_pack + hi = cfg.owner_pack_range(r)[1] * per_pack + d = out.view(-1)[lo:hi].float() - ref.view(-1)[lo:hi].float() + if d.norm().item() > tol * denom: + bad.append(r) + print( + f"[rank {rank}] VALIDATION FAILED relL2={rel_l2:.3e} " + f"bad_owner_slices={bad} mine_ok={rank not in bad}", + flush=True, + ) + + comm.barrier() + elapsed_us = _median_us(once, args.warmup, args.iters, graph=not args.eager) + comm.barrier() + + stats = torch.tensor([elapsed_us], dtype=torch.float64) + gathered = [torch.zeros_like(stats) for _ in range(world_size)] + dist.all_gather(gathered, stats) + per_rank = [float(t[0]) for t in gathered] + max_rank = max(range(world_size), key=lambda r: per_rank[r]) + max_us = per_rank[max_rank] + + ok = torch.tensor([1 if validated else 0], dtype=torch.int32) + dist.all_reduce(ok, op=dist.ReduceOp.MIN) + validated = bool(ok.item()) + + if rank == 0: + gbps = cfg.remote_bytes_per_rank * world_size / (max_us * 1e-6) / 1e9 + result = { + "backend": args.backend, + "stage": stage, + "gather": args.gather if args.backend == "lsa" else None, + "world_size": world_size, + "m": cfg.m, + "n": cfg.n, + "dtype": "bf16", + "payload_bytes": cfg.nbytes, + "slice_bytes": cfg.slice_bytes, + "blocks": cfg.blocks, + "threads": cfg.threads, + "max_rank_time_ms": max_us / 1000.0, + "max_rank_time_us": max_us, + "critical_rank": max_rank, + "per_rank_time_us": per_rank, + "remote_bytes_per_rank": cfg.remote_bytes_per_rank, + "aggregate_gbps": gbps, + "rel_l2": rel_l2, + "validated": validated, + "timing": "eager" if args.eager else "graph", + } + print("RESULT_JSON " + json.dumps(result, sort_keys=True), flush=True) + print( + f"[ar] backend={args.backend} stage={stage} m={cfg.m} n={cfg.n} " + f"blocks={cfg.blocks} max_rank_time={max_us:.2f}us " + f"aggregate={gbps:.1f}GB/s validated={validated}", + flush=True, + ) + if args.json_out: + with open(args.json_out, "w") as fh: + json.dump(result, fh, indent=2, sort_keys=True) + + dist.barrier() + return 0 if (result is None or result["validated"]) else 1 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser() + p.add_argument("--backend", choices=("lsa", "sdma"), default="lsa") + p.add_argument( + "--comm-dtype", + choices=("bf16", "fp8"), + default="bf16", + help="what travels on the wire, both legs. fp8 quantises to e4m3 with " + "fp32 scales, halving the bytes; the reduce still accumulates in fp32 " + "and the output is still bf16. sdma only.", + ) + # The legs are worth setting apart: gather is ~40% of a fused layer and its + # payload is final values, while scatter is already hidden behind the GEMM + # and its payload is partial sums whose fp8 error compounds across ranks. + p.add_argument("--scatter-dtype", choices=("bf16", "fp8"), default=None) + p.add_argument("--gather-dtype", choices=("bf16", "fp8"), default=None) + p.add_argument("-m", type=int, default=4096) + p.add_argument("-n", type=int, default=7168) + p.add_argument( + "--force-stage", + type=int, + choices=(1, 2), + default=None, + help="override the size-based 1-stage/2-stage choice (LSA only)", + ) + p.add_argument( + "--gather", + choices=("interleaved", "sequential"), + default="interleaved", + help="2-stage all-gather loop nesting (LSA only): 'interleaved' unrolls " + "the peer loop inside the index loop so all links stream at once; " + "'sequential' drains one peer at a time and is ~3x slower", + ) + p.add_argument( + "--blocks", + type=int, + default=None, + help="override the grid size (default: aiter's formula, capped at 80)", + ) + p.add_argument( + "--max-blocks", + type=int, + default=80, + help="signal-array rows; must be >= --blocks", + ) + p.add_argument( + "--reduce-blocks", + type=int, + default=0, + help="grid for the SDMA reduce (0 = layout default). Separate from " + "--blocks because that kernel is local-HBM bound, not xGMI bound", + ) + p.add_argument("--sdma-queues", type=int, default=8) + p.add_argument( + "--sdma-signal", + action="store_true", + help="attach put's trailing local ATOMIC. Nothing polls it and quiet is " + "documented as signal-independent, so this only exists to measure the " + "claim in FlyDSL's Sdma docstring that no-signal puts cannot be drained", + ) + p.add_argument("--warmup", type=int, default=5) + p.add_argument("--iters", type=int, default=51) + p.add_argument( + "--eager", + action="store_true", + help="time direct launches instead of graph replay (includes flydsl's " + "Python launch overhead)", + ) + p.add_argument("--skip-validation", action="store_true") + p.add_argument("--json-out") + return p + + +if __name__ == "__main__": + sys.exit(run(build_parser().parse_args())) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py new file mode 100644 index 000000000..806a6178f --- /dev/null +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -0,0 +1,738 @@ +#!/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. +"""GEMM + all-reduce at the DSV4-Pro ``wo_b`` shape: fused vs split. + + MORI_SOCKET_IFNAME=lo MORI_ENABLE_SDMA=1 \\ + torchrun --standalone --nproc_per_node=8 bench_gemm_ar.py \\ + --mode fused-sdma -m 4096 -n 7168 -k 1024 + +Modes, all measuring the same end state (``output`` holds the all-reduced GEMM +result on every rank): + + split-sdma gemm -> scatter -> reduce -> gather (4 kernels) + fused-sdma gemm+scatter -> drain -> reduce -> gather (4 kernels) + fused-lsa gemm writing straight into the peers -> barrier -> reduce + -> gather. gcnasm's "Direct LSA" shape: no staging, no copy + engine, and nothing to publish mid-kernel + split-lsa gemm -> LSA 2-stage all-reduce (2 kernels) + gemm-only the GEMM alone, to size the ceiling + +``split-*`` uses the *same* kernel as ``fused-sdma`` with the epilogue tail +compiled out, so the two differ in one thing only. ``split-lsa`` is included +because LSA is the faster collective, and a fused SDMA path has to beat +``gemm + LSA``, not just ``gemm + SDMA``, to be worth anything. + +Headline at the shape the model actually runs -- ``[16384, 7168]`` out, K=2048, +which is a ``--chunked-prefill-size 16384`` TP8 prefill chunk of DSV4-Pro +``wo_b`` (1792 tiles of 256x256 per rank). 8x MI355X, graph replay, median of +31, max over ranks, on an idle box, with the current defaults: + + --quant ptpc --quant blockscale + fused-sdma (chunks=8) 1114.5us 1146.2us + split-sdma 1261.7 1465.5 + split-lsa 1262.4 1469.9 + fused-lsa (n_stripe=2) 1365.6 1699.6 + gemm-only 228.9 388.5 + +``blockscale`` is the quantisation the model actually runs -- A 1x128 and B +128x128 with fp32 scales -- and is the column to read. ``ptpc`` is the 8-wave +kernel's native per-token/per-channel form, kept because the swap_ab bitwise +test and every earlier measurement are on it. + +For scale, the same layer in the running model costs 1491.3us (GEMM 348.0 + +NCCL 1143.3, medians over its 61 layers), and the collective on its own is +1043us for LSA / 1047 for SDMA against 1113 for NCCL. So at the model's own +quantisation ``fused-sdma`` is **21.8% under split and 23.1% under what the +model runs today**. + +The margin *grows* with blockscale (11.7% -> 21.8%) for the reason the whole +exercise was about: the GEMM goes 228.9 -> 388.5us while the 489us of link time +does not move, so there is more compute to hide the transfer behind. Our blockscale +GEMM is *not* competitive standalone: 370.5us against CK's 300.4 for its fastest +blockscale instance (64x256, Intrawave v1) measured the same way. A port of +that shape reached CK's instruction mix and not its speed; see the closing +section of ``python/mori/ops/gemm_ar/README.md`` for what that ruled out. + +Fusing over SDMA wins, and only because of ``--chunks``; see the table at its +definition in ``run()``. It was pinned to 1 while the GEMM's under-counted +``s_waitcnt`` made every chunked run intermittently wrong, and +with one chunk the fused path overlaps nothing at all. Fusing over LSA still +loses, though by much less since ``--n-stripe 2``: it spends 560us of its GEMM +pushing C over xGMI, where the copy engines move the same bytes in 499. + +Read ``python/mori/ops/gemm_ar/kernels_fused.py`` before trusting any fused +number from this benchmark: several of its options are intermittently wrong, and a single +passing run proves nothing. Every ``raw-wt`` fence mode is known-racy and kept +only to reproduce that. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys + +import flydsl.expr as fx +import torch +import torch.distributed as dist + +from mori.cco import ( + CCODevCommRequirements, + Communicator, + GDA_CONNECTION_NONE, + UniqueId, +) +from mori.tensor_utils import from_gpu_ptr + +from mori.ops.gemm_ar import ( + ArConfig, + build_lsa_ar, + build_sdma_phases, + compile_fused_gemm_scatter, + preshuffle_b, +) + +VMM_SLACK = 512 * 1024 * 1024 +MODES = ("gemm-only", "split-sdma", "fused-sdma", "fused-lsa", "split-lsa") + + +def _setup_distributed(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="cpu:gloo") + rank, world_size = dist.get_rank(), dist.get_world_size() + payload = [bytes(Communicator.get_unique_id()) if rank == 0 else None] + dist.broadcast_object_list(payload, src=0) + return local_rank, rank, world_size, UniqueId.from_bytes(payload[0]) + + +#: fp8 block-scale group size along K. Fixed by the model's quantiser +#: (``aiter_per1x128_quant``) and by the kernel, whose BLOCK_K is already 128. +SCALE_BK = 128 + + +def make_operands(rank: int, m: int, n: int, k: int, quant: str = "ptpc"): + """Deterministic per-rank fp8 operands and their scales. + + Values are kept small so the fp8 rounding is the only error source and the + all-reduce reference can be built on the host from the same recipe. + + ``quant="ptpc"`` is a8w8 per-token/per-channel: ``sa[M]``, ``sb[N]``, one + scale for a whole row of A and a whole column of B. + + ``quant="blockscale"`` is what the model actually runs -- A quantised 1x128 + and B 128x128, so both scales vary along K: ``sa[M, K/128]`` and + ``sb[N/128, K/128]``. The layouts match + ``aiter.gemm_a8w8_blockscale_bpreshuffle`` so the two can be compared + directly: **sa is column-major** (sglang's + ``materialize_bpreshuffle_fp8_scale`` does ``t().contiguous().t()``) and sb + is row-major. Established by sweeping the four combinations against that + kernel: only (sa column-major, sb row-major) lands at the fp8 floor + (1.66e-3); the other three give 0.17-0.32. + """ + g = torch.Generator(device="cuda").manual_seed(1234 + rank) + a = (torch.randn(m, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + b = (torch.randn(n, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + if quant == "ptpc": + sa = ( + torch.rand(m, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + ) + sb = ( + torch.rand(n, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + ) + return a, b, sa, sb + if quant != "blockscale": + raise ValueError(f"quant must be ptpc or blockscale, got {quant!r}") + kb = k // SCALE_BK + sa = ( + torch.rand(m, kb, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + ) + sb = ( + torch.rand(n // SCALE_BK, kb, generator=g, device="cuda", dtype=torch.float32) + * 0.01 + + 0.01 + ) + return a, b, sa.t().contiguous().t(), sb.contiguous() + + +def reference_partial(a, b, sa, sb, quant: str) -> torch.Tensor: + """One rank's fp32 GEMM reference, matching ``make_operands``' quantisation.""" + af, bf = a.float(), b.float() + if quant == "ptpc": + return (af @ bf.T) * sa[:, None] * sb[None, :] + out = torch.zeros(a.shape[0], b.shape[0], device=a.device, dtype=torch.float32) + for i in range(a.shape[1] // SCALE_BK): + ks = slice(i * SCALE_BK, (i + 1) * SCALE_BK) + out += ( + (af[:, ks] @ bf[:, ks].T) + * sa[:, i][:, None] + * sb[:, i].repeat_interleave(SCALE_BK)[None, :] + ) + return out + + +def _median_us(fn, warmup: int, iters: int, *, graph: bool = True) -> float: + if graph: + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(warmup): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + torch.cuda.synchronize() + return _median_us(g.replay, warmup, iters, graph=False) + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples = [] + for _ in range(iters): + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0) + return statistics.median(samples) + + +def run(args) -> int: + local_rank, rank, world_size, uid = _setup_distributed() + # blockscale keeps a second fp32 accumulator for the per-K-block promotion, + # which doubles the accumulator VGPRs; 256x256 needs 256 of them and the + # kernel already runs at ~254 with zero spill, so the tile has to halve. + if not args.block_m: + args.block_m = 128 if args.quant == "blockscale" else 256 + direct_lsa = args.mode == "fused-lsa" + fused = args.mode in ("fused-sdma", "fused-lsa") + # fused-lsa still uses the SDMA reduce/gather tail, so it wants queues too. + needs_sdma = args.mode in ("split-sdma", "fused-sdma", "fused-lsa") + # Pushes per destination. This is the whole overlap mechanism: with one + # chunk the tile counter only fires when a destination's entire slice is + # done, which under the rotated tile order is the end of the GEMM, so + # nothing overlaps at all -- measured, the drain then costs the same as the + # split scatter (493.8 vs 499.0us). With eight, a destination's first chunk + # leaves while the GEMM is still computing its later ones. + # + # Per-kernel medians at [16384, 7168] K=2048, 8x MI355X: + # + # chunks GEMM drain reduce gather total end-to-end + # 1 257.5 493.8 44.0 496.4 1291.7 1299.6 + # 2 266.5 379.4 44.2 496.4 1186.6 1194.9 + # 8 265.0 303.7 43.8 496.6 1109.1 1115.7 + # + # The drain falls 38% while the GEMM grows 7.5us for the extra counter + # atomics and the submit lock. 8 hides 190 of the 265us of GEMM it could + # possibly hide. More is worse: at 8 each PUT is 3.5 MiB, and 16 (via + # --block-m 128) drops that to 1.75 MiB, below the knee in the SDMA + # bandwidth curve -- 1130.8us, slower than 8. + # + # This was pinned to 1 because >1 produced wrong output 3 runs in 10. That + # was the aiter GEMM's under-counted s_waitcnt, not the chunk protocol: + # since that fix, chunks 1/2/4/8 are 40/40 correct and chunks=8 alone is + # 25/25, with no hang. Two hangs were seen at 8 ranks *before* the sweep + # settled and never reproduced in 100+ runs after; the submit lock is a + # plain test-and-set spin, so treat a hang as possible and use `timeout` + # when sweeping. + DEFAULT_CHUNKS = 8 + chunks = args.chunks if args.chunks else DEFAULT_CHUNKS + if fused: + # Fall back rather than fail: chunks has to divide the M-tiles per + # destination, which at small M can be fewer than 8. + m_tiles_per_peer = max( + 1, (args.m + args.block_m - 1) // args.block_m // world_size + ) + while chunks > 1 and m_tiles_per_peer % chunks: + chunks //= 2 + cfg = ArConfig( + world_size=world_size, + m=args.m, + n=args.n, + recv_slots=world_size if needs_sdma else 0, + counter_chunks=chunks, + gather_dtype=args.gather_dtype, + gather_transport=args.gather_transport, + gather_bands=args.gather_bands, + ) + cfg.validate() + + a, b, sa, sb = make_operands(rank, args.m, args.n, args.k, args.quant) + b_shuf = preshuffle_b(b) + + vmm = max(4 * cfg.window_bytes + VMM_SLACK, VMM_SLACK) + result = None + + with Communicator.init(world_size, rank, uid, per_rank_vmm=vmm) as comm: + mem = comm.alloc_mem(cfg.window_bytes) + win = comm.register_window(mem.ptr, mem.size) + from_gpu_ptr(mem.ptr, (cfg.window_bytes,), torch.uint8).zero_() + # C *is* the all-reduce's input region: the GEMM writes its partial + # straight into the symmetric window, which is the whole host-side change + # the SDMA fusion needs. + c = from_gpu_ptr(mem.ptr + cfg.input_off, (args.m, args.n), torch.bfloat16) + out = from_gpu_ptr(mem.ptr + cfg.output_off, (args.m, args.n), torch.bfloat16) + torch.cuda.synchronize() + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + if needs_sdma: + reqs.sdma_queue_count = args.sdma_queues + dc = comm.create_dev_comm(reqs) + + gemm = compile_fused_gemm_scatter( + cfg, + rank, + K=args.k, + BLOCK_M=args.block_m, + BLOCK_N=args.block_n, + b_preshuffled=True, + quant=args.quant, + sdma_queues=args.sdma_queues, + waves_per_eu=args.waves_per_eu, + xcd_swizzle=args.xcd_swizzle, + fuse=fused, + transport="lsa" if direct_lsa else "sdma", + swap_ab=args.swap_ab, + store_probe=args.store_probe, + permlane=args.permlane, + lane_transpose=args.lane_transpose, + hoist_scales=args.hoist_scales, + peer_uncached=args.peer_uncached, + direct_fence=args.direct_fence, + rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", + n_stripe=args.n_stripe or None, + fence=args.fence, + emit_put=not args.no_put, + atomic_order=args.atomic, + ) + a_i8 = a.contiguous().view(torch.int8).view(-1) + b_i8 = b_shuf.contiguous().view(torch.int8).view(-1) + c_flat = c.view(-1) + + # The kernel indexes both scale buffers linearly, so hand it the + # *physical* element order as 1-D contiguous tensors. sa is logically + # [M, K/128] but column-major (the layout aiter's blockscale GEMM + # consumes), so its physical order is sa.t(); passing the 2-D + # non-contiguous view through DLPack would give the kernel the wrong + # strides. + if args.quant == "blockscale": + sa_arg = sa.t().reshape(-1).contiguous() + sb_arg = sb.reshape(-1).contiguous() + else: + sa_arg, sb_arg = sa, sb + + def run_gemm(stream): + gemm( + a_i8, + b_i8, + c_flat, + sa_arg, + sb_arg, + args.m, + args.n, + dc.ptr, + win.handle, + stream=stream, + ) + + if needs_sdma: + parts = build_sdma_phases( + cfg, + rank, + queues=args.sdma_queues, + reduce_self_from_recv=direct_lsa, + recv_uncached=direct_lsa, + fuse_quantize=args.fuse_quantize, + fuse_reduce_push=args.fuse_reduce_push, + publish=args.publish, + ) + if args.mode == "split-lsa": + lsa_ar, _ = build_lsa_ar(cfg, rank) + + def once(): + stream = fx.Stream(torch.cuda.current_stream()) + run_gemm(stream) + if args.mode == "gemm-only": + return + if args.mode == "split-lsa": + lsa_ar(dc.ptr, win.handle, stream=stream) + return + # fused-sdma: the puts were issued from the epilogue, so only the + # drain runs. fused-lsa: the epilogue already wrote into the peers, + # so this is a pure cross-rank barrier (its quiet drains nothing). + # split: the full scatter kernel. + # parts["order"]/["fused_order"] rather than a literal list: the + # fp8 gather inserts a quantise and a dequantise, and hardcoding the + # three phases would skip them and all-reduce into zeros. + order = parts["fused_order" if fused else "order"] + parts[order[0]](dc.ptr, win.handle, stream=stream) + if args.stop_after == "scatter": + return + for phase in order[1:]: + parts[phase](dc.ptr, win.handle, stream=stream) + + comm.barrier() + once() + torch.cuda.synchronize() + comm.barrier() + + rel_l2 = float("nan") + validated = True + if args.mode == "gemm-only" and not args.skip_validation: + # The GEMM on its own, against this rank's partial. Without this + # control a GEMM bug reads as an all-reduce bug: every other mode + # validates the *sum*, so a corrupt partial and a corrupt collective + # are indistinguishable from the reported relL2. + ref = reference_partial(a, b, sa, sb, args.quant) + diff = (c.float() - ref).norm().item() + denom = ref.norm().item() + rel_l2 = diff / denom if denom else diff + validated = rel_l2 < 3e-3 + if not validated: + print( + f"[rank {rank}] GEMM VALIDATION FAILED relL2={rel_l2:.3e}", + flush=True, + ) + del ref + elif ( + args.mode != "gemm-only" + and args.stop_after == "all" + and not args.skip_validation + ): + # Reference: every rank's partial, summed in fp32 on the host side. + acc = torch.zeros(args.m, args.n, device="cuda", dtype=torch.float32) + for r in range(world_size): + ar, br, sar, sbr = make_operands(r, args.m, args.n, args.k, args.quant) + acc += reference_partial(ar, br, sar, sbr, args.quant) + diff = (out.float() - acc).norm().item() + denom = acc.norm().item() + rel_l2 = diff / denom if denom else diff + # fp8 inputs, so this is quantization error, not a collective error; + # the collective itself is checked bit-exactly in test_flydsl_ar.py. + # 3e-3, not 5e-3. The fp8 quantisation floor for these inputs is + # 2.35e-3 and the corruption this pipeline produces lands at + # 4-9e-3, so a looser gate reports a corrupt run as validated -- one + # did, at 3.97e-3, while the fused path was racing. + # + # --gather-dtype fp8 adds a second, much larger rounding on top: the + # all-gather carries e4m3, which costs ~2.1e-2 on its own. Its gate + # has to clear that, and is two-sided -- below 4e-2 says the kernel + # is right (a broken one lands near 0.9), above 5e-3 says the fp8 + # wire was actually taken and not silently skipped. + if cfg.fp8_gather: + validated = 5e-3 < rel_l2 < 4e-2 + else: + validated = rel_l2 < 3e-3 + if not validated: + print(f"[rank {rank}] VALIDATION FAILED relL2={rel_l2:.3e}", flush=True) + del acc + + comm.barrier() + elapsed_us = _median_us(once, args.warmup, args.iters, graph=not args.eager) + comm.barrier() + + stats = torch.tensor([elapsed_us], dtype=torch.float64) + gathered = [torch.zeros_like(stats) for _ in range(world_size)] + dist.all_gather(gathered, stats) + per_rank = [float(t[0]) for t in gathered] + max_rank = max(range(world_size), key=lambda r: per_rank[r]) + + ok = torch.tensor([1 if validated else 0], dtype=torch.int32) + dist.all_reduce(ok, op=dist.ReduceOp.MIN) + validated = bool(ok.item()) + + if rank == 0: + result = { + "mode": args.mode, + "world_size": world_size, + "m": args.m, + "n": args.n, + "k": args.k, + "quant": args.quant, + "block_m": args.block_m, + "block_n": args.block_n, + "chunks": chunks if fused else None, + "tile_order": args.tile_order, + "swap_ab": args.swap_ab, + "permlane": args.permlane, + "lane_transpose": args.lane_transpose, + "fence": args.fence if fused else None, + "stop_after": args.stop_after, + # The wire and transport, which "mode" does not distinguish: an + # fp8/lsa and an fp8/sdma run both say mode="fused-sdma", and a + # record separated from its command line cannot say which ran. + "gather_dtype": args.gather_dtype, + "gather_transport": args.gather_transport, + "fuse_quantize": args.fuse_quantize, + "fuse_reduce_push": args.fuse_reduce_push, + "publish": args.publish, + "gather_bands": args.gather_bands, + "sdma_queues": args.sdma_queues, + "atomic_order": args.atomic, + # "auto" resolves to the fused path's rotation + # (kernels_fused: `rotated = fuse if rotated is None`), so record + # what ran rather than the word that selected it. + "resolved_tile_order": ( + ("rotated" if fused else "linear") + if args.tile_order == "auto" + else args.tile_order + ), + "max_rank_time_us": per_rank[max_rank], + "critical_rank": max_rank, + "per_rank_time_us": per_rank, + "rel_l2": rel_l2, + "validated": validated, + "timing": "eager" if args.eager else "graph", + } + print("RESULT_JSON " + json.dumps(result, sort_keys=True), flush=True) + print( + f"[gemm_ar] mode={args.mode} m={args.m} n={args.n} k={args.k} " + f"max_rank_time={per_rank[max_rank]:.2f}us " + f"relL2={rel_l2:.2e} validated={validated}", + flush=True, + ) + if args.json_out: + with open(args.json_out, "w") as fh: + json.dump(result, fh, indent=2, sort_keys=True) + + dist.barrier() + return 0 if (result is None or result["validated"]) else 1 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=MODES, default="fused-sdma") + p.add_argument("-m", type=int, default=4096) + p.add_argument("-n", type=int, default=7168) + p.add_argument("-k", type=int, default=1024) + p.add_argument( + "--quant", + choices=("ptpc", "blockscale"), + default="ptpc", + help="a8w8 per-token/per-channel (the aiter 8wave kernel's native form) " + "or the model's 1x128 / 128x128 block scale. blockscale applies the " + "scales per K-block in the mainloop, which needs a second accumulator " + "and therefore --block-m 128", + ) + p.add_argument( + "--block-m", + type=int, + default=0, + help="0 = 256 for --quant ptpc, 128 for --quant blockscale", + ) + p.add_argument("--block-n", type=int, default=256) + p.add_argument("--waves-per-eu", type=int, default=2) + p.add_argument("--xcd-swizzle", type=int, default=0) + p.add_argument( + "--chunks", + type=int, + default=0, + help="pushes per destination (0 = the default 8, halved until it " + "divides the M-tiles per destination). >1 is what produces the " + "overlap, and requires the submit lock", + ) + p.add_argument( + "--n-stripe", + type=int, + default=0, + help="with the rotated order: rotate the destination every N N-tiles " + "instead of after a whole chunk (gcnasm opus_direct_stripe_tile). " + "0 = per-mode default: 2 for fused-lsa, chunk-major elsewhere. 1 rotates " + "the destination on every block", + ) + p.add_argument( + "--tile-order", + choices=("auto", "rotated", "linear"), + default="auto", + help="auto = rotated when fusing, linear otherwise", + ) + p.add_argument( + "--fence", + choices=( + "all", + "agent", + "agent-leader", + "nt-agent", + "leader", + "none", + "writethrough", + "wt-agent", + "raw-wt", + "raw-wt-agent", + "raw-wt-leader", + ), + default="none", + help="release issued before the epilogue hands a range to the copy " + "engines. The default is 'none'. Measured at m=16384 blockscale on 8 " + "ranks, twice each: 'none' 1143.5/1144.8us, 'writethrough' " + "1148.5/1145.4, 'agent' 1797.0/1798.5, 'all' 1936.1/1937.5 -- all four " + "at relL2 2.350e-3, i.e. the explicit fences cost 650-790us here and " + "did not move accuracy. Every 'raw-wt*' mode is known-racy and kept " + "only to reproduce that. See docs/MORI-GEMM-AR-BENCHMARK.md for how a " + "release can be paid in the store policy instead of in a fence", + ) + p.add_argument( + "--stop-after", + choices=("all", "scatter"), + default="all", + help="'scatter' stops after the reduce-scatter transfer has landed, " + "which is the only phase fusion can affect; use it to attribute a " + "fused-vs-split difference instead of inferring it from the total", + ) + p.add_argument( + "--no-put", + action="store_true", + help="keep the epilogue but drop the transfer, to price how much of its " + "cost is the copy engine reading C while the GEMM writes it. Output is " + "wrong by construction; use with --skip-validation", + ) + p.add_argument( + "--atomic", + choices=("acq_rel", "acquire", "release", "monotonic"), + default="acq_rel", + help="ordering on the tile counter. The winner needs the acquire half to " + "order the other blocks' releases before its put; 'monotonic' drops it " + "and races", + ) + # The three C-store stages are ON by default. They compose as a ladder -- + # each needs the one before it -- and measured together at [16384, 7168] + # K=2048 on 8x MI355X they help every mode and change no output: + # gemm-only -5.0%, fused-sdma -3.2%, fused-lsa -2.4%, split-lsa -0.9%, + # split-sdma -0.7%, all still at relL2 2.350e-3. Turn one off with its + # --no- form (and the ones above it in the ladder). + # + # Note where the win is: the stages coalesce the *store issue*, so they pay + # off when C lands in local memory (fused-sdma's GEMM 297.3 -> 256.6 us) + # and not when it goes straight to a peer (fused-lsa's GEMM 954.5 -> 952.1, + # i.e. nothing -- that path is xGMI-bound, and ATT shows 99% of its store + # time is stall, not issue). + # + # ``compile_fused_gemm_scatter`` still defaults them to False, so that + # its default stays the stock pipeline and these stages are opt-in. + p.add_argument( + "--swap-ab", + action=argparse.BooleanOptionalAction, + default=True, + help="exchange the MFMA operands so each lane owns 4 consecutive N, " + "letting C be stored 64 bits at a time instead of 16 " + "(gcnasm mfma_adaptor_swap_ab). On by default", + ) + p.add_argument( + "--store-probe", + action="store_true", + help="PERF PROBE, output is wrong by construction: emit the 16B-per-lane " + "store pattern the permlane stage would produce, without the shuffle, to " + "price it before building it. Requires --swap-ab; use --skip-validation", + ) + p.add_argument( + "--permlane", + action=argparse.BooleanOptionalAction, + default=True, + help="with --swap-ab: two permlane16_swap per M-tile so each lane owns 16 " + "contiguous bytes and each row gets 64. This is the real version of what " + "--store-probe prices. On by default", + ) + p.add_argument( + "--lane-transpose", + action=argparse.BooleanOptionalAction, + default=True, + help="with --swap-ab --permlane: one ds_bpermute per dword so adjacent " + "lanes cover one C row (gcnasm kernel_template.hpp:491-510). On by " + "default", + ) + p.add_argument( + "--hoist-scales", + action="store_true", + help="load each A/B scale once for the four half-tile stores instead of " + "twice (needs --permlane)", + ) + p.add_argument( + "--peer-uncached", + action="store_true", + help="fused-lsa: store to the peer with sc0|sc1 so nothing is left dirty " + "in a local L2 the barrier cannot reach", + ) + p.add_argument( + "--direct-fence", + choices=("all", "leader"), + default="leader", + help="fused-lsa: whether every lane or only thread 0 publishes the " + "block's peer stores. leader is legal because the barrier pair is closed " + "and is worth 80us (355 vs 434)", + ) + p.add_argument("--sdma-queues", type=int, default=8) + p.add_argument("--gather-bands", type=int, default=8) + p.add_argument( + "--publish", choices=("fence", "writethrough"), default="writethrough" + ) + p.add_argument( + "--fuse-reduce-push", + action=argparse.BooleanOptionalAction, + default=False, + help="fire the gather's SDMA puts from inside the reduce, band by band, " + "instead of from a separate kernel after it. bf16 gather only.", + ) + p.add_argument( + "--fuse-quantize", + action=argparse.BooleanOptionalAction, + default=False, + help="fold the fp8 narrowing into the reduce instead of running it as " + "its own kernel. Costs ~23us at M=16384 despite saving a read of the " + "bf16: off, matching build_sdma_phases' own default.", + ) + p.add_argument( + "--gather-transport", + choices=("sdma", "lsa"), + default="sdma", + help="who moves the fp8 gather. sdma pushes with the copy engines and " + "widens in a second kernel; lsa pulls over xGMI into registers and " + "widens on the way in, one kernel, at the cost of CU time.", + ) + p.add_argument( + "--gather-dtype", + choices=("bf16", "fp8"), + default="bf16", + help="what the all-gather leg carries. fp8 halves its bytes -- it is " + "~40%% of a fused layer and runs at the xGMI ceiling -- at a cost of " + "relL2 ~2.1e-2 against a bf16 wire that is exact.", + ) + p.add_argument("--warmup", type=int, default=10) + p.add_argument("--iters", type=int, default=51) + p.add_argument("--eager", action="store_true") + p.add_argument("--skip-validation", action="store_true") + p.add_argument("--json-out") + return p + + +if __name__ == "__main__": + sys.exit(run(build_parser().parse_args())) diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md new file mode 100644 index 000000000..a5601cea4 --- /dev/null +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -0,0 +1,437 @@ +# MORI GEMM + All-Reduce Benchmark + +Measurements for `mori.ops.gemm_ar`, the fused fp8 GEMM + all-reduce. The design +and the API live next to the code in +[`python/mori/ops/gemm_ar/README.md`](../python/mori/ops/gemm_ar/README.md); +this file is how to reproduce the numbers and what they were. + +Every number below was taken on **8x MI355X (gfx950)**, one node, at the shape +`[M, 7168]` with `K=2048` — DeepSeek-V4-Pro's `wo_b` under TP8 with +`--chunked-prefill-size 16384`. Kernel timings are the median over 11 +graph-replayed iterations, maximum over ranks, on an otherwise idle box. +Run-to-run spread at this shape is about **2%**, so differences below that are +not differences. + +## Table of Contents + +- [Running the benchmark](#running-the-benchmark) +- [Headline](#headline) +- [Where the time goes](#where-the-time-goes) +- [The fp8 wire](#the-fp8-wire) + - [Who moves the gather](#who-moves-the-gather) + - [The pull grid](#the-pull-grid) + - [What fp8 costs, numerically](#what-fp8-costs-numerically) +- [Model-level evaluation](#model-level-evaluation) +- [Negative results](#negative-results) +- [Reproducing the end-to-end numbers](#reproducing-the-end-to-end-numbers) + +## Running the benchmark + +Needs a mori built with `BUILD_CCO_SDMA=ON`. Setting `MORI_ENABLE_SDMA` in the +environment only rebuilds the *device* bitcode — a host library built without +the flag has no SDMA queues, every put silently does nothing, and the all-reduce +quietly produces zeros. + +```bash +cd /path/to/mori +BUILD_CCO_SDMA=ON pip install . + +MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo \ + torchrun --standalone --nproc_per_node=8 \ + benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py \ + --mode fused-sdma --quant blockscale -m 16384 -n 7168 -k 2048 +``` + +`--mode` selects what is measured, all reaching the same end state: + +| mode | what runs | +|---|---| +| `gemm-only` | the GEMM alone, to size the ceiling | +| `split-sdma` | `gemm` then a 4-kernel SDMA all-reduce | +| `fused-sdma` | GEMM with the scatter fused into its epilogue | +| `split-lsa` | `gemm` then the 2-kernel LSA all-reduce | +| `fused-lsa` | GEMM storing straight into peers | + +`--quant blockscale` is the model's own quantisation (A 1x128, B 128x128, fp32 +scales) and is the column to read. `--gather-dtype fp8` and +`--gather-transport {sdma,lsa}` select the wire; see [the fp8 wire](#the-fp8-wire). + +Each run prints a `RESULT_JSON` line with `max_rank_time_us`, `rel_l2` and +`validated`, so a sweep can be parsed rather than eyeballed. + +> **Between runs, give the SDMA queues time to drain.** Back-to-back 8-rank runs +> hit `hsaKmtCreateQueueExt` failures (`anvil.cpp:237`) if a previous run's ranks +> have not exited. ~10s is enough; a leftover server holding queues is not. + +## Headline + +`--quant blockscale`, median of 11: + +| M | `split-sdma` | `fused-sdma` | `fused-sdma` + fp8 gather | +|---|---:|---:|---:| +| 4096 | 398.1 us | 351.0 us | **329.5 us** | +| 8192 | 722.0 | 621.1 | **539.3** | +| 16384 | 1472.9 | 1148.8 | **979.3** | + +Fusing is worth **-22%** at M=16384; the fp8 gather a further **-15%**. + +> Both survive to the server, in the same order. See +> [the end-to-end numbers](#reproducing-the-end-to-end-numbers). + +For scale, the same layer as the model runs it today (a separate GEMM then an +NCCL all-reduce) measures **1419.5 us** at M=16384, and the GEMM alone is +**369.4 us**. + +## Where the time goes + +Per layer, captured in SGLang over one 20000-token prefill (not the standalone +benchmark — this is the pipeline as the model drives it): + +| phase | bf16 | fp8 / sdma | fp8 / lsa | +|---|---:|---:|---:| +| gemm | 441.9 us | 438.2 us | 441.3 us | +| drain | 180.2 | 170.9 | 204.8 | +| reduce | 41.7 | 42.5 | 43.1 | +| quantize | — | 11.8 | 11.9 | +| gather | 437.8 | 256.4 | 7.5 (barrier only) | +| dequantize | — | 61.0 | — | +| pull | — | — | 230.5 | +| **wo_b layer** | **1101.6** | **980.8** | **939.1** | +| | | -11.0% | **-14.7%** | + +Two things to read out of the bf16 column: + +* **`gather` is the bottleneck**, not `drain`. It moves 196 MiB at 470 GB/s, + which is 7 xGMI links flat out, so halving its bytes halves its time. +* **`drain`'s apparent 1140 GB/s is not a bandwidth.** Seven links cannot do + that. It is the tell that the scatter's pushes already went out from the GEMM + epilogue and the drain is only waiting for the tail — which is why the scatter + leg has far less to give than its byte count suggests, and why it is still + bf16. + +## The fp8 wire + +`--gather-dtype fp8` sends the all-gather leg as e4m3 with one fp32 scale per +row. The reduce still accumulates in fp32 and `output` is still bf16; only the +wire changes. The scatter leg is unchanged — it carries partial sums that are +then added across every rank, so its error would compound rather than being a +single rounding. + +### Who moves the gather + +| `--gather-transport` | how | +|---|---| +| `sdma` | copy engines push, a second kernel widens | +| `lsa` (default) | CUs pull over xGMI and widen on the way in | + +A copy engine has no ALU, so for SDMA the widening *cannot* be the same step: it +is a second kernel that reads the landed fp8 back out of local HBM, 98 MiB a +layer. A CU pull has those bytes in registers already. + +| gather | M=16384 | +|---|---:| +| bf16 / sdma | 1150.7 us | +| fp8 / sdma | 1018.9 | +| **fp8 / lsa** | **957.3** | + +The ordering holds in the server too: 1070.3ms of GPU busy for bf16 against +1052.1 for fp8/sdma and 1041.2 for fp8/lsa. + +The pull also removes fp8's small-M penalty. With the SDMA gather the two +conversion kernels were a fixed cost against a transfer that shrinks with M, so +fp8 measured **+2.3% (slower)** at M=4096. With the pull there is no fixed +conversion cost left and fp8 wins wherever fusing does. + +### The pull grid + +The single most important tuning parameter, and it is not obvious from the +source. These are *xGMI* reads, so the grid throttles outstanding remote +requests rather than covering HBM latency — it wants roughly a **tenth** of what +the local conversion kernels want. + +| blocks | 16 | 24 | 32 | 48 | 64 | 80 | 128 | 256 | 512 | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| us | 1261 | 1091 | 1006 | 962 | **959** | 963 | 1018 | 1138 | 1184 | + +Flat from 48 to 80, steep either side. The first implementation launched 512 — +the grid the local quantize kernel uses — and **lost to SDMA by 17%**, which +looked like "LSA is the wrong transport" rather than "the grid is wrong". + +Note this is also not `LSA_BLOCK_CAP`'s 24: that cap is for a kernel moving bf16 +with no arithmetic, while this one moves half the bytes and dequantises them, so +it needs more waves in flight to keep the links fed. + +Sweep it with `MORI_GEMM_AR_PULL_BLOCKS`: + +```bash +for b in 16 24 32 48 64 80 128 256; do + MORI_GEMM_AR_PULL_BLOCKS=$b MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo \ + torchrun --standalone --nproc_per_node=8 \ + benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py --mode fused-sdma \ + --quant blockscale --gather-dtype fp8 --gather-transport lsa \ + -m 16384 -n 7168 -k 2048 + sleep 10 +done +``` + +### What fp8 costs, numerically + +relL2 against an fp32 host reference goes from **2.35e-3** (bf16 wire, which is +bitwise exact through the collective) to **2.49e-2**. + +That is a floor, not a tuning problem. e4m3 carries 3 mantissa bits, and scale +granularity barely moves it — measured in torch on a `[2048, 7168]` standard +normal payload: + +| scale granularity | relL2 | scale bytes | +|---|---:|---:| +| per row (7168) | 2.646e-2 | 0.06% | +| per 512 | 2.631e-2 | 0.78% | +| per 256 | 2.609e-2 | 1.56% | +| per 128 | 2.572e-2 | 3.12% | +| per 32 | 2.399e-2 | 12.5% | + +**200x the scale bytes buys 9%.** Per-row is therefore the right choice, and +~2.5e-2 is what fp8 costs. + +## Model-level evaluation + +The kernel-level cost above is large — 6.7x the bf16 wire. Whether it *matters* +is a different question, and it needs the model, so this section was measured in +SGLang on DeepSeek-V4-Pro at TP8. The fp8 path was verified live throughout: +`SGLANG_DEBUG_FUSED_WO_B_AR=1` logs relL2 per layer call during the very +requests being scored. + +**At the layer**, relL2 against the unfused path, 488 layer calls: + +| wire | min | median | max | +|---|---:|---:|---:| +| bf16 | 3.706e-3 | — | 4.046e-3 | +| fp8 / sdma | 1.435e-2 | **2.505e-2** | 2.654e-2 | +| fp8 / lsa | 2.153e-2 | **2.496e-2** | 2.683e-2 | + +**At the model output**, it is not detectable. Scoring 10941 tokens of real +source text in a single prefill (mean logprob; lower is a worse model): + +| run | mean logprob | ppl | +|---|---:|---:| +| bf16 | -1.184169 | 3.2680 | +| bf16, rerun | -1.182018 | 3.2609 | +| bf16, again | -1.180106 | 3.2547 | +| **fp8** | **-1.183895** | **3.2671** | +| fp8, with debug | -1.183358 | 3.2653 | + +fp8 lands **inside the bf16 run-to-run band**. Paired per token against the same +bf16 run: + +| pair | mean d | sd d | max abs d | +|---|---:|---:|---:| +| CONTROL bf16 rerun | +0.002151 | 0.229 | 3.42 | +| CONTROL bf16 again | +0.004063 | 0.230 | 4.05 | +| **TEST fp8** | **+0.000274** | **0.219** | **2.93** | + +On every statistic, fp8 is closer to bf16 than bf16 is to itself. + +That band is wide because **this model is already strongly non-deterministic**: +two bf16 runs disagree on ~48% of tokens by more than 0.01 logprob, and greedy +decode diverges within 10-20 tokens. The likely cause is the MoE stage-2 +epilogue, which accumulates with `atomic_fadd`. This is also why greedy token +agreement is useless as a metric here — the bf16-vs-bf16 control is as divergent +as bf16-vs-fp8: + +| pair | ~12k tok | ~24k tok | +|---|---:|---:| +| CONTROL bf16 vs bf16 | 35.9% | 15.6% | +| TEST bf16 vs fp8 | 28.1% | 23.4% | + +Needle-in-a-haystack retrieval at 15140 tokens is **24/24 on both wires** — +saturated, so it bounds gross damage without resolving anything finer. + +**What this does and does not say.** It says fp8 causes no gross degradation and +no measurable shift in next-token distribution on one scoring task. It does not +say quality is unaffected on long-chain reasoning, code or maths — that needs a +task benchmark, which has not been run. Note also that short prompts and decode +never reach this path at all (it engages only at M >= 4096), so only +long-prefill workloads are affected. + +## Negative results + +Kept because each reads as obviously right and the reason it is not cannot be +seen from the source. + +**Folding the narrowing into the reduce** (`fuse_quantize=True`) saves a 28 MiB +re-read and a kernel launch, and costs 20 us: + +| | us | +|---|---:| +| split reduce + quantize | 957.4 | +| fused, row stashed in registers | 977.1 | +| fused, row re-read | 982.0 | + +Not register pressure — the re-reading variant keeps no stash and is no better. +It is the thread map: a per-row amax cannot be taken by a block holding only part +of a row, so fusing forces one-wave-per-row, where `sdma_reduce` walks packs with +a flat grid stride and streams a block through all 8 source slices at once. + +Re-measured with three repeats each, the cost holds and is if anything larger: +926.6 against 903.2 on `fp8/lsa`, 980.4 against 961.9 on `fp8/sdma`. +`build_sdma_phases` has always defaulted it off, so the op and SGLang take the +fast path; the *benchmark* defaulted it on until `96bd9522`, and any number +produced by an older driver without an explicit `--no-fuse-quantize` reads about +23 us slow. + +**Firing the gather's puts from inside the reduce** (`fuse_reduce_push=True`) +looked like the safest of the three: the push sends this rank's *own* slice, so +unlike the pull it has no cross-rank dependency, and SDMA is a copy engine so it +costs no CU time. It reaches parity and not a win — against 1148.7us unfused: + +| bands | 1 | 4 | 8 | 16 | 32 | +|---|---:|---:|---:|---:|---:| +| `publish="writethrough"` | 1162.1 | **1157.2** | 1160.3 | 1190.1 | 1323.5 | +| `publish="fence"` | 1227.3 | 1380.1 | 1630.3 | 2137.2 | 3026.7 | + +The gap between those rows is the useful result, and it is a lesson about how to +pay for a release rather than whether to. + +Handing a range to a copy engine *does* need one: the engine reads over the +fabric, not through a CU's cache, so `s_waitcnt vmcnt(0)` alone is not enough — +it only retires the stores as far as this XCD's L2. But the release can be paid +two ways. Releasing to system scope **after** the stores is L2-writeback work +charged once per block per band (256 x bands of it, ~61us per band, against a +reduce that is only 42us in total — unrepayable). Storing with `sc0+sc1` so the +bytes never stop in L1 or L2 makes the `waitcnt` itself the release, and that is +**free here**: applying the same store policy to the plain unfused reduce moves +it 1153.4 -> 1148.7us, i.e. nothing. This output is written once and nothing +local reads it again before the gather, so holding it in L2 bought nothing. + +What remains after that fix is small on both sides and nearly cancels. At +bands=1 the publish carries the mechanism's cost with none of its benefit — +1162.1 vs 1148.7, about 13us for the per-band 256-block `wait_barrier`, the +counter atomic and the elected block's locked puts. Four bands buy back about +5us of overlap before the sync cost takes over again. + +Dropping the release entirely is not an option even though it briefly looks like +one: with cached stores and no fence the kernel reaches 1157.9us, but at 32 bands +it produced relL2 1.8e-2 against the 2.35e-3 floor, differing per rank. The same +32 bands are exact under either correct publish mode, which rules out an indexing +bug. + +The contrast with the GEMM's fused scatter is the transferable part: there the +publish is amortised against a 437us transfer hidden behind a compute-bound GEMM; +here against 42us of bandwidth-saturated reduce. The mechanism pays when what is +hidden is much larger than the cost of publishing it. + +Re-measured after the window-geometry work below, with three alternating +repeats rather than a sweep, it is a clearer loss than the table suggests: +1128.4 us on against 1112.0 off, +16.4 us, against spreads of 2.2 and 1.5 us. + +**Hoisting the window geometry out of `lsa_ptr`.** `cco_lsa_ptr` is +`winBase + peer*stride + offset` and loads both fields on every call, through a +*generic* pointer -- which has to be a `flat_load`, since the compiler cannot +rule out LDS, so it counts against `lgkmcnt` as well as `vmcnt`. FlyDSL emits it +as an opaque extern call, and a kernel storing through addresses derived from +that base gives LLVM no way to prove the loads are not clobbered. + +Reading the geometry once and doing the arithmetic in the DSL removes all of +that. It was tried four ways -- hoisting out the band loop, a `lsa_geometry()` +API, `global_load` accessors in C++ (`cco_lsa_win_base` / `cco_lsa_stride`, which +take an `address_space(1)` pointer so each is a single `global_load`), and +finally `cco.CachedWindow`, which reads both in its constructor so a kernel +changes by one line. All four measured nothing on `kernels_sdma` (21 call sites) +and `kernels_fused`, in every wire configuration. + +It pays in exactly one place (`ptpc`, three alternating repeats): + +| | Window | CachedWindow | +|---|---|---| +| `split-lsa` | 1264.69 1264.23 1264.85 | 1250.25 1256.85 1254.53 | +| `fused-sdma` | 1110.45 1109.77 1112.85 | 1110.53 1113.69 1112.61 | + +-10.7 us on `split-lsa`, against a 0.6 us spread; nothing on `fused-sdma`. +`ar_1stage`/`ar_2stage` build nine peer addresses in *every block* of a short +kernel; everywhere else the addresses are built once per launch against a body +that runs for a millisecond. **Count address constructions per launch, not +`grep -c lsa_ptr`.** + +The same holds against PR #662's branch in `blockscale`, two alternating +repeats: `split-lsa` 1463.6 -> 1455.6 us, while `split-sdma` (1461.2 -> 1459.9), +`fused-sdma` bf16 (1144.9 -> 1145.8) and fp8/lsa (949.2 -> 949.8) do not move. + +Two things worth carrying. A `CachedWindow` cannot cross an `scf.if` -- FlyDSL +captures every variable an if body reads as state and requires single MLIR +values, which `Window` satisfies only by having exactly one field -- and the way +out is to compute the addresses before the branch, which is what the offsets +usually allow. And pin `--quant` when comparing against anything: the benchmark +defaults to `ptpc`, ~3% faster than the `blockscale` every number on this page is +quoted in, and reading one against the other looks exactly like a machine that +drifts overnight. + +**A CK-shaped 4-wave GEMM**, chasing a 22% gap against CK's block-scale kernel +at the same shape, reached CK's instruction mix and not its speed. Ten hypotheses +were falsified by measurement; hardware counters show identical `SQ_INSTS_MFMA` +(7,340,032) and `SQ_VALU_MFMA_BUSY_CYCLES` (234,881,024), VALU within 1%, +`MemUnitStalled` at approximately zero — but `SQ_WAIT_ANY` 156.0M against 120.4M. +A K-sweep puts the whole difference per-iteration: our fixed cost is *lower* +(51.6 us against 66.0), while each K-block costs 20.9 us against 14.5. The gap is +wait, not work. See commits `cc696762`, `54fef960`, `9f990637`, `1dead3fc`. + +## Reproducing the end-to-end numbers + +The SGLang numbers need the integration branch and a mori built with +`BUILD_CCO_SDMA=ON` on `PYTHONPATH`. + +```bash +export MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo +export PYTHONPATH=/path/to/mori-with-sdma +export SGLANG_OPT_FUSED_WO_B_AR=1 +export SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER=1 # optional, the fp8 wire +export SGLANG_DEBUG_FUSED_WO_B_AR=1 # optional, logs per-layer relL2 + +sglang serve --model-path --tp 8 \ + --attention-backend dsv4 --page-size 256 --chunked-prefill-size 16384 \ + --mem-fraction-static 0.88 --kv-cache-dtype fp8_e4m3 \ + --enforce-shared-experts-fusion +``` + +`--mem-fraction-static` has to leave room for the symmetric window, which is VMM +memory **outside** torch's allocator: 700 MiB on the bf16 wire and 812 MiB on +fp8, which needs the extra staging region. + +The capture protocol matters more than it looks. Profile a *different* prompt +than the one used to warm up, `flush_cache` between them, and compare the same +request on both sides — an earlier A/B without those controls reported a **+8.3% +regression** that did not exist. + +```python +warm = "Pack my box with five dozen liquor jugs. " * 2800 +main = "The quick brown fox jumps over the lazy dog. " * 2800 +gen(warm); post("/flush_cache"); post("/start_profile") +gen(main); post("/stop_profile") +``` + +GPU busy time over that capture: + +GPU busy over that capture: + +| | GPU busy | wall | vs unfused | +|---|---:|---:|---:| +| unfused (GEMM + NCCL) | 1096.0 ms | 1.1969 s | — | +| fused, bf16 wire | 1070.3 | 1.1728 | -2.3% | +| fused, fp8 / sdma | 1052.1 | 1.1455 | -4.0% | +| **fused, fp8 / lsa** | **1041.2** | **1.1385** | **-5.0%** | + +An earlier capture of the same four read 1101.9 / 1077.2 / 1050.4 / 1048.2, so +this reproduces to about half a percent. + +> **Check `BUILD_CCO_SDMA=ON` before believing any end-to-end number.** With it +> off every put silently does nothing: the all-reduce returns mostly the local +> slice, the model still answers fluently, every mori kernel still appears in +> the profile, and the fused path measures **faster** than it is because it is +> not moving data -- -6.8% instead of -2.3%, with fp8/lsa appearing *worst* of +> the three rather than best, since the pull is the one leg that does not go +> through SDMA. Perplexity catches it and nothing cheaper does: 862511 against +> 3.26 on the same text. A short prompt cannot catch it either, because fusing +> needs M >= 4096. + +The layer-level win is larger than the end-to-end one because `wo_b` is about +12% of the profile. diff --git a/docs/index.rst b/docs/index.rst index dd13be026..9d70c8ce2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -26,6 +26,7 @@ MORI Documentation :caption: Benchmarks MORI-EP-BENCHMARK + MORI-GEMM-AR-BENCHMARK MORI-IO-BENCHMARK .. toctree:: diff --git a/python/mori/cco/device/flydsl/__init__.py b/python/mori/cco/device/flydsl/__init__.py index c42dbd759..a38209379 100644 --- a/python/mori/cco/device/flydsl/__init__.py +++ b/python/mori/cco/device/flydsl/__init__.py @@ -43,12 +43,22 @@ from mori.cco.device.bitcode import find_cco_bitcode, get_bitcode_path from mori.cco.device.ops import SdmaOptFlags -from .handles import DevComm, Window, Gda, CoopScope, SignalOp, ThreadMode, Sdma +from .handles import ( + DevComm, + Window, + CachedWindow, + Gda, + CoopScope, + SignalOp, + ThreadMode, + Sdma, +) from . import _bindings __all__ = [ "DevComm", "Window", + "CachedWindow", "Gda", "CoopScope", "SignalOp", diff --git a/python/mori/cco/device/flydsl/_bindings.py b/python/mori/cco/device/flydsl/_bindings.py index dff227c02..af6bbe29b 100644 --- a/python/mori/cco/device/flydsl/_bindings.py +++ b/python/mori/cco/device/flydsl/_bindings.py @@ -98,6 +98,9 @@ def _binding(name): # ── axis-free symbols ── # cco_lsa_ptr(window, peerLsaRank, offset) -> peer's load/store-accessible VA. cco_lsa_ptr = _binding("lsa_ptr") +# Window geometry: one global_load each, for callers that take it once. +cco_lsa_win_base = _binding("lsa_win_base") +cco_lsa_stride = _binding("lsa_stride") cco_system_fence = _binding("system_fence") cco_devcomm_rank = _binding("devcomm_rank") diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index a8f917d8a..70b34d3b0 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -111,6 +111,75 @@ def lsa_ptr(self, peer_lsa_rank, offset=0): return raw.cco_lsa_ptr(self.handle, peer_lsa_rank, offset) +class CachedWindow: + """A :class:`Window` that reads its geometry once, at construction. + + Drop-in for ``Window`` at the top of a kernel:: + + w = cco.CachedWindow(win) # instead of cco.Window(win) + + and every ``w.lsa_ptr(...)`` after it becomes arithmetic instead of an + extern call. Nothing else about the kernel changes. + + ``cco_lsa_ptr`` is ``winBase + peer*stride + offset`` and loads *both* + fields out of the window descriptor on every call. It is opaque to FlyDSL's + tracer, and although it is ``always_inline`` -- so LLVM does see the loads -- + a kernel that stores through addresses derived from that same base gives + LLVM no way to prove they are not clobbered, so it reloads them. + + The two reads here go through ``cco_lsa_win_base`` / ``cco_lsa_stride``, + which take an ``address_space(1)`` pointer so each is a single + ``global_load``. ``cco_lsa_ptr`` casts to a *generic* pointer, and a generic + access has to be a ``flat_load`` -- the compiler cannot rule out LDS -- so it + counts against ``lgkmcnt`` as well as ``vmcnt``, giving any following + ``s_waitcnt`` one more counter to wait on. + + **What it is worth depends on how often the kernel builds an address, not on + how many call sites it has.** A/B at ``m=16384 n=7168 k=2048`` on 8x MI355X, + max_rank_time_us, three alternating repeats:: + + mode Window CachedWindow + split-lsa 1264.69 1264.23 1264.85 1250.25 1256.85 1254.53 -10.7us + fused-sdma 1110.45 1109.77 1112.85 1110.53 1113.69 1112.61 null + + ``split-lsa`` wins 0.85%, against a Window spread of 0.6us across its three + runs. Its ``ar_1stage``/``ar_2stage`` build nine peer addresses in *every + block* of a short kernel, so the extern calls are a real fraction of it. + ``fused-sdma`` gains nothing: its ``lsa_ptr`` calls are per kernel launch, + once, against a body that runs for a millisecond -- and the same was true of + ``kernels_sdma``, where converting all 21 sites moved nothing. + + So: count address constructions per launch, not ``grep -c lsa_ptr``. + + Measure A/B in one session if you revisit this, and pin ``--quant``. The + benchmark defaults to ``ptpc``, which is ~3% faster than the ``blockscale`` + the model actually uses and that the headline numbers are quoted in; reading + one against the other looks exactly like a machine that drifts overnight. + + Unlike :class:`Window` this is a plain Python object, not a ``cco_struct``, + so it cannot cross an ``scf.if``/``scf.for`` boundary -- and neither would a + ``cco_struct`` version, because FlyDSL captures every variable a dynamic + ``if`` body reads as state and requires each to be a *single* MLIR value. + ``Window`` qualifies only by having one field; a three-field struct fails the + same check ("state variable 'w' is list, not an MLIR Value"). + + Usually the way out is neither: compute the addresses *before* the branch and + let the ``if`` capture plain ``Int64``. That works whenever the offsets do + not depend on the branch, which in practice they rarely do. + """ + + __slots__ = ("handle", "base", "stride") + + def __init__(self, win): + self.handle = fx.Int64(win) + self.base = fx.Int64(raw.cco_lsa_win_base(self.handle)) + self.stride = fx.Int64(raw.cco_lsa_stride(self.handle)) + + def lsa_ptr(self, peer_lsa_rank, offset=0): + """Same contract as :meth:`Window.lsa_ptr`, with no extern call.""" + return self.base + fx.Int64(peer_lsa_rank) * self.stride + fx.Int64(offset) + + @cco_struct class Gda: """GDA handle bound to a ccoDevComm device pointer + context index. @@ -259,8 +328,9 @@ def _xfer( signal, aggregate, ): - # coop/signal are compile-time; "_ns" (no-signal) is fire-and-forget and - # cannot be drained by quiet. aggregate skips the doorbell — call commit(). + # coop/signal are compile-time. "_ns" (no-signal) only drops the trailing + # ATOMIC; quiet()/quietQueue() drain the queue's read pointer and work on + # it either way (cco.hpp:1707). aggregate skips the doorbell — call commit(). tag = _COOP_TAG[_const(coop, "coop")] if not signal: tag += "_ns" diff --git a/python/mori/cco/device/ops.py b/python/mori/cco/device/ops.py index f2b0b6663..21acb2a23 100644 --- a/python/mori/cco/device/ops.py +++ b/python/mori/cco/device/ops.py @@ -125,6 +125,13 @@ def _add( pure=True, family="lsa", ) +# Window geometry, for callers that need many addresses from one window. +# Both read through an address_space(1) pointer so the access is a global_load +# (vmcnt only) rather than the flat_load a generic pointer forces (vmcnt and +# lgkmcnt). See cco_lsa_win_base in the device wrapper. +_add("lsa_win_base", "cco_lsa_win_base", [_U64], _U64, pure=True, family="lsa") +_add("lsa_stride", "cco_lsa_stride", [_U64], _U64, pure=True, family="lsa") + # Benchmark-only publication primitive. It is not a cooperative barrier: # leaderOnly=1 fences thread 0 only, so producers must synchronize first and # must not use it as a substitute for a per-lane release fence. diff --git a/python/mori/ops/__init__.py b/python/mori/ops/__init__.py index fa7aebd7d..66913ad4d 100644 --- a/python/mori/ops/__init__.py +++ b/python/mori/ops/__init__.py @@ -30,10 +30,10 @@ launch_local_expert_count, ) -# dispatch_combine_v2 (FlyDSL) requires the optional `flydsl` dependency, so it -# is imported lazily: `import mori.ops` stays usable without flydsl installed, -# and `mori.ops.dispatch_combine_v2` resolves on first access. -_LAZY_SUBMODULES = {"dispatch_combine_v2"} +# dispatch_combine_v2 and gemm_ar (FlyDSL) require the optional `flydsl` +# dependency, so they are imported lazily: `import mori.ops` stays usable +# without flydsl installed, and each submodule resolves on first access. +_LAZY_SUBMODULES = {"dispatch_combine_v2", "gemm_ar"} def __getattr__(name: str): diff --git a/python/mori/ops/gemm_ar/README.md b/python/mori/ops/gemm_ar/README.md new file mode 100644 index 000000000..eaeddf60e --- /dev/null +++ b/python/mori/ops/gemm_ar/README.md @@ -0,0 +1,614 @@ +# Fused GEMM + all-reduce (`mori.ops.gemm_ar`) + +Fuses a `RowParallelLinear`'s fp8 GEMM with the all-reduce that follows it, over +cco's SDMA copy engines. Built for DeepSeek-V4-Pro's `wo_b` and measured there: +**1146.2us against the model's 1419.5us, -19.4%**, at `[16384, 7168] K=2048` on +8x MI355X. + +The public API is `GemmAllReduceOp` in `op.py`. `kernels_fused.py`, +`kernels_sdma.py` and `kernels_lsa.py` hold the kernels, `layout.py` every +window offset and count, and `_gemm_a8w8_8wave.py` / `_shuffle.py` the two +pieces vendored from aiter so mori does not depend on it. + +Full measurements, how to reproduce them, and the model-level evaluation of the +fp8 wire are in +[`docs/MORI-GEMM-AR-BENCHMARK.md`](../../../../docs/MORI-GEMM-AR-BENCHMARK.md). + +## What it does + +aiter's 8-wave fp8 GEMM with the all-reduce's scatter fused into its epilogue. + +The target is DeepSeek-V4-Pro's ``wo_b``: a RowParallelLinear whose per-rank GEMM +is ``[M,1024] x [7168,1024]`` fp8 -> ``[M,7168]`` bf16, immediately all-reduced. +Split, that is ``gemm(); all_reduce()``. Fused, the GEMM's C lands directly in a +registered cco window and each destination's slice is pushed by the copy engine +as soon as its last tile is written, so the reduce-scatter transfer overlaps the +rest of the GEMM instead of following it. + +Only the *scatter* half of the all-reduce is absorbed. The reduce and all-gather +phases still run as their own kernels, reused verbatim from ``ar.kernels_sdma`` +(``build_sdma_phases``), with ``scatter`` swapped for its drain-only twin. + +**Nothing here is a production kernel.** aiter's tuned CSV picks a ck/asm/cktile +backend for (N=7168, K=1024), not this one, so the fused-vs-split comparison is +internally valid but is not a claim about DSV4 as shipped. + +## Using it + +Needs a mori built with `BUILD_CCO_SDMA=ON` and `MORI_ENABLE_SDMA=1` in the +environment. Setting the variable alone is not enough: a host library built +without the flag has no queues, so every put silently does nothing and the +all-reduce quietly produces zeros. + +```python +import torch, torch.distributed as dist +from mori.cco import Communicator, UniqueId +from mori.ops.gemm_ar import GemmAllReduceOp, preshuffle_b + +rank, world = dist.get_rank(), dist.get_world_size() +torch.cuda.set_device(rank) + +# The window comes out of the communicator's VMM reservation, so size that +# first -- the op cannot grow it later. +N, K, M_MAX = 7168, 2048, 16384 +vmm = 2 * GemmAllReduceOp.window_bytes_for(world, m_max=M_MAX, n=N) + (64 << 20) + +uid = [bytes(Communicator.get_unique_id()) if rank == 0 else None] +dist.broadcast_object_list(uid, src=0) + +# init takes the UniqueId wrapper, not the bytes that survive a broadcast. +with Communicator.init( + world, rank, UniqueId.from_bytes(uid[0]), per_rank_vmm=vmm +) as comm: + with GemmAllReduceOp(comm, n=N, k=K, m_max=M_MAX) as op: + b_shuffled = preshuffle_b(b_fp8) # [N, K], once per weight + + m_pad = op.padded_m(x.shape[0]) # instance method: M only + a_fp8, a_scale = quantize_per_1x128(op.pad_rows(x, m_pad)) + # The result aliases the window, which the `with` frees on exit -- + # so clone inside it, after the work it was issued behind has run. + out = op(a_fp8, b_shuffled, a_scale, b_scale)[: x.shape[0]] + torch.cuda.current_stream().synchronize() + out = out.clone() +``` + +**Pad before quantising.** A padded row produces a zero output row, which the +caller slices off; quantising first and padding after would also have to extend +the scales, and in the layout below that is not a row append. + +**Scales.** `a_scale` is the A 1x128 block scale as fp32. The kernel reads +element `(row, kb)` at `kb * M + row`, so three spellings are accepted because +each already *is* that order: flat 1-D, `[K/128, M]`, and `[M, K/128]` +**column-major** (stride `(1, M)`), which is transposed for free. + +`[M, K/128]` **row-major is rejected** with a `ValueError`. It is not that the +op cannot transpose it -- it is that the two `[M, K/128]` tensors are +indistinguishable by shape, and they need opposite treatment: the column-major +one must be read flat, the row-major one must be transposed. Guessing gets one +of them silently wrong (relL2 0.26 either way). Pass `.reshape(-1)` or `.t()` +yourself to say which you have. `aiter_per1x128_quant(transpose_scale=True)` +returns the column-major one, so `.reshape(-1)` is right for it. + +`b_scale` is `[N/128, K/128]` fp32 row-major. + +**Shapes.** `supports(m, n, k, world_size)` answers whether a shape is +expressible. `K >= 256` and a multiple of 128; `N` a multiple of 256, and of +1024 if the gather is fp8; `world_size` in [2, 8] and dividing 512. + +**Collective contract.** Every rank must construct the op and call it the same +number of times in the same order -- the phases carry device-side barriers. +Work is issued on the current stream and completes asynchronously, so the usual +stream rules apply to the result. + +**The result aliases the window** and is overwritten by the next call; clone it +to keep it. `pad_rows` returns a reused buffer with the same caveat. One +instance is not usable concurrently. + +**Distinct M values.** Each M compiles its own kernel and takes its own counter +set, capped by `max_shapes`. That defaults to *every* legal M under `m_max` +(`m_max / (world_size * block_m)`), so the cap is not something a caller +normally meets; a ninth distinct shape used to raise. Warm each M up once before +capturing a graph. + +**Cleanup.** `close()` (or the `with` above) releases the window and its memory +-- at the model shape about 700 MiB per rank. The communicator holds a strong +reference to each, so dropping the op is not enough. The dev-comm is +deliberately *not* destroyed here: its SDMA queues are communicator-scoped, and +tearing them down ahead of the communicator aborts the process. They go with the +communicator. Synchronise first if anything may still be in flight. + +**dtype.** A and B must be `torch.float8_e4m3fn`. `float8_e4m3fnuz` is the same +width but a different exponent bias, and the MMA atom implements OCP's, so it is +rejected rather than silently returning a result 4x too large. + +## fp8 on the wire + +`gather_dtype="fp8"` sends the all-gather leg as e4m3 with one fp32 scale per +row, halving its bytes. That leg is ~40% of a fused layer and already runs at +the xGMI ceiling (470 GB/s over 7 links), so halving the bytes halves the time. + +Measured, fused-sdma at `[16384, 7168]` K=2048 on 8x MI355X: + +| gather wire | us | relL2 | +|---|---:|---:| +| bf16 | 1151.5 | 2.35e-3 | +| fp8 | **1033.9** (-10.2%) | **2.49e-2** | + +### Who moves the fp8 gather + +`gather_transport="lsa"` (the default for fp8) pulls each peer's slice over +xGMI into registers and widens it on the way to memory. `"sdma"` pushes with +the copy engines and widens in a second kernel -- a copy engine has no ALU, so +for SDMA those cannot be one step, and the second kernel has to read the landed +fp8 back out of local HBM (98 MiB a layer). + +| gather | us | +|---|---:| +| bf16 / sdma | 1150.7 | +| fp8 / sdma | 1018.9 | +| **fp8 / lsa** | **957.3** | + +The pull's grid is the whole story and is not obvious: these are xGMI reads, so +the grid throttles outstanding remote requests rather than covering HBM latency, +and it wants roughly a tenth of what the local conversion kernels want. + +| blocks | 16 | 24 | 32 | 48 | 64 | 80 | 128 | 256 | 512 | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| us | 1261 | 1091 | 1006 | 962 | **959** | 963 | 1018 | 1138 | 1184 | + +The first version launched 512 -- the quantize grid -- and lost to SDMA by 17%. + +### Folding the narrowing into the reduce: measured, and it loses + +`fuse_quantize=True` makes the reduce write its bf16 and narrow to fp8 from the +same accumulators, saving a 28 MiB re-read and a launch. It is **off**, because +it costs 20 us rather than saving 12: + +| | us | +|---|---:| +| split reduce + quantize | 957.4 | +| fused, row stashed in registers | 977.1 | +| fused, row re-read | 982.0 | + +Not register pressure -- the re-reading variant keeps no stash and is no better. +It is the thread map: a per-row amax cannot be taken by a block holding only +part of a row, so fusing forces one-wave-per-row, where `sdma_reduce` walks +packs with a flat grid stride and streams a block through all 8 source slices at +once. Confining a wave to 14 KiB at a time costs the reduce more than the +re-read saves. + +It is not free and the cost is not a tuning problem. e4m3 carries 3 mantissa +bits, so one rounding costs ~2.1e-2 on a normal payload whatever the scale +granularity -- per-row measures 2.65e-2 and per-32 measures 2.40e-2, 9% better +for 200x the scale bytes. Going from 2.35e-3 to 2.49e-2 is the price of the +10%, and whether that is payable is a model-level question, not a kernel one. + +It also only pays at large M. The two conversion kernels are a fixed cost +against a transfer that shrinks with M, so on the standalone all-reduce it is ++2.4% at M=4096 and -11.2% at M=16384. Per phase at M=16384: quantize 12.4us, +dequantize 88.2us, against ~218us saved on the push. + +The scatter leg stays bf16. It carries partial sums that are then added across +every rank, so its fp8 error compounds rather than being a single rounding, and +it is already mostly hidden behind the GEMM -- its 1140 GB/s is not a bandwidth, +it is the tell that the pushes went out from the epilogue and the drain is only +waiting for the tail. `scatter_dtype="fp8"` sizes its regions but raises +`NotImplementedError`. + +## Benchmarks and tests + +```bash +# numerics +MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo \ + pytest tests/python/cco/test_gemm_ar.py tests/python/cco/test_flydsl_ar.py + +# the full mode comparison at the model's shape +MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo python -m torch.distributed.run \ + --standalone --nproc_per_node=8 benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py \ + --mode fused-sdma --quant blockscale -m 16384 -n 7168 -k 2048 +``` + +Requires a mori built with `BUILD_CCO_SDMA=ON`. Setting `MORI_ENABLE_SDMA` in +the environment only rebuilds the device bitcode -- a host library built without +the flag has no queues, so every put silently does nothing and the all-reduce +quietly produces zeros. + +## A measured negative result, for the record + +An earlier revision carried `kernels_preshuffle4w.py`, a port of CK's 4-wave +B-out-of-LDS shape, chasing a 22% gap between this GEMM and CK's at the same +shape. It reached CK's instruction mix and not its speed, and it was deleted +rather than carried (it needed ~1000 lines of further aiter vendoring to serve +a kernel nothing calls). What the investigation ruled out, since the same +ground should not be walked twice: + +Ten hypotheses were falsified by measurement -- promote arithmetic, register +spill, store width, VALU scheduling groups, MFMA batch size, scale loads, +address hoisting, load-to-use distance, occupancy, and dependency structure. +Hardware counters (`rocprofv3 --pmc`) show *identical* `SQ_INSTS_MFMA` +(7,340,032) and `SQ_VALU_MFMA_BUSY_CYCLES` (234,881,024), VALU within 1%, and +`MemUnitStalled` at approximately zero -- but `SQ_WAIT_ANY` at 156.0M against +CK's 120.4M. A K-sweep puts the whole difference per-iteration: our fixed cost +is *lower* (51.6us against 66.0us), while each K-block costs 20.9us against +14.5us. The gap is wait, not work, and it is not in any of the places listed +above. See commits `cc696762`, `54fef960`, `9f990637`, `1dead3fc` for the +traces. + +## Why the SDMA transport needs no epilogue change + +``C`` is an ordinary tensor argument, so pointing it at a cco window is a host-side +change; ``StoreC`` is untouched. That is the whole reason SDMA is the cheap +transport to fuse. The LSA alternative -- the epilogue storing straight into a +peer -- is implemented as ``fused-lsa``, and it is the slower of the two for the +reason that was predicted: ``StoreC._store_bf16`` writes one bf16 at a time +through ``BufferCopy16b``, and gcnasm measured that lane-scatter at 0.26x when +the destination is a peer. It spends ~560us of its GEMM pushing C over xGMI +where the copy engines move the same bytes in ~499. + +## Completion protocol + +One monotonic counter per (destination, chunk) in the window (``cfg.counter_off``). +Every block, after its four ``store_c.store`` calls: + + s_waitcnt vmcnt(0) ; s_barrier -- this block's C tile has retired + __threadfence_system() -- ...and is visible to the copy engine + thread 0: prev = atomic_add(counter[dest][chunk], 1) + if (prev + 1) % tiles_per_chunk == 0 -- I am the last tile of that chunk + sdma.put(dest, ...) -- fire and forget, no quiet + +Counters are never reset, so the modulo test works on every launch and the kernel +stays CUDA-graph-safe, exactly like the barrier flags in ``ar.kernels_lsa``. + +``quiet`` is deliberately *not* called here: gcnasm measured fusing it into the +GEMM as a 1.8-10.7% regression, so the drain kernel does it afterwards. No +per-destination spin lock either -- gcnasm needed one because several CTAs could +submit for one destination, whereas the modulo test elects exactly one. + +Tiles are walked (chunk, destination, n) with the destination rotated by rank, +rather than in aiter's linear order. Linear order finishes destination 0's whole +slice, then 1's, and so on, so the last destination's link only starts at the end +of the GEMM and nothing overlaps. The rotation is gcnasm's +``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for +the GEMM alone. + +## C store: three stages off gcnasm, 6.7% off the GEMM, bit-exact + +``StoreC`` emits 128 ``buffer_store_short`` per block -- one bf16 at a time -- +because of the MFMA accumulator layout, not a missed vectorization. Lane ``l`` +holds ``D[4*(l/16)+i][l%16]``: four consecutive *rows*, stride ``c_cols``, so its +four values are 14336 bytes apart in a row-major C, and the eight bf16 that would +make a 16-byte store live in eight different lanes. + + stage GEMM kernel store instrs store cycles + (aiter as-is) 38.30us 128 short 14,608 + --swap-ab slower 32 dwordx2 32,720 + --swap-ab --permlane 37.04us 16 dwordx4 14,604 + --swap-ab --permlane --lane-transpose 35.74us 16 dwordx4 3,852 + +Median of 42 dispatches, 8 ranks, [4096,7168] K=1024. Bit-identical to aiter at +[512,512,256], [1024,768,512] and wo_b (``test_swap_ab_is_bitwise_identical``), +and registers are unchanged throughout (VGPR 128 / SGPR 112 / 0 scratch). + +**1. ``--swap-ab``** -- ``mfma_adaptor_swap_ab`` (opus.hpp:2064, literally +``base::operator()(b, a, c)`` with ``dim_c()`` redefined). Computing +``B^T A^T = (A B)^T`` in the accumulator's own layout moves a lane to +``D[l%16][4*(l/16)+k]``: four consecutive *columns*, 8 contiguous bytes. Alone it +is *slower* -- a row still only gets 32 bytes, so the same 16 transactions are +squeezed onto a quarter as many instructions and per-instruction address fan-out +quadruples (1022 cycles each against 114). + +**2. ``--permlane``** -- two ``v_permlane16_swap_b32`` per M-tile. Semantics +measured rather than assumed: ``vdst' = [X.r0, Y.r0, X.r2, Y.r2]``, +``vsrc' = [X.r1, Y.r1, X.r3, Y.r3]``, so:: + + (A, B) = permlane16_swap(tile0.d0, tile1.d0) + (C, D) = permlane16_swap(tile0.d1, tile1.d1) + lane group g stores (A, C, B, D) + +lands g = 0,1,2,3 on columns 0-7, 16-23, 8-15, 24-31 -- together columns 0..31 +contiguously, 16 bytes per lane and 64 per row. No ``ds_bpermute``, no LDS; the +column permutation is absorbed into the address. + +This does **not** speed the store up (14,608 -> 14,604). What it pays for is +everything a 2-byte store drags along: 128 stores need 128 addresses, 128 bounds +predicates and 128 scalar multiplies, 16 need 16, and the swapped layout makes +B's scale a vec4 so the scaling packs into ``v_pk_mul_f32``:: + + v_mul_f32_e32 257 -> 1 v_lshlrev_b32_e32 159 -> 45 + v_pk_mul_f32 0 -> 128 v_add_u32_e32 133 -> 21 + v_cvt_pk_bf16_f32 128 -> 64 v_cndmask_b32_e64 96 -> 8 + +**3. ``--lane-transpose``** -- gcnasm's second stage +(kernel_template.hpp:491-510), one ``ds_bpermute`` per dword. After the permlane +stage a row's four 8-column chunks sit in lanes 16 apart, so a 16-lane group +touches 16 rows at 16 bytes each. Transposing the lane index -- lane +``l' = 4r'+q'`` pulls from the lane holding ``(row r', chunk q')``, with +``g = q``'s two bits swapped, 0,1,2,3 -> 0,2,1,3 -- puts adjacent lanes on one +row, so lanes 0-3 write 64 contiguous bytes and a group covers 4 rows. + +Same instructions, same 64 addresses, only which lane holds which. The store's +own latency falls **14,604 -> 3,852**, which is the coalescer being sensitive to +lane adjacency and not just to the address set -- exactly what gcnasm's +"pair-coalesced" comment is about. The 64 ``ds_bpermute`` cost 9,604 cycles. + +### ``--hoist-scales``: a real redundancy that does not pay to remove + +The epilogue calls ``store_c.store`` four times, and those calls share base rows +pairwise and base columns pairwise, so every scale is fetched twice. The +source-attributed thread trace counts it exactly: 16 ``buffer_load_dwordx4`` at +``gemm_a8w8_8wave.py:191`` where only 8 addresses are distinct, and 8 +``buffer_load_dword`` at :199 where only 4 are -- 12 of 24 loads redundant, plus +12 redundant address computations, 4,276 cycles or 0.8% of the kernel. The +compiler cannot merge them because all four calls write the same ``reg_f32_*`` +register buffer, which makes them a chain of overwrites rather than pure loads. + +``store_all`` loads each scale once. It removes exactly the predicted loads -- +A-scale 16 -> 8, B-scale 8 -> 4 -- and is slightly **slower**: + + variant VGPR scratch instrs dwordx4 dword gemm + --permlane --lane-transpose 254 0B 1581 72 16 33.32us + + --hoist-scales 256 0B 1595 68 8 33.96us + +Keeping both scale sets live across all four stores costs more in register +moves and re-materialised addresses than the twelve loads were worth, and it +spends the last 2 VGPRs of headroom (254 -> 256). Bit-exact either way. + +Kept as a switch rather than deleted: the redundant fraction grows with +``N_TILES_B``, so a larger ``BLOCK_N`` would change the arithmetic, and if +register pressure ever loosens this flips sign. It is also cheaper to re-measure +a flag than to re-derive why it was rejected. + +Two limits. **It is worth nothing on ``fused-lsa``**, where the store goes to a +peer over xGMI rather than to local memory:: + + gemm barrier reduce gather sum + fused-lsa 177.4us 5.4 11.1 135.3 329.1 + + all three 175.6us 7.8 11.1 135.0 329.4 + +against 39.0 -> 35.7us for the same three stages on the split path's local +store. Stage 3 buys the local memory pipeline's sensitivity to lane adjacency; +a peer store crosses the fabric instead, and that phase is limited by ~44 GB/s +of link bandwidth, which no amount of transaction shaping changes. (This was +first "measured" with only stages 1 and 2 wired into the direct path -- which +are exactly the two that do *not* speed the store up -- and the conclusion was +right by accident. The direct-LSA branch for stage 3 exists now, and the +argument-validation above is there so a missing branch fails loudly instead of +silently running a weaker variant.) + +And 64 bytes is one wave's ceiling here, since a wave owns +``N_TILES_B * 16 = 32`` columns; gcnasm reaches a full 128-byte line only by +having four ``wave_id_n`` waves tile adjacent 16-column runs. + +All three stages are candidates to push back into aiter's ``StoreC``: bit-exact, +no register cost, and the 6.7% is on the GEMM itself, independent of any +all-reduce. + +## Persistent tiles: tried, and there is no register budget for it + +gcnasm builds this kernel family both ways (``PERSISTENT=1|0``) and its README +has a tail-balance sweep: the win is entirely a function of the remainder after +whole 256-CTA batches -- +9.97% at remainder 8, +8.41% at 32, and only +0.77% at +192. wo_b is 16 x 28 = 448 tiles on 256 CUs, i.e. remainder **192**, the benign +end of that curve. + +The tail itself is real here, and large. Sweeping N at M=4096, K=1024 +(gemm-only, best of 3, 8 ranks): + + N tiles remainder time + 4096 256 0 32.64us + 4352 272 16 45.88us <- 6% more work, 40% more time + 4608 288 32 44.68 + 5120 320 64 45.24 + 6144 384 128 47.40 + 7168 448 192 49.80 + +So it was implemented (one workgroup per CU, striding over tiles) and made +bit-exact. It is a ~2x regression, and the reason is a hard resource wall rather +than anything to do with scheduling. From the kernel metadata in the final ISA: + + variant VGPR AGPR SGPR V-spill scratch instrs gemm + as committed 256 0 50 0 0B 2052 38.54us + + 3-stage C store 254 0 51 0 0B 1573 35.74us + body inside an scf.for 256 0 68 38 156B 2260 74.72us + persistent, grid 256 256 0 67 38 156B 2264 72.22us + +**The baseline already uses all 256 VGPRs with zero spill.** Wrapping the +pipeline in a runtime loop asks for 38 more -- the K pipeline's accumulators and +operand fragments die at the end of a tile in the flat version, but a loop makes +the compiler assume they may be live across the back-edge -- and there is nowhere +to put them. ``--amdgpu-num-vgpr 256/512`` changes nothing, because the cap was +never the constraint. Persistent and non-persistent spill identically, which +confirms the cost is ``scf.for`` itself and not the scheduling idea. + +Note ``fly-promote-regmem-to-vectorssa`` is **not** the problem, contrary to what +an earlier version of this note claimed. It handles ``scf::ForOp``, it promoted +all 451 register allocas here (zero left afterwards, zero ``llvm.alloca`` in the +final IR), and the emitted loop carries no ``iter_args`` at all. The 156 bytes +are ordinary register-allocator spill: ``.vgpr_spill_count: 38``, 38 +``scratch_store_dword`` / 38 ``scratch_load_dword``. + +The pass does have one real inefficiency, just not one we hit: an alloca declared +*outside* a loop is carried as an ``iter_arg`` unconditionally, even when it is +fully overwritten every iteration, because ``collectTouchedRegAllocaInRegion`` +records on any load *or* store with no liveness test. Allocas declared *inside* +the loop are correctly materialised as ``ub.poison`` and not carried, and ours +are all inside. + +Two bugs the attempt surfaced, worth knowing if anyone loops this pipeline: + +* **The half-wave barrier does not survive a loop.** The prologue's + ``if wave_m == 1: rocdl.s_barrier()`` deliberately runs the two half-waves one + barrier out of phase. Fine once; in a loop the offset accumulates by one per + tile, so from the second tile on the halves rendezvous at mismatched program + points and the LDS double-buffering races -- silently, as partial corruption + inside otherwise-correct tiles. A compensating ``if wave_m == 0: s_barrier()`` + at the end of each tile fixes it exactly. +* **The LDS handles must be rebound per tile.** The pipeline swaps those Python + bindings as it advances, so hoisting them makes eight shared-address-space + pointers loop-carried, which fails to legalize. + +And a FlyDSL gotcha: ``range(...)`` must appear literally in the ``for`` +statement or the AST rewriter does not see it and Python evaluates it eagerly +("dynamic 'ArithValue' has no Python integer representation"). Assigning it to a +variable first does not work, so a kernel cannot cheaply offer both a looped and +a flat form from one body. + +Reviving this needs 38 VGPRs from somewhere. The 3-stage C store is the only +change measured to *reduce* pressure (256 -> 254, and 2052 -> 1573 instructions, +since 128 two-byte stores need 128 addresses and 128 predicates live at once), +and it is nowhere near enough. Halving BLOCK_M would free roughly 16 by halving +the accumulator count, but moves the tile count to 896 -- remainder 128, where +gcnasm measured +0.41%. There is no version of this that pays. + +## Direct LSA (``--mode fused-lsa``): correct now, still slower than split + +gcnasm's best mode has the GEMM epilogue store *straight into the destination +rank's window*: no staging buffer, no copy engine, nothing to publish mid-kernel. +Ported here, the structural half is emphatic -- the scatter collapses from +136.8us to 10.1us, i.e. the transfer is entirely absorbed, which fused-sdma never +managed. And the un-coalesced store is nowhere near as bad as feared: 7.34MB per +link in 166.6us is ~44 GB/s, 82% of the SDMA scatter's 53.7, against the 0.26x +gcnasm measured for a lane-scatter pushed to a peer. + +It was wrong for a long time, and the fix is one line in the right place. The LSA +2-stage all-reduce publishes from the kernel that produced the data -- +``ar/kernels_lsa.py`` fences right after its ``tmp`` stores, in every block. +Direct LSA's producer is the GEMM, and the only fence was in the *separate* +barrier kernel: one block, therefore one XCD's L2 out of eight. The other seven +kept the peer-homed lines dirty, and the LSA flag, being a system-scope atomic, +overtook them. Moving the fence into the GEMM fixes it: **10/10 runs bit-correct** +where it had been 2-3 in 6. + +``--direct-fence leader`` (thread 0 only, the default) rather than every lane is +worth 80us, 434 -> 355. It is legal only because the half-wave barrier pair is +now closed -- ``wait_barrier(0)`` really does mean every wave's stores have +retired into this CU's L2, so one wave writing it back covers all eight. Measured +as incorrect before that fix, which is what made a per-wave release look +mandatory. + +Two things that did *not* work, and are worth not re-trying: + +* **Uncached (sc0|sc1) peer stores**, so there would be nothing to publish. Wrong + consistently, ~1.7e-2, at every store width. The first time this was tried the + store was one bf16 and the explanation looked like partial-line writes losing + updates over the fabric; with ``--permlane --lane-transpose`` making it 16 + bytes per lane and 64 contiguous per row it is *still* wrong, so that + explanation was not it and the real one is unknown. +* An uncached recv load in the reduce. No effect, which is what rules out a stale + read on the consumer side. + +It does not beat the split path. Per-kernel (rocprofv3, 8 ranks, [4096,7168] +K=1024, with all three C-store stages): + + fused-lsa gemm 189.4 barrier 10.1 reduce 10.9 gather 135.6 sum 346.0 + split-lsa gemm 35.7 + one LSA all-reduce kernel sum ~310 + + end to end, median of 4: split-lsa 318.8us fused-lsa 359.0us + +The 41us gap is structural, not tuning. LSA's 2-stage does read + reduce + write +in a single pass over the wire, so its reduce is free; Direct LSA writes to the +peer and then pays a separate 10.9us local reduce pass, its store rate is ~14% +under LSA's read rate, and it adds the publishing fence LSA gets from being +local. Absorbing the scatter completely still does not cover that. + +## What gcnasm does differently, and why its GEMM+a2a wins + +``/workspace/gcnasm/opus_gemm_dist/opus_gemm_a2a_lsa`` fuses a GEMM with an +all-to-all and gets 16-27% out of it. Reading it explains most of why this does +not, and one of its lessons was worth ~90us here. + +1. **Its collective is one phase; this one is three.** An a2a scatters the GEMM + output once. An all-reduce is scatter + reduce + all-gather, and fusion only + touches the scatter -- reduce + gather is 150us of the 327us baseline, 46%, + untouchable by construction. +2. **Its ratio is inverted.** M=2048 N=18432 K=8192 gives ~518us of GEMM against + ~200us of comm, 2.6:1. wo_b is 39us against 137us per phase, 0.29:1. Overlap + can hide at most the smaller of the two, so theirs hides most of the comm and + this hides at most one GEMM. +3. **Its best mode has no producer/consumer handoff at all.** "Direct LSA" has + the GEMM epilogue store *straight into the destination rank's buffer*. No + staging, no copy engine, nothing to publish mid-kernel -- the only sync is the + barrier at the end. The publication problem this file spends all its time on + simply does not exist there. +4. **Its fused-SDMA path uses no cache fence.** ``opus_chunk_sdma_submit`` is + ``s_waitcnt vmcnt(0)``, ``s_barrier``, an ``__ATOMIC_ACQ_REL`` counter, and a + per-destination submit spin lock -- no ``__threadfence_system``, no + ``buffer_wbl2``. That is the lesson that transferred: the acq_rel counter *is* + the release, and the explicit fence added here was 90us of pure waste + (fused 440us -> 350us on removing it). Its ISA shows the atomic already emits + its own ``buffer_wbl2``/``buffer_inv`` pair, on thread 0 only. + +## The chunks race: it was the GEMM, and it is gone + +``--chunks`` > 1 produced wrong output about 3 runs in 10 and was pinned to 1 +for that reason. The cause was not the chunk protocol at all: it was aiter's +8-wave GEMM under-counting one ``s_waitcnt`` in its main loop, which corrupted +output non-deterministically on large grids whatever the epilogue did (see the +comment at that ``wait_barrier`` below). Since that fix, at [16384, 7168] +K=2048 on 8 ranks: + +* chunks 1/2/4/8, 10 runs each -- 40/40 correct, no hang; +* chunks=8 alone, 25 more runs -- 25/25 correct, no hang. + +Two hangs were seen at 8 ranks while the sweep was still being set up and never +reproduced in the 100+ runs after. The submit lock is a plain test-and-set spin +(``_acquire_peer_lock``), so a hang is not impossible; use ``timeout`` when +sweeping and treat one as a finding rather than a flake. + +Two things in this file survived that misdiagnosis and are worth keeping +straight: + +* The half-wave barrier pairing (``if wave_m == 0: s_barrier()`` before + ``store_c``) is still needed and still right -- gcnasm does it + unconditionally at kernel_template.hpp:693. It took ``--chunks 2`` from + 1-in-3 failing to 3-in-10, which at the time read as "better but not fixed"; + the residual 3-in-10 was the GEMM. +* Two *other* diagnoses were wrong and are recorded so they are not retried. + **A shared SDMA queue**: gcnasm's per-destination submit lock was ported and + ISA-verified; it fixed nothing, and is kept only because cco's + one-issuing-warp-per-queue rule still applies. **The counter atomic's + ordering**: ``acq_rel`` appeared to beat ``monotonic``, which is how the + acquire half got justified; against a 1-in-3 intermittent failure that + comparison was noise. ``acq_rel`` stays because release/acquire is right for + a producer handing tiles to a consumer, not because it was measured -- and at + chunks=1 it demonstrably orders nothing, since 20 runs of ``monotonic`` pass + too. + +Repeated runs are still the only way to judge any of this, and the gate has to +be tight: the corruption landed at 4-9e-3 against an fp8 floor of 2.35e-3, so a +5e-3 threshold reported a corrupt run as validated (it did, at 3.97e-3). The +bench gates at 3e-3, and ``test_fused_is_stable_across_repeats`` requires three +runs to be *identical* rather than each small. + +## Result: fused-sdma wins, once the chunks are unblocked + +8 ranks, [16384, 7168] K=2048 -- the real prefill shape -- graph replay, all +three C-store stages on (now the benchmark default), median of 31, max over +ranks: + + mode time + fused-sdma chunks=8 1114.1us <- best + split-sdma 1261.7 + split-lsa 1262.4 + fused-lsa 1571.4 + gemm-only 228.9 + +Per-kernel, the overlap is visible directly: + + chunks GEMM drain reduce gather total + 1 257.5 493.8 44.0 496.4 1291.7 + 2 266.5 379.4 44.2 496.4 1186.6 + 8 265.0 303.7 43.8 496.6 1109.1 + +The drain falls 38% while the GEMM grows 7.5us for the extra counter atomics +and the lock. Going finer is worse: at chunks=8 each PUT is 3.5 MiB, and 16 +(via ``--block-m 128``) halves that to 1.75 MiB, under the knee in the SDMA +bandwidth curve, for 1130.8us. + +fused-lsa still loses, and for a reason that is not going away: it spends 730us +of its GEMM pushing C over xGMI where the copy engines move the same bytes in +499, and ATT shows 99% of that store time is *stall*, so coalescing the stores +(the three C-store stages) buys it nothing -- 954.5 -> 952.1us. + +The remaining floor is the tail: reduce plus all-gather is 540us of the 1114, +and neither is touched by fusing the GEMM. diff --git a/python/mori/ops/gemm_ar/__init__.py b/python/mori/ops/gemm_ar/__init__.py new file mode 100644 index 000000000..ea81b00a8 --- /dev/null +++ b/python/mori/ops/gemm_ar/__init__.py @@ -0,0 +1,75 @@ +# 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. +# +# Fused fp8 GEMM + all-reduce over cco SDMA. See README.md for the design and +# the measurements. +# +# `layout` is pure arithmetic and imports nothing, so it stays eagerly +# available; everything else needs FlyDSL (`pip install amd_mori[flydsl]`) and +# is imported lazily, so `import mori.ops.gemm_ar` works without it and only +# touching the op or the kernels raises. +import importlib + +from .layout import ArConfig, MAX_WORLD, ar_config, select_stage +from ._shuffle import preshuffle_b + +_LAZY = { + "GemmAllReduceOp": "op", + "counter_chunks": "op", + "padded_m": "op", + "supports": "op", + "DEFAULT_BLOCK_M": "op", + "DEFAULT_BLOCK_N": "op", + "MAX_CHUNKS": "op", + "SCALE_BLOCK_K": "op", + "compile_fused_gemm_scatter": "kernels_fused", + "build_sdma_phases": "kernels_sdma", + "build_sdma_ar": "kernels_sdma", + "build_lsa_ar": "kernels_lsa", +} + +__all__ = [ + "ArConfig", + "ar_config", + "GemmAllReduceOp", + "MAX_WORLD", + "build_lsa_ar", + "build_sdma_ar", + "build_sdma_phases", + "compile_fused_gemm_scatter", + "counter_chunks", + "padded_m", + "preshuffle_b", + "select_stage", + "supports", +] + + +def __getattr__(name: str): + module = _LAZY.get(name) + if module is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(f".{module}", __name__), name) + + +def __dir__(): + return sorted(set(globals()) | set(_LAZY)) diff --git a/python/mori/ops/gemm_ar/_compat.py b/python/mori/ops/gemm_ar/_compat.py new file mode 100644 index 000000000..ff60b4632 --- /dev/null +++ b/python/mori/ops/gemm_ar/_compat.py @@ -0,0 +1,387 @@ +# 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. +"""Buffer load/store with working cache policy, on flydsl 0.2.x and 0.3.x. + +flydsl 0.3.0 removed ``flydsl.expr.buffer_ops``. ``mori.ops.dispatch_combine_v2 +.flydsl_compat`` already bridges that for EPv2, but it documents that +``cache_modifier`` is *accepted and ignored* on 0.3.0. For EPv2 that is a hint; +here it is a correctness requirement: + +* the cross-GPU barrier spins on a slot a **peer** writes, so a cached load can + sit on a stale L2 line forever, and +* the signal store must be visible past L1+L2 or no peer ever leaves the spin. + +So the 0.3.0 path is built directly on ``rocdl.MakeBufferRsrcOp`` + +``rocdl.raw_ptr_buffer_{load,store}``, which do still carry an explicit ``aux`` +cache-policy operand. Only the handful of primitives the all-reduce needs are +implemented; ``aiter/ops/flydsl/kernels/buffer_ops.py`` is the full-featured +version of the same idea and was the reference for the descriptor encoding. + +Cache-policy (``aux``) bits on CDNA: SC0 = bypass L1, SC1 = bypass L2 -- the +same constants aiter's ``custom_all_reduce.cuh`` uses for ``start_sync`` / +``end_sync``. +""" + +from __future__ import annotations + +import inspect + +import flydsl +import flydsl.expr as fx +from flydsl._mlir import ir + +FLYDSL_VERSION = getattr(flydsl, "__version__", "unknown") + +# Cache-modifier (``aux``) values. This is the LLVM ``CPol`` operand, **not** the +# hardware bit positions, and the two do not line up -- which is worth stating +# because getting it wrong is silent. +# +# Established by assembling each policy and diffing the encoding:: +# +# echo 'buffer_store_short v0, v1, s[0:3], 0 offen ' \ +# | llvm-mc -triple=amdgcn-amd-amdhsa -mcpu=gfx950 -show-encoding +# +# pol encoding byte1/byte2 delta +# (none) 0x10 0x68 -- +# sc0 0x50 0x68 byte1 bit6 -> insn bit 14 +# sc1 0x90 0x68 byte1 bit7 -> insn bit 15 +# nt 0x10 0x6a byte2 bit1 -> insn bit 17 +# +# Insn bit 14 is GLC, 15 is SCC, 17 is SLC, and CPol spells those 1, 16, 2. So +# SC1 is **16**, not 4 -- 4 is DLC, which is RDNA-only and silently ignored on +# CDNA. An earlier version of this file had ``CM_SC1 = 4`` and a comment +# claiming aux=5 emitted ``sc0 sc1``; it emits ``sc0`` alone, leaving the line +# dirty in L2. Nothing depended on it (the barrier uses atomics and the payload +# stores use CM_CACHED), but two experiments were invalidated before it showed +# up in the ISA. +# +# Note also that ``fx.rocdl.BufferCopy*(cache_modifier)`` does **not** take these +# values: its modifier is a two-value enum (0=cached, 2=nt). Reaching sc0/sc1 +# needs ``buffer_store`` below, which passes aux through. +CM_CACHED = 0 +CM_SC0 = 1 # bypass L1 (GLC) +CM_NT = 2 # non-temporal (SLC) +CM_SC1 = 16 # bypass L2 -- required to observe a peer's fresh write (SCC) +CM_SC0_SC1 = CM_SC0 | CM_SC1 # 17: publish past L1+L2 + +# V# flags word for CDNA (gfx9xx): DATA_FORMAT=7, NUM_FORMAT=4. Bit 24 and +# OOB_SELECT are RDNA-only; this benchmark is gfx95x/gfx94x. +_CDNA_BUFFER_FLAGS = (7 << 12) | (4 << 15) +_NO_OOB_LIMIT = 0xFFFFFFFF + +# Needed on both sides of the split below -- `wave_uniform_i64` and the fp8 +# conversions use it whichever buffer-op surface is available -- so it is +# imported once here rather than inside a branch. It was in the 0.3 branch only, +# which made every 0.2.x run die at trace time on `_rocdl is not defined`. +from flydsl.expr import rocdl as _rocdl # noqa: E402 + +try: # flydsl <= 0.2.x + # ImportError, not ModuleNotFoundError: on 0.3.0 ``buffer_ops`` is neither a + # submodule nor a name in ``flydsl.expr``, so this form raises the parent. + from flydsl.expr import buffer_ops as _bo + from flydsl.expr.typing import T as _T + + HAS_BUFFER_OPS = True + + def _dtype(name): + return getattr(_T, name) + + create_buffer_resource_from_addr = _bo.create_buffer_resource_from_addr + buffer_load = _bo.buffer_load + buffer_store = _bo.buffer_store + +except ImportError: # flydsl >= 0.3.0 + from flydsl._mlir.dialects import arith as _arith + from flydsl._mlir.dialects import llvm as _llvm + from flydsl.expr.typing import T as _T + + HAS_BUFFER_OPS = False + + def _dtype(name): + # 0.3.0 turned T. from a factory into a property. + attr = getattr(_T, name) + return attr() if callable(attr) else attr + + def _raw(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + def _const(value, width): + ty = ir.IntegerType.get_signless(width) + return _arith.ConstantOp(ty, ir.IntegerAttr.get(ty, int(value))).result + + def _as_i32(v): + if isinstance(v, int): # literal element offsets, e.g. the signal slot + return _const(v, 32) + v = _raw(v) + ty = v.type + if isinstance(ty, ir.IndexType): + return _arith.IndexCastOp(ir.IntegerType.get_signless(32), v).result + if isinstance(ty, ir.IntegerType) and ty.width != 32: + i32 = ir.IntegerType.get_signless(32) + if ty.width > 32: + return _arith.TruncIOp(i32, v).result + return _arith.ExtSIOp(i32, v).result + return v + + def create_buffer_resource_from_addr(addr_i64, *, num_records_bytes=None): + """Raw i64 device address -> ``!llvm.ptr<8>`` buffer descriptor (V#).""" + base = _llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr"), _raw(addr_i64)).result + nrec = ( + _NO_OOB_LIMIT + if num_records_bytes is None + else max(0, min(int(num_records_bytes), _NO_OOB_LIMIT)) + ) + return _rocdl.MakeBufferRsrcOp( + ir.Type.parse("!llvm.ptr<8>"), + base, + _const(0, 16), # stride: raw (non-structured) buffer + _const(nrec, 64), + _const(_CDNA_BUFFER_FLAGS, 32), + ).result + + def _byte_offset(offset, elem_ty): + elem_bytes = elem_ty.width // 8 + return _arith.MulIOp(_as_i32(offset), _const(elem_bytes, 32)).result + + # FlyDSL 0.3.2 moved `aux` on the raw buffer ops from a positional i32 + # *value* to a keyword-only *attribute*. Both spellings are in use here -- + # aiter's venv has 0.3.0, sglang's has 0.3.2 -- so probe once and adapt + # rather than pinning either. + _AUX_IS_KWARG = ( + "aux" + in getattr(inspect.signature(_rocdl.raw_ptr_buffer_store), "parameters", {}) + and inspect.signature(_rocdl.raw_ptr_buffer_store).parameters["aux"].kind + is inspect.Parameter.KEYWORD_ONLY + ) + + def _aux_args(cache_modifier): + """Positional tail / kwargs for the raw buffer ops' ``aux`` operand.""" + if _AUX_IS_KWARG: + return (), { + "aux": ir.IntegerAttr.get( + ir.IntegerType.get_signless(32), int(cache_modifier) + ) + } + return (_const(cache_modifier, 32),), {} + + def buffer_load(rsrc, offset, vec_width=4, dtype=None, cache_modifier=CM_CACHED): + """Load ``vec_width`` x ``dtype`` at ELEMENT ``offset`` of ``rsrc``.""" + elem = dtype if dtype is not None else _dtype("f32") + if hasattr(elem, "ir_type"): + elem = elem.ir_type + res = elem if vec_width == 1 else ir.VectorType.get([vec_width], elem) + tail, kw = _aux_args(cache_modifier) + return _rocdl.raw_ptr_buffer_load( + res, rsrc, _byte_offset(offset, elem), _const(0, 32), *tail, **kw + ) + + def buffer_store(data, rsrc, offset, cache_modifier=CM_CACHED): + """Store ``data`` at ELEMENT ``offset`` of ``rsrc``.""" + value = _raw(data) + ty = value.type + try: # this binding has no VectorType.isinstance; construction is the probe + elem = ir.VectorType(ty).element_type + except (ValueError, TypeError): + elem = ty + tail, kw = _aux_args(cache_modifier) + return _rocdl.raw_ptr_buffer_store( + value, rsrc, _byte_offset(offset, elem), _const(0, 32), *tail, **kw + ) + + +# --- signal slot accessors ------------------------------------------------- +# +# These are **atomic**, not plain buffer accesses with a cache modifier, for two +# independent reasons -- both found by reading the emitted ISA: +# +# 1. A ``rocdl.raw.ptr.buffer.load`` is side-effect-free, so the spin loop that +# reloads it was dead-code-eliminated outright: the barrier compiled down to +# a signal store and nothing else. An atomic is never removed. +# 2. The syncscope tells the AMDGPU backend which cache levels to bypass, so we +# stop hand-encoding ``aux`` bits whose meaning changes between CDNA3 and +# CDNA4 (see the CM_* note above). +# +# Scopes mirror aiter's start_sync/end_sync: publish at **system** scope (the +# store must cross xGMI), poll at **agent** scope (we only read our own HBM). +from flydsl._mlir.dialects import llvm as _llvm_d # noqa: E402 + +_SYSTEM_SCOPE = "" # LLVM AMDGPU: the empty syncscope is system-wide +_AGENT_SCOPE = "agent" + + +def wave_uniform_i64(addr): + """Force a 64-bit address into scalar registers. + + ``Window.lsa_ptr`` is an opaque extern call, so even for a compile-time peer + index the compiler cannot prove its result is wave-uniform. Building a + buffer descriptor from a value it believes divergent makes the backend wrap + **every** load and store in a readfirstlane waterfall loop, serialising each + 16-byte access lane by lane -- worth ~2.5x on this kernel, visible in the ISA + as ``s_cbranch_execnz`` around each ``buffer_load_dwordx4``. + + Same idiom as aiter's + ``kernels/flydsl_dispatch_combine_intranode_kernel.py::_wave_uniform_i64``. + """ + v = fx.Uint64(addr) + lo = _rocdl.readfirstlane(_dtype("i32"), fx.Uint32(v)) + hi = _rocdl.readfirstlane(_dtype("i32"), fx.Uint32(v >> 32)) + return (fx.Uint64(hi) << 32) | fx.Uint64(lo) + + +#: AMDGPU's LLVM address space for global memory. `fx.to_llvm_ptr` would +#: resolve this from the active backend, but it does not exist before flydsl +#: 0.3 and the CI image pins `flydsl<0.3` (the cco Python examples still import +#: `flydsl.expr.buffer_ops`, which 0.3 dropped). `llvm.inttoptr` is identical in +#: both -- same signature, same dialect module -- so going straight to an +#: `!llvm.ptr` keeps one code path instead of a version probe. +_GLOBAL_LLVM_AS = 1 +_LLVM_GLOBAL_PTR = f"!llvm.ptr<{_GLOBAL_LLVM_AS}>" + + +def signal_ptr(addr_i64): + """Raw i64 device address -> ``!llvm.ptr`` over one u32 signal slot. + + The llvm dialect ops below (atomic store/load/cmpxchg) need a lowered + ``!llvm.ptr``, not the ``!fly.ptr`` that ``fly``'s own ``inttoptr`` yields. + """ + return _llvm_d.inttoptr( + ir.Type.parse(_LLVM_GLOBAL_PTR), fx.Int64(addr_i64).ir_value() + ) + + +def signal_store_u32(ptr, value): + _llvm_d.store( + fx.Int32(value).ir_value(), + ptr, + alignment=4, + ordering=_llvm_d.AtomicOrdering.monotonic, + syncscope=_SYSTEM_SCOPE, + ) + + +def signal_load_u32(ptr): + return _llvm_d.load( + _dtype("i32"), + ptr, + alignment=4, + ordering=_llvm_d.AtomicOrdering.monotonic, + syncscope=_AGENT_SCOPE, + ) + + +# The per-block ``_flag`` counter is private to this rank: only our own block +# reads and writes it, so it needs no atomic and no scope. Making it atomic had +# all 512 lanes issue an agent-scope atomic against one address every launch, +# which measurably serialised the prologue. ``volatile`` is still required so the +# load is not hoisted above the barrier that publishes the previous value. +def local_load_u32(ptr): + return _llvm_d.load(_dtype("i32"), ptr, alignment=4, volatile_=True) + + +def local_store_u32(ptr, value): + _llvm_d.store(fx.Int32(value).ir_value(), ptr, alignment=4, volatile_=True) + + +def atomic_add_u32(ptr, value, *, ordering="monotonic"): + """Device-scope ``fetch_add``; returns the value *before* the add. + + Agent scope, because the tile counters it serves are only ever touched by + blocks of the same kernel on this GPU -- a system-scope RMW would be a fabric + round trip per tile. + + ``monotonic`` (relaxed), not ``acq_rel``, and that is worth 5% of the fused + GEMM. An ordered RMW is bracketed by its own ``buffer_wbl2`` / + ``buffer_inv``, and a thread trace showed those costing 91.6k cycles -- more + per wave than the explicit release fence itself. Both halves are redundant + here: the caller has already issued that release, and the winner never reads + the data it is counting (the copy engine does), so it needs no acquire. Pass + ``ordering="acq_rel"`` if a caller ever does need the fence. + """ + return _llvm_d.atomicrmw( + _llvm_d.AtomicBinOp.add, + ptr, + fx.Int32(value).ir_value(), + getattr(_llvm_d.AtomicOrdering, ordering), + syncscope=_AGENT_SCOPE, + alignment=4, + ) + + +def atomic_xchg_u32(ptr, value, *, ordering="acquire"): + """Device-scope exchange; returns the previous value. Used as a test-and-set.""" + return _llvm_d.atomicrmw( + _llvm_d.AtomicBinOp.xchg, + ptr, + fx.Int32(value).ir_value(), + getattr(_llvm_d.AtomicOrdering, ordering), + syncscope=_AGENT_SCOPE, + alignment=4, + ) + + +def atomic_store_u32(ptr, value, *, ordering="release"): + _llvm_d.store( + fx.Int32(value).ir_value(), + ptr, + alignment=4, + ordering=getattr(_llvm_d.AtomicOrdering, ordering), + syncscope=_AGENT_SCOPE, + ) + + +def release_fence(scope="agent"): + """Release fence at an explicit scope, for producers of copy-engine input. + + cco only exposes ``cco_system_fence`` (``__threadfence_system``), which on + gfx950 writes back through to memory. An SDMA engine reading a buffer we just + wrote lives on the *same* device, so agent scope is the scope that actually + describes the handoff; system scope over-synchronises against host and peers + that are not the reader here. + """ + _llvm_d.fence(_llvm_d.AtomicOrdering.release, syncscope=scope) + + +def i32_type(): + return _dtype("i32") + + +__all__ = [ + "FLYDSL_VERSION", + "HAS_BUFFER_OPS", + "buffer_load", + "buffer_store", + "create_buffer_resource_from_addr", + "signal_ptr", + "wave_uniform_i64", + "signal_store_u32", + "signal_load_u32", + "local_load_u32", + "local_store_u32", + "atomic_add_u32", + "release_fence", + "atomic_xchg_u32", + "atomic_store_u32", + "i32_type", + "CM_CACHED", + "CM_SC1", + "CM_SC0_SC1", +] diff --git a/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py b/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py new file mode 100644 index 000000000..c4e09cf93 --- /dev/null +++ b/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py @@ -0,0 +1,625 @@ +# 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. +# +# This file is vendored from aiter, where it carries the notice below. It is +# reproduced here unchanged as the upstream attribution; the MIT grant above is +# mori's, and does not restate or replace it. +# +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +# Modifications Copyright (C) 2026 Advanced Micro Devices, Inc. + +"""8-wave FP8 matmul with row-wise scaling for AMD CDNA4 (gfx950 only). + +Vendored from aiter ``aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py``. +``kernels_fused.py`` builds its fused GEMM out of these pieces, and mori does +not depend on aiter, so the file is copied rather than imported. + +Verbatim except for one thing: ``split_row_major_2d`` is inlined below instead +of imported from ``mfma_preshuffle_pipeline`` -- three lines against a +1230-line module. + +Note that ``kernels_fused.py`` does *not* call the pipeline in this file. It +re-implements the main loop so it can fuse a scatter into the epilogue, and in +doing so corrects a ``s_waitcnt`` count that is wrong here and still wrong +upstream (see the comment at that site). The classes below -- the loaders, the +MMA atom, ``StoreC`` -- are what it reuses. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace +from flydsl.expr import arith, const_expr, range_constexpr, rocdl +from flydsl.expr.typing import Vector as Vec + + +def split_row_major_2d(index, minor_extent): + """Split a linear row-major index into (major, minor). + + Inlined from aiter ``mfma_preshuffle_pipeline.split_row_major_2d``. + """ + return index // minor_extent, index % minor_extent + + +def ceildiv(numer, denom): + """Ceiling division; works on Python ints and on DSL scalars alike. + + Unlike an int-only ``cdiv``, this one is also applied to runtime + ``fx.Int32`` values, through operator overloading. + """ + return (numer + denom - 1) // denom + + +def _min(a, b): + return arith.select(a < b, a, b) + + +def make_fp8_buffer_tensor(arg_i8, fp8_ir_t): + # max_size=False with no num_records_bytes: cosize(layout) becomes a + # runtime expression because TensorAdaptor defaults to layout-dynamic + # memref (post #554), so the descriptor adapts to the actual tensor + # extent and no longer bakes the first-call's shape into IR. + t_i8 = fx.rocdl.make_buffer_tensor(arg_i8, max_size=False) + iter_i8 = fx.get_iter(t_i8) + f8_buf_ptr_ty = fx.PointerType.get( + elem_ty=fp8_ir_t, + address_space=TargetAddressSpace.BufferDesc, + alignment=fx.PointerType(iter_i8.type).alignment, + ) + iter_f8 = fx.recast_iter(f8_buf_ptr_ty, iter_i8) + return fx.Tensor(fx.make_view(iter_f8, fx.get_layout(t_i8))) + + +def swizzle_128(row, col): + offset = row * 128 + col + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def compute_global_swizzle(lane_id, wave_id, K, n_rounds, preshuffled): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + row = lane_id % 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id // 8) * 16 + offsets.append( + (row // 16) * (K * 16) + + (row % 16) * 16 + + (col // 64) * 1024 + + ((col % 64) // 16) * 256 + + (col % 16) + ) + else: + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + r, c = swizzle_128(row, col) + offsets.append(r * K + c) + return offsets + + +def pack_i32x4_i32x8(lo, hi): + # Pack two i32x4 as one i32x8 + return lo.shuffle(hi, list(range(8))) + + +def wait_barrier(count): + _llvm.inline_asm( + res=None, + operands_=[], + asm_string=f"s_waitcnt vmcnt({count})\ns_barrier", + constraints="", + has_side_effects=True, + ) + + +class G2SLoader: + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + sum_i32 = base_i32 + fx.Int32(step_off) + lds_ptr = fx.inttoptr(self.LdsPtr_t, sum_i32) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, k_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + def load_one(self, lds_dst, k_offset, step): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + +class S2RLoader: + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16xf8(self, lds_src, offset): + off_tup = fx.make_int_tuple(offset) + ptr_off = fx.add_offset(lds_src.ptr, off_tup) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + view = fx.make_view(i8_iter, fx.make_layout(16, 1)) + return view.load() + + def load(self, lds_src, preshuffled=False): + frag = [] + for i in range_constexpr(self.n_tiles): + halves = [] + row = self.wave_idx * (self.n_tiles * 16) + i * 16 + self.lane_id % 16 + for step in range_constexpr(2): + col = (self.lane_id // 16) * 16 + step * 64 + if const_expr(preshuffled): + offset = (row // 8) * 1024 + (row % 8) * 16 + (col // 16) * 128 + else: + row_swz, col_swz = swizzle_128(row, col) + offset = row_swz * 128 + col_swz + v = self._vec_load_16xf8(lds_src, offset) + halves.append(v.bitcast(fx.Int32)) + frag.append(pack_i32x4_i32x8(halves[0], halves[1])) + return frag + + def load_one(self, lds_src, lds_offset): + v = self._vec_load_16xf8(lds_src, lds_offset) + return v.bitcast(fx.Int32) + + +class StoreC: + def __init__( + self, A_scale, B_scale, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b + ): + self.c_rows = c_rows + self.c_cols = c_cols + self.lane_id = fx.thread_idx.x % 64 + self.c_idx_fn = c_idx_fn + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + # Exact byte counts from compile-time shape (BF16 C output, FP32 scales). + # ``num_records_bytes`` is required when ``max_size=False`` -- see + # ``make_buffer_tensor`` docstring for the silent-OOB rationale. + c_nbytes = c_rows * c_cols * 2 # BFloat16 = 2 bytes + sa_nbytes = c_rows * 4 # Float32 row-wise scale + sb_nbytes = c_cols * 4 # Float32 col-wise scale + gC = fx.rocdl.make_buffer_tensor(C, max_size=False, num_records_bytes=c_nbytes) + gSA = fx.rocdl.make_buffer_tensor( + A_scale, max_size=False, num_records_bytes=sa_nbytes + ) + gSB = fx.rocdl.make_buffer_tensor( + B_scale, max_size=False, num_records_bytes=sb_nbytes + ) + self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + self.sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) + self.sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + + self.scale_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) + self.scale_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), fx.BFloat16) + self.reg_f32_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) + self.reg_f32_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + self.reg_bf16_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.BFloat16) + + def _load_scale_vec4(self, row): + fx.copy( + self.scale_atom_4, + fx.slice(self.sa_div, (None, fx.Int32(row))), + self.reg_f32_4, + ) + return Vec(fx.memref_load_vec(self.reg_f32_4)) + + def _load_scale_scalar(self, col): + fx.copy( + self.scale_atom_1, + fx.slice(self.sb_div, (None, fx.Int32(col))), + self.reg_f32_1, + ) + return Vec(fx.memref_load_vec(self.reg_f32_1))[0] + + def _store_bf16(self, value_bf16, c_index): + fx.memref_store_vec(Vec.filled(1, value_bf16, fx.BFloat16), self.reg_bf16_1) + fx.copy( + self.out_atom_1, + self.reg_bf16_1, + fx.slice(self.c_div, (None, fx.Int32(c_index))), + ) + + def store(self, c_frag, base_row, base_col): + a_scales = [ + self._load_scale_vec4(base_row + i * 16 + (self.lane_id // 16) * 4) + for i in range_constexpr(self.n_tiles_a) + ] + b_scales = [ + self._load_scale_scalar(base_col + i * 16 + self.lane_id % 16) + for i in range_constexpr(self.n_tiles_b) + ] + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + (self.lane_id // 16) * 4 + for tj in range_constexpr(self.n_tiles_b): + col = base_col + tj * 16 + self.lane_id % 16 + col_valid = col < self.c_cols + oob = fx.Int32(self.c_rows * self.c_cols) + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + for i in range_constexpr(4): + scaled = (vec_f32[i] * (a_scales[ti][i] * b_scales[tj])).to( + fx.BFloat16 + ) + c_index = (row + i) * self.c_cols + col + self._store_bf16(scaled, arith.select(col_valid, c_index, oob)) + + +class Mfma16x16x128: + def __init__(self, n_tiles_a, n_tiles_b): + self.atom = fx.make_mma_atom( + fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, fx.Float8E4M3FN) + ) + self.zero_value = Vec.filled(4, 0.0, fx.Float32) + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + + def idx(self, i, j): + return i * self.n_tiles_b + j + + def _make_operand_frag(self, value): + frag = fx.make_rmem_tensor(8, fx.Int32) + frag.store(Vec(value)) + return frag + + def _make_accum_frag(self, value): + frag = fx.make_rmem_tensor(4, fx.Float32) + frag.store(Vec(value)) + return frag + + def _do_mma(self, a, b, c): + a_frag = self._make_operand_frag(a) + b_frag = self._make_operand_frag(b) + c_frag = self._make_accum_frag(c) + fx.gemm(self.atom, c_frag, a_frag, b_frag, c_frag) + return c_frag.load().ir_value() + + def call(self, a, b, c, *, set_prio=True): + assert len(a) == self.n_tiles_a + assert len(b) == self.n_tiles_b + assert len(c) == self.n_tiles_a * self.n_tiles_b + + a_frags = [ + self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a) + ] + b_frags = [ + self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b) + ] + c_frags = [ + self._make_accum_frag(c[idx]) + for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b) + ] + if const_expr(set_prio): + rocdl.s_setprio(1) + for i in range_constexpr(self.n_tiles_a): + for j in range_constexpr(self.n_tiles_b): + cf = c_frags[self.idx(i, j)] + fx.gemm(self.atom, cf, a_frags[i], b_frags[j], cf) + if const_expr(set_prio): + rocdl.s_setprio(0) + rocdl.s_barrier() + return [ + c_frags[idx].load().ir_value() + for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b) + ] + + def call_one(self, a, b, c, i, j): + assert i < self.n_tiles_a and j < self.n_tiles_b + + return self._do_mma(a[i], b[j], c[self.idx(i, j)]) + + +def _xcd_swizzle_any(num_pid_m, num_pid_n, wgm): + """XCD-aware tile remap that works for any workgroup count. This function is used to swizzle the tile indices to improve the L2 reuse.""" + NUM_XCDS = 8 + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Hardware dispatches workgroup i to XCD (i % NUM_XCDS). Invert that so + # logically adjacent tiles share an XCD, and thus its L2. + xcd = wgid % NUM_XCDS + intra_xcd = wgid // NUM_XCDS + base = num_wg // NUM_XCDS + extra = num_wg % NUM_XCDS + wgid_remap = xcd * base + _min(xcd, extra) + intra_xcd + + # Group-of-WGM row swizzle for L2 reuse within one XCD. + num_wgid_in_group = wgm * num_pid_n + group_id, intra_group = split_row_major_2d(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * wgm + group_size_m = _min(num_pid_m - first_pid_m, wgm) + pid_n, intra_group_m = split_row_major_2d(intra_group, group_size_m) + return first_pid_m + intra_group_m, pid_n + + +def compile_fp8_gemm_8w( + *, + K: int, + BLOCK_M: int = 256, + BLOCK_N: int = 256, + b_preshuffled: bool = False, + waves_per_eu: int = 2, + xcd_swizzle: int = 0, +): + BLOCK_K = 128 + + assert ( + BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + ) + assert K % BLOCK_K == 0 + + K_ITERS = K // BLOCK_K + + N_TILES_A = BLOCK_M // 64 + N_TILES_B = BLOCK_N // 128 + N_ACCUMS = N_TILES_A * N_TILES_B + assert N_ACCUMS > 0 + + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + # half size + a_lds_size = LDS_BLOCK_M * BLOCK_K + b_lds_size = LDS_BLOCK_N * BLOCK_K + + _layout_tag = "bpreshuffle" if b_preshuffled else "rowmajor" + _kname = ( + f"flydsl_{_layout_tag}_8w_{BLOCK_M}x{BLOCK_N}x{BLOCK_K}_F8_F8_B16_" + f"{waves_per_eu}x{xcd_swizzle}_k{K}" + ) + + @fx.struct + class SharedStorage: + A_lds_cur_0: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + A_lds_cur_1: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + A_lds_next_0: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + A_lds_next_1: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + B_lds_cur_0: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + B_lds_cur_1: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + B_lds_next_0: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + B_lds_next_1: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + + @flyc.kernel(name=_kname, known_block_size=[512, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B_T: fx.Tensor, + C: fx.Tensor, + A_scale: fx.Tensor, + B_scale: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + F8_IR_t = fx.Float8E4M3FN.ir_type + + n_blocks = ceildiv(c_n, BLOCK_N) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + a_cur0 = lds.A_lds_cur_0 + a_cur1 = lds.A_lds_cur_1 + a_next0 = lds.A_lds_next_0 + a_next1 = lds.A_lds_next_1 + b_cur0 = lds.B_lds_cur_0 + b_cur1 = lds.B_lds_cur_1 + b_next0 = lds.B_lds_next_0 + b_next1 = lds.B_lds_next_1 + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + if const_expr(xcd_swizzle > 0): + block_m, block_n = _xcd_swizzle_any( + ceildiv(c_m, BLOCK_M), n_blocks, wgm=xcd_swizzle + ) + else: + block_m, block_n = split_row_major_2d(fx.block_idx.x, n_blocks) + + A0_gl_offset = (block_m * BLOCK_M) * K + A1_gl_offset = (block_m * BLOCK_M + LDS_BLOCK_M) * K + B_K_STEP = (2 * 1024) if b_preshuffled else BLOCK_K + B0_gl_offset = (block_n * BLOCK_N) * K + B1_gl_offset = (block_n * BLOCK_N + LDS_BLOCK_N) * K + + gA = make_fp8_buffer_tensor(A, F8_IR_t) + gB = make_fp8_buffer_tensor(B_T, F8_IR_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + + gl_off_a = compute_global_swizzle( + lane_id, wave_id, K, N_LDS_ROUNDS, preshuffled=False + ) + gl_off_b = compute_global_swizzle( + lane_id, wave_id, K, N_LDS_ROUNDS, preshuffled=b_preshuffled + ) + + mfma = Mfma16x16x128(N_TILES_A, N_TILES_B) + + a_g2s = G2SLoader(a_div, gl_off_a, N_LDS_STEPS_A, F8_IR_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, F8_IR_t, wave_id) + a_s2r = S2RLoader(wave_m, N_TILES_A) + b_s2r = S2RLoader(wave_n, N_TILES_B) + store_c = StoreC(A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B) + + # 2x2 config of 4x2 (instead of 4x4 in 4wave) 16x16 sub-tiles + c00_frag = [mfma.zero_value] * N_ACCUMS + c01_frag = [mfma.zero_value] * N_ACCUMS + c10_frag = [mfma.zero_value] * N_ACCUMS + c11_frag = [mfma.zero_value] * N_ACCUMS + + b_g2s.load(b_cur0, B0_gl_offset + 0 * B_K_STEP) + a_g2s.load(a_cur0, A0_gl_offset + 0 * BLOCK_K) + b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) + a_g2s.load(a_cur1, A1_gl_offset + 0 * BLOCK_K) + + if wave_m == 1: + rocdl.s_barrier() + + wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B) + + b_g2s.load(b_next0, B0_gl_offset + 1 * B_K_STEP) + a_g2s.load(a_next0, A0_gl_offset + 1 * BLOCK_K) + b_g2s.load(b_next1, B1_gl_offset + 1 * B_K_STEP) + + wait_barrier(N_LDS_STEPS_A + 2 * N_LDS_STEPS_B) + + for k in range_constexpr(K_ITERS - 2): + b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) + a0_frag = a_s2r.load(a_cur0) + a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) + rocdl.s_barrier() + + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + + b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) + b_g2s.load(b_cur0, B0_gl_offset + (k + 2) * B_K_STEP) + rocdl.s_barrier() + + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + + a1_frag = a_s2r.load(a_cur1) + a_g2s.load(a_cur0, A0_gl_offset + (k + 2) * BLOCK_K) + rocdl.s_barrier() + + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + + b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * B_K_STEP) + wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) + + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + + # Swap cur and next + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + # Step k = K_ITERS - 2 + k = K_ITERS - 2 + b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) + a0_frag = a_s2r.load(a_cur0) + rocdl.s_barrier() + + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + + b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) + rocdl.s_barrier() + + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + + a1_frag = a_s2r.load(a_cur1) + # Main loop prefetches a_next1 one step behind; issue the final + # K_ITERS - 1 tile here, otherwise c10 / c11 read stale A1 data. + a_g2s.load(a_next1, A1_gl_offset + (K_ITERS - 1) * BLOCK_K) + rocdl.s_barrier() + + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + + b0_frag = b_s2r.load(b_next0, preshuffled=b_preshuffled) + rocdl.s_barrier() + + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + # Swap cur and next + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + # Step k = K_ITERS - 1 + k = K_ITERS - 1 + a0_frag = a_s2r.load(a_cur0) + wait_barrier(0) + + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + + b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) + rocdl.s_barrier() + + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + + a1_frag = a_s2r.load(a_cur1) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, set_prio=False) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, set_prio=False) + rocdl.s_setprio(0) + rocdl.s_barrier() + + # Scale and store back to gmem + wave_n_offset = wave_n * (N_TILES_B * 16) + wave_m_offset = wave_m * (N_TILES_A * 16) + base_row = block_m * BLOCK_M + wave_m_offset + base_col = block_n * BLOCK_N + wave_n_offset + + store_c.store(c00_frag, base_row + 0, base_col + 0) + store_c.store(c01_frag, base_row + 0, base_col + LDS_BLOCK_N) + store_c.store(c10_frag, base_row + LDS_BLOCK_M, base_col + 0) + store_c.store(c11_frag, base_row + LDS_BLOCK_M, base_col + LDS_BLOCK_N) + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B_T: fx.Tensor, + C: fx.Tensor, + A_scale: fx.Tensor, + B_scale: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream, + ): + grid_x = ceildiv(c_m, BLOCK_M) * ceildiv(c_n, BLOCK_N) + kernel_gemm( + A, + B_T, + C, + A_scale, + B_scale, + c_m, + c_n, + value_attrs={ + "rocdl.waves_per_eu": waves_per_eu, + "rocdl.flat_work_group_size": "512,512", + }, + ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) + + return launch_gemm diff --git a/python/mori/ops/gemm_ar/_shuffle.py b/python/mori/ops/gemm_ar/_shuffle.py new file mode 100644 index 000000000..b5fb20349 --- /dev/null +++ b/python/mori/ops/gemm_ar/_shuffle.py @@ -0,0 +1,65 @@ +# 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. +"""Put B in the layout the 8-wave GEMM's ``b_preshuffled=True`` path reads. + +Equivalent to aiter's ``shuffle_weight(x, layout=(16, 16))`` for the one case +this op needs -- a 2-D fp8 weight, no int4, no gate/up interleave -- vendored +so mori does not depend on aiter. +""" + +import torch + + +def preshuffle_b(w: torch.Tensor) -> torch.Tensor: + """Reorder ``[N, K]`` into the MFMA 16x16x128 preshuffled B layout. + + The emitted order is ``N0 K0 KLane(4) NLane(16) KPack(16)``: lane ``l`` of a + wave reads ``n1 = l % 16`` and ``k = (l // 16) * 16 + [0..15]``, so a whole + 16x128 B tile is one ``buffer_load_dwordx4`` per lane. This is the same byte + order CK's ``preShuffleBuffer`` emits with ``NXdl=16``, which is why a + weight shuffled here can be handed to either kernel. + + Shape is preserved; only the element order changes. ``K`` counts *elements*, + so for fp8 (1 byte) the ``BK = 32`` and ``KPack = 16`` below are in + elements too. + """ + if w.ndim != 2: + raise ValueError(f"expected a 2-D [N, K] weight, got {w.ndim}-D") + dtype = w.dtype + n, k = w.shape + if w.element_size() != 1: + raise ValueError( + f"expected a 1-byte element type (fp8), got {dtype} " + f"({w.element_size()} bytes)" + ) + # NLane = 16 lanes over N, KLane = 64 // NLane = 4 lanes over K, and + # KPack = 16 elements per lane -- so a lane group spans BK = KLane * KPack. + n_lane, k_pack = 16, 16 + bk = 64 // n_lane * k_pack + if n % n_lane: + raise ValueError(f"N={n} must be a multiple of {n_lane}") + if k % bk: + raise ValueError(f"K={k} must be a multiple of {bk}") + + out = w.view(n // n_lane, n_lane, k // bk, bk // k_pack, k_pack) + out = out.permute(0, 2, 3, 1, 4).contiguous() + return out.view(n, k).view(dtype) diff --git a/python/mori/ops/gemm_ar/kernels_fused.py b/python/mori/ops/gemm_ar/kernels_fused.py new file mode 100644 index 000000000..08fb2874d --- /dev/null +++ b/python/mori/ops/gemm_ar/kernels_fused.py @@ -0,0 +1,1582 @@ +# 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. +"""The 8-wave fp8 GEMM with the all-reduce's scatter fused into its epilogue. + +Built for DeepSeek-V4-Pro's ``wo_b``: a RowParallelLinear whose per-rank GEMM is +``[M,K] x [N,K]`` fp8 -> ``[M,N]`` bf16, immediately all-reduced. Split, that is +``gemm(); all_reduce()``. Fused, the GEMM's C lands directly in a registered cco +window and each destination's slice is pushed by the copy engine as soon as its +last tile is written, so the reduce-scatter transfer overlaps the rest of the +GEMM instead of following it. + +Only the *scatter* half of the all-reduce is absorbed. The reduce and all-gather +phases still run as their own kernels, reused verbatim from ``kernels_sdma`` +(``build_sdma_phases``), with ``scatter`` swapped for its drain-only twin. + +``README.md`` carries the design narrative and the measurements: why SDMA needs +no epilogue change, the completion protocol, the three C-store stages, and what +was tried and rejected. This docstring covers only what a caller must know. + +## Completion protocol + +One monotonic counter per (destination, chunk) in the window +(``cfg.counter_off``). Every block, after its four ``store_c.store`` calls:: + + s_waitcnt vmcnt(0) ; s_barrier -- this block's C tile has retired + __threadfence_system() -- ...and is visible to the copy engine + thread 0: prev = atomic_add(counter[dest][chunk], 1) + if (prev + 1) % tiles_per_chunk == 0 -- I am the last tile of that chunk + sdma.put(dest, ...) -- fire and forget, no quiet + +Counters are never reset, so the modulo test works on every launch and the +kernel stays CUDA-graph-safe, exactly like the barrier flags in +``kernels_lsa``. + +## Relationship to the vendored pipeline + +``_fused_kernel_body`` is a copy of ``compile_fp8_gemm_8w``'s ``kernel_gemm`` +from ``_gemm_a8w8_8wave.py`` with the epilogue tail added. It is a copy rather +than a hook because the original's main loop contains a runtime +``if wave_m == 1``, and FlyDSL only rewrites ``if`` inside the +``@flyc.kernel`` function's own AST -- factoring the body into a shared helper +makes that line raise at trace time. Every reusable piece (``G2SLoader``, +``S2RLoader``, ``StoreC``, ``Mfma16x16x128``, the swizzles) is imported rather +than copied, so the duplication is the ~90-line pipeline only. + +The copy carries **one deliberate divergence**, marked at its site in the main +loop: the final ``wait_barrier`` count. The upstream count is too permissive and +corrupts the output non-deterministically on large grids. Do not "restore" it -- +see the comment at that line. The bug is still live in aiter. +""" + +# NOTE: no `from __future__ import annotations` here, deliberately, and +# _gemm_a8w8_8wave.py omits it for the same reason: `fx.struct` reads the +# `SharedStorage` field annotations as live objects, and PEP 563 would hand it +# the string "fx.Array[fx.Float8E4M3FN, a_lds_size, 16]" -- whose size operands +# are locals of this factory and so cannot be resolved from the module globals. + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm_d +from flydsl._mlir.dialects import scf +from flydsl.expr import arith, const_expr, range_constexpr, rocdl +from flydsl.expr import gpu as fgpu +from flydsl.expr.typing import Int64 +from flydsl.expr.typing import Vector as Vec + +import mori.cco.device.flydsl as cco +from mori.cco.device.flydsl import _bindings as raw_cco + +from ._compat import ( + CM_CACHED, + CM_SC0_SC1, + atomic_add_u32, + atomic_store_u32, + atomic_xchg_u32, + buffer_store, + create_buffer_resource_from_addr, + i32_type, + release_fence, + signal_ptr, + wave_uniform_i64, +) +from ._gemm_a8w8_8wave import ( + G2SLoader, + Mfma16x16x128, + S2RLoader, + StoreC, + _xcd_swizzle_any, + ceildiv, + compute_global_swizzle, + make_fp8_buffer_tensor, + split_row_major_2d, + wait_barrier, +) + + +def _acquire_peer_lock(lock_ptr): + """Test-and-set spin on one destination's submit lock. + + Transferred from gcnasm ``opus_chunk_sdma_submit``: with more than one chunk + per destination the tile counter elects one block per *chunk*, so two of them + can reach the SDMA submit for the same queue at once. Serialising only the + submit -- not the GEMM work -- is the point. + + Termination: the contending blocks are elected in chunk order, and chunk c's + tiles are dispatched before chunk c+1's, so a spinning block is always + waiting on one that has already run. + """ + i32 = i32_type() + first = atomic_xchg_u32(lock_ptr, 1) + loop = scf.WhileOp([i32], [first]) + cond = ir.Block.create_at_start(loop.before, [i32]) + body = ir.Block.create_at_start(loop.after, [i32]) + with ir.InsertionPoint(cond): + held = fx.Uint32(fx.Int32(cond.arguments[0])) > fx.Uint32(fx.Int32(0)) + scf.ConditionOp(held.ir_value(), [cond.arguments[0]]) + with ir.InsertionPoint(body): + fx.rocdl.s_sleep(1) + nxt = atomic_xchg_u32(lock_ptr, 1) + scf.YieldOp([nxt]) + + +class _RawWriteThroughStoreC(StoreC): + """``StoreC`` whose C stores really do bypass L1 *and* L2. + + ``StoreC`` stores through a copy atom whose ``cache_modifier`` is a two-value + enum (0=cached, 2=nt), so ``sc0|sc1`` is not expressible there -- asking for + it emits a plain ``sc0`` and leaves the line dirty in L2, which is why the + earlier ``--fence writethrough`` probe measured nothing. This goes around the + atom to ``raw_ptr_buffer_store``, whose ``aux`` operand is the real cache + policy, so the data reaches memory at the store and the per-block + ``buffer_wbl2`` release becomes unnecessary. + + The descriptor is built from the window rather than from the ``C`` argument + because in the fused kernel they are the same bytes: C *is* the all-reduce's + input region. ``num_records`` still bounds it, so ``store``'s out-of-bounds + index for masked columns is dropped exactly as before. + + Deliberately *not* coalesced: the point of this variant is to find out + whether the coalescing rewrite is needed before paying for it. + """ + + def __init__(self, *args, c_rsrc=None, **kwargs): + super().__init__(*args, **kwargs) + self._c_rsrc = c_rsrc + + def _store_bf16(self, value_bf16, c_index): + buffer_store(value_bf16, self._c_rsrc, c_index, cache_modifier=CM_SC0_SC1) + + +class _SwappedMfma: + """``Mfma16x16x128`` with the A/B operands exchanged at every call. + + Thin on purpose: the accumulator registers, the atom and ``zero_value`` are + the wrapped object's, only the operand order and the ``idx`` argument order + change. ``idx(ti, tj)`` here forwards to the wrapped ``idx(tj, ti)`` so the + store can keep indexing in (M-tile, N-tile) order. + """ + + def __init__(self, inner): + self._inner = inner + self.zero_value = inner.zero_value + + def idx(self, i, j): + return self._inner.idx(j, i) + + def call(self, a, b, c, *, set_prio=True): + return self._inner.call(b, a, c, set_prio=set_prio) + + +class _BlockScaleK: + """The model's 1x128 / 128x128 fp8 block scales, applied per K-block. + + ``sum_k s_a[m,k] * s_b[nb,k] * d_k`` does not factor, so unlike the + per-token/per-channel form these scales cannot wait for the epilogue. + ``BLOCK_K`` is already 128, so a mainloop iteration *is* a scale block and + there is no sub-block bookkeeping. + + The obvious form -- promote each K-block into a second fp32 accumulator and + zero the MFMA one -- does not work here, and the reason is worth recording. + Zeroing removes the only dependency between consecutive K-blocks, and the K + loop is ``range_constexpr`` (fully unrolled), so the scheduler hoists every + block's MFMAs above every promotion: measured at K=512, all 64 MFMAs land in + the first 400 instructions and all the promotion arithmetic after, keeping + ``K_ITERS`` accumulator sets live at once -- 256 VGPR with **81 spilled**, + and the GEMM 10x slower (4052us against a 374us target). + ``rocdl.sched_barrier(0)`` between them does not stop it. + + So instead the accumulator carries the running sum *rescaled*, which keeps a + real data dependency the scheduler cannot break and needs no second + accumulator at all:: + + t_0 = d_0 + t_k = t_{k-1} * (s_{k-1} / s_k) + d_k + out = t_{K-1} * s_{K-1} + + with the invariant ``t_k = (sum_{j<=k} s_j d_j) / s_k``. The divisions cost + ~1e-7 relative each and there are ``K/128`` of them, three orders of + magnitude under the fp8 floor of 1.7e-3. It does assume no scale is exactly + zero, which a real quantiser never emits. + + Two properties of the tile geometry make the loads cheap, both checked + against ``_PermlaneStoreC._emit``'s indexing: + + * every accumulator set lies inside a single 128-column block + (``wave_n*32 + tj*16 + 15 <= 127``), so the B scale is **one scalar for + the whole set**, uniform across the wave; + * ``sa`` is column-major -- the layout + ``gemm_a8w8_blockscale_bpreshuffle`` consumes, via sglang's + ``materialize_bpreshuffle_fp8_scale`` -- so element ``(row, kb)`` sits at + ``kb * M + row`` and a lane's rows are contiguous. With ``swap_ab`` a lane + owns one row per M-tile (a scalar load); without it, four consecutive rows + (one vec4). + """ + + def __init__(self, A_scale, B_scale, m, n, k, *, swap_ab, n_tiles_a): + self.kb_count = k // 128 + self.m = m + self.swap_ab = swap_ab + self.n_tiles_a = n_tiles_a + self.lane = fx.thread_idx.x % 64 + gSA = fx.rocdl.make_buffer_tensor( + A_scale, max_size=False, num_records_bytes=m * self.kb_count * 4 + ) + gSB = fx.rocdl.make_buffer_tensor( + B_scale, max_size=False, num_records_bytes=(n // 128) * self.kb_count * 4 + ) + self.sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) + self.sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + self.atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + self.atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) + self.reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + self.reg_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) + + def _load1(self, div, index): + fx.copy(self.atom_1, fx.slice(div, (None, fx.Int32(index))), self.reg_1) + return Vec(fx.memref_load_vec(self.reg_1))[0] + + def _load4(self, div, index): + fx.copy(self.atom_4, fx.slice(div, (None, fx.Int32(index))), self.reg_4) + return Vec(fx.memref_load_vec(self.reg_4)) + + def a_scales(self, base_row, kb): + """Per M-tile A scale for this lane's rows at K-block ``kb``.""" + col = fx.Int32(kb) * fx.Int32(self.m) + lane = self.lane + if const_expr(self.swap_ab): + return [ + self._load1(self.sa_div, col + base_row + ti * 16 + lane % 16) + for ti in range_constexpr(self.n_tiles_a) + ] + return [ + self._load4(self.sa_div, col + base_row + ti * 16 + (lane // 16) * 4) + for ti in range_constexpr(self.n_tiles_a) + ] + + def b_scale(self, n_block, kb): + """The single B scale of a 128-column block at K-block ``kb``.""" + return self._load1( + self.sb_div, n_block * fx.Int32(self.kb_count) + fx.Int32(kb) + ) + + def scale_acc(self, acc, a_sc, b_sc, idx_fn, n_tiles_b): + """``acc *= s_a * s_b`` elementwise; returns the new accumulator. + + Used both for the running rescale between K-blocks and for the final + multiply, which are the same operation with different scale pairs. + """ + res = list(acc) + for ti in range_constexpr(self.n_tiles_a): + for tj in range_constexpr(n_tiles_b): + i = idx_fn(ti, tj) + v = Vec(res[i]) + if const_expr(self.swap_ab): + s = a_sc[ti] * b_sc + vals = [v[e] * s for e in range_constexpr(4)] + else: + vals = [v[e] * (a_sc[ti][e] * b_sc) for e in range_constexpr(4)] + res[i] = Vec.from_elements(vals, fx.Float32) + return res + + @staticmethod + def ratio(prev, cur): + """``prev / cur`` for a scale pair, elementwise over the M-tile list.""" + if isinstance(prev, list): + return [p / c for p, c in zip(prev, cur)] + return prev / cur + + def rescale_for( + self, kb, accs, prev, *, base_row, nb0, idx_fn, n_tiles_b, lds_block_m + ): + """Rebase the four accumulators from K-block ``kb-1``'s scales to kb's. + + A method rather than a closure inside the kernel: FlyDSL rewrites the + AST of every ``def`` nested in a ``@flyc.kernel`` function, and the + rewrite of an ``if`` inside such a nested def turned the captured + ``_BlockScaleK`` instance into a local, so reading it raised + UnboundLocalError. Module-level methods are ordinary Python at trace + time and are left alone. + """ + c00, c01, c10, c11 = accs + cur = ( + self.a_scales(base_row + 0 * lds_block_m, kb), + self.a_scales(base_row + 1 * lds_block_m, kb), + self.b_scale(nb0 + fx.Int32(0), kb), + self.b_scale(nb0 + fx.Int32(1), kb), + ) + if prev is not None: + pa0, pa1, pb0, pb1 = prev + ca0, ca1, cb0, cb1 = cur + ra0, ra1 = self.ratio(pa0, ca0), self.ratio(pa1, ca1) + rb0, rb1 = self.ratio(pb0, cb0), self.ratio(pb1, cb1) + c00 = self.scale_acc(c00, ra0, rb0, idx_fn, n_tiles_b) + c01 = self.scale_acc(c01, ra0, rb1, idx_fn, n_tiles_b) + c10 = self.scale_acc(c10, ra1, rb0, idx_fn, n_tiles_b) + c11 = self.scale_acc(c11, ra1, rb1, idx_fn, n_tiles_b) + return (c00, c01, c10, c11), cur + + def final_scale(self, accs, prev, *, idx_fn, n_tiles_b): + """Undo the invariant: multiply by the last K-block's scales.""" + fa0, fa1, fb0, fb1 = prev + c00, c01, c10, c11 = accs + return ( + self.scale_acc(c00, fa0, fb0, idx_fn, n_tiles_b), + self.scale_acc(c01, fa0, fb1, idx_fn, n_tiles_b), + self.scale_acc(c10, fa1, fb0, idx_fn, n_tiles_b), + self.scale_acc(c11, fa1, fb1, idx_fn, n_tiles_b), + ) + + +class _SwapABStoreC(StoreC): + """C store for an A/B-swapped MFMA: 4 consecutive N per lane -> 64-bit store. + + ``fx.gemm(atom, c, b, a, c)`` computes ``B^T A^T = (A B)^T`` in the + accumulator's native layout, so lane ``l`` value ``k`` moves from + ``D[4*(l/16)+k][l%16]`` (4 consecutive *rows*, stride c_cols) to + ``D[l%16][4*(l/16)+k]`` (4 consecutive *columns*). In a row-major C that is + 8 contiguous bytes, which is the whole point -- the unswapped layout can only + ever emit ``buffer_store_short``. This is gcnasm's ``mfma_adaptor_swap_ab``, + which does literally ``base::operator()(b, a, c)`` and redefines ``dim_c()`` + to match (opus.hpp:2064). + + The scale loads swap with it, and stay equally well vectorized: A's scale is + indexed by row, which is now ``lane % 16`` -- one row, so a scalar load -- + while B's is indexed by column, now 4 consecutive, so a vec4. + """ + + def __init__( + self, *args, peer_rsrc=None, elem_base=None, scales_preapplied=False, **kwargs + ): + super().__init__(*args, **kwargs) + # blockscale applies the scales per K-block in the mainloop, so by the + # time C reaches here it is already scaled and the epilogue is a plain + # fp32 -> bf16 convert. Skipping the loads rather than multiplying by a + # constant 1.0 keeps them out of the IR instead of trusting a fold. + self._preapplied = scales_preapplied + self.out_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.BFloat16) + self.reg_bf16_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.BFloat16) + self.out_atom_8 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.BFloat16) + self.reg_bf16_8 = fx.make_rmem_tensor(fx.make_layout(8, 1), fx.BFloat16) + self._peer_rsrc = peer_rsrc + self._elem_base = elem_base + + def _scaled(self, v, a, b): + """``v * a * b``, or ``v`` when the mainloop already applied them.""" + if const_expr(self._preapplied): + return v + return v * (a * b) + + def _load_a_scale_scalar(self, row): + if const_expr(self._preapplied): + return None + fx.copy( + self.scale_atom_1, + fx.slice(self.sa_div, (None, fx.Int32(row))), + self.reg_f32_1, + ) + return Vec(fx.memref_load_vec(self.reg_f32_1))[0] + + def _load_b_scale_vec4(self, col): + if const_expr(self._preapplied): + return [None] * 4 + fx.copy( + self.scale_atom_4, + fx.slice(self.sb_div, (None, fx.Int32(col))), + self.reg_f32_4, + ) + return Vec(fx.memref_load_vec(self.reg_f32_4)) + + def _store_bf16x4(self, values, c_index): + fx.memref_store_vec(Vec.from_elements(values, fx.BFloat16), self.reg_bf16_4) + fx.copy( + self.out_atom_4, + self.reg_bf16_4, + fx.slice(self.c_div, (None, fx.Int32(c_index))), + ) + + def store(self, c_frag, base_row, base_col): + lane = self.lane_id + a_scales = [ + self._load_a_scale_scalar(base_row + ti * 16 + lane % 16) + for ti in range_constexpr(self.n_tiles_a) + ] + b_scales = [ + self._load_b_scale_vec4(base_col + tj * 16 + (lane // 16) * 4) + for tj in range_constexpr(self.n_tiles_b) + ] + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + lane % 16 + for tj in range_constexpr(self.n_tiles_b): + col = base_col + tj * 16 + (lane // 16) * 4 + # BLOCK_N divides N in every shape this builds for, so the four + # columns are either all in range or all out; one predicate. + col_valid = col + 3 < self.c_cols + oob = fx.Int32(self.c_rows * self.c_cols) + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + vals = [ + self._scaled(vec_f32[k], a_scales[ti], b_scales[tj][k]).to( + fx.BFloat16 + ) + for k in range_constexpr(4) + ] + c_index = row * self.c_cols + col + self._store_bf16x4(vals, arith.select(col_valid, c_index, oob)) + + +def _raw(v): + """Unwrap a DSL value to its raw ``ir.Value``, for ops that take one.""" + return v.ir_value() if hasattr(v, "ir_value") else v + + +def _permlane16_swap(x, y): + """gfx950 ``v_permlane16_swap_b32``: exchange 16-lane rows between two VGPRs. + + Verified on hardware rather than assumed:: + + vdst_new = [X.r0, Y.r0, X.r2, Y.r2] + vsrc_new = [X.r1, Y.r1, X.r3, Y.r3] + + Position *within* a row is preserved, so a lane keeps its ``lane % 16`` and + therefore its C row -- which is what lets the A-scale stay correct across the + shuffle. + """ + st = ir.Type.parse("!llvm.struct<(i32, i32)>") + res = fx.rocdl.permlane16_swap(st, _raw(x), _raw(y), False, True) + i32 = ir.IntegerType.get_signless(32) + return ( + fx.Int32( + _llvm_d.ExtractValueOp(i32, res, ir.DenseI64ArrayAttr.get([0])).result + ), + fx.Int32( + _llvm_d.ExtractValueOp(i32, res, ir.DenseI64ArrayAttr.get([1])).result + ), + ) + + +def _ds_bpermute(value, src_lane): + """Pull ``value`` from ``src_lane``. Byte-addressed, hence the <<2.""" + i32 = ir.IntegerType.get_signless(32) + return fx.Int32( + fx.rocdl.ds_bpermute(i32, _raw(fx.Int32(src_lane) * fx.Int32(4)), _raw(value)) + ) + + +class _PermlaneStoreC(_SwapABStoreC): + """A/B-swapped MFMA + one cross-lane step, giving 16B per lane and 64B rows. + + After the swap a lane owns 4 consecutive columns of one N-tile, i.e. 8 bytes, + and the two N-tiles it holds are 16 columns apart -- so a row still only gets + 32 bytes per instruction. Two ``permlane16_swap`` fix that exactly, with no + ``ds_bpermute`` and no LDS: + + (A, B) = permlane16_swap(tile0.d0, tile1.d0) + (C, D) = permlane16_swap(tile0.d1, tile1.d1) + lane group g stores (A, C, B, D) + + which lands g = 0,1,2,3 on columns 0-7, 16-23, 8-15, 24-31. The four groups + together cover columns 0..31 contiguously: 64 bytes per row, from 16 bytes per + lane. The column permutation is absorbed into the address, so it is free. + + Scaling happens before the shuffle, in the lane that owns the value; the + shuffle only moves whole 16-lane rows, so every lane keeps its own C row. + """ + + _adjacent_probe = False + _lane_transpose = False + _peer_uncached = False + + def store(self, c_frag, base_row, base_col): + self._emit( + c_frag, + base_row, + base_col, + self._a_scales(base_row), + self._b_scales(base_col), + ) + + def store_all(self, frags, base_row, base_col, row_step, col_step): + """All four half-tiles, loading each scale once instead of twice. + + The four ``store`` calls the epilogue makes share base rows pairwise and + base columns pairwise, so calling them individually issues every scale + load twice -- 16 ``buffer_load_dwordx4`` for A where 8 addresses are + distinct, 8 ``buffer_load_dword`` for B where 4 are. The compiler cannot + merge them because they all write the same ``reg_f32_*`` register buffer, + which turns them into a chain of overwrites rather than pure loads. + """ + a = [self._a_scales(base_row + r * row_step) for r in range_constexpr(2)] + b = [self._b_scales(base_col + c * col_step) for c in range_constexpr(2)] + for frag, r, c in frags: + self._emit( + frag, base_row + r * row_step, base_col + c * col_step, a[r], b[c] + ) + + def _a_scales(self, base_row): + lane = self.lane_id + return [ + self._load_a_scale_scalar(base_row + ti * 16 + lane % 16) + for ti in range_constexpr(self.n_tiles_a) + ] + + def _b_scales(self, base_col): + grp = self.lane_id // 16 + return [ + self._load_b_scale_vec4(base_col + tj * 16 + grp * 4) + for tj in range_constexpr(self.n_tiles_b) + ] + + def _emit(self, c_frag, base_row, base_col, a_scales, b_scales): + assert self.n_tiles_b == 2, ( + "the permlane mapping pairs exactly two N-tiles (BLOCK_N == 256); " + f"got n_tiles_b={self.n_tiles_b}" + ) + lane = self.lane_id + grp = lane // 16 + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + lane % 16 + dwords = [] + for tj in range_constexpr(self.n_tiles_b): + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + packed = Vec.from_elements( + [ + self._scaled(vec_f32[k], a_scales[ti], b_scales[tj][k]).to( + fx.BFloat16 + ) + for k in range_constexpr(4) + ], + fx.BFloat16, + ).bitcast(fx.Int32) + dwords.append((packed[0], packed[1])) + a, b = _permlane16_swap(dwords[0][0], dwords[1][0]) + c, d = _permlane16_swap(dwords[0][1], dwords[1][1]) + out8 = Vec.from_elements([a, c, b, d], fx.Int32).bitcast(fx.BFloat16) + if const_expr(self._lane_transpose): + # gcnasm's second stage (kernel_template.hpp:491-510), as one + # ds_bpermute per dword. The permlane stage leaves a row's four + # 8-column chunks in lanes 16 apart, so a 16-lane group touches + # 16 rows at 16 bytes each. Transposing the lane index -- lane + # l' = 4r'+q' pulls from the lane holding (row r', chunk q') -- + # puts adjacent lanes on one row, so lanes 0-3 write 64 + # contiguous bytes and a group covers 4 rows instead of 16. + # + # chunk q' lives at column q'*8, and the permlane stage put + # column (g%2)*16 + (g//2)*8 in group g, so g = swap of q's two + # bits: 0,1,2,3 -> 0,2,1,3. + q = lane % 4 + src_lane = ((q % 2) * 2 + q // 2) * 16 + lane // 4 + a, c, b, d = (_ds_bpermute(v, src_lane) for v in (a, c, b, d)) + out8 = Vec.from_elements([a, c, b, d], fx.Int32).bitcast(fx.BFloat16) + row = base_row + ti * 16 + lane // 4 + col = base_col + q * 8 + elif const_expr(self._adjacent_probe): + # PERF PROBE, output wrong by construction. Same 64 addresses, + # reassigned so *adjacent* lanes cover one row (lanes 0-3 -> row 0, + # 64 contiguous bytes) instead of lanes 16 apart. Prices gcnasm's + # ds_bpermute lane transpose before writing it: if the address + # coalescer works per 16-lane group rather than per wave, this + # takes a group from 16 sectors to 4. + row = base_row + ti * 16 + lane // 4 + col = base_col + (lane % 4) * 8 + else: + col = base_col + (grp % 2) * 16 + (grp // 2) * 8 + oob = fx.Int32(self.c_rows * self.c_cols) + idx = arith.select(col + 7 < self.c_cols, row * self.c_cols + col, oob) + if self._peer_rsrc is not None: + # The store goes to a *peer*, so a cached one leaves the line + # dirty in whichever XCD's L2 this block ran on, where nothing + # downstream can reach it -- the barrier kernel is one block and + # buffer_wbl2 writes back only its own XCD's L2, one of eight. + # sc0|sc1 skips both. That was not usable while the store was one + # bf16 (a 2-byte partial-line write over the fabric loses + # updates), but the permlane stage makes it 16 bytes per lane and + # 64 contiguous per row, which is what makes it viable now. + buffer_store( + out8, + self._peer_rsrc, + fx.Int32(idx) - self._elem_base, + cache_modifier=CM_SC0_SC1 if self._peer_uncached else CM_CACHED, + ) + else: + fx.memref_store_vec(out8, self.reg_bf16_8) + fx.copy( + self.out_atom_8, + self.reg_bf16_8, + fx.slice(self.c_div, (None, fx.Int32(idx))), + ) + + +class _LaneTransposeStoreC(_PermlaneStoreC): + """``_PermlaneStoreC`` plus gcnasm's ds_bpermute lane transpose.""" + + _lane_transpose = True + + +class _LaneTransposeUncachedPeerStoreC(_LaneTransposeStoreC): + """Same, but the peer store bypasses L1 and L2 (Direct LSA only).""" + + _peer_uncached = True + + +class _AdjacentLaneProbeC(_PermlaneStoreC): + """PERF PROBE ONLY -- see the note in ``_PermlaneStoreC.store``.""" + + _adjacent_probe = True + + +class _WideStoreProbeC(_SwapABStoreC): + """PERF PROBE ONLY -- the output is wrong on purpose. + + Answers "is 16B per lane / 64B per row actually faster?" before anyone builds + the cross-lane shuffle that would make it correct. It emits exactly the access + pattern the shuffled version would have -- one 128-bit store per lane at + ``col = (lane//16)*8``, so 4 lanes cover a row's 32 columns as 64 contiguous + bytes -- but feeds it the two N-tiles' values concatenated, which is not what + belongs at those addresses. + + Correct would need lane ``(a, r)`` to gather cols ``8a..8a+7``, which live in + two *other* lanes' registers (``a'=2a`` and ``2a+1`` of one tile). That is the + ``permlane16_swap`` step; this probe skips it and keeps only its cost profile. + """ + + def store(self, c_frag, base_row, base_col): + lane = self.lane_id + a_scales = [ + self._load_a_scale_scalar(base_row + ti * 16 + lane % 16) + for ti in range_constexpr(self.n_tiles_a) + ] + b_scales = [ + self._load_b_scale_vec4(base_col + tj * 16 + (lane // 16) * 4) + for tj in range_constexpr(self.n_tiles_b) + ] + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + lane % 16 + vals = [] + for tj in range_constexpr(self.n_tiles_b): + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + vals += [ + self._scaled(vec_f32[k], a_scales[ti], b_scales[tj][k]).to( + fx.BFloat16 + ) + for k in range_constexpr(4) + ] + col = base_col + (lane // 16) * 8 + oob = fx.Int32(self.c_rows * self.c_cols) + c_index = row * self.c_cols + col + idx = arith.select(col + 7 < self.c_cols, c_index, oob) + if self._peer_rsrc is not None: + buffer_store( + Vec.from_elements(vals, fx.BFloat16), + self._peer_rsrc, + fx.Int32(idx) - self._elem_base, + ) + else: + fx.memref_store_vec( + Vec.from_elements(vals, fx.BFloat16), self.reg_bf16_8 + ) + fx.copy( + self.out_atom_8, + self.reg_bf16_8, + fx.slice(self.c_div, (None, fx.Int32(idx))), + ) + + +class _PeerDirectStoreC(StoreC): + """``StoreC`` that writes its tile straight into the owning peer's window. + + This is gcnasm's "Direct LSA" shape (``opus_gemm_a2a_lsa``), and its whole + point is that there is no producer/consumer handoff inside the kernel: the + bytes go to the destination as they are produced, so nothing has to be + published mid-kernel, there is no completion counter, no submit lock and no + release fence. The only synchronisation left is the cross-rank barrier after + the kernel, where a kernel-end release has already happened for free. + + A block's rows all belong to one destination (``block_m // m_tiles_per_peer``), + so the peer descriptor is uniform per block and is built once. + + ``elem_base`` rebases the global row index onto the destination's slice. + ``StoreC.store`` masks out-of-range columns by redirecting them to + ``c_rows * c_cols``, and that stays out of range after rebasing: the largest + ``elem_base`` is ``(world-1) * slice_rows * c_cols``, which leaves exactly + ``slice_rows * c_cols`` elements -- one past the end of the descriptor. + """ + + def __init__(self, *args, peer_rsrc=None, elem_base=None, **kwargs): + super().__init__(*args, **kwargs) + self._peer_rsrc = peer_rsrc + self._elem_base = elem_base + + def _store_bf16(self, value_bf16, c_index): + # Cached, and that is the unresolved half of this variant. See the + # "Direct LSA" note in the module docstring: uncached (sc0|sc1) is + # *consistently* wrong here because a 2-byte store bypassing L2 is a + # partial-line write, and cached is only intermittently right because the + # dirty peer-homed lines cannot be flushed out of eight XCDs' L2s + # afterwards. Coalescing the store first resolves both. + buffer_store(value_bf16, self._peer_rsrc, c_index - self._elem_base) + + +class _NonTemporalStoreC(StoreC): + """``StoreC`` with non-temporal C stores (the copy atom's only other mode). + + Probes whether ``buffer_wbl2`` gets cheaper when the lines it is asked to + flush have already been evicted. The atom's ``cache_modifier`` is a two-value + enum, 0=cached / 2=nt -- it is *not* the aux bitmask that + ``raw_ptr_buffer_store`` takes, so sc0|sc1 cannot be requested this way. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(2), fx.BFloat16) + + +class _WriteThroughStoreC(StoreC): + """``StoreC`` whose C stores bypass L1 and L2 instead of being written back. + + The only change is the copy atom's cache modifier. It exists because a + per-block release fence is quadratic: ``buffer_wbl2`` flushes the *whole* L2, + so block k redundantly writes back every tile blocks 1..k-1 already wrote, + 448 times over. Writing C through in the first place makes the release free -- + ``s_waitcnt vmcnt(0)`` is then all the copy engine needs -- and adds no + traffic, since C has to reach memory regardless. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.out_atom_1 = fx.make_copy_atom( + fx.rocdl.BufferCopy16b(CM_SC0_SC1), fx.BFloat16 + ) + + +BLOCK_K = 128 + + +def compile_fused_gemm_scatter( + cfg, + rank: int, + *, + K: int, + BLOCK_M: int = 256, + BLOCK_N: int = 256, + b_preshuffled: bool = True, + waves_per_eu: int = 2, + xcd_swizzle: int = 0, + fuse: bool = True, + rotated: bool | None = None, + n_stripe: int | None = None, + quant: str = "ptpc", + sdma_queues: int = 8, + swap_ab: bool = False, + store_probe: bool = False, + permlane: bool = False, + lane_transpose: bool = False, + hoist_scales: bool = False, + peer_uncached: bool = False, + direct_fence: str = "leader", + transport: str = "sdma", + fence: str = "none", + emit_put: bool = True, + atomic_order: str = "acq_rel", +): + """Compile the GEMM, with (``fuse=True``) or without the scatter epilogue. + + ``fuse=False`` emits the identical kernel minus the epilogue tail, so the + split baseline and the fused kernel differ in exactly one thing. That is why + the baseline is built here rather than called through a separate wrapper, which + would also change the launch path and the C tensor. + + ``rotated`` selects the destination-rotated chunk-major tile order; it + defaults to ``fuse``, since it only matters when something is watching tiles + complete. It is separately settable so the benchmark can charge the split + baseline the same tile order and show that the ordering itself is neutral. + + ``n_stripe`` is how many N-tiles a block walks before the destination + rotates -- gcnasm's ``opus_direct_stripe_tile``, transposed (its + destination is on N and ours is on M, so its ``m_stripe`` is our + ``n_stripe``). It spans the whole range: 1 rotates on every block, and + ``N // BLOCK_N`` is chunk-major, where a destination's whole chunk is + finished before moving on. ``None`` picks 2 for ``fused-lsa`` and + chunk-major elsewhere. + + Swept at [16384, 7168] K=2048 on 8 ranks (28 N-tiles, so the divisors are + 1, 2, 4, 7, 14, 28): + + n_stripe 1 2 4 7 14 28 + fused-lsa 1370.9 1362.5 1404.9 1402.9 1437.0 1569.0 + fused-sdma 1117.9 1117.1 1114.8 1124.0 1117.5 1115.7 + + ``fused-lsa`` is monotone from 28 down to 2 and then flat -- 1 and 2 are the + same within noise -- for 13% end to end. ``fused-sdma`` is flat across the + whole range: its bytes leave from a staging buffer on the copy engines, so + only the moment a chunk *completes* matters, and with 256 resident blocks + out of 1792 the first chunk completes inside the first wave of blocks + whatever the order. + + Two hypotheses this sweep kills. It is not XCD locality: with + ``dest = (block / n_stripe) % 8`` and blocks going round-robin over the 8 + XCDs, a destination is fed by exactly ``min(n_stripe, 8)`` XCDs, so 14 and + 28 spread over all eight and 2 over only two -- and 2 is the fast one. And + it is not chunk-completion timing, or ``fused-sdma`` would prefer + chunk-major, which it does not. + + Returns ``launch(A, B_T, C, A_scale, B_scale, c_m, c_n, dev_comm, win, + stream=...)``. + """ + cfg.validate() + ws = cfg.world_size + M, N = cfg.m, cfg.n + if transport not in ("sdma", "lsa"): + raise ValueError(f"transport must be sdma or lsa, got {transport!r}") + # The C-store variant is picked by a chain of const_expr branches, so an + # unimplemented combination does not fail -- it quietly falls through to a + # weaker one. That already cost a wrong conclusion once: --lane-transpose had + # no direct-LSA branch, so `--mode fused-lsa --lane-transpose` silently ran + # stages 1+2 and was measured as "no benefit from stage 3". + if permlane and not swap_ab: + raise ValueError( + "--permlane requires --swap-ab. Both are on by default in " + "bench_gemm_ar.py, so pass --no-permlane alongside --no-swap-ab" + ) + if lane_transpose and not permlane: + raise ValueError( + "--lane-transpose requires --permlane. Both are on by default in " + "bench_gemm_ar.py, so pass --no-lane-transpose alongside " + "--no-permlane" + ) + if store_probe and not swap_ab: + raise ValueError("--store-probe requires --swap-ab (on by default)") + if store_probe and lane_transpose: + raise ValueError( + "--store-probe and --lane-transpose are alternatives: the probe emits " + "the adjacent-lane access pattern without the shuffle that makes it " + "correct, which is what --lane-transpose then implements" + ) + if transport == "lsa" and store_probe and permlane: + raise ValueError( + "no peer-direct variant of the adjacent-lane probe; use " + "--lane-transpose for the real thing on --mode fused-lsa" + ) + direct_lsa = fuse and transport == "lsa" + if quant not in ("ptpc", "blockscale"): + raise ValueError(f"quant must be ptpc or blockscale, got {quant!r}") + blockscale = quant == "blockscale" + if blockscale: + if not swap_ab: + raise ValueError( + "--quant blockscale needs --swap-ab: the unswapped epilogue " + "applies its scales inside the stock StoreC, which this copy does " + "not override" + ) + if K % 128: + raise ValueError(f"blockscale needs K % 128 == 0, got K={K}") + if N % 128: + raise ValueError(f"blockscale needs N % 128 == 0, got N={N}") + rotated = fuse if rotated is None else rotated + # n_stripe spans 1..N//BLOCK_N; the top of the range *is* chunk-major, so + # 0 ("per-mode default") resolves to it rather than to a separate branch. + if n_stripe is None: + n_stripe = 2 if direct_lsa else 0 + if n_stripe == 0: + n_stripe = N // BLOCK_N + if n_stripe < 1 or n_stripe > N // BLOCK_N: + raise ValueError(f"n_stripe must be in [1, {N // BLOCK_N}], got {n_stripe}") + if (N // BLOCK_N) % n_stripe: + raise ValueError(f"n_stripe={n_stripe} must divide the {N // BLOCK_N} N-tiles") + # Only a *real* stripe needs the rotated order. Chunk-major is the + # unstriped case and is what the linear order already does per destination, + # so leaving it on a non-rotated build is not an error -- the split + # baseline compiles that way. + if n_stripe < N // BLOCK_N and not rotated: + raise ValueError("a striped tile order needs --tile-order rotated") + + assert BLOCK_M >= 128 and BLOCK_N >= 256 + assert BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert K % BLOCK_K == 0 + if N % BLOCK_N: + raise ValueError( + f"N={N} must be a multiple of BLOCK_N={BLOCK_N}: the tile count per " + "destination has to be a compile-time constant for the modulo test" + ) + slice_rows = M // ws + if rotated and slice_rows % BLOCK_M: + raise ValueError( + f"each peer's row slice ({slice_rows}) must be a whole number of " + f"BLOCK_M={BLOCK_M} tiles, so M must be a multiple of " + f"world_size*BLOCK_M={ws * BLOCK_M} (got M={M}). Decode shapes are " + "out of scope for this host anyway -- it pads M below 128." + ) + + K_ITERS = K // BLOCK_K + N_TILES_A = BLOCK_M // 64 + N_TILES_B = BLOCK_N // 128 + N_ACCUMS = N_TILES_A * N_TILES_B + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + a_lds_size = LDS_BLOCK_M * BLOCK_K + b_lds_size = LDS_BLOCK_N * BLOCK_K + + # Compile-time tile accounting for the completion counters. + n_blocks_const = N // BLOCK_N + m_tiles_per_peer = slice_rows // BLOCK_M if slice_rows % BLOCK_M == 0 else 0 + + # Push granularity. A destination's slice is `m_tiles_per_peer` row-bands of + # BLOCK_M rows; pushing each band as it completes is what lets every link + # start streaming early instead of waiting for its whole slice. + chunks = cfg.counter_chunks if fuse else 1 + if fuse and m_tiles_per_peer % chunks: + raise ValueError( + f"counter_chunks={chunks} must divide the {m_tiles_per_peer} " + f"BLOCK_M={BLOCK_M} row-bands in each peer's slice" + ) + m_tiles_per_chunk = max(1, m_tiles_per_peer // chunks) + tiles_per_chunk = m_tiles_per_chunk * n_blocks_const + chunk_bytes = cfg.slice_bytes // chunks + + if fence not in ( + "all", + "agent", + "agent-leader", + "nt-agent", + "leader", + "none", + "writethrough", + "wt-agent", + "raw-wt", + "raw-wt-agent", + "raw-wt-leader", + ): + raise ValueError( + f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" + ) + _kname_tag = f"{'B' if blockscale else ''}{'S' if swap_ab else ''}{'P' if permlane else ''}{'T' if lane_transpose else ''}{'H' if hoist_scales else ''}{'U' if peer_uncached else ''}{direct_fence[0]}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}q{sdma_queues}{'' if n_stripe == 1 else f's{n_stripe}'}{fence[0]}" + counter_off = cfg.counter_off + lock_off = cfg.lock_off + in_off = cfg.input_off + my_recv_slot = cfg.recv_slot_off(rank) if cfg.recv_slots else 0 + slice_bytes = cfg.slice_bytes + + _kname = ( + f"mori_fused_{(transport if fuse else 'split')}_8w_" + f"{BLOCK_M}x{BLOCK_N}x{BLOCK_K}_k{K}_{_kname_tag}" + f"{'p' if emit_put else 'x'}{atomic_order[0]}_r{rank}" + ) + + @fx.struct + class SharedStorage: + A_lds_cur_0: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + A_lds_cur_1: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + A_lds_next_0: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + A_lds_next_1: fx.Array[fx.Float8E4M3FN, a_lds_size, 16] + B_lds_cur_0: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + B_lds_cur_1: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + B_lds_next_0: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + B_lds_next_1: fx.Array[fx.Float8E4M3FN, b_lds_size, 16] + + @flyc.kernel(name=_kname, known_block_size=[512, 1, 1]) + def kernel_gemm_scatter( + A: fx.Tensor, + B_T: fx.Tensor, + C: fx.Tensor, + A_scale: fx.Tensor, + B_scale: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + dev_comm: Int64, + win: Int64, + ): + # ---- begin pinned copy of _gemm_a8w8_8wave kernel_gemm ---- + F8_IR_t = fx.Float8E4M3FN.ir_type + + n_blocks = ceildiv(c_n, BLOCK_N) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + a_cur0 = lds.A_lds_cur_0 + a_cur1 = lds.A_lds_cur_1 + a_next0 = lds.A_lds_next_0 + a_next1 = lds.A_lds_next_1 + b_cur0 = lds.B_lds_cur_0 + b_cur1 = lds.B_lds_cur_1 + b_next0 = lds.B_lds_next_0 + b_next1 = lds.B_lds_next_1 + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + if const_expr(rotated): + # Destination-rotated, chunk-major tile order. Linear order finishes + # destination 0's whole slice, then 1's, ... -- so the last + # destination's link only starts at the end of the GEMM and the + # transfer cannot overlap anything. Walking (chunk, dest, n) instead + # gets every link streaming within the first 1/(chunks*world) of the + # GEMM. The `+ rank` rotation keeps all 8 ranks from pushing at the + # same destination simultaneously (gcnasm's opus_direct_stripe_tile). + # split_row_major_2d(i, n) -> (i // n, i % n) + if const_expr(n_stripe >= n_blocks_const): + rest, bn = split_row_major_2d(fx.block_idx.x, n_blocks) + tile_i, dest_seq = split_row_major_2d(rest, ws) + else: + # gcnasm's opus_direct_stripe_tile: rotate the destination + # every `n_stripe` tiles of the other axis rather than after a + # whole chunk, so the peer stores of the resident blocks are + # spread over all links instead of arriving in same-destination + # bursts. It delays chunk completion by n_groups, which is why + # gcnasm keeps it out of its chunked path and so do we by + # leaving the default at 1. + rest, n_in = split_row_major_2d(fx.block_idx.x, n_stripe) + rest, dest_seq = split_row_major_2d(rest, ws) + tile_i, n_grp = split_row_major_2d(rest, n_blocks // n_stripe) + bn = n_grp * fx.Int32(n_stripe) + n_in + dest_i = (dest_seq + fx.Int32(rank)) % fx.Int32(ws) + block_m = dest_i * m_tiles_per_peer + tile_i + block_n = bn + elif const_expr(xcd_swizzle > 0): + block_m, block_n = _xcd_swizzle_any( + ceildiv(c_m, BLOCK_M), n_blocks, wgm=xcd_swizzle + ) + else: + block_m, block_n = split_row_major_2d(fx.block_idx.x, n_blocks) + + A0_gl_offset = (block_m * BLOCK_M) * K + A1_gl_offset = (block_m * BLOCK_M + LDS_BLOCK_M) * K + B_K_STEP = (2 * 1024) if b_preshuffled else BLOCK_K + B0_gl_offset = (block_n * BLOCK_N) * K + B1_gl_offset = (block_n * BLOCK_N + LDS_BLOCK_N) * K + + gA = make_fp8_buffer_tensor(A, F8_IR_t) + gB = make_fp8_buffer_tensor(B_T, F8_IR_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + + gl_off_a = compute_global_swizzle( + lane_id, wave_id, K, N_LDS_ROUNDS, preshuffled=False + ) + gl_off_b = compute_global_swizzle( + lane_id, wave_id, K, N_LDS_ROUNDS, preshuffled=b_preshuffled + ) + + if const_expr(swap_ab): + # Tile counts swap with the operands so Mfma's own asserts and its + # idx() line up; the store then addresses the accumulator as + # idx(tj, ti). + mfma_raw = Mfma16x16x128(N_TILES_B, N_TILES_A) + mfma = _SwappedMfma(mfma_raw) + else: + mfma = Mfma16x16x128(N_TILES_A, N_TILES_B) + w_pre = cco.CachedWindow(win) + + a_g2s = G2SLoader(a_div, gl_off_a, N_LDS_STEPS_A, F8_IR_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, F8_IR_t, wave_id) + a_s2r = S2RLoader(wave_m, N_TILES_A) + b_s2r = S2RLoader(wave_n, N_TILES_B) + if const_expr(swap_ab and permlane and lane_transpose and not direct_lsa): + store_c = _LaneTransposeStoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + elif const_expr(swap_ab and permlane and store_probe and not direct_lsa): + store_c = _AdjacentLaneProbeC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + elif const_expr(swap_ab and permlane and not direct_lsa): + store_c = _PermlaneStoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + elif const_expr(swap_ab and store_probe and not direct_lsa): + store_c = _WideStoreProbeC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + elif const_expr(swap_ab and not direct_lsa): + store_c = _SwapABStoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + elif const_expr(direct_lsa and swap_ab and permlane and lane_transpose): + dest_blk = block_m // fx.Int32(m_tiles_per_peer) + _cls = ( + _LaneTransposeUncachedPeerStoreC + if const_expr(peer_uncached) + else _LaneTransposeStoreC + ) + store_c = _cls( + A_scale, + B_scale, + C, + c_m, + c_n, + mfma.idx, + N_TILES_A, + N_TILES_B, + peer_rsrc=create_buffer_resource_from_addr( + wave_uniform_i64(w_pre.lsa_ptr(dest_blk, my_recv_slot)), + num_records_bytes=cfg.slice_bytes, + ), + elem_base=dest_blk * fx.Int32(slice_rows) * c_n, + ) + elif const_expr(direct_lsa and swap_ab and permlane): + dest_blk = block_m // fx.Int32(m_tiles_per_peer) + store_c = _PermlaneStoreC( + A_scale, + B_scale, + C, + c_m, + c_n, + mfma.idx, + N_TILES_A, + N_TILES_B, + peer_rsrc=create_buffer_resource_from_addr( + wave_uniform_i64(w_pre.lsa_ptr(dest_blk, my_recv_slot)), + num_records_bytes=cfg.slice_bytes, + ), + elem_base=dest_blk * fx.Int32(slice_rows) * c_n, + ) + elif const_expr(direct_lsa and swap_ab and store_probe): + dest_blk = block_m // fx.Int32(m_tiles_per_peer) + store_c = _WideStoreProbeC( + A_scale, + B_scale, + C, + c_m, + c_n, + mfma.idx, + N_TILES_A, + N_TILES_B, + peer_rsrc=create_buffer_resource_from_addr( + wave_uniform_i64(w_pre.lsa_ptr(dest_blk, my_recv_slot)), + num_records_bytes=cfg.slice_bytes, + ), + elem_base=dest_blk * fx.Int32(slice_rows) * c_n, + ) + elif const_expr(direct_lsa): + # dest is uniform across the block; readfirstlane keeps the peer + # descriptor scalar and avoids a waterfall around every store. + dest_blk = block_m // fx.Int32(m_tiles_per_peer) + store_c = _PeerDirectStoreC( + A_scale, + B_scale, + C, + c_m, + c_n, + mfma.idx, + N_TILES_A, + N_TILES_B, + peer_rsrc=create_buffer_resource_from_addr( + wave_uniform_i64(w_pre.lsa_ptr(dest_blk, my_recv_slot)), + num_records_bytes=cfg.slice_bytes, + ), + elem_base=dest_blk * fx.Int32(slice_rows) * c_n, + ) + elif const_expr(fence in ("writethrough", "wt-agent")): + store_c = _WriteThroughStoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + elif const_expr(fence in ("raw-wt", "raw-wt-agent", "raw-wt-leader")): + store_c = _RawWriteThroughStoreC( + A_scale, + B_scale, + C, + c_m, + c_n, + mfma.idx, + N_TILES_A, + N_TILES_B, + c_rsrc=create_buffer_resource_from_addr( + fx.Int64(w_pre.lsa_ptr(rank, in_off)), + num_records_bytes=cfg.nbytes, + ), + ) + elif const_expr(fence == "nt-agent"): + store_c = _NonTemporalStoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + else: + store_c = StoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) + # Bound unconditionally: the kernel body is re-parsed by FlyDSL's AST + # rewriter, and names that only exist inside a branch are not reliably + # visible to a nested def afterwards (a `nonlocal` on one raised "no + # binding" and reading one raised UnboundLocalError). + bsk = nb0 = base_row_pre = None + if blockscale: + # Set after construction rather than threading a kwarg through + # thirteen call sites: it is a trace-time Python bool read by + # ``const_expr`` and nothing looks at it before the first store. + store_c._preapplied = True + bsk = _BlockScaleK( + A_scale, B_scale, c_m, N, K, swap_ab=swap_ab, n_tiles_a=N_TILES_A + ) + # The B scale is one scalar per 128-column block, and each + # accumulator set sits inside exactly one: c*0 is block + # ``block_n * (BLOCK_N // 128)``, c*1 the next. + nb0 = block_n * fx.Int32(BLOCK_N // 128) + # Same expression the epilogue rebuilds at ``base_row`` below; the + # rescale needs it one loop earlier. + base_row_pre = block_m * BLOCK_M + wave_m * (N_TILES_A * 16) + # No second accumulator: the running sum lives in the MFMA registers, + # rescaled between K-blocks (see _BlockScaleK). + prev_scales = None + + c00_frag = [mfma.zero_value] * N_ACCUMS + c01_frag = [mfma.zero_value] * N_ACCUMS + c10_frag = [mfma.zero_value] * N_ACCUMS + c11_frag = [mfma.zero_value] * N_ACCUMS + + b_g2s.load(b_cur0, B0_gl_offset + 0 * B_K_STEP) + a_g2s.load(a_cur0, A0_gl_offset + 0 * BLOCK_K) + b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) + a_g2s.load(a_cur1, A1_gl_offset + 0 * BLOCK_K) + + if wave_m == 1: + rocdl.s_barrier() + + wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B) + + b_g2s.load(b_next0, B0_gl_offset + 1 * B_K_STEP) + a_g2s.load(a_next0, A0_gl_offset + 1 * BLOCK_K) + b_g2s.load(b_next1, B1_gl_offset + 1 * B_K_STEP) + + wait_barrier(N_LDS_STEPS_A + 2 * N_LDS_STEPS_B) + + for k in range_constexpr(K_ITERS - 2): + if blockscale: + (c00_frag, c01_frag, c10_frag, c11_frag), prev_scales = bsk.rescale_for( + k, + (c00_frag, c01_frag, c10_frag, c11_frag), + prev_scales, + base_row=base_row_pre, + nb0=nb0, + idx_fn=mfma.idx, + n_tiles_b=N_TILES_B, + lds_block_m=LDS_BLOCK_M, + ) + b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) + a0_frag = a_s2r.load(a_cur0) + a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) + rocdl.s_barrier() + + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + + b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) + b_g2s.load(b_cur0, B0_gl_offset + (k + 2) * B_K_STEP) + rocdl.s_barrier() + + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + + a1_frag = a_s2r.load(a_cur1) + a_g2s.load(a_cur0, A0_gl_offset + (k + 2) * BLOCK_K) + rocdl.s_barrier() + + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + + b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * B_K_STEP) + # aiter has `2 * N_LDS_STEPS_A + N_LDS_STEPS_B` here, which lets + # too many global->LDS prefetches stay in flight across the barrier: + # a wave passes it and refills an LDS buffer that another wave has + # not finished consuming. The result is a silent, non-deterministic + # partial corruption -- always a whole (wave_m sub-tile x all four + # wave_n) region, which is the signature of a shared A-side LDS + # buffer -- that only shows up on large grids. Measured boundary at + # [16384, 7168] K=512 on 8x MI355X, 4-5 runs each: at 256x256 the + # counts 6 (aiter's), 5 and 4 all corrupt while 3, 2 and 0 are + # clean; at 128x256 the counts 4 (aiter's) and 3 corrupt while 2, 1 + # and 0 are clean. Both thresholds are N_LDS_STEPS_A + + # N_LDS_STEPS_B - 1. Costs nothing: 241/240/245 us against + # 240/246/237 us for aiter's count at [16384, 7168] K=2048. + # + # Not a complete theory -- 128x512 does not reproduce at all, so the + # count is not the only variable, and the same under-wait may exist + # at the other sync points (forcing *any* of the 14 in this loop to + # vmcnt(0) also fixes it). This is the smallest change that is + # measurably correct here, not a proven-general formula. + wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B - 1) + + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + # Step k = K_ITERS - 2 + if blockscale: + (c00_frag, c01_frag, c10_frag, c11_frag), prev_scales = bsk.rescale_for( + K_ITERS - 2, + (c00_frag, c01_frag, c10_frag, c11_frag), + prev_scales, + base_row=base_row_pre, + nb0=nb0, + idx_fn=mfma.idx, + n_tiles_b=N_TILES_B, + lds_block_m=LDS_BLOCK_M, + ) + b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) + a0_frag = a_s2r.load(a_cur0) + rocdl.s_barrier() + + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + + b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) + rocdl.s_barrier() + + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + + a1_frag = a_s2r.load(a_cur1) + a_g2s.load(a_next1, A1_gl_offset + (K_ITERS - 1) * BLOCK_K) + rocdl.s_barrier() + + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + + b0_frag = b_s2r.load(b_next0, preshuffled=b_preshuffled) + rocdl.s_barrier() + + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + # Step k = K_ITERS - 1 + if blockscale: + (c00_frag, c01_frag, c10_frag, c11_frag), prev_scales = bsk.rescale_for( + K_ITERS - 1, + (c00_frag, c01_frag, c10_frag, c11_frag), + prev_scales, + base_row=base_row_pre, + nb0=nb0, + idx_fn=mfma.idx, + n_tiles_b=N_TILES_B, + lds_block_m=LDS_BLOCK_M, + ) + a0_frag = a_s2r.load(a_cur0) + wait_barrier(0) + + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + + b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) + rocdl.s_barrier() + + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + + a1_frag = a_s2r.load(a_cur1) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, set_prio=False) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, set_prio=False) + rocdl.s_setprio(0) + rocdl.s_barrier() + + if blockscale: + # Undo the rescaling invariant: the accumulator holds the true sum + # divided by the last K-block's scale. store_c._preapplied makes the + # epilogue skip its own scale loads. + c00_frag, c01_frag, c10_frag, c11_frag = bsk.final_scale( + (c00_frag, c01_frag, c10_frag, c11_frag), + prev_scales, + idx_fn=mfma.idx, + n_tiles_b=N_TILES_B, + ) + + wave_n_offset = wave_n * (N_TILES_B * 16) + wave_m_offset = wave_m * (N_TILES_A * 16) + base_row = block_m * BLOCK_M + wave_m_offset + base_col = block_n * BLOCK_N + wave_n_offset + + # Close the half-wave barrier pairing the prologue opened. Its + # `if wave_m == 1: s_barrier()` gives waves 4-7 one extra barrier, and + # s_barrier is a counting rendezvous, so from then on waves 0-3 run one + # phase ahead -- which is the intended stagger for the double-buffered + # main loop, and wrong for anything that comes after it. gcnasm closes + # the pair here too and does it unconditionally + # (opus_gemm_a2a_lsa/gemm_a16w16_quad_subtile_kernel_template.hpp:693, + # outside its `if constexpr (ChunkFused)`); aiter's 8-wave kernel opens + # the stagger at gemm_a8w8_8wave.py:449 and never closes it. + # + # Not gated on `fuse`: with the offset live, `wait_barrier(0)` after + # store_c does not mean "every wave's C tile has retired" -- it + # rendezvouses waves 0-3, which hold thread 0 and therefore the counter + # and the transfer, with waves 4-7 still sitting at the previous barrier, + # before their stores. That is what made --chunks 2 fail one run in + # three. The split path has nothing after store_c so it cannot observe + # the imbalance today, but leaving the counts unbalanced makes the next + # thing added after the epilogue silently wrong, which is exactly how + # this bug got here. + if wave_m == 0: + rocdl.s_barrier() + + if const_expr(hoist_scales) and hasattr(store_c, "store_all"): + store_c.store_all( + [ + (c00_frag, 0, 0), + (c01_frag, 0, 1), + (c10_frag, 1, 0), + (c11_frag, 1, 1), + ], + base_row, + base_col, + LDS_BLOCK_M, + LDS_BLOCK_N, + ) + else: + store_c.store(c00_frag, base_row + 0, base_col + 0) + store_c.store(c01_frag, base_row + 0, base_col + LDS_BLOCK_N) + store_c.store(c10_frag, base_row + LDS_BLOCK_M, base_col + 0) + store_c.store(c11_frag, base_row + LDS_BLOCK_M, base_col + LDS_BLOCK_N) + # ---- end pinned copy ---- + + if const_expr(direct_lsa): + # Publish the peer stores from the blocks that made them. This is the + # one thing Direct LSA cannot inherit from the split path: in the LSA + # 2-stage all-reduce the producing kernel fences (ar/kernels_lsa.py, + # every lane of every block, right after the tmp stores), whereas + # here the producer is the GEMM and the only fence was in the + # separate barrier kernel -- one block, hence one XCD's L2 out of + # eight. The other seven XCDs kept the peer-homed lines dirty and the + # LSA flag, being a system-scope atomic, overtook them. + wait_barrier(0) + # leaderOnly is legal here only because the barrier pair above is + # now closed: wait_barrier(0) really does mean every wave's stores + # have retired into this CU's L2, so one wave writing it back covers + # all eight. It was measured as incorrect before that fix, which is + # what made it look like a per-wave release was mandatory. + raw_cco.cco_system_fence(fx.Int32(1 if direct_fence == "leader" else 0)) + + if const_expr(fuse and not direct_lsa): + # Retire every lane's C stores, then agree block-wide that they are + # retired. `wait_barrier` is aiter's own `s_waitcnt vmcnt(0); + # s_barrier` pair, reused so the tail matches the main loop's idiom. + wait_barrier(0) + # Then publish them. This fence is not optional: dropping it -- on the + # theory that the copy engine reads through our own L2 -- still + # validates on 2 ranks but takes relL2 from 2.35e-3 to 3.76e-3 on 8, + # i.e. the engine reads some tiles stale. Every lane fences, which is + # the form cco_device_wrapper.cpp:105-111 documents for producer + # lanes; leaderOnly orders thread 0 only and is explicitly not a + # substitute for a release issued by every producer. + if const_expr(fence == "all"): + raw_cco.cco_system_fence(fx.Int32(0)) + elif const_expr(fence == "leader"): + raw_cco.cco_system_fence(fx.Int32(1)) + elif const_expr(fence in ("agent", "wt-agent", "nt-agent", "raw-wt-agent")): + release_fence("agent") + # "writethrough": nothing to do -- the stores already went to memory, + # and wait_barrier(0) above retired them. + + # Both bases are rank-local and constant-offset, so they are the same + # for every thread; taking them here rather than inside the `if` also + # keeps the window out of the if's captured state, which has to be + # single MLIR values (see cco.CachedWindow). + ctr_base = fx.Int64(w_pre.lsa_ptr(rank, counter_off)) + lock_base = fx.Int64(w_pre.lsa_ptr(rank, lock_off)) + sdma = cco.DevComm(dev_comm).sdma() + if fx.thread_idx.x == 0: + # One release per *block* instead of per wave. cco's leaderOnly + # form was already shown incorrect, but that went through an + # extern wrapper; emitting the fence inline rules out the + # compiler having sunk it past the atomic. + if const_expr(fence in ("agent-leader", "raw-wt-leader")): + release_fence("agent") + dest = block_m // fx.Int32(m_tiles_per_peer) + chunk = (block_m % fx.Int32(m_tiles_per_peer)) // fx.Int32( + m_tiles_per_chunk + ) + slot = dest * fx.Int32(chunks) + chunk + ctr = signal_ptr(ctr_base + fx.Int64(slot) * fx.Int64(4)) + seq = fx.Int32( + atomic_add_u32(ctr, 1, ordering=atomic_order) + ) + fx.Int32(1) + # Monotonic: the counters are never reset, so "last tile of this + # chunk, this epoch" is a modulo test rather than a compare. Same + # reason the barrier flags in ar/kernels_lsa are never reset -- it + # is what makes graph replay behave like a fresh launch. + if seq % fx.Int32(tiles_per_chunk) == fx.Int32(0): + if const_expr(emit_put) and dest != fx.Int32(rank): + # One queue per destination. Two chunks of the same + # destination can be elected at nearly the same moment and + # will then post to the same queue; that serialises the + # short issue only, and they share one xGMI link either + # way, so a queue each would buy nothing. + off = fx.Int64(chunk) * fx.Int64(chunk_bytes) + lock = signal_ptr(lock_base + fx.Int64(dest) * fx.Int64(4)) + if const_expr(chunks > 1): + _acquire_peer_lock(lock) + sdma.put( + dest, + win, + fx.Int64(my_recv_slot) + off, + win, + fx.Int64(in_off) + + fx.Int64(dest) * fx.Int64(slice_bytes) + + off, + fx.Int64(chunk_bytes), + # Queues are per (source, destination) pair, so + # every peer gets its own even at sdma_queues=1; + # allocating one per peer *per pair* only wastes + # hardware queues. See build_sdma_phases. + dest % fx.Int32(sdma_queues), + coop=cco.CoopScope.THREAD, + signal=False, + ) + if const_expr(chunks > 1): + atomic_store_u32(lock, 0) + # gcnasm closes its ChunkFused epilogue with a barrier here; its + # README lists removing it as a rejected experiment that deadlocked. + fgpu.barrier() + + @flyc.jit + def launch( + A: fx.Tensor, + B_T: fx.Tensor, + C: fx.Tensor, + A_scale: fx.Tensor, + B_scale: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + dev_comm: Int64, + win: Int64, + stream: fx.Stream = fx.Stream(None), + ): + grid_x = ceildiv(c_m, BLOCK_M) * ceildiv(c_n, BLOCK_N) + kernel_gemm_scatter( + A, + B_T, + C, + A_scale, + B_scale, + c_m, + c_n, + dev_comm, + win, + value_attrs={ + "rocdl.waves_per_eu": waves_per_eu, + "rocdl.flat_work_group_size": "512,512", + }, + ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) + + return launch + + +__all__ = ["compile_fused_gemm_scatter", "BLOCK_K"] diff --git a/python/mori/ops/gemm_ar/kernels_lsa.py b/python/mori/ops/gemm_ar/kernels_lsa.py new file mode 100644 index 000000000..de169f09a --- /dev/null +++ b/python/mori/ops/gemm_ar/kernels_lsa.py @@ -0,0 +1,348 @@ +# 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. +"""LSA all-reduce in FlyDSL: 1-stage and 2-stage, mirroring aiter's kernels. + +Transport is plain load/store against peer pointers obtained from cco's flat +symmetric VA (``Window.lsa_ptr``) -- the same mechanism aiter's custom all-reduce +uses via HIP IPC handles, so this is a substrate swap, not an algorithm change. + +Protocol parity with ``aiter/csrc/include/custom_all_reduce.cuh``: + +* per-(block, peer) signal slots and a **monotonic** ``_flag[block]`` that is + never reset, which is what makes CUDA-graph replay safe (``start_sync`` + :160-201, ``end_sync`` :204-240); +* the signal store is uncached (SC0|SC1, i.e. system-visible) and the spin load + bypasses L2 (SC1) -- aiter uses ``__MEMORY_SCOPE_SYSTEM`` / ``_DEVICE`` for the + same reason; +* accumulation order matches aiter so results are bitwise comparable: 1-stage + walks peers 0..P-1 unrotated (its source comment: "we don't reorder the + address so the accumulation order is the same for all ranks"), 2-stage walks + ``(rank + i) % P`` since each element is reduced exactly once by its owner; +* stage 2 reuses stage 1's thread->index mapping, because cross-device + visibility is only guaranteed between threads with the same tid + (``custom_all_reduce.cuh:556-560``). + +Deliberate divergence: aiter's 2-stage stage 1 assigns one peer per warp +(``tid = blockIdx.x * tnum_gpu + lane_id``, ``warp_id`` selects the peer), stages +each pack through LDS and reduces on ``warp_id == 0`` only, at the cost of two +``__syncthreads()`` per iteration. Here each thread reads all peers straight into +registers, which needs no LDS and no block-wide sync and keeps 8x more loads in +flight per thread. That extra concurrency is why the grid has to be much smaller +than aiter's -- see ``layout.LSA_BLOCK_CAP``. Measured within 5-10% of aiter at +M >= 2048 and faster at M = 8192; see ``bench_ar.py`` for the table. + +What is *not* a legitimate divergence is the all-gather loop nesting -- see the +comment in ``ar_2stage``. Getting that wrong costs ~3x and is invisible in the +output, so it is a selectable variant rather than a deleted mistake. + +Each kernel is specialised on the Python ``rank`` so every window offset folds +to a constant, as ``examples/cco/python/07_flydsl_sdma`` does. +""" + +from __future__ import annotations + + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf +from flydsl.expr import gpu as fgpu +from flydsl.expr import range_constexpr +from flydsl.expr.typing import Int64 + +import mori.cco.device.flydsl as cco +from mori.cco.device.flydsl import _bindings as raw_cco + +from ._compat import ( + CM_CACHED, + buffer_load, + buffer_store, + create_buffer_resource_from_addr, + i32_type, + signal_load_u32, + local_load_u32, + local_store_u32, + signal_ptr, + wave_uniform_i64, + signal_store_u32, +) +from .layout import MAX_WORLD + + +def _spin_until(rsrc, flag): + """Block until the u32 at ``rsrc`` reaches ``flag`` (unsigned compare). + + Unsigned so a wrapped counter still terminates; ``>=`` so an increment can + never be missed, matching aiter's ``while (load < flag);``. + """ + i32 = i32_type() + first = signal_load_u32(rsrc) + first_v = first.ir_value() if hasattr(first, "ir_value") else first + loop = scf.WhileOp([i32], [first_v]) + cond = ir.Block.create_at_start(loop.before, [i32]) + body = ir.Block.create_at_start(loop.after, [i32]) + with ir.InsertionPoint(cond): + cur = fx.Int32(cond.arguments[0]) + should_wait = fx.Uint32(cur) < fx.Uint32(flag) + scf.ConditionOp(should_wait.ir_value(), [cond.arguments[0]]) + with ir.InsertionPoint(body): + nxt = signal_load_u32(rsrc) + scf.YieldOp([nxt.ir_value() if hasattr(nxt, "ir_value") else nxt]) + + +#: All-gather loop nestings that ``build_lsa_ar`` can emit. See ``ar_2stage``. +GATHER_VARIANTS = ("interleaved", "sequential") + + +def build_lsa_ar( + cfg, + rank: int, + *, + force_stage: int | None = None, + gather: str = "interleaved", +): + """Compile an LSA all-reduce launcher for ``cfg`` on ``rank``. + + Returns ``(run, stage)`` where ``run(dev_comm, win, stream=...)`` performs one + in-place all-reduce of the window's input region into its output region. + + ``gather`` selects the 2-stage all-gather loop nesting; see ``ar_2stage``. + """ + cfg.validate() + if gather not in GATHER_VARIANTS: + raise ValueError(f"gather must be one of {GATHER_VARIANTS}, got {gather!r}") + ws = cfg.world_size + stage = force_stage if force_stage is not None else cfg.stage + if stage not in (1, 2): + raise ValueError(f"stage must be 1 or 2, got {stage}") + + threads = cfg.threads + blocks = cfg.blocks + elem_dtype = fx.BFloat16 if cfg.elem_bytes == 2 else fx.Float32 + num_packs = cfg.num_packs + stride_packs = blocks * threads + + start_off, end_off, flag_off = cfg.start_off, cfg.end_off, cfg.flag_off + in_off, out_off, tmp_off = cfg.input_off, cfg.output_off, cfg.tmp_off + own_start, own_end = cfg.owner_pack_range(rank) + own_packs = own_end - own_start + part = cfg.packs_per_rank + + def _signal_and_wait(w, bid, arr_off, flag, tid): + """Barrier body for the ``tid < ws`` lanes: publish, then wait. + + Branch-free on purpose. FlyDSL only rewrites ``if`` inside the + ``@flyc.kernel`` function's own AST, so a runtime ``if`` in a helper is + evaluated as a Python bool and raises during tracing; the guard has to + stay in the kernel body. + """ + peer_arr = fx.Int64(w.lsa_ptr(tid, arr_off)) + mine = fx.Int64((bid * MAX_WORLD + rank) * 4) + signal_store_u32(signal_ptr(peer_arr + mine), flag) + + self_arr = fx.Int64(w.lsa_ptr(rank, arr_off)) + theirs = fx.Int64(bid * MAX_WORLD * 4) + fx.Int64(tid) * fx.Int64(4) + _spin_until(signal_ptr(self_arr + theirs), flag) + + def _next_flag(w, bid): + """``_flag[block] + 1`` -- monotonic, never reset (graph-replay safe).""" + base = fx.Int64(w.lsa_ptr(rank, flag_off)) + fx.Int64(bid) * fx.Int64(4) + rsrc = signal_ptr(base) + return fx.Int32(local_load_u32(rsrc)) + fx.Int32(1), rsrc + + # A 16B pack is always moved as 4 x i32, so buffer offsets are in i32 units + # regardless of the payload dtype -- for bf16 that pack holds 8 values, and + # indexing it in payload elements would stride twice as far as intended. + I32_PER_PACK = 4 + + def _unpack(rsrc, i32_off): + """One 16B pack -> a float32 vector of ``elems_per_pack`` lanes.""" + raw = fx.Vector(buffer_load(rsrc, i32_off, vec_width=4, dtype=i32_type())) + if elem_dtype is fx.Float32: + return raw.bitcast(fx.Float32) + return raw.bitcast(elem_dtype).to(fx.Float32) + + def _pack(acc): + """float32 vector -> the 4 x i32 that the payload dtype occupies.""" + if elem_dtype is fx.Float32: + return acc.bitcast(fx.Int32) + return acc.to(elem_dtype).bitcast(fx.Int32) + + def _acc_peers(rsrcs, pack_idx): + """fp32 sum of one 16B pack across every peer, in ``rsrcs`` order.""" + i32_off = pack_idx * I32_PER_PACK + acc = None + for i in range_constexpr(len(rsrcs)): + v = _unpack(rsrcs[i], i32_off) + acc = v if acc is None else acc + v + return acc + + @flyc.kernel(known_block_size=[threads, 1, 1]) + def ar_1stage(dev_comm: Int64, win: Int64): + bid = fx.block_idx.x + tid = fx.thread_idx.x + w = cco.CachedWindow(win) + + flag, flag_rsrc = _next_flag(w, bid) + if tid < ws: + _signal_and_wait(w, bid, start_off, flag, tid) + fgpu.barrier() + if tid == 0: + local_store_u32(flag_rsrc, flag) + + # Unrotated peer order: every rank accumulates identically, so all ranks + # produce bitwise identical output (aiter's 1-stage does the same). + ins = [ + create_buffer_resource_from_addr(wave_uniform_i64(w.lsa_ptr(p, in_off))) + for p in range(ws) + ] + out = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, out_off)) + ) + + gtid = bid * threads + tid + for pk in range(gtid, num_packs, stride_packs): + acc = _acc_peers(ins, pk) + buffer_store(_pack(acc), out, pk * I32_PER_PACK, cache_modifier=CM_CACHED) + + @flyc.kernel(known_block_size=[threads, 1, 1]) + def ar_2stage(dev_comm: Int64, win: Int64): + bid = fx.block_idx.x + tid = fx.thread_idx.x + w = cco.CachedWindow(win) + + flag_a, flag_rsrc = _next_flag(w, bid) + if tid < ws: + _signal_and_wait(w, bid, start_off, flag_a, tid) + fgpu.barrier() + if tid == 0: + local_store_u32(flag_rsrc, flag_a) + + # --- stage 1: reduce-scatter. I own packs [own_start, own_end) and read + # that slice from every peer, rotated so the 8 links are not all pointed + # at rank 0 at once. Rotation is safe for bitwise agreement because each + # pack has exactly one owner. + ins = [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr((rank + i) % ws, in_off)) + ) + for i in range(ws) + ] + tmp = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, tmp_off)) + ) + + gtid = bid * threads + tid + for pk in range(gtid, own_packs, stride_packs): + acc = _acc_peers(ins, own_start + pk) + buffer_store(_pack(acc), tmp, pk * I32_PER_PACK, cache_modifier=CM_CACHED) + + # Release barrier. Three things must happen in this order, and each was + # a real bug when missing: + # 1. block-wide barrier so every lane's tmp store is issued, AND so + # thread 0's `_flag` update from the first barrier is visible -- + # without it the other lanes re-read the stale flag, compute the + # same value again, and the second barrier is already satisfied + # (aiter's end_sync opens with __syncthreads() for exactly this); + # 2. a system fence, the release that flushes those cached tmp stores + # out where peers reading over xGMI can see them (aiter gets this + # from __ATOMIC_RELEASE at __MEMORY_SCOPE_SYSTEM on the signal); + # 3. only then publish the signal. + # leaderOnly=0: every lane fences. The cco ABI notes leaderOnly=1 fences + # thread 0 only and is not a substitute for a per-lane release. + fgpu.barrier() + raw_cco.cco_system_fence(fx.Int32(0)) + flag_b, flag_rsrc_b = _next_flag(w, bid) + if tid < ws: + _signal_and_wait(w, bid, end_off, flag_b, tid) + # Acquire half of the pair: invalidate lines this rank may still hold + # for a peer's tmp. The barrier's atomic load is agent-scope monotonic, + # which orders but does not invalidate, so without this the gather can + # read stale L2. aiter gets it from __ATOMIC_ACQUIRE on its spin load. + raw_cco.cco_system_fence(fx.Int32(0)) + fgpu.barrier() + if tid == 0: + local_store_u32(flag_rsrc_b, flag_b) + + # --- stage 2: all-gather. Thread `gtid` reads exactly the tmp indices it + # wrote in stage 1 -- cross-device visibility only holds between threads + # with the same tid. + # + # The loop nesting is the whole ballgame here. `interleaved` puts the + # peer loop *inside* the index loop and unrolls it (aiter's + # `#pragma unroll for i < ngpus` at custom_all_reduce.cuh:566), so one + # iteration has 8 loads outstanding against 8 different peers and all 8 + # xGMI links stream at once. `sequential` drains one peer completely + # before starting the next, which pins the all-gather to a single link: + # 7 remote passes x (nbytes/8 / 55GB/s) instead of one. At M=4096 that is + # the difference between ~133us and ~660us for this half of the kernel. + # Kept selectable because it is the cleaner-looking code and the trap is + # not obvious from reading it. + out = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, out_off)) + ) + srcs = [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr((rank + j) % ws, tmp_off)) + ) + for j in range(ws) + ] + # validate() forces num_packs % world == 0, so every rank owns exactly + # `part` packs and aiter's `idx < part || from == ngpus-1` guard is moot. + if gather == "interleaved": + for pk in range(gtid, part, stride_packs): + # Issue every load before any store so the 8 links overlap; a + # load/store pair per peer would serialise on each s_waitcnt. + vals = [ + buffer_load( + srcs[j], pk * I32_PER_PACK, vec_width=4, dtype=i32_type() + ) + for j in range_constexpr(ws) + ] + for j in range_constexpr(ws): + dst = ((rank + j) % ws) * part + pk + buffer_store( + vals[j], out, dst * I32_PER_PACK, cache_modifier=CM_CACHED + ) + else: + for j in range_constexpr(ws): + dst_base = ((rank + j) % ws) * part + for pk in range(gtid, part, stride_packs): + v = buffer_load( + srcs[j], pk * I32_PER_PACK, vec_width=4, dtype=i32_type() + ) + buffer_store( + v, out, (dst_base + pk) * I32_PER_PACK, cache_modifier=CM_CACHED + ) + + kernel = ar_1stage if stage == 1 else ar_2stage + + @flyc.jit + def run(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): + kernel(dev_comm, win).launch( + grid=(blocks, 1, 1), block=[threads, 1, 1], stream=stream + ) + + return run, stage + + +__all__ = ["build_lsa_ar"] diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py new file mode 100644 index 000000000..db375446f --- /dev/null +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -0,0 +1,1182 @@ +# 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. +"""SDMA all-reduce in FlyDSL: reduce-scatter push, local reduce, all-gather push. + +Same algorithm and the same window as ``kernels_lsa``'s 2-stage, and the same +``2(P-1)/P`` bytes on the wire. The only difference is who moves them: here the +copy engines do, so the CUs are free for compute during the transfers. That is +the property the GEMM fusion in step 3 is after -- raw bandwidth is not expected +to beat LSA. + +Three kernels rather than one, because both transfers need a *device-wide* +ordering point (every block must see the landed data), and a kernel boundary is +the only one that is free and graph-capturable: + + 1. scatter lane p pushes my rows of peer p's slice into p's recv slot, + drains queue p, then bumps p's `start[0][rank]` flag + 2. reduce fp32-sum my own slice plus the P-1 landed slices into + `output + rank*slice` + 3. gather lane p pushes my reduced slice into peer p's output at the same + offset, drains, then bumps p's `end[0][rank]` flag + +Rank-to-rank arrival is signalled over LSA, not SDMA, because FlyDSL's ``Sdma`` +exposes no ``wait_signal`` and its ``put`` wrapper hard-codes ``remoteSignal = +false`` (``src/cco/device/cco_device_wrapper.cpp:140``), so a peer cannot observe +an SDMA signal at all. This is the ``opus_sdma_a2a_quiet_notify_kernel`` pattern +from gcnasm: no per-PUT signal, ``quiet_queue`` on the sender, then an LSA atomic +into the peer's flag. The flags, their monotonic epoch and the spin loop are all +reused verbatim from ``kernels_lsa`` -- only row 0 of the signal array is used, +since the pushes are issued by one warp. + +Transfer sizing follows the copy engine's one real lever: bytes per op. Each peer +gets exactly one ``put`` of ``slice_bytes``, never split across queues -- one peer +is one xGMI link, so splitting only pays the ~2us per-packet cost again. At +[4096, 7168] bf16 a slice is 7.3MB (~59 GB/s territory); at [64, 7168] it is +114KB, below the ~256KB knee, so decode sizes are expected to lose to LSA. +cco's own chunking is not a factor: ``CCO_SDMA_MAX_COPY_BYTES`` is 1GB, so every +slice we push is a single packet. + +Measured on 8x MI355X, [M, 7168] bf16, against the LSA kernel (full table in +``bench_ar.py``): within 5% of LSA for M >= 2048, and faster than aiter at +M = 8192. What is left is mostly a fixed ~40us floor -- two transfer phases times +the engine's ~6us dispatch -- so decode sizes still lose by ~1.5x, as expected +for 114KB slices. + +Getting there needed one non-obvious thing: the reduce kernel must **not** share +the LSA all-reduce's grid. Its loads are local HBM, not xGMI, so the cap that +throttles outstanding xGMI requests starves it instead -- 24 blocks left it at +1.11 TB/s on an ~8 TB/s part. Sized independently it runs at 4.9 TB/s, worth 48us +at M=4096 (342 -> 294) and turning a 1.20x deficit against LSA into 1.03x. See +``SDMA_REDUCE_BLOCK_CAP``. + +The other reason this backend exists is that its transfers occupy one warp of one +block instead of the whole grid, which is what makes an in-GEMM epilogue possible +in step 3. +""" + +from __future__ import annotations + +import os + + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr import gpu as fgpu +from flydsl.expr.typing import Int64 + +import mori.cco.device.flydsl as cco +from mori.cco.device.flydsl import _bindings as raw_cco + +from flydsl._mlir import ir + +from ._compat import ( + CM_CACHED, + CM_SC0_SC1, + CM_SC1, + atomic_add_u32, + atomic_store_u32, + buffer_load, + buffer_store, + create_buffer_resource_from_addr, + i32_type, + local_load_u32, + local_store_u32, + signal_ptr, + signal_store_u32, + wave_uniform_i64, +) +from .kernels_lsa import _spin_until +from ._gemm_a8w8_8wave import wait_barrier +from .layout import PACK_BYTES +from .kernels_fused import _acquire_peer_lock + +#: e4m3's largest finite magnitude. ``scale = amax / FP8_E4M3_MAX`` puts a row's +#: biggest element exactly at the top of the range. +FP8_E4M3_MAX = 448.0 + +#: Grid for the LSA pull gather. See the sweep at its use site; it is +#: deliberately not LSA_BLOCK_CAP. +PULL_BLOCKS = 64 + + +#: Wave width, and the width of the amax butterfly below (2**6). +WAVE = 64 +WAVE_LOG2 = 6 + +#: Elements one lane converts per chunk. 16, so that the *fp8* side is a full +#: 16B access (one dwordx4) rather than 8B: at 8 the dequantise read was half +#: width while its bf16 write was full width, and the kernel ran at 3.4 TB/s +#: against roughly 8 TB/s of HBM. The bf16 side is then 32B, issued as two +#: dwordx4. Both stay contiguous across the wave. +#: +#: The conversion helpers still work 8 at a time -- that is the width of +#: `v_cvt_pk_fp8_f32`'s natural grouping here -- so a chunk is two of them. +CHUNK_ELEMS = 16 +CVT_ELEMS = 8 + + +def _raw_value(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + +def _bpermute_f32(value, src_lane): + """``ds_bpermute`` on an f32, through its i32 bit pattern. Byte-addressed.""" + i32 = ir.IntegerType.get_signless(32) + raw = fx.rocdl.ds_bpermute( + i32, + _raw_value(fx.Int32(src_lane) * fx.Int32(4)), + _raw_value(value.bitcast(fx.Int32)), + ) + return fx.Int32(raw).bitcast(fx.Float32) + + +def _pack8_fp8(v8): + """Eight scaled f32 -> eight e4m3 bytes, as a ``vector<2xi32>``. + + ``v_cvt_pk_fp8_f32`` rather than the one-shot ``pk8`` form: the latter is + gfx1250, and this has to run on gfx950. Four instructions instead of one, + which is nothing in a kernel this memory-bound. Note that + ``arith.truncf`` to fp8 is *not* an option -- it has no LLVM lowering and + fails in the translation pass, not at trace time. + + ``old`` threads the two halves of a dword through one register: word_sel=0 + writes the low 16 bits, word_sel=1 the high. + """ + i32 = ir.IntegerType.get_signless(32) + poison = fx.Int32(0) + words = [] + for half in range_constexpr(2): + acc = _raw_value(poison) + for pair in range_constexpr(2): + i = half * 4 + pair * 2 + acc = fx.rocdl.cvt_pk_fp8_f32( + i32, + _raw_value(v8[i]), + _raw_value(v8[i + 1]), + acc, + pair == 1, + ) + words.append(fx.Int32(acc)) + return fx.Vector.from_elements(words, fx.Int32) + + +def _unpack8_fp8(v2i32): + """Eight e4m3 bytes in a ``vector<2xi32>`` -> eight f32. The inverse.""" + f32x2 = ir.VectorType.get([2], ir.F32Type.get()) + out = [] + for half in range_constexpr(2): + word = _raw_value(v2i32[half]) + for sel in range_constexpr(2): + pair = fx.Vector(fx.rocdl.cvt_pk_f32_fp8(f32x2, word, sel == 1)) + out.append(pair[0]) + out.append(pair[1]) + return out + + +def _wave_amax(value, lane): + """Max of ``value`` over the wave, as a butterfly -- no LDS, no barrier. + + Every lane ends with the same result, which is what lets each of them scale + its own chunk without a second broadcast. + """ + acc = value + for k in range_constexpr(WAVE_LOG2): + acc = acc.maximumf(_bpermute_f32(acc, lane ^ fx.Int32(1 << k))) + return acc + + +#: The push warp. One lane per peer, each on its own queue, so the lanes post in +#: parallel (cco's ccoSdmaThreadIndependent shape). 64 is the wave size; only the +#: first `world_size` lanes do anything. +PUSH_THREADS = 64 + + +def build_sdma_phases( + cfg, + rank: int, + *, + queues: int = 8, + signal: bool = False, + reduce_blocks=None, + reduce_self_from_recv: bool = False, + recv_uncached: bool = False, + fuse_quantize: bool = False, + fuse_reduce_push: bool = False, + publish: str = "writethrough", +): + """Compile the phases separately, so the fused GEMM can reuse the tail. + + Returns ``{"scatter", "drain", "reduce", "gather"}``, each a launcher taking + ``(dev_comm, win, stream=...)``. ``scatter`` pushes *and* drains; ``drain`` + only drains and barriers, and exists for the fused path, where the pushes + were already issued from inside the GEMM epilogue. + + ``queues`` must match ``reqs.sdma_queue_count``. Queue ids are taken modulo + it, and the hardware queues are per **(source, destination) pair**, so even + ``queues=1`` still gives every peer its own queue -- which is all this + pipeline needs, since one destination is one xGMI link and cco's rule is + about concurrently-issuing *warps* per queue, of which there is one. + + Prefer ``queues=1`` when sharing the process with something else. With + ``queues=world_size`` a rank creates ``world_size`` hardware queues per peer + and touches exactly one of them: 56 queues to use 7. Standalone that is + merely wasteful, but inside a server that already holds SDMA engines for its + own copies, ``hsaKmtCreateQueueExt`` starts failing + (``anvil.cpp:237``). + + ``signal`` selects ``put``'s trailing local ATOMIC. It is off by default: + ``quiet``/``quietQueue`` drain the queue's read pointer and are documented as + "independent of the signals" (``include/mori/cco/cco.hpp:1707``), and nothing + here polls a signal, so the packet would be pure overhead. Exposed because + FlyDSL's own ``Sdma._xfer`` docstring claims the opposite -- that a no-signal + put "cannot be drained by quiet" -- and that claim is worth being able to + disprove by measurement rather than by reading. + """ + cfg.validate() + ws = cfg.world_size + if cfg.recv_slots < ws: + raise ValueError( + f"SDMA needs a landing slot per peer: build the ArConfig with " + f"recv_slots={ws} (got {cfg.recv_slots})" + ) + if queues < 1: + raise ValueError(f"queues must be >= 1, got {queues}") + if cfg.fp8_scatter: + # The regions exist (layout.py sizes them), but nothing writes the + # quantised payload yet: on the fused path that is the GEMM epilogue's + # job, and on the standalone path it needs its own kernel. Refuse + # rather than read an fp8-sized region that still holds bf16, which + # faults somewhere unrelated. + raise NotImplementedError( + "scatter_dtype='fp8' is not wired yet; the gather leg is " + "(gather_dtype='fp8'), and it is where the time is" + ) + + threads = cfg.threads + # The reduce is a *local HBM* kernel, so it must not inherit the grid the LSA + # all-reduce uses. LSA_BLOCK_CAP is small on purpose -- it caps the number of + # outstanding xGMI requests -- but here every load is local, and 24x512 + # threads x 8 packs is only 1.6MB in flight, about a fifth of what it takes to + # cover HBM latency. Sized independently, and swept by bench_ar.py. + red_blocks = reduce_blocks if reduce_blocks else cfg.reduce_blocks + red_stride = red_blocks * threads + elem_dtype = fx.BFloat16 if cfg.elem_bytes == 2 else fx.Float32 + part = cfg.packs_per_rank + slice_bytes = cfg.slice_bytes + + start_off, end_off, flag_off = cfg.start_off, cfg.end_off, cfg.flag_off + lock_off = cfg.lock_off + in_off, out_off = cfg.input_off, cfg.output_off + # Constant-folded because `rank` is a Python int: the slice I own, and the + # slot I occupy in every peer's recv region. + my_slice_off = rank * slice_bytes + my_recv_slot = cfg.recv_slot_off(rank) + + I32_PER_PACK = 4 # a 16B pack always moves as 4 x i32, whatever the payload + + def _next_flag(w, bid): + """``_flag[block] + 1`` -- monotonic, never reset (graph-replay safe).""" + base = fx.Int64(w.lsa_ptr(rank, flag_off)) + fx.Int64(bid) * fx.Int64(4) + rsrc = signal_ptr(base) + return fx.Int32(local_load_u32(rsrc)) + fx.Int32(1), rsrc + + def _signal_and_wait(w, arr_off, flag, tid): + """Row-0 barrier: publish to peer ``tid``, then wait on ``tid``'s slot.""" + peer_arr = fx.Int64(w.lsa_ptr(tid, arr_off)) + signal_store_u32(signal_ptr(peer_arr + fx.Int64(rank * 4)), flag) + self_arr = fx.Int64(w.lsa_ptr(rank, arr_off)) + _spin_until(signal_ptr(self_arr + fx.Int64(tid) * fx.Int64(4)), flag) + + def _push_kernel( + arr_off, + dst_off_of_peer=None, + src_off_expr=None, + nbytes=None, + extra=None, + ): + """A push phase: one put per peer, drain, then the cross-rank barrier. + + ``dst_off_of_peer`` is a constant byte offset *in the destination's* + window; ``src_off_expr(tid)`` builds the offset in mine, which may depend + on which peer the lane is serving. Passing neither emits a drain-only + kernel -- the pushes are assumed already issued elsewhere. + + ``extra`` is an optional second ``(dst_off, src_off_expr, nbytes)`` put on + the same queue, for the fp8 wire's scales. It rides the same queue rather + than its own so it cannot overtake the payload, and it is ~8 KiB against + 14 MiB -- about 2us of packet cost, against the ~210us the halved payload + saves. + """ + pushes = dst_off_of_peer is not None + put_bytes = slice_bytes if nbytes is None else nbytes + + @flyc.kernel(known_block_size=[PUSH_THREADS, 1, 1]) + def push(dev_comm: Int64, win: Int64): + tid = fx.thread_idx.x + w = cco.CachedWindow(win) + sdma = cco.DevComm(dev_comm).sdma() + + flag, flag_rsrc = _next_flag(w, 0) + if tid < ws: + if tid != rank: + if const_expr(pushes): + # Lane `tid` owns peer `tid` and queue `tid`: distinct + # queues per lane, so the posts do not serialise on one + # commit chain. + sdma.put( + tid, + win, + fx.Int64(dst_off_of_peer), + win, + src_off_expr(tid), + fx.Int64(put_bytes), + tid % fx.Int32(queues), + coop=cco.CoopScope.THREAD, + signal=signal, + ) + if const_expr(extra is not None): + sdma.put( + tid, + win, + fx.Int64(extra[0]), + win, + extra[1](tid), + fx.Int64(extra[2]), + tid % fx.Int32(queues), + coop=cco.CoopScope.THREAD, + signal=signal, + ) + sdma.quiet_queue(tid, tid % fx.Int32(queues)) + # Release before publishing arrival. The bytes were moved by the + # copy engine rather than by this CU, so there is nothing of ours + # to flush; the fence is here to keep the flag store from being + # hoisted above the drain, which is the one ordering the receiver + # depends on. + raw_cco.cco_system_fence(fx.Int32(0)) + _signal_and_wait(w, arr_off, flag, tid) + fgpu.barrier() + if tid == 0: + local_store_u32(flag_rsrc, flag) + + return push + + # Scatter. Peer p owns slice p, so what p needs from me is *p's* slice range + # of my input -- the source offset moves with the lane. It lands in the slot p + # reserves for me, which is at the same byte offset in every rank's window. + scatter = _push_kernel( + start_off, + dst_off_of_peer=my_recv_slot, + src_off_expr=lambda tid: fx.Int64(in_off) + + fx.Int64(tid) * fx.Int64(slice_bytes), + ) + + @flyc.kernel(known_block_size=[threads, 1, 1]) + def sdma_reduce(dev_comm: Int64, win: Int64): + """Sum my own slice plus the P-1 landed slices into my part of output. + + Peer order is ``(rank + j) % ws`` to match ``kernels_lsa``'s 2-stage, so + the two backends are bitwise identical and the test can assert + ``rel_l2 == 0`` rather than a tolerance. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.CachedWindow(win) + + # j=0 is my own contribution. Normally it is still sitting in my input + # region; with the Direct-LSA fused GEMM the epilogue wrote it into my own + # recv slot instead, along with everyone else's. + self_off = ( + cfg.recv_slot_off(rank) if reduce_self_from_recv else in_off + my_slice_off + ) + srcs = [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, self_off)) + ) + ] + [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, cfg.recv_slot_off((rank + j) % ws))) + ) + for j in range(1, ws) + ] + out = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, out_off + my_slice_off)) + ) + + gtid = bid * threads + tid + for pk in range(gtid, part, red_stride): + i32_off = pk * I32_PER_PACK + acc = None + for j in range_constexpr(ws): + # recv is filled by a *peer*. When the copy engine wrote it the + # lines are coherent, but when a peer's CUs stored into it over + # xGMI our own L2 is never invalidated, so a cached load here can + # return the previous iteration's bytes. SC1 skips L2 for those. + raw = fx.Vector( + buffer_load( + srcs[j], + i32_off, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_SC1 if recv_uncached else CM_CACHED, + ) + ) + v = ( + raw.bitcast(fx.Float32) + if elem_dtype is fx.Float32 + else raw.bitcast(elem_dtype).to(fx.Float32) + ) + acc = v if acc is None else acc + v + packed = ( + acc.bitcast(fx.Int32) + if elem_dtype is fx.Float32 + else acc.to(elem_dtype).bitcast(fx.Int32) + ) + buffer_store(packed, out, i32_off, cache_modifier=CM_CACHED) + + # --- reduce with the gather's puts fired from inside it ------------------- + # + # The gather push sends *my own* reduced slice, so unlike the pull it has no + # cross-rank dependency: nothing outside this rank has to happen before the + # bytes can leave. That makes it the same shape as the GEMM's fused scatter + # -- produce a contiguous range, elect one block, hand it to a copy engine -- + # and it reuses that machinery: the submit lock and a monotonic counter. + # + # The traversal is band-major rather than flat grid-strided, and that is the + # whole trick. Under a flat stride every block touches every band, so no + # band's counter fills until the kernel is nearly done and there is nothing + # to publish early. Sweeping band by band keeps the *inner* map -- and so the + # coalescing -- exactly as it was, and only reorders which region the grid + # covers first. + # + # Cost of election is one atomic per (block, band), i.e. red_blocks * + # gather_bands per launch. Counting per pack instead would be ~1.8M atomics. + + if fuse_reduce_push: + if publish not in ("fence", "writethrough"): + raise ValueError(f"publish must be fence or writethrough, got {publish!r}") + if cfg.fp8_gather: + raise ValueError( + "fuse_reduce_push is for the bf16 gather. The fp8 gather either " + "pushes from `gout` after a separate quantise, or is an LSA pull " + "reading peers -- neither is this shape." + ) + if part % cfg.gather_bands: + raise ValueError( + f"gather_bands={cfg.gather_bands} must divide the {part} packs " + f"per rank" + ) + packs_per_band = part // cfg.gather_bands + band_bytes = packs_per_band * PACK_BYTES + gather_counter_off = cfg.gather_counter_off + + @flyc.kernel(known_block_size=[threads, 1, 1]) + def sdma_reduce_push(dev_comm: Int64, win: Int64): + """`sdma_reduce`, publishing each band to the copy engines as it lands. + + Identical arithmetic and identical output to `sdma_reduce`; the only + difference is when the bytes start moving. The following `gather` phase + becomes drain-only. + + **Off by default: it reaches parity, not a win.** At [16384, 7168] + K=2048 on 8 ranks, whole fused layer, against 1148.7us unfused: + + bands 1 4 8 16 32 + publish=wt 1162.1 1157.2 1160.3 1190.1 1323.5 + publish=fence 1227.3 1380.1 1630.3 2137.2 3026.7 + + The interesting part is the gap between those two rows, because it is + not what it first looks like. + + Handing a range to a copy engine needs a release: the engine reads over + the fabric, not through a CU's cache, so `s_waitcnt vmcnt(0)` alone is + not enough -- that only retires the stores as far as *this XCD's* L2. + But there are two ways to pay for it: + + publish="fence" leave the bytes in L2 and release to system + scope afterwards. That is L2 writeback work + charged once per block per band -- 256 x bands + of it, ~61us per band, and the reduce is only + 42us in total, so it can never be repaid. + publish="writethrough" store with sc0+sc1 so the bytes never stop in + L1 or L2, and the waitcnt above *is* the + release. Same trick as the GEMM's + _WriteThroughStoreC under fence="writethrough". + + Write-through is free here, measured on its own: applying that store + policy to the plain unfused reduce moves it 1153.4 -> 1148.7us, i.e. + nothing. This output is written once and nothing local reads it again + before the gather, so holding it in L2 was buying nothing. + + What is left after that fix is small on both sides and nearly cancels. + At bands=1 the publish happens only after the whole slice is done, so it + carries the mechanism's cost with none of its benefit: 1162.1 against + 1148.7, about 13us for the per-band 256-block `wait_barrier`, the + counter atomic and the elected block's locked puts. Going to 4 bands + buys back about 5us of overlap and then the sync cost takes over again. + + Dropping the release entirely is not an option even though it briefly + looks like one: with cached stores and no fence the kernel reaches + 1157.9us, but at 32 bands it produced relL2 1.8e-2 against the 2.35e-3 + floor, differing per rank. The same 32 bands are exact under either + correct publish mode, which is what rules out an indexing bug. + + Why the GEMM's fused scatter wins where this does not: there the + publish is amortised against a 437us transfer hidden behind a + compute-bound GEMM, here against 42us of bandwidth-saturated reduce. + The mechanism pays when what is hidden is much larger than the cost of + publishing it. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.CachedWindow(win) + sdma = cco.DevComm(dev_comm).sdma() + + self_off = ( + cfg.recv_slot_off(rank) if reduce_self_from_recv else in_off + my_slice_off + ) + srcs = [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, self_off)) + ) + ] + [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, cfg.recv_slot_off((rank + j) % ws))) + ) + for j in range(1, ws) + ] + out = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, out_off + my_slice_off)) + ) + + # Hoisted above the band loop, and out of the `if tid == 0` below: both + # are rank-local at a constant offset, so they are invariant across bands + # and identical on every thread. It also keeps the window out of the if's + # captured state, which FlyDSL requires to be single MLIR values. + ctr_base = fx.Int64(w.lsa_ptr(rank, gather_counter_off)) + lock_base = fx.Int64(w.lsa_ptr(rank, lock_off)) + + gtid = bid * threads + tid + for band in range_constexpr(cfg.gather_bands): + base = band * packs_per_band + for pk in range(gtid + base, base + packs_per_band, red_stride): + i32_off = pk * I32_PER_PACK + acc = None + for j in range_constexpr(ws): + raw = fx.Vector( + buffer_load( + srcs[j], + i32_off, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_SC1 if recv_uncached else CM_CACHED, + ) + ) + v = ( + raw.bitcast(fx.Float32) + if elem_dtype is fx.Float32 + else raw.bitcast(elem_dtype).to(fx.Float32) + ) + acc = v if acc is None else acc + v + packed = ( + acc.bitcast(fx.Int32) + if elem_dtype is fx.Float32 + else acc.to(elem_dtype).bitcast(fx.Int32) + ) + buffer_store( + packed, + out, + i32_off, + cache_modifier=( + CM_SC0_SC1 if publish == "writethrough" else CM_CACHED + ), + ) + + # `wait_barrier(0)` is s_waitcnt vmcnt(0) + s_barrier: it retires + # this block's stores, where a bare s_barrier would only prove they + # were issued. But "retired" means reached the coherence point the + # store's own cache policy targets, and for an ordinary cached store + # that is just *this XCD's* L2 -- which a copy engine, reading over + # the fabric rather than through a CU's cache, cannot see. + # + # Two ways to close that gap, and they cost very differently: + # + # "fence" leave the bytes in L2 and release to system scope + # afterwards. That is L2 writeback work, and it is + # charged once per block per band. + # "writethrough" store with sc0+sc1 so they never stop in L1 or L2, + # and the waitcnt above is then the whole release. + # Same trick as the GEMM's _WriteThroughStoreC + # under fence="writethrough". + # + # Writethrough should be the right shape here: this output is + # written once and nothing local reads it again before the gather, + # so holding it in L2 buys nothing to begin with. + wait_barrier(0) + if const_expr(publish == "fence"): + raw_cco.cco_system_fence(fx.Int32(0)) + if tid == fx.Int32(0): + ctr = signal_ptr(ctr_base + fx.Int64(band) * fx.Int64(4)) + seq = fx.Int32(atomic_add_u32(ctr, 1)) + fx.Int32(1) + # Monotonic, never reset, so "last block of this band, this + # epoch" is a modulo rather than a compare -- the property that + # makes graph replay behave like a fresh launch. + if seq % fx.Int32(red_blocks) == fx.Int32(0): + for j in range_constexpr(1, ws): + dest = (rank + j) % ws + # Two bands can be elected at nearly the same moment and + # would then post to the same per-destination queue. + lock = signal_ptr(lock_base + fx.Int64(dest) * fx.Int64(4)) + _acquire_peer_lock(lock) + sdma.put( + dest, + win, + fx.Int64(out_off + my_slice_off + band * band_bytes), + win, + fx.Int64(out_off + my_slice_off + band * band_bytes), + fx.Int64(band_bytes), + dest % fx.Int32(queues), + coop=cco.CoopScope.THREAD, + signal=signal, + ) + atomic_store_u32(lock, 0) + + # --- fp8 wire: the two conversions that bracket the gather --------------- + # + # quantise runs after the reduce on my own slice; dequantise runs after the + # gather on everybody else's. My own rows never travel, so they stay the + # exact bf16 the reduce produced -- only the 7/8 that cross a link are + # rounded, and each of them exactly once. + + QUANT_WAVES = 4 # waves per block; one wave owns one row + QUANT_THREADS = QUANT_WAVES * WAVE + + if cfg.fp8_gather: + if cfg.n % (CHUNK_ELEMS * WAVE): + raise ValueError( + f"the fp8 wire gives one wave one row and converts " + f"{CHUNK_ELEMS} elements per lane per chunk, so n={cfg.n} must " + f"be a multiple of {CHUNK_ELEMS * WAVE}" + ) + chunks_per_row = cfg.n // (CHUNK_ELEMS * WAVE) + slice_rows = cfg.slice_rows + quant_blocks = max(1, min(1024, (slice_rows + QUANT_WAVES - 1) // QUANT_WAVES)) + # The pull gets its own, much smaller grid. These are *xGMI* reads, so + # the grid throttles outstanding remote requests rather than covering + # HBM latency -- the opposite of what quantize wants, which is why that + # one keeps the big grid. + # + # Swept at [16384, 7168] on 8 ranks, whole fused layer: + # + # blocks 16 24 32 48 64 80 128 256 + # us 1261 1091 1006 962 959 963 1018 1138 + # + # Flat from 48 to 80 and steep either side: 512 blocks (the quantize + # grid) costs 1184, and 16 costs 1261. Note this is *not* + # LSA_BLOCK_CAP's 24 -- that cap is for a kernel moving bf16 with no + # arithmetic, while this one moves half the bytes and dequantises them, + # so it needs more waves in flight to keep the links fed. + pull_blocks = max( + 1, + min( + int(os.environ.get("MORI_GEMM_AR_PULL_BLOCKS", 0)) or PULL_BLOCKS, + quant_blocks, + ), + ) + + def _row_addr_at(w, peer, byte_off): + return create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(peer, byte_off)) + ) + + def _row_addr(w, byte_off): + return create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, byte_off)) + ) + + @flyc.kernel(known_block_size=[QUANT_THREADS, 1, 1]) + def sdma_reduce_quant(dev_comm: Int64, win: Int64): + """The reduce, with the fp8 narrowing folded in. + + Same arithmetic as `sdma_reduce` followed by `quantize_gather`, and + the same numbers out -- it quantises from the bf16 it just wrote, so + the rounding is identical. What it saves is reading that bf16 back + out of memory, 28 MiB a layer at the model's shape, plus a kernel. + + The per-row amax is why this needs a different thread map from + `sdma_reduce`: that one is flat and grid-strided, so a row is spread + over many blocks and no single block can take its maximum. Here a + wave owns a row, which makes the amax a butterfly (below) and keeps + every lane's own 112 elements addressable. + + Those 112 are held in registers between the two passes, as bf16 + packed into 56 i32 rather than 112 f32. That is the whole reason + this fits without LDS: at f32 it would not. + + **This loses, and is off by default.** At [16384, 7168] on 8 ranks, + whole fused layer with the LSA pull gather: + + split reduce + quantize 957.4us + fused, row stashed in regs 977.1us + fused, row re-read 982.0us + + It is not register pressure -- the re-reading variant, which keeps + no stash at all, is no better. It is the thread map. `sdma_reduce` + walks packs with a flat grid stride, so a block streams straight + through all 8 source slices at once; one-wave-per-row confines a + wave to 14 KiB at a time and scatters the block's footprint across + the slice. The reduce loses more to that map than folding the + narrowing in saves by skipping the 28 MiB re-read. + + Kept, off, because the shape of the idea is right and the reason it + fails is not visible from reading it. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.CachedWindow(win) + lane = tid % fx.Int32(WAVE) + wave = tid // fx.Int32(WAVE) + + self_off = ( + cfg.recv_slot_off(rank) + if reduce_self_from_recv + else in_off + my_slice_off + ) + srcs = [_row_addr(w, self_off)] + [ + _row_addr(w, cfg.recv_slot_off((rank + j) % ws)) for j in range(1, ws) + ] + out = _row_addr(w, out_off + my_slice_off) + gdst = _row_addr(w, cfg.gout_off + rank * cfg.slice_rows * cfg.n) + gsca = _row_addr(w, cfg.gout_scale_slice_off(rank)) + + for row in range( + bid * QUANT_WAVES + wave, slice_rows, quant_blocks * QUANT_WAVES + ): + local = fx.Float32(0.0) + stash = [] + for c in range_constexpr(chunks_per_row): + off = row * cfg.n + (c * WAVE + lane) * CHUNK_ELEMS + for h in range_constexpr(2): + i32_off = off // 2 + h * 4 + acc = None + for j in range_constexpr(ws): + # recv is filled by a peer. A cached load can return + # the previous iteration's bytes when the peer's CUs + # wrote it, since our L2 is never invalidated by a + # remote store; SC1 skips L2 for those. + raw = fx.Vector( + buffer_load( + srcs[j], + i32_off, + vec_width=4, + dtype=i32_type(), + cache_modifier=( + CM_SC1 if recv_uncached else CM_CACHED + ), + ) + ) + v = raw.bitcast(fx.BFloat16).to(fx.Float32) + acc = v if acc is None else acc + v + for e in range_constexpr(CVT_ELEMS): + a = acc[e] + local = local.maximumf((-a).maximumf(a)) + packed = acc.to(fx.BFloat16).bitcast(fx.Int32) + buffer_store(packed, out, i32_off, cache_modifier=CM_CACHED) + stash.append(packed) + + amax = _wave_amax(local, lane) + is_zero = amax == fx.Float32(0.0) + scale = is_zero.select( + fx.Float32(1.0), amax * fx.Float32(1.0 / FP8_E4M3_MAX) + ) + inv = is_zero.select(fx.Float32(1.0), fx.Float32(FP8_E4M3_MAX) / amax) + if lane == fx.Int32(0): + buffer_store(scale, gsca, row, cache_modifier=CM_CACHED) + + for c in range_constexpr(chunks_per_row): + off = row * cfg.n + (c * WAVE + lane) * CHUNK_ELEMS + words = [] + for h in range_constexpr(2): + v = ( + fx.Vector(stash[c * 2 + h]) + .bitcast(fx.BFloat16) + .to(fx.Float32) + ) + scaled = v * inv + q = _pack8_fp8([scaled[e] for e in range_constexpr(CVT_ELEMS)]) + words += [q[0], q[1]] + buffer_store( + fx.Vector.from_elements(words, fx.Int32), + gdst, + off // 4, + cache_modifier=CM_CACHED, + ) + + @flyc.kernel(known_block_size=[QUANT_THREADS, 1, 1]) + def quantize_gather(dev_comm: Int64, win: Int64): + """My reduced slice, bf16 -> fp8 + one fp32 scale per row. + + Reads `output + my_slice` (what the reduce just wrote) and writes + `gout + my_slice` plus `gout_scale`. Two passes over the row rather than + holding it in registers: 112 fp32 per lane would not fit, and the row is + 14 KiB and still in L2 from the reduce. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.CachedWindow(win) + lane = tid % fx.Int32(WAVE) + wave = tid // fx.Int32(WAVE) + + src = _row_addr(w, out_off + my_slice_off) + dst = _row_addr(w, cfg.gout_off + rank * cfg.slice_rows * cfg.n) + sca = _row_addr(w, cfg.gout_scale_slice_off(rank)) + + for row in range( + bid * QUANT_WAVES + wave, slice_rows, quant_blocks * QUANT_WAVES + ): + # pass 1: this row's amax, over the whole wave + local = fx.Float32(0.0) + for c in range_constexpr(chunks_per_row): + off = row * cfg.n + (c * WAVE + lane) * CHUNK_ELEMS + for h in range_constexpr(2): + v = ( + fx.Vector( + buffer_load( + src, + off // 2 + h * 4, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_CACHED, + ) + ) + .bitcast(fx.BFloat16) + .to(fx.Float32) + ) + for e in range_constexpr(CVT_ELEMS): + a = v[e] + local = local.maximumf((-a).maximumf(a)) + amax = _wave_amax(local, lane) + + # A row of exact zeros would divide by zero; 1.0 keeps it exact. + is_zero = amax == fx.Float32(0.0) + scale = is_zero.select( + fx.Float32(1.0), amax * fx.Float32(1.0 / FP8_E4M3_MAX) + ) + inv = is_zero.select(fx.Float32(1.0), fx.Float32(FP8_E4M3_MAX) / amax) + if lane == fx.Int32(0): + buffer_store(scale, sca, row, cache_modifier=CM_CACHED) + + # pass 2: scale and narrow + for c in range_constexpr(chunks_per_row): + off = row * cfg.n + (c * WAVE + lane) * CHUNK_ELEMS + words = [] + for h in range_constexpr(2): + v = ( + fx.Vector( + buffer_load( + src, + off // 2 + h * 4, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_CACHED, + ) + ) + .bitcast(fx.BFloat16) + .to(fx.Float32) + ) + scaled = v * inv + q = _pack8_fp8([scaled[e] for e in range_constexpr(CVT_ELEMS)]) + words += [q[0], q[1]] + # One dwordx4 out, matching the dequantise side. + buffer_store( + fx.Vector.from_elements(words, fx.Int32), + dst, + off // 4, + cache_modifier=CM_CACHED, + ) + + @flyc.kernel(known_block_size=[QUANT_THREADS, 1, 1]) + def dequantize_gather(dev_comm: Int64, win: Int64): + """Every peer's gathered slice, fp8 * scale -> bf16 into `output`. + + Skips my own slice: it was never quantised and is already in `output`. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.CachedWindow(win) + lane = tid % fx.Int32(WAVE) + wave = tid // fx.Int32(WAVE) + + for j in range_constexpr(1, ws): + peer = (rank + j) % ws + src = _row_addr(w, cfg.gout_off + peer * cfg.slice_rows * cfg.n) + sca = _row_addr(w, cfg.gout_scale_slice_off(peer)) + dst = _row_addr(w, out_off + peer * cfg.slice_rows * cfg.n * 2) + + for row in range( + bid * QUANT_WAVES + wave, slice_rows, quant_blocks * QUANT_WAVES + ): + # Written by a peer's copy engine. SC1 for the same reason the + # reduce uses it on `recv`: our L2 is never invalidated by a + # remote store, so a cached load can return stale bytes. + scale = fx.Float32( + buffer_load( + sca, + row, + vec_width=1, + dtype=fx.Float32, + cache_modifier=CM_SC1, + ) + ) + for c in range_constexpr(chunks_per_row): + off = row * cfg.n + (c * WAVE + lane) * CHUNK_ELEMS + # One dwordx4: 16 fp8. At 8 this was a half-width read + # against a full-width write, and the kernel sat at + # 3.4 TB/s. + packed = fx.Vector( + buffer_load( + src, + off // 4, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_SC1, + ) + ) + for h in range_constexpr(2): + half = fx.Vector.from_elements( + [packed[2 * h], packed[2 * h + 1]], fx.Int32 + ) + wide = [e * scale for e in _unpack8_fp8(half)] + v = fx.Vector.from_elements(wide, fx.Float32) + buffer_store( + v.to(fx.BFloat16).bitcast(fx.Int32), + dst, + off // 2 + h * 4, + cache_modifier=CM_CACHED, + ) + + @flyc.kernel(known_block_size=[QUANT_THREADS, 1, 1]) + def lsa_pull_gather(dev_comm: Int64, win: Int64): + """Pull every peer's fp8 slice over xGMI and widen it on the way in. + + The alternative to `gather` + `dequantize_gather`. A copy engine has + no ALU, so an SDMA push can never scale anything and the widening + has to be a second kernel reading the landed fp8 back out of local + HBM. A CU pull has the bytes in registers already, so that read -- + 98 MiB a layer at the model's shape -- simply does not happen. + + What it costs is CU time during the transfer, which SDMA does not + spend. Whether that trade pays is a measurement, which is why both + transports are kept. + + Two things here are load-bearing, and `kernels_lsa.ar_2stage` + learned both the hard way: + + * the acquire fence. The barrier before this kernel orders, but does + not invalidate, so without it we can read our own stale L2 lines + for a peer's gout. Belt and braces with SC1 below, which also + costs nothing here: peer bytes are read exactly once, so there is + no reuse for L2 to capture. + * the loop nesting. Every peer's load is issued before any store, so + all 7 links stream at once. Draining one peer at a time pins the + whole gather to a single link -- measured there as ~133us against + ~660us, and it is invisible in the output. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.CachedWindow(win) + lane = tid % fx.Int32(WAVE) + wave = tid // fx.Int32(WAVE) + + raw_cco.cco_system_fence(fx.Int32(0)) + + # Peer p reduced slice p, so slice p is read out of p's own window. + peers = [(rank + j) % ws for j in range(1, ws)] + srcs = [ + _row_addr_at(w, p, cfg.gout_off + p * cfg.slice_rows * cfg.n) + for p in peers + ] + scas = [_row_addr_at(w, p, cfg.gout_scale_slice_off(p)) for p in peers] + # One descriptor for the destination, not one per peer: the slices + # are contiguous in my own output, so the peer index folds into the + # offset and 7 SGPR quads stay free. + out = _row_addr(w, out_off) + + for row in range( + bid * QUANT_WAVES + wave, slice_rows, pull_blocks * QUANT_WAVES + ): + scales = [ + fx.Float32( + buffer_load( + scas[i], + row, + vec_width=1, + dtype=fx.Float32, + cache_modifier=CM_SC1, + ) + ) + for i in range_constexpr(ws - 1) + ] + for c in range_constexpr(chunks_per_row): + off = row * cfg.n + (c * WAVE + lane) * CHUNK_ELEMS + vals = [ + buffer_load( + srcs[i], + off // 4, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_SC1, + ) + for i in range_constexpr(ws - 1) + ] + for i in range_constexpr(ws - 1): + packed = fx.Vector(vals[i]) + base = peers[i] * cfg.slice_rows * cfg.n + off + for h in range_constexpr(2): + half = fx.Vector.from_elements( + [packed[2 * h], packed[2 * h + 1]], fx.Int32 + ) + wide = [e * scales[i] for e in _unpack8_fp8(half)] + v = fx.Vector.from_elements(wide, fx.Float32) + buffer_store( + v.to(fx.BFloat16).bitcast(fx.Int32), + out, + base // 2 + h * 4, + cache_modifier=CM_CACHED, + ) + + # Gather. My reduced slice goes to the same offset in every peer's window. + # On the bf16 wire that offset is in `output` and there is nothing to stage; + # on the fp8 wire it is `gout`, which `quantize_gather` filled, and the + # scales ride along on the same queue. + if cfg.fp8_gather: + my_gout_off = cfg.gout_off + rank * cfg.slice_rows * cfg.n + my_gscale_off = cfg.gout_scale_slice_off(rank) + gather = _push_kernel( + end_off, + dst_off_of_peer=my_gout_off, + src_off_expr=lambda tid: fx.Int64(my_gout_off), + nbytes=cfg.gather_slice_bytes, + extra=( + my_gscale_off, + lambda tid: fx.Int64(my_gscale_off), + cfg.gather_scale_slice_bytes, + ), + ) + else: + gather = _push_kernel( + end_off, + dst_off_of_peer=out_off + my_slice_off, + src_off_expr=lambda tid: fx.Int64(out_off + my_slice_off), + ) + + # Barrier on the gather's signal row, no puts. The LSA pull needs every + # peer's quantise to have landed before it reads, and a push kernel with no + # puts is exactly that barrier. + gather_barrier = _push_kernel(end_off) + + # Drain-only twin of `scatter`: same barrier, no puts. The fused GEMM issues + # the puts from its epilogue, but still has to drain the queues and tell the + # peers their slices landed. + drain = _push_kernel(start_off) + + def _phase(kern, grid, blk): + @flyc.jit + def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): + kern(dev_comm, win).launch( + grid=(grid, 1, 1), block=[blk, 1, 1], stream=stream + ) + + return go + + phases = { + "scatter": _phase(scatter, 1, PUSH_THREADS), + "drain": _phase(drain, 1, PUSH_THREADS), + "reduce": _phase(sdma_reduce, red_blocks, threads), + "gather": _phase(gather, 1, PUSH_THREADS), + } + if fuse_reduce_push: + phases["reduce_push"] = _phase(sdma_reduce_push, red_blocks, threads) + # The puts are already out; this only drains the queues and runs the + # cross-rank barrier that tells peers their slices landed. + phases["gather_drain"] = _phase(_push_kernel(end_off), 1, PUSH_THREADS) + if cfg.fp8_gather: + # `quantize` has to follow the reduce and precede whichever transport + # moves the bytes -- unless it is folded into the reduce, which is the + # default, in which case there is no separate phase at all. + phases["quantize"] = _phase(quantize_gather, quant_blocks, QUANT_THREADS) + phases["reduce_quant"] = _phase(sdma_reduce_quant, quant_blocks, QUANT_THREADS) + if cfg.lsa_gather: + # Pull: barrier first (the peers must have finished quantising), + # then one kernel that transfers and widens together. + phases["gather_barrier"] = _phase(gather_barrier, 1, PUSH_THREADS) + phases["pull"] = _phase(lsa_pull_gather, pull_blocks, QUANT_THREADS) + else: + # Push: the copy engine cannot widen, so that is a second kernel, + # and it runs after the barrier the push ends with. + phases["dequantize"] = _phase( + dequantize_gather, quant_blocks, QUANT_THREADS + ) + # The order a caller must run them in. Carried with the phases rather than + # left to each caller to hardcode: the fp8 wire adds two, and a caller that + # kept its own ("scatter", "reduce", "gather") list would silently skip them + # and all-reduce into zeros. + if cfg.fp8_gather: + # Folding the narrowing into the reduce replaces two phases with one. + # `fuse_quantize=False` keeps them split, which is how the two are + # compared. + head = ("reduce_quant",) if fuse_quantize else ("reduce", "quantize") + tail = head + ( + ("gather_barrier", "pull") if cfg.lsa_gather else ("gather", "dequantize") + ) + elif fuse_reduce_push: + tail = ("reduce_push", "gather_drain") + else: + tail = ("reduce", "gather") + phases["order"] = ("scatter",) + tail + #: The same sequence for the fused GEMM, which issued the scatter's puts from + #: its own epilogue and only has to drain them. + phases["fused_order"] = ("drain",) + tail + return phases + + +def build_sdma_ar(cfg, rank: int, *, queues: int = 8, signal: bool = False): + """Compile a whole SDMA all-reduce for ``cfg`` on ``rank``. + + Returns ``(run, stage)``; ``stage`` is always 2 -- there is no 1-stage SDMA + variant, since a copy engine gains nothing from broadcasting the whole + payload to every peer. + """ + parts = build_sdma_phases(cfg, rank, queues=queues, signal=signal) + # On the fp8 wire this is five kernels, not three: the reduced slice has to + # be narrowed before the push and widened after it, and each needs its own + # device-wide ordering point for the same reason the original three do. + seq = [parts[k] for k in parts["order"]] + + def run(dev_comm, win, stream=fx.Stream(None)): + for phase in seq: + phase(dev_comm, win, stream=stream) + + return run, 2 + + +__all__ = ["build_sdma_ar", "build_sdma_phases", "PUSH_THREADS"] diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py new file mode 100644 index 000000000..2744f7831 --- /dev/null +++ b/python/mori/ops/gemm_ar/layout.py @@ -0,0 +1,729 @@ +# 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. +"""Window layout and launch geometry for the FlyDSL cco all-reduce. + +All of the arithmetic that the kernels depend on lives here so it can be unit +tested without a GPU: symmetric-window offsets, the per-block signal slot map, +the 1-stage/2-stage dispatch thresholds, and the launch geometry. + +Everything mirrors aiter's ``custom_all_reduce.cuh`` so the two are comparable +kernel-for-kernel: + +* ``Signal { start[kMaxBlocks][8]; end[kMaxBlocks][8]; _flag[kMaxBlocks] }``, + each region 128-byte aligned -- the slots are per (block, peer) so blocks + synchronise independently (``custom_all_reduce.cuh:43-51``). +* 1-stage below the size thresholds, 2-stage above + (``custom_all_reduce.cuh:3771-3788``). + +The launch geometry deliberately does *not* follow aiter -- see ``LSA_BLOCK_CAP``. + +The transport is the only thing this file does not model: LSA and SDMA share the +window layout, they differ only in how bytes move. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace as _dc_replace + +# --- aiter parity constants (custom_all_reduce.cuh) --- +K_MAX_BLOCKS = 80 # kMaxBlocks; also the number of signal slot rows +THREADS = 512 # THREAD_NUM / __launch_bounds__(512, 1) +MAX_WORLD = 8 # Signal::start[..][8] fixes the peer fan-out at 8 +PACK_BYTES = 16 # one 16B pack is the atomic transfer unit +SIGNAL_ALIGN = 128 # alignas(128) on each Signal region + +# Grid cap for the LSA kernels. NOT aiter's 80: aiter gives each thread one pack +# from one peer (its `warp_id` selects the peer), while ours gives each thread one +# index across all `world` peers -- 8x the in-flight bytes per thread. Matching +# aiter's grid therefore over-subscribes the xGMI request queues, and measurably +# so: at [4096, 7168] bf16 on 8x MI355X the time rises monotonically past ~32 +# blocks (286us @28 -> 361 @80 -> 583 @256). Measured plateau is 24-36; 24 is +# chosen because it is also the optimum at [512, 7168], where the curve is +# sharper (64.0us @24 vs 68.5 @32). +# +# The signal array still has K_MAX_BLOCKS rows, so this cap can be raised for +# experiments without touching the window layout. +LSA_BLOCK_CAP = 24 + +# Grid cap for the SDMA all-reduce's local reduce kernel. Deliberately not +# LSA_BLOCK_CAP: that one throttles outstanding xGMI requests, whereas this +# kernel's loads are all local HBM and want the opposite. See the sweep in +# bench_ar.py. +SDMA_REDUCE_BLOCK_CAP = 256 + +# 1-stage thresholds, in bytes (custom_all_reduce.cuh:3779). +ONE_STAGE_MAX_BYTES_LE4 = 160 * 1024 +ONE_STAGE_MAX_BYTES_LE8 = 80 * 1024 + +# Legal mode strings. Both are selected by equality in the kernels, so an +# unrecognised value does not fail -- it silently takes the other branch. +WIRE_DTYPES = frozenset({"bf16", "fp8"}) +GATHER_TRANSPORTS = frozenset({"sdma", "lsa"}) + +# The fp8 conversion gives one wave one row and converts CHUNK_ELEMS (16) +# elements per lane per chunk, so a row must be a whole number of wave-chunks. +FP8_GATHER_N_MULTIPLE = 16 * 64 + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def select_stage(world_size: int, nbytes: int) -> int: + """Return 1 or 2 -- aiter's dispatch rule, reproduced exactly. + + world==2 is always 1-stage: with a single peer the reduce-scatter half of + the 2-stage algorithm saves nothing and costs an extra barrier. + """ + if world_size == 2: + return 1 + if world_size <= 4 and nbytes < ONE_STAGE_MAX_BYTES_LE4: + return 1 + if world_size <= 8 and nbytes < ONE_STAGE_MAX_BYTES_LE8: + return 1 + return 2 + + +def ar_config(*, comm_dtype: str | None = None, **kwargs) -> "ArConfig": + """``ArConfig`` with ``comm_dtype`` as a shorthand for both wire legs. + + The two legs are separate fields because they are separately useful -- the + gather leg is where nearly all of the time is and where fp8 is one rounding, + while the scatter leg is mostly hidden behind the GEMM already and its fp8 + error compounds across ranks. Setting both at once is still the common case. + """ + if comm_dtype is not None: + if kwargs.get("scatter_dtype") or kwargs.get("gather_dtype"): + raise ValueError("pass comm_dtype or the per-leg dtypes, not both") + kwargs["scatter_dtype"] = comm_dtype + kwargs["gather_dtype"] = comm_dtype + return ArConfig(**kwargs) + + +@dataclass(frozen=True) +class ArConfig: + """Shape + world size -> every offset and count the kernels need. + + ``m``/``n`` describe the logical bf16 tensor being reduced; ``n`` is the row + length (7168 for DSV4-Pro ``wo_b``) and is never sharded -- the reduce-scatter + shards along ``m`` so every peer slice is one contiguous byte range. Sharding + along ``n`` would make each slice ``m`` separate ``n/world*2``-byte pieces, + which is fatal for SDMA (~2us per packet regardless of size). + """ + + world_size: int + m: int + n: int + elem_bytes: int = 2 # bf16 + threads: int = THREADS + max_blocks: int = K_MAX_BLOCKS # signal-array rows, never the launched grid + block_cap: int = LSA_BLOCK_CAP + #: SDMA landing slots, one slice per peer. LSA needs none -- it reduces + #: straight out of the peers' input regions -- but a copy engine has to be + #: given somewhere to write, so the SDMA backend reserves ``world_size`` + #: slices (slot ``rank`` goes unused; keeping it makes the index be the + #: peer rank, which is worth one slice of VRAM). + recv_slots: int = 0 + #: Sub-slices each peer's landing region is pushed in, for the fused GEMM. + #: 1 means "one put per destination"; more lets a destination's first chunk + #: leave while the GEMM is still computing its later ones. Only the fused + #: path reads this; the standalone SDMA all-reduce always pushes whole slices. + counter_chunks: int = 1 + #: Chunks the counter *region* is sized for, as opposed to the number this + #: config uses. ``counter_chunks`` is the largest divisor of the band count + #: and so changes with ``m``; sizing the region from it moves ``lock_off`` + #: and everything after it, which silently reinterprets the previous shape's + #: payload as control state when one window serves several ``m``. 0 means + #: "same as ``counter_chunks``", which is right for a single-shape config. + counter_capacity: int = 0 + #: Independent counter sets the region holds, and which one this config + #: uses. The election test is ``seq % tiles_per_chunk == 0`` on a counter + #: that is never reset, and ``tiles_per_chunk`` depends on ``m`` -- so two + #: shapes sharing a set elect on a residue left by the other. One set per + #: shape keeps them independent without resetting anything, which is what + #: lets graphs for different ``m`` be replayed in any order. + counter_shape_slots: int = 1 + counter_shape_index: int = 0 + #: Lay the payload regions out for this ``m`` instead of for ``m``. + #: + #: Every offset past the control area is a running sum of m-sized regions, + #: so two shapes sharing a window put different regions at the same bytes: + #: at world=2 n=1024, m=1024's output rows 512..767 land exactly where + #: m=512's recv slot for peer 1 was, and read that peer's previous scatter + #: payload instead of its gathered slice. Pinning the offsets to the largest + #: m an instance will serve makes every shape's map identical, so a region + #: only ever aliases *itself*. Sizes stay m-based -- this moves where things + #: are, not how much is moved. 0 means "same as m". + capacity_m: int = 0 + #: What travels on the wire, per leg. ``"bf16"`` sends the payload as-is; + #: ``"fp8"`` quantises to e4m3 with fp32 scales, halving the bytes. The + #: reduce always accumulates in fp32 and ``output`` is always bf16 -- only + #: the wire changes. + #: + #: The legs are separate because they are not equivalent. ``gather`` carries + #: *finished* values, so fp8 is one rounding; ``scatter`` carries partial + #: sums that are then added across every rank, so its error compounds. They + #: also quantise in different places, hence the two scale granularities + #: below. ``ArConfig(comm_dtype=...)`` sets both at once. + scatter_dtype: str = "bf16" + gather_dtype: str = "bf16" + #: Bands the reduce publishes its slice in, when the gather's puts are + #: fired from inside it. One SDMA put per band per peer, so this trades + #: ramp-up (the first put waits for one band) against per-packet cost (the + #: engine charges ~2us per packet whatever its size). + gather_bands: int = 8 + #: Who moves the gather leg's bytes, when it is fp8. + #: + #: ``"sdma"`` pushes with the copy engines and then dequantises in a second + #: kernel -- a copy engine has no ALU, so the two cannot be one step. + #: ``"lsa"`` instead *pulls* each peer's slice over xGMI into registers and + #: dequantises on the way to memory, which is one kernel and saves reading + #: the landed fp8 back out of local HBM. It costs CU time during the + #: transfer, where SDMA costs none. + gather_transport: str = "sdma" + #: N-tile the scatter leg's scales are taken over. Must match the fused + #: GEMM's ``BLOCK_N``: a block owns ``BLOCK_M x BLOCK_N`` of C, so this is + #: the widest span of a row it can take an amax over. + scatter_scale_n: int = 256 + #: Grid for the SDMA reduce. Independent of ``block_cap`` because that cap + #: throttles *xGMI* requests, while this kernel reads local HBM and wants as + #: much in flight as it can get. + force_reduce_blocks: int = 0 + #: Override the computed grid. Only the signal array bounds the grid for + #: correctness (``start[max_blocks][8]``), so raising this also needs + #: ``max_blocks`` raised. Exists so the benchmark can sweep occupancy. + force_blocks: int | None = None + + # --- basic sizes --- + + @property + def num_elems(self) -> int: + return self.m * self.n + + @property + def nbytes(self) -> int: + return self.num_elems * self.elem_bytes + + @property + def elems_per_pack(self) -> int: + return PACK_BYTES // self.elem_bytes + + @property + def num_packs(self) -> int: + """Payload measured in 16B packs -- the unit the kernels index in.""" + return self.nbytes // PACK_BYTES + + @property + def stage(self) -> int: + return select_stage(self.world_size, self.nbytes) + + # --- wire format ------------------------------------------------------- + # + # Only the two transfer legs change dtype. ``output`` stays bf16 because it + # is the consumer's tensor, and the reduce still accumulates in fp32. + + @property + def fp8_scatter(self) -> bool: + return self.scatter_dtype == "fp8" + + @property + def fp8_gather(self) -> bool: + return self.gather_dtype == "fp8" + + @property + def lsa_gather(self) -> bool: + """fp8 gather moved by CU pull rather than by the copy engines.""" + return self.fp8_gather and self.gather_transport == "lsa" + + @property + def fp8_wire(self) -> bool: + """Either leg quantised. Use the per-leg flags for sizing.""" + return self.fp8_scatter or self.fp8_gather + + @property + def scatter_elem_bytes(self) -> int: + return 1 if self.fp8_scatter else self.elem_bytes + + @property + def gather_elem_bytes(self) -> int: + return 1 if self.fp8_gather else self.elem_bytes + + @property + def scatter_nbytes(self) -> int: + """``input``: the GEMM's C, in the scatter leg's dtype.""" + return self.num_elems * self.scatter_elem_bytes + + @property + def gather_nbytes(self) -> int: + return self.num_elems * self.gather_elem_bytes + + @property + def scatter_slice_bytes(self) -> int: + """One peer slice on the scatter leg. The unit of a single SDMA put.""" + return self.slice_rows * self.n * self.scatter_elem_bytes + + @property + def gather_slice_bytes(self) -> int: + return self.slice_rows * self.n * self.gather_elem_bytes + + @property + def slice_rows(self) -> int: + """Rows in one peer's slice. The reduce-scatter shards along m only.""" + if self.m % self.world_size: + raise ValueError(f"m={self.m} must divide by world_size={self.world_size}") + return self.m // self.world_size + + # --- scale regions ----------------------------------------------------- + # + # Two granularities, because the two legs quantise in different places. + # + # scatter: the GEMM epilogue quantises, and a block owns BLOCK_M x BLOCK_N + # of C. With N=7168 and BLOCK_N=256 a row spans 28 blocks, so a per-row amax + # is not available there at all -- it would need a cross-block reduction, + # which serialises exactly what the fusion overlaps. One scale per (row, + # N-tile) is what a block can compute, and being 28 scales per row instead + # of 1 it is *finer* than per-row, not coarser. + # + # gather: the reduce kernel owns whole rows of its slice, so amax over all + # of n is free and one scale per row is enough. + + @property + def scatter_tiles_per_row(self) -> int: + if self.n % self.scatter_scale_n: + raise ValueError( + f"n={self.n} must be a multiple of scatter_scale_n=" + f"{self.scatter_scale_n}" + ) + return self.n // self.scatter_scale_n + + @property + def scatter_scale_bytes(self) -> int: + """fp32 scales for the whole ``[m, n]`` payload, scatter granularity.""" + if not self.fp8_scatter: + return 0 + return _align_up(self.m * self.scatter_tiles_per_row * 4, SIGNAL_ALIGN) + + @property + def scatter_scale_slice_bytes(self) -> int: + """The scales that travel with one peer slice.""" + if not self.fp8_scatter: + return 0 + return self.slice_rows * self.scatter_tiles_per_row * 4 + + @property + def gather_scale_bytes(self) -> int: + """fp32 scales for the whole payload, one per row.""" + if not self.fp8_gather: + return 0 + return _align_up(self.m * 4, SIGNAL_ALIGN) + + @property + def gather_scale_slice_bytes(self) -> int: + if not self.fp8_gather: + return 0 + return self.slice_rows * 4 + + # --- reduce-scatter partition (2-stage), sharded along m --- + + @property + def packs_per_rank(self) -> int: + """Packs each rank owns and reduces. Trailing packs go to the last rank.""" + return self.num_packs // self.world_size + + def owner_pack_range(self, rank: int) -> tuple[int, int]: + start = rank * self.packs_per_rank + end = ( + self.num_packs + if rank == self.world_size - 1 + else start + self.packs_per_rank + ) + return start, end + + @property + def largest_pack_part(self) -> int: + """``part + size % ngpus`` -- the all-gather loop bound in aiter.""" + return self.packs_per_rank + self.num_packs % self.world_size + + @property + def slice_bytes(self) -> int: + """Contiguous bytes in one peer slice; drives SDMA bandwidth.""" + return self.packs_per_rank * PACK_BYTES + + # --- launch geometry --- + + @property + def threads_per_peer(self) -> int: + """``tnum_gpu``: the thread group that services one peer.""" + return self.threads // self.world_size + + @property + def blocks(self) -> int: + # One thread owns one pack index and reduces it across every peer, so a + # block advances the index space by `threads` per iteration -- not by + # `threads_per_peer` as in aiter, whose threads are split across peers. + if self.force_blocks is not None: + return self.force_blocks + per_block = self.threads + covered = self.num_packs if self.stage == 1 else self.packs_per_rank + need = (covered + per_block - 1) // per_block + return max(1, min(self.block_cap, need)) + + # --- symmetric window layout: [ signal | input | output | tmp | recv ] --- + # ``tmp`` holds each rank's reduced slice during 2-stage; peers all-gather + # straight out of it, which is why it lives inside the registered window. + # ``recv`` is SDMA-only and empty unless ``recv_slots`` is set. + + @property + def signal_bytes(self) -> int: + # Everything before `input`: the barrier arrays, the scatter's tile + # counters, the submit locks, and the gather's band counters. + # The region bases, not this config's slot: the control area has to be + # the same size whichever shape slot a config happens to use. + return self.gather_counter_region_off + self.gather_counter_bytes + + @property + def start_off(self) -> int: + return 0 + + @property + def end_off(self) -> int: + return _align_up(self.max_blocks * MAX_WORLD * 4, SIGNAL_ALIGN) + + @property + def flag_off(self) -> int: + return self.end_off * 2 + + @property + def reduce_blocks(self) -> int: + if self.force_reduce_blocks: + return self.force_reduce_blocks + need = (self.packs_per_rank + self.threads - 1) // self.threads + return max(1, min(SDMA_REDUCE_BLOCK_CAP, need)) + + @property + def counter_capacity_eff(self) -> int: + """Chunks per set the region is sized for; ``counter_chunks`` if unset.""" + return self.counter_capacity or self.counter_chunks + + @property + def counter_region_off(self) -> int: + """Monotonic tile counters for the fused GEMM, one per (dest, chunk). + + Its own 128-byte-aligned region rather than slack in ``_flag``: the count + is ``world_size * counter_capacity`` and grows with M, so borrowing the + 16 spare ``_flag`` slots would silently stop fitting at M >= 8192. + """ + return _align_up(self.flag_off + self.max_blocks * 4, SIGNAL_ALIGN) + + @property + def counter_set_bytes(self) -> int: + """One shape's set of ``(dest, chunk)`` counters.""" + return self.world_size * self.counter_capacity_eff * 4 + + @property + def counter_off(self) -> int: + """This config's counter set. Sized by capacity, so it does not move.""" + return ( + self.counter_region_off + self.counter_shape_index * self.counter_set_bytes + ) + + @property + def counter_bytes(self) -> int: + return _align_up( + self.counter_set_bytes * self.counter_shape_slots, SIGNAL_ALIGN + ) + + @property + def lock_off(self) -> int: + """One submit lock per destination, for the fused GEMM's epilogue. + + Needed as soon as a destination has more than one chunk: the tile + counter elects one block per *chunk*, so two of them can reach the SDMA + submit for the same queue at once. + + Shared across shapes, unlike the counters: a lock is taken and released + inside one kernel, so it carries nothing between calls. + """ + return self.counter_region_off + self.counter_bytes + + @property + def lock_bytes(self) -> int: + return _align_up(self.world_size * 4, SIGNAL_ALIGN) + + @property + def gather_counter_region_off(self) -> int: + """Band counters for the gather, when its puts are fired by the reduce. + + Its own region rather than a corner of ``counter_off``: the fused GEMM + bumps every ``(dest, chunk)`` slot including its own rank's -- the + ``dest != rank`` test skips the *put*, not the count -- so there is no + free slot to borrow, and silently sharing one would elect on a mixture + of two different epochs. + """ + return self.lock_off + self.lock_bytes + + @property + def gather_counter_off(self) -> int: + """This config's band counters, one set per shape like the tile ones.""" + return ( + self.gather_counter_region_off + + self.counter_shape_index * self.gather_bands * 4 + ) + + @property + def gather_counter_bytes(self) -> int: + return _align_up(self.gather_bands * 4 * self.counter_shape_slots, SIGNAL_ALIGN) + + def counter_slot(self, dest: int, chunk: int) -> int: + if not 0 <= dest < self.world_size: + raise IndexError(f"dest {dest} outside [0, {self.world_size})") + if not 0 <= chunk < self.counter_chunks: + raise IndexError(f"chunk {chunk} outside [0, {self.counter_chunks})") + return dest * self.counter_chunks + chunk + + def signal_slot(self, block: int, peer: int) -> int: + """Element index into ``start``/``end`` for (block, peer).""" + if not 0 <= block < self.max_blocks: + raise IndexError(f"block {block} outside [0, {self.max_blocks})") + if not 0 <= peer < MAX_WORLD: + raise IndexError(f"peer {peer} outside [0, {MAX_WORLD})") + return block * MAX_WORLD + peer + + @property + def _cap(self) -> "ArConfig": + """This config sized for ``capacity_m``; ``self`` when there is none. + + Every offset below is taken from here rather than from ``self``, which + is what makes the map identical across shapes. It has ``capacity_m=0``, + so the recursion stops immediately. + """ + if not self.capacity_m or self.capacity_m == self.m: + return self + return _dc_replace(self, m=self.capacity_m, capacity_m=0) + + @property + def input_off(self) -> int: + """The GEMM's C, in the wire dtype. Source of the scatter leg.""" + return _align_up(self.signal_bytes, SIGNAL_ALIGN) + + @property + def input_scale_off(self) -> int: + """Scatter-leg scales for ``input``. Zero-sized on the bf16 wire.""" + c = self._cap + return _align_up(c.input_off + c.scatter_nbytes, SIGNAL_ALIGN) + + def input_scale_slice_off(self, peer: int) -> int: + """Where peer ``peer``'s row band's scales start within ``input_scale``.""" + c = self._cap + return c.input_scale_off + peer * c.slice_rows * c.scatter_tiles_per_row * 4 + + @property + def output_off(self) -> int: + """The consumer's tensor. Always bf16 -- only the wire changes dtype.""" + c = self._cap + return c.input_scale_off + c.scatter_scale_bytes + + @property + def output_end(self) -> int: + c = self._cap + return _align_up(c.output_off + c.nbytes, SIGNAL_ALIGN) + + @property + def gout_off(self) -> int: + """The reduced result in the wire dtype. Source of the gather leg. + + On the bf16 wire this *aliases* ``output``: the gather pushes the bf16 + result straight out of it and there is nothing to stage, so the region + costs nothing. On the fp8 wire it is its own region, because ``output`` + has to stay bf16 for the consumer while the wire carries fp8. + """ + c = self._cap + return c.output_end if c.fp8_gather else c.output_off + + @property + def gout_bytes(self) -> int: + return self.gather_nbytes if self.fp8_gather else 0 + + @property + def gout_scale_off(self) -> int: + c = self._cap + return _align_up(c.gout_off + c.gout_bytes, SIGNAL_ALIGN) + + def gout_scale_slice_off(self, peer: int) -> int: + c = self._cap + return c.gout_scale_off + peer * c.slice_rows * 4 + + @property + def tmp_off(self) -> int: + # Not simply "after gout_scale": on the bf16 wire gout aliases output and + # contributes nothing, so the running offset has to be taken past + # output's own bytes explicitly or tmp lands on top of it. + c = self._cap + return max(c.output_end, c.gout_scale_off + c.gather_scale_bytes) + + @property + def tmp_bytes(self) -> int: + """Room for the largest owned slice, so rank world-1's tail fits too.""" + return self.largest_pack_part * PACK_BYTES + + @property + def recv_off(self) -> int: + c = self._cap + return c.tmp_off + c.tmp_bytes + + @property + def recv_bytes(self) -> int: + c = self._cap + return c.recv_slots * c.scatter_slice_bytes + + def recv_slot_off(self, peer: int) -> int: + """Where peer ``peer``'s contribution to *my* slice lands.""" + if not 0 <= peer < self.world_size: + raise IndexError(f"peer {peer} outside [0, {self.world_size})") + if peer >= self.recv_slots: + raise IndexError( + f"peer {peer} has no landing slot; recv_slots={self.recv_slots}" + ) + c = self._cap + return c.recv_off + peer * c.scatter_slice_bytes + + @property + def recv_scale_off(self) -> int: + c = self._cap + return _align_up(c.recv_off + c.recv_bytes, SIGNAL_ALIGN) + + @property + def recv_scale_bytes(self) -> int: + if not self.fp8_wire: + return 0 + c = self._cap + return _align_up(c.recv_slots * c.scatter_scale_slice_bytes, SIGNAL_ALIGN) + + def recv_scale_slot_off(self, peer: int) -> int: + """Where peer ``peer``'s scales for my slice land.""" + if not 0 <= peer < self.recv_slots: + raise IndexError( + f"peer {peer} has no landing slot; recv_slots={self.recv_slots}" + ) + c = self._cap + return c.recv_scale_off + peer * c.scatter_scale_slice_bytes + + @property + def window_bytes(self) -> int: + c = self._cap + return c.recv_scale_off + c.recv_scale_bytes + + # --- traffic model, for reporting alongside measured time --- + + @property + def remote_bytes_per_rank(self) -> int: + """Bytes crossing xGMI per rank: (P-1)/P*2 for 2-stage, (P-1) for 1-stage.""" + p = self.world_size + if self.stage == 1: + return (p - 1) * self.nbytes + return 2 * (p - 1) * self.nbytes // p + + def validate(self) -> None: + if self.world_size < 2 or self.world_size > MAX_WORLD: + raise ValueError( + f"world_size must be in [2, {MAX_WORLD}], got {self.world_size}" + ) + if self.threads % self.world_size: + raise ValueError( + f"threads ({self.threads}) must be divisible by world_size " + f"({self.world_size}); aiter falls back to the naive kernel for " + "world_size=6 for exactly this reason" + ) + if self.nbytes % (self.world_size * PACK_BYTES): + raise ValueError( + f"payload {self.nbytes}B must be a multiple of " + f"world_size*{PACK_BYTES}={self.world_size * PACK_BYTES} so the " + "vectorized path applies (aiter: DISPATCH_REDUCE falls back to " + "_naive otherwise)" + ) + if ( + self.force_blocks is not None + and not 1 <= self.force_blocks <= self.max_blocks + ): + raise ValueError( + f"force_blocks={self.force_blocks} outside [1, {self.max_blocks}]; " + "the signal array has one row per block, so raise max_blocks too" + ) + if self.counter_chunks < 1: + raise ValueError(f"counter_chunks must be >= 1, got {self.counter_chunks}") + if self.counter_capacity and self.counter_capacity < self.counter_chunks: + raise ValueError( + f"counter_capacity ({self.counter_capacity}) must be >= " + f"counter_chunks ({self.counter_chunks}); the region has to hold " + "the set this config indexes into" + ) + if self.counter_shape_slots < 1: + raise ValueError( + f"counter_shape_slots must be >= 1, got {self.counter_shape_slots}" + ) + if not 0 <= self.counter_shape_index < self.counter_shape_slots: + raise ValueError( + f"counter_shape_index must be in [0, {self.counter_shape_slots}), " + f"got {self.counter_shape_index}" + ) + # Mode strings select a code path by equality, so a typo does not fail -- + # it silently picks the default. "FP8" sent bf16; "sla" sent SDMA. + for name, allowed in ( + ("scatter_dtype", WIRE_DTYPES), + ("gather_dtype", WIRE_DTYPES), + ("gather_transport", GATHER_TRANSPORTS), + ): + got = getattr(self, name) + if got not in allowed: + raise ValueError(f"{name}={got!r} is not one of {sorted(allowed)}") + if self.gather_bands < 1: + raise ValueError(f"gather_bands must be >= 1, got {self.gather_bands}") + if self.m < 1 or self.n < 1: + raise ValueError(f"m and n must be positive, got m={self.m} n={self.n}") + if self.fp8_gather and self.n % FP8_GATHER_N_MULTIPLE: + raise ValueError( + f"the fp8 wire gives one wave one row and converts a fixed number " + f"of elements per lane per chunk, so n ({self.n}) must be a " + f"multiple of {FP8_GATHER_N_MULTIPLE}" + ) + if not 0 <= self.recv_slots <= self.world_size: + raise ValueError( + f"recv_slots must be in [0, world_size={self.world_size}], " + f"got {self.recv_slots}" + ) + if self.elem_bytes not in (2, 4): + raise ValueError(f"elem_bytes must be 2 or 4, got {self.elem_bytes}") + + +__all__ = [ + "ArConfig", + "WIRE_DTYPES", + "GATHER_TRANSPORTS", + "FP8_GATHER_N_MULTIPLE", + "select_stage", + "K_MAX_BLOCKS", + "LSA_BLOCK_CAP", + "SDMA_REDUCE_BLOCK_CAP", + "THREADS", + "MAX_WORLD", + "PACK_BYTES", +] diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py new file mode 100644 index 000000000..134aab8a5 --- /dev/null +++ b/python/mori/ops/gemm_ar/op.py @@ -0,0 +1,776 @@ +# 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. +"""Host op-layer for the fused GEMM + all-reduce. See ``README.md``.""" + +from __future__ import annotations + +from typing import NamedTuple, Optional + +import torch + +import flydsl.expr as fx +from mori.cco import CCODevCommRequirements, GDA_CONNECTION_NONE +from mori.tensor_utils import from_gpu_ptr + +from .kernels_fused import BLOCK_K, compile_fused_gemm_scatter +from .kernels_sdma import build_sdma_phases +from .layout import GATHER_TRANSPORTS, WIRE_DTYPES, ArConfig + +# A destination's row band is BLOCK_M rows, so M has to divide into whole bands +# per destination. Both are also the tile the kernel is compiled for: 128x256 is +# what the block-scale path needs, because its per-K-block promotion doubles the +# accumulator VGPRs and 256x256 would spill. +DEFAULT_BLOCK_M = 128 +DEFAULT_BLOCK_N = 256 + +# The tile granules the kernel is written for. Its mainloop builds N_TILES_A = +# BLOCK_M/64 and N_TILES_B = BLOCK_N/128 MFMA accumulators and the LDS staging +# halves each dimension, so a tile that is not a whole multiple of these does +# not describe a real schedule -- the kernel asserts it (`_gemm_a8w8_8wave.py`, +# the `BLOCK_M >= 128 and ... % 128 == 0` assert). These are the same numbers as +# DEFAULT_BLOCK_M/N, but for a different reason: those are the measured best +# tile, these are the only shapes that exist at all. +TILE_M_GRANULE = 128 +TILE_N_GRANULE = 256 + +# The block-scale group, on both operands: A is 1x128, B is 128x128. +SCALE_BLOCK_K = 128 + +# Chunks are how many separate pushes a destination receives, and so how early +# the first bytes leave. More is better until the pieces get small enough that +# the SDMA per-packet cost shows; 8 is the measured knee. +MAX_CHUNKS = 8 + +# The mainloop prefetches K block 1 and then runs tail steps K_ITERS-2 and +# K_ITERS-1, so a single-K-block pipeline does not exist: K_ITERS must be >= 2. +# K=128 otherwise passes the divisibility check and returns finite garbage +# (relL2 1.41 against 2.4e-3 at K=256). +MIN_K = 2 * BLOCK_K + + +# Distinct M values one instance can serve, when not given explicitly. +# +# Each needs its own counter set (ArConfig.counter_shape_slots), and the default +# is *every* M the window can legally take: a padded M is a multiple of +# world_size*block_m and at most m_max, so there are m_max/(world_size*block_m) +# of them and a legal call can never exhaust the table. +# +# It is sized that way because running out was not graceful. A server's M +# changes with the batch -- 4096, 5120, 6144, 7168, ... at this shape -- so a +# fixed 8 ran out within seconds of real traffic, and SGLang turns any exception +# from this path into a permanent fallback, so the whole optimisation switched +# itself off. A counter set is world_size*MAX_CHUNKS*4 bytes, 256 B at tp8, so +# 16 of them cost 4 KiB of a 700 MiB window; rationing them bought nothing. +def default_max_shapes(m_pad_max: int, world_size: int, block_m: int) -> int: + """How many distinct padded M values fit under ``m_pad_max``.""" + return max(1, m_pad_max // (world_size * block_m)) + + +def padded_m(m: int, world_size: int, block_m: int = DEFAULT_BLOCK_M) -> int: + """``m`` rounded up to a whole number of row bands per destination.""" + granule = world_size * block_m + return (m + granule - 1) // granule * granule + + +def counter_chunks(m_pad: int, world_size: int, block_m: int = DEFAULT_BLOCK_M) -> int: + """The largest chunk count that divides the row bands per destination. + + Chunks must divide evenly or the completion counter's modulo test never + fires and the push is never issued. At small M there are fewer bands than + ``MAX_CHUNKS``, so take the largest divisor rather than failing: a fixed 8 + has nothing to divide at ``m_pad = 4096`` (4 bands). + + Worth taking the largest: at M=15360 the divisor rule measures 1104.8us + against 1434.0us for a chunk count of 1. + """ + bands = m_pad // (world_size * block_m) + if bands < 1: + raise ValueError( + f"m_pad={m_pad} is smaller than one row band per destination " + f"(world_size*block_m = {world_size * block_m}), so there is nothing " + f"to chunk" + ) + return max(c for c in range(1, min(MAX_CHUNKS, bands) + 1) if bands % c == 0) + + +def _tile_constraints(block_m: int, block_n: int) -> Optional[str]: + """Why this tile is not one the kernel can build, or None. + + Checked before anything divides by a tile size. ``padded_m`` and + ``default_max_shapes`` both do, so a ``block_m`` of 0 raises + ZeroDivisionError from inside the padding arithmetic if this runs later -- + and the kernel's own assert is bare, so reaching it gives an empty + AssertionError on the first real call rather than a reason at the boundary. + """ + for name, value, granule in ( + ("block_m", block_m, TILE_M_GRANULE), + ("block_n", block_n, TILE_N_GRANULE), + ): + if value < granule or value % granule: + return ( + f"{name}={value} must be a positive multiple of {granule}: the " + f"mainloop's MFMA tiling and LDS staging are written for that " + f"granule" + ) + return None + + +def _gemm_constraints(n: int, k: int, block_n: int) -> Optional[str]: + """Why the GEMM cannot take this shape, or None. See :func:`supports`.""" + if k % SCALE_BLOCK_K: + return f"K={k} must be a multiple of {SCALE_BLOCK_K} (the block-scale group)" + if k < MIN_K: + return ( + f"K={k} is below the minimum {MIN_K}: the mainloop prefetches a " + f"second K block and runs two tail steps, so K/{BLOCK_K} must be >= 2" + ) + if n % block_n: + return f"N={n} must be a multiple of block_n={block_n}" + if n % SCALE_BLOCK_K: + return f"N={n} must be a multiple of {SCALE_BLOCK_K} (the B scale group)" + return None + + +def _build_cfg( + *, + world_size: int, + m_pad: int, + n: int, + block_m: int, + gather_dtype: str, + gather_transport: str, + counter_capacity: int, + shape_slots: int, + shape_index: int, + capacity_m: int = 0, +) -> ArConfig: + """The one place an ``ArConfig`` is built, so every caller agrees. + + ``supports``, ``window_bytes_for`` and the op itself all come through here, + which is what keeps the support predicate from drifting away from what + execution actually accepts. + """ + cfg = ArConfig( + world_size=world_size, + m=m_pad, + n=n, + recv_slots=world_size, + counter_chunks=counter_chunks(m_pad, world_size, block_m), + counter_capacity=counter_capacity, + counter_shape_slots=shape_slots, + counter_shape_index=shape_index, + capacity_m=capacity_m, + gather_dtype=gather_dtype, + gather_transport=gather_transport, + ) + cfg.validate() + return cfg + + +def supports( + m: int, + n: int, + k: int, + world_size: int, + *, + block_m: int = DEFAULT_BLOCK_M, + block_n: int = DEFAULT_BLOCK_N, + gather_dtype: str = "bf16", + gather_transport: str = "lsa", +) -> bool: + """Whether this shape is *expressible*, which is not whether it is faster. + + Profitability depends on how much GEMM there is to hide the transfer behind + and is the caller's call -- see the measured curve in ``README.md``. + + The collective half is answered by building the config and validating it, + rather than by a second copy of its rules: a predicate that says yes where + construction raises, or vice versa, is worse than no predicate. + """ + # Tile and world size first: everything below divides by their product. + if world_size < 1 or _tile_constraints(block_m, block_n) is not None: + return False + if m <= 0 or _gemm_constraints(n, k, block_n) is not None: + return False + try: + _build_cfg( + world_size=world_size, + m_pad=padded_m(m, world_size, block_m), + n=n, + block_m=block_m, + gather_dtype=gather_dtype, + gather_transport=gather_transport, + counter_capacity=MAX_CHUNKS, + shape_slots=default_max_shapes( + padded_m(m, world_size, block_m), world_size, block_m + ), + shape_index=0, + ) + except ValueError: + return False + return True + + +#: The one fp8 encoding the kernel implements. The launch reinterprets the +#: tensor as bytes, so the check has to be on the *encoding*, not on the width: +#: +#: * ``float8_e4m3fnuz`` is also one byte, but its exponent bias is 8 where OCP's +#: is 7, so the same byte means twice the value. The MMA atom is fixed to +#: ``fx.Float8E4M3FN`` (``_gemm_a8w8_8wave.py``), so admitting FNUZ would read +#: every A and every B element at 2x and return a result 4x too large -- +#: finite, plausible, and wrong. sglang reaches this on gfx942, where +#: ``normalize_e4m3fn_to_e4m3fnuz`` converts weights to FNUZ. +#: * A wider dtype reinterprets two rows as one, likewise finite nonsense. +#: +#: Widening this needs a kernel path for the new encoding and a numerical test +#: per accepted dtype -- matching element size is not evidence of anything. +FP8_DTYPES = (torch.float8_e4m3fn,) + + +def _flatten_a_scale(a_scale: torch.Tensor, m: int, kb: int) -> torch.Tensor: + """A's block scales in the physical order the kernel indexes. + + The kernel reads element ``(row, kb)`` at ``kb * M + row``. This does not + guess which spelling of that a caller meant, because two of them are + indistinguishable: + + * **1-D** -- already in physical order. Taken as is. + * **[K/128, M]** -- the physical shape. Taken as is. + * **[M, K/128] column-major** (stride ``(1, M)``) -- the logical shape whose + storage is *already* ``kb``-major. Transposed here, which is free. + ``reshape(-1)`` on it would copy in logical row order and pair every scale + with the wrong K block (relL2 0.26 against 2.4e-3). + * **[M, K/128] row-major** -- **rejected**. It is either a genuinely logical + ``[M, K/128]`` tensor, which needs transposing, or a ``kb``-major buffer + wearing a shape that does not describe it, which must be read flat. + ``aiter_per1x128_quant(transpose_scale=True)`` returns the second: it + writes transposed data and keeps the ``(M, K/128)`` shape. Shape and + stride are identical in both cases, so the caller has to say which by + passing ``.reshape(-1)`` or ``.t()``. + + Guessing here is what broke SGLang: treating the row-major case as logical + took the model's perplexity from 3.26 to 862511 while every shape-level test + stayed green, because the tests build the logical tensor the rule assumes. + """ + if a_scale.dim() == 1: + if a_scale.numel() != m * kb: + raise ValueError( + f"a_scale has {a_scale.numel()} elements, expected M*K/128 = {m * kb}" + ) + return a_scale + if a_scale.dim() != 2: + raise ValueError( + f"a_scale must be [K/128, M], [M, K/128] column-major, or flat; got " + f"shape {tuple(a_scale.shape)}" + ) + shape = tuple(a_scale.shape) + if shape == (kb, m): + return a_scale.reshape(-1) + if shape == (m, kb): + if a_scale.stride() == (1, m): + return a_scale.t().reshape(-1) + raise ValueError( + f"a_scale is {shape} row-major, which is ambiguous: a logical " + f"[M, K/128] needs transposing, while a K/128-major buffer with this " + f"shape -- what aiter_per1x128_quant(transpose_scale=True) returns -- " + f"must be read flat. Pass a_scale.reshape(-1) for the second or " + f"a_scale.t() for the first." + ) + raise ValueError( + f"a_scale shape {shape} matches neither [M, K/128] = {(m, kb)} nor its " + f"transpose" + ) + + +class _Plan(NamedTuple): + """One M's compiled pipeline. + + ``tail`` is already resolved from ``parts["fused_order"]`` into the + launchers to call, in order -- the hot path should not be looking phase + names up in a dict that also holds the order under a string key. + """ + + gemm: object + tail: tuple + input: torch.Tensor + output: torch.Tensor + #: The unresolved phase table, kept for :meth:`GemmAllReduceOp.self_test`, + #: which needs the standalone ``scatter`` that ``fused_order`` omits. + parts: dict + + +class GemmAllReduceOp: + """Fused fp8 block-scale GEMM + all-reduce over cco SDMA. + + Owns the symmetric window and a per-M kernel cache; one instance serves one + ``(n, k)`` weight shape and any ``M`` up to ``m_max``. + + The window cannot grow once allocated, so size ``m_max`` for the largest M + the deployment will see, not for the first one it does. + """ + + def __init__( + self, + comm, + *, + n: int, + k: int, + m_max: int, + block_m: int = DEFAULT_BLOCK_M, + block_n: int = DEFAULT_BLOCK_N, + sdma_queues: int = 1, + gather_dtype: str = "bf16", + gather_transport: str = "lsa", + max_shapes: Optional[int] = None, + ): + # Cleanup state before anything can raise. The rollback below calls + # close(), which clears these; initialising them after the try block + # meant any acquisition failure died in cleanup with an AttributeError, + # masking the real error and skipping the handle release it exists for. + self._closed = True # so a failed constructor leaves close() a no-op + self.mem = self.win = self.dev_comm = None + self._cache: dict[int, _Plan] = {} + self._pad_in: Optional[torch.Tensor] = None + + if gather_dtype not in WIRE_DTYPES: + raise ValueError( + f"gather_dtype={gather_dtype!r} is not one of {sorted(WIRE_DTYPES)}" + ) + if gather_transport not in GATHER_TRANSPORTS: + raise ValueError( + f"gather_transport={gather_transport!r} is not one of " + f"{sorted(GATHER_TRANSPORTS)}" + ) + if block_n != DEFAULT_BLOCK_N: + raise ValueError( + f"block_n must be {DEFAULT_BLOCK_N}: the op always compiles with " + f"permlane, whose lane transpose is written for that width" + ) + # The tile has to be legal before the padding arithmetic runs: both + # padded_m and default_max_shapes divide by world_size * block_m. + why = _tile_constraints(block_m, block_n) + if why is not None: + raise ValueError(f"unsupported tile: {why}") + if m_max < 1 or sdma_queues < 1: + raise ValueError( + f"m_max and sdma_queues must both be >= 1; got m_max={m_max} " + f"sdma_queues={sdma_queues}" + ) + if max_shapes is None: + max_shapes = default_max_shapes( + padded_m(m_max, comm.nranks, block_m), comm.nranks, block_m + ) + if max_shapes < 1: + raise ValueError(f"max_shapes must be >= 1; got {max_shapes}") + why = _gemm_constraints(n, k, block_n) + if why is not None: + raise ValueError(f"unsupported shape: {why}") + self.comm = comm + self.rank = comm.rank + # Communicator calls it nranks; DevCommHandle calls the same thing + # world_size. Normalise here so the op's own surface has one name. + self.world_size = comm.nranks + self.n, self.k = n, k + self.block_m, self.block_n = block_m, block_n + self.sdma_queues = sdma_queues + #: fp8 halves the all-gather's bytes, which is ~40% of a fused layer. + #: It costs relL2 ~2.1e-2 against a bf16 wire that is exact -- e4m3's + #: mantissa, not a tuning knob -- so it is off unless asked for. + self.gather_dtype = gather_dtype + #: Only meaningful with fp8. "lsa" pulls each peer's slice over xGMI and + #: widens it on the way in -- one kernel instead of an SDMA push plus a + #: dequantise, and it does not read the landed fp8 back out of HBM. + #: Measured 957us against SDMA's 1019 on the fused layer. + self.gather_transport = gather_transport + self.max_shapes = max_shapes + self.m_max = padded_m(m_max, self.world_size, block_m) + #: Which counter set each M uses, assigned in first-seen order. The sets + #: are disjoint, so two M values never elect on each other's residue. + self._shape_slot: dict[int, int] = {} + + self.window_bytes = self._make_cfg(self.m_max).window_bytes + try: + self.mem = comm.alloc_mem(self.window_bytes) + self.win = comm.register_window(self.mem.ptr, self.mem.size) + # Counters live in this window and are read before they are first + # written, so it has to start zeroed. + from_gpu_ptr(self.mem.ptr, (self.window_bytes,), torch.uint8).zero_() + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + # One queue per (source, destination) pair is all the pipeline uses + # -- a destination is one xGMI link, and splitting a transfer across + # queues only costs bandwidth. Asking for world_size would create + # world_size queues *per peer* (56 on an 8-rank node, to use 7); + # inside a process that already holds SDMA engines that overruns the + # per-engine queue slots and hsaKmtCreateQueueExt fails. + reqs.sdma_queue_count = sdma_queues + self.dev_comm = comm.create_dev_comm(reqs) + except Exception: + # Roll back whatever was acquired: the communicator holds a strong + # reference to each handle, so a half-built op leaks the window for + # as long as the communicator lives. + self._closed = False + self.close() + raise + self._closed = False + + @staticmethod + def window_bytes_for( + world_size: int, + *, + m_max: int, + n: int, + block_m: int = DEFAULT_BLOCK_M, + gather_dtype: str = "bf16", + gather_transport: str = "lsa", + max_shapes: Optional[int] = None, + ) -> int: + """Symmetric-window bytes an op for this shape will allocate. + + Needed before the op exists: the window comes out of the communicator's + VMM reservation, so ``Communicator.init(per_rank_vmm=...)`` has to be + sized first. Same arithmetic the constructor uses. + """ + return _build_cfg( + world_size=world_size, + m_pad=padded_m(m_max, world_size, block_m), + n=n, + block_m=block_m, + gather_dtype=gather_dtype, + gather_transport=gather_transport, + counter_capacity=MAX_CHUNKS, + shape_slots=max_shapes + or default_max_shapes( + padded_m(m_max, world_size, block_m), world_size, block_m + ), + shape_index=0, + ).window_bytes + + def padded_m(self, m: int) -> int: + """``m`` rounded up to what :meth:`__call__` accepts.""" + return padded_m(m, self.world_size, self.block_m) + + def _slot_for(self, m: int) -> int: + """This M's counter set, assigned on first sight and then fixed.""" + slot = self._shape_slot.get(m) + if slot is None: + if len(self._shape_slot) >= self.max_shapes: + raise ValueError( + f"this op already serves {self.max_shapes} distinct M values " + f"{sorted(self._shape_slot)} and M={m} would need another " + f"counter set; raise max_shapes, or pad to fewer distinct M" + ) + slot = len(self._shape_slot) + self._shape_slot[m] = slot + return slot + + def _make_cfg(self, m: int) -> ArConfig: + return _build_cfg( + world_size=self.world_size, + m_pad=m, + n=self.n, + block_m=self.block_m, + gather_dtype=self.gather_dtype, + gather_transport=self.gather_transport, + # Pinned to the capacity, not to this M's chunk count: the control + # region has to sit at the same offsets for every M the instance + # serves, or one shape reads the previous shape's payload as locks. + counter_capacity=MAX_CHUNKS, + shape_slots=self.max_shapes, + shape_index=self._slot_for(m), + # Every M this instance serves gets the same map. Without it the + # regions of two shapes overlap and a call reads the previous + # shape's leftovers -- see ArConfig.capacity_m. + capacity_m=self.m_max, + ) + + def _compiled(self, m: int): + """Kernel plus SDMA phases for this M, compiled on first sight.""" + hit = self._cache.get(m) + if hit is not None: + return hit + cfg = self._make_cfg(m) + gemm = compile_fused_gemm_scatter( + cfg, + self.rank, + K=self.k, + BLOCK_M=self.block_m, + BLOCK_N=self.block_n, + b_preshuffled=True, + fuse=True, + transport="sdma", + quant="blockscale", + sdma_queues=self.sdma_queues, + # The three C-store stages are off by default in + # compile_fused_gemm_scatter, and blockscale requires swap_ab. + swap_ab=True, + permlane=True, + lane_transpose=True, + ) + parts = build_sdma_phases(cfg, self.rank, queues=self.sdma_queues) + # fused_order rather than a literal list: the fp8 wire inserts a + # quantise and a dequantise around the gather, and a hardcoded + # drain/reduce/gather would skip them and reduce into zeros. + tail = tuple(parts[name] for name in parts["fused_order"]) + hit = _Plan( + gemm=gemm, + tail=tail, + parts=parts, + input=from_gpu_ptr( + self.mem.ptr + cfg.input_off, (m, self.n), torch.bfloat16 + ), + output=from_gpu_ptr( + self.mem.ptr + cfg.output_off, (m, self.n), torch.bfloat16 + ), + ) + self._cache[m] = hit + return hit + + def self_test(self, m: Optional[int] = None) -> None: + """Verify that the collective actually moves bytes. Raises if it does not. + + The failure this exists for is silent. A mori built without + ``BUILD_CCO_SDMA=ON`` -- which is the *default*, and what SGLang's CI + image ships -- still has every symbol, still launches every kernel, and + still returns from every put. The puts simply do nothing. The all-reduce + then returns mostly the local slice, the model keeps answering fluently, + the profile keeps showing all the kernels, and the fused path measures + **faster** than it is because it is not moving any data. An end-to-end + run read -6.8% instead of -2.3% that way, and the only signal that caught + it was perplexity: 862511 against 3.26. + + So: write ``rank + 1`` over the whole input region, run the *split* order + (``scatter -> reduce -> gather``), and check every element came back as + ``sum(1..world_size)``. With the puts inert each rank sees only its own + contribution and the check fails. + + Uses ``parts["order"]`` rather than ``fused_order`` on purpose: + ``fused_order`` opens with ``drain``, which assumes the GEMM epilogue has + already pushed, so without a GEMM it would drain nothing and reduce + whatever the window happened to hold. + + Runs at ``m_max`` by default, which compiles the shape the first real + call needs anyway rather than burning a second one. + """ + if self.mem is None: + raise RuntimeError("this GemmAllReduceOp has been closed") + m = m or self.m_max + + try: + from mori.cco.device._build_flags import BUILD_CCO_SDMA + except ImportError: + BUILD_CCO_SDMA = None + built_without_sdma = BUILD_CCO_SDMA is False + + plan = self._compiled(m) + want = self.world_size * (self.world_size + 1) / 2 + plan.input.fill_(float(self.rank + 1)) + plan.output.zero_() + + stream = fx.Stream(torch.cuda.current_stream()) + for name in plan.parts["order"]: + plan.parts[name](self.dev_comm.ptr, self.win.handle, stream=stream) + torch.cuda.current_stream().synchronize() + + got = plan.output.float() + worst = (got - want).abs().max().item() / want + # The bf16 wire is exact on small integers; the fp8 gather rounds, but + # e4m3 represents integers this small exactly once the row scale is + # applied, so a percent is generous either way. + if worst > 1e-2: + hint = ( + " mori was built with BUILD_CCO_SDMA=OFF (the default), so the " + "SDMA puts are compiled out and do nothing." + if built_without_sdma + else " Check that mori was built with BUILD_CCO_SDMA=ON and that " + "MORI_ENABLE_SDMA=1 is set: without either, every put silently " + "does nothing." + ) + raise RuntimeError( + f"gemm_ar self-test failed: the all-reduce returned " + f"{got.flatten()[0].item():g} where {want:g} was expected " + f"(worst relative error {worst:.3g} at M={m}, world_size=" + f"{self.world_size}).{hint}" + ) + + def close(self) -> None: + """Release this op's window and memory. Idempotent. + + The communicator holds a strong reference to each handle, so dropping + the op is not enough -- without this the window (about 700 MiB at the + model shape) lives until the communicator does. Ownership of ``comm`` + stays with the caller. + + Synchronise first if anything may still be in flight: the phases are + launched on the current stream and this frees memory they address. + """ + if self._closed: + return + self._closed = True + self._cache.clear() + self._pad_in = None + # Window first, then the memory behind it: deregister before free. + # + # The dev-comm is deliberately left to the communicator, which is also + # what EpDispatchCombineOp does. Destroying it here and then letting the + # communicator tear down aborts the process (exit 255, after this op's + # own work has finished) -- its SDMA queues are comm-scoped, and the + # queue count is fixed by the first ccoDevCommCreate on a comm anyway. + # The window is what this exists to release: ~700 MiB at the model shape, + # against a handful of queues. + for handle in (self.win, self.mem): + if handle is not None: + handle.close() + self.win = self.mem = None + + def __enter__(self) -> "GemmAllReduceOp": + return self + + def __exit__(self, *exc) -> None: + self.close() + + def pad_rows(self, x: torch.Tensor, m_pad: int) -> torch.Tensor: + """Zero-extend ``x`` to ``m_pad`` rows, in a buffer reused across calls. + + A GEMM row and a reduce-scatter row both depend only on the same input + row, so the added rows produce zeros that the caller slices off. Pad + *before* quantising: that is what keeps the A scale's column-major + ``[K/128, M]`` layout intact without touching it. + + The returned tensor is a view of a buffer this op reuses, so a later + ``pad_rows`` call overwrites it. Clone it to keep it. + """ + if x.dim() != 2: + raise ValueError( + f"expected a 2-D [M, K] tensor, got shape {tuple(x.shape)}" + ) + m, k = x.shape + if m_pad < m: + raise ValueError(f"m_pad={m_pad} is smaller than the input's M={m}") + # Key on dtype and device as well as capacity: reusing a bf16 buffer for + # an fp32 input would silently round it, and this promises zero-extension + # and nothing else. + buf = self._pad_in + if ( + buf is None + or buf.shape[0] < m_pad + or buf.shape[1] != k + or buf.dtype != x.dtype + or buf.device != x.device + ): + buf = self._pad_in = torch.zeros((m_pad, k), dtype=x.dtype, device=x.device) + buf = buf[:m_pad] + buf[:m].copy_(x) + buf[m:].zero_() + return buf + + def __call__( + self, + a_fp8: torch.Tensor, + b_preshuffled: torch.Tensor, + a_scale: torch.Tensor, + b_scale: torch.Tensor, + ) -> torch.Tensor: + """One fused GEMM + all-reduce; returns a view of the window's output. + + ``a_fp8`` is ``[M, K]`` with ``M`` a multiple of ``world_size * + block_m`` (see :meth:`padded_m`), ``b_preshuffled`` is ``[N, K]`` + through :func:`preshuffle_b`. + + ``a_scale`` may be ``[M, K/128]`` (the logical shape, either memory + order), ``[K/128, M]``, or already flat in physical order -- see + :func:`_flatten_a_scale`, which is where the distinction is made rather + than assumed. ``b_scale`` is ``[N/128, K/128]`` row-major. + + The returned tensor aliases the window and is overwritten by the next + call. Clone it to keep it. One instance is not usable concurrently: the + window, the kernel cache and the padding buffer are all shared, and the + phases run on the current stream. + """ + m = a_fp8.shape[0] + if m > self.m_max: + raise ValueError( + f"M={m} exceeds the window's m_max={self.m_max}; the window " + f"cannot grow after construction" + ) + if m % (self.world_size * self.block_m): + raise ValueError( + f"M={m} must be a multiple of world_size * block_m = " + f"{self.world_size * self.block_m}; use padded_m()/pad_rows()" + ) + self._check_operands(a_fp8, b_preshuffled, a_scale, b_scale, m) + plan = self._compiled(m) + stream = fx.Stream(torch.cuda.current_stream()) + plan.gemm( + a_fp8.contiguous().view(torch.int8).view(-1), + b_preshuffled.contiguous().view(torch.int8).view(-1), + plan.input.view(-1), + _flatten_a_scale(a_scale, m, self.k // SCALE_BLOCK_K), + b_scale.reshape(-1), + m, + self.n, + self.dev_comm.ptr, + self.win.handle, + stream=stream, + ) + for phase in plan.tail: + phase(self.dev_comm.ptr, self.win.handle, stream=stream) + return plan.output + + def _check_operands(self, a_fp8, b_preshuffled, a_scale, b_scale, m: int) -> None: + """Reject what the launch boundary would otherwise reinterpret as bytes. + + ``a_fp8``/``b_preshuffled`` go in as ``int8``, so a bf16 tensor of the + right shape reaches the kernel and comes back finite and wrong. These + are all metadata checks -- no device work, nothing that shows in a + profile. + """ + if self.mem is None: + raise RuntimeError("this GemmAllReduceOp has been closed") + kb = self.k // SCALE_BLOCK_K + for name, t, shape in ( + ("a_fp8", a_fp8, (m, self.k)), + ("b_preshuffled", b_preshuffled, (self.n, self.k)), + ): + if t.dim() != 2 or tuple(t.shape) != shape: + raise ValueError(f"{name} must be {shape}, got {tuple(t.shape)}") + if t.dtype not in FP8_DTYPES: + raise ValueError( + f"{name} must be one of {[str(d) for d in FP8_DTYPES]}, got " + f"{t.dtype}; the launch reinterprets it as raw bytes, so a " + f"wider dtype runs and returns nonsense rather than failing" + ) + if t.device.type != "cuda": + raise ValueError(f"{name} must be on a GPU, got {t.device}") + for name, t, numel in ( + ("a_scale", a_scale, m * kb), + ("b_scale", b_scale, (self.n // SCALE_BLOCK_K) * kb), + ): + if t.dtype != torch.float32: + raise ValueError(f"{name} must be float32, got {t.dtype}") + if t.numel() != numel: + raise ValueError(f"{name} must have {numel} elements, got {t.numel()}") + if t.device.type != "cuda": + raise ValueError(f"{name} must be on a GPU, got {t.device}") diff --git a/src/cco/device/cco_device_wrapper.cpp b/src/cco/device/cco_device_wrapper.cpp index a139a0662..b48b96229 100644 --- a/src/cco/device/cco_device_wrapper.cpp +++ b/src/cco/device/cco_device_wrapper.cpp @@ -102,6 +102,31 @@ CCO_DEV uint64_t cco_lsa_ptr(uint64_t window, int peer, uint64_t offset) { return reinterpret_cast(w->winBase) + static_cast(peer) * stride + offset; } +// ── LSA window geometry ─────────────────────────────────────────────────────── +// `cco_lsa_ptr` above loads winBase and stride4G on *every* call. A kernel that +// needs many addresses from one window can take the geometry once with these two +// and compute `base + peer*stride + offset` itself. +// +// The pointer is explicitly address_space(1). `AsWindow` casts a uint64 to a +// *generic* pointer, and a generic access must be a flat_load: the compiler +// cannot prove the target is not LDS, so it counts against lgkmcnt as well as +// vmcnt. Naming the address space lets it emit global_load, which is vmcnt-only +// -- one fewer counter for any following s_waitcnt to wait on. +typedef const ccoWindowDevice __attribute__((address_space(1))) * ccoWindowGlobalPtr; + +__device__ __attribute__((always_inline)) inline ccoWindowGlobalPtr AsWindowGlobal(uint64_t h) { + return reinterpret_cast(h); +} + +CCO_DEV uint64_t cco_lsa_win_base(uint64_t window) { + return reinterpret_cast(AsWindowGlobal(window)->winBase); +} + +// Returns the *byte* stride between peers, i.e. stride4G already shifted. +CCO_DEV uint64_t cco_lsa_stride(uint64_t window) { + return static_cast(AsWindowGlobal(window)->stride4G) << 32; +} + // System-scope publication fence used by direct LSA benchmark/store paths. // Latency kernels fence from block lane 0; bandwidth kernels fence every lane. // This primitive is NOT a barrier: callers must synchronize participating diff --git a/tests/python/cco/gemm_ar_op_worker.py b/tests/python/cco/gemm_ar_op_worker.py new file mode 100644 index 000000000..a758cfcff --- /dev/null +++ b/tests/python/cco/gemm_ar_op_worker.py @@ -0,0 +1,377 @@ +#!/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. +"""Ranks for ``test_gemm_ar_op.py``: exercises ``GemmAllReduceOp`` itself. + +Separate from ``bench_gemm_ar.py`` on purpose. That driver calls the low-level +builders, so every number it produces is silent about the public wrapper -- the +column-major A scale, the per-M control layout and the operand checks all live +between the two, and had no coverage at all until this existed. + +Prints one ``RESULT_JSON`` line per case from rank 0. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +import torch +import torch.distributed as dist + +from mori.cco import Communicator, UniqueId +from mori.ops.gemm_ar import GemmAllReduceOp, preshuffle_b + +SCALE_BK = 128 + + +def _setup(): + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) + if not dist.is_initialized(): + dist.init_process_group(backend="cpu:gloo") + rank, world = dist.get_rank(), dist.get_world_size() + payload = [bytes(Communicator.get_unique_id()) if rank == 0 else None] + dist.broadcast_object_list(payload, src=0) + return rank, world, UniqueId.from_bytes(payload[0]) + + +def _operands(rank: int, m: int, n: int, k: int, salt: int = 0): + """fp8 operands with block scales, in the layouts the model produces. + + ``sa`` is built column-major, which is what + ``aiter_per1x128_quant(transpose_scale=True)`` and sglang's + ``materialize_bpreshuffle_fp8_scale`` return, and is the case the public op + used to flatten in the wrong order. + """ + g = torch.Generator(device="cuda").manual_seed(1234 + rank + 7919 * salt) + a = (torch.randn(m, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + b = (torch.randn(n, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + kb = k // SCALE_BK + sa = ( + torch.rand(m, kb, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + ) + sa = sa.t().contiguous().t() # logical [M, kb], physically [kb, M] + sb = ( + torch.rand(n // SCALE_BK, kb, generator=g, device="cuda", dtype=torch.float32) + * 0.01 + + 0.01 + ) + return a, b, sa, sb + + +def _reference(world: int, m: int, n: int, k: int, salt: int = 0) -> torch.Tensor: + """All-reduced fp32 reference, built from the same recipe on every rank.""" + acc = torch.zeros(m, n, device="cuda", dtype=torch.float32) + kb = k // SCALE_BK + for r in range(world): + a, b, sa, sb = _operands(r, m, n, k, salt) + af, bf = a.float(), b.float() + part = torch.zeros(m, n, device="cuda", dtype=torch.float32) + for j in range(kb): + sl = slice(j * SCALE_BK, (j + 1) * SCALE_BK) + blk = af[:, sl] @ bf[:, sl].t() + part += blk * sa[:, j : j + 1] * sb[:, j].repeat_interleave(SCALE_BK)[:n] + acc += part + return acc + + +def _rel_l2(got: torch.Tensor, ref: torch.Tensor) -> float: + return ( + torch.linalg.vector_norm(got.float() - ref) / torch.linalg.vector_norm(ref) + ).item() + + +def _emit(rank: int, **record) -> None: + if rank == 0: + print("RESULT_JSON " + json.dumps(record, sort_keys=True), flush=True) + + +def case_scale_order(op, rank, world, m, n, k): + """The documented [M, K/128] column-major scale must give the fp8 floor. + + Passing the same tensor through reshape(-1) -- the old behaviour -- walks the + logical rows and pairs each scale with the wrong K block. + """ + a, b, sa, sb = _operands(rank, m, n, k) + ref = _reference(world, m, n, k) + bs = preshuffle_b(b) + + got = op(a, bs, sa, sb).clone() + logical = _rel_l2(got, ref) + + # The two other accepted spellings must agree with it bitwise. + flat = op(a, bs, sa.t().reshape(-1), sb).clone() + physical = op(a, bs, sa.t(), sb).clone() + + _emit( + rank, + case="scale_order", + rel_l2=logical, + flat_matches=bool(torch.equal(got, flat)), + physical_matches=bool(torch.equal(got, physical)), + ) + + +def _peer_c(peer: int, m: int, n: int, k: int) -> torch.Tensor: + """Peer's *unreduced* GEMM output at this shape, in bf16. + + What the scatter pushes into my recv slot for that peer, so it is also what + a stale copy of that push would leave behind. + """ + a, b, sa, sb = _operands(peer, m, n, k) + af, bf = a.float(), b.float() + out = torch.zeros(m, n, device="cuda", dtype=torch.float32) + kb = k // SCALE_BK + for j in range(kb): + sl = slice(j * SCALE_BK, (j + 1) * SCALE_BK) + out += ( + (af[:, sl] @ bf[:, sl].t()) + * sa[:, j : j + 1] + * sb[:, j].repeat_interleave(SCALE_BK)[:n] + ) + return out.to(torch.bfloat16) + + +def _row_report(got, ref, m): + """Which rows are wrong, as a contiguous-block summary. + + A fixed wrong value hit intermittently means the race decides *whether* a + region is read stale, not *what* is in it -- so the row map is what points + at the region. + """ + err = (got.float() - ref).abs().amax(dim=1) + scale = ref.abs().amax(dim=1).clamp_min(1e-6) + bad = (err / scale > 0.05).nonzero().flatten() + if bad.numel() == 0: + return {"bad_rows": 0} + return { + "bad_rows": int(bad.numel()), + "m": m, + "first_bad": int(bad[0]), + "last_bad": int(bad[-1]), + "contiguous": bool(bad.numel() == int(bad[-1]) - int(bad[0]) + 1), + } + + +def case_alternating_m(op, rank, world, m_small, m_large, n, k): + """Two M values through one instance, alternating, then repeated. + + The control region moves with counter_chunks unless it is sized for the + capacity, and the election counters carry a residue between shapes unless + each shape owns a set. Either one hangs or corrupts the second call. + """ + results = {} + for label, m in ( + ("first_small", m_small), + ("first_large", m_large), + ("again_small", m_small), + ("again_large", m_large), + ): + a, b, sa, sb = _operands(rank, m, n, k) + got = op(a, preshuffle_b(b), sa, sb).clone() + ref = _reference(world, m, n, k) + results[label] = _rel_l2(got, ref) + if results[label] > 0.01: + rep = _row_report(got, ref, m) + # Hypothesis: the bad range is exactly where the *previous* shape's + # recv slot for peer 1 sat, and still holds that peer's scatter + # payload from the previous call. + if rep["bad_rows"] and m_small != m: + stale = _peer_c(1, m_small, n, k)[: rep["bad_rows"]] + seen = got[rep["first_bad"] : rep["last_bad"] + 1] + rep["matches_previous_scatter"] = bool(torch.equal(seen, stale)) + rep["max_abs_diff_vs_stale"] = float( + (seen.float() - stale.float()).abs().max() + ) + results[f"{label}_rows"] = rep + _emit( + rank, + case=f"alternating_m_{m_small}_{m_large}", + **results, + ) + + +def case_self_test(op, rank, world, n, k): + """self_test must pass on a working stack, and its check must have teeth. + + The second half matters more than the first: a check that cannot fail is + exactly what let a whole end-to-end campaign run against a mori whose SDMA + puts were compiled out. + """ + op.self_test() + ok = True + # Now break it the way a dead transport would: run `reduce -> gather` with no + # scatter, which is what an inert put amounts to. + # + # The landing slots have to be zeroed first, and that is not a detail. This + # op has already served several cases, so the slots still hold a previous + # call's peer data -- reducing that produces a plausible sum and the check + # passes, which is the same stale-read effect that made an inert transport + # look correct in the first place. Without this the teeth test has none. + from mori.tensor_utils import from_gpu_ptr + + cfg = op._make_cfg(op.m_max) + from_gpu_ptr(op.mem.ptr + cfg.recv_off, (cfg.recv_bytes,), torch.uint8).zero_() + plan = op._compiled(op.m_max) + plan.input.fill_(float(rank + 1)) + plan.output.zero_() + import flydsl.expr as _fx + + stream = _fx.Stream(torch.cuda.current_stream()) + for name in plan.parts["order"]: + if name == "scatter": + continue + plan.parts[name](op.dev_comm.ptr, op.win.handle, stream=stream) + torch.cuda.current_stream().synchronize() + want = world * (world + 1) / 2 + detected = bool((plan.output.float() - want).abs().max().item() / want > 1e-2) + _emit( + rank, + case="self_test", + passes_on_working_stack=ok, + detects_missing_scatter=detected, + ) + + +def case_close_is_idempotent(comm, rank, n, k, m_max): + """close() releases, twice is a no-op, and a closed op refuses to run.""" + op = GemmAllReduceOp(comm, n=n, k=k, m_max=m_max) + op.close() + op.close() + refused = False + try: + op( + torch.zeros(m_max, k, device="cuda", dtype=torch.float8_e4m3fn), + torch.zeros(n, k, device="cuda", dtype=torch.float8_e4m3fn), + torch.zeros(m_max * (k // SCALE_BK), device="cuda", dtype=torch.float32), + torch.zeros( + (n // SCALE_BK) * (k // SCALE_BK), device="cuda", dtype=torch.float32 + ), + ) + except RuntimeError: + refused = True + with GemmAllReduceOp(comm, n=n, k=k, m_max=m_max): + pass # context manager closes on exit + _emit(rank, case="close", refused_after_close=refused) + + +def case_changing_data(op, rank, world, m, n, k, calls): + """Same M, different operands each call. + + Every existing test and benchmark feeds identical operands on every + iteration, so a phase that reads a peer's slice before it has landed returns + the *previous* call's bytes -- which are bit-identical to the right answer. + Changing the data between calls is what makes such a read visible. + """ + out = {} + for i in range(calls): + a, b, sa, sb = _operands(rank, m, n, k, salt=i) + got = op(a, preshuffle_b(b), sa, sb).clone() + out[f"call{i}"] = _rel_l2(got, _reference(world, m, n, k, salt=i)) + _emit(rank, case="changing_data", **out) + + +#: Every case this worker knows, and the operands each wants, so the whole file +#: can be served by one spawn. +#: +#: Six separate 2-rank launches is what it cost before, and the queues do not +#: recycle that fast: back-to-back 8-rank jobs fail in hsaKmtCreateQueueExt +#: (anvil.cpp:237), and the ones that fell over were as often the *neighbouring* +#: file's as this one's. One process, one communicator, every case. +ALL_CASES = [ + ("scale_order", {"m": 512}), + ("alternating_m", {"m_small": 512, "m_large": 1024}), + ("alternating_m", {"m_small": 2048, "m_large": 4096}), + ("changing_data", {"m": 512, "calls": 5}), + ("self_test", {"m": 1024}), +] + + +def run_all(comm, rank, world, n, k): + """Every case in one process, sharing one communicator. + + m_max covers the largest M any case asks for, so one op serves them all -- + which is also closer to how a server uses it than a fresh op per shape. + """ + m_max = max(max(kw.get("m", 0), kw.get("m_large", 0)) for _, kw in ALL_CASES) + with GemmAllReduceOp(comm, n=n, k=k, m_max=m_max) as op: + for case, kw in ALL_CASES: + if case == "scale_order": + case_scale_order(op, rank, world, kw["m"], n, k) + elif case == "alternating_m": + case_alternating_m(op, rank, world, kw["m_small"], kw["m_large"], n, k) + elif case == "changing_data": + case_changing_data(op, rank, world, kw["m"], n, k, kw["calls"]) + elif case == "self_test": + case_self_test(op, rank, world, n, k) + # close() needs its own ops, so it comes after the shared one is released. + case_close_is_idempotent(comm, rank, n, k, 512) + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--case", required=True) + p.add_argument("-n", type=int, default=1024) + p.add_argument("-k", type=int, default=512) + p.add_argument("-m", type=int, default=0) + p.add_argument("--m-small", type=int, default=0) + p.add_argument("--m-large", type=int, default=0) + p.add_argument("--calls", type=int, default=4) + args = p.parse_args() + + rank, world, uid = _setup() + if args.case == "all": + # The cases carry their own shapes; the command line has none. + m_max = max(max(kw.get("m", 0), kw.get("m_large", 0)) for _, kw in ALL_CASES) + else: + m_max = max(args.m, args.m_large, args.m_small) + vmm = 2 * GemmAllReduceOp.window_bytes_for(world, m_max=m_max, n=args.n) + ( + 256 << 20 + ) + with Communicator.init(world, rank, uid, per_rank_vmm=vmm) as comm: + if args.case == "all": + run_all(comm, rank, world, args.n, args.k) + return 0 + if args.case == "close": + case_close_is_idempotent(comm, rank, args.n, args.k, m_max) + return 0 + with GemmAllReduceOp(comm, n=args.n, k=args.k, m_max=m_max) as op: + if args.case == "scale_order": + case_scale_order(op, rank, world, args.m, args.n, args.k) + elif args.case == "self_test": + case_self_test(op, rank, world, args.n, args.k) + elif args.case == "changing_data": + case_changing_data(op, rank, world, args.m, args.n, args.k, args.calls) + elif args.case == "alternating_m": + case_alternating_m( + op, rank, world, args.m_small, args.m_large, args.n, args.k + ) + else: + raise SystemExit(f"unknown case {args.case}") + dist.barrier() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/python/cco/test_flydsl_ar.py b/tests/python/cco/test_flydsl_ar.py new file mode 100644 index 000000000..bd1d262cf --- /dev/null +++ b/tests/python/cco/test_flydsl_ar.py @@ -0,0 +1,282 @@ +# 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. +"""FlyDSL cco all-reduce: layout unit tests + multi-GPU smoke. + +Mirrors ``test_triton_gemm_a2a.py``: the pure-Python half checks the arithmetic +the kernels depend on without a GPU, the GPU half runs ``bench_ar.py`` under +torchrun and parses its ``RESULT_JSON`` line. +""" + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +from mori.ops.gemm_ar import layout + +REPO_ROOT = Path(__file__).resolve().parents[3] +AR_DIR = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "ar" + + +# -------------------------------------------------------------------------- +# layout / dispatch arithmetic (no GPU) +# -------------------------------------------------------------------------- + + +def _cfg(**kw): + base = dict(world_size=8, m=4096, n=7168) + base.update(kw) + return layout.ArConfig(**base) + + +def test_window_regions_are_disjoint_and_ordered(): + c = _cfg() + c.validate() + assert c.input_off >= c.signal_bytes + assert c.output_off == c.input_off + c.nbytes + assert c.tmp_off == c.output_off + c.nbytes + assert c.window_bytes == c.tmp_off + c.tmp_bytes + # signal sub-regions must not overlap the payload + assert c.flag_off + c.max_blocks * 4 <= c.input_off + + +def test_recv_region_is_absent_unless_sdma_asks_for_it(): + """LSA reduces out of the peers' inputs, so it must not pay for landing slots.""" + lsa = _cfg() + assert lsa.recv_bytes == 0 + assert lsa.window_bytes == lsa.tmp_off + lsa.tmp_bytes + with pytest.raises(IndexError, match="no landing slot"): + lsa.recv_slot_off(0) + + +def test_recv_slots_are_one_contiguous_slice_per_peer(): + c = _cfg(recv_slots=8) + assert c.recv_off == c.tmp_off + c.tmp_bytes + assert c.recv_bytes == 8 * c.slice_bytes + assert c.window_bytes == c.recv_off + c.recv_bytes + offs = [c.recv_slot_off(p) for p in range(8)] + assert offs == sorted(offs) + assert all(b - a == c.slice_bytes for a, b in zip(offs, offs[1:])) + with pytest.raises(IndexError): + c.recv_slot_off(8) + + +def test_validate_rejects_more_recv_slots_than_peers(): + with pytest.raises(ValueError, match="recv_slots"): + _cfg(recv_slots=9).validate() + + +def test_signal_slot_is_block_major_over_peers(): + c = _cfg() + assert c.signal_slot(0, 0) == 0 + assert c.signal_slot(0, 7) == 7 + assert c.signal_slot(1, 0) == layout.MAX_WORLD + assert c.signal_slot(79, 7) == 79 * layout.MAX_WORLD + 7 + with pytest.raises(IndexError): + c.signal_slot(layout.K_MAX_BLOCKS, 0) + with pytest.raises(IndexError): + c.signal_slot(0, layout.MAX_WORLD) + + +def test_pack_and_shard_math(): + c = _cfg(m=4096, n=7168) # 4096*7168*2 = 58,720,256 B + assert c.nbytes == 4096 * 7168 * 2 + assert c.num_packs == c.nbytes // 16 + assert c.packs_per_rank == c.num_packs // 8 + start, end = c.owner_pack_range(0) + assert (start, end) == (0, c.packs_per_rank) + # the last rank absorbs the remainder so no pack is dropped + assert c.owner_pack_range(7)[1] == c.num_packs + covered = sum(e - s for s, e in (c.owner_pack_range(r) for r in range(8))) + assert covered == c.num_packs + + +def test_slice_is_contiguous_and_large_enough_for_sdma(): + """M-sharding must give one contiguous run per peer, not m strided pieces.""" + c = _cfg(m=4096, n=7168) + assert c.slice_bytes == (4096 // 8) * 7168 * 2 + # >=1MB is where SDMA leaves the ~6us dispatch floor behind + assert c.slice_bytes > 1 << 20 + + +@pytest.mark.parametrize( + "world_size,nbytes,expected", + [ + (2, 1 << 30, 1), # world==2 is always 1-stage + (4, 100 * 1024, 1), # < 160KB + (4, 200 * 1024, 2), + (8, 64 * 1024, 1), # < 80KB + (8, 96 * 1024, 2), + (8, 80 * 1024, 2), # threshold is strict < + ], +) +def test_stage_selection_matches_aiter(world_size, nbytes, expected): + assert layout.select_stage(world_size, nbytes) == expected + + +def test_decode_shape_picks_two_stage(): + """[M=64, 7168] bf16 = 896KB -> above the 80KB gate, so 2-stage.""" + c = _cfg(m=64) + assert c.nbytes == 64 * 7168 * 2 + assert c.stage == 2 + + +def test_blocks_are_capped_and_never_zero(): + assert _cfg(m=16384).blocks == layout.LSA_BLOCK_CAP + assert _cfg(m=1).blocks >= 1 + assert _cfg(m=1).blocks <= layout.LSA_BLOCK_CAP + # the cap is a tuning knob, the signal array is the correctness bound + assert layout.LSA_BLOCK_CAP <= layout.K_MAX_BLOCKS + + +def test_blocks_stride_the_index_space_by_threads_not_by_peer_group(): + """One thread owns one index across all peers, so a block covers `threads`.""" + c = _cfg(m=4096) + small = layout.ArConfig(world_size=8, m=8, n=7168) + assert small.packs_per_rank == 896 + assert small.blocks == 2 # ceil(896/512), not aiter's ceil(896/64) = 14 + assert c.blocks == layout.LSA_BLOCK_CAP + + +def test_force_blocks_must_fit_the_signal_array(): + with pytest.raises(ValueError, match="force_blocks"): + layout.ArConfig( + world_size=8, m=4096, n=7168, force_blocks=layout.K_MAX_BLOCKS + 1 + ).validate() + + +def test_remote_traffic_model(): + c = _cfg(m=4096) + assert c.stage == 2 + # 2-stage moves 2*(P-1)/P of the payload, i.e. 1.75x at P=8 + assert c.remote_bytes_per_rank == 2 * 7 * c.nbytes // 8 + + +def test_validate_rejects_indivisible_world(): + with pytest.raises(ValueError, match="divisible"): + _cfg(world_size=6).validate() + + +def test_validate_rejects_unvectorizable_payload(): + with pytest.raises(ValueError, match="multiple"): + # 7 elements of bf16 is not a multiple of world*16 + layout.ArConfig(world_size=8, m=1, n=7).validate() + + +# -------------------------------------------------------------------------- +# GPU smoke +# -------------------------------------------------------------------------- + + +def _run_bench(world_size, backend, m, n, extra_env=None, extra=()): + env = os.environ.copy() + env.setdefault("MORI_SOCKET_IFNAME", "lo") + env.setdefault("MORI_ENABLE_SDMA", "1") + if extra_env: + env.update(extra_env) + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc_per_node={world_size}", + str(AR_DIR / "bench_ar.py"), + "--backend", + backend, + "-m", + str(m), + "-n", + str(n), + "--warmup", + "1", + "--iters", + "3", + *extra, + ] + result = subprocess.run( + command, cwd=REPO_ROOT, env=env, capture_output=True, text=True, timeout=600 + ) + output = result.stdout + result.stderr + assert result.returncode == 0, output + records = [ + json.loads(line.removeprefix("RESULT_JSON ")) + for line in output.splitlines() + if line.startswith("RESULT_JSON ") + ] + assert len(records) == 1, output + return records[0] + + +@pytest.mark.parametrize("backend", ["lsa", "sdma"]) +@pytest.mark.parametrize("world_size,m,n", [(2, 256, 1024), (4, 256, 1024)]) +def test_flydsl_ar_gpu_smoke(backend, world_size, m, n): + if torch.cuda.device_count() < world_size: + pytest.skip(f"requires {world_size} GPUs") + record = _run_bench(world_size, backend, m, n) + assert record["validated"] is True + assert record["max_rank_time_ms"] > 0 + assert record["backend"] == backend + + +@pytest.mark.parametrize("world_size,m,n", [(8, 4096, 7168)]) +def test_fp8_gather_is_correct_and_stays_at_the_e4m3_floor(world_size, m, n): + """The fp8 gather leg, against the same host reference the bf16 wire uses. + + Two assertions, and the upper one matters as much as the lower. Below 4e-2 + says the kernel is right: a structurally broken quantise/gather/dequantise + lands near 0.9 (that is sqrt(7/8) -- every slice that travelled arriving as + zeros). *Above* 5e-3 says fp8 is actually being used, so a silent fallback + to the bf16 path cannot pass this as a success. + + The floor itself is e4m3's 3-bit mantissa and not the scale granularity: + per-row costs 2.65e-2 on a normal payload and per-32 costs 2.40e-2, 9% + better for 200x the scale bytes. There is no granularity that avoids it. + """ + if torch.cuda.device_count() < world_size: + pytest.skip(f"requires {world_size} GPUs") + record = _run_bench( + world_size, "sdma", m, n, {"AR_EXTRA": ""}, extra=["--gather-dtype", "fp8"] + ) + assert record["validated"] is True + assert record["rel_l2"] < 4e-2, ( + f"relL2 {record['rel_l2']:.3e}: the fp8 gather is structurally wrong, " + f"not merely quantised" + ) + assert record["rel_l2"] > 5e-3, ( + f"relL2 {record['rel_l2']:.3e} is at the bf16 path's exactness -- the " + f"fp8 wire cannot have been taken" + ) + + +def test_flydsl_ar_matches_aiter_bitwise(): + """Every element is reduced once on its owner rank, so the two agree exactly.""" + world_size = 2 + if torch.cuda.device_count() < world_size: + pytest.skip(f"requires {world_size} GPUs") + pytest.importorskip("aiter", reason="aiter baseline not installed in this venv") + record = _run_bench(world_size, "lsa", 256, 1024, {"AR_CHECK_VS_AITER": "1"}) + assert record["validated"] is True + assert record.get("rel_l2_vs_aiter") == 0.0 diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py new file mode 100644 index 000000000..882d7c0fc --- /dev/null +++ b/tests/python/cco/test_gemm_ar.py @@ -0,0 +1,556 @@ +# 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. +"""Numerics for ``mori.ops.gemm_ar``: the GEMM, the epilogue, and the layout math. + +Three layers, cheapest first. The counter-region tests are pure arithmetic and +need no GPU. ``test_blockscale_*`` and ``test_swap_ab_*`` run the GEMM on one +GPU against an fp32 reference, with no cco at all -- they are what catches the +kernel drifting. The ``_run_bench`` tests spawn ranks and check the modes agree. +""" + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +from mori.ops.gemm_ar import layout + +REPO_ROOT = Path(__file__).resolve().parents[3] +BENCH = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "gemm_ar" / "bench_gemm_ar.py" + + +# -------------------------------------------------------------------------- +# counter-region layout (no GPU) +# -------------------------------------------------------------------------- + + +def test_counter_region_sizes_with_chunks_and_never_overlaps_payload(): + a = layout.ArConfig(world_size=8, m=4096, n=7168, counter_chunks=1) + b = layout.ArConfig(world_size=8, m=4096, n=7168, counter_chunks=4) + a.validate() + b.validate() + # own region, past the last _flag slot + assert a.counter_off >= a.flag_off + a.max_blocks * 4 + assert a.counter_off % layout.SIGNAL_ALIGN == 0 + # grows with chunks, and the payload is pushed out to match + assert b.counter_bytes >= a.counter_bytes + assert b.input_off >= a.input_off + assert a.counter_off + a.counter_bytes <= a.input_off + assert b.counter_off + b.counter_bytes <= b.input_off + + +def test_counter_slot_is_dest_major(): + c = layout.ArConfig(world_size=8, m=4096, n=7168, counter_chunks=2) + assert c.counter_slot(0, 0) == 0 + assert c.counter_slot(0, 1) == 1 + assert c.counter_slot(1, 0) == 2 + assert c.counter_slot(7, 1) == 15 + with pytest.raises(IndexError): + c.counter_slot(8, 0) + with pytest.raises(IndexError): + c.counter_slot(0, 2) + + +def test_validate_rejects_zero_chunks(): + with pytest.raises(ValueError, match="counter_chunks"): + layout.ArConfig(world_size=8, m=4096, n=7168, counter_chunks=0).validate() + + +def test_bf16_wire_layout_is_unchanged_by_the_fp8_option(): + """Adding the fp8 wire must not move a single bf16 offset. + + The bf16 path is what ships today and what every recorded measurement is + on, so the fp8 regions have to be additive. On bf16 ``gout`` aliases + ``output`` and every scale region is zero-sized. + """ + c = layout.ArConfig(world_size=8, m=16384, n=7168, recv_slots=8, counter_chunks=8) + c.validate() + assert c.scatter_dtype == "bf16" and c.gather_dtype == "bf16" + # The point of pinning this is to catch layout drift nobody meant, so it + # moves only with a stated reason. 734009088 was the size the running + # server logged before the fp8 regions existed; +128 is the gather band + # counters, one SIGNAL_ALIGN region added when the reduce learned to fire + # the gather's puts itself. + assert c.window_bytes == 734009088 + 128 + assert c.gout_off == c.output_off + assert c.gout_bytes == 0 + assert c.scatter_scale_bytes == 0 + assert c.gather_scale_bytes == 0 + assert c.recv_scale_bytes == 0 + assert c.scatter_slice_bytes == c.slice_bytes + assert c.gather_slice_bytes == c.slice_bytes + + +def test_fp8_wire_halves_the_payload_and_every_region_is_disjoint(): + c = layout.ar_config( + world_size=8, + m=16384, + n=7168, + recv_slots=8, + counter_chunks=8, + comm_dtype="fp8", + ) + c.validate() + assert c.scatter_slice_bytes * 2 == c.slice_bytes + assert c.gather_slice_bytes * 2 == c.slice_bytes + assert c.scatter_nbytes * 2 == c.nbytes + # output stays bf16: it is the consumer's tensor, not a wire format. + assert c.nbytes == c.m * c.n * 2 + + regions = [ + ("input", c.input_off, c.scatter_nbytes), + ("input_scale", c.input_scale_off, c.scatter_scale_bytes), + ("output", c.output_off, c.nbytes), + ("gout", c.gout_off, c.gout_bytes), + ("gout_scale", c.gout_scale_off, c.gather_scale_bytes), + ("tmp", c.tmp_off, c.tmp_bytes), + ("recv", c.recv_off, c.recv_bytes), + ("recv_scale", c.recv_scale_off, c.recv_scale_bytes), + ] + regions.sort(key=lambda r: r[1]) + for (an, ao, asz), (bn, bo, _) in zip(regions, regions[1:]): + assert ao + asz <= bo, f"{an} [{ao}, {ao + asz}) overruns {bn} at {bo}" + last_n, last_o, last_sz = regions[-1] + assert last_o + last_sz <= c.window_bytes, f"{last_n} past the window" + + +def test_fp8_scale_slots_are_peer_major_and_stay_inside_their_region(): + c = layout.ar_config( + world_size=8, + m=16384, + n=7168, + recv_slots=8, + counter_chunks=8, + comm_dtype="fp8", + ) + per_peer = c.scatter_scale_slice_bytes + assert per_peer == c.slice_rows * c.scatter_tiles_per_row * 4 + for peer in range(8): + off = c.recv_scale_slot_off(peer) + assert off == c.recv_scale_off + peer * per_peer + assert off + per_peer <= c.recv_scale_off + c.recv_scale_bytes + # the gather leg is one scale per row, so a slice's scales are slice_rows + assert c.gather_scale_slice_bytes == c.slice_rows * 4 + for peer in range(8): + off = c.gout_scale_slice_off(peer) + assert ( + off + c.gather_scale_slice_bytes <= c.gout_scale_off + c.gather_scale_bytes + ) + + +def test_the_two_wire_legs_are_independent(): + """Each leg sizes only its own regions. + + They are separate because the gather leg is where nearly all the time is and + where fp8 costs one rounding, while the scatter leg is already hidden behind + the GEMM and its fp8 error compounds across ranks -- so being able to take + one without the other is the point, not a convenience. + """ + gather_only = layout.ArConfig( + world_size=8, m=4096, n=7168, recv_slots=8, gather_dtype="fp8" + ) + gather_only.validate() + assert gather_only.scatter_slice_bytes == gather_only.slice_bytes # untouched + assert gather_only.gather_slice_bytes * 2 == gather_only.slice_bytes + assert gather_only.scatter_scale_bytes == 0 + assert gather_only.gather_scale_bytes > 0 + + scatter_only = layout.ArConfig( + world_size=8, m=4096, n=7168, recv_slots=8, scatter_dtype="fp8" + ) + scatter_only.validate() + assert scatter_only.scatter_slice_bytes * 2 == scatter_only.slice_bytes + assert scatter_only.gather_slice_bytes == scatter_only.slice_bytes + assert scatter_only.gather_scale_bytes == 0 + assert scatter_only.gout_off == scatter_only.output_off # nothing to stage + + +def test_scatter_scale_granularity_is_finer_than_per_row(): + """The substitution that makes the scatter leg implementable at all. + + A GEMM block owns BLOCK_M x BLOCK_N of C, so with N=7168 and BLOCK_N=256 a + row spans 28 blocks and a per-row amax is not available in the epilogue. + Per-(row, N-tile) is what a block can compute, and it is 28 scales per row + rather than 1 -- finer, so it cannot be less accurate than what was asked + for. + """ + c = layout.ar_config( + world_size=8, + m=16384, + n=7168, + recv_slots=8, + counter_chunks=8, + comm_dtype="fp8", + ) + assert c.scatter_tiles_per_row == 28 + assert c.scatter_scale_bytes >= c.gather_scale_bytes * 28 + # and it is cheap: scales are a few percent of the payload they describe + assert c.scatter_scale_bytes < c.scatter_nbytes * 0.02 + + +# -------------------------------------------------------------------------- +# the anti-rot guard: the pinned copy still computes what the pipeline does +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("m,n,k", [(512, 512, 256), (1024, 768, 512)]) +def test_ptpc_matches_an_fp32_reference(m, n, k): + """The per-token/per-channel path, against fp32 done in torch. + + ``ptpc`` is the vendored 8-wave pipeline's native quantisation and the form + every earlier measurement is on, so it needs a guard of its own even though + the model runs ``blockscale``. Single GPU, no cco: this is what catches the + copied pipeline in ``kernels_fused.py`` drifting from + ``_gemm_a8w8_8wave.py``. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + pytest.importorskip("flydsl") + import flydsl.expr as fx + + from mori.ops.gemm_ar import compile_fused_gemm_scatter, preshuffle_b + + g = torch.Generator(device="cuda").manual_seed(7) + a = (torch.randn(m, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + b = (torch.randn(n, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + sa = torch.rand(m, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + sb = torch.rand(n, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + b_shuf = preshuffle_b(b) + + ref = (a.float() @ b.float().T) * sa[:, None] * sb[None, :] + + # world_size=1 is not a legal ArConfig, and it does not need to be: with + # fuse=False nothing in the epilogue reads the window, so any valid cfg of + # the right (m, n) picks the same tile shape. + cfg = layout.ArConfig(world_size=2, m=m, n=n) + gemm = compile_fused_gemm_scatter( + cfg, 0, K=k, BLOCK_M=256, BLOCK_N=256, b_preshuffled=True, fuse=False + ) + got = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) + gemm( + a.contiguous().view(torch.int8).view(-1), + b_shuf.contiguous().view(torch.int8).view(-1), + got.view(-1), + sa, + sb, + m, + n, + 0, # dev_comm: unused when fuse=False + 0, # win: ditto + stream=fx.Stream(torch.cuda.current_stream()), + ) + torch.cuda.synchronize() + + rel = ((got.float() - ref).norm() / ref.norm()).item() + assert rel < 3e-3, ( + f"ptpc relL2 {rel:.3e} against the fp32 reference; the pinned copy in " + f"kernels_fused.py may have drifted from _gemm_a8w8_8wave.py" + ) + + +@pytest.mark.parametrize("m,n,k", [(1024, 512, 512), (2048, 7168, 2048)]) +def test_blockscale_lands_at_the_fp8_floor(m, n, k): + """``quant="blockscale"`` -- the model's quantisation -- against fp32. + + The model's ROCm path quantises A 1x128 and B 128x128 with **fp32** scales, + so this is the operating point the fused-vs-split numbers are drawn at. Not + a bitwise check against anything: the reference accumulates in a different + order, so the gate is the fp8 floor. + + The scale layouts are the ones the model's own blockscale GEMM consumes -- + ``sa`` logically ``[M, K/128]`` but **column-major** (sglang's + ``materialize_bpreshuffle_fp8_scale`` is ``t().contiguous().t()``), ``sb`` + ``[N/128, K/128]`` row-major. Sweeping the four combinations, only this one + lands at the floor; the others give 0.17-0.32. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + pytest.importorskip("flydsl") + import flydsl.expr as fx + + from mori.ops.gemm_ar import compile_fused_gemm_scatter, preshuffle_b + + BK = 128 + g = torch.Generator(device="cuda").manual_seed(11) + a = (torch.randn(m, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + b = (torch.randn(n, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + sa = ( + torch.rand(m, k // BK, generator=g, device="cuda", dtype=torch.float32) * 0.01 + + 0.01 + ) + sb = ( + torch.rand(n // BK, k // BK, generator=g, device="cuda", dtype=torch.float32) + * 0.01 + + 0.01 + ) + sa = sa.t().contiguous().t() + b_shuf = preshuffle_b(b) + + af, bf = a.float(), b.float() + ref = torch.zeros(m, n, device="cuda", dtype=torch.float32) + for i in range(k // BK): + ks = slice(i * BK, (i + 1) * BK) + ref += ( + (af[:, ks] @ bf[:, ks].T) + * sa[:, i][:, None] + * sb[:, i].repeat_interleave(BK)[None, :] + ) + rn = ref.norm() + + cfg = layout.ArConfig(world_size=2, m=m, n=n) + gemm = compile_fused_gemm_scatter( + cfg, + 0, + K=k, + BLOCK_M=128, + BLOCK_N=256, + b_preshuffled=True, + fuse=False, + swap_ab=True, + permlane=True, + lane_transpose=True, + quant="blockscale", + ) + ours = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) + gemm( + a.contiguous().view(torch.int8).view(-1), + b_shuf.contiguous().view(torch.int8).view(-1), + ours.view(-1), + sa.t().reshape(-1).contiguous(), + sb.reshape(-1).contiguous(), + m, + n, + 0, + 0, + stream=fx.Stream(torch.cuda.current_stream()), + ) + torch.cuda.synchronize() + + rel = ((ours.float() - ref).norm() / rn).item() + assert rel < 3e-3, f"blockscale relL2 {rel:.3e} against the fp32 reference" + + +@pytest.mark.parametrize("m,n,k", [(512, 512, 256), (4096, 7168, 1024)]) +def test_swap_ab_is_bitwise_identical(m, n, k): + """Exchanging the MFMA operands must not change a single bit. + + ``fx.gemm(atom, c, b, a, c)`` computes ``(A B)^T`` in the accumulator's own + layout: same products, same fp32 accumulation order along k, only the + register-to-(row, col) mapping differs. Adding ``permlane`` only moves whole + 16-lane rows between registers. So both are exactly reproducible, and + anything less than bit-equality means an index or a lane mapping drifted, + not that the arithmetic changed. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + pytest.importorskip("flydsl") + import flydsl.expr as fx + + from mori.ops.gemm_ar import compile_fused_gemm_scatter, preshuffle_b + + g = torch.Generator(device="cuda").manual_seed(7) + a = (torch.randn(m, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + b = (torch.randn(n, k, generator=g, device="cuda") / 8).to(torch.float8_e4m3fn) + sa = torch.rand(m, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + sb = torch.rand(n, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.01 + b_shuf = preshuffle_b(b) + + cfg = layout.ArConfig(world_size=2, m=m, n=n) + + # The reference is the same kernel with the operands unswapped, which is the + # property under test: swapping must be a pure register relabelling. + ref_gemm = compile_fused_gemm_scatter( + cfg, + 0, + K=k, + BLOCK_M=256, + BLOCK_N=256, + b_preshuffled=True, + fuse=False, + swap_ab=False, + ) + ref = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) + ref_gemm( + a.contiguous().view(torch.int8).view(-1), + b_shuf.contiguous().view(torch.int8).view(-1), + ref.view(-1), + sa, + sb, + m, + n, + 0, + 0, + stream=fx.Stream(torch.cuda.current_stream()), + ) + torch.cuda.synchronize() + + for extra in ( + {}, + {"permlane": True}, + {"permlane": True, "lane_transpose": True}, + {"permlane": True, "lane_transpose": True, "hoist_scales": True}, + ): + gemm = compile_fused_gemm_scatter( + cfg, + 0, + K=k, + BLOCK_M=256, + BLOCK_N=256, + b_preshuffled=True, + fuse=False, + swap_ab=True, + **extra, + ) + got = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) + gemm( + a.contiguous().view(torch.int8).view(-1), + b_shuf.contiguous().view(torch.int8).view(-1), + got.view(-1), + sa, + sb, + m, + n, + 0, + 0, + stream=fx.Stream(torch.cuda.current_stream()), + ) + torch.cuda.synchronize() + assert torch.equal( + got, ref + ), f"swap_ab {extra} is not bit-identical to the unswapped path" + + +def test_rotated_tile_order_is_a_permutation_of_the_linear_one(): + """Reordering tiles must not change which tiles exist -- checked on the host. + + Reproduces the kernel's index math in Python so a mapping bug is caught + without an 8-GPU run. + """ + ws, m, n, block_m, block_n = 8, 4096, 7168, 256, 256 + n_blocks = n // block_n + m_tiles_per_peer = (m // ws) // block_m + total = (m // block_m) * n_blocks + for rank in range(ws): + seen = set() + for idx in range(total): + rest, bn = divmod(idx, n_blocks) + tile_i, dest_seq = divmod(rest, ws) + dest_i = (dest_seq + rank) % ws + seen.add((dest_i * m_tiles_per_peer + tile_i, bn)) + assert len(seen) == total + assert seen == { + (bm, bn) for bm in range(m // block_m) for bn in range(n_blocks) + } + + +# -------------------------------------------------------------------------- +# GPU smoke: every mode reaches the same answer +# -------------------------------------------------------------------------- + + +def _run_bench(world_size, mode, m, n, k, extra=()): + env = os.environ.copy() + env.setdefault("MORI_SOCKET_IFNAME", "lo") + env.setdefault("MORI_ENABLE_SDMA", "1") + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc_per_node={world_size}", + str(BENCH), + "--mode", + mode, + "-m", + str(m), + "-n", + str(n), + "-k", + str(k), + "--warmup", + "1", + "--iters", + "3", + *extra, + ] + result = subprocess.run( + command, cwd=REPO_ROOT, env=env, capture_output=True, text=True, timeout=1200 + ) + output = result.stdout + result.stderr + assert result.returncode == 0, output + records = [ + json.loads(line.removeprefix("RESULT_JSON ")) + for line in output.splitlines() + if line.startswith("RESULT_JSON ") + ] + assert len(records) == 1, output + return records[0] + + +def test_fused_is_stable_across_repeats(): + """A single passing run proves nothing about the fused path. + + Its epilogue publishes tiles from 448 blocks and elects one of them to issue + the transfer, so every ordering bug in it is intermittent -- and each + individual relL2 looks "small" on its own. Two real races were shipped here + behind exactly that: a shared SDMA queue when a destination has more than one + chunk, and a counter atomic that had lost its acquire half. Both passed a + one-shot check repeatedly before failing. + + Needs the full world size: the race needs more than one chunk per + destination, and at 2 ranks the wo_b shape only has one. + """ + if torch.cuda.device_count() < 8: + pytest.skip("the fused race only reproduces at world_size=8") + pytest.importorskip("flydsl") + seen = [] + for _ in range(3): + record = _run_bench(8, "fused-sdma", 4096, 7168, 1024) + seen.append(record["rel_l2"]) + assert all(v == pytest.approx(seen[0], rel=1e-6) for v in seen), ( + f"fused all-reduce is not deterministic across runs: relL2 = {seen}. " + "That is an ordering bug in the epilogue, not numerical noise." + ) + assert all(v < 5e-3 for v in seen), seen + + +@pytest.mark.parametrize("mode", ["split-sdma", "fused-sdma", "split-lsa"]) +def test_gemm_ar_modes_agree(mode): + """All three paths must land on the same all-reduced result. + + The tolerance is fp8 quantization error, not collective error -- the + collective itself is asserted bit-exact in ``test_flydsl_ar.py``. The point + here is that the three modes agree with *each other*, which is what caught + the missing release fence in the fused epilogue (it showed up as 3.8e-3 + against split's 2.35e-3, both individually "small"). + """ + if torch.cuda.device_count() < 2: + pytest.skip("requires 2 GPUs") + pytest.importorskip("flydsl") + record = _run_bench(2, mode, 512, 1024, 256) + assert record["validated"] is True + assert record["max_rank_time_us"] > 0 + assert record["rel_l2"] < 5e-3 diff --git a/tests/python/cco/test_gemm_ar_bench.py b/tests/python/cco/test_gemm_ar_bench.py new file mode 100644 index 000000000..b0407468a --- /dev/null +++ b/tests/python/cco/test_gemm_ar_bench.py @@ -0,0 +1,146 @@ +# 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. +"""The fused GEMM+AR benchmark, run as a test: accuracy gates, timings recorded. + +What is asserted and what is not, deliberately: + +* **Accuracy is a gate.** ``validated`` and ``rel_l2`` come out of + ``bench_gemm_ar.py``'s own host-side check and are hard assertions. A kernel + that is fast and wrong fails here. +* **Absolute timings are only recorded.** CI runs on a shared runner, so a + microsecond threshold would be flaky in both directions -- it would fail on a + busy box and pass a real regression on an idle one. The numbers are printed + (run with ``-s``) so a bisect has something to read. +* **One relative ordering is asserted**: ``fused-sdma`` must beat + ``split-sdma``. That is the entire claim of the op, it is a 20%+ margin at the + model's shape, and both sides are measured in the same process on the same + box within seconds of each other -- so load affects them together. +""" + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +BENCH = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "gemm_ar" / "bench_gemm_ar.py" + +# Small enough that FlyDSL compiles it quickly, large enough that the fused path +# has more than one chunk per destination to publish (M / (ws * 128) = 4 bands). +SMOKE = dict(world_size=8, m=4096, n=1024, k=512) + +# The DSV4-Pro wo_b prefill chunk the op was built for: a --chunked-prefill-size +# 16384 TP8 chunk, N=7168 K=2048. +MODEL_SHAPE = dict(world_size=8, m=16384, n=7168, k=2048) + + +def _run(world_size, mode, m, n, k, quant="blockscale", iters=11): + env = os.environ.copy() + env.setdefault("MORI_SOCKET_IFNAME", "lo") + env.setdefault("MORI_ENABLE_SDMA", "1") + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc_per_node={world_size}", + str(BENCH), + "--mode", + mode, + "--quant", + quant, + "-m", + str(m), + "-n", + str(n), + "-k", + str(k), + "--warmup", + "3", + "--iters", + str(iters), + ] + result = subprocess.run( + command, cwd=REPO_ROOT, env=env, capture_output=True, text=True, timeout=1800 + ) + output = result.stdout + result.stderr + assert result.returncode == 0, output + records = [ + json.loads(line.removeprefix("RESULT_JSON ")) + for line in output.splitlines() + if line.startswith("RESULT_JSON ") + ] + assert len(records) == 1, output + return records[0] + + +def _require(world_size): + if torch.cuda.device_count() < world_size: + pytest.skip(f"requires {world_size} GPUs") + + +@pytest.mark.parametrize("mode", ["gemm-only", "split-sdma", "fused-sdma"]) +def test_mode_is_accurate_and_reports_a_time(mode): + """Every mode computes the right thing, and says how long it took.""" + _require(SMOKE["world_size"]) + r = _run(mode=mode, **SMOKE) + print( + f"\n[gemm_ar] {mode:<11s} m={r['m']} n={r['n']} k={r['k']} " + f"quant={r['quant']} {r['max_rank_time_us']:9.1f} us " + f"relL2={r['rel_l2']:.2e}" + ) + assert r["validated"] is True, f"{mode} failed the benchmark's own check" + assert r["rel_l2"] < 3e-3, f"{mode} relL2 {r['rel_l2']:.3e} is past the fp8 floor" + assert r["max_rank_time_us"] > 0 + + +def test_fusing_beats_splitting_at_the_model_shape(): + """The op's whole claim, measured back to back on one box. + + Recorded on an idle 8x MI355X at this shape: fused-sdma 1146.2us against + split-sdma 1465.5us (-21.8%), with the GEMM alone at 388.5us. The assertion + is only the ordering -- see this module's docstring for why. + """ + _require(MODEL_SHAPE["world_size"]) + fused = _run(mode="fused-sdma", **MODEL_SHAPE) + split = _run(mode="split-sdma", **MODEL_SHAPE) + + f, s = fused["max_rank_time_us"], split["max_rank_time_us"] + print( + f"\n[gemm_ar] model shape m={MODEL_SHAPE['m']} n={MODEL_SHAPE['n']} " + f"k={MODEL_SHAPE['k']}\n" + f" fused-sdma {f:9.1f} us relL2={fused['rel_l2']:.2e}\n" + f" split-sdma {s:9.1f} us relL2={split['rel_l2']:.2e}\n" + f" fused is {(1 - f / s) * 100:+.1f}% against split" + ) + for name, r in (("fused-sdma", fused), ("split-sdma", split)): + assert r["validated"] is True, f"{name} failed the benchmark's own check" + assert r["rel_l2"] < 3e-3, f"{name} relL2 {r['rel_l2']:.3e}" + assert f < s, ( + f"fused-sdma ({f:.1f} us) did not beat split-sdma ({s:.1f} us). " + f"The recorded margin is 21.8%, so this is a real regression unless the " + f"runner was heavily contended." + ) diff --git a/tests/python/cco/test_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py new file mode 100644 index 000000000..08d7e5599 --- /dev/null +++ b/tests/python/cco/test_gemm_ar_op.py @@ -0,0 +1,514 @@ +# 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. +"""``GemmAllReduceOp``, the public wrapper. + +``test_gemm_ar.py`` covers the kernels through ``bench_gemm_ar.py``, which calls +the low-level builders directly -- so it says nothing about this layer, and three +correctness bugs lived here behind a green suite: the A scale flattened in the +wrong order, a control region that moved when one instance served a second M, and +a support predicate that admitted K=128. +""" + +from __future__ import annotations + +import json +import os +import pathlib +import subprocess +import sys +import time + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("flydsl") + +from mori.ops.gemm_ar import ArConfig # noqa: E402 +from mori.ops.gemm_ar.op import ( # noqa: E402 + FP8_DTYPES, + MAX_CHUNKS, + MIN_K, + _flatten_a_scale, + counter_chunks, + padded_m, + supports, +) + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +WORKER = pathlib.Path(__file__).with_name("gemm_ar_op_worker.py") + +# The fp8 floor for these operands; the kernel-level tests gate on the same value. +FP8_FLOOR = 3e-3 + + +# --- host-only: the predicate and the layout ------------------------------ + + +@pytest.mark.parametrize("k", [128, MIN_K - 128, 64]) +def test_supports_rejects_k_below_two_blocks(k): + """The mainloop prefetches block 1 and runs two tail steps. + + K=128 passed the divisibility check, constructed, ran, and returned finite + output at relL2 1.41 against 2.4e-3 for the otherwise identical K=256. + """ + assert supports(512, 1024, k, 2) is False + + +def test_supports_accepts_the_smallest_real_k(): + assert supports(512, 1024, MIN_K, 2) is True + + +@pytest.mark.parametrize("world_size", [3, 5, 6, 7]) +def test_supports_agrees_with_config_validation_on_world_size(world_size): + """``supports`` said yes where construction raised. + + 512 threads must divide by the world size, which ``ArConfig.validate`` + enforces and the predicate used to ignore. + """ + assert supports(512, 1024, 512, world_size) is False + with pytest.raises(ValueError): + ArConfig(world_size=world_size, m=512, n=1024, recv_slots=world_size).validate() + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_supports_accepts_world_sizes_that_divide_the_block(world_size): + assert supports(512, 1024, 512, world_size) is True + + +def test_supports_requires_the_fp8_row_multiple(): + """The fp8 conversion gives one wave a row, so N must be a whole wave-chunk.""" + assert supports(512, 7168, 512, 8, gather_dtype="fp8") is True + assert supports(512, 1536, 512, 8, gather_dtype="fp8") is False + assert supports(512, 1536, 512, 8, gather_dtype="bf16") is True + + +def test_control_region_does_not_move_between_m_values(): + """The bug: counter_chunks changes with M and everything after it shifts. + + At capacity, the payload starts at the same offset whatever M is, which is + what makes one window safe for several shapes. + """ + offsets = set() + for m in (4096, 8192, 16384): + cfg = ArConfig( + world_size=8, + m=m, + n=7168, + recv_slots=8, + counter_chunks=counter_chunks(m, 8), + counter_capacity=MAX_CHUNKS, + counter_shape_slots=4, + ) + cfg.validate() + # Only the control region has to be stable. output_off sits after the + # payload and moves with M by construction. + offsets.add((cfg.counter_region_off, cfg.lock_off, cfg.input_off)) + assert len(offsets) == 1, offsets + + +def test_each_shape_slot_gets_its_own_counters(): + """Counters are never reset and the election modulus depends on M. + + Two shapes sharing a set elect on the residue the other left, so the sets + have to be disjoint -- while everything after them stays put. + """ + cfgs = [ + ArConfig( + world_size=8, + m=4096, + n=7168, + recv_slots=8, + counter_chunks=4, + counter_capacity=MAX_CHUNKS, + counter_shape_slots=4, + counter_shape_index=i, + ) + for i in range(4) + ] + starts = [c.counter_off for c in cfgs] + assert len(set(starts)) == 4 + span = 8 * MAX_CHUNKS * 4 + for a, b in zip(starts, starts[1:]): + assert b - a == span + assert starts[-1] + span <= cfgs[0].lock_off + assert len({c.input_off for c in cfgs}) == 1 + + +@pytest.mark.parametrize( + "field,value", + [ + ("gather_dtype", "FP8"), + ("gather_dtype", "fp16"), + ("gather_transport", "sla"), + ("scatter_dtype", "int8"), + ], +) +def test_validate_rejects_misspelled_modes(field, value): + """A typo used to select the default silently: "FP8" sent bf16.""" + cfg = ArConfig(world_size=8, m=4096, n=7168, recv_slots=8, **{field: value}) + with pytest.raises(ValueError, match=field): + cfg.validate() + + +# --- host-only: the scale flattening -------------------------------------- + + +def test_flatten_a_scale_preserves_physical_order(): + """The kernel reads (row, kb) at kb*M + row, for the unambiguous spellings.""" + m, kb = 8, 3 + logical = torch.arange(m * kb, dtype=torch.float32).reshape(m, kb) + want = logical.t().reshape(-1) + + col_major = logical.t().contiguous().t() # [M, kb], storage already kb-major + assert col_major.stride() == (1, m) + + assert torch.equal(_flatten_a_scale(col_major, m, kb), want) + assert torch.equal(_flatten_a_scale(logical.t().contiguous(), m, kb), want) + assert torch.equal(_flatten_a_scale(want, m, kb), want) + + +def test_flatten_a_scale_is_not_a_reshape(): + """Guards the review's finding: reshape(-1) on the column-major tensor.""" + m, kb = 8, 3 + col_major = ( + torch.arange(m * kb, dtype=torch.float32).reshape(m, kb).t().contiguous().t() + ) + assert not torch.equal(_flatten_a_scale(col_major, m, kb), col_major.reshape(-1)) + + +def test_flatten_a_scale_refuses_the_ambiguous_row_major_case(): + """[M, K/128] row-major means two different things and must not be guessed. + + A logical [M, K/128] needs transposing; a K/128-major buffer carrying that + shape -- which is what aiter_per1x128_quant(transpose_scale=True) returns -- + must be read flat. Shape and stride are identical in both cases, and + guessing "transpose" took SGLang's perplexity from 3.26 to 862511 while + every shape-level test here stayed green, because they build the logical + tensor the guess assumes. + """ + m, kb = 8, 3 + row_major = torch.arange(m * kb, dtype=torch.float32).reshape(m, kb) + assert row_major.stride() == (kb, 1) + with pytest.raises(ValueError, match="ambiguous"): + _flatten_a_scale(row_major, m, kb) + # Both meanings stay expressible, and they differ. + flat = _flatten_a_scale(row_major.reshape(-1), m, kb) + transposed = _flatten_a_scale(row_major.t(), m, kb) + assert torch.equal(flat, row_major.reshape(-1)) + assert torch.equal(transposed, row_major.t().reshape(-1)) + assert not torch.equal(flat, transposed) + + +@pytest.mark.parametrize("shape", [(7, 3), (8, 4), (8,)]) +def test_flatten_a_scale_rejects_shapes_that_are_neither(shape): + with pytest.raises(ValueError): + _flatten_a_scale(torch.zeros(shape), 8, 3) + + +def test_padded_m_rounds_to_whole_bands(): + assert padded_m(1, 8, 128) == 1024 + assert padded_m(1024, 8, 128) == 1024 + assert padded_m(1025, 8, 128) == 2048 + + +# --- host-only: the tile and dtype boundaries ----------------------------- + + +@pytest.mark.parametrize("block_m", [64, 1, 192, 0, -128]) +def test_supports_rejects_tiles_the_kernel_cannot_build(block_m): + """The kernel's assert is bare, so the boundary has to answer instead. + + block_m=64 used to return True here, construct, acquire all three + resources, and only then raise an empty AssertionError from inside the + mainloop on the first real call. block_m=0 used to raise ZeroDivisionError + out of the padding arithmetic before reaching any check at all. + """ + assert supports(4096, 7168, 2048, 8, block_m=block_m) is False + + +@pytest.mark.parametrize("block_n", [128, 384, 0]) +def test_supports_rejects_n_tiles_the_kernel_cannot_build(block_n): + assert supports(4096, 7168, 2048, 8, block_n=block_n) is False + + +def test_supports_accepts_the_legal_tile_multiples(): + """256x512 is legal even though it is not the default 128x256.""" + assert supports(4096, 7168, 2048, 8, block_m=256, block_n=512) is True + + +@pytest.mark.parametrize("world_size", [0, -1]) +def test_supports_rejects_nonsense_world_size(world_size): + """padded_m divides by world_size * block_m, so this has to be caught first.""" + assert supports(4096, 7168, 2048, world_size) is False + + +def test_fnuz_is_not_an_accepted_dtype(): + """Same width, different exponent bias -- see FP8_DTYPES. + + The MMA atom is fixed to OCP e4m3. Reading FNUZ bytes through it scales + every A and every B element by 2, so the product comes back 4x too large: + finite, plausible and wrong. Matching element size is not evidence. + """ + assert torch.float8_e4m3fn in FP8_DTYPES + assert torch.float8_e4m3fnuz not in FP8_DTYPES + + +def test_constructor_rolls_back_cleanly_when_a_resource_fails(monkeypatch): + """Every acquisition step, failed in turn. + + The rollback calls close(), which clears _cache and _pad_in. Those used to + be assigned after the try block, so any failure here died with + AttributeError -- masking the real error and skipping the release loop that + is the whole point of the rollback. + + Recording fakes, not a real communicator: this is about the ordering of the + constructor's own state, which needs no GPU. The fake handles carry a null + pointer, so the window-zeroing write is stubbed out -- left in, it faults + the context and the fault surfaces in whatever runs next. + """ + from mori.ops.gemm_ar import op as op_module + from mori.ops.gemm_ar.op import GemmAllReduceOp + + monkeypatch.setattr(op_module, "from_gpu_ptr", lambda *a, **kw: torch.zeros(1)) + + class _Handle: + def __init__(self, log, name): + self._log, self._name = log, name + self.ptr, self.size = 0, 0 + + def close(self): + self._log.append(f"close:{self._name}") + + class _Comm: + """Fails at `fail_at`; records what was acquired and what was closed.""" + + nranks, rank = 8, 0 + + def __init__(self, fail_at): + self.fail_at, self.log = fail_at, [] + + def _step(self, name): + self.log.append(f"acquire:{name}") + if name == self.fail_at: + raise RuntimeError(f"injected failure in {name}") + return _Handle(self.log, name) + + def alloc_mem(self, nbytes): + return self._step("alloc_mem") + + def register_window(self, ptr, size): + return self._step("register_window") + + def create_dev_comm(self, reqs): + return self._step("create_dev_comm") + + for fail_at in ("alloc_mem", "register_window", "create_dev_comm"): + comm = _Comm(fail_at) + with pytest.raises(RuntimeError) as excinfo: + GemmAllReduceOp(comm, n=7168, k=2048, m_max=4096) + + # The injected error propagates, rather than an AttributeError from + # cleanup with the real one demoted to __context__. + assert fail_at in str(excinfo.value) + + # Everything acquired before the failure was released. The dev-comm is + # not in that set by design -- its queues are communicator-scoped and go + # with the communicator, so it is never the last-but-one here. + acquired = [x[len("acquire:") :] for x in comm.log if x.startswith("acquire:")] + closed = {x[len("close:") :] for x in comm.log if x.startswith("close:")} + assert acquired[-1] == fail_at, f"stopped at {acquired[-1]}, not {fail_at}" + assert ( + set(acquired[:-1]) == closed + ), f"failing at {fail_at}: acquired {acquired}, closed {closed}" + + +# --- multi-rank: the op itself -------------------------------------------- + + +#: Back-to-back 8-rank launches outrun SDMA queue teardown -- the next process +#: then fails in hsaKmtCreateQueueExt (anvil.cpp:237). Each of these cases passes +#: alone and they failed only when spawned in sequence, so let the previous one +#: release before starting the next. +_SETTLE_SECONDS = 20 + + +#: hsaKmtCreateQueueExt failing because the previous run's SDMA queues are not +#: reclaimed yet. It is a resource race, not a defect in what is under test, and +#: it is the only failure worth retrying -- everything else is reported as is. +_QUEUE_EXHAUSTED = "anvil.cpp" + + +def _run_worker_multi(world_size: int, case: str, *extra: str, timeout: int = 900): + result = _spawn_worker(world_size, case, *extra, timeout=timeout) + if result is None: + # One retry, with a longer settle. Two files' worth of 8-rank runs come + # before this one, and the queues do not always come back in 20s. + time.sleep(_SETTLE_SECONDS * 3) + result = _spawn_worker(world_size, case, *extra, timeout=timeout, last=True) + return result + + +def _spawn_worker( + world_size: int, case: str, *extra: str, timeout: int = 900, last: bool = False +): + time.sleep(_SETTLE_SECONDS) + env = os.environ.copy() + env.setdefault("MORI_SOCKET_IFNAME", "lo") + env.setdefault("MORI_ENABLE_SDMA", "1") + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc_per_node={world_size}", + str(WORKER), + "--case", + case, + *extra, + ], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + output = result.stdout + result.stderr + if result.returncode != 0: + if _QUEUE_EXHAUSTED in output and not last: + return None + raise AssertionError(output) + records = [ + json.loads(line.removeprefix("RESULT_JSON ")) + for line in output.splitlines() + if line.startswith("RESULT_JSON ") + ] + assert records, output + return records + + +def _run_worker(world_size: int, case: str, *extra: str, timeout: int = 900): + """One case, one record.""" + return _run_worker_multi(world_size, case, *extra, timeout=timeout)[0] + + +requires_two_gpus = pytest.mark.skipif( + torch.cuda.device_count() < 2, reason="needs 2 GPUs" +) + + +@pytest.fixture(scope="module") +def worker_results(): + """Every GPU case, from one 2-rank spawn. + + Each case used to get its own launch. Six of them back to back outran the + SDMA queue reclaim (hsaKmtCreateQueueExt, anvil.cpp:237) and the casualties + were as often the neighbouring file's tests as this file's, because the cost + is paid by whoever launches next. One spawn, results keyed by case. + """ + records = _run_worker_multi(2, "all", "-n", "1024", "-k", "512") + return records + + +def _case(records, name): + for r in records: + if r.get("case") == name: + return r + raise AssertionError(f"worker produced no record for {name!r}: {records}") + + +@requires_two_gpus +def test_public_op_takes_the_column_major_scale(worker_results): + """End to end through ``GemmAllReduceOp``, with the scale the model emits. + + Before the fix this returned relL2 0.26: ``reshape(-1)`` on a column-major + ``[M, K/128]`` walks the logical rows, so every scale met the wrong K block. + """ + r = _case(worker_results, "scale_order") + assert r["rel_l2"] < FP8_FLOOR, r + assert r["flat_matches"], r + assert r["physical_matches"], r + + +@requires_two_gpus +@pytest.mark.parametrize("m_small,m_large", [(512, 1024), (2048, 4096)]) +def test_one_instance_serves_two_m_values_in_any_order( + worker_results, m_small, m_large +): + """Two M values through one instance, alternating, then repeated. + + Two separate defects lived here. The reviewer's: counter_chunks changes with + M, which moved lock_off and everything after it, so the second shape read the + first one's payload as control state and hung. Then, once that was fixed, + this: the payload offsets are a running sum of m-sized regions, so at + world=2 n=1024 m=1024's output rows 512..767 land exactly where m=512's recv + slot for peer 1 was. The first call at the second shape came back with that + peer's previous *scatter* payload in those rows -- 4 runs in 8, always the + same value, with the repeat correct because by then the bytes were its own. + + Both are fixed by pinning the map: the control region to MAX_CHUNKS, the + payload to the instance's m_max. Every shape now has the same map, so a + region only ever aliases itself. + """ + r = _case(worker_results, f"alternating_m_{m_small}_{m_large}") + for key in ("first_small", "first_large", "again_small", "again_large"): + assert r[key] < FP8_FLOOR, (key, r) + + +@requires_two_gpus +def test_changing_operands_between_calls_stays_correct(worker_results): + """Five calls at one M, different data each time. + + Every other test and every benchmark feeds identical operands on every + iteration, which makes a read of a peer's not-yet-landed slice return the + previous call's bytes -- bit-identical to the right answer. This is the only + thing in the suite that would notice. + """ + r = _case(worker_results, "changing_data") + for key in (f"call{i}" for i in range(5)): + assert r[key] < FP8_FLOOR, (key, r) + + +@requires_two_gpus +def test_self_test_passes_and_can_fail(worker_results): + """The guard against a mori whose SDMA puts were compiled out. + + That build is the default (`BUILD_CCO_SDMA` is OFF unless BUILD_BENCHMARK is + ON) and it fails silently: every symbol is there, every kernel launches, + every put returns, and nothing moves. The all-reduce then yields mostly the + local slice, the model still answers fluently, and the fused path measures + *faster* than it is. An end-to-end campaign read -6.8% instead of -2.3% that + way and only perplexity caught it. + + So this asserts both halves: the check passes on a working stack, and it + fails when the scatter does not happen -- which is what an inert put is. + """ + r = _case(worker_results, "self_test") + assert r["passes_on_working_stack"], r + assert r["detects_missing_scatter"], r + + +@requires_two_gpus +def test_close_releases_and_is_idempotent(worker_results): + """The communicator holds the handles, so dropping the op is not enough.""" + r = _case(worker_results, "close") + assert r["refused_after_close"], r diff --git a/tests/python/cco/test_triton_bindings.py b/tests/python/cco/test_triton_bindings.py index 1911dfc06..84c02eaca 100644 --- a/tests/python/cco/test_triton_bindings.py +++ b/tests/python/cco/test_triton_bindings.py @@ -55,7 +55,7 @@ def invalid_constexpr_axis_kernel( def test_all_cco_device_symbols_are_exported_to_triton(): - assert len(CCO_DEVICE_FUNCTIONS) == 68 + assert len(CCO_DEVICE_FUNCTIONS) == 70 assert set(CCO_DEVICE_FUNCTIONS).issubset(cco.__all__) for name in CCO_DEVICE_FUNCTIONS: assert callable(getattr(cco, name)) @@ -180,6 +180,8 @@ def test_flydsl_and_triton_share_the_same_scalar_abi(): *_bindings.SDMA_COMMIT.values(), _bindings.cco_sdma_quiet_queue, _bindings.cco_lsa_ptr, + _bindings.cco_lsa_win_base, + _bindings.cco_lsa_stride, _bindings.cco_system_fence, _bindings.cco_devcomm_rank, _bindings.cco_devcomm_world_size,