From f48a595d4444dbb350b245d118638f708bc07c3b Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 7 Sep 2026 12:25:07 +0000 Subject: [PATCH 01/79] feat(cco): FlyDSL LSA all-reduce benchmarked against aiter custom AR Adds a FlyDSL 1-stage/2-stage all-reduce over cco symmetric windows, plus the harness and tests to compare it with aiter's custom all-reduce. Target shape is DeepSeek-V4-Pro wo_b on 8x MI355X: [M, 7168] bf16, TP8. Structure follows benchmark/cco/triton/gemm_a2a: pure-Python layout math that is unit-testable without a GPU, a torchrun benchmark emitting one RESULT_JSON line with critical-rank decomposition, and a pytest that runs it on 2/4 ranks. Protocol parity with aiter's custom_all_reduce.cuh: per-(block, peer) signal slots with a monotonic _flag[block] that is never reset (graph-replay safe), peer rotation in the reduce-scatter, stage 2 reusing stage 1's thread->index mapping (cross-device visibility only holds between threads with the same tid), and matching accumulation order so results are bitwise comparable. Status: correctness done -- rel_l2 == 0.0 against a float32 reference on 2 and 8 ranks at M=64/512/4096, n=7168. Performance is not yet at parity (M=4096: 798us vs aiter 265.7us). The likely cause is the one deliberate divergence from aiter, noted in kernels_lsa.py: each thread reads all peers straight into registers instead of staging through LDS and reducing on warp 0. A/B of the two is next, along with the SDMA backend and folding the aiter baseline into this harness. Two bugs in examples/cco/python/05_flydsl_lsa_allreduce surfaced while porting it, both of which make its cross-GPU barrier a no-op. The example still passes because its inputs are host-written before the kernel, so nothing in it actually depends on the barrier; the 2-stage all-reduce here does, which is how they showed up. Both are worked around in _compat.py rather than in the example: 1. Its cache-modifier constants are gfx942-specific. On gfx950 the aux encoding is bit0=SC0, bit1=NT, bit2=SC1, so SC0|SC1=3 assembles to "sc0 nt" -- a non-temporal hint that does not bypass L2. Confirmed by reading the emitted ISA (aux=3 -> sc0 nt, aux=5 -> sc0 sc1). 2. rocdl.raw.ptr.buffer.load is side-effect-free, so the scf.While spin loop is dead-code-eliminated outright: the barrier compiles down to a signal store and nothing else. The signal path here is therefore atomic load/store with an explicit syncscope (system to publish, agent to poll, mirroring aiter's __scoped_atomic_*), which is neither removable nor dependent on hand-encoded cache bits. Also carried over from aiter: _wave_uniform_i64 on peer base addresses. Window .lsa_ptr is an opaque extern call, so without it the backend cannot prove the buffer descriptor is wave-uniform and wraps every 16-byte access in a readfirstlane waterfall loop. --- benchmark/cco/flydsl/ar/_compat.py | 280 +++++++++++++++++++++++ benchmark/cco/flydsl/ar/_probe_lsa.py | 185 ++++++++++++++++ benchmark/cco/flydsl/ar/bench_ar.py | 293 ++++++++++++++++++++++++ benchmark/cco/flydsl/ar/kernels_lsa.py | 296 +++++++++++++++++++++++++ benchmark/cco/flydsl/ar/layout.py | 251 +++++++++++++++++++++ tests/python/cco/test_flydsl_ar.py | 219 ++++++++++++++++++ 6 files changed, 1524 insertions(+) create mode 100644 benchmark/cco/flydsl/ar/_compat.py create mode 100644 benchmark/cco/flydsl/ar/_probe_lsa.py create mode 100644 benchmark/cco/flydsl/ar/bench_ar.py create mode 100644 benchmark/cco/flydsl/ar/kernels_lsa.py create mode 100644 benchmark/cco/flydsl/ar/layout.py create mode 100644 tests/python/cco/test_flydsl_ar.py diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py new file mode 100644 index 000000000..0dbd64fcb --- /dev/null +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -0,0 +1,280 @@ +# 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 flydsl +import flydsl.expr as fx +from flydsl._mlir import ir + +FLYDSL_VERSION = getattr(flydsl, "__version__", "unknown") + +# Cache-modifier (``aux``) bits. **These are CDNA3/CDNA4-dependent.** +# +# On gfx950 the encoding is bit0=SC0, bit1=NT, bit2=SC1, so the gfx942 spelling +# used by ``examples/cco/python/05_flydsl_lsa_allreduce`` (SC1=2, SC0|SC1=3) +# silently compiles to ``nt`` and ``sc0 nt`` here -- non-temporal hints that do +# **not** bypass L2. Verified by reading the emitted ISA: aux=3 produced +# ``buffer_store_dwordx4 ... sc0 nt``, aux=5 produces ``... sc0 sc1``. +# +# The examples get away with it because their inputs are host-written before the +# kernel, so nothing there actually depends on the barrier ordering. +CM_CACHED = 0 +CM_SC0 = 1 # bypass L1 +CM_NT = 2 # non-temporal +CM_SC1 = 4 # bypass L2 -- required to observe a peer's fresh write +CM_SC0_SC1 = CM_SC0 | CM_SC1 # 5: 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 + +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 import rocdl as _rocdl + 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 + + 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) + return _rocdl.raw_ptr_buffer_load( + res, + rsrc, + _byte_offset(offset, elem), + _const(0, 32), + _const(cache_modifier, 32), + ) + + 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 + return _rocdl.raw_ptr_buffer_store( + value, + rsrc, + _byte_offset(offset, elem), + _const(0, 32), + _const(cache_modifier, 32), + ) + + +# --- 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 +from flydsl.expr.typing import AddressSpace as _AS # noqa: E402 +from flydsl.expr.typing import PointerType as _PT # noqa: E402 +from flydsl.expr.typing import inttoptr as _inttoptr # 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) + + +def signal_ptr(addr_i64): + """Raw i64 device address -> ``!llvm.ptr`` over one u32 signal slot. + + ``inttoptr`` yields a ``!fly.ptr``; the llvm dialect ops need the lowered + ``!llvm.ptr``, hence the extra hop (same shape as aiter's + ``kernels/mxfp4_gemm_common.py::_global_base_ptr1``). + """ + return fx.to_llvm_ptr( + _inttoptr(_PT.get(_dtype("i32"), _AS.Global), fx.Int64(addr_i64)) + ) + + +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 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", + "i32_type", + "CM_CACHED", + "CM_SC1", + "CM_SC0_SC1", +] diff --git a/benchmark/cco/flydsl/ar/_probe_lsa.py b/benchmark/cco/flydsl/ar/_probe_lsa.py new file mode 100644 index 000000000..df171857d --- /dev/null +++ b/benchmark/cco/flydsl/ar/_probe_lsa.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +# 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 + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from _compat import ( # noqa: E402 + 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..b48141303 --- /dev/null +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -0,0 +1,293 @@ +#!/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 +""" + +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 + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from layout import ArConfig # noqa: E402 + +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) + 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 kernels_lsa import build_lsa_ar + + launch, stage = build_lsa_ar(cfg, rank, force_stage=args.force_stage) + elif args.backend == "sdma": + from kernels_sdma import build_sdma_ar + + launch, stage = build_sdma_ar(cfg, rank, queues=args.sdma_queues) + 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 + validated = rel_l2 == 0.0 + 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() != 0.0: + 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, + "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("-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("--sdma-queues", type=int, default=8) + 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/ar/kernels_lsa.py b/benchmark/cco/flydsl/ar/kernels_lsa.py new file mode 100644 index 000000000..c69bae49f --- /dev/null +++ b/benchmark/cco/flydsl/ar/kernels_lsa.py @@ -0,0 +1,296 @@ +# 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 stages every peer's pack through LDS and +reduces on ``warp_id == 0`` only, costing two ``__syncthreads()`` per iteration +at 1/8 occupancy. Here each thread reads all peers straight into registers. The +benchmark A/Bs the two. + +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 os +import sys + +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 + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from _compat import ( # noqa: E402 + CM_CACHED, + CM_SC0_SC1, + CM_SC1, + 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 # noqa: E402 + + +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]) + + +def build_lsa_ar(cfg, rank: int, *, force_stage: int | None = None): + """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. + """ + cfg.validate() + 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 + elems_per_pack = cfg.elems_per_pack + 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.Window(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.Window(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. + out = create_buffer_resource_from_addr(wave_uniform_i64(w.lsa_ptr(rank, out_off))) + for j in range_constexpr(ws): + src_rank = (rank + j) % ws + src = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(src_rank, tmp_off)) + ) + dst_base = src_rank * part + n_here = cfg.owner_pack_range(src_rank)[1] - dst_base + for pk in range(gtid, n_here, stride_packs): + v = buffer_load( + src, 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/benchmark/cco/flydsl/ar/layout.py b/benchmark/cco/flydsl/ar/layout.py new file mode 100644 index 000000000..9bb9d7001 --- /dev/null +++ b/benchmark/cco/flydsl/ar/layout.py @@ -0,0 +1,251 @@ +# 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``). +* ``blocks = min(80, ceil((size_packs / world) / (threads / world)))`` with + ``threads = 512`` (``custom_all_reduce.cuh:3790-3796``). +* 1-stage below the size thresholds, 2-stage above + (``custom_all_reduce.cuh:3771-3788``). + +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 + +# --- 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 + +# 1-stage thresholds, in bytes (custom_all_reduce.cuh:3779). +ONE_STAGE_MAX_BYTES_LE4 = 160 * 1024 +ONE_STAGE_MAX_BYTES_LE8 = 80 * 1024 + + +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 + + +@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 + + # --- 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) + + # --- 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: + per_block = self.threads_per_peer + if self.stage == 1: + need = (self.num_packs + per_block - 1) // per_block + else: + need = (self.packs_per_rank + per_block - 1) // per_block + return max(1, min(self.max_blocks, need)) + + # --- symmetric window layout: [ signal | input | output | tmp ] --- + # ``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. + + @property + def signal_bytes(self) -> int: + row = _align_up(self.max_blocks * MAX_WORLD * 4, SIGNAL_ALIGN) + flag = _align_up(self.max_blocks * 4, SIGNAL_ALIGN) + return row * 2 + flag # start[] + end[] + _flag[] + + @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 + + 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 input_off(self) -> int: + return _align_up(self.signal_bytes, SIGNAL_ALIGN) + + @property + def output_off(self) -> int: + return self.input_off + self.nbytes + + @property + def tmp_off(self) -> int: + return self.output_off + self.nbytes + + @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 window_bytes(self) -> int: + return self.tmp_off + self.tmp_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.elem_bytes not in (2, 4): + raise ValueError(f"elem_bytes must be 2 or 4, got {self.elem_bytes}") + + +__all__ = [ + "ArConfig", + "select_stage", + "K_MAX_BLOCKS", + "THREADS", + "MAX_WORLD", + "PACK_BYTES", +] diff --git a/tests/python/cco/test_flydsl_ar.py b/tests/python/cco/test_flydsl_ar.py new file mode 100644 index 000000000..636de5f45 --- /dev/null +++ b/tests/python/cco/test_flydsl_ar.py @@ -0,0 +1,219 @@ +# 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 importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +AR_DIR = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "ar" +sys.path.insert(0, str(AR_DIR)) + + +def _load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +layout = _load("flydsl_ar_layout", AR_DIR / "layout.py") + + +# -------------------------------------------------------------------------- +# 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_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.K_MAX_BLOCKS + assert _cfg(m=1).blocks >= 1 + assert _cfg(m=1).blocks <= layout.K_MAX_BLOCKS + + +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): + 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", + ] + 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 + + +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 From 30ddad9ac9f0afdc3738d990833ff0659e06e372 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 02:48:46 +0000 Subject: [PATCH 02/79] perf(cco): bring the FlyDSL LSA all-reduce to parity with aiter Two fixes, found by A/B rather than by inspection. At [4096, 7168] bf16 on 8x MI355X the kernel goes from 798us to 286us; aiter is 271us. 1. All-gather loop nesting (the 2.2x). Stage 2 had the peer loop *outside* the index loop, so it drained one peer completely before starting the next -- pinning the all-gather to a single xGMI link for 7 consecutive passes. aiter unrolls the peer loop inside the index loop (custom_all_reduce.cuh:566) so all 8 links stream at once. The arithmetic matches the measurement: one 7.34MB slice over one 55GB/s link is ~133us, and 133 + 7*95 is the 798us we were seeing. The old nesting is kept as `--gather sequential`. It is the more natural way to write the loop and the penalty is invisible in the output, so it is worth being able to demonstrate rather than just warn about. 2. Grid size (a further 1.25x). We were using aiter's 80 blocks, but the two kernels put a different amount of work on each thread: aiter's thread takes one pack from one peer (its warp_id picks the peer), ours takes one index across all 8. Matching aiter's grid therefore over-subscribes the xGMI request queues by 8x, and the time rises monotonically past ~32 blocks (286us @28 -> 361 @80 -> 583 @256). Cap is now 24, and the per-block stride is `threads` rather than `threads / world`. K_MAX_BLOCKS stays 80: it sizes the signal array, which is a correctness bound, not a tuning knob. The two are now separate. Best-of-3, 101 iterations, CUDA-graph replay, vs aiter custom AR: M aiter LSA ratio 64 25.6us 30.7us 1.20 128 27.2 41.8 1.54 512 53.6 66.5 1.24 1024 85.6 100.6 1.17 2048 146.1 160.3 1.10 4096 271.1 285.9 1.05 8192 558.8 531.7 0.95 Parity (+-15%) holds for M >= 1024. Below that both implementations sit on a payload-independent floor -- ours is ~28us and does not move between M=8 and M=64 -- so those ratios are two constants divided, not throughput. Best-of-3 is not defensive: single runs of *either* implementation vary by up to 1.7x (aiter's M=64 spanned 25.6-44.0us), and an earlier single-run sweep reported 204us at M=2048 where the stable figure is 160us. Two things checked and left alone: - aiter's 1-stage/2-stage threshold is right for us too. Forcing 1-stage is worse at every small size (M=512: 150us vs 68us), so the extra barrier costs less than the 4x traffic. - This box is affected by known-issue 3 (ROCm 7.2 puts uncached VMM allocations in the coarse-grained pool): CCO_UNCACHED_WINDOW=0 changes nothing, which is the documented tell. aiter uses IPC handles and does not pay it. Not quantified here, and it does not account for the whole small-M gap. The earlier note blaming the gap on our register-direct reduce vs aiter's LDS-staged one was wrong -- that divergence is real but costs nothing, and the register version is what makes the smaller grid necessary. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/bench_ar.py | 56 +++++++++++++++- benchmark/cco/flydsl/ar/kernels_lsa.py | 88 ++++++++++++++++++++------ benchmark/cco/flydsl/ar/layout.py | 45 ++++++++++--- tests/python/cco/test_flydsl_ar.py | 22 ++++++- 4 files changed, 180 insertions(+), 31 deletions(-) diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py index b48141303..bff35c5ce 100644 --- a/benchmark/cco/flydsl/ar/bench_ar.py +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -29,6 +29,29 @@ 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 ratio + 64 25.6us 30.7us 1.20 + 128 27.2 41.8 1.54 + 512 53.6 66.5 1.24 + 1024 85.6 100.6 1.17 + 2048 146.1 160.3 1.10 + 4096 271.1 285.9 1.05 + 8192 558.8 531.7 0.95 + +Read those with two caveats. First, both implementations have a payload-independent +floor -- ours is ~28us and does not move between M=8 and M=64 -- so the small-M +ratios are two constants divided, not a throughput result. Second, single runs of +either implementation 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. + +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 @@ -127,7 +150,13 @@ def _median_us(fn, warmup: int, iters: int, *, graph: bool = True) -> float: def run(args) -> int: local_rank, rank, world_size, uid = _setup_distributed() - cfg = ArConfig(world_size=world_size, m=args.m, n=args.n) + 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, + ) cfg.validate() vmm = max(4 * cfg.window_bytes + VMM_SLACK, VMM_SLACK) @@ -155,7 +184,9 @@ def run(args) -> int: if args.backend == "lsa": from kernels_lsa import build_lsa_ar - launch, stage = build_lsa_ar(cfg, rank, force_stage=args.force_stage) + launch, stage = build_lsa_ar( + cfg, rank, force_stage=args.force_stage, gather=args.gather + ) elif args.backend == "sdma": from kernels_sdma import build_sdma_ar @@ -230,6 +261,7 @@ def once(): 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, @@ -275,6 +307,26 @@ def build_parser() -> argparse.ArgumentParser: 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("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=5) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/ar/kernels_lsa.py b/benchmark/cco/flydsl/ar/kernels_lsa.py index c69bae49f..18b14aa41 100644 --- a/benchmark/cco/flydsl/ar/kernels_lsa.py +++ b/benchmark/cco/flydsl/ar/kernels_lsa.py @@ -41,10 +41,18 @@ visibility is only guaranteed between threads with the same tid (``custom_all_reduce.cuh:556-560``). -Deliberate divergence: aiter's 2-stage stages every peer's pack through LDS and -reduces on ``warp_id == 0`` only, costing two ``__syncthreads()`` per iteration -at 1/8 occupancy. Here each thread reads all peers straight into registers. The -benchmark A/Bs the two. +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. @@ -106,13 +114,27 @@ def _spin_until(rsrc, flag): scf.YieldOp([nxt.ir_value() if hasattr(nxt, "ir_value") else nxt]) -def build_lsa_ar(cfg, rank: int, *, force_stage: int | None = None): +#: 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): @@ -266,21 +288,51 @@ def ar_2stage(dev_comm: Int64, win: Int64): # --- 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))) - for j in range_constexpr(ws): - src_rank = (rank + j) % ws - src = create_buffer_resource_from_addr( - wave_uniform_i64(w.lsa_ptr(src_rank, tmp_off)) + srcs = [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr((rank + j) % ws, tmp_off)) ) - dst_base = src_rank * part - n_here = cfg.owner_pack_range(src_rank)[1] - dst_base - for pk in range(gtid, n_here, stride_packs): - v = buffer_load( - src, pk * I32_PER_PACK, vec_width=4, dtype=i32_type() - ) - buffer_store( - v, out, (dst_base + pk) * I32_PER_PACK, cache_modifier=CM_CACHED - ) + 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 diff --git a/benchmark/cco/flydsl/ar/layout.py b/benchmark/cco/flydsl/ar/layout.py index 9bb9d7001..18a5890d1 100644 --- a/benchmark/cco/flydsl/ar/layout.py +++ b/benchmark/cco/flydsl/ar/layout.py @@ -31,11 +31,11 @@ * ``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``). -* ``blocks = min(80, ceil((size_packs / world) / (threads / world)))`` with - ``threads = 512`` (``custom_all_reduce.cuh:3790-3796``). * 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. """ @@ -51,6 +51,19 @@ 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 + # 1-stage thresholds, in bytes (custom_all_reduce.cuh:3779). ONE_STAGE_MAX_BYTES_LE4 = 160 * 1024 ONE_STAGE_MAX_BYTES_LE8 = 80 * 1024 @@ -91,7 +104,12 @@ class ArConfig: n: int elem_bytes: int = 2 # bf16 threads: int = THREADS - max_blocks: int = K_MAX_BLOCKS + max_blocks: int = K_MAX_BLOCKS # signal-array rows, never the launched grid + block_cap: int = LSA_BLOCK_CAP + #: 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 --- @@ -151,12 +169,15 @@ def threads_per_peer(self) -> int: @property def blocks(self) -> int: - per_block = self.threads_per_peer - if self.stage == 1: - need = (self.num_packs + per_block - 1) // per_block - else: - need = (self.packs_per_rank + per_block - 1) // per_block - return max(1, min(self.max_blocks, need)) + # 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 ] --- # ``tmp`` holds each rank's reduced slice during 2-stage; peers all-gather @@ -237,6 +258,11 @@ def validate(self) -> None: "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.elem_bytes not in (2, 4): raise ValueError(f"elem_bytes must be 2 or 4, got {self.elem_bytes}") @@ -245,6 +271,7 @@ def validate(self) -> None: "ArConfig", "select_stage", "K_MAX_BLOCKS", + "LSA_BLOCK_CAP", "THREADS", "MAX_WORLD", "PACK_BYTES", diff --git a/tests/python/cco/test_flydsl_ar.py b/tests/python/cco/test_flydsl_ar.py index 636de5f45..28ab0fda2 100644 --- a/tests/python/cco/test_flydsl_ar.py +++ b/tests/python/cco/test_flydsl_ar.py @@ -131,9 +131,27 @@ def test_decode_shape_picks_two_stage(): def test_blocks_are_capped_and_never_zero(): - assert _cfg(m=16384).blocks == layout.K_MAX_BLOCKS + assert _cfg(m=16384).blocks == layout.LSA_BLOCK_CAP assert _cfg(m=1).blocks >= 1 - assert _cfg(m=1).blocks <= layout.K_MAX_BLOCKS + 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(): From 0f4d2776aaee752857b80539364ea9cea23b259a Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 03:10:00 +0000 Subject: [PATCH 03/79] feat(cco): SDMA all-reduce in FlyDSL, and settle the _ns/quiet question Same algorithm, window and wire traffic as the LSA 2-stage kernel; the copy engines move the bytes instead of the CUs. Three kernels (scatter push / local reduce / gather push) because both transfers need a device-wide ordering point and a kernel boundary is the only one that is free and graph-capturable. Arrival is signalled over LSA, not SDMA. This is forced, not a preference: FlyDSL's Sdma exposes no wait_signal, and its put wrapper hard-codes remoteSignal=false (cco_device_wrapper.cpp:140), so a peer cannot observe an SDMA signal at all. That leaves gcnasm's opus_sdma_a2a_quiet_notify shape -- no per-PUT signal, quiet_queue on the sender, then an LSA atomic into the peer's flag. The flag array, its monotonic epoch and the spin loop are reused from kernels_lsa; only row 0 is needed, since one warp issues the pushes. Resolves the open question the plan flagged. FlyDSL's Sdma._xfer docstring said a no-signal put "cannot be drained by quiet"; the gcnasm README said the opposite. The README is right -- cco.hpp:1707 states quiet()/quietQueue() are "independent of the signals", and quietQueue drains the queue read pointer (cco.hpp:1790). Verified rather than assumed: --sdma-signal A/Bs the trailing ATOMIC, and the no-signal default validates rel_l2=0 on 2/4/8 ranks while being slightly faster (M=512: 79.2us vs 84.0). The misleading docstring is corrected here. Measured on 8x MI355X, [M, 7168] bf16, best of 3 x 101 iterations: M aiter LSA (/aiter) SDMA (/LSA) 64 25.6us 30.7us 1.20 47.4us 1.54 128 27.2 41.8 1.54 51.7 1.24 512 53.6 66.5 1.24 79.2 1.19 1024 85.6 100.6 1.17 116.4 1.16 2048 146.1 160.3 1.10 194.2 1.21 4096 271.1 285.9 1.05 342.4 1.20 8192 558.8 531.7 0.95 637.8 1.20 SDMA does not reach the plan's +-15%-of-aiter bar (it lands at 1.26-1.48x for M >= 512) and on this evidence it will not: the gap decomposes into a ~12us fixed floor -- SDMA's payload-independent time is ~40us against LSA's ~28us, which is two transfer phases times the engine's ~6us dispatch cost -- plus ~18% lower sustained bandwidth than CU-issued loads, still present at M=8192. Neither is a tuning knob. cco's own chunking is not implicated: CCO_SDMA_MAX_COPY_BYTES is 1GB, so each 7.3MB slice is a single packet. This backend is not here to be faster. It is here because the transfers occupy one warp of one block rather than the whole grid, which is the property step 3 needs to hide an all-reduce inside a GEMM epilogue. Reporting it as a standalone collective is the wrong scoreboard, so the numbers above are recorded as a baseline for that comparison rather than as a result. layout.py gains an optional `recv` region (one landing slice per peer, slot `rank` unused so the index is the peer rank). LSA passes recv_slots=0 and its window is unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/bench_ar.py | 46 ++-- benchmark/cco/flydsl/ar/kernels_sdma.py | 283 +++++++++++++++++++++++ benchmark/cco/flydsl/ar/layout.py | 34 ++- python/mori/cco/device/flydsl/handles.py | 5 +- tests/python/cco/test_flydsl_ar.py | 26 +++ 5 files changed, 375 insertions(+), 19 deletions(-) create mode 100644 benchmark/cco/flydsl/ar/kernels_sdma.py diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py index bff35c5ce..a600b0f69 100644 --- a/benchmark/cco/flydsl/ar/bench_ar.py +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -34,20 +34,24 @@ 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 ratio - 64 25.6us 30.7us 1.20 - 128 27.2 41.8 1.54 - 512 53.6 66.5 1.24 - 1024 85.6 100.6 1.17 - 2048 146.1 160.3 1.10 - 4096 271.1 285.9 1.05 - 8192 558.8 531.7 0.95 - -Read those with two caveats. First, both implementations have a payload-independent -floor -- ours is ~28us and does not move between M=8 and M=64 -- so the small-M -ratios are two constants divided, not a throughput result. Second, single runs of -either implementation 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. + M aiter LSA (/aiter) SDMA (/LSA) + 64 25.6us 30.7us 1.20 47.4us 1.54 + 128 27.2 41.8 1.54 51.7 1.24 + 512 53.6 66.5 1.24 79.2 1.19 + 1024 85.6 100.6 1.17 116.4 1.16 + 2048 146.1 160.3 1.10 194.2 1.21 + 4096 271.1 285.9 1.05 342.4 1.20 + 8192 558.8 531.7 0.95 637.8 1.20 + +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 is ~18% slower than LSA even after removing the floor, +so it does not win on bandwidth and is not meant to -- it wins by leaving the CUs +free, which only pays off once a GEMM is running underneath it. 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, @@ -156,6 +160,9 @@ def run(args) -> int: n=args.n, max_blocks=max(args.max_blocks, args.blocks or 0), force_blocks=args.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, ) cfg.validate() @@ -190,7 +197,9 @@ def run(args) -> int: elif args.backend == "sdma": from kernels_sdma import build_sdma_ar - launch, stage = build_sdma_ar(cfg, rank, queues=args.sdma_queues) + launch, stage = build_sdma_ar( + cfg, rank, queues=args.sdma_queues, signal=args.sdma_signal + ) else: raise ValueError(f"unknown backend {args.backend!r}") @@ -328,6 +337,13 @@ def build_parser() -> argparse.ArgumentParser: help="signal-array rows; must be >= --blocks", ) 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( diff --git a/benchmark/cco/flydsl/ar/kernels_sdma.py b/benchmark/cco/flydsl/ar/kernels_sdma.py new file mode 100644 index 000000000..0c4f88d20 --- /dev/null +++ b/benchmark/cco/flydsl/ar/kernels_sdma.py @@ -0,0 +1,283 @@ +# 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``): ~1.20x slower across M >= 128, of which ~12us is a fixed floor +(two transfer phases x the engine's ~6us dispatch cost) and the rest is ~18% +lower sustained bandwidth than CU-issued loads over the same links. Both are +expected. This backend is not here to be faster -- it is here because the +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 sys + +import flydsl.compiler as flyc +import flydsl.expr as fx +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 + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from _compat import ( # noqa: E402 + CM_CACHED, + 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 # noqa: E402 +from layout import MAX_WORLD # noqa: E402 + +#: 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_ar(cfg, rank: int, *, queues: int = 8, signal: bool = False): + """Compile an SDMA all-reduce launcher 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. + + ``queues`` must match ``reqs.sdma_queue_count``; one queue per peer keeps + concurrently-issuing warps per queue at 1, which is cco's stated rule. + + ``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 < ws: + raise ValueError( + f"need one SDMA queue per peer: queues={queues} < world_size={ws}" + ) + + threads, blocks = cfg.threads, cfg.blocks + elem_dtype = fx.BFloat16 if cfg.elem_bytes == 2 else fx.Float32 + stride_packs = blocks * threads + 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 + 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, src_off_expr): + """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. + """ + + @flyc.kernel(known_block_size=[PUSH_THREADS, 1, 1]) + def push(dev_comm: Int64, win: Int64): + tid = fx.thread_idx.x + w = cco.Window(win) + sdma = cco.DevComm(dev_comm).sdma() + + flag, flag_rsrc = _next_flag(w, 0) + if tid < ws: + if tid != rank: + # 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(slice_bytes), + tid, + coop=cco.CoopScope.THREAD, + signal=signal, + ) + sdma.quiet_queue(tid, tid) + # 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.Window(win) + + srcs = [ + create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, in_off + my_slice_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, stride_packs): + 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()) + ) + 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) + + # Gather. My reduced slice goes to the same offset in every peer's output. + 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), + ) + + @flyc.jit + def run(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): + scatter(dev_comm, win).launch( + grid=(1, 1, 1), block=[PUSH_THREADS, 1, 1], stream=stream + ) + sdma_reduce(dev_comm, win).launch( + grid=(blocks, 1, 1), block=[threads, 1, 1], stream=stream + ) + gather(dev_comm, win).launch( + grid=(1, 1, 1), block=[PUSH_THREADS, 1, 1], stream=stream + ) + + return run, 2 + + +__all__ = ["build_sdma_ar", "PUSH_THREADS"] diff --git a/benchmark/cco/flydsl/ar/layout.py b/benchmark/cco/flydsl/ar/layout.py index 18a5890d1..0724485aa 100644 --- a/benchmark/cco/flydsl/ar/layout.py +++ b/benchmark/cco/flydsl/ar/layout.py @@ -106,6 +106,12 @@ class ArConfig: 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 #: 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. @@ -179,9 +185,10 @@ def blocks(self) -> int: need = (covered + per_block - 1) // per_block return max(1, min(self.block_cap, need)) - # --- symmetric window layout: [ signal | input | output | tmp ] --- + # --- 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: @@ -227,9 +234,27 @@ def tmp_bytes(self) -> int: return self.largest_pack_part * PACK_BYTES @property - def window_bytes(self) -> int: + def recv_off(self) -> int: return self.tmp_off + self.tmp_bytes + @property + def recv_bytes(self) -> int: + return self.recv_slots * self.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}" + ) + return self.recv_off + peer * self.slice_bytes + + @property + def window_bytes(self) -> int: + return self.recv_off + self.recv_bytes + # --- traffic model, for reporting alongside measured time --- @property @@ -263,6 +288,11 @@ def validate(self) -> None: 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 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}") diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index a8f917d8a..30dce38d1 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -259,8 +259,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/tests/python/cco/test_flydsl_ar.py b/tests/python/cco/test_flydsl_ar.py index 28ab0fda2..ae55287a8 100644 --- a/tests/python/cco/test_flydsl_ar.py +++ b/tests/python/cco/test_flydsl_ar.py @@ -75,6 +75,32 @@ def test_window_regions_are_disjoint_and_ordered(): 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 From 3bdc32a4101ce2a421979821ff4a45f1f8fb54d9 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 03:51:24 +0000 Subject: [PATCH 04/79] feat(cco): fuse the AR scatter into aiter's 8-wave GEMM epilogue Step 3 of the wo_b plan. 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 overlaps the rest of the GEMM. Reduce and all-gather are unchanged, reused from ar/kernels_sdma via the new build_sdma_phases with scatter swapped for a drain-only twin. The result is negative, and cleanly attributed. 8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51: gemm-only 49.6us split-lsa 327.0 split-sdma 382.0 fused-sdma 450.2 (chunks=2) Fused is 66-68us worse, and by the same 66-68us at K=4096 and K=8192 -- a fixed cost, not a scaling one. It is the release fence: * --fence all (every lane) costs 70us over the 448 blocks: 452.0 vs 382.4. * --fence leader is *wrong*, measured not inferred: relL2 goes 2.35e-3 -> 8.9e-3 on 8 ranks. s_waitcnt vmcnt(0) plus s_barrier retires every lane's stores and is still not enough, which is exactly what cco_device_wrapper.cpp:105-111 says -- leaderOnly orders thread 0 only and does not substitute for a release from every producer lane. I had assumed the opposite (that the engine reads through our own L2 so no fence was needed at all); the first 8-rank run disagreed at 3.76e-3 vs split's 2.35e-3, which is why the modes now cross-check each other in the tests rather than each being judged "small enough" alone. * --fence none is incorrect too and exists only to price the ceiling: 373.0 vs split's 382.0. So even with a free fence the overlap is worth ~9us, against a GEMM of 49.6us. Why ~9us and not ~50: 448 blocks over 256 CUs is 1.75 waves, so no destination's first chunk completes before ~57% of the GEMM has elapsed, and each push then pays the engine's ~6us dispatch. Only the tail is available to hide anything in. Finer chunking helps as far as it can (chunks=2 is 450.2 vs chunks=1's 505.6) but the slice is only 2 BLOCK_M row-bands here, so 2 is the cap. The obstacle underneath is the ratio the plan flagged as risk 6: the transfer is ~2.7x the GEMM (7.3MB per link at ~50GB/s is ~147us, twice, against 49.6us of compute), so the links are the critical path with or without overlap and perfect fusion could remove at most the GEMM. Raising K does not rescue it -- K and the collective are independent axes. Two things that did work and are worth keeping: * Destination-rotated, chunk-major tile order (gcnasm's opus_direct_stripe_tile). aiter's linear order finishes destination 0's whole slice first, so the last destination's link only starts at the end of the GEMM. Rotating costs nothing measurable: 49.72us vs 49.80. * The completion-counter protocol itself. Monotonic per (dest, chunk) counters with a modulo test elect exactly one CTA per chunk, so no per-destination spin lock is needed (gcnasm needed one because several CTAs could submit), and never resetting them keeps the kernel graph-safe. The LSA transport is deliberately not implemented; LSA_NOTE in kernels_fused records why. It needs StoreC coalesced first (its _store_bf16 is a per-lane scatter that gcnasm measured at 0.26x when pushed to a peer), and the ceiling that rewrite competes for is the ~50us just shown to be unreachable. kernels_fused.py pins a copy of aiter's kernel_gemm rather than adding a hook, because the original's main loop has 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. Everything reusable is imported; the copy is the ~90-line pipeline only, and test_gemm_ar.py asserts it is bit-identical to aiter's kernel so it cannot rot silently. Not a production claim: aiter's tuned CSV picks ck/asm/cktile for (N=7168, K=1024), not this host. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_compat.py | 20 + benchmark/cco/flydsl/ar/kernels_sdma.py | 96 ++- benchmark/cco/flydsl/ar/layout.py | 32 +- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 355 +++++++++++ benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 572 ++++++++++++++++++ tests/python/cco/test_gemm_ar.py | 231 +++++++ 6 files changed, 1270 insertions(+), 36 deletions(-) create mode 100644 benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py create mode 100644 benchmark/cco/flydsl/gemm_ar/kernels_fused.py create mode 100644 tests/python/cco/test_gemm_ar.py diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py index 0dbd64fcb..d6daa2824 100644 --- a/benchmark/cco/flydsl/ar/_compat.py +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -257,6 +257,25 @@ 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): + """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. ``acq_rel``: the release half publishes this + block's C stores to whoever observes the count, and the acquire half is what + lets the winning block read every other block's tile. + """ + return _llvm_d.atomicrmw( + _llvm_d.AtomicBinOp.add, + ptr, + fx.Int32(value).ir_value(), + _llvm_d.AtomicOrdering.acq_rel, + syncscope=_AGENT_SCOPE, + alignment=4, + ) + + def i32_type(): return _dtype("i32") @@ -273,6 +292,7 @@ def i32_type(): "signal_load_u32", "local_load_u32", "local_store_u32", + "atomic_add_u32", "i32_type", "CM_CACHED", "CM_SC1", diff --git a/benchmark/cco/flydsl/ar/kernels_sdma.py b/benchmark/cco/flydsl/ar/kernels_sdma.py index 0c4f88d20..c4bd42611 100644 --- a/benchmark/cco/flydsl/ar/kernels_sdma.py +++ b/benchmark/cco/flydsl/ar/kernels_sdma.py @@ -71,8 +71,8 @@ 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 import range_constexpr from flydsl.expr.typing import Int64 import mori.cco.device.flydsl as cco @@ -100,12 +100,13 @@ PUSH_THREADS = 64 -def build_sdma_ar(cfg, rank: int, *, queues: int = 8, signal: bool = False): - """Compile an SDMA all-reduce launcher for ``cfg`` on ``rank``. +def build_sdma_phases(cfg, rank: int, *, queues: int = 8, signal: bool = False): + """Compile the phases separately, so the fused GEMM can reuse the tail. - 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. + 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``; one queue per peer keeps concurrently-issuing warps per queue at 1, which is cco's stated rule. @@ -158,13 +159,15 @@ def _signal_and_wait(w, arr_off, flag, tid): 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, src_off_expr): + def _push_kernel(arr_off, dst_off_of_peer=None, src_off_expr=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. + on which peer the lane is serving. Passing neither emits a drain-only + kernel -- the pushes are assumed already issued elsewhere. """ + pushes = dst_off_of_peer is not None @flyc.kernel(known_block_size=[PUSH_THREADS, 1, 1]) def push(dev_comm: Int64, win: Int64): @@ -175,19 +178,21 @@ def push(dev_comm: Int64, win: Int64): flag, flag_rsrc = _next_flag(w, 0) if tid < ws: if tid != rank: - # 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(slice_bytes), - tid, - coop=cco.CoopScope.THREAD, - signal=signal, - ) + 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(slice_bytes), + tid, + coop=cco.CoopScope.THREAD, + signal=signal, + ) sdma.quiet_queue(tid, tid) # Release before publishing arrival. The bytes were moved by the # copy engine rather than by this CU, so there is nothing of ours @@ -265,19 +270,44 @@ def sdma_reduce(dev_comm: Int64, win: Int64): src_off_expr=lambda tid: fx.Int64(out_off + my_slice_off), ) - @flyc.jit - def run(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): - scatter(dev_comm, win).launch( - grid=(1, 1, 1), block=[PUSH_THREADS, 1, 1], stream=stream - ) - sdma_reduce(dev_comm, win).launch( - grid=(blocks, 1, 1), block=[threads, 1, 1], stream=stream - ) - gather(dev_comm, win).launch( - grid=(1, 1, 1), block=[PUSH_THREADS, 1, 1], stream=stream - ) + # 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 + + return { + "scatter": _phase(scatter, 1, PUSH_THREADS), + "drain": _phase(drain, 1, PUSH_THREADS), + "reduce": _phase(sdma_reduce, blocks, threads), + "gather": _phase(gather, 1, PUSH_THREADS), + } + + +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) + scatter, reduce_, gather = parts["scatter"], parts["reduce"], parts["gather"] + + def run(dev_comm, win, stream=fx.Stream(None)): + scatter(dev_comm, win, stream=stream) + reduce_(dev_comm, win, stream=stream) + gather(dev_comm, win, stream=stream) return run, 2 -__all__ = ["build_sdma_ar", "PUSH_THREADS"] +__all__ = ["build_sdma_ar", "build_sdma_phases", "PUSH_THREADS"] diff --git a/benchmark/cco/flydsl/ar/layout.py b/benchmark/cco/flydsl/ar/layout.py index 0724485aa..ceeea4e01 100644 --- a/benchmark/cco/flydsl/ar/layout.py +++ b/benchmark/cco/flydsl/ar/layout.py @@ -112,6 +112,11 @@ class ArConfig: #: 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 #: 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. @@ -192,9 +197,7 @@ def blocks(self) -> int: @property def signal_bytes(self) -> int: - row = _align_up(self.max_blocks * MAX_WORLD * 4, SIGNAL_ALIGN) - flag = _align_up(self.max_blocks * 4, SIGNAL_ALIGN) - return row * 2 + flag # start[] + end[] + _flag[] + return self.counter_off + self.counter_bytes @property def start_off(self) -> int: @@ -208,6 +211,27 @@ def end_off(self) -> int: def flag_off(self) -> int: return self.end_off * 2 + @property + def counter_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_chunks`` 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_bytes(self) -> int: + return _align_up(self.world_size * self.counter_chunks * 4, 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: @@ -288,6 +312,8 @@ def validate(self) -> None: 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 not 0 <= self.recv_slots <= self.world_size: raise ValueError( f"recv_slots must be in [0, world_size={self.world_size}], " 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..bbe628fc2 --- /dev/null +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -0,0 +1,355 @@ +#!/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 PYTHONPATH=/path/to/aiter \\ + 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) + 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, 8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51: + + gemm-only 49.6us + split-lsa 327.0 + split-sdma 382.0 + fused-sdma 450.2 (chunks=2) + +Fusing loses, by a fixed 66-68us that is the mandatory all-lane release fence. +``kernels_fused.py`` has the decomposition; the short version is that the fence +costs 70us while the overlap available is ~9us. ``--fence`` and ``--stop-after`` +exist to reproduce that attribution, not because either has a correct non-default +setting. +""" + +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 + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "ar")) +from kernels_fused import compile_fused_gemm_scatter # noqa: E402 +from kernels_lsa import build_lsa_ar # noqa: E402 +from kernels_sdma import build_sdma_phases # noqa: E402 +from layout import ArConfig # noqa: E402 + +VMM_SLACK = 512 * 1024 * 1024 +MODES = ("gemm-only", "split-sdma", "fused-sdma", "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]) + + +def make_operands(rank: int, m: int, n: int, k: int): + """Deterministic per-rank fp8 operands, plus the bf16 reference partial. + + 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. + """ + 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) + 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 + + +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: + from aiter.ops.shuffle import shuffle_weight + + local_rank, rank, world_size, uid = _setup_distributed() + fused = args.mode == "fused-sdma" + needs_sdma = args.mode in ("split-sdma", "fused-sdma") + # One push per BLOCK_M row-band by default: the finest granularity the tile + # order can signal, and at the wo_b shape still a 3.7MB transfer, well past + # the ~1MB where SDMA bandwidth flattens out. + m_tiles_per_peer = max(1, (args.m // world_size) // args.block_m) + chunks = args.chunks if args.chunks else m_tiles_per_peer + cfg = ArConfig( + world_size=world_size, + m=args.m, + n=args.n, + recv_slots=world_size if needs_sdma else 0, + counter_chunks=chunks, + ) + cfg.validate() + + a, b, sa, sb = make_operands(rank, args.m, args.n, args.k) + b_shuf = shuffle_weight(b, layout=(16, 16)) + + 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, + waves_per_eu=args.waves_per_eu, + xcd_swizzle=args.xcd_swizzle, + fuse=fused, + rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", + fence=args.fence, + ) + a_i8 = a.contiguous().view(torch.int8).view(-1) + b_i8 = b_shuf.contiguous().view(torch.int8).view(-1) + c_flat = c.view(-1) + + def run_gemm(stream): + gemm(a_i8, b_i8, c_flat, sa, sb, args.m, args.n, dc.ptr, win.handle, + stream=stream) + + if needs_sdma: + parts = build_sdma_phases(cfg, rank, queues=args.sdma_queues) + 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: the scatter puts were already issued from the epilogue, so + # only the drain runs here; split: the full scatter kernel. + parts["drain" if fused else "scatter"](dc.ptr, win.handle, stream=stream) + if args.stop_after == "scatter": + return + parts["reduce"](dc.ptr, win.handle, stream=stream) + parts["gather"](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 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) + acc += (ar.float() @ br.float().T) * sar[:, None] * sbr[None, :] + 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. + validated = rel_l2 < 5e-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, + "block_m": args.block_m, + "block_n": args.block_n, + "chunks": chunks if fused else None, + "tile_order": args.tile_order, + "fence": args.fence if fused else None, + "stop_after": args.stop_after, + "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("--block-m", type=int, default=256) + 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 = one per BLOCK_M row-band, the finest " + "the tile order can signal)", + ) + p.add_argument( + "--tile-order", + choices=("auto", "rotated", "linear"), + default="auto", + help="auto = rotated when fusing, linear otherwise", + ) + p.add_argument( + "--fence", + choices=("all", "leader", "none"), + default="all", + help="release before the epilogue push. Only 'all' is correct: 'leader' " + "measurably corrupts (relL2 8.9e-3 vs 2.35e-3 on 8 ranks) and 'none' more " + "so. Both exist only to price the 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("--sdma-queues", type=int, default=8) + 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/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py new file mode 100644 index 000000000..ca503a342 --- /dev/null +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -0,0 +1,572 @@ +# 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. +"""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. + +## 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 -- having the epilogue store straight into +a peer -- would first require coalescing ``StoreC._store_bf16``, which writes one +bf16 at a time through ``BufferCopy16b``; gcnasm measured that lane-scatter at +0.26x when pushed to a peer. See ``LSA_NOTE`` at the bottom. + +## 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. + +## Result: this fusion does not pay at the wo_b shape + +8x MI355X, [4096, 7168] bf16 out, K=1024, CUDA-graph replay, median of 51: + + gemm-only 49.6us + split-lsa gemm + LSA AR 327.0 + split-sdma gemm + SDMA AR 382.0 + fused-sdma chunks=2 450.2 + +Fused is *worse*, by 66-68us -- and by the same 66-68us at K=4096 and K=8192, so +it is a fixed cost, not a scaling one. It is the release fence, and the fence is +not negotiable: + +* ``--fence all`` (every lane, the only correct form) costs 70us across the 448 + blocks: 452.0us against 382.4 for ``--fence leader``. +* ``--fence leader`` is *wrong*, which is a measurement and not an inference: + relL2 goes from 2.35e-3 to 8.9e-3 on 8 ranks. `s_waitcnt vmcnt(0)` plus + `s_barrier` retires every lane's stores, and it is still not enough -- exactly + what ``cco_device_wrapper.cpp:105-111`` warns, that leaderOnly orders thread 0 + only and does not substitute for a release from every producer lane. +* ``--fence none`` is incorrect too, and exists only to price the ceiling: it + gives 373.0us against split's 382.0. So **even with a free fence the overlap is + worth ~9us**, against a GEMM of 49.6us. + +Why the overlap is ~9us and not ~50us: 448 blocks over 256 CUs is 1.75 waves, so +no destination's first chunk can complete before ~57% of the GEMM has elapsed, +and each push then pays the copy engine's ~6us dispatch. Only the tail of the +GEMM is available to hide anything in. Finer chunking helps as far as it can go +(chunks=2 is 450.2us against chunks=1's 505.6), but at this shape the slice is +only 2 BLOCK_M row-bands, so 2 is the maximum. + +The underlying ratio is the real obstacle, and it is the one the plan flagged: +the transfer is ~2.7x the GEMM (a 7.3MB slice per link at ~50GB/s is ~147us, +twice, against 49.6us of compute), so the links are the critical path with or +without overlap and perfect fusion could remove at most the GEMM. Raising K does +not rescue it -- K=4096 and K=8192 show the same constant deficit -- because N +and M set both the C traffic and the collective, and K only buys compute the +collective still has to wait for. + +Read the numbers against ``split-lsa`` (327.0us), not ``split-sdma``: the +question a fused SDMA path has to answer is whether it beats the *best* split +pipeline, and it is 38% behind that. + +## Pinned copy + +``_fused_kernel_body`` is a copy of ``compile_fp8_gemm_8w``'s ``kernel_gemm`` +(aiter ``aiter/ops/flydsl/kernels/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. + +``tests/python/cco/test_gemm_ar.py`` asserts this kernel's C output is +bit-identical to aiter's unfused kernel on the same inputs, which is what stops +the copy from rotting silently. +""" + +# NOTE: no `from __future__ import annotations` here, deliberately, and aiter's +# 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 os +import sys + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr.typing import Int64 + +from aiter.ops.flydsl.kernels.gemm_a8w8_8wave import ( + G2SLoader, + Mfma16x16x128, + S2RLoader, + StoreC, + _xcd_swizzle_any, + ceildiv, + compute_global_swizzle, + make_fp8_buffer_tensor, + wait_barrier, +) +from aiter.ops.flydsl.kernels.mfma_preshuffle_pipeline import split_row_major_2d + +import mori.cco.device.flydsl as cco +from mori.cco.device.flydsl import _bindings as raw_cco + +_AR_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ar") +sys.path.insert(0, _AR_DIR) +from _compat import atomic_add_u32, signal_ptr # noqa: E402 + +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, + fence: str = "all", +): + """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 aiter's 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. + + 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 + rotated = fuse if rotated is None else 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 + tiles_per_peer = m_tiles_per_peer * n_blocks_const + + # 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", "leader", "none"): + raise ValueError(f"fence must be all/leader/none, got {fence!r}") + _kname_tag = f"c{chunks}{'r' if rotated else 'l'}{fence[0]}" + counter_off = cfg.counter_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_{'sdma' if fuse else 'split'}_8w_" + f"{BLOCK_M}x{BLOCK_N}x{BLOCK_K}_k{K}_{_kname_tag}_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 aiter 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) + rest, bn = split_row_major_2d(fx.block_idx.x, n_blocks) + tile_i, dest_seq = split_row_major_2d(rest, ws) + 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 + ) + + 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) + + 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) + + 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 + 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 + 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() + + 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) + # ---- end pinned copy ---- + + if const_expr(fuse): + # 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)) + + w = cco.Window(win) + sdma = cco.DevComm(dev_comm).sdma() + if fx.thread_idx.x == 0: + 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( + fx.Int64(w.lsa_ptr(rank, counter_off)) + + fx.Int64(slot) * fx.Int64(4) + ) + seq = fx.Int32(atomic_add_u32(ctr, 1)) + 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 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) + 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), + dest, + coop=cco.CoopScope.THREAD, + signal=False, + ) + + @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 + + +LSA_NOTE = """\ +The LSA transport is not implemented, and the reason is a measurement rather than +a preference. + +Fusing LSA means the epilogue stores into a peer instead of into local C. But +`StoreC._store_bf16` writes **one bf16 at a time** through `BufferCopy16b`, at +`c_index = (row + i) * c_cols + col` with `row = base_row + ti*16 + lane//16*4 + i` +and `col = base_col + tj*16 + lane%16` -- a lane-scatter. gcnasm measured that +pattern at 0.26x when the destination is a peer, so the variant is worthless +until `StoreC` is coalesced (their winning form was a wave-local `ds_bpermute` +pair-coalesced store; staging C through LDS regressed 1-3%). + +That rewrite is the expensive half of this step, and the ceiling it is competing +for is small: at [4096, 7168] the GEMM is ~40us against an all-reduce of ~290us +(LSA) or ~340us (SDMA), so even perfect overlap removes at most the GEMM, ~12%. +Fusing the cheap transport first establishes whether the completion-counter +structure works at all; if the answer is no, the `StoreC` rewrite is wasted +regardless of transport. +""" + +__all__ = ["compile_fused_gemm_scatter", "LSA_NOTE", "BLOCK_K"] diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py new file mode 100644 index 000000000..59de2c660 --- /dev/null +++ b/tests/python/cco/test_gemm_ar.py @@ -0,0 +1,231 @@ +# 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 GEMM + all-reduce: anti-rot guard for the pinned kernel copy, plus smoke. + +``benchmark/cco/flydsl/gemm_ar/kernels_fused.py`` holds a copy of aiter's +``kernel_gemm`` with an epilogue added. The single most valuable test here is not +that the fused path is correct end to end -- ``bench_gemm_ar.py`` checks that +against a host reference -- but that the *copy* still computes what aiter's +original computes. That one runs on a single GPU and needs no cco at all. +""" + +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +GEMM_AR_DIR = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "gemm_ar" +AR_DIR = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "ar" +for _d in (str(AR_DIR), str(GEMM_AR_DIR)): + if _d not in sys.path: + sys.path.insert(0, _d) + + +def _load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +layout = _load("flydsl_ar_layout_gemm", AR_DIR / "layout.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() + + +# -------------------------------------------------------------------------- +# the anti-rot guard: our pinned copy must match aiter's original bit-for-bit +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("m,n,k", [(512, 512, 256), (1024, 768, 512)]) +def test_pinned_copy_matches_aiter_kernel_bitwise(m, n, k): + """The unfused copy must be indistinguishable from aiter's kernel. + + Bitwise, not approximate: the two are meant to be the same instruction + sequence over the same inputs, so any drift in the copied pipeline shows up + here rather than as a mystery in a fused benchmark. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + pytest.importorskip("aiter", reason="aiter not importable (set PYTHONPATH)") + pytest.importorskip("flydsl") + from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 + from aiter.ops.shuffle import shuffle_weight + + import flydsl.expr as fx + + from kernels_fused import compile_fused_gemm_scatter + + 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 = shuffle_weight(b, layout=(16, 16)) + + ref = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) + flydsl_8wave_gemm_a8(a, b_shuf, sa, sb, ref, 256, 256) + + # 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() + assert torch.equal(got, ref), ( + "the pinned copy of aiter's kernel_gemm has drifted from the original; " + "re-sync benchmark/cco/flydsl/gemm_ar/kernels_fused.py" + ) + + +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(GEMM_AR_DIR / "bench_gemm_ar.py"), + "--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] + + +@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("aiter", reason="aiter not importable (set PYTHONPATH)") + 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 From 16ddbb9454ca390f0fe8a5cf111e407c0ce0c96b Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 04:16:26 +0000 Subject: [PATCH 05/79] perf(cco): profile the fused GEMM+AR; the cost is a quadratic release fence rocprofv3 --kernel-trace contradicts the attribution in the previous commit. I had blamed a flat ~70us of fence spread across the pipeline, inferred from end-to-end totals. Per-kernel it is nothing like that (8x MI355X, [4096, 7168] out, K=1024, median of 11, us): config gemm scatter/drain reduce gather sum correct split 39.1 136.8 59.0 137.3 372.3 yes fused --fence none 64.5 103.4 59.8 136.6 364.2 NO fused --fence leader 83.5 96.9 59.9 136.2 376.5 NO fused --fence agent 156.1 67.3 59.8 137.0 420.1 yes fused --fence all 175.9 71.1 59.3 137.3 443.6 yes Two corrections to what I said before. The overlap is not "worth ~9us" -- it is worth 69.5us, the scatter dropping from 136.8 to 67.3, more than the whole 39.1us GEMM. And the cost is not distributed: all of it is inside the GEMM kernel, +117us. Registers are unchanged (VGPR=128, SGPR=112, scratch=0), so it is not occupancy. The +117us splits three ways: * +25us for the tail with no fence at all (barrier, atomic, and the copy engine reading C while the GEMM still writes it). * +92us for `buffer_wbl2 sc1` and the `s_waitcnt vmcnt(0)` on it. This is quadratic: buffer_wbl2 flushes the *entire* L2, so block k writes back every tile blocks 1..k-1 already wrote, 448 blocks over. The split path gets one bulk release at kernel end for free. * +20us more for system scope over agent scope. The ISA shows why: __threadfence_system is bidirectional and emits `buffer_wbl2 sc0 sc1 ; s_waitcnt ; buffer_inv sc0 sc1`, and that buffer_inv discards the B weights the GEMM is streaming, 3584 times. Three attempts to make the release cheap, all failed, all now selectable so the measurements can be reproduced: * agent scope (new default): correct, drops the buffer_inv, buys 20us of 112. cco exposes no agent-scope release -- only cco_system_fence -- so this is emitted as a raw llvm.fence via a new _compat.release_fence. That gap is worth closing in the device API regardless of this benchmark. * one fence per block instead of per wave (--fence leader): incorrect, relL2 9.6e-3, indistinguishable from no fence, even though every lane has passed s_waitcnt vmcnt(0) and a block barrier. So the release is genuinely per-wave, which is the factor of 8 in 3584. * write-through C stores (--fence writethrough, a one-argument cache modifier on StoreC's copy atom): not a release on its own (relL2 1.72e-2), and combined with the agent fence (--fence wt-agent, 432.5us) it is indistinguishable from the agent fence alone (432.4). buffer_wbl2 costs the same whether or not the lines are dirty. So the conclusion is narrower and more useful than "the compute:comm ratio is wrong". What is missing is a **per-tile release**: every primitive available publishes the whole cache, while the GEMM needs to publish 448 tiles. The counter protocol and the destination-rotated tile order are sound -- they are what produced the 69.5us of real overlap -- and would pay immediately on a transport whose producer-side release is per-tile. The bottom line is unchanged: fused 432.4us loses to split-sdma 382.0 and to gemm+LSA 327.0, which is still the number to beat. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_compat.py | 13 ++ benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 23 +-- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 146 ++++++++++++------ 3 files changed, 124 insertions(+), 58 deletions(-) diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py index d6daa2824..5d40318f9 100644 --- a/benchmark/cco/flydsl/ar/_compat.py +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -276,6 +276,18 @@ def atomic_add_u32(ptr, value): ) +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") @@ -293,6 +305,7 @@ def i32_type(): "local_load_u32", "local_store_u32", "atomic_add_u32", + "release_fence", "i32_type", "CM_CACHED", "CM_SC1", diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index bbe628fc2..40799a7ad 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -44,13 +44,13 @@ gemm-only 49.6us split-lsa 327.0 split-sdma 382.0 - fused-sdma 450.2 (chunks=2) + fused-sdma 432.4 (chunks=2, --fence agent) -Fusing loses, by a fixed 66-68us that is the mandatory all-lane release fence. -``kernels_fused.py`` has the decomposition; the short version is that the fence -costs 70us while the overlap available is ~9us. ``--fence`` and ``--stop-after`` -exist to reproduce that attribution, not because either has a correct non-default -setting. +Fusing loses. The overlap is real -- a kernel trace shows the scatter dropping +from 136.8us to 67.3us -- but the release fence the epilogue needs adds more than +that back inside the GEMM, which goes from 39.1us to 156.1us. ``kernels_fused.py`` +has the full per-kernel decomposition. ``--fence`` and ``--stop-after`` exist to +reproduce that attribution. """ from __future__ import annotations @@ -328,11 +328,12 @@ def build_parser() -> argparse.ArgumentParser: ) p.add_argument( "--fence", - choices=("all", "leader", "none"), - default="all", - help="release before the epilogue push. Only 'all' is correct: 'leader' " - "measurably corrupts (relL2 8.9e-3 vs 2.35e-3 on 8 ranks) and 'none' more " - "so. Both exist only to price the fence", + choices=("all", "agent", "leader", "none", "writethrough", "wt-agent"), + default="agent", + help="release before the epilogue push. Correct: 'agent' (default, " + "cheapest) and 'all' (system scope, adds a buffer_inv that costs 20us). " + "Incorrect, and present only to price the fence: 'leader', 'none', " + "'writethrough'. 'wt-agent' is correct but no faster than 'agent'", ) p.add_argument( "--stop-after", diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index ca503a342..88e0e5d6a 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,48 +71,63 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## Result: this fusion does not pay at the wo_b shape - -8x MI355X, [4096, 7168] bf16 out, K=1024, CUDA-graph replay, median of 51: - - gemm-only 49.6us - split-lsa gemm + LSA AR 327.0 - split-sdma gemm + SDMA AR 382.0 - fused-sdma chunks=2 450.2 - -Fused is *worse*, by 66-68us -- and by the same 66-68us at K=4096 and K=8192, so -it is a fixed cost, not a scaling one. It is the release fence, and the fence is -not negotiable: - -* ``--fence all`` (every lane, the only correct form) costs 70us across the 448 - blocks: 452.0us against 382.4 for ``--fence leader``. -* ``--fence leader`` is *wrong*, which is a measurement and not an inference: - relL2 goes from 2.35e-3 to 8.9e-3 on 8 ranks. `s_waitcnt vmcnt(0)` plus - `s_barrier` retires every lane's stores, and it is still not enough -- exactly - what ``cco_device_wrapper.cpp:105-111`` warns, that leaderOnly orders thread 0 - only and does not substitute for a release from every producer lane. -* ``--fence none`` is incorrect too, and exists only to price the ceiling: it - gives 373.0us against split's 382.0. So **even with a free fence the overlap is - worth ~9us**, against a GEMM of 49.6us. - -Why the overlap is ~9us and not ~50us: 448 blocks over 256 CUs is 1.75 waves, so -no destination's first chunk can complete before ~57% of the GEMM has elapsed, -and each push then pays the copy engine's ~6us dispatch. Only the tail of the -GEMM is available to hide anything in. Finer chunking helps as far as it can go -(chunks=2 is 450.2us against chunks=1's 505.6), but at this shape the slice is -only 2 BLOCK_M row-bands, so 2 is the maximum. - -The underlying ratio is the real obstacle, and it is the one the plan flagged: -the transfer is ~2.7x the GEMM (a 7.3MB slice per link at ~50GB/s is ~147us, -twice, against 49.6us of compute), so the links are the critical path with or -without overlap and perfect fusion could remove at most the GEMM. Raising K does -not rescue it -- K=4096 and K=8192 show the same constant deficit -- because N -and M set both the C traffic and the collective, and K only buys compute the -collective still has to wait for. - -Read the numbers against ``split-lsa`` (327.0us), not ``split-sdma``: the -question a fused SDMA path has to answer is whether it beats the *best* split -pipeline, and it is 38% behind that. +## Result: the fusion works, and the release fence costs more than it saves + +Per-kernel, from ``rocprofv3 --kernel-trace`` on 8x MI355X, [4096, 7168] out, +K=1024, median of 11 (times in us): + + config gemm scatter/drain reduce gather sum correct + split 39.1 136.8 59.0 137.3 372.3 yes + fused --fence none 64.5 103.4 59.8 136.6 364.2 NO + fused --fence leader 83.5 96.9 59.9 136.2 376.5 NO + fused --fence agent 156.1 67.3 59.8 137.0 420.1 yes + fused --fence all 175.9 71.1 59.3 137.3 443.6 yes + +The overlap is real and large: the scatter collapses from 136.8us to 67.3us. But +every microsecond of that, and more, reappears inside the GEMM, which goes from +39.1us to 156.1us. Splitting that +117us three ways: + +* **+25us** for the tail itself with no fence at all -- ``s_waitcnt vmcnt(0)``, + the block-wide barrier, the atomic, and the copy engine reading C while the + GEMM is still writing it. +* **+92us** for ``buffer_wbl2 sc1`` and the ``s_waitcnt vmcnt(0)`` that waits on + it. This is the whole problem, and it is **quadratic**: ``buffer_wbl2`` flushes + the entire L2, so block k writes back every tile blocks 1..k-1 already wrote, + 448 blocks over. The split path gets one bulk release at kernel end instead, + for free. +* **+20us** more if the fence is system-scoped rather than agent-scoped. The ISA + says why: ``__threadfence_system`` is bidirectional, so it emits + ``buffer_wbl2 sc0 sc1 ; s_waitcnt ; buffer_inv sc0 sc1`` -- and that + ``buffer_inv`` discards the B weights the GEMM is streaming, 3584 times. + +Three things were tried to make the release cheap, and all three failed: + +* **Agent scope instead of system.** Correct (relL2 2.35e-3, same as split) and + it does drop the ``buffer_inv``, but it only buys 20us of the 112. Worth + keeping as the default anyway -- and note cco exposes no agent-scope release, + only ``cco_system_fence``, so this had to be emitted as a raw ``llvm.fence``. +* **One fence per block instead of per wave** (``--fence leader``). Incorrect: + relL2 9.6e-3, indistinguishable from no fence at all, even though every lane + has already passed ``s_waitcnt vmcnt(0)`` and a block barrier. So the release + really is per-wave, which is where the factor of 8 in 3584 comes from. +* **Write-through C stores** (``--fence writethrough``), so no writeback would be + needed. Not a release on its own (relL2 1.72e-2), and pairing it with the agent + fence (``--fence wt-agent``, 432.5us) is indistinguishable from the agent fence + alone (432.4us) -- ``buffer_wbl2`` costs the same whether or not the lines it + is asked to flush are dirty. + +So the honest conclusion is narrower than "the ratio is wrong", though the ratio +does cap the prize: what is missing is a **per-tile release**. Every primitive +available here publishes the whole cache, and the GEMM needs to publish 448 +tiles, so the release is O(blocks^2) against an overlap worth at most one GEMM. +Even the incorrect ``--fence none`` configuration only reaches 364.2us against +split's 372.3 -- and both lose to ``gemm + LSA`` at 327.0us, which remains the +number any fused SDMA path has to beat. + +The completion-counter machinery and the destination-rotated tile order are both +sound and are what produced the 69.5us of genuine overlap; they would pay +immediately on a transport whose producer-side release is per-tile rather than +per-cache. ## Pinned copy @@ -162,7 +177,31 @@ _AR_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ar") sys.path.insert(0, _AR_DIR) -from _compat import atomic_add_u32, signal_ptr # noqa: E402 +from _compat import ( # noqa: E402 + CM_SC0_SC1, + atomic_add_u32, + release_fence, + signal_ptr, +) + + +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 @@ -179,7 +218,7 @@ def compile_fused_gemm_scatter( xcd_swizzle: int = 0, fuse: bool = True, rotated: bool | None = None, - fence: str = "all", + fence: str = "agent", ): """Compile the GEMM, with (``fuse=True``) or without the scatter epilogue. @@ -248,8 +287,10 @@ def compile_fused_gemm_scatter( tiles_per_chunk = m_tiles_per_chunk * n_blocks_const chunk_bytes = cfg.slice_bytes // chunks - if fence not in ("all", "leader", "none"): - raise ValueError(f"fence must be all/leader/none, got {fence!r}") + if fence not in ("all", "agent", "leader", "none", "writethrough", "wt-agent"): + raise ValueError( + f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" + ) _kname_tag = f"c{chunks}{'r' if rotated else 'l'}{fence[0]}" counter_off = cfg.counter_off in_off = cfg.input_off @@ -348,7 +389,14 @@ def kernel_gemm_scatter( 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) + if 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 + ) + else: + store_c = StoreC( + A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B + ) c00_frag = [mfma.zero_value] * N_ACCUMS c01_frag = [mfma.zero_value] * N_ACCUMS @@ -476,6 +524,10 @@ def kernel_gemm_scatter( 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")): + release_fence("agent") + # "writethrough": nothing to do -- the stores already went to memory, + # and wait_barrier(0) above retired them. w = cco.Window(win) sdma = cco.DevComm(dev_comm).sdma() From 1b5b4278186c2d506935e8f27a8ea0ab80b1c630 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 06:01:34 +0000 Subject: [PATCH 06/79] perf(cco): ATT-guided fixes to the fused epilogue, 450us -> 413us The trace decoder *is* installed -- /opt/rocm-7.2.4/lib/librocprof-trace- decoder.so. My earlier "not available" was a bad search: /opt/rocm is a symlink and find(1) does not follow those without -L. With --att-library-path the decode works and attributes the epilogue exactly. The tail is 26% of the kernel's latency, and almost all of it is waiting: vaddr cycles % kernel instruction 19396 91336 5.1% s_waitcnt vmcnt(0) retire this block's C 19400 22984 1.3% s_barrier 19404 292 0.0% buffer_wbl2 sc1 <- the op itself is free 19412 246804 13.9% s_waitcnt vmcnt(0) <- waiting on the writeback 19608+ 91612 5.1% waits bracketing the counter atomic The writeback wait is the largest single line item in the whole kernel -- larger than any main-loop stall -- and it is quadratic, since buffer_wbl2 flushes the entire L2 and every one of the 448 blocks waits for every earlier block's tile. The last row was my own mistake and is now fixed. I had given the counter atomic acq_rel, so it emitted its own buffer_wbl2/buffer_inv pair costing more per wave than the explicit fence. Both halves were redundant: the caller has already issued the release, and the winning block never reads what it counts (the copy engine does), so it needs no acquire. Relaxed: fused --fence agent gemm 156.1 -> 136.9us, total 432.4 -> 413.4us still relL2 2.35e-3, matching split exactly. Cumulative from profiling this path: 450.2 -> 413.4us, of which 20us was agent-vs-system scope (dropping a buffer_inv that discarded the B weights the GEMM streams) and 19us this atomic. Conclusion is unchanged and now has instruction-level evidence: what is missing is a per-tile release. Every primitive available publishes the whole cache while the GEMM has 448 tiles to publish. Fused 413.4us still loses to split-sdma 372.3 and to gemm+LSA 327.0. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_compat.py | 18 ++- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 6 +- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 108 ++++++++++-------- 3 files changed, 74 insertions(+), 58 deletions(-) diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py index 5d40318f9..9f81157b4 100644 --- a/benchmark/cco/flydsl/ar/_compat.py +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -257,20 +257,26 @@ 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): +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. ``acq_rel``: the release half publishes this - block's C stores to whoever observes the count, and the acquire half is what - lets the winning block read every other block's tile. + 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(), - _llvm_d.AtomicOrdering.acq_rel, + getattr(_llvm_d.AtomicOrdering, ordering), syncscope=_AGENT_SCOPE, alignment=4, ) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 40799a7ad..4430f8004 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -44,11 +44,11 @@ gemm-only 49.6us split-lsa 327.0 split-sdma 382.0 - fused-sdma 432.4 (chunks=2, --fence agent) + fused-sdma 413.4 (chunks=2, --fence agent) Fusing loses. The overlap is real -- a kernel trace shows the scatter dropping -from 136.8us to 67.3us -- but the release fence the epilogue needs adds more than -that back inside the GEMM, which goes from 39.1us to 156.1us. ``kernels_fused.py`` +from 136.8us to 72.3us -- but the release fence the epilogue needs adds more than +that back inside the GEMM, which goes from 39.1us to 136.9us. ``kernels_fused.py`` has the full per-kernel decomposition. ``--fence`` and ``--stop-after`` exist to reproduce that attribution. """ diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 88e0e5d6a..b96168916 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -73,61 +73,71 @@ ## Result: the fusion works, and the release fence costs more than it saves -Per-kernel, from ``rocprofv3 --kernel-trace`` on 8x MI355X, [4096, 7168] out, -K=1024, median of 11 (times in us): - - config gemm scatter/drain reduce gather sum correct - split 39.1 136.8 59.0 137.3 372.3 yes - fused --fence none 64.5 103.4 59.8 136.6 364.2 NO - fused --fence leader 83.5 96.9 59.9 136.2 376.5 NO - fused --fence agent 156.1 67.3 59.8 137.0 420.1 yes - fused --fence all 175.9 71.1 59.3 137.3 443.6 yes - -The overlap is real and large: the scatter collapses from 136.8us to 67.3us. But -every microsecond of that, and more, reappears inside the GEMM, which goes from -39.1us to 156.1us. Splitting that +117us three ways: - -* **+25us** for the tail itself with no fence at all -- ``s_waitcnt vmcnt(0)``, - the block-wide barrier, the atomic, and the copy engine reading C while the - GEMM is still writing it. -* **+92us** for ``buffer_wbl2 sc1`` and the ``s_waitcnt vmcnt(0)`` that waits on - it. This is the whole problem, and it is **quadratic**: ``buffer_wbl2`` flushes - the entire L2, so block k writes back every tile blocks 1..k-1 already wrote, - 448 blocks over. The split path gets one bulk release at kernel end instead, - for free. -* **+20us** more if the fence is system-scoped rather than agent-scoped. The ISA - says why: ``__threadfence_system`` is bidirectional, so it emits - ``buffer_wbl2 sc0 sc1 ; s_waitcnt ; buffer_inv sc0 sc1`` -- and that - ``buffer_inv`` discards the B weights the GEMM is streaming, 3584 times. - -Three things were tried to make the release cheap, and all three failed: - -* **Agent scope instead of system.** Correct (relL2 2.35e-3, same as split) and - it does drop the ``buffer_inv``, but it only buys 20us of the 112. Worth - keeping as the default anyway -- and note cco exposes no agent-scope release, - only ``cco_system_fence``, so this had to be emitted as a raw ``llvm.fence``. +Per-kernel, ``rocprofv3 --kernel-trace``, 8x MI355X, [4096, 7168] out, K=1024, +median of 11 (us): + + config gemm scatter/drain reduce gather sum correct + split 39.1 136.8 59.0 137.3 372.3 yes + fused --fence none 64.5 103.4 59.8 136.6 364.2 NO + fused --fence leader 83.5 96.9 59.9 136.2 376.5 NO + fused --fence agent 136.9 72.3 59.3 137.7 406.2 yes + fused --fence all 175.9 71.1 59.3 137.3 443.6 yes + +The overlap is real and large: the scatter collapses from 136.8us to 72.3us, a +saving bigger than the entire 39.1us GEMM. But more than that reappears *inside* +the GEMM, which goes from 39.1us to 136.9us. + +An ATT thread trace (``rocprofv3 --att --att-library-path /opt/rocm/lib``) +attributes the epilogue tail exactly -- it is 26% of the kernel's latency, and +essentially all of it is one instruction waiting: + + vaddr cycles % kernel instruction + 19396 91336 5.1% s_waitcnt vmcnt(0) retire this block's C + 19400 22984 1.3% s_barrier + 19404 292 0.0% buffer_wbl2 sc1 <- the op itself is free + 19412 246804 13.9% s_waitcnt vmcnt(0) <- waiting on the writeback + 19608+ 91612 5.1% waits bracketing the counter atomic + +The writeback wait is the single largest line item in the whole kernel, larger +than any main-loop stall, and it is **quadratic**: ``buffer_wbl2`` flushes the +entire L2, so block k waits for every tile blocks 1..k-1 already wrote, 448 +blocks over. The split path gets one bulk release at kernel end instead, free. + +Two costs the trace found that *were* removable, both now fixed: + +* The counter atomic was ``acq_rel``, so it emitted its own ``buffer_wbl2`` / + ``buffer_inv`` pair -- 91.6k cycles, more per wave than the explicit fence. + Both halves were redundant (the caller already released; the winner never + reads what it counts, the copy engine does). Relaxed: 156.1us -> 136.9us. +* System scope emitted ``buffer_inv sc0 sc1`` on top of the writeback, and that + invalidate discards the B weights the GEMM streams, 3584 times. Agent scope + drops it: 175.9us -> 156.1us. cco exposes no agent-scope release, only + ``cco_system_fence``, so it is emitted as a raw ``llvm.fence`` + (``_compat.release_fence``) -- a gap worth closing in the device API. + +Two more attempts that did *not* work, kept selectable so they stay disprovable: + * **One fence per block instead of per wave** (``--fence leader``). Incorrect: relL2 9.6e-3, indistinguishable from no fence at all, even though every lane has already passed ``s_waitcnt vmcnt(0)`` and a block barrier. So the release - really is per-wave, which is where the factor of 8 in 3584 comes from. + is genuinely per-wave -- that is the factor of 8 in 3584. * **Write-through C stores** (``--fence writethrough``), so no writeback would be - needed. Not a release on its own (relL2 1.72e-2), and pairing it with the agent - fence (``--fence wt-agent``, 432.5us) is indistinguishable from the agent fence - alone (432.4us) -- ``buffer_wbl2`` costs the same whether or not the lines it - is asked to flush are dirty. - -So the honest conclusion is narrower than "the ratio is wrong", though the ratio -does cap the prize: what is missing is a **per-tile release**. Every primitive -available here publishes the whole cache, and the GEMM needs to publish 448 -tiles, so the release is O(blocks^2) against an overlap worth at most one GEMM. -Even the incorrect ``--fence none`` configuration only reaches 364.2us against -split's 372.3 -- and both lose to ``gemm + LSA`` at 327.0us, which remains the -number any fused SDMA path has to beat. + needed. Not a release on its own (relL2 1.72e-2), and paired with the agent + fence (``--fence wt-agent``, 432.5us) it is indistinguishable from the agent + fence alone -- ``buffer_wbl2`` costs the same whether or not the lines it is + asked to flush are dirty. + +So the conclusion is narrower and more useful than "the compute:comm ratio is +wrong", though the ratio does cap the prize: what is missing is a **per-tile +release**. Every primitive available publishes the whole cache, and the GEMM has +448 tiles to publish, so the release is O(blocks^2) against an overlap worth at +most one GEMM. Profiling took the fused path from 450.2us to 406-413us, and it +still loses to split-sdma at 372.3 and to ``gemm + LSA`` at 327.0, which remains +the number any fused SDMA path has to beat. The completion-counter machinery and the destination-rotated tile order are both -sound and are what produced the 69.5us of genuine overlap; they would pay -immediately on a transport whose producer-side release is per-tile rather than -per-cache. +sound -- they are what produced the 64.5us of genuine overlap -- and would pay +immediately on a transport whose producer-side release is per-tile. ## Pinned copy From 60996cf770383ee9bab0e7be5462bb6d24e6610c Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 06:22:35 +0000 Subject: [PATCH 07/79] perf(cco): size the SDMA reduce grid independently; 342us -> 294us Profiling the fusion turned up a much larger miss in the plain SDMA all-reduce. Its reduce kernel was inheriting cfg.blocks, i.e. LSA_BLOCK_CAP, which is 24. That cap exists to throttle outstanding *xGMI* requests, and it is right for the LSA kernel. The SDMA reduce reads local HBM. Same constant, opposite regime. At [4096, 7168] it left the kernel with 24 x 512 threads x 8 packs = 1.6MB in flight, about a fifth of what covers HBM latency: 66.1MB moved in 59.3us, or 1.11 TB/s on an ~8 TB/s part. Sized on its own (SDMA_REDUCE_BLOCK_CAP = 256, flat from 256 to 512 in a sweep) it is 13.5us, 4.9 TB/s. standalone SDMA all-reduce, [M, 7168] bf16, best of 2 x 101 iters: M aiter LSA (/aiter) SDMA before SDMA after (/LSA) 64 25.6us 30.7us 1.20 47.4us 45.6us 1.49 512 53.6 66.5 1.24 79.2 73.6 1.11 2048 146.1 160.3 1.10 194.2 167.6 1.05 4096 271.1 285.9 1.05 342.4 294.5 1.03 8192 558.8 531.7 0.95 637.8 544.6 1.02 So the Step 2 conclusion needs correcting: SDMA is not "~1.20x slower than LSA and not meant to win". Past M=2048 it is within 5% of LSA, and at M=8192 it beats aiter. What remains is almost entirely the ~40us dispatch floor, which is why decode sizes still lose ~1.5x. Same fix flows into the fused GEMM path: split-sdma 382.0 -> 334.3us fused-sdma 413.4 -> 365.4us split-lsa 327.0us (unchanged) Also settles, negatively, whether the fused epilogue's release can be per-block instead of per-wave. `--fence agent-leader` emits the fence inline in our own kernel rather than through cco's leaderOnly wrapper, ruling out the compiler having sunk it past the atomic. Still incorrect: relL2 1.36e-2. Two independent implementations agree, so one wave's buffer_wbl2 genuinely does not publish its sibling waves' stores even after s_waitcnt vmcnt(0) and s_barrier. Its *time* is the useful part though -- 322.5us, which would beat split-lsa -- so the 8x reduction in fence count is worth 41us and the prize for a cheaper release is now priced. Retracting one claim from the previous commit: "buffer_wbl2 costs the same whether or not the lines are dirty" was based on a broken experiment. --fence writethrough asks for sc0|sc1 but the ISA shows the store emitted with `sc0` only, so the data stayed dirty in L2 and nothing was tested. The copy atom's cache_modifier is a two-value enum (0=cached, 2=nt), not the aux bitmask that raw_ptr_buffer_store takes, so sc0|sc1 is not expressible through StoreC at all. Whether write-through would help is still open. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/bench_ar.py | 45 +++++++++++++------ benchmark/cco/flydsl/ar/kernels_sdma.py | 36 +++++++++++---- benchmark/cco/flydsl/ar/layout.py | 18 ++++++++ benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 2 +- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 8 +++- 5 files changed, 84 insertions(+), 25 deletions(-) diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py index a600b0f69..dc88c5a19 100644 --- a/benchmark/cco/flydsl/ar/bench_ar.py +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -35,13 +35,11 @@ ``tensor_model_parallel_all_reduce`` with ``set_custom_all_reduce(True)``: M aiter LSA (/aiter) SDMA (/LSA) - 64 25.6us 30.7us 1.20 47.4us 1.54 - 128 27.2 41.8 1.54 51.7 1.24 - 512 53.6 66.5 1.24 79.2 1.19 - 1024 85.6 100.6 1.17 116.4 1.16 - 2048 146.1 160.3 1.10 194.2 1.21 - 4096 271.1 285.9 1.05 342.4 1.20 - 8192 558.8 531.7 0.95 637.8 1.20 + 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 @@ -49,9 +47,10 @@ 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 is ~18% slower than LSA even after removing the floor, -so it does not win on bandwidth and is not meant to -- it wins by leaving the CUs -free, which only pays off once a GEMM is running underneath it. +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, @@ -160,6 +159,7 @@ def run(args) -> int: 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, @@ -195,11 +195,21 @@ def run(args) -> int: cfg, rank, force_stage=args.force_stage, gather=args.gather ) elif args.backend == "sdma": - from kernels_sdma import build_sdma_ar - - launch, stage = build_sdma_ar( - cfg, rank, queues=args.sdma_queues, signal=args.sdma_signal + from kernels_sdma import build_sdma_phases + + parts = build_sdma_phases( + cfg, + rank, + queues=args.sdma_queues, + signal=args.sdma_signal, + reduce_blocks=args.reduce_blocks, ) + + def launch(dc_ptr, win_h, stream=None): + for k in ("scatter", "reduce", "gather"): + parts[k](dc_ptr, win_h, stream=stream) + + stage = 2 else: raise ValueError(f"unknown backend {args.backend!r}") @@ -336,6 +346,13 @@ def build_parser() -> argparse.ArgumentParser: 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", diff --git a/benchmark/cco/flydsl/ar/kernels_sdma.py b/benchmark/cco/flydsl/ar/kernels_sdma.py index c4bd42611..44450c105 100644 --- a/benchmark/cco/flydsl/ar/kernels_sdma.py +++ b/benchmark/cco/flydsl/ar/kernels_sdma.py @@ -56,12 +56,21 @@ slice we push is a single packet. Measured on 8x MI355X, [M, 7168] bf16, against the LSA kernel (full table in -``bench_ar.py``): ~1.20x slower across M >= 128, of which ~12us is a fixed floor -(two transfer phases x the engine's ~6us dispatch cost) and the rest is ~18% -lower sustained bandwidth than CU-issued loads over the same links. Both are -expected. This backend is not here to be faster -- it is here because the -transfers occupy one warp of one block instead of the whole grid, which is what -makes an in-GEMM epilogue possible in step 3. +``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 @@ -100,7 +109,9 @@ PUSH_THREADS = 64 -def build_sdma_phases(cfg, rank: int, *, queues: int = 8, signal: bool = False): +def build_sdma_phases( + cfg, rank: int, *, queues: int = 8, signal: bool = False, reduce_blocks=None +): """Compile the phases separately, so the fused GEMM can reuse the tail. Returns ``{"scatter", "drain", "reduce", "gather"}``, each a launcher taking @@ -132,6 +143,13 @@ def build_sdma_phases(cfg, rank: int, *, queues: int = 8, signal: bool = False): ) threads, blocks = cfg.threads, cfg.blocks + # 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 stride_packs = blocks * threads part = cfg.packs_per_rank @@ -243,7 +261,7 @@ def sdma_reduce(dev_comm: Int64, win: Int64): ) gtid = bid * threads + tid - for pk in range(gtid, part, stride_packs): + for pk in range(gtid, part, red_stride): i32_off = pk * I32_PER_PACK acc = None for j in range_constexpr(ws): @@ -287,7 +305,7 @@ def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): return { "scatter": _phase(scatter, 1, PUSH_THREADS), "drain": _phase(drain, 1, PUSH_THREADS), - "reduce": _phase(sdma_reduce, blocks, threads), + "reduce": _phase(sdma_reduce, red_blocks, threads), "gather": _phase(gather, 1, PUSH_THREADS), } diff --git a/benchmark/cco/flydsl/ar/layout.py b/benchmark/cco/flydsl/ar/layout.py index ceeea4e01..bb4d4fd48 100644 --- a/benchmark/cco/flydsl/ar/layout.py +++ b/benchmark/cco/flydsl/ar/layout.py @@ -64,6 +64,12 @@ # 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 @@ -117,6 +123,10 @@ class ArConfig: #: 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 + #: 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. @@ -211,6 +221,13 @@ def end_off(self) -> int: 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_off(self) -> int: """Monotonic tile counters for the fused GEMM, one per (dest, chunk). @@ -328,6 +345,7 @@ def validate(self) -> None: "select_stage", "K_MAX_BLOCKS", "LSA_BLOCK_CAP", + "SDMA_REDUCE_BLOCK_CAP", "THREADS", "MAX_WORLD", "PACK_BYTES", diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 4430f8004..5457c7085 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -328,7 +328,7 @@ def build_parser() -> argparse.ArgumentParser: ) p.add_argument( "--fence", - choices=("all", "agent", "leader", "none", "writethrough", "wt-agent"), + choices=("all", "agent", "agent-leader", "leader", "none", "writethrough", "wt-agent"), default="agent", help="release before the epilogue push. Correct: 'agent' (default, " "cheapest) and 'all' (system scope, adds a buffer_inv that costs 20us). " diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index b96168916..37fd88dde 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -297,7 +297,7 @@ def compile_fused_gemm_scatter( tiles_per_chunk = m_tiles_per_chunk * n_blocks_const chunk_bytes = cfg.slice_bytes // chunks - if fence not in ("all", "agent", "leader", "none", "writethrough", "wt-agent"): + if fence not in ("all", "agent", "agent-leader", "leader", "none", "writethrough", "wt-agent"): raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) @@ -542,6 +542,12 @@ def kernel_gemm_scatter( w = cco.Window(win) 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 == "agent-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 From b4ed569a7487f2f395d369743a623aab3911afe1 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 06:26:45 +0000 Subject: [PATCH 08/79] perf(cco): price the fused epilogue's fence in isolation (--no-put, nt stores) Adds two probes that separate the three things the fused epilogue costs, which previously could only be inferred from differences of totals. --no-put keeps the counter protocol and the fence but drops the transfer nt-agent non-temporal C stores + agent fence At [4096, 7168], K=1024, 8 ranks: none --no-put 206.7us epilogue without a fence, no transfer agent --no-put 284.9us +78.2us <- the release fence, alone agent 364.3us +79.4us <- the transfers nt-agent 411.2us Three results: 1. The barrier and atomic are nearly free. A GEMM with no epilogue at all is ~197us in this pipeline, so the counter protocol costs ~9us, not the ~25us I had inferred. Most of what I attributed to "the tail itself" was actually the copy engine reading C while the GEMM writes it. 2. The release fence is 78.2us, now measured with no SDMA in the picture at all rather than derived from a difference of totals. 3. The overlap is 57.4us: the transfers add 79.4us on top of the fenced no-put run, against 136.8us for the standalone scatter kernel. Which prices the whole thing: 334.3 (split) - 57.4 (overlap) + 9 (protocol) = ~284us if the release were free, against split-lsa's 327. That is a 13% win, materially better than the "~10-20us" I estimated from totals, and it makes the StoreC coalescing rewrite worth costing out rather than dismissing. nt-agent is a negative result and stays in as one: non-temporal C stores make it *worse* (411.2 vs 364.3). They do not make buffer_wbl2 cheaper, and they cost the GEMM's stores. The only untried way to avoid the writeback is real write-through (sc0|sc1), which the copy atom cannot express -- it needs a raw raw_ptr_buffer_store, and probably needs StoreC coalesced first, since its per-lane 2-byte scatter bypassing L2 would write 32B partial lines to HBM. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 13 +++++++- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 30 ++++++++++++++++--- 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 5457c7085..7913065e1 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -197,6 +197,7 @@ def run(args) -> int: fuse=fused, rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", fence=args.fence, + emit_put=not args.no_put, ) a_i8 = a.contiguous().view(torch.int8).view(-1) b_i8 = b_shuf.contiguous().view(torch.int8).view(-1) @@ -328,7 +329,10 @@ def build_parser() -> argparse.ArgumentParser: ) p.add_argument( "--fence", - choices=("all", "agent", "agent-leader", "leader", "none", "writethrough", "wt-agent"), + choices=( + "all", "agent", "agent-leader", "nt-agent", "leader", "none", + "writethrough", "wt-agent", + ), default="agent", help="release before the epilogue push. Correct: 'agent' (default, " "cheapest) and 'all' (system scope, adds a buffer_inv that costs 20us). " @@ -343,6 +347,13 @@ def build_parser() -> argparse.ArgumentParser: "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("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 37fd88dde..6e643c3d7 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -195,6 +195,20 @@ ) +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. @@ -229,6 +243,7 @@ def compile_fused_gemm_scatter( fuse: bool = True, rotated: bool | None = None, fence: str = "agent", + emit_put: bool = True, ): """Compile the GEMM, with (``fuse=True``) or without the scatter epilogue. @@ -297,7 +312,10 @@ def compile_fused_gemm_scatter( tiles_per_chunk = m_tiles_per_chunk * n_blocks_const chunk_bytes = cfg.slice_bytes // chunks - if fence not in ("all", "agent", "agent-leader", "leader", "none", "writethrough", "wt-agent"): + if fence not in ( + "all", "agent", "agent-leader", "nt-agent", "leader", "none", + "writethrough", "wt-agent", + ): raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) @@ -309,7 +327,7 @@ def compile_fused_gemm_scatter( _kname = ( f"mori_fused_{'sdma' if fuse else 'split'}_8w_" - f"{BLOCK_M}x{BLOCK_N}x{BLOCK_K}_k{K}_{_kname_tag}_r{rank}" + f"{BLOCK_M}x{BLOCK_N}x{BLOCK_K}_k{K}_{_kname_tag}{'p' if emit_put else 'x'}_r{rank}" ) @fx.struct @@ -403,6 +421,10 @@ def kernel_gemm_scatter( store_c = _WriteThroughStoreC( A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B ) + 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 @@ -534,7 +556,7 @@ def kernel_gemm_scatter( 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")): + elif const_expr(fence in ("agent", "wt-agent", "nt-agent")): release_fence("agent") # "writethrough": nothing to do -- the stores already went to memory, # and wait_barrier(0) above retired them. @@ -563,7 +585,7 @@ def kernel_gemm_scatter( # 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 dest != fx.Int32(rank): + 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 From c73efe160522043294abce6dbc77b99806577d70 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 07:06:21 +0000 Subject: [PATCH 09/79] fix(cco): the fused GEMM+AR had two races; correct the record and the default Running the X1 precursor experiment turned up something more important than the experiment: the fused path is intermittently wrong, and has been. Every "validated=True" I reported for it was a single sample. Repeating the runs: config passes note chunks=1, --fence agent, acq_rel atomic 6/6 the only stable one chunks=2, --fence agent, acq_rel atomic 5/6 RACY chunks=2, --fence agent, monotonic atomic 3/5 RACY, worse chunks=1, --fence raw-wt-agent 2/4 RACY standalone SDMA all-reduce (control) 5/5 stable Two races, both mine: 1. chunks > 1 shares an SDMA queue. The modulo test elects one block per *chunk*, not one per *queue*, so two winners post to queue `dest` at once. The plan called for a per-destination submit lock and I skipped it, arguing the counter already elects a single CTA -- true, and about the wrong thing. Default is now chunks=1; >1 stays settable to reproduce it. 2. The counter atomic needed its acquire half. I made it relaxed two commits ago for a 19us win, reasoning the winner never reads what it counts. Wrong: the winner needs the acquire so the other blocks' releases are ordered before *its put*. acq_rel is the default again, and it measurably lowers the failure rate (5/6 vs 3/5) though it does not fix race 1. The X1 answer is no. Write-through C stores looked like the fix -- they make buffer_wbl2 cheap because there is nothing dirty to drain, GEMM 139.1 -> 97.6us and total 333.8 against split-lsa's 327 -- and they are racy at chunks=1 too, so that number describes a kernel that does not reliably compute the right answer. Not worth building the StoreC coalescing rewrite on. Getting there needed two ISA checks, because both earlier write-through attempts silently did nothing: - fx.rocdl.BufferCopy*(cache_modifier) is a two-value enum (0=cached, 2=nt), not an aux bitmask, so sc0|sc1 cannot be asked for through the copy atom. - _compat's CM_SC1 was 4. Assembling each policy and diffing the encoding (llvm-mc, gfx950) puts SC0 at insn bit 14, SC1 at bit 15, NT at bit 17 = CPol 1, 16, 2. So SC1 is 16; 4 is DLC, RDNA-only, silently ignored on CDNA. Latent -- the barrier uses atomics and payload stores use CM_CACHED -- but it invalidated two experiments before showing up in a disassembly. One genuinely useful sub-result survives: the un-coalesced per-lane 2-byte C store does *not* collapse when it bypasses L2. That was the open question about whether StoreC must be coalesced first, and the answer is no. Corrected standings, correct configurations only: split-lsa 324.4us split-sdma 338.0 fused-sdma chunks=1, agent 440.0 Adds test_fused_is_stable_across_repeats, which runs the fused path three times and requires the results to be *identical*, not merely each small. It skips below 8 GPUs, since the queue race needs more than one chunk per destination and the wo_b shape only has one at 2 ranks. A one-shot check would not have caught either race. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_compat.py | 40 +++- benchmark/cco/flydsl/ar/kernels_lsa.py | 2 - benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 42 ++-- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 205 ++++++++++++------ tests/python/cco/test_gemm_ar.py | 27 +++ 5 files changed, 221 insertions(+), 95 deletions(-) diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py index 9f81157b4..05a1b7c8c 100644 --- a/benchmark/cco/flydsl/ar/_compat.py +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -49,21 +49,37 @@ FLYDSL_VERSION = getattr(flydsl, "__version__", "unknown") -# Cache-modifier (``aux``) bits. **These are CDNA3/CDNA4-dependent.** +# 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. # -# On gfx950 the encoding is bit0=SC0, bit1=NT, bit2=SC1, so the gfx942 spelling -# used by ``examples/cco/python/05_flydsl_lsa_allreduce`` (SC1=2, SC0|SC1=3) -# silently compiles to ``nt`` and ``sc0 nt`` here -- non-temporal hints that do -# **not** bypass L2. Verified by reading the emitted ISA: aux=3 produced -# ``buffer_store_dwordx4 ... sc0 nt``, aux=5 produces ``... sc0 sc1``. +# Established by assembling each policy and diffing the encoding:: # -# The examples get away with it because their inputs are host-written before the -# kernel, so nothing there actually depends on the barrier ordering. +# 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 -CM_NT = 2 # non-temporal -CM_SC1 = 4 # bypass L2 -- required to observe a peer's fresh write -CM_SC0_SC1 = CM_SC0 | CM_SC1 # 5: publish past L1+L2 +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. diff --git a/benchmark/cco/flydsl/ar/kernels_lsa.py b/benchmark/cco/flydsl/ar/kernels_lsa.py index 18b14aa41..8da12a041 100644 --- a/benchmark/cco/flydsl/ar/kernels_lsa.py +++ b/benchmark/cco/flydsl/ar/kernels_lsa.py @@ -77,8 +77,6 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from _compat import ( # noqa: E402 CM_CACHED, - CM_SC0_SC1, - CM_SC1, buffer_load, buffer_store, create_buffer_resource_from_addr, diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 7913065e1..d5b0b891f 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -41,16 +41,14 @@ Headline, 8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51: - gemm-only 49.6us - split-lsa 327.0 - split-sdma 382.0 - fused-sdma 413.4 (chunks=2, --fence agent) - -Fusing loses. The overlap is real -- a kernel trace shows the scatter dropping -from 136.8us to 72.3us -- but the release fence the epilogue needs adds more than -that back inside the GEMM, which goes from 39.1us to 136.9us. ``kernels_fused.py`` -has the full per-kernel decomposition. ``--fence`` and ``--stop-after`` exist to -reproduce that attribution. + split-lsa 324.4us + split-sdma 338.0 + fused-sdma chunks=1, agent 440.0 + +Fusing loses. Read ``kernels_fused.py`` before trusting any faster fused number +from this benchmark: several of its options are intermittently wrong, and a +single passing run proves nothing. ``--chunks`` > 1 and every ``raw-wt`` fence +mode are known-racy and kept only to reproduce that. """ from __future__ import annotations @@ -149,8 +147,12 @@ def run(args) -> int: # One push per BLOCK_M row-band by default: the finest granularity the tile # order can signal, and at the wo_b shape still a 3.7MB transfer, well past # the ~1MB where SDMA bandwidth flattens out. - m_tiles_per_peer = max(1, (args.m // world_size) // args.block_m) - chunks = args.chunks if args.chunks else m_tiles_per_peer + # Default 1, NOT one push per row-band. With more than one chunk per + # destination, two winning blocks post to the same SDMA queue concurrently + # and the result is intermittently wrong -- see the race note in + # kernels_fused. Fixing it needs the per-destination submit lock that the + # gcnasm protocol has and this does not. + chunks = args.chunks if args.chunks else 1 cfg = ArConfig( world_size=world_size, m=args.m, @@ -198,6 +200,7 @@ def run(args) -> int: rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", 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) @@ -318,8 +321,8 @@ def build_parser() -> argparse.ArgumentParser: "--chunks", type=int, default=0, - help="pushes per destination (0 = one per BLOCK_M row-band, the finest " - "the tile order can signal)", + help="pushes per destination (0 = 1). Values > 1 are RACY: two winners " + "then share one SDMA queue. Kept settable only to reproduce that", ) p.add_argument( "--tile-order", @@ -331,7 +334,8 @@ def build_parser() -> argparse.ArgumentParser: "--fence", choices=( "all", "agent", "agent-leader", "nt-agent", "leader", "none", - "writethrough", "wt-agent", + "writethrough", "wt-agent", "raw-wt", "raw-wt-agent", + "raw-wt-leader", ), default="agent", help="release before the epilogue push. Correct: 'agent' (default, " @@ -354,6 +358,14 @@ def build_parser() -> argparse.ArgumentParser: "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", + ) p.add_argument("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 6e643c3d7..ff45d01ae 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,73 +71,97 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## Result: the fusion works, and the release fence costs more than it saves - -Per-kernel, ``rocprofv3 --kernel-trace``, 8x MI355X, [4096, 7168] out, K=1024, -median of 11 (us): - - config gemm scatter/drain reduce gather sum correct - split 39.1 136.8 59.0 137.3 372.3 yes - fused --fence none 64.5 103.4 59.8 136.6 364.2 NO - fused --fence leader 83.5 96.9 59.9 136.2 376.5 NO - fused --fence agent 136.9 72.3 59.3 137.7 406.2 yes - fused --fence all 175.9 71.1 59.3 137.3 443.6 yes - -The overlap is real and large: the scatter collapses from 136.8us to 72.3us, a -saving bigger than the entire 39.1us GEMM. But more than that reappears *inside* -the GEMM, which goes from 39.1us to 136.9us. - -An ATT thread trace (``rocprofv3 --att --att-library-path /opt/rocm/lib``) -attributes the epilogue tail exactly -- it is 26% of the kernel's latency, and -essentially all of it is one instruction waiting: +## Correctness first: two races, and what is actually demonstrable + +Repeated runs matter more than any timing here. Validating once and reporting the +number was how several wrong results in this file's history got published; the +failure mode is intermittent and every individual relL2 looks "small". + +Measured over repeated 8-rank runs at [4096, 7168], K=1024: + + config passes note + chunks=1, --fence agent, acq_rel atomic 6/6 the only stable one + chunks=2, --fence agent, acq_rel atomic 5/6 RACY + chunks=2, --fence agent, monotonic atomic 3/5 RACY, worse + chunks=1, --fence raw-wt-agent 2/4 RACY + standalone SDMA all-reduce (control) 5/5 stable + +Two independent races, both mine: + +1. **``chunks > 1`` shares an SDMA queue.** The modulo test elects exactly one + block per *chunk*, not one per *queue*, so with two chunks per destination two + winning blocks post to queue ``dest`` at the same time. The plan called for a + per-destination submit lock (gcnasm has one) and this skipped it on the + argument that the counter already elects a single CTA. That argument was about + the wrong thing. Default is now 1; ``--chunks 2`` reproduces the race. + +2. **The counter atomic needs its acquire half.** It was changed to relaxed for a + 19us win on the reasoning that the winner never reads the data it counts (the + copy engine does). Wrong: the winner needs the acquire so the other blocks' + releases are ordered before *its put*. ``acq_rel`` is the default again, and it + measurably reduces the failure rate. + +Write-through C stores (``--fence raw-wt``, ``raw-wt-agent``) are kept only as +disproved options. The idea was to make the per-block ``buffer_wbl2`` unnecessary +by having C reach memory at the store. It looked excellent -- GEMM 139.1us -> +97.6us, total 333.8 against split-lsa's 327 -- and it is racy even at chunks=1, +so those numbers describe a kernel that does not reliably compute the right +answer. Note the sub-experiment was worth running anyway: it showed the +un-coalesced per-lane 2-byte store does **not** collapse when it bypasses L2, +which was the open question about whether ``StoreC`` has to be coalesced first. + +## Result: the fusion does not pay + +Correct configurations only, 8 ranks, [4096, 7168], K=1024, graph replay: + + split-lsa 324.4us + split-sdma 338.0us + fused-sdma chunks=1, agent 440.0us + +Per-kernel for the racy chunks=2 variant, which is the best case the structure +can aspire to if the queue lock were built (rocprofv3 --kernel-trace, us): + + config gemm scatter/drain reduce gather sum + split 39.1 136.8 13.5 137.3 326.7 + fused --fence none 64.5 103.4 13.5 136.6 318.0 + fused --fence agent 139.1 66.3 13.5 136.5 355.4 + +The overlap is real -- the scatter collapses from 136.8us to 66.3us -- and it is +still not enough, because the release fence adds more than that back inside the +GEMM. An ATT thread trace puts the epilogue tail at 26% of the kernel's latency, +almost all of it one instruction: vaddr cycles % kernel instruction 19396 91336 5.1% s_waitcnt vmcnt(0) retire this block's C 19400 22984 1.3% s_barrier 19404 292 0.0% buffer_wbl2 sc1 <- the op itself is free 19412 246804 13.9% s_waitcnt vmcnt(0) <- waiting on the writeback - 19608+ 91612 5.1% waits bracketing the counter atomic - -The writeback wait is the single largest line item in the whole kernel, larger -than any main-loop stall, and it is **quadratic**: ``buffer_wbl2`` flushes the -entire L2, so block k waits for every tile blocks 1..k-1 already wrote, 448 -blocks over. The split path gets one bulk release at kernel end instead, free. - -Two costs the trace found that *were* removable, both now fixed: - -* The counter atomic was ``acq_rel``, so it emitted its own ``buffer_wbl2`` / - ``buffer_inv`` pair -- 91.6k cycles, more per wave than the explicit fence. - Both halves were redundant (the caller already released; the winner never - reads what it counts, the copy engine does). Relaxed: 156.1us -> 136.9us. -* System scope emitted ``buffer_inv sc0 sc1`` on top of the writeback, and that - invalidate discards the B weights the GEMM streams, 3584 times. Agent scope - drops it: 175.9us -> 156.1us. cco exposes no agent-scope release, only - ``cco_system_fence``, so it is emitted as a raw ``llvm.fence`` - (``_compat.release_fence``) -- a gap worth closing in the device API. - -Two more attempts that did *not* work, kept selectable so they stay disprovable: - -* **One fence per block instead of per wave** (``--fence leader``). Incorrect: - relL2 9.6e-3, indistinguishable from no fence at all, even though every lane - has already passed ``s_waitcnt vmcnt(0)`` and a block barrier. So the release - is genuinely per-wave -- that is the factor of 8 in 3584. -* **Write-through C stores** (``--fence writethrough``), so no writeback would be - needed. Not a release on its own (relL2 1.72e-2), and paired with the agent - fence (``--fence wt-agent``, 432.5us) it is indistinguishable from the agent - fence alone -- ``buffer_wbl2`` costs the same whether or not the lines it is - asked to flush are dirty. - -So the conclusion is narrower and more useful than "the compute:comm ratio is -wrong", though the ratio does cap the prize: what is missing is a **per-tile -release**. Every primitive available publishes the whole cache, and the GEMM has -448 tiles to publish, so the release is O(blocks^2) against an overlap worth at -most one GEMM. Profiling took the fused path from 450.2us to 406-413us, and it -still loses to split-sdma at 372.3 and to ``gemm + LSA`` at 327.0, which remains -the number any fused SDMA path has to beat. - -The completion-counter machinery and the destination-rotated tile order are both -sound -- they are what produced the 64.5us of genuine overlap -- and would pay -immediately on a transport whose producer-side release is per-tile. + +That wait is the largest single item in the whole kernel, larger than any +main-loop stall, and it is **quadratic**: ``buffer_wbl2`` flushes the entire L2, +so block k waits for every tile blocks 1..k-1 already wrote, 448 blocks over. The +split path gets one bulk release at kernel end, free. + +Everything tried to make it cheaper, and what happened: + +* **Agent scope instead of system.** Correct and worth 20us: ``__threadfence_system`` + is bidirectional and emits ``buffer_wbl2 sc0 sc1 ; s_waitcnt ; buffer_inv sc0 + sc1``, and that ``buffer_inv`` discards the B weights the GEMM streams, 3584 + times. cco exposes no agent-scope release, only ``cco_system_fence``, so it is + emitted as a raw ``llvm.fence`` (``_compat.release_fence``). Kept. +* **One fence per block instead of per wave** (``--fence agent-leader``, and cco's + ``leaderOnly``). Incorrect in both spellings, including emitted inline in this + kernel to rule out the compiler sinking it. So the release is genuinely + per-wave. Its *time* is the useful part: 322.5us, i.e. an 8x cut in fence count + would be worth 41us and would win -- if it were legal. +* **Non-temporal C stores** (``--fence nt-agent``). Worse, 411.2us. +* **Write-through C stores.** Racy, see above. + +So what is missing is a **per-tile release**. Every primitive here publishes the +whole cache, and the GEMM has 448 tiles to publish, so the cost is O(blocks^2) +against an overlap worth at most one GEMM. The completion-counter machinery and +the destination-rotated tile order are sound and produced the real overlap; they +would pay immediately on a transport whose producer-side release is per-tile. ## Pinned copy @@ -190,11 +214,41 @@ from _compat import ( # noqa: E402 CM_SC0_SC1, atomic_add_u32, + buffer_store, + create_buffer_resource_from_addr, release_fence, signal_ptr, ) +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 _NonTemporalStoreC(StoreC): """``StoreC`` with non-temporal C stores (the copy atom's only other mode). @@ -244,6 +298,7 @@ def compile_fused_gemm_scatter( rotated: bool | None = None, fence: str = "agent", emit_put: bool = True, + atomic_order: str = "acq_rel", ): """Compile the GEMM, with (``fuse=True``) or without the scatter epilogue. @@ -314,7 +369,7 @@ def compile_fused_gemm_scatter( if fence not in ( "all", "agent", "agent-leader", "nt-agent", "leader", "none", - "writethrough", "wt-agent", + "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}" @@ -327,7 +382,8 @@ def compile_fused_gemm_scatter( _kname = ( f"mori_fused_{'sdma' if fuse else 'split'}_8w_" - f"{BLOCK_M}x{BLOCK_N}x{BLOCK_K}_k{K}_{_kname_tag}{'p' if emit_put else 'x'}_r{rank}" + 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 @@ -421,6 +477,21 @@ def kernel_gemm_scatter( 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(cco.Window(win).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 @@ -556,7 +627,7 @@ def kernel_gemm_scatter( 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")): + 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. @@ -568,7 +639,7 @@ def kernel_gemm_scatter( # 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 == "agent-leader"): + 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( @@ -579,7 +650,9 @@ def kernel_gemm_scatter( fx.Int64(w.lsa_ptr(rank, counter_off)) + fx.Int64(slot) * fx.Int64(4) ) - seq = fx.Int32(atomic_add_u32(ctr, 1)) + fx.Int32(1) + 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 diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 59de2c660..01951bc1b 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -212,6 +212,33 @@ def _run_bench(world_size, mode, m, n, k, extra=()): 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("aiter", reason="aiter not importable (set PYTHONPATH)") + 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. From d6fb537afa16315b8ec328d24bf99c344ec10c61 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 07:51:08 +0000 Subject: [PATCH 10/79] perf(cco): the fused epilogue never needed a release fence; 440us -> 348us From reading gcnasm's opus_gemm_a2a_lsa, which fuses a GEMM with an a2a and gets 16-27% out of it. Its ChunkFused epilogue (gemm_a16w16_quad_subtile_kernel_template.hpp:701-741) uses **no cache fence at all** -- 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 correct and I had missed it: the acq_rel counter *is* the release. Its ISA already emits a buffer_wbl2/buffer_inv pair, on thread 0 only -- once per block, not once per wave. The explicit per-wave fence I had added on top was redundant, and it was 90us: fused chunks=1, agent fence 440.0us gemm 152.4 fused chunks=1, no fence 348.5us gemm 62.7 stable 6/6 The earlier conclusion that "the release fence is mandatory and costs 78us" was wrong. --fence none had only ever been measured at chunks=2, where a separate bug corrupts the result, so the fence was being credited for fixing something it had nothing to do with. Also ports gcnasm's per-destination submit lock (atomic test-and-set with s_sleep backoff, ISA-verified) and its post-submit barrier. Neither fixes chunks>1, which stays racy in roughly 1 run in 3 and is now documented as unexplained after ruling out, by measurement: the lock, the barrier, the fence at both scopes, and the atomic ordering. The only factor that separates working from broken is whether the put is issued mid-GEMM or at the end. Default stays chunks=1; note that makes the overlap ~zero (scatter 129.0 vs split's 136.8), which is why fused still loses. Corrected standings: split-lsa 324.4us split-sdma 338.0 fused-sdma chunks=1, no fence 348.5 With the fence gone the entire remaining deficit is +23.6us inside the GEMM -- barrier, counter, and the copy engine reading C while the GEMM writes it -- against a 7.8us overlap gain. Four other things from that README worth having on record here: its collective is one phase where an all-reduce is three, so fusion can only ever touch a third of the traffic; its compute:comm is 2.6:1 where wo_b is 0.29:1; its best mode (Direct LSA) has the epilogue store straight into the peer, so it has no mid-kernel publication problem at all; and it also found chunked PUTs gave no end-to-end benefit on the non-fused path. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_compat.py | 24 +++ benchmark/cco/flydsl/ar/layout.py | 16 +- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 24 +-- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 204 ++++++++++-------- 4 files changed, 164 insertions(+), 104 deletions(-) diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py index 05a1b7c8c..c6fc16b9c 100644 --- a/benchmark/cco/flydsl/ar/_compat.py +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -298,6 +298,28 @@ def atomic_add_u32(ptr, value, *, ordering="monotonic"): ) +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. @@ -328,6 +350,8 @@ def i32_type(): "local_store_u32", "atomic_add_u32", "release_fence", + "atomic_xchg_u32", + "atomic_store_u32", "i32_type", "CM_CACHED", "CM_SC1", diff --git a/benchmark/cco/flydsl/ar/layout.py b/benchmark/cco/flydsl/ar/layout.py index bb4d4fd48..52c3d78c9 100644 --- a/benchmark/cco/flydsl/ar/layout.py +++ b/benchmark/cco/flydsl/ar/layout.py @@ -207,7 +207,7 @@ def blocks(self) -> int: @property def signal_bytes(self) -> int: - return self.counter_off + self.counter_bytes + return self.lock_off + self.lock_bytes @property def start_off(self) -> int: @@ -242,6 +242,20 @@ def counter_off(self) -> int: def counter_bytes(self) -> int: return _align_up(self.world_size * self.counter_chunks * 4, 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. + """ + return self.counter_off + self.counter_bytes + + @property + def lock_bytes(self) -> int: + return _align_up(self.world_size * 4, 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})") diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index d5b0b891f..82c5070bb 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -41,9 +41,9 @@ Headline, 8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51: - split-lsa 324.4us - split-sdma 338.0 - fused-sdma chunks=1, agent 440.0 + split-lsa 324.4us + split-sdma 338.0 + fused-sdma chunks=1, no fence 348.5 Fusing loses. Read ``kernels_fused.py`` before trusting any faster fused number from this benchmark: several of its options are intermittently wrong, and a @@ -144,14 +144,10 @@ def run(args) -> int: local_rank, rank, world_size, uid = _setup_distributed() fused = args.mode == "fused-sdma" needs_sdma = args.mode in ("split-sdma", "fused-sdma") - # One push per BLOCK_M row-band by default: the finest granularity the tile - # order can signal, and at the wo_b shape still a 3.7MB transfer, well past - # the ~1MB where SDMA bandwidth flattens out. - # Default 1, NOT one push per row-band. With more than one chunk per - # destination, two winning blocks post to the same SDMA queue concurrently - # and the result is intermittently wrong -- see the race note in - # kernels_fused. Fixing it needs the per-destination submit lock that the - # gcnasm protocol has and this does not. + # 1, even though >1 is what would create the overlap. A chunk boundary means + # the copy engine starts reading C *while the GEMM is still running*, and no + # combination of submit lock, post-submit barrier, release fence or atomic + # ordering made that reliable here -- see the race note in kernels_fused. chunks = args.chunks if args.chunks else 1 cfg = ArConfig( world_size=world_size, @@ -321,8 +317,8 @@ def build_parser() -> argparse.ArgumentParser: "--chunks", type=int, default=0, - help="pushes per destination (0 = 1). Values > 1 are RACY: two winners " - "then share one SDMA queue. Kept settable only to reproduce that", + help="pushes per destination (0 = one per BLOCK_M row-band). >1 is what " + "produces the overlap, and requires the submit lock", ) p.add_argument( "--tile-order", @@ -337,7 +333,7 @@ def build_parser() -> argparse.ArgumentParser: "writethrough", "wt-agent", "raw-wt", "raw-wt-agent", "raw-wt-leader", ), - default="agent", + default="none", help="release before the epilogue push. Correct: 'agent' (default, " "cheapest) and 'all' (system scope, adds a buffer_inv that costs 20us). " "Incorrect, and present only to price the fence: 'leader', 'none', " diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index ff45d01ae..2ff6a559d 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,97 +71,78 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## Correctness first: two races, and what is actually demonstrable - -Repeated runs matter more than any timing here. Validating once and reporting the -number was how several wrong results in this file's history got published; the -failure mode is intermittent and every individual relL2 looks "small". - -Measured over repeated 8-rank runs at [4096, 7168], K=1024: - - config passes note - chunks=1, --fence agent, acq_rel atomic 6/6 the only stable one - chunks=2, --fence agent, acq_rel atomic 5/6 RACY - chunks=2, --fence agent, monotonic atomic 3/5 RACY, worse - chunks=1, --fence raw-wt-agent 2/4 RACY - standalone SDMA all-reduce (control) 5/5 stable - -Two independent races, both mine: - -1. **``chunks > 1`` shares an SDMA queue.** The modulo test elects exactly one - block per *chunk*, not one per *queue*, so with two chunks per destination two - winning blocks post to queue ``dest`` at the same time. The plan called for a - per-destination submit lock (gcnasm has one) and this skipped it on the - argument that the counter already elects a single CTA. That argument was about - the wrong thing. Default is now 1; ``--chunks 2`` reproduces the race. - -2. **The counter atomic needs its acquire half.** It was changed to relaxed for a - 19us win on the reasoning that the winner never reads the data it counts (the - copy engine does). Wrong: the winner needs the acquire so the other blocks' - releases are ordered before *its put*. ``acq_rel`` is the default again, and it - measurably reduces the failure rate. - -Write-through C stores (``--fence raw-wt``, ``raw-wt-agent``) are kept only as -disproved options. The idea was to make the per-block ``buffer_wbl2`` unnecessary -by having C reach memory at the store. It looked excellent -- GEMM 139.1us -> -97.6us, total 333.8 against split-lsa's 327 -- and it is racy even at chunks=1, -so those numbers describe a kernel that does not reliably compute the right -answer. Note the sub-experiment was worth running anyway: it showed the -un-coalesced per-lane 2-byte store does **not** collapse when it bypasses L2, -which was the open question about whether ``StoreC`` has to be coalesced first. +## 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 chunking race, unexplained + +Overlap requires more than one chunk per destination: with one, the counter only +fires when the whole slice is done, which under the rotated tile order is the end +of the GEMM, so nothing overlaps (measured: scatter 129.0us against split's +136.8 -- a 7.8us gain, i.e. none). + +``--chunks 2`` is worth it on paper -- 328us, which would beat split-sdma's 338 +and tie split-lsa's 324 -- and it is **intermittently wrong**, in roughly 1 run in +3. Ruled out, each by measurement rather than argument: + +* the per-destination submit lock (gcnasm's, ISA-verified: test-and-set, + ``s_sleep`` backoff, correct spin) -- still 2/6 wrong; +* the post-submit barrier gcnasm keeps -- still 4/6 wrong; +* the release fence, at both agent and system scope -- still 3/6 wrong; +* the counter atomic's ordering, ``acq_rel`` vs relaxed -- helps, does not fix. + +The one factor that separates working from broken is *when* the put is issued. At +chunks=1 it lands at the very end of the GEMM; at chunks=2 the first one is +issued mid-kernel and the engine reads C while the rest of the GEMM is still +running. gcnasm's ChunkFused path does the same thing successfully, but its GEMM +is 13x longer, so its chunk boundaries are far apart in time. Whether chunks=1 +here is *correct* or merely always-lucky is not established -- the mechanism +suggests the latter, so treat the fused path as unproven either way. ## Result: the fusion does not pay -Correct configurations only, 8 ranks, [4096, 7168], K=1024, graph replay: - - split-lsa 324.4us - split-sdma 338.0us - fused-sdma chunks=1, agent 440.0us - -Per-kernel for the racy chunks=2 variant, which is the best case the structure -can aspire to if the queue lock were built (rocprofv3 --kernel-trace, us): - - config gemm scatter/drain reduce gather sum - split 39.1 136.8 13.5 137.3 326.7 - fused --fence none 64.5 103.4 13.5 136.6 318.0 - fused --fence agent 139.1 66.3 13.5 136.5 355.4 - -The overlap is real -- the scatter collapses from 136.8us to 66.3us -- and it is -still not enough, because the release fence adds more than that back inside the -GEMM. An ATT thread trace puts the epilogue tail at 26% of the kernel's latency, -almost all of it one instruction: - - vaddr cycles % kernel instruction - 19396 91336 5.1% s_waitcnt vmcnt(0) retire this block's C - 19400 22984 1.3% s_barrier - 19404 292 0.0% buffer_wbl2 sc1 <- the op itself is free - 19412 246804 13.9% s_waitcnt vmcnt(0) <- waiting on the writeback - -That wait is the largest single item in the whole kernel, larger than any -main-loop stall, and it is **quadratic**: ``buffer_wbl2`` flushes the entire L2, -so block k waits for every tile blocks 1..k-1 already wrote, 448 blocks over. The -split path gets one bulk release at kernel end, free. - -Everything tried to make it cheaper, and what happened: - -* **Agent scope instead of system.** Correct and worth 20us: ``__threadfence_system`` - is bidirectional and emits ``buffer_wbl2 sc0 sc1 ; s_waitcnt ; buffer_inv sc0 - sc1``, and that ``buffer_inv`` discards the B weights the GEMM streams, 3584 - times. cco exposes no agent-scope release, only ``cco_system_fence``, so it is - emitted as a raw ``llvm.fence`` (``_compat.release_fence``). Kept. -* **One fence per block instead of per wave** (``--fence agent-leader``, and cco's - ``leaderOnly``). Incorrect in both spellings, including emitted inline in this - kernel to rule out the compiler sinking it. So the release is genuinely - per-wave. Its *time* is the useful part: 322.5us, i.e. an 8x cut in fence count - would be worth 41us and would win -- if it were legal. -* **Non-temporal C stores** (``--fence nt-agent``). Worse, 411.2us. -* **Write-through C stores.** Racy, see above. - -So what is missing is a **per-tile release**. Every primitive here publishes the -whole cache, and the GEMM has 448 tiles to publish, so the cost is O(blocks^2) -against an overlap worth at most one GEMM. The completion-counter machinery and -the destination-rotated tile order are sound and produced the real overlap; they -would pay immediately on a transport whose producer-side release is per-tile. +Correct-as-measured configurations only, 8 ranks, [4096, 7168], K=1024: + + split-lsa 324.4us + split-sdma 338.0 + fused-sdma chunks=1, no fence 348.5 (stable 6/6) + fused-sdma chunks=1, agent fence 440.0 (the fence is 90us of waste) + +Per-kernel (rocprofv3 --kernel-trace, us): + + config gemm scatter reduce gather sum + split-sdma 39.1 136.8 13.5 137.3 326.8 + fused c=1, no fence 62.7 129.0 13.0 136.2 341.0 + fused c=1, agent fence 152.4 132.7 13.6 135.4 434.1 + +So with the fence removed the whole remaining deficit is +23.6us inside the GEMM +-- the block-wide barrier, the acq_rel counter, and the copy engine reading C +while the GEMM writes it -- against a 7.8us overlap gain. The structure is sound; +the shape is wrong for it. ## Pinned copy @@ -191,6 +172,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr import gpu as fgpu from flydsl.expr.typing import Int64 from aiter.ops.flydsl.kernels.gemm_a8w8_8wave import ( @@ -211,9 +193,15 @@ _AR_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ar") sys.path.insert(0, _AR_DIR) +from flydsl._mlir import ir # noqa: E402 +from flydsl._mlir.dialects import scf # noqa: E402 + from _compat import ( # noqa: E402 CM_SC0_SC1, atomic_add_u32, + atomic_store_u32, + atomic_xchg_u32, + i32_type, buffer_store, create_buffer_resource_from_addr, release_fence, @@ -221,6 +209,32 @@ ) +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. @@ -296,7 +310,7 @@ def compile_fused_gemm_scatter( xcd_swizzle: int = 0, fuse: bool = True, rotated: bool | None = None, - fence: str = "agent", + fence: str = "none", emit_put: bool = True, atomic_order: str = "acq_rel", ): @@ -376,6 +390,7 @@ def compile_fused_gemm_scatter( ) _kname_tag = f"c{chunks}{'r' if rotated else 'l'}{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 @@ -665,6 +680,12 @@ def kernel_gemm_scatter( # 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( + fx.Int64(w.lsa_ptr(rank, lock_off)) + + fx.Int64(dest) * fx.Int64(4) + ) + if const_expr(chunks > 1): + _acquire_peer_lock(lock) sdma.put( dest, win, @@ -678,6 +699,11 @@ def kernel_gemm_scatter( 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( From b5e5916955845b3f20c6a04b0813e26e2b4ac2c9 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 08:15:50 +0000 Subject: [PATCH 11/79] feat(cco): port gcnasm's Direct LSA fusion; it absorbs the transfer but is racy --mode fused-lsa has the GEMM epilogue store straight into the destination rank's window, gcnasm's best a2a mode. No staging, no copy engine, no completion counter, no submit lock, no release fence -- the only sync left is the cross-rank barrier after the kernel. Structurally it does what fused-sdma never could (rocprofv3, us): config gemm scatter/barrier reduce gather sum split-sdma 39.1 136.8 13.5 137.3 326.8 fused-sdma c=1 62.7 129.0 13.0 136.2 341.0 fused-lsa (Direct) 166.6 13.3 11.1 135.5 326.6 The scatter collapses 136.8 -> 13.3us; the transfer is entirely absorbed. And the un-coalesced store is far better than 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 is not correct, and finding out why turns the StoreC coalescing rewrite from an optimisation into a prerequisite: - cached peer stores validate 2-3 runs in 6. The lines are homed in the peer but sit dirty in whichever XCD's L2 the block ran on, and nothing downstream can clean that: the barrier kernel is one block and buffer_wbl2 writes back only its own XCD's L2, one of eight. - uncached sc0|sc1 peer stores are *consistently* wrong (relL2 ~3e-2) and slower (426 vs 341us). _store_bf16 writes one bf16, so bypassing L2 makes every store a 2-byte partial-line write and concurrent sub-32B writes into a sector across the fabric lose updates. An uncached recv load in the reduce does not change it, so it is the store, not a stale read. L2 merging is what makes the narrow store work at all; L2 residency is what makes it invisible. Only a wider store separates them -- which is exactly why gcnasm needed its ds_bpermute pair-coalesced C store. Even coalesced it likely ties rather than wins: 326.6 against split-lsa's ~324. LSA's 2-stage does read+reduce+write in one pass over the wire so its reduce is free, while Direct LSA writes to the peer and then pays a separate local reduce, with a store rate ~18% under LSA's read rate. Also plumbs two switches the mode needs: build_sdma_phases gains reduce_self_from_recv (the epilogue writes my own slice into my recv slot, so j=0 comes from there) and recv_uncached. fused-lsa is deliberately left out of test_gemm_ar_modes_agree while it is racy. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/kernels_sdma.py | 32 ++++- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 27 ++++- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 114 +++++++++++++++++- 3 files changed, 160 insertions(+), 13 deletions(-) diff --git a/benchmark/cco/flydsl/ar/kernels_sdma.py b/benchmark/cco/flydsl/ar/kernels_sdma.py index 44450c105..ac97f52ae 100644 --- a/benchmark/cco/flydsl/ar/kernels_sdma.py +++ b/benchmark/cco/flydsl/ar/kernels_sdma.py @@ -90,6 +90,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from _compat import ( # noqa: E402 CM_CACHED, + CM_SC1, buffer_load, buffer_store, create_buffer_resource_from_addr, @@ -110,7 +111,14 @@ def build_sdma_phases( - cfg, rank: int, *, queues: int = 8, signal: bool = False, reduce_blocks=None + cfg, + rank: int, + *, + queues: int = 8, + signal: bool = False, + reduce_blocks=None, + reduce_self_from_recv: bool = False, + recv_uncached: bool = False, ): """Compile the phases separately, so the fused GEMM can reuse the tail. @@ -246,9 +254,17 @@ def sdma_reduce(dev_comm: Int64, win: Int64): bid = fx.block_idx.x w = cco.Window(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, in_off + my_slice_off)) + wave_uniform_i64(w.lsa_ptr(rank, self_off)) ) ] + [ create_buffer_resource_from_addr( @@ -265,8 +281,18 @@ def sdma_reduce(dev_comm: Int64, win: Int64): 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()) + 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) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 82c5070bb..22cfd4d03 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -31,6 +31,9 @@ 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 @@ -44,6 +47,7 @@ split-lsa 324.4us split-sdma 338.0 fused-sdma chunks=1, no fence 348.5 + fused-lsa 341.0 RACY, see kernels_fused Fusing loses. Read ``kernels_fused.py`` before trusting any faster fused number from this benchmark: several of its options are intermittently wrong, and a @@ -80,7 +84,7 @@ from layout import ArConfig # noqa: E402 VMM_SLACK = 512 * 1024 * 1024 -MODES = ("gemm-only", "split-sdma", "fused-sdma", "split-lsa") +MODES = ("gemm-only", "split-sdma", "fused-sdma", "fused-lsa", "split-lsa") def _setup_distributed(): @@ -142,8 +146,10 @@ def run(args) -> int: from aiter.ops.shuffle import shuffle_weight local_rank, rank, world_size, uid = _setup_distributed() - fused = args.mode == "fused-sdma" - needs_sdma = args.mode in ("split-sdma", "fused-sdma") + 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") # 1, even though >1 is what would create the overlap. A chunk boundary means # the copy engine starts reading C *while the GEMM is still running*, and no # combination of submit lock, post-submit barrier, release fence or atomic @@ -193,6 +199,7 @@ def run(args) -> int: waves_per_eu=args.waves_per_eu, xcd_swizzle=args.xcd_swizzle, fuse=fused, + transport="lsa" if direct_lsa else "sdma", rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", fence=args.fence, emit_put=not args.no_put, @@ -207,7 +214,13 @@ def run_gemm(stream): stream=stream) if needs_sdma: - parts = build_sdma_phases(cfg, rank, queues=args.sdma_queues) + parts = build_sdma_phases( + cfg, + rank, + queues=args.sdma_queues, + reduce_self_from_recv=direct_lsa, + recv_uncached=direct_lsa, + ) if args.mode == "split-lsa": lsa_ar, _ = build_lsa_ar(cfg, rank) @@ -219,8 +232,10 @@ def once(): if args.mode == "split-lsa": lsa_ar(dc.ptr, win.handle, stream=stream) return - # fused: the scatter puts were already issued from the epilogue, so - # only the drain runs here; split: the full scatter kernel. + # 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["drain" if fused else "scatter"](dc.ptr, win.handle, stream=stream) if args.stop_after == "scatter": return diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 2ff6a559d..403e4806e 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,7 +71,51 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## What gcnasm does differently, and why its GEMM+a2a wins +## Direct LSA (``--mode fused-lsa``): the structure works, the store does not + +gcnasm's best mode has the GEMM epilogue store *straight into the destination +rank's window*, so there is no staging buffer, no copy engine and nothing to +publish mid-kernel. Ported here, and the structural half is emphatic +(rocprofv3 --kernel-trace, us): + + config gemm scatter/barrier reduce gather sum + split-sdma 39.1 136.8 13.5 137.3 326.8 + fused-sdma c=1 62.7 129.0 13.0 136.2 341.0 + fused-lsa (Direct) 166.6 13.3 11.1 135.5 326.6 + +The scatter collapses from 136.8us to 13.3us -- the transfer is *entirely* +absorbed, which is what fused-sdma never managed. And the un-coalesced store is +nowhere near as bad as expected: 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 is nonetheless **not correct**, and the reason makes the ``StoreC`` coalescing +rewrite a correctness prerequisite rather than an optimisation: + +* **Cached peer stores validate 2-3 runs in 6.** The lines are homed in the + peer's memory but sit dirty in whichever XCD's L2 the block ran on. Nothing + downstream can clean that up -- the barrier kernel is one block, and + ``buffer_wbl2`` writes back only the L2 of the XCD it runs on, one of eight. +* **Uncached (sc0|sc1) peer stores are consistently wrong** (relL2 ~3e-2, never + right) **and slower** (426us against 341). ``_store_bf16`` writes one bf16, so + bypassing L2 makes every store a 2-byte partial-line write; concurrent sub-32B + writes into the same sector across the fabric lose updates. Adding an uncached + recv load in the reduce does not change it, so this is the store, not a stale + read. + +So L2 merging is what makes the narrow store work at all, and L2 residency is +what makes it invisible. The two are only separable by widening the store -- +gcnasm's wave-local ``ds_bpermute`` pair-coalesced C store (their "C-store mode +2"; staging C through LDS regressed 1-3%). That is why they needed it. + +Even coalesced, this does not obviously win. The fused sum is 326.6us against +split-lsa's ~324: LSA's 2-stage does read + reduce + write in a single pass over +the wire, so its reduce is free, whereas Direct LSA writes to the peer and then +pays a separate local reduce pass -- and its store rate is ~18% under LSA's read +rate. The prize for coalescing is correctness plus maybe a few percent, not a +step change. + +## What gcnasm does differently## 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 @@ -201,7 +245,10 @@ atomic_add_u32, atomic_store_u32, atomic_xchg_u32, + buffer_store, + create_buffer_resource_from_addr, i32_type, + wave_uniform_i64, buffer_store, create_buffer_resource_from_addr, release_fence, @@ -263,6 +310,41 @@ def _store_bf16(self, value_bf16, c_index): buffer_store(value_bf16, self._c_rsrc, c_index, cache_modifier=CM_SC0_SC1) +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). @@ -310,6 +392,7 @@ def compile_fused_gemm_scatter( xcd_swizzle: int = 0, fuse: bool = True, rotated: bool | None = None, + transport: str = "sdma", fence: str = "none", emit_put: bool = True, atomic_order: str = "acq_rel", @@ -332,6 +415,9 @@ def compile_fused_gemm_scatter( 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}") + direct_lsa = fuse and transport == "lsa" rotated = fuse if rotated is None else rotated assert BLOCK_M >= 128 and BLOCK_N >= 256 @@ -396,7 +482,7 @@ def compile_fused_gemm_scatter( slice_bytes = cfg.slice_bytes _kname = ( - f"mori_fused_{'sdma' if fuse else 'split'}_8w_" + 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}" ) @@ -483,12 +569,32 @@ def kernel_gemm_scatter( ) mfma = Mfma16x16x128(N_TILES_A, N_TILES_B) + w_pre = cco.Window(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(fence in ("writethrough", "wt-agent")): + if 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 ) @@ -626,7 +732,7 @@ def kernel_gemm_scatter( store_c.store(c11_frag, base_row + LDS_BLOCK_M, base_col + LDS_BLOCK_N) # ---- end pinned copy ---- - if const_expr(fuse): + 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. From 72c47dee6294508a3ff84467cfe978f93dcd5a6f Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 10:07:42 +0000 Subject: [PATCH 12/79] perf(cco): MFMA operand swap for a wider C store -- correct, and not worth it StoreC emits 128 buffer_store_short per block because the MFMA accumulator layout gives each lane four consecutive *rows* (stride c_cols = 14336B), not columns. The eight bf16 that would make a 16-byte store live in eight different lanes. Ports gcnasm's mfma_adaptor_swap_ab (opus.hpp:2064, literally base::operator()(b, a, c) with dim_c() redefined) as --swap-ab: computing (A B)^T in the accumulator's own layout puts four consecutive *columns* in a lane. A/B scale loads swap with it, scalar for A and vec4 for B. Bit-identical to aiter's kernel at three shapes including wo_b, guarded by test_swap_ab_is_bitwise_identical. And slower: variant per lane lanes/row bytes/row rows/instr gemm-only unswapped (128 short) 2B 16 32B 4 50.4us --swap-ab (32 dwordx2) 8B 4 32B 16 55.6us + 16B/lane (probe) 16B 4 64B 16 48.1us Store instructions drop 4x and total kernel instructions 30%, and it does not help, because the transaction count is unchanged -- both layouts put 32 bytes in a row and need 16 of them. The swap just moves those 16 from four instructions onto one, so per-instruction address fan-out goes 4 -> 16. ATT: buffer_store_dwordx2 costs 1022 cycles each against buffer_store_short's 114. --store-probe answers whether the permlane stage is worth building before building it: it emits the 16B-per-lane / 64B-per-row pattern with deliberately wrong data, so it prices the access pattern with the shuffle's cost removed. 48.1us against 50.4 is a -4.6% ceiling locally, and on fused-lsa -- where the store is 14% of the kernel with 60x the stall -- it is a wash (340.8/338.7/333.1 vs 336.4/335.3/337.7), because that phase is xGMI bandwidth-limited, not transaction-limited. So the permlane stage is not built. The floor is the 16-column MFMA tile: 16 bf16 is 32 bytes, a quarter line, whatever the operand order. 64B needs two N-tiles merged, which needs a shuffle in *both* layouts; 128B needs merging across waves through LDS, which gcnasm measured at a 1-3% regression. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 17 ++ benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 245 +++++++++++++++++- tests/python/cco/test_gemm_ar.py | 46 ++++ 3 files changed, 303 insertions(+), 5 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 22cfd4d03..7f0015e86 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -200,6 +200,8 @@ def run(args) -> int: xcd_swizzle=args.xcd_swizzle, fuse=fused, transport="lsa" if direct_lsa else "sdma", + swap_ab=args.swap_ab, + store_probe=args.store_probe, rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", fence=args.fence, emit_put=not args.no_put, @@ -294,6 +296,7 @@ def once(): "block_n": args.block_n, "chunks": chunks if fused else None, "tile_order": args.tile_order, + "swap_ab": args.swap_ab, "fence": args.fence if fused else None, "stop_after": args.stop_after, "max_rank_time_us": per_rank[max_rank], @@ -377,6 +380,20 @@ def build_parser() -> argparse.ArgumentParser: "order the other blocks' releases before its put; 'monotonic' drops it " "and races", ) + p.add_argument( + "--swap-ab", + action="store_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)", + ) + 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("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 403e4806e..2453ebd3a 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,7 +71,56 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## Direct LSA (``--mode fused-lsa``): the structure works, the store does not +## Store width: the operand swap works and does not pay + +``StoreC`` emits 128 ``buffer_store_short`` per block -- one bf16 at a time -- and +the reason is 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 a +lane's four values are 14336 bytes apart in a row-major C. The eight bf16 that +would make a 16-byte store, ``C[r][c..c+7]``, live in eight *different* lanes. +(The same class loads A's scale 128 bits at a time, because that one is indexed +by row and a lane owns four consecutive rows.) + +``--swap-ab`` is gcnasm's ``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. The +A and B scale loads swap with it, scalar for A and vec4 for B. + +It is bit-identical to aiter's kernel at [512,512,256], [1024,768,512] and the +wo_b shape (``test_swap_ab_is_bitwise_identical``), and it is **slower**: + + variant per lane lanes/row bytes/row rows/instr gemm-only + unswapped (128 short) 2B 16 32B 4 50.4us + --swap-ab (32 dwordx2) 8B 4 32B 16 55.6us + + 16B/lane (probe) 16B 4 64B 16 48.1us + +Store instructions drop 4x and total kernel instructions 30%, and none of that +matters, because the number of memory transactions does not change: both layouts +put 32 bytes in a row and need 16 of them. The swap only moves those 16 +transactions from four instructions onto one, quadrupling the per-instruction +address fan-out -- ATT puts ``buffer_store_dwordx2`` at 1022 cycles each against +``buffer_store_short``'s 114. + +What actually helps is 64 bytes per row, and only a little. ``--store-probe`` +emits that access pattern with deliberately wrong data, so it is an upper bound +with the cross-lane shuffle's cost removed: 48.1us against 50.4, i.e. **-4.6% at +best**, before paying for the ``permlane16_swap`` + ``shfl`` that would make it +correct. Not worth building. + +And it does nothing at all on the path where the store dominates. On +``fused-lsa``, where the store is 14% of the kernel with 60x the stall of the +local case, the same probe measures 340.8/338.7/333.1us against the current +336.4/335.3/337.7 -- a wash. That phase is limited by xGMI bandwidth, not by +transaction count, so halving the transactions buys nothing. + +The 16-column MFMA tile is the real floor here: 16 bf16 is 32 bytes, a quarter of +a 128-byte line, whatever the operand order. Reaching 64B needs two N-tiles +merged (they are 16 columns apart in *both* layouts, so it needs a shuffle +either way), and 128B needs merging across waves through LDS -- which gcnasm +measured at a 1-3% regression. + +## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not gcnasm's best mode has the GEMM epilogue store *straight into the destination rank's window*, so there is no staging buffer, no copy engine and nothing to @@ -215,9 +264,10 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl.expr import const_expr, range_constexpr, rocdl +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 from aiter.ops.flydsl.kernels.gemm_a8w8_8wave import ( G2SLoader, @@ -310,6 +360,164 @@ 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 _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, **kwargs): + super().__init__(*args, **kwargs) + 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) + + def _load_a_scale_scalar(self, row): + 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): + 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 = [ + (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)) + + +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 __init__(self, *args, peer_rsrc=None, elem_base=None, **kwargs): + super().__init__(*args, **kwargs) + 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) + # When set, the 16B stores go straight to the owning peer (Direct LSA) + # rather than to local C, so the probe can price the same access pattern + # on the path where the store actually dominates. + self._peer_rsrc = peer_rsrc + self._elem_base = elem_base + + 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 += [ + (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. @@ -392,6 +600,8 @@ def compile_fused_gemm_scatter( xcd_swizzle: int = 0, fuse: bool = True, rotated: bool | None = None, + swap_ab: bool = False, + store_probe: bool = False, transport: str = "sdma", fence: str = "none", emit_put: bool = True, @@ -474,7 +684,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"c{chunks}{'r' if rotated else 'l'}{fence[0]}" + _kname_tag = f"{'S' if swap_ab else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" counter_off = cfg.counter_off lock_off = cfg.lock_off in_off = cfg.input_off @@ -568,14 +778,39 @@ def kernel_gemm_scatter( lane_id, wave_id, K, N_LDS_ROUNDS, preshuffled=b_preshuffled ) - mfma = Mfma16x16x128(N_TILES_A, N_TILES_B) + 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.Window(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(direct_lsa): + if 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 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) diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 01951bc1b..a6446d96c 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -156,6 +156,52 @@ def test_pinned_copy_matches_aiter_kernel_bitwise(m, n, k): ) +@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. So this is exactly reproducible, and + anything less than bit-equality means the store's index math drifted rather + than that the arithmetic changed. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + pytest.importorskip("aiter", reason="aiter not importable (set PYTHONPATH)") + from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 + from aiter.ops.shuffle import shuffle_weight + + import flydsl.expr as fx + + from kernels_fused import compile_fused_gemm_scatter + + 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 = shuffle_weight(b, layout=(16, 16)) + + ref = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) + flydsl_8wave_gemm_a8(a, b_shuf, sa, sb, ref, 256, 256) + + 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, swap_ab=True, + ) + 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) + + def test_rotated_tile_order_is_a_permutation_of_the_linear_one(): """Reordering tiles must not change which tiles exist -- checked on the host. From f555292d8ea9d9198ba290d23cf9f11c42f72464 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 10:29:25 +0000 Subject: [PATCH 13/79] perf(cco): --swap-ab --permlane, bit-exact and 3.3% off the GEMM I dismissed this last commit on a guess that permlane16_swap + shfl would eat the 4.6% the --store-probe ceiling showed. The guess was wrong: it needs two permlane16_swap per M-tile, no ds_bpermute and no LDS, and they cost 4,260 cycles to save about 33,000. The mapping falls out once permlane16_swap's semantics are measured rather than assumed -- vdst' = [X.r0, Y.r0, X.r2, Y.r2], vsrc' = [X.r1, Y.r1, X.r3, Y.r3]: (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 -- together columns 0..31 contiguously. 16 bytes per lane, 64 per row, and the column permutation is absorbed into the address. Bit-identical to aiter's kernel at three shapes including wo_b. GEMM kernel median of 42 dispatches, 8 ranks: variant bytes/row store instrs store cycles gemm kernel unswapped 32B 128 short 14,608 38.30us --swap-ab 32B 32 dwordx2 32,720 slower --swap-ab --permlane 64B 16 dwordx4 14,604 37.04us The store is not where the time goes, which was the part I had wrong in both directions. Its total latency is unchanged, 14,608 vs 14,604, and --swap-ab alone is slower because at 32B/row the same 16 transactions get squeezed onto a quarter as many instructions (1022 cycles each vs 114). permlane's 64B rows just put the store back where it started. What pays is what a 2-byte store drags along. Thread trace, unswapped -> permlane: v_mul_f32_e32 257 -> 1 17,664 -> 68 cyc v_pk_mul_f32 0 -> 128 0 -> 9,276 v_lshlrev_b32_e32 159 -> 45 11,304 -> 3,828 v_add_u32_e32 133 -> 21 8,944 -> 1,604 v_cndmask_b32_e64 96 -> 8 6,368 -> 656 v_cvt_pk_bf16_f32 128 -> 64 8,384 -> 4,244 v_permlane16_swap_b32_e64 0 -> 32 0 -> 4,260 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. Kernel instructions 2062 -> 1421, registers unchanged (VGPR 128 / SGPR 112 / 0 scratch). Two limits: it does nothing on fused-lsa (335-338us either way), where the store phase is xGMI bandwidth-limited rather than issue-limited; and 64B is one wave's ceiling, since a wave owns N_TILES_B*16 = 32 columns. Worth pushing back into aiter's StoreC: bit-exact, no register cost, and the 3.3% is on the GEMM itself, independent of any all-reduce. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 8 + benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 244 +++++++++++++----- tests/python/cco/test_gemm_ar.py | 34 +-- 3 files changed, 208 insertions(+), 78 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 7f0015e86..c8669f940 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -202,6 +202,7 @@ def run(args) -> int: transport="lsa" if direct_lsa else "sdma", swap_ab=args.swap_ab, store_probe=args.store_probe, + permlane=args.permlane, rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", fence=args.fence, emit_put=not args.no_put, @@ -394,6 +395,13 @@ def build_parser() -> argparse.ArgumentParser: "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="store_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", + ) p.add_argument("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 2453ebd3a..97d1ad948 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,56 +71,78 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## Store width: the operand swap works and does not pay - -``StoreC`` emits 128 ``buffer_store_short`` per block -- one bf16 at a time -- and -the reason is 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 a -lane's four values are 14336 bytes apart in a row-major C. The eight bf16 that -would make a 16-byte store, ``C[r][c..c+7]``, live in eight *different* lanes. -(The same class loads A's scale 128 bits at a time, because that one is indexed -by row and a lane owns four consecutive rows.) - -``--swap-ab`` is gcnasm's ``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. The -A and B scale loads swap with it, scalar for A and vec4 for B. - -It is bit-identical to aiter's kernel at [512,512,256], [1024,768,512] and the -wo_b shape (``test_swap_ab_is_bitwise_identical``), and it is **slower**: - - variant per lane lanes/row bytes/row rows/instr gemm-only - unswapped (128 short) 2B 16 32B 4 50.4us - --swap-ab (32 dwordx2) 8B 4 32B 16 55.6us - + 16B/lane (probe) 16B 4 64B 16 48.1us - -Store instructions drop 4x and total kernel instructions 30%, and none of that -matters, because the number of memory transactions does not change: both layouts -put 32 bytes in a row and need 16 of them. The swap only moves those 16 -transactions from four instructions onto one, quadrupling the per-instruction -address fan-out -- ATT puts ``buffer_store_dwordx2`` at 1022 cycles each against -``buffer_store_short``'s 114. - -What actually helps is 64 bytes per row, and only a little. ``--store-probe`` -emits that access pattern with deliberately wrong data, so it is an upper bound -with the cross-lane shuffle's cost removed: 48.1us against 50.4, i.e. **-4.6% at -best**, before paying for the ``permlane16_swap`` + ``shfl`` that would make it -correct. Not worth building. - -And it does nothing at all on the path where the store dominates. On -``fused-lsa``, where the store is 14% of the kernel with 60x the stall of the -local case, the same probe measures 340.8/338.7/333.1us against the current -336.4/335.3/337.7 -- a wash. That phase is limited by xGMI bandwidth, not by -transaction count, so halving the transactions buys nothing. - -The 16-column MFMA tile is the real floor here: 16 bf16 is 32 bytes, a quarter of -a 128-byte line, whatever the operand order. Reaching 64B needs two N-tiles -merged (they are 16 columns apart in *both* layouts, so it needs a shuffle -either way), and 128B needs merging across waves through LDS -- which gcnasm -measured at a 1-3% regression. - -## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not +## Store width: --swap-ab --permlane is worth 3.3%, for a reason that is not width + +``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. (The same class already loads +A's scale 128 bits at a time -- that one is indexed by row, and a lane owns four +consecutive rows.) + +Two steps fix it, both from gcnasm: + +* ``--swap-ab`` is ``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. +* ``--permlane`` adds two ``v_permlane16_swap_b32`` per M-tile. Its semantics, + measured on hardware rather than assumed, are + ``vdst' = [X.r0, Y.r0, X.r2, Y.r2]`` and ``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. The four groups together + cover columns 0..31 contiguously: 16 bytes per lane, 64 per row. The column + permutation is absorbed into the address, so no ``ds_bpermute`` and no LDS. + +Both are bit-identical to aiter's kernel at [512,512,256], [1024,768,512] and the +wo_b shape (``test_swap_ab_is_bitwise_identical``). Measured, 8 ranks, GEMM kernel +median over 42 dispatches: + + variant bytes/row store instrs store cycles gemm kernel + unswapped 32B 128 short 14,608 38.30us + --swap-ab 32B 32 dwordx2 32,720 slower + --swap-ab --permlane 64B 16 dwordx4 14,604 37.04us + +**The store is not where the time goes.** Its total latency is the same before and +after -- 14,608 against 14,604 -- and ``--swap-ab`` alone is *slower* because at +32 bytes per row the same 16 transactions get squeezed onto a quarter as many +instructions, quadrupling per-instruction address fan-out (1022 cycles each +against 114). ``permlane``'s 64-byte rows halve the transactions, which is only +enough to put the store back where it started. + +What actually pays is everything a 2-byte store drags along with it. From the +thread trace, unswapped -> permlane: + + v_mul_f32_e32 257 -> 1 17,664 -> 68 cyc + v_pk_mul_f32 0 -> 128 0 -> 9,276 + v_lshlrev_b32_e32 159 -> 45 11,304 -> 3,828 + v_add_u32_e32 133 -> 21 8,944 -> 1,604 + v_cndmask_b32_e64 96 -> 8 6,368 -> 656 + v_cvt_pk_bf16_f32 128 -> 64 8,384 -> 4,244 + v_permlane16_swap_b32_e64 0 -> 32 0 -> 4,260 + +128 stores need 128 addresses, 128 bounds predicates and 128 scalar multiplies; +16 stores need 16. And the swapped layout makes B's scale a vec4, so the scaling +packs into ``v_pk_mul_f32``. Total kernel instructions drop 2062 -> 1421 and the +permlanes cost 4,260 cycles to save roughly 33,000. + +Two limits worth recording. This does nothing on ``fused-lsa`` +(335-338us either way), where the store phase is xGMI bandwidth-limited rather +than transaction- or issue-limited. And 64 bytes is the ceiling for one wave: a +wave owns ``N_TILES_B * 16 = 32`` columns, so 128-byte lines would need merging +across waves through LDS, which gcnasm measured at a 1-3% regression. + +This is a candidate to push back into aiter's own ``StoreC`` -- it is bit-exact, +costs no registers (VGPR 128 / SGPR 112 / 0 scratch, unchanged), and the 3.3% is +on the GEMM itself rather than on anything to do with the all-reduce. + +## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not gcnasm's best mode has the GEMM epilogue store *straight into the destination rank's window*, so there is no staging buffer, no copy engine and nothing to @@ -288,6 +310,7 @@ _AR_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ar") sys.path.insert(0, _AR_DIR) from flydsl._mlir import ir # noqa: E402 +from flydsl._mlir.dialects import llvm as _llvm_d # noqa: E402 from flydsl._mlir.dialects import scf # noqa: E402 from _compat import ( # noqa: E402 @@ -397,10 +420,14 @@ class _SwapABStoreC(StoreC): while B's is indexed by column, now 4 consecutive, so a vec4. """ - def __init__(self, *args, **kwargs): + def __init__(self, *args, peer_rsrc=None, elem_base=None, **kwargs): super().__init__(*args, **kwargs) 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 _load_a_scale_scalar(self, row): fx.copy( @@ -453,6 +480,92 @@ def store(self, c_frag, base_row, base_col): self._store_bf16x4(vals, arith.select(col_valid, c_index, oob)) +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)>") + raw = lambda v: v.ir_value() if hasattr(v, "ir_value") else v + 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), + ) + + +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. + """ + + def store(self, c_frag, base_row, base_col): + 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 + 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 + grp * 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 + 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( + [ + (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) + 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: + buffer_store(out8, self._peer_rsrc, fx.Int32(idx) - self._elem_base) + 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 _WideStoreProbeC(_SwapABStoreC): """PERF PROBE ONLY -- the output is wrong on purpose. @@ -468,15 +581,7 @@ class _WideStoreProbeC(_SwapABStoreC): ``permlane16_swap`` step; this probe skips it and keeps only its cost profile. """ - def __init__(self, *args, peer_rsrc=None, elem_base=None, **kwargs): - super().__init__(*args, **kwargs) - 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) - # When set, the 16B stores go straight to the owning peer (Direct LSA) - # rather than to local C, so the probe can price the same access pattern - # on the path where the store actually dominates. - self._peer_rsrc = peer_rsrc - self._elem_base = elem_base + def store(self, c_frag, base_row, base_col): lane = self.lane_id @@ -602,6 +707,7 @@ def compile_fused_gemm_scatter( rotated: bool | None = None, swap_ab: bool = False, store_probe: bool = False, + permlane: bool = False, transport: str = "sdma", fence: str = "none", emit_put: bool = True, @@ -684,7 +790,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"{'S' if swap_ab else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" + _kname_tag = f"{'S' if swap_ab else ''}{'P' if permlane else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" counter_off = cfg.counter_off lock_off = cfg.lock_off in_off = cfg.input_off @@ -792,7 +898,11 @@ def kernel_gemm_scatter( 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 store_probe and not direct_lsa): + if 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 ) @@ -800,6 +910,16 @@ def kernel_gemm_scatter( 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): + 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( diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index a6446d96c..b659e98fd 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -162,9 +162,10 @@ def test_swap_ab_is_bitwise_identical(m, n, k): ``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. So this is exactly reproducible, and - anything less than bit-equality means the store's index math drifted rather - than that the arithmetic changed. + 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") @@ -187,19 +188,20 @@ def test_swap_ab_is_bitwise_identical(m, n, k): flydsl_8wave_gemm_a8(a, b_shuf, sa, sb, ref, 256, 256) 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, swap_ab=True, - ) - 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) + for extra in ({}, {"permlane": 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} drifted from aiter" def test_rotated_tile_order_is_a_permutation_of_the_linear_one(): From 0b26ebf8f2278e82ea3522f2708e7d693e18f487 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 10:49:38 +0000 Subject: [PATCH 14/79] perf(cco): port gcnasm's lane transpose too; C store now 6.7% off the GEMM Third stage, from the ds_bpermute block at opus_gemm_a2a_lsa/gemm_a16w16_quad_subtile_kernel_template.hpp:491-510. 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 being 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 instead of 16. GEMM kernel, median of 42 dispatches, 8 ranks, [4096,7168] K=1024: stage gemm 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 Bit-identical to aiter at all three shapes including wo_b, registers unchanged (VGPR 128 / SGPR 112 / 0 scratch) at every stage. The striking part is that the third stage changes neither the instructions nor the addresses -- same 16 buffer_store_dwordx4 over the same 64 addresses, only which lane holds which -- and the store's own latency falls 14,604 -> 3,852. So the coalescer is sensitive to lane adjacency and not just to the address set, which is what gcnasm's "pair-coalesced" comment is really about. The 64 ds_bpermute cost 9,604 cycles. Priced before building, as with the previous stage: a --store-probe variant emits the adjacent-lane access pattern with deliberately wrong data, and measured 45.4 against 48.1us, so the real version at 46.3 lands close to its ceiling. Note the three stages pay for different things. --swap-ab alone is slower. --permlane does not speed the store up at all (14,608 -> 14,604); it pays by deleting what a 2-byte store drags along -- 128 addresses, 128 bounds predicates and 128 scalar multiplies become 16, and B's scale becoming a vec4 packs the scaling into v_pk_mul_f32 (v_mul_f32_e32 257 -> 1). Only --lane-transpose actually makes the store faster. Still nothing on fused-lsa (335-338us either way): that store phase is xGMI bandwidth-limited. And 64 bytes is one wave's ceiling here, since a wave owns N_TILES_B*16 = 32 columns; gcnasm reaches a 128-byte line only by tiling four wave_id_n waves across adjacent 16-column runs. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 7 + benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 200 ++++++++++++------ tests/python/cco/test_gemm_ar.py | 6 +- 3 files changed, 142 insertions(+), 71 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index c8669f940..adc5bdf43 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -203,6 +203,7 @@ def run(args) -> int: swap_ab=args.swap_ab, store_probe=args.store_probe, permlane=args.permlane, + lane_transpose=args.lane_transpose, rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", fence=args.fence, emit_put=not args.no_put, @@ -402,6 +403,12 @@ def build_parser() -> argparse.ArgumentParser: "contiguous bytes and each row gets 64. This is the real version of what " "--store-probe prices", ) + p.add_argument( + "--lane-transpose", + action="store_true", + help="with --swap-ab --permlane: one ds_bpermute per dword so adjacent " + "lanes cover one C row (gcnasm kernel_template.hpp:491-510)", + ) p.add_argument("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 97d1ad948..dbc8e7d24 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -71,78 +71,77 @@ ``opus_direct_stripe_tile`` idea. It costs nothing: 49.72us against 49.80us for the GEMM alone. -## Store width: --swap-ab --permlane is worth 3.3%, for a reason that is not width +## 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. (The same class already loads -A's scale 128 bits at a time -- that one is indexed by row, and a lane owns four -consecutive rows.) - -Two steps fix it, both from gcnasm: - -* ``--swap-ab`` is ``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. -* ``--permlane`` adds two ``v_permlane16_swap_b32`` per M-tile. Its semantics, - measured on hardware rather than assumed, are - ``vdst' = [X.r0, Y.r0, X.r2, Y.r2]`` and ``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. The four groups together - cover columns 0..31 contiguously: 16 bytes per lane, 64 per row. The column - permutation is absorbed into the address, so no ``ds_bpermute`` and no LDS. - -Both are bit-identical to aiter's kernel at [512,512,256], [1024,768,512] and the -wo_b shape (``test_swap_ab_is_bitwise_identical``). Measured, 8 ranks, GEMM kernel -median over 42 dispatches: - - variant bytes/row store instrs store cycles gemm kernel - unswapped 32B 128 short 14,608 38.30us - --swap-ab 32B 32 dwordx2 32,720 slower - --swap-ab --permlane 64B 16 dwordx4 14,604 37.04us - -**The store is not where the time goes.** Its total latency is the same before and -after -- 14,608 against 14,604 -- and ``--swap-ab`` alone is *slower* because at -32 bytes per row the same 16 transactions get squeezed onto a quarter as many -instructions, quadrupling per-instruction address fan-out (1022 cycles each -against 114). ``permlane``'s 64-byte rows halve the transactions, which is only -enough to put the store back where it started. - -What actually pays is everything a 2-byte store drags along with it. From the -thread trace, unswapped -> permlane: - - v_mul_f32_e32 257 -> 1 17,664 -> 68 cyc - v_pk_mul_f32 0 -> 128 0 -> 9,276 - v_lshlrev_b32_e32 159 -> 45 11,304 -> 3,828 - v_add_u32_e32 133 -> 21 8,944 -> 1,604 - v_cndmask_b32_e64 96 -> 8 6,368 -> 656 - v_cvt_pk_bf16_f32 128 -> 64 8,384 -> 4,244 - v_permlane16_swap_b32_e64 0 -> 32 0 -> 4,260 - -128 stores need 128 addresses, 128 bounds predicates and 128 scalar multiplies; -16 stores need 16. And the swapped layout makes B's scale a vec4, so the scaling -packs into ``v_pk_mul_f32``. Total kernel instructions drop 2062 -> 1421 and the -permlanes cost 4,260 cycles to save roughly 33,000. - -Two limits worth recording. This does nothing on ``fused-lsa`` -(335-338us either way), where the store phase is xGMI bandwidth-limited rather -than transaction- or issue-limited. And 64 bytes is the ceiling for one wave: a -wave owns ``N_TILES_B * 16 = 32`` columns, so 128-byte lines would need merging -across waves through LDS, which gcnasm measured at a 1-3% regression. - -This is a candidate to push back into aiter's own ``StoreC`` -- it is bit-exact, -costs no registers (VGPR 128 / SGPR 112 / 0 scratch, unchanged), and the 3.3% is -on the GEMM itself rather than on anything to do with the all-reduce. - -## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not +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. + +Two limits. It does nothing on ``fused-lsa`` (335-338us either way), where the +store phase is xGMI bandwidth-limited rather than issue- or coalescing-limited. +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. + +## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not gcnasm's best mode has the GEMM epilogue store *straight into the destination rank's window*, so there is no staging buffer, no copy engine and nothing to @@ -502,6 +501,15 @@ def _permlane16_swap(x, y): ) +def _ds_bpermute(value, src_lane): + """Pull ``value`` from ``src_lane``. Byte-addressed, hence the <<2.""" + i32 = ir.IntegerType.get_signless(32) + raw = lambda v: v.ir_value() if hasattr(v, "ir_value") else v + 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. @@ -522,6 +530,9 @@ class _PermlaneStoreC(_SwapABStoreC): shuffle only moves whole 16-lane rows, so every lane keeps its own C row. """ + _adjacent_probe = False + _lane_transpose = False + def store(self, c_frag, base_row, base_col): assert self.n_tiles_b == 2, ( "the permlane mapping pairs exactly two N-tiles (BLOCK_N == 256); " @@ -553,7 +564,35 @@ def store(self, c_frag, base_row, base_col): 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) - col = base_col + (grp % 2) * 16 + (grp // 2) * 8 + 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: @@ -566,6 +605,18 @@ def store(self, c_frag, base_row, base_col): ) +class _LaneTransposeStoreC(_PermlaneStoreC): + """``_PermlaneStoreC`` plus gcnasm's ds_bpermute lane transpose.""" + + _lane_transpose = 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. @@ -708,6 +759,7 @@ def compile_fused_gemm_scatter( swap_ab: bool = False, store_probe: bool = False, permlane: bool = False, + lane_transpose: bool = False, transport: str = "sdma", fence: str = "none", emit_put: bool = True, @@ -790,7 +842,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"{'S' if swap_ab else ''}{'P' if permlane else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" + _kname_tag = f"{'S' if swap_ab else ''}{'P' if permlane else ''}{'T' if lane_transpose else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" counter_off = cfg.counter_off lock_off = cfg.lock_off in_off = cfg.input_off @@ -898,7 +950,15 @@ def kernel_gemm_scatter( 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 not direct_lsa): + 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 ) diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index b659e98fd..23f4f25ed 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -188,7 +188,11 @@ def test_swap_ab_is_bitwise_identical(m, n, k): flydsl_8wave_gemm_a8(a, b_shuf, sa, sb, ref, 256, 256) cfg = layout.ArConfig(world_size=2, m=m, n=n) - for extra in ({}, {"permlane": True}): + for extra in ( + {}, + {"permlane": True}, + {"permlane": True, "lane_transpose": 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, From 76e285bebc9b0e8b24fce010acc20d3e2bf82058 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 13:02:29 +0000 Subject: [PATCH 15/79] fix(cco): wire the lane transpose into fused-lsa, and fail loudly on bad combos The C-store variant is picked by a chain of const_expr branches, and there was no `direct_lsa and ... lane_transpose` branch. So `--mode fused-lsa --lane-transpose` did not error -- it fell through and quietly ran stages 1+2. That produced a wrong conclusion. I reported "the store optimization does nothing on fused-lsa, that phase is xGMI bandwidth-limited" on the strength of a measurement that only ever exercised the two stages which do *not* speed the store up (permlane leaves it at 14,608 -> 14,604); stage 3, the one that takes it to 3,852, was never in the binary. With the branch actually wired, the conclusion survives, now for a reason that was measured rather than assumed: 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, and that phase is bounded by ~44 GB/s of link bandwidth, which transaction shaping does not move. Also adds argument validation for the store flags -- permlane requires swap_ab, lane_transpose requires permlane, store_probe requires swap_ab, and the two combinations that have no branch are rejected by name. Silent fallthrough is what made the original mistake possible and it should not be possible again. Current standings at [4096,7168] K=1024, 8 ranks, graph replay, median of 3, default C-store / with all three stages: gemm-only 49.6 / 46.8us split-lsa 324.8 / 322.8 correct split-sdma 333.2 / 331.4 correct fused-sdma 348.6 / 340.2 correct fused-lsa 334.2 / -- RACY, 1 run in 3 The C-store work barely moves the collectives -- the GEMM is ~15% of the pipeline, so 6.7% of it is ~3us. Its value is on the GEMM itself and is independent of any all-reduce. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 51 ++++++++++++++++++- 1 file changed, 49 insertions(+), 2 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index dbc8e7d24..5b8196a46 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -131,8 +131,23 @@ 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. -Two limits. It does nothing on ``fused-lsa`` (335-338us either way), where the -store phase is xGMI bandwidth-limited rather than issue- or coalescing-limited. +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. @@ -785,6 +800,28 @@ def compile_fused_gemm_scatter( 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") + if lane_transpose and not permlane: + raise ValueError("--lane-transpose requires --permlane") + if store_probe and not swap_ab: + raise ValueError("--store-probe requires --swap-ab") + 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" rotated = fuse if rotated is None else rotated @@ -970,6 +1007,16 @@ def kernel_gemm_scatter( 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) + store_c = _LaneTransposeStoreC( + 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( From 8bd495d6541e3153bf4d7af646fcb08c8a573e31 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 13:48:04 +0000 Subject: [PATCH 16/79] docs(cco): persistent tiles tried and rejected -- FlyDSL spills through scf.for Implemented gcnasm's persistent scheduler (one workgroup per CU striding over tiles), got it bit-exact, measured it, and reverted it. Recording the result and the two non-obvious bugs it surfaced, because the next person to read gcnasm's PERSISTENT=1 will want to try the same thing. The tail it targets is real and large. Sweeping N at M=4096, K=1024: N tiles remainder time 4096 256 0 32.64us 4352 272 16 45.88us <- 6% more work, 40% more time 6144 384 128 47.40 7168 448 192 49.80 <- wo_b But wo_b sits at remainder 192, the benign end of gcnasm's own sweep (+0.77% there, against +9.97% at remainder 8). And the implementation is a ~2x regression for a reason unrelated to scheduling: variant gemm scratch one workgroup per tile 38.54us 0 B body wrapped in an scf.for 74.48us 156 B/thread fly-promote-regmem-to-vectorssa does not promote make_rmem_tensor values declared inside an scf.for region, so the MFMA accumulators spill. Persistent does beat the *looped* baseline by ~3% (82.3 vs 84.8us at remainder 192, 76.8 vs 79.8 at remainder 16), consistent with gcnasm -- but that is a gain over an already 94%-regressed kernel. Two findings kept in the docstring: - The prologue's `if wave_m == 1: s_barrier()` runs the half-waves one barrier out of phase on purpose. In a loop that offset accumulates by one per tile, and 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()` per tile fixes it exactly; that is how the persistent version was made bit-exact. - The LDS handles must be rebound per tile. Hoisting them makes eight shared-address-space pointers loop-carried and fails to legalize. Also worth knowing for anyone wrapping a FlyDSL kernel in a runtime loop: `range(...)` has to appear literally in the `for` statement or the AST rewriter does not see it and Python evaluates it eagerly. Assigning it to a variable first raises "dynamic 'ArithValue' has no Python integer representation", and the workaround of always emitting the loop (step = tile count when not persistent) makes the default path pay the spill too. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 5b8196a46..0c6ccb10c 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -156,6 +156,57 @@ no register cost, and the 6.7% is on the GEMM itself, independent of any all-reduce. +## Persistent tiles: tried, and not viable on FlyDSL today + +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 a persistent scheduler was worth trying rather than arguing about. It was +implemented (one workgroup per CU, striding over tiles) and made bit-exact, and +it is a **~2x regression**, for a reason that has nothing to do with scheduling: + + variant gemm scratch + one workgroup per tile 38.54us 0 B + body wrapped in an scf.for 74.48us 156 B/thread + +Wrapping the pipeline in a runtime loop makes FlyDSL spill the MFMA +accumulators. ``fly-promote-regmem-to-vectorssa`` does not promote +``make_rmem_tensor`` values declared inside an ``scf.for`` region, and 156 bytes +per thread of scratch in the inner pipeline costs more than the entire tail. The +persistent version does beat the *looped* baseline by ~3% (82.3 vs 84.8us at +remainder 192, 76.8 vs 79.8 at remainder 16), which matches gcnasm's numbers -- +but that is a gain over an already 94%-regressed kernel. + +Two things worth keeping from the attempt: + +* **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. That is fine once, but in a persistent 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, + 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**, not hoisted: the pipeline swaps + those Python bindings as it advances, and hoisting makes eight + shared-address-space pointers loop-carried, which fails to legalize. + +Reviving this needs the FlyDSL promotion fix first. Until then the ceiling is +gcnasm's +0.77% at our remainder, against a 94% floor. + ## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not gcnasm's best mode has the GEMM epilogue store *straight into the destination From 4ec4e3ba88440f7660919b74d7ba5e4bcbaec734 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 14:20:38 +0000 Subject: [PATCH 17/79] docs(cco): correct the persistent-tile diagnosis -- it is VGPRs, not the promoter The previous commit blamed fly-promote-regmem-to-vectorssa for not promoting through scf.for. That is wrong, and checking it properly says so: - the pass handles scf::ForOp (rewriteForOp) - it promoted all 451 register allocas in the persistent kernel, zero left afterwards and zero llvm.alloca in the final IR - the emitted loop carries no iter_args at all - the 156 bytes are ordinary register-allocator spill: .vgpr_spill_count = 38, with 38 scratch_store_dword / 38 scratch_load_dword The real wall is the register budget, from the ISA kernel metadata: 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. 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 them potentially 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 is what pins the cost on scf.for rather than on the scheduling idea. The pass does have one real inefficiency, just not one we hit: an alloca declared *outside* a loop becomes an iter_arg unconditionally even when fully overwritten each iteration, because collectTouchedRegAllocaInRegion records on any load or store with no liveness test. Allocas declared inside the loop are correctly materialised as ub.poison; ours are all inside. Minimal repro and the IR diff are described in the docstring. Also notes that the 3-stage C store is the only change measured to reduce register pressure (256 -> 254 VGPR, 2052 -> 1573 instructions, because 128 two-byte stores need 128 addresses and 128 predicates live at once) -- relevant if anyone revisits this, though 2 VGPRs against 38 is not close. Working tree reverted; the persistent path is not kept. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 84 ++++++++++++------- 1 file changed, 55 insertions(+), 29 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 0c6ccb10c..4610a8001 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -156,7 +156,7 @@ no register cost, and the 6.7% is on the GEMM itself, independent of any all-reduce. -## Persistent tiles: tried, and not viable on FlyDSL today +## 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 @@ -175,37 +175,63 @@ 6144 384 128 47.40 7168 448 192 49.80 -So a persistent scheduler was worth trying rather than arguing about. It was -implemented (one workgroup per CU, striding over tiles) and made bit-exact, and -it is a **~2x regression**, for a reason that has nothing to do with scheduling: - - variant gemm scratch - one workgroup per tile 38.54us 0 B - body wrapped in an scf.for 74.48us 156 B/thread - -Wrapping the pipeline in a runtime loop makes FlyDSL spill the MFMA -accumulators. ``fly-promote-regmem-to-vectorssa`` does not promote -``make_rmem_tensor`` values declared inside an ``scf.for`` region, and 156 bytes -per thread of scratch in the inner pipeline costs more than the entire tail. The -persistent version does beat the *looped* baseline by ~3% (82.3 vs 84.8us at -remainder 192, 76.8 vs 79.8 at remainder 16), which matches gcnasm's numbers -- -but that is a gain over an already 94%-regressed kernel. - -Two things worth keeping from the attempt: +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. That is fine once, but in a persistent 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, - 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**, not hoisted: the pipeline swaps - those Python bindings as it advances, and hoisting makes eight - shared-address-space pointers loop-carried, which fails to legalize. - -Reviving this needs the FlyDSL promotion fix first. Until then the ceiling is -gcnasm's +0.77% at our remainder, against a 94% floor. + 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``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not From a03d5a0048539b2667ae2394f6e7ae813b3e3f3b Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 14:46:45 +0000 Subject: [PATCH 18/79] fix(cco): the fused epilogue raced on a barrier phase, not a queue Found by review, not by me: the prologue's `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 for the rest of the kernel. A one-shot GEMM does not care -- the trailing unmatched barrier is released when the other half exits, and the offset is deliberate, it staggers the two halves of M. The fused epilogue does care. `wait_barrier(0)` after store_c is supposed to mean "every wave's C tile has retired". Under the offset it rendezvouses waves 0-3 -- which contain thread 0, hence the counter and the transfer -- with waves 4-7 sitting at the *previous* barrier, before their stores. Thread 0 then counts the tile complete and pushes it while half of it is unwritten. One `if wave_m == 0: s_barrier()` before store_c re-balances the counts. --chunks 2 goes from failing about one run in three to 12 of 12, and it is the configuration that actually overlaps, so it is the default again. That changes the Step 3 conclusion. 8 ranks, [4096,7168], K=1024, median of 3-4: default C-store + 3-stage C-store split-lsa 326.8us 321.8us split-sdma 334.6 330.9 fused-sdma chunks=2 326.2 323.0 Fusing is no longer a loss: it beats split-sdma by ~2.5% and ties split-lsa. It still does not win, because the overlap (scatter 136.8 -> 126.9us) is roughly cancelled by what the epilogue costs the GEMM (39.1 -> 60.9), and reduce plus all-gather -- 46% of the pipeline -- are untouched by construction. Retracting two earlier diagnoses of this race, both wrong: - "chunks > 1 shares an SDMA queue". gcnasm's per-destination submit lock was ported and ISA-verified and did not fix anything. It is kept because cco's one-issuing-warp-per-queue rule still applies, but it was not the bug. - "the counter atomic needs its acquire half". acq_rel did appear to lower the failure rate against monotonic, and with a one-in-three 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. fused-lsa is still racy and is a separate problem -- its peer stores sit dirty in eight XCDs' L2s, which no barrier fixes. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 21 +++++++++++-------- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 15 +++++++++++++ 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index adc5bdf43..0fa3cc3de 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -44,10 +44,11 @@ Headline, 8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51: - split-lsa 324.4us - split-sdma 338.0 - fused-sdma chunks=1, no fence 348.5 - fused-lsa 341.0 RACY, see kernels_fused + default C-store + 3-stage C-store + split-lsa 326.8us 321.8us + split-sdma 334.6 330.9 + fused-sdma chunks=2 326.2 323.0 + fused-lsa 341.0 RACY, see kernels_fused Fusing loses. Read ``kernels_fused.py`` before trusting any faster fused number from this benchmark: several of its options are intermittently wrong, and a @@ -150,11 +151,13 @@ def run(args) -> int: 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") - # 1, even though >1 is what would create the overlap. A chunk boundary means - # the copy engine starts reading C *while the GEMM is still running*, and no - # combination of submit lock, post-submit barrier, release fence or atomic - # ordering made that reliable here -- see the race note in kernels_fused. - chunks = args.chunks if args.chunks else 1 + # One push per BLOCK_M row-band. This is what creates the overlap: with a + # single chunk per destination the counter only fires once the whole slice is + # done, which under the rotated tile order is the end of the GEMM, so nothing + # overlaps. It was pinned to 1 for a while because it raced; the cause was the + # half-wave barrier offset, fixed in kernels_fused's epilogue. + m_tiles_per_peer = max(1, (args.m // world_size) // args.block_m) + chunks = args.chunks if args.chunks else m_tiles_per_peer cfg = ArConfig( world_size=world_size, m=args.m, diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 4610a8001..6cc88acde 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -1265,6 +1265,21 @@ def kernel_gemm_scatter( base_row = block_m * BLOCK_M + wave_m_offset base_col = block_n * BLOCK_N + wave_n_offset + if const_expr(fuse): + # Re-balance the half-wave barrier counts before the epilogue. + # The prologue's `if wave_m == 1: s_barrier()` gives waves 4-7 one + # extra barrier, so every later rendezvous pairs w1's k-th barrier + # with w0's (k+1)-th and waves 0-3 run one phase ahead. A one-shot + # GEMM does not care -- the trailing barrier is released when the + # other half exits. The fused epilogue does: `wait_barrier(0)` after + # store_c is supposed to mean "every wave's C tile has retired", and + # under the offset it instead rendezvouses waves 0-3 (which hold + # thread 0, hence the counter and the put) with waves 4-7 sitting at + # the *previous* barrier -- before their stores. Thread 0 then counts + # the tile and can issue the transfer while half the tile is unwritten. + if wave_m == 0: + rocdl.s_barrier() + 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) From 40f4d39f6c133f07cb139fc535e86e817044b900 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 14:50:30 +0000 Subject: [PATCH 19/79] docs(cco): rewrite the race section for the barrier-phase cause; repair headers The code fix landed in the previous commit but its docstring did not -- the edit's anchor had already been mangled by an earlier replacement that left duplicated section headers behind. Both are fixed here. Replaces "The chunking race, unexplained" and "Result: the fusion does not pay" with the actual cause and the new standings, and retracts the two wrong diagnoses (shared SDMA queue, missing acquire on the counter) in place rather than deleting them, since both are things a reader would otherwise try again. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 111 ++++++++++-------- 1 file changed, 64 insertions(+), 47 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 6cc88acde..d8200a1de 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -233,7 +233,7 @@ 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``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``)## Direct LSA (``--mode fused-lsa``): the structure works, the store does not +## Direct LSA (``--mode fused-lsa``): the structure works, the store does not gcnasm's best mode has the GEMM epilogue store *straight into the destination rank's window*, so there is no staging buffer, no copy engine and nothing to @@ -277,7 +277,7 @@ rate. The prize for coalescing is correctness plus maybe a few percent, not a step change. -## What gcnasm does differently## What gcnasm does differently, and why its GEMM+a2a wins +## 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 @@ -304,51 +304,68 @@ (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 chunking race, unexplained - -Overlap requires more than one chunk per destination: with one, the counter only -fires when the whole slice is done, which under the rotated tile order is the end -of the GEMM, so nothing overlaps (measured: scatter 129.0us against split's -136.8 -- a 7.8us gain, i.e. none). - -``--chunks 2`` is worth it on paper -- 328us, which would beat split-sdma's 338 -and tie split-lsa's 324 -- and it is **intermittently wrong**, in roughly 1 run in -3. Ruled out, each by measurement rather than argument: - -* the per-destination submit lock (gcnasm's, ISA-verified: test-and-set, - ``s_sleep`` backoff, correct spin) -- still 2/6 wrong; -* the post-submit barrier gcnasm keeps -- still 4/6 wrong; -* the release fence, at both agent and system scope -- still 3/6 wrong; -* the counter atomic's ordering, ``acq_rel`` vs relaxed -- helps, does not fix. - -The one factor that separates working from broken is *when* the put is issued. At -chunks=1 it lands at the very end of the GEMM; at chunks=2 the first one is -issued mid-kernel and the engine reads C while the rest of the GEMM is still -running. gcnasm's ChunkFused path does the same thing successfully, but its GEMM -is 13x longer, so its chunk boundaries are far apart in time. Whether chunks=1 -here is *correct* or merely always-lucky is not established -- the mechanism -suggests the latter, so treat the fused path as unproven either way. - -## Result: the fusion does not pay - -Correct-as-measured configurations only, 8 ranks, [4096, 7168], K=1024: - - split-lsa 324.4us - split-sdma 338.0 - fused-sdma chunks=1, no fence 348.5 (stable 6/6) - fused-sdma chunks=1, agent fence 440.0 (the fence is 90us of waste) - -Per-kernel (rocprofv3 --kernel-trace, us): - - config gemm scatter reduce gather sum - split-sdma 39.1 136.8 13.5 137.3 326.8 - fused c=1, no fence 62.7 129.0 13.0 136.2 341.0 - fused c=1, agent fence 152.4 132.7 13.6 135.4 434.1 - -So with the fence removed the whole remaining deficit is +23.6us inside the GEMM --- the block-wide barrier, the acq_rel counter, and the copy engine reading C -while the GEMM writes it -- against a 7.8us overlap gain. The structure is sound; -the shape is wrong for it. +## The race that pinned chunks=1: a barrier phase, not a queue + +The prologue's ``if wave_m == 1: rocdl.s_barrier()`` gives waves 4-7 one *extra* +barrier. ``s_barrier`` is a counting rendezvous, so from then on every arrival +pairs waves 4-7's k-th barrier with waves 0-3's (k+1)-th, and waves 0-3 run one +phase ahead for the rest of the kernel. A one-shot GEMM does not care: the +trailing unmatched barrier is released when the other half exits, and the offset +is the point -- it staggers the two halves of the M dimension. + +The fused epilogue does care. ``wait_barrier(0)`` after ``store_c`` is supposed +to mean "every wave's C tile has retired". Under the offset it instead +rendezvouses waves 0-3 -- which contain thread 0, hence the counter and the +transfer -- with waves 4-7 sitting at the *previous* barrier, before their +stores. Thread 0 then counts the tile complete and can push it while half of it +is still unwritten. + +One ``if wave_m == 0: rocdl.s_barrier()`` before ``store_c`` re-balances the +counts and fixes it: ``--chunks 2`` goes from failing about one run in three to +12 of 12, and it is the configuration that actually overlaps. + +Two things were blamed for this before it was found, and both were wrong: + +* **A shared SDMA queue.** With two chunks per destination the modulo test elects + two winners that both post to queue ``dest``, so gcnasm's per-destination + submit lock was ported (ISA-verified: test-and-set with ``s_sleep`` backoff). + It fixed nothing. It is kept because cco's rule of at most one issuing warp per + queue still applies, but it was not the bug. +* **The counter atomic's ordering.** ``acq_rel`` appeared to lower the failure + rate against ``monotonic``, which is how the acquire half got justified. With a + one-in-three intermittent failure that comparison was noise. ``acq_rel`` stays + because release/acquire on the counter is right for a producer handing tiles to + a consumer, not because it was measured. + +Repeated runs remain the only way to judge any of this; +``test_fused_is_stable_across_repeats`` requires three runs to be *identical*, +not merely each small, because every individual relL2 here looks plausible. + +## Result: fused-sdma reaches parity, and does not beat LSA + +8 ranks, [4096, 7168], K=1024, graph replay, median of 3-4 runs: + + default C-store + 3-stage C-store + split-lsa 326.8us 321.8us + split-sdma 334.6 330.9 + fused-sdma chunks=2 326.2 323.0 + fused-lsa 341.0 racy, see above + +Per-kernel for the SDMA paths (rocprofv3 --kernel-trace, us): + + config gemm scatter/drain reduce gather sum + split-sdma 39.1 136.8 13.5 137.3 326.8 + fused-sdma chunks=2 60.9 126.9 13.1 135.1 336.1 + +So the fusion is no longer a loss -- it beats split-sdma by ~2.5% and ties +split-lsa -- but it does not win. The overlap it buys (scatter 136.8 -> 126.9) +is roughly cancelled by what the epilogue costs the GEMM (39.1 -> 60.9), and the +reduce and all-gather, 46% of the pipeline, are untouched by construction. + +An ATT thread trace puts the epilogue at 26% of the kernel's latency, almost all +of it one instruction waiting on the release fence's writeback -- 13.9% of the +kernel, more than any main-loop stall. The fence discussion under ``--fence`` +covers why that is quadratic and what did and did not help. ## Pinned copy From ddb48d834c9a91b03624ceee3a50ccead746a128 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 15:02:28 +0000 Subject: [PATCH 20/79] fix(cco): close the half-wave barrier pair unconditionally, as gcnasm does The compensating `if wave_m == 0: s_barrier()` added last commit was gated on `fuse`. It should not be. gcnasm closes the pair at opus_gemm_a2a_lsa/gemm_a16w16_quad_subtile_kernel_template.hpp:693, before store_c and *outside* its `if constexpr (ChunkFused)` -- unconditional. aiter's 8-wave kernel opens the stagger at gemm_a8w8_8wave.py:449 and never closes it, which is where the pinned copy inherited the imbalance. Gating it on `fuse` was reasoning from the symptom rather than the invariant. The split path has nothing after store_c so it cannot observe the imbalance today -- but leaving the barrier counts unbalanced is what makes the next thing added after the epilogue silently wrong, which is exactly how this bug arrived, and it is also what the persistent-tile experiment had to work around independently. Costs nothing measurable on the split path: +8 instructions, VGPR 256 -> 256, SGPR 50 -> 42, zero spill either way, gemm-only unchanged within noise. fused-sdma chunks=2 remains 6/6 stable at ~326us, and the full suite is 28/28. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 35 +++++++++++-------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index d8200a1de..4f1b07199 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -1282,20 +1282,27 @@ def kernel_gemm_scatter( base_row = block_m * BLOCK_M + wave_m_offset base_col = block_n * BLOCK_N + wave_n_offset - if const_expr(fuse): - # Re-balance the half-wave barrier counts before the epilogue. - # The prologue's `if wave_m == 1: s_barrier()` gives waves 4-7 one - # extra barrier, so every later rendezvous pairs w1's k-th barrier - # with w0's (k+1)-th and waves 0-3 run one phase ahead. A one-shot - # GEMM does not care -- the trailing barrier is released when the - # other half exits. The fused epilogue does: `wait_barrier(0)` after - # store_c is supposed to mean "every wave's C tile has retired", and - # under the offset it instead rendezvouses waves 0-3 (which hold - # thread 0, hence the counter and the put) with waves 4-7 sitting at - # the *previous* barrier -- before their stores. Thread 0 then counts - # the tile and can issue the transfer while half the tile is unwritten. - if wave_m == 0: - rocdl.s_barrier() + # 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() store_c.store(c00_frag, base_row + 0, base_col + 0) store_c.store(c01_frag, base_row + 0, base_col + LDS_BLOCK_N) From 9f8e0e7f141ec016175ca08af83400eeeb3b7fde Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 15:35:56 +0000 Subject: [PATCH 21/79] perf(cco): --hoist-scales -- the redundancy is real, removing it is not a win Review spotted that the epilogue's four store_c calls each reload the scales, and they do: the 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 8 addresses are distinct, 8 buffer_load_dword at :199 where 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, so they are a chain of overwrites rather than pure loads. store_all() loads each once. It removes exactly the predicted loads 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 two VGPRs of headroom. Bit-exact either way, and the bitwise test now covers it. Kept as a default-off switch rather than deleted: the redundant fraction grows with N_TILES_B, so a larger BLOCK_N changes the arithmetic, and it flips sign if register pressure ever loosens. Re-measuring a flag is cheaper than re-deriving why it was rejected. One honest note on the test runs: a single unidentified failure appeared once in the combined suite and did not reproduce in seven subsequent runs (full suite x3, fused stability and modes-agree x3, AR suite x2, all green). It is recorded here rather than dismissed, since an intermittent failure is exactly what this epilogue has produced twice before. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 7 ++ benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 85 ++++++++++++++++--- tests/python/cco/test_gemm_ar.py | 1 + 3 files changed, 81 insertions(+), 12 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 0fa3cc3de..edea27d2a 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -207,6 +207,7 @@ def run(args) -> int: store_probe=args.store_probe, permlane=args.permlane, lane_transpose=args.lane_transpose, + hoist_scales=args.hoist_scales, rotated=None if args.tile_order == "auto" else args.tile_order == "rotated", fence=args.fence, emit_put=not args.no_put, @@ -412,6 +413,12 @@ def build_parser() -> argparse.ArgumentParser: help="with --swap-ab --permlane: one ds_bpermute per dword so adjacent " "lanes cover one C row (gcnasm kernel_template.hpp:491-510)", ) + 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("--sdma-queues", type=int, default=8) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 4f1b07199..6c52e73bc 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -131,6 +131,33 @@ 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:: @@ -643,20 +670,46 @@ class _PermlaneStoreC(_SwapABStoreC): _lane_transpose = False def store(self, c_frag, base_row, base_col): - 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}" - ) + 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 - grp = lane // 16 - a_scales = [ + return [ self._load_a_scale_scalar(base_row + ti * 16 + lane % 16) for ti in range_constexpr(self.n_tiles_a) ] - b_scales = [ + + 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 = [] @@ -869,6 +922,7 @@ def compile_fused_gemm_scatter( store_probe: bool = False, permlane: bool = False, lane_transpose: bool = False, + hoist_scales: bool = False, transport: str = "sdma", fence: str = "none", emit_put: bool = True, @@ -973,7 +1027,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"{'S' if swap_ab else ''}{'P' if permlane else ''}{'T' if lane_transpose else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" + _kname_tag = f"{'S' if swap_ab else ''}{'P' if permlane else ''}{'T' if lane_transpose else ''}{'H' if hoist_scales else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" counter_off = cfg.counter_off lock_off = cfg.lock_off in_off = cfg.input_off @@ -1304,10 +1358,17 @@ def kernel_gemm_scatter( if wave_m == 0: rocdl.s_barrier() - 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) + 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(fuse and not direct_lsa): diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 23f4f25ed..5a1046327 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -192,6 +192,7 @@ def test_swap_ab_is_bitwise_identical(m, n, k): {}, {"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, From 0f2aa7477b0f19acf2b75a5e1bd1ef8025abf7c9 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 16:22:02 +0000 Subject: [PATCH 22/79] fix(cco): fused-lsa is correct now -- the producing kernel has to publish fused-lsa was wrong roughly half the time. The cause: in the LSA 2-stage all-reduce the producer fences (ar/kernels_lsa.py, every lane of every block, right after the tmp stores). Here the producer is the GEMM and the only fence was in the separate barrier kernel -- one block, so one XCD's L2 out of eight. The other seven kept the peer-homed lines dirty while the LSA flag, a system-scope atomic, overtook them. Fencing at the end of the GEMM fixes it: 10 of 10 at the fp8 floor. The fence can be leaderOnly, worth 80us (434 -> 355), and that is only legal 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. It measured as incorrect before that fix, which is what made a per-wave release look mandatory. Two things that did not work and are recorded rather than dropped: - uncached sc0|sc1 peer stores, so there would be nothing to flush. Consistently wrong (~1.7e-2) at 2 bytes per lane and still at 16 after the permlane stage, so the "2-byte partial-line writes lose updates" theory was wrong and the cause is unknown. - the 3-stage C store, worth 6.7% on a local store, is worth nothing here (189.4us of GEMM either way) -- the peer store is xGMI bandwidth-bound. Standings, 8 ranks, [4096,7168] K=1024, 3-stage C store, failures out of 10: split-lsa 318.8us 0/10 fused-sdma chunks=1 342 0/10 fused-sdma chunks=2 325 3/10 fastest, still racy split-sdma 329.7 0/10 fused-lsa 3-stage, leader 359.0 0/10 split-lsa still wins. Both fusions pay a separate local reduce that LSA's stage 1 gets for free, and reduce plus all-gather are 46% of the pipeline and untouched by construction. Two corrections to the previous commit, from wider sampling: - chunks=2 is NOT fixed. The barrier-phase fix took it from about 1 in 3 to 3 in 10; a second, rarer bug remains. Default goes back to 1. Adding a fence there makes it worse, not better (leader 3/6 wrong, agent 5/6), which is unexplained. - the bench's validation gate was 5e-3, loose enough to pass a corrupt run at 3.97e-3 while the fused path was racing. Tightened to 3e-3: the fp8 floor is 2.35e-3 and the corruption lands at 4-9e-3. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 29 ++- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 182 +++++++++++------- 2 files changed, 143 insertions(+), 68 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index edea27d2a..5548d1700 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -156,8 +156,11 @@ def run(args) -> int: # done, which under the rotated tile order is the end of the GEMM, so nothing # overlaps. It was pinned to 1 for a while because it raced; the cause was the # half-wave barrier offset, fixed in kernels_fused's epilogue. - m_tiles_per_peer = max(1, (args.m // world_size) // args.block_m) - chunks = args.chunks if args.chunks else m_tiles_per_peer + # 1. >1 is what creates the overlap and is measurably faster (~325us vs + # ~342), and it is still wrong 3 runs in 10 -- see the race note in + # kernels_fused. The barrier-phase fix took it from 1-in-3 to 3-in-10, not + # to zero. + chunks = args.chunks if args.chunks else 1 cfg = ArConfig( world_size=world_size, m=args.m, @@ -208,6 +211,8 @@ def run(args) -> int: 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", fence=args.fence, emit_put=not args.no_put, @@ -272,7 +277,11 @@ def once(): 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. - validated = rel_l2 < 5e-3 + # 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. + validated = rel_l2 < 3e-3 if not validated: print(f"[rank {rank}] VALIDATION FAILED relL2={rel_l2:.3e}", flush=True) del acc @@ -419,6 +428,20 @@ def build_parser() -> argparse.ArgumentParser: 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("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 6c52e73bc..79db5ed35 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -260,51 +260,50 @@ 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``): the structure works, the store does not +## Direct LSA (``--mode fused-lsa``): correct now, and still slower than split gcnasm's best mode has the GEMM epilogue store *straight into the destination -rank's window*, so there is no staging buffer, no copy engine and nothing to -publish mid-kernel. Ported here, and the structural half is emphatic -(rocprofv3 --kernel-trace, us): - - config gemm scatter/barrier reduce gather sum - split-sdma 39.1 136.8 13.5 137.3 326.8 - fused-sdma c=1 62.7 129.0 13.0 136.2 341.0 - fused-lsa (Direct) 166.6 13.3 11.1 135.5 326.6 - -The scatter collapses from 136.8us to 13.3us -- the transfer is *entirely* -absorbed, which is what fused-sdma never managed. And the un-coalesced store is -nowhere near as bad as expected: 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 is nonetheless **not correct**, and the reason makes the ``StoreC`` coalescing -rewrite a correctness prerequisite rather than an optimisation: - -* **Cached peer stores validate 2-3 runs in 6.** The lines are homed in the - peer's memory but sit dirty in whichever XCD's L2 the block ran on. Nothing - downstream can clean that up -- the barrier kernel is one block, and - ``buffer_wbl2`` writes back only the L2 of the XCD it runs on, one of eight. -* **Uncached (sc0|sc1) peer stores are consistently wrong** (relL2 ~3e-2, never - right) **and slower** (426us against 341). ``_store_bf16`` writes one bf16, so - bypassing L2 makes every store a 2-byte partial-line write; concurrent sub-32B - writes into the same sector across the fabric lose updates. Adding an uncached - recv load in the reduce does not change it, so this is the store, not a stale - read. - -So L2 merging is what makes the narrow store work at all, and L2 residency is -what makes it invisible. The two are only separable by widening the store -- -gcnasm's wave-local ``ds_bpermute`` pair-coalesced C store (their "C-store mode -2"; staging C through LDS regressed 1-3%). That is why they needed it. - -Even coalesced, this does not obviously win. The fused sum is 326.6us against -split-lsa's ~324: LSA's 2-stage does read + reduce + write in a single pass over -the wire, so its reduce is free, whereas Direct LSA writes to the peer and then -pays a separate local reduce pass -- and its store rate is ~18% under LSA's read -rate. The prize for coalescing is correctness plus maybe a few percent, not a -step change. - -## What gcnasm does differently, and why its GEMM+a2a wins +rank's window*, so there is no staging buffer, no copy engine and nothing to hand +to one. The structural half is emphatic -- the scatter collapses from 136.8us to +about 10us, the transfer is entirely absorbed -- and it was wrong for a long time. + +What was missing: **the producing kernel has to publish.** In the LSA 2-stage +all-reduce the producer fences (``ar/kernels_lsa.py``, every lane of every block, +right after the ``tmp`` stores). Here the producer is the GEMM, and the only +fence was in the separate barrier kernel -- one block, so one XCD's L2 out of +eight. The other seven kept the peer-homed lines dirty while the LSA flag, being +a system-scope atomic, overtook them. Adding ``cco_system_fence`` at the end of +the GEMM fixes it: 10 runs of 10 at the fp8 floor. + +The fence can be ``leaderOnly``, which is worth 80us (434us -> 355). That is only +legal because the half-wave barrier pair is closed (see below): ``wait_barrier(0)`` +now 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, both worth recording because they look obvious: + +* **Uncached (``sc0|sc1``) peer stores**, so there would be nothing to flush. + Consistently wrong, ~1.7e-2, at 2 bytes per lane and still at 16 after the + permlane stage -- so the earlier "2-byte partial-line writes lose updates" + explanation was not it, and the cause is unknown. +* **The 3-stage C store**, which is worth 6.7% on a local store, is worth nothing + here: 189.4us of GEMM either way. The peer store crosses the fabric and that + phase is bounded by ~44 GB/s of link bandwidth, not by issue rate or + coalescing. + +Per-kernel, 8 ranks, [4096, 7168], K=1024, with the 3-stage C store: + + config gemm barrier reduce gather sum e2e + split-lsa 35.7 (one all-reduce kernel, 274.3) 318.8us + fused-lsa 189.4 10.1 10.9 135.6 346.0 359.0us + +So it is correct and 12% slower. The reason is structural: LSA's stage 1 does +read + reduce + write in a single pass over the wire, so its reduce is free, +while Direct LSA writes to the peer and then pays a separate local reduce pass -- +and its store rate is ~44 GB/s against LSA's ~51 GB/s read rate. + +## What gcnasm does differently## 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 @@ -368,31 +367,39 @@ ``test_fused_is_stable_across_repeats`` requires three runs to be *identical*, not merely each small, because every individual relL2 here looks plausible. -## Result: fused-sdma reaches parity, and does not beat LSA +## Result: split-lsa still wins -8 ranks, [4096, 7168], K=1024, graph replay, median of 3-4 runs: +8 ranks, [4096, 7168], K=1024, graph replay, with the 3-stage C store. Failures +are out of 10 runs, counting any result that is not the 2.35e-3 fp8 floor: - default C-store + 3-stage C-store - split-lsa 326.8us 321.8us - split-sdma 334.6 330.9 - fused-sdma chunks=2 326.2 323.0 - fused-lsa 341.0 racy, see above + mode time wrong + split-lsa 318.8us 0/10 + fused-sdma chunks=1 342 0/10 + fused-sdma chunks=2 325 3/10 <- fastest, still racy + split-sdma 329.7 0/10 + fused-lsa 3-stage, leader 359.0 0/10 -Per-kernel for the SDMA paths (rocprofv3 --kernel-trace, us): +Per-kernel: - config gemm scatter/drain reduce gather sum - split-sdma 39.1 136.8 13.5 137.3 326.8 - fused-sdma chunks=2 60.9 126.9 13.1 135.1 336.1 + config gemm scatter/barrier reduce gather sum + split-sdma 35.7 136.8 13.5 135.5 321.5 + fused-sdma chunks=1 57.6 130.4 13.1 135.1 336.2 + fused-lsa 189.4 10.1 10.9 135.6 346.0 -So the fusion is no longer a loss -- it beats split-sdma by ~2.5% and ties -split-lsa -- but it does not win. The overlap it buys (scatter 136.8 -> 126.9) -is roughly cancelled by what the epilogue costs the GEMM (39.1 -> 60.9), and the -reduce and all-gather, 46% of the pipeline, are untouched by construction. +Neither fusion beats ``split-lsa``, and the reason is the same for both: LSA's +2-stage does read + reduce + write in one pass over the wire, so its reduce is +free, while every fused variant here writes somewhere and then pays a separate +local reduce. On top of that ``fused-lsa``'s peer store runs at ~44 GB/s against +LSA's ~51 GB/s read, and ``fused-sdma``'s overlap (scatter 136.8 -> 130.4 at +chunks=1) is cancelled by what the epilogue costs the GEMM (35.7 -> 57.6). -An ATT thread trace puts the epilogue at 26% of the kernel's latency, almost all -of it one instruction waiting on the release fence's writeback -- 13.9% of the -kernel, more than any main-loop stall. The fence discussion under ``--fence`` -covers why that is quadratic and what did and did not help. +``--chunks 2`` is the one configuration that is genuinely faster than the best +split path would suggest -- 325us against split-sdma's 329.7 -- and it is wrong +3 runs in 10. The barrier-phase fix took it from about 1 in 3 to 3 in 10, so +whatever remains is a second, rarer bug. Default is 1. + +Reduce and all-gather are 46% of the pipeline and untouched by any of this, +which caps what fusing the scatter can ever be worth. ## Pinned copy @@ -449,6 +456,7 @@ from flydsl._mlir.dialects import scf # noqa: E402 from _compat import ( # noqa: E402 + CM_CACHED, CM_SC0_SC1, atomic_add_u32, atomic_store_u32, @@ -668,6 +676,7 @@ class _PermlaneStoreC(_SwapABStoreC): _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, @@ -758,7 +767,20 @@ def _emit(self, c_frag, base_row, base_col, a_scales, b_scales): 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: - buffer_store(out8, self._peer_rsrc, fx.Int32(idx) - self._elem_base) + # 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( @@ -773,6 +795,12 @@ class _LaneTransposeStoreC(_PermlaneStoreC): _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``.""" @@ -923,6 +951,8 @@ def compile_fused_gemm_scatter( 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, @@ -1027,7 +1057,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"{'S' if swap_ab else ''}{'P' if permlane else ''}{'T' if lane_transpose else ''}{'H' if hoist_scales else ''}{'W' if store_probe else ''}c{chunks}{'r' if rotated else 'l'}{fence[0]}" + _kname_tag = f"{'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'}{fence[0]}" counter_off = cfg.counter_off lock_off = cfg.lock_off in_off = cfg.input_off @@ -1157,7 +1187,12 @@ def kernel_gemm_scatter( ) elif const_expr(direct_lsa and swap_ab and permlane and lane_transpose): dest_blk = block_m // fx.Int32(m_tiles_per_peer) - store_c = _LaneTransposeStoreC( + _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)), @@ -1371,6 +1406,23 @@ def kernel_gemm_scatter( 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); From 1c661cb89eff3514f7b4544aa32527df9cf0126e Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 8 Sep 2026 16:25:18 +0000 Subject: [PATCH 23/79] fix(cco): fused-lsa is correct now -- publish from the producing kernel Direct LSA had been wrong 2-3 runs in 6 since it was written. The fix is one fence 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 lived in the separate barrier kernel -- one block, so one XCD's L2 out of eight. The other seven kept the peer-homed lines dirty and the LSA flag, a system-scope atomic, overtook them. Fencing in the GEMM gives 10/10 bit-correct runs. --direct-fence leader (thread 0 only) 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. It measured as incorrect before that fix, which is what made a per-wave release look mandatory. Two things that did not work, recorded so they are not retried: uncached (sc0|sc1) peer stores are consistently wrong at *every* store width, including 16B/lane with --permlane --lane-transpose, which kills the partial-line-write explanation and leaves the real one unknown; and an uncached recv load in the reduce changes nothing, which rules out a stale read on the consumer side. It still does not beat split-lsa. 8 ranks, [4096,7168], K=1024, all three C-store stages, median of 4, correctness over 10 runs: mode time correct split-lsa 318.8us 10/10 <- best fused-sdma chunks=2 325.0 7/10 racy split-sdma 329.7 10/10 fused-sdma chunks=1 342.0 10/10 fused-lsa 359.0 10/10 The 41us gap for fused-lsa is structural: LSA's 2-stage does read+reduce+write in one pass over the wire so its reduce is free, while Direct LSA writes to the peer and then pays a separate 10.9us reduce, runs its stores ~14% under LSA's read rate, and adds the publishing fence LSA gets for free by being local. Absorbing the scatter completely (136.8 -> 10.1us) still does not cover that. Also corrects two things about fused-sdma. --chunks 2 is *not* fixed: the barrier-phase fix took it from 1-in-3 to 3-in-10, so the default goes back to 1 (10/10, ~17us slower). And the bench's validation gate was too loose -- the corruption lands at 4-9e-3 against an fp8 floor of 2.35e-3, so 5e-3 reported a corrupt run as validated, which happened at 3.97e-3. Gate is 3e-3 now. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 210 +++++++++--------- 1 file changed, 104 insertions(+), 106 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 79db5ed35..29625f530 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -260,48 +260,56 @@ 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, and still slower than split +## 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*, so there is no staging buffer, no copy engine and nothing to hand -to one. The structural half is emphatic -- the scatter collapses from 136.8us to -about 10us, the transfer is entirely absorbed -- and it was wrong for a long time. - -What was missing: **the producing kernel has to publish.** In the LSA 2-stage -all-reduce the producer fences (``ar/kernels_lsa.py``, every lane of every block, -right after the ``tmp`` stores). Here the producer is the GEMM, and the only -fence was in the separate barrier kernel -- one block, so one XCD's L2 out of -eight. The other seven kept the peer-homed lines dirty while the LSA flag, being -a system-scope atomic, overtook them. Adding ``cco_system_fence`` at the end of -the GEMM fixes it: 10 runs of 10 at the fp8 floor. - -The fence can be ``leaderOnly``, which is worth 80us (434us -> 355). That is only -legal because the half-wave barrier pair is closed (see below): ``wait_barrier(0)`` -now 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, both worth recording because they look obvious: - -* **Uncached (``sc0|sc1``) peer stores**, so there would be nothing to flush. - Consistently wrong, ~1.7e-2, at 2 bytes per lane and still at 16 after the - permlane stage -- so the earlier "2-byte partial-line writes lose updates" - explanation was not it, and the cause is unknown. -* **The 3-stage C store**, which is worth 6.7% on a local store, is worth nothing - here: 189.4us of GEMM either way. The peer store crosses the fabric and that - phase is bounded by ~44 GB/s of link bandwidth, not by issue rate or - coalescing. - -Per-kernel, 8 ranks, [4096, 7168], K=1024, with the 3-stage C store: - - config gemm barrier reduce gather sum e2e - split-lsa 35.7 (one all-reduce kernel, 274.3) 318.8us - fused-lsa 189.4 10.1 10.9 135.6 346.0 359.0us - -So it is correct and 12% slower. The reason is structural: LSA's stage 1 does -read + reduce + write in a single pass over the wire, so its reduce is free, -while Direct LSA writes to the peer and then pays a separate local reduce pass -- -and its store rate is ~44 GB/s against LSA's ~51 GB/s read rate. +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## What gcnasm does differently, and why its GEMM+a2a wins @@ -330,76 +338,66 @@ (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 race that pinned chunks=1: a barrier phase, not a queue +## The chunks race: much better understood, still not gone The prologue's ``if wave_m == 1: rocdl.s_barrier()`` gives waves 4-7 one *extra* barrier. ``s_barrier`` is a counting rendezvous, so from then on every arrival pairs waves 4-7's k-th barrier with waves 0-3's (k+1)-th, and waves 0-3 run one -phase ahead for the rest of the kernel. A one-shot GEMM does not care: the -trailing unmatched barrier is released when the other half exits, and the offset -is the point -- it staggers the two halves of the M dimension. - -The fused epilogue does care. ``wait_barrier(0)`` after ``store_c`` is supposed -to mean "every wave's C tile has retired". Under the offset it instead -rendezvouses waves 0-3 -- which contain thread 0, hence the counter and the -transfer -- with waves 4-7 sitting at the *previous* barrier, before their -stores. Thread 0 then counts the tile complete and can push it while half of it -is still unwritten. - -One ``if wave_m == 0: rocdl.s_barrier()`` before ``store_c`` re-balances the -counts and fixes it: ``--chunks 2`` goes from failing about one run in three to -12 of 12, and it is the configuration that actually overlaps. - -Two things were blamed for this before it was found, and both were wrong: - -* **A shared SDMA queue.** With two chunks per destination the modulo test elects - two winners that both post to queue ``dest``, so gcnasm's per-destination - submit lock was ported (ISA-verified: test-and-set with ``s_sleep`` backoff). - It fixed nothing. It is kept because cco's rule of at most one issuing warp per - queue still applies, but it was not the bug. -* **The counter atomic's ordering.** ``acq_rel`` appeared to lower the failure - rate against ``monotonic``, which is how the acquire half got justified. With a - one-in-three intermittent failure that comparison was noise. ``acq_rel`` stays - because release/acquire on the counter is right for a producer handing tiles to - a consumer, not because it was measured. - -Repeated runs remain the only way to judge any of this; -``test_fused_is_stable_across_repeats`` requires three runs to be *identical*, -not merely each small, because every individual relL2 here looks plausible. - -## Result: split-lsa still wins - -8 ranks, [4096, 7168], K=1024, graph replay, with the 3-stage C store. Failures -are out of 10 runs, counting any result that is not the 2.35e-3 fp8 floor: - - mode time wrong - split-lsa 318.8us 0/10 - fused-sdma chunks=1 342 0/10 - fused-sdma chunks=2 325 3/10 <- fastest, still racy - split-sdma 329.7 0/10 - fused-lsa 3-stage, leader 359.0 0/10 - -Per-kernel: - - config gemm scatter/barrier reduce gather sum - split-sdma 35.7 136.8 13.5 135.5 321.5 - fused-sdma chunks=1 57.6 130.4 13.1 135.1 336.2 - fused-lsa 189.4 10.1 10.9 135.6 346.0 - -Neither fusion beats ``split-lsa``, and the reason is the same for both: LSA's -2-stage does read + reduce + write in one pass over the wire, so its reduce is -free, while every fused variant here writes somewhere and then pays a separate -local reduce. On top of that ``fused-lsa``'s peer store runs at ~44 GB/s against -LSA's ~51 GB/s read, and ``fused-sdma``'s overlap (scatter 136.8 -> 130.4 at -chunks=1) is cancelled by what the epilogue costs the GEMM (35.7 -> 57.6). - -``--chunks 2`` is the one configuration that is genuinely faster than the best -split path would suggest -- 325us against split-sdma's 329.7 -- and it is wrong -3 runs in 10. The barrier-phase fix took it from about 1 in 3 to 3 in 10, so -whatever remains is a second, rarer bug. Default is 1. - -Reduce and all-gather are 46% of the pipeline and untouched by any of this, -which caps what fusing the scatter can ever be worth. +phase ahead. A one-shot GEMM does not care -- the trailing unmatched barrier is +released when the other half exits, and the offset is the point, it staggers the +two halves of M. The fused epilogue does care: ``wait_barrier(0)`` after +``store_c`` is supposed to mean "every wave's C tile has retired", and under the +offset it rendezvouses waves 0-3, which hold thread 0 and therefore the counter +and the transfer, with waves 4-7 still at the *previous* barrier, before their +stores. + +Closing the pair (``if wave_m == 0: s_barrier()`` before ``store_c``, which is +what gcnasm does at kernel_template.hpp:693, unconditionally) took ``--chunks 2`` +from failing about 1 run in 3 to **3 in 10**. Better, not fixed, so the default +stays 1 -- which is 10/10 over the same sample and costs ~17us, since with one +chunk the counter only fires when the whole slice is done and nothing overlaps. + +What is left is not the fence. Adding one makes it *worse*, reproducibly: + + --fence none (default) 6/6 correct + --fence leader 3/6 wrong + --fence agent 5/6 wrong + +Two earlier diagnoses of this race 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 (test-and-set, ``s_sleep`` backoff). It fixed nothing. Kept + 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. + +Repeated runs are the only way to judge any of this, and the gate has to be +tight: the corruption lands 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 now gates at 3e-3, and ``test_fused_is_stable_across_repeats`` requires +three runs to be *identical* rather than each small. + +## Result: nothing beats split-lsa + +8 ranks, [4096, 7168], K=1024, graph replay, all three C-store stages on, +median of 4, correctness over 10 runs: + + mode time correct + split-lsa 318.8us 10/10 <- best + fused-sdma chunks=2 325.0 7/10 racy + split-sdma 329.7 10/10 + fused-sdma chunks=1 342.0 10/10 + fused-lsa 359.0 10/10 + +Both transports were taken to correctness and neither wins. fused-sdma's overlap +is real but small (scatter 136.8 -> 126.9) and roughly cancelled by what the +epilogue costs the GEMM; fused-lsa absorbs the scatter completely (136.8 -> 10.1) +and loses it again to a separate reduce pass, a lower store rate and the +publishing fence. And in both, reduce plus all-gather -- 46% of the pipeline -- +are untouchable by construction. ## Pinned copy From 69c6896e47e6dd0edbb997a3847a49a05ff0d083 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 09:41:36 +0000 Subject: [PATCH 24/79] fix(cco): the 8wave GEMM under-waits on one s_waitcnt, corrupting large grids At the real DSV4-Pro prefill shape ([16384, 7168] K=2048) none of the GEMM+AR modes validated -- relL2 4e-3 to 2e-2, different on every run. The all-reduce was innocent: bench_ar.py is bit-exact at the same 224 MiB payload for both LSA and SDMA. The GEMM is what is wrong, and aiter's own flydsl_8wave_gemm_a8 reproduces it, so this is upstream and not a transcription error in the pinned copy. The main loop closes with wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) # s_waitcnt vmcnt(n); s_barrier which leaves too many global->LDS prefetches in flight across the barrier, so a wave passes it and refills an LDS buffer another wave has not finished consuming. The corruption is always a whole (wave_m sub-tile x all four wave_n) region -- the signature of a shared A-side LDS buffer -- and the values are wrong rather than zero, so it does not look like missing output. The trigger is grid size, not kernel duration: clean at 896 blocks even with K=16384, corrupt at 1792 blocks even with K=512. Measured boundary at [16384, 7168] K=512 on 8x MI355X, 4-5 runs per count: 256x256 corrupts at 6 (aiter's), 5 and 4, and is clean at 3, 2 and 0; 128x256 corrupts at 4 (aiter's) and 3, and is clean at 2, 1 and 0. Both thresholds are N_LDS_STEPS_A + N_LDS_STEPS_B - 1, which is what this commit uses. It costs nothing -- 241/240/245 us against 240/246/237 us for aiter's count at [16384, 7168] K=2048 -- and takes all four modes to 5/5 correct at relL2 2.350e-3, the fp8 floor. Not a complete theory, and the comment at the site says so: 128x512 never reproduces even with aiter's count, and forcing *any* of the 14 sync points in the loop to vmcnt(0) also fixes it, so the same under-wait may exist elsewhere. This is the smallest change that is measurably correct here. Still to do: push it upstream to aiter. Also add validation to --mode gemm-only. Without that control a GEMM bug reads as an all-reduce bug, because every other mode validates the *sum* -- which is why this took as long as it did to find. --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 15 +++++++++- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 28 ++++++++++++++++++- 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 5548d1700..2273c0b2e 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -262,7 +262,20 @@ def once(): rel_l2 = float("nan") validated = True - if ( + 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 = (a.float() @ b.float().T) * sa[:, None] * sb[None, :] + 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 diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 29625f530..895525f8c 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -413,6 +413,13 @@ ``tests/python/cco/test_gemm_ar.py`` asserts this kernel's C output is bit-identical to aiter's unfused kernel on the same inputs, which is what stops the copy from rotting silently. + +The copy carries **one deliberate divergence** from aiter, marked at its site in +the main loop: the final ``wait_barrier`` count. aiter's is too permissive and +corrupts the output non-deterministically on large grids; the test above still +passes because it runs shapes small enough that aiter is correct there too. Do +not "restore" it when re-syncing the copy -- see the comment at that line, and +push the fix upstream instead. """ # NOTE: no `from __future__ import annotations` here, deliberately, and aiter's @@ -1307,7 +1314,26 @@ def kernel_gemm_scatter( 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) + # 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) From 73c08659e4cb3a73fe52a53eaa0f73d5fc018fa7 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 10:31:18 +0000 Subject: [PATCH 25/79] perf(cco): turn the three C-store stages on by default in the benchmark --swap-ab, --permlane and --lane-transpose were opt-in because they were added one rung at a time and each needed its own A/B. Measured together at the shape the model actually runs -- [16384, 7168] K=2048, 8x MI355X, median of 31 -- they help every mode and change no output (all still relL2 2.350e-3): gemm-only 243.0 -> 230.8us -5.0% fused-sdma 1342.6 -> 1299.9 -3.2% fused-lsa 1600.9 -> 1563.2 -2.4% split-lsa 1275.2 -> 1263.8 -0.9% split-sdma 1269.8 -> 1260.8 -0.7% They are now BooleanOptionalAction with default True, so a rung can still be taken out with its --no- form. The dependency errors name that form, since hitting them now means "--no-swap-ab without --no-permlane" rather than the old "--permlane without --swap-ab". Where the win is, from the per-kernel trace, is worth writing down because it is not where the ATT profile first suggests. The stages coalesce the store *issue*, so they pay off when C lands in local memory -- fused-sdma's GEMM goes 297.3 -> 256.6us -- and do nothing when it goes straight to a peer: fused-lsa's GEMM is 954.5 -> 952.1, i.e. noise, even though 56% of that kernel is in the C-store path. That path is xGMI-bound, and ATT already said so: 99% of its store time is stall, not issue. Fewer, wider stores do not shorten a wait for the link. compile_fused_gemm_scatter keeps all three defaulting to False. Its default has to stay "aiter's kernel verbatim" -- that is the premise of test_pinned_copy_matches_aiter_kernel_bitwise, and the ladder is enumerated explicitly by test_swap_ab_is_bitwise_identical. Also refresh the module headline, which was still quoting [4096, 7168] K=1024 -- a shape that does not occur in the model. --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 72 ++++++++++++++----- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 13 +++- 2 files changed, 64 insertions(+), 21 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 2273c0b2e..8536f4152 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -42,18 +42,32 @@ 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, 8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51: - - default C-store + 3-stage C-store - split-lsa 326.8us 321.8us - split-sdma 334.6 330.9 - fused-sdma chunks=2 326.2 323.0 - fused-lsa 341.0 RACY, see kernels_fused - -Fusing loses. Read ``kernels_fused.py`` before trusting any faster fused number -from this benchmark: several of its options are intermittently wrong, and a -single passing run proves nothing. ``--chunks`` > 1 and every ``raw-wt`` fence -mode are known-racy and kept only to reproduce that. +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: + + 3-stage C-store without it (--no-swap-ab ...) + split-sdma 1260.8us 1269.8us + split-lsa 1263.8 1275.2 + fused-sdma 1299.9 1342.6 + fused-lsa 1563.2 1600.9 + gemm-only 230.8 243.0 + +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. + +Fusing loses, in both transports. The per-kernel breakdown says why: at this +shape ``fused-sdma`` overlaps *nothing* -- its drain costs 502.6us against the +split scatter's 499.0 -- and it pays 76us to issue the puts, while +``fused-lsa`` spends 730us of its GEMM pushing C over xGMI where the copy +engines do the same bytes in 499. + +Read ``kernels_fused.py`` before trusting any faster fused number from this +benchmark: several of its options are intermittently wrong, and a single +passing run proves nothing. ``--chunks`` > 1 and every ``raw-wt`` fence mode +are known-racy and kept only to reproduce that. """ from __future__ import annotations @@ -325,6 +339,8 @@ def once(): "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, "max_rank_time_us": per_rank[max_rank], @@ -408,12 +424,29 @@ def build_parser() -> argparse.ArgumentParser: "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: there the + # default has to stay "aiter's kernel verbatim", which is what + # ``test_pinned_copy_matches_aiter_kernel_bitwise`` checks. p.add_argument( "--swap-ab", - action="store_true", + 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)", + "(gcnasm mfma_adaptor_swap_ab). On by default", ) p.add_argument( "--store-probe", @@ -424,16 +457,19 @@ def build_parser() -> argparse.ArgumentParser: ) p.add_argument( "--permlane", - action="store_true", + 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", + "--store-probe prices. On by default", ) p.add_argument( "--lane-transpose", - action="store_true", + 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)", + "lanes cover one C row (gcnasm kernel_template.hpp:491-510). On by " + "default", ) p.add_argument( "--hoist-scales", diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 895525f8c..bd177cef0 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -989,11 +989,18 @@ def compile_fused_gemm_scatter( # 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") + 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") + 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") + 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 " From 2cd72de10b79e592ec69c0ed3314d6b3c274db97 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 13:40:54 +0000 Subject: [PATCH 26/79] perf(cco): unblock --chunks, and fused-sdma finally beats split --chunks > 1 is the whole overlap mechanism -- with one chunk per destination the tile counter only fires when the destination's entire slice is done, which under the rotated tile order is the end of the GEMM, so nothing overlaps. It was pinned to 1 because it produced wrong output about 3 runs in 10. That was never the chunk protocol. It was aiter's 8-wave GEMM under-counting one s_waitcnt in its main loop (fixed two commits ago), which corrupted output non-deterministically on large grids whatever the epilogue did. At [16384, 7168] K=2048 on 8 ranks, chunks 1/2/4/8 are now 40/40 correct with no hang, and chunks=8 alone is 25/25. Default is now 8, halved until it divides the M-tiles per destination so small M still runs (M=2048 -> 1, M=4096 -> 2, M=16384 -> 8). Per-kernel medians, [16384, 7168] K=2048: 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. Finer is worse: at 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. Which flips the headline. At the real prefill shape: fused-sdma chunks=8 1114.1us <- best split-sdma 1261.7 split-lsa 1262.4 fused-lsa 1571.4 gemm-only 228.9 11.7% under split, and 25.3% under the 1491.3us the same layer costs in the running model (GEMM 348.0 + NCCL 1143.3, medians over its 61 layers). fused-lsa is unchanged and still loses: 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 the C-store coalescing buys it nothing. Two hangs were seen at 8 ranks while the sweep was being set up and never reproduced in the 100+ runs after. The submit lock is a plain test-and-set spin, so this is not proof it cannot happen; the notes say to use timeout when sweeping and to treat a hang as a finding. Also rewrite the two stale sections this invalidates: the "chunks race" note in kernels_fused, which concluded the race was unsolved and the default had to stay 1, and the headline in bench_gemm_ar, which concluded fusing loses. --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 83 +++++++---- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 132 ++++++++++-------- 2 files changed, 128 insertions(+), 87 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 8536f4152..4fab770e3 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -45,29 +45,30 @@ 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: +31, max over ranks, on an idle box, with the current defaults: - 3-stage C-store without it (--no-swap-ab ...) - split-sdma 1260.8us 1269.8us - split-lsa 1263.8 1275.2 - fused-sdma 1299.9 1342.6 - fused-lsa 1563.2 1600.9 - gemm-only 230.8 243.0 + fused-sdma (chunks=8) 1114.1us + split-sdma 1261.7 + split-lsa 1262.4 + fused-lsa 1571.4 + gemm-only 228.9 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. +1043us for LSA / 1047 for SDMA against 1113 for NCCL. So ``fused-sdma`` is +**11.7% under split and 25.3% under what the model runs today**. -Fusing loses, in both transports. The per-kernel breakdown says why: at this -shape ``fused-sdma`` overlaps *nothing* -- its drain costs 502.6us against the -split scatter's 499.0 -- and it pays 76us to issue the puts, while -``fused-lsa`` spends 730us of its GEMM pushing C over xGMI where the copy -engines do the same bytes in 499. +Fusing over SDMA wins, and only because of ``--chunks``; see the table at its +definition in ``run()``. It was pinned to 1 while the aiter 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 badly: it spends 730us of its GEMM pushing C over xGMI, where the copy +engines move the same bytes in 499. -Read ``kernels_fused.py`` before trusting any faster fused number from this +Read ``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. ``--chunks`` > 1 and every ``raw-wt`` fence mode -are known-racy and kept only to reproduce that. +passing run proves nothing. Every ``raw-wt`` fence mode is known-racy and kept +only to reproduce that. """ from __future__ import annotations @@ -165,16 +166,41 @@ def run(args) -> int: 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") - # One push per BLOCK_M row-band. This is what creates the overlap: with a - # single chunk per destination the counter only fires once the whole slice is - # done, which under the rotated tile order is the end of the GEMM, so nothing - # overlaps. It was pinned to 1 for a while because it raced; the cause was the - # half-wave barrier offset, fixed in kernels_fused's epilogue. - # 1. >1 is what creates the overlap and is measurably faster (~325us vs - # ~342), and it is still wrong 3 runs in 10 -- see the race note in - # kernels_fused. The barrier-phase fix took it from 1-in-3 to 3-in-10, not - # to zero. - chunks = args.chunks if args.chunks else 1 + # 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, @@ -379,8 +405,9 @@ def build_parser() -> argparse.ArgumentParser: "--chunks", type=int, default=0, - help="pushes per destination (0 = one per BLOCK_M row-band). >1 is what " - "produces the overlap, and requires the submit lock", + 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( "--tile-order", diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index bd177cef0..0ad72a640 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -338,66 +338,80 @@ (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: much better understood, still not gone - -The prologue's ``if wave_m == 1: rocdl.s_barrier()`` gives waves 4-7 one *extra* -barrier. ``s_barrier`` is a counting rendezvous, so from then on every arrival -pairs waves 4-7's k-th barrier with waves 0-3's (k+1)-th, and waves 0-3 run one -phase ahead. A one-shot GEMM does not care -- the trailing unmatched barrier is -released when the other half exits, and the offset is the point, it staggers the -two halves of M. The fused epilogue does care: ``wait_barrier(0)`` after -``store_c`` is supposed to mean "every wave's C tile has retired", and under the -offset it rendezvouses waves 0-3, which hold thread 0 and therefore the counter -and the transfer, with waves 4-7 still at the *previous* barrier, before their -stores. - -Closing the pair (``if wave_m == 0: s_barrier()`` before ``store_c``, which is -what gcnasm does at kernel_template.hpp:693, unconditionally) took ``--chunks 2`` -from failing about 1 run in 3 to **3 in 10**. Better, not fixed, so the default -stays 1 -- which is 10/10 over the same sample and costs ~17us, since with one -chunk the counter only fires when the whole slice is done and nothing overlaps. - -What is left is not the fence. Adding one makes it *worse*, reproducibly: - - --fence none (default) 6/6 correct - --fence leader 3/6 wrong - --fence agent 5/6 wrong - -Two earlier diagnoses of this race 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 (test-and-set, ``s_sleep`` backoff). It fixed nothing. Kept - 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. - -Repeated runs are the only way to judge any of this, and the gate has to be -tight: the corruption lands at 4-9e-3 against an fp8 floor of 2.35e-3, so a +## 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 now gates at 3e-3, and ``test_fused_is_stable_across_repeats`` requires -three runs to be *identical* rather than each small. - -## Result: nothing beats split-lsa - -8 ranks, [4096, 7168], K=1024, graph replay, all three C-store stages on, -median of 4, correctness over 10 runs: - - mode time correct - split-lsa 318.8us 10/10 <- best - fused-sdma chunks=2 325.0 7/10 racy - split-sdma 329.7 10/10 - fused-sdma chunks=1 342.0 10/10 - fused-lsa 359.0 10/10 - -Both transports were taken to correctness and neither wins. fused-sdma's overlap -is real but small (scatter 136.8 -> 126.9) and roughly cancelled by what the -epilogue costs the GEMM; fused-lsa absorbs the scatter completely (136.8 -> 10.1) -and loses it again to a separate reduce pass, a lower store rate and the -publishing fence. And in both, reduce plus all-gather -- 46% of the pipeline -- -are untouchable by construction. +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. ## Pinned copy From 5cd52507657e7dbe6c112bfa1e719a7030776177 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 14:47:55 +0000 Subject: [PATCH 27/79] perf(cco): port gcnasm's striped tile order, -12.5% on fused-lsa The destination-rotated chunk-major order we already had is gcnasm's ChunkFused ordering, transposed -- its destination is on N and ours is on M, so its `m_in_chunk -> peer -> inner` is our `n -> dest -> m_tile`. Worth confirming what that is worth now that chunks actually overlap: at [16384, 7168] K=2048, rotated 1114.7us against linear 1323.0, i.e. 18.6%. At chunks=1 it was worth almost nothing, which is the same statement as "without rotation the chunks buy nothing" -- linear at chunks=8 (1323.0) is no better than chunks=1 (1299.6). What we did *not* have is gcnasm's other variant, opus_direct_stripe_tile: it rotates the destination every `m_stripe` tiles of the other axis (2) rather than after a whole chunk. gcnasm applies it only to Direct LSA and keeps chunk-major for its chunked-SDMA path, and that split reproduces exactly: mode n_stripe=1 n_stripe=2 fused-lsa 1560.7us 1366.3 -12.5% fused-sdma 1116.0 1113.7 noise Per kernel, all of fused-lsa's gain is in the GEMM -- 937.9 -> 788.8us, -15.9% -- plus 36us of rank skew off the barrier that follows it. reduce and gather do not move. The mechanism is what the comment says: with chunk-major, the 256 resident blocks retire their peer stores in bursts of 28 to the same destination; striping by 2 spreads them across all seven links. The cost is that a chunk completes later, which is exactly why it is wrong for the chunked path and why gcnasm does not use it there either. 4 is worse than 2 (1392.1). So `n_stripe` defaults to 2 for fused-lsa and 1 elsewhere, mirroring gcnasm. `--n-stripe N` overrides; 0 means the per-mode default. fused-lsa is still last (1365.6 against split-lsa's 1262.4 and fused-sdma's 1114.5), but the gap to split is down from +24% to +8%. 10/10 correct at the new default, and the test suite passes. --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 17 ++++++--- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 36 +++++++++++++++++-- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 4fab770e3..9e6d1720e 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -47,10 +47,10 @@ ``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: - fused-sdma (chunks=8) 1114.1us + fused-sdma (chunks=8) 1114.5us split-sdma 1261.7 split-lsa 1262.4 - fused-lsa 1571.4 + fused-lsa (n_stripe=2) 1365.6 gemm-only 228.9 For scale, the same layer in the running model costs 1491.3us (GEMM 348.0 + @@ -62,8 +62,8 @@ definition in ``run()``. It was pinned to 1 while the aiter 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 badly: it spends 730us of its GEMM pushing C over xGMI, where the copy -engines move the same bytes in 499. +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 ``kernels_fused.py`` before trusting any fused number from this benchmark: several of its options are intermittently wrong, and a single @@ -254,6 +254,7 @@ def run(args) -> int: 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, @@ -409,6 +410,14 @@ def build_parser() -> argparse.ArgumentParser: "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, which is 2 for fused-lsa and 1 elsewhere", + ) p.add_argument( "--tile-order", choices=("auto", "rotated", "linear"), diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 0ad72a640..0e1eea17b 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -965,6 +965,7 @@ def compile_fused_gemm_scatter( xcd_swizzle: int = 0, fuse: bool = True, rotated: bool | None = None, + n_stripe: int | None = None, swap_ab: bool = False, store_probe: bool = False, permlane: bool = False, @@ -989,6 +990,16 @@ def compile_fused_gemm_scatter( 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`` > 1 rotates the destination every ``n_stripe`` N-tiles instead + of after a whole chunk -- 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 defaults to 2 for ``fused-lsa`` and 1 everywhere else, + which is the same split gcnasm makes: with a PUT to trigger you want a + chunk finished as early as possible, and with peer stores you want them + spread evenly over the links rather than arriving in same-destination + bursts. Measured at [16384, 7168] K=2048 on 8 ranks: fused-lsa 1560.7 -> + 1366.3us (its GEMM 937.9 -> 788.8), fused-sdma unchanged within noise. + Returns ``launch(A, B_T, C, A_scale, B_scale, c_m, c_n, dev_comm, win, stream=...)``. """ @@ -1028,6 +1039,12 @@ def compile_fused_gemm_scatter( ) direct_lsa = fuse and transport == "lsa" rotated = fuse if rotated is None else rotated + if n_stripe is None: + n_stripe = 2 if (fuse and transport == "lsa") else 1 + if n_stripe < 1: + raise ValueError(f"n_stripe must be >= 1, got {n_stripe}") + if n_stripe > 1 and not rotated: + raise ValueError("n_stripe > 1 needs the rotated tile order") assert BLOCK_M >= 128 and BLOCK_N >= 256 assert BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 @@ -1083,7 +1100,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"{'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'}{fence[0]}" + _kname_tag = f"{'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'}{'' 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 @@ -1147,8 +1164,21 @@ def kernel_gemm_scatter( # 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) - rest, bn = split_row_major_2d(fx.block_idx.x, n_blocks) - tile_i, dest_seq = split_row_major_2d(rest, ws) + if const_expr(n_stripe == 1): + 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 From 253c71b0123388af9dbdc85cf09ebc4dfbb5220d Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 15:09:54 +0000 Subject: [PATCH 28/79] perf(cco): sweep the whole stripe range, and drop the separate chunk-major path n_stripe was two branches -- 1 meant chunk-major and >1 meant the striped order -- which left the finest setting untestable and hid that the two are the same family. Chunk-major *is* a stripe of N // BLOCK_N, and measuring it both ways agrees (1569.0 against 1573.9). So the branch now keys on `n_stripe >= n_blocks_const` and the parameter spans 1..n_blocks. Full sweep at [16384, 7168] K=2048, 8 ranks, the divisors of 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 flat below that; 2 stays the default. fused-sdma is flat across the whole range, so its default is now spelled "chunk-major" rather than 1, which is what it always was. The sweep kills two plausible explanations, both worth recording so they are not proposed again: * Not XCD locality. dest = (block / n_stripe) % 8 and blocks go round-robin over the 8 XCDs, so a destination is fed by exactly min(n_stripe, 8) XCDs. 14 and 28 spread a destination's traffic over all eight XCDs and 2 confines it to two -- and 2 is the fast one. More XCD coverage is monotonically worse. * Not chunk-completion timing. If it were, fused-sdma would prefer chunk-major; it is indifferent. With 256 resident blocks out of 1792 the first chunk completes inside the first wave of blocks whatever the order. What is left is the simple reading: for peer stores, rotating the destination often keeps all seven links busy, and 2 is already often enough. --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 3 +- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 47 ++++++++++++++----- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 9e6d1720e..26cb98c18 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -416,7 +416,8 @@ def build_parser() -> argparse.ArgumentParser: 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, which is 2 for fused-lsa and 1 elsewhere", + "0 = per-mode default: 2 for fused-lsa, chunk-major elsewhere. 1 rotates " + "the destination on every block", ) p.add_argument( "--tile-order", diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 0e1eea17b..85cc96ff9 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -990,15 +990,34 @@ def compile_fused_gemm_scatter( 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`` > 1 rotates the destination every ``n_stripe`` N-tiles instead - of after a whole chunk -- 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 defaults to 2 for ``fused-lsa`` and 1 everywhere else, - which is the same split gcnasm makes: with a PUT to trigger you want a - chunk finished as early as possible, and with peer stores you want them - spread evenly over the links rather than arriving in same-destination - bursts. Measured at [16384, 7168] K=2048 on 8 ranks: fused-lsa 1560.7 -> - 1366.3us (its GEMM 937.9 -> 788.8), fused-sdma unchanged within noise. + ``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=...)``. @@ -1040,11 +1059,13 @@ def compile_fused_gemm_scatter( direct_lsa = fuse and transport == "lsa" rotated = fuse if rotated is None else rotated if n_stripe is None: - n_stripe = 2 if (fuse and transport == "lsa") else 1 + n_stripe = 2 if (fuse and transport == "lsa") else 0 # 0 -> chunk-major + if n_stripe == 0: + n_stripe = N // BLOCK_N # chunk-major == stripe of a whole chunk if n_stripe < 1: raise ValueError(f"n_stripe must be >= 1, got {n_stripe}") - if n_stripe > 1 and not rotated: - raise ValueError("n_stripe > 1 needs the rotated tile order") + if n_stripe != 1 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 @@ -1164,7 +1185,7 @@ def kernel_gemm_scatter( # 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 == 1): + 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: From 3175289dfbfe39c70cd94d01bacd2bfdf86564af Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 15:15:35 +0000 Subject: [PATCH 29/79] fix(cco): the n_stripe guard rejected chunk-major on non-rotated builds Folding chunk-major into the n_stripe range made `n_stripe == 1` stop meaning "unstriped", but the guard still read `n_stripe != 1 and not rotated`. With `fuse=False` the default resolves to chunk-major (N // BLOCK_N, e.g. 2 at N=512), so every non-rotated build -- the split baseline, and test_pinned_copy_matches_aiter_kernel_bitwise -- now raised. Committed in 253c71b0 with the test failure visible; this is the fix. The guard is only about a *real* stripe, so it is now `n_stripe < N // BLOCK_N and not rotated`. Also bound n_stripe to [1, N // BLOCK_N] and require it to divide the N-tile count, which the striped index math assumed silently. Test suite passes (12/12) and the five modes are unchanged at [16384, 7168] K=2048: gemm-only 233.9, fused-sdma 1114.6, split-sdma 1264.2, split-lsa 1264.4, fused-lsa 1359.0, all validating. --- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 85cc96ff9..3b53631f8 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -1058,13 +1058,25 @@ def compile_fused_gemm_scatter( ) direct_lsa = fuse and transport == "lsa" 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 (fuse and transport == "lsa") else 0 # 0 -> chunk-major + n_stripe = 2 if direct_lsa else 0 if n_stripe == 0: - n_stripe = N // BLOCK_N # chunk-major == stripe of a whole chunk - if n_stripe < 1: - raise ValueError(f"n_stripe must be >= 1, got {n_stripe}") - if n_stripe != 1 and not rotated: + 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 From ef76030001920546e4db99f3260100f532e48551 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 16:23:36 +0000 Subject: [PATCH 30/79] feat(cco): support the model's fp8 block scale, and the fusion margin doubles The benchmark's GEMM was a8w8 per-token/per-channel while the model dispatches a block-scale kernel whose scales vary along K, so every fused-vs-split number was being drawn at a compute:comm of 0.44 where the real one is 0.76 -- our GEMM was 1.7x cheaper than the one the fusion actually has to hide behind. `--quant blockscale` adds the model's form: A quantised 1x128, B 128x128, fp32 scales. `--quant ptpc` keeps the old path and stays the default. The scale layouts are the ones aiter's gemm_a8w8_blockscale_bpreshuffle 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. Established by sweeping the four combinations against that kernel -- only this one lands at the fp8 floor, the others give 0.17 to 0.32. Applying the scales needs work in the mainloop, since `sum_k s_a[m,k] * s_b[nb,k] * d_k` does not factor. The obvious form -- promote each K-block into a second accumulator and zero the MFMA one -- does not work, and the reason is worth recording: zeroing removes the only dependency between K-blocks, the K loop is `range_constexpr` so fully unrolled, and the scheduler then hoists every block's MFMAs above every promotion. Measured at K=512, all 64 MFMAs in the first 400 instructions and all the arithmetic after, keeping K_ITERS accumulator sets live at once: 256 VGPR with 81 spilled and the GEMM 10x slow (4052us). `rocdl.sched_barrier(0)` between them does not stop it. So the accumulator instead carries the running sum rescaled, which keeps a real data dependency and needs no second accumulator: t_0 = d_0; t_k = t_{k-1} * (s_{k-1}/s_k) + d_k; out = t_{K-1} * s_{K-1} 170 VGPR, 0 spill, against 160 for ptpc at the same tile. BLOCK_M defaults to 128 for blockscale, which is why the tile is 128x256 there. Results at [16384, 7168] K=2048 on 8x MI355X, median of 31: --quant ptpc --quant blockscale fused-sdma 1114.5us 1146.2us split-sdma 1261.7 1465.5 split-lsa 1262.4 1469.9 fused-lsa 1365.6 1699.6 gemm-only 228.9 388.5 The fusion margin over split goes 11.7% -> 21.8%, which is the point: the GEMM grows 228.9 -> 388.5us while the 489us of link time does not, so there is more compute to hide the transfer behind. Against the model's own 1491.3us per layer that is 23.1% under. And our blockscale GEMM at 388.5 is within 4% of the 373.9us aiter's gemm_a8w8_blockscale_bpreshuffle measures standalone here, so the operating point is now the model's. Verification: a new test compares our kernel against the kernel the model dispatches and against an fp32 reference (not bitwise -- different accumulation order); 14/14 tests pass; fused-sdma, fused-lsa and split-sdma each validate 8/8 at the production shape. One FlyDSL trap hit on the way, noted at the site: a `def` nested inside a `@flyc.kernel` function gets its AST rewritten, and the rewrite of an `if` inside it turned a captured object into a local, so reading it raised UnboundLocalError and a `nonlocal` on another raised "no binding". Module-level methods are left alone, so the rescale lives on _BlockScaleK. --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 125 +++++++-- benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 264 +++++++++++++++++- tests/python/cco/test_gemm_ar.py | 80 ++++++ 3 files changed, 445 insertions(+), 24 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 26cb98c18..1026675b5 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -47,16 +47,29 @@ ``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: - fused-sdma (chunks=8) 1114.5us - split-sdma 1261.7 - split-lsa 1262.4 - fused-lsa (n_stripe=2) 1365.6 - gemm-only 228.9 + --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 aiter +8-wave kernel's native per-token/per-channel form, kept because the two bitwise +tests 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 ``fused-sdma`` is -**11.7% under split and 25.3% under what the model runs today**. +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 at 388.5 is within 4% of the 373.9us the model's own +``gemm_a8w8_blockscale_bpreshuffle`` measures standalone on this box. Fusing over SDMA wins, and only because of ``--chunks``; see the table at its definition in ``run()``. It was pinned to 1 while the aiter GEMM's @@ -114,18 +127,63 @@ def _setup_distributed(): return local_rank, rank, world_size, UniqueId.from_bytes(payload[0]) -def make_operands(rank: int, m: int, n: int, k: int): - """Deterministic per-rank fp8 operands, plus the bf16 reference partial. +#: 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) - 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 == "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: @@ -162,6 +220,11 @@ def run(args) -> int: from aiter.ops.shuffle import shuffle_weight 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. @@ -210,7 +273,7 @@ def run(args) -> int: ) cfg.validate() - a, b, sa, sb = make_operands(rank, args.m, args.n, args.k) + a, b, sa, sb = make_operands(rank, args.m, args.n, args.k, args.quant) b_shuf = shuffle_weight(b, layout=(16, 16)) vmm = max(4 * cfg.window_bytes + VMM_SLACK, VMM_SLACK) @@ -242,6 +305,7 @@ def run(args) -> int: BLOCK_M=args.block_m, BLOCK_N=args.block_n, b_preshuffled=True, + quant=args.quant, waves_per_eu=args.waves_per_eu, xcd_swizzle=args.xcd_swizzle, fuse=fused, @@ -263,9 +327,21 @@ def run(args) -> int: 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, sb, args.m, args.n, dc.ptr, win.handle, - stream=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( @@ -308,7 +384,7 @@ def once(): # 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 = (a.float() @ b.float().T) * sa[:, None] * sb[None, :] + 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 @@ -324,8 +400,8 @@ def once(): # 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) - acc += (ar.float() @ br.float().T) * sar[:, None] * sbr[None, :] + 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 @@ -361,6 +437,7 @@ def once(): "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, @@ -398,7 +475,17 @@ def build_parser() -> argparse.ArgumentParser: 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("--block-m", type=int, default=256) + 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) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 3b53631f8..477ef31b7 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -565,6 +565,164 @@ 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 + ``aiter.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. @@ -582,8 +740,14 @@ class _SwapABStoreC(StoreC): while B's is indexed by column, now 4 consecutive, so a vec4. """ - def __init__(self, *args, peer_rsrc=None, elem_base=None, **kwargs): + 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) @@ -591,7 +755,15 @@ def __init__(self, *args, peer_rsrc=None, elem_base=None, **kwargs): 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))), @@ -600,6 +772,8 @@ def _load_a_scale_scalar(self, row): 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))), @@ -635,7 +809,7 @@ def store(self, c_frag, base_row, base_col): oob = fx.Int32(self.c_rows * self.c_cols) vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) vals = [ - (vec_f32[k] * (a_scales[ti] * b_scales[tj][k])).to(fx.BFloat16) + 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 @@ -745,7 +919,7 @@ def _emit(self, c_frag, base_row, base_col, a_scales, b_scales): vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) packed = Vec.from_elements( [ - (vec_f32[k] * (a_scales[ti] * b_scales[tj][k])).to(fx.BFloat16) + self._scaled(vec_f32[k], a_scales[ti], b_scales[tj][k]).to(fx.BFloat16) for k in range_constexpr(4) ], fx.BFloat16, @@ -859,7 +1033,7 @@ def store(self, c_frag, base_row, base_col): for tj in range_constexpr(self.n_tiles_b): vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) vals += [ - (vec_f32[k] * (a_scales[ti] * b_scales[tj][k])).to(fx.BFloat16) + 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 @@ -966,6 +1140,7 @@ def compile_fused_gemm_scatter( fuse: bool = True, rotated: bool | None = None, n_stripe: int | None = None, + quant: str = "ptpc", swap_ab: bool = False, store_probe: bool = False, permlane: bool = False, @@ -1057,6 +1232,20 @@ def compile_fused_gemm_scatter( "--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 aiter's 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. @@ -1133,7 +1322,7 @@ def compile_fused_gemm_scatter( raise ValueError( f"fence must be all/agent/leader/none/writethrough/wt-agent, got {fence!r}" ) - _kname_tag = f"{'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'}{'' if n_stripe == 1 else f's{n_stripe}'}{fence[0]}" + _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'}{'' 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 @@ -1355,6 +1544,29 @@ def kernel_gemm_scatter( 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 @@ -1378,6 +1590,17 @@ def kernel_gemm_scatter( 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) @@ -1427,6 +1650,17 @@ def kernel_gemm_scatter( 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() @@ -1455,6 +1689,17 @@ def kernel_gemm_scatter( 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) @@ -1474,6 +1719,15 @@ def kernel_gemm_scatter( 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 diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 5a1046327..b05e766c2 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -156,6 +156,86 @@ def test_pinned_copy_matches_aiter_kernel_bitwise(m, n, k): ) +@pytest.mark.parametrize("m,n,k", [(1024, 512, 512), (2048, 7168, 2048)]) +def test_blockscale_agrees_with_the_kernel_the_model_dispatches(m, n, k): + """``--quant blockscale`` against ``gemm_a8w8_blockscale_bpreshuffle``. + + The model's ROCm path quantises A 1x128 and B 128x128 with **fp32** scales + and dispatches that aiter kernel; this benchmark's own GEMM has to agree + with it, or the fused-vs-split numbers are being drawn at the wrong + operating point. Not a bitwise check: the two accumulate in different + orders, so both are compared to an fp32 reference and to each other. + + The scale layouts are the ones that kernel 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 against it, + only this one lands at the fp8 floor; the others give 0.17-0.32. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + pytest.importorskip("aiter", reason="aiter not importable (set PYTHONPATH)") + import flydsl.expr as fx + from aiter.ops.gemm_op_a8w8 import gemm_a8w8_blockscale_bpreshuffle + from aiter.ops.shuffle import shuffle_weight + + sys.path.insert(0, str(GEMM_AR_DIR)) + sys.path.insert(0, str(AR_DIR)) + from kernels_fused import compile_fused_gemm_scatter + + 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 = shuffle_weight(b, layout=(16, 16)) + + 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() + + theirs = torch.empty(m, n, device="cuda", dtype=torch.bfloat16) + gemm_a8w8_blockscale_bpreshuffle(a, b_shuf, sa, sb, out=theirs) + + 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 = ((ours.float() - ref).norm() / rn).item() + rel_theirs = ((theirs.float() - ref).norm() / rn).item() + rel_pair = ((ours.float() - theirs.float()).norm() / rn).item() + assert rel_ours < 3e-3, f"ours {rel_ours:.3e} vs the fp32 reference" + assert rel_theirs < 3e-3, f"aiter {rel_theirs:.3e} vs the fp32 reference" + assert rel_pair < 3e-3, f"ours vs aiter {rel_pair:.3e}" + + @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. From fd291387958b88010814fc15ac2aa4a9f9f4f162 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 9 Sep 2026 17:22:57 +0000 Subject: [PATCH 31/79] fix(cco): adapt the raw buffer ops to FlyDSL 0.3.2's aux signature 0.3.2 moved `aux` on raw_ptr_buffer_load/store from a positional i32 *value* to a keyword-only *attribute*. Both spellings are in use on this box -- aiter's venv has 0.3.0, sglang's has 0.3.2 -- so probe the signature once at import and adapt, rather than pinning either. Found while checking whether the fused GEMM+AR runs under sglang's interpreter, which is the next step for measuring it end to end. Without this every kernel in ar/ and gemm_ar/ dies at trace time with "raw_ptr_buffer_store() takes 4 positional arguments but 5 were given"; with it, LSA is bit-exact and SDMA matches the 0.3.0 numbers (1048.5us against 1046.8 at [16384, 7168]). The 14 tests still pass under 0.3.0. Unrelated finding, recorded because it cost time: SDMA is a silent no-op under sglang's mori (`/root/mori-sglang-12d1bc32`, BUILD_CCO_SDMA=OFF at C++ build time). `BUILD_CCO_SDMA=1` in the environment only rebuilds the *device* bitcode; the host library still has no queues, so every put returns having done nothing and the all-reduce silently produces zeros. LSA is unaffected, which is what makes it easy to misread. Point PYTHONPATH at a mori built with the flag on. --- benchmark/cco/flydsl/ar/_compat.py | 36 +++++++++++++++++++++--------- 1 file changed, 26 insertions(+), 10 deletions(-) diff --git a/benchmark/cco/flydsl/ar/_compat.py b/benchmark/cco/flydsl/ar/_compat.py index c6fc16b9c..610efb576 100644 --- a/benchmark/cco/flydsl/ar/_compat.py +++ b/benchmark/cco/flydsl/ar/_compat.py @@ -43,6 +43,8 @@ from __future__ import annotations +import inspect + import flydsl import flydsl.expr as fx from flydsl._mlir import ir @@ -155,18 +157,35 @@ 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), - _const(cache_modifier, 32), + res, rsrc, _byte_offset(offset, elem), _const(0, 32), *tail, **kw ) def buffer_store(data, rsrc, offset, cache_modifier=CM_CACHED): @@ -177,12 +196,9 @@ def buffer_store(data, rsrc, offset, cache_modifier=CM_CACHED): 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), - _const(cache_modifier, 32), + value, rsrc, _byte_offset(offset, elem), _const(0, 32), *tail, **kw ) From d823dae109ceeb29179bc8c7c9a48a00b2f12d55 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 01:53:41 +0000 Subject: [PATCH 32/79] feat(cco): make the SDMA queue count a parameter, and default the pipeline to 1 Queue ids were the destination rank, which forced sdma_queue_count=world_size. The hardware queues are per (source, destination) pair, so that allocates world_size queues *per peer* -- 56 on an 8-rank node -- and touches exactly one of them. Standalone that is only wasteful; inside a server that already holds SDMA engines for its own copies, hsaKmtCreateQueueExt starts failing (anvil.cpp:237). Take the id modulo the queue count instead, in both the standalone phases and the fused GEMM epilogue, and plumb --sdma-queues through to the fused kernel (it is now part of the kernel name tag, so the two cannot silently disagree). One queue per peer is all the pipeline needs: a destination is one xGMI link, and cco's rule is about concurrently-issuing warps per queue, of which there is one. Verified equivalent at the production shape -- queues=1 1147.7us against queues=8 1148.4us, [16384, 7168] K=2048, 8 ranks. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/kernels_sdma.py | 24 ++++++++++++------- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 1 + benchmark/cco/flydsl/gemm_ar/kernels_fused.py | 9 +++++-- 3 files changed, 24 insertions(+), 10 deletions(-) diff --git a/benchmark/cco/flydsl/ar/kernels_sdma.py b/benchmark/cco/flydsl/ar/kernels_sdma.py index ac97f52ae..8f36755b8 100644 --- a/benchmark/cco/flydsl/ar/kernels_sdma.py +++ b/benchmark/cco/flydsl/ar/kernels_sdma.py @@ -127,8 +127,18 @@ def build_sdma_phases( 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``; one queue per peer keeps - concurrently-issuing warps per queue at 1, which is cco's stated rule. + ``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 @@ -145,10 +155,8 @@ def build_sdma_phases( f"SDMA needs a landing slot per peer: build the ArConfig with " f"recv_slots={ws} (got {cfg.recv_slots})" ) - if queues < ws: - raise ValueError( - f"need one SDMA queue per peer: queues={queues} < world_size={ws}" - ) + if queues < 1: + raise ValueError(f"queues must be >= 1, got {queues}") threads, blocks = cfg.threads, cfg.blocks # The reduce is a *local HBM* kernel, so it must not inherit the grid the LSA @@ -215,11 +223,11 @@ def push(dev_comm: Int64, win: Int64): win, src_off_expr(tid), fx.Int64(slice_bytes), - tid, + tid % fx.Int32(queues), coop=cco.CoopScope.THREAD, signal=signal, ) - sdma.quiet_queue(tid, tid) + 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 diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 1026675b5..91e157054 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -306,6 +306,7 @@ def run(args) -> int: 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, diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py index 477ef31b7..f0d2b43aa 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_fused.py @@ -1141,6 +1141,7 @@ def compile_fused_gemm_scatter( 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, @@ -1322,7 +1323,7 @@ def compile_fused_gemm_scatter( 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'}{'' if n_stripe == 1 else f's{n_stripe}'}{fence[0]}" + _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 @@ -1854,7 +1855,11 @@ def kernel_gemm_scatter( + fx.Int64(dest) * fx.Int64(slice_bytes) + off, fx.Int64(chunk_bytes), - dest, + # 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, ) From e675079a751b3b0eac237682fa1efae92f083123 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 09:24:06 +0000 Subject: [PATCH 33/79] bench(cco): port CK's 4-wave B-out-of-LDS GEMM, and find the gap is elsewhere Our blockscale GEMM measures 370.5us at [16384, 7168] K=2048 against 300.4 for CK's fastest instance measured the same way, and the structural difference looked like the explanation: CK keeps B entirely out of LDS (global->VGPR, double-buffered in registers, 16 KB of LDS for A alone) where kernels_fused stages it (96 KB per block, one block per CU), and CK's winning tile is 64x256 on 4 waves where ours is 128x256 on 8. aiter's preshuffle_gemm is already that architecture -- 4 waves, 256 threads, A the only thing in LDS, B pulled straight to registers from the same N0/K0/KLane/NLane/KPack preshuffle CK uses. Pin it and add the model's quantisation, which upstream does not have (its fp8 scales are per-row x per-col, applied once in the epilogue). The shape is not the problem: kernel no scale blockscale scale cost 8-wave 128x256 (kernels_fused) 224.0 370.5 +146.5 preshuffle4w 128x128x256 290.9 857.3 +566.4 preshuffle4w 64x128x128 (best) -- 532.4 -- CK 64x256 Intrawave v1 -- 300.4 ~0 Unscaled we are already faster than CK's scaled kernel. What costs us is applying the per-K-block scale, and this kernel makes it worse. CK gets it nearly free by promoting one MFMA tile at a time -- loop order (m0, n0, kscale) with k innermost, partial in a single GetRegSizePerXdlops() buffer, 4 VGPRs. fx.gemm issues a whole fragment's MFMAs at once, so only whole-fragment forms are expressible and both are expensive: a running rescale costs a divide per element per K-block (660us, of which 88 came back with v_rcp_f32), and whole-fragment promotion needs a duplicate accumulator and spills -- 2753us, the same failure the 8-wave kernel hit at 4052us. So closing the gap means driving MFMAs per accumulator tile, not swapping the pipeline. Kept as a measured negative result, gemm-only, behind --gemm-impl preshuffle4w; nothing else changes path. Numerics are right throughout: relL2 1.658e-3, the fp8 floor the 8-wave kernel also hits. One upstream note: every buffer_ops.buffer_load in preshuffle_gemm raises on this FlyDSL (RawPtrBufferLoadOp wants aux as an operand, buffer_ops passes None when no cache modifier is set), so --quant ptpc does not run here. The blockscale path goes through the layout API instead and is unaffected. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 53 +- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 1109 +++++++++++++++++ 2 files changed, 1159 insertions(+), 3 deletions(-) create mode 100644 benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 91e157054..79ad79a45 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -67,9 +67,12 @@ 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 at 388.5 is within 4% of the 373.9us the model's own -``gemm_a8w8_blockscale_bpreshuffle`` measures standalone on this box. +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. The gap is the +per-K-block scale, not the GEMM -- unscaled we are 224.0us. See +``kernels_preshuffle4w.py``, which ports CK's 4-wave B-out-of-LDS shape and +lands at 532.4us for exactly that reason. Fusing over SDMA wins, and only because of ``--chunks``; see the table at its definition in ``run()``. It was pinned to 1 while the aiter GEMM's @@ -108,6 +111,7 @@ sys.path.insert(0, _HERE) sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "ar")) from kernels_fused import compile_fused_gemm_scatter # noqa: E402 +from kernels_preshuffle4w import compile_preshuffle_gemm # noqa: E402 from kernels_lsa import build_lsa_ar # noqa: E402 from kernels_sdma import build_sdma_phases # noqa: E402 from layout import ArConfig # noqa: E402 @@ -220,6 +224,10 @@ def run(args) -> int: from aiter.ops.shuffle import shuffle_weight local_rank, rank, world_size, uid = _setup_distributed() + if args.gemm_impl == "preshuffle4w" and args.mode != "gemm-only": + raise SystemExit( + "--gemm-impl preshuffle4w has no fused epilogue; use --mode gemm-only" + ) # 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. @@ -298,6 +306,24 @@ def run(args) -> int: reqs.sdma_queue_count = args.sdma_queues dc = comm.create_dev_comm(reqs) + if args.gemm_impl == "preshuffle4w": + # CK's shape: 4 waves, B never staged in LDS. gemm-only, so the + # window's input region is just an ordinary output buffer here. + launch4w = compile_preshuffle_gemm( + N=args.n, + K=args.k, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + in_dtype="fp8", + out_dtype="bf16", + quant=args.quant, + waves_per_eu=args.waves_per_eu, + xcd_swizzle=args.xcd_swizzle, + ) + semaphore = torch.zeros(1, device="cuda", dtype=torch.int32) + bias_unused = torch.zeros(1, device="cuda", dtype=torch.bfloat16) + gemm = compile_fused_gemm_scatter( cfg, rank, @@ -341,6 +367,10 @@ def run(args) -> int: sa_arg, sb_arg = sa, sb def run_gemm(stream): + if args.gemm_impl == "preshuffle4w": + launch4w(c, c, semaphore, a, b_shuf, sa_arg, sb_arg, bias_unused, + args.m, args.n, stream=stream) + return gemm(a_i8, b_i8, c_flat, sa_arg, sb_arg, args.m, args.n, dc.ptr, win.handle, stream=stream) @@ -454,6 +484,12 @@ def once(): "rel_l2": rel_l2, "validated": validated, "timing": "eager" if args.eager else "graph", + "gemm_impl": args.gemm_impl, + "tile": ( + [args.tile_m, args.tile_n, args.tile_k] + if args.gemm_impl == "preshuffle4w" + else None + ), } print("RESULT_JSON " + json.dumps(result, sort_keys=True), flush=True) print( @@ -617,6 +653,17 @@ def build_parser() -> argparse.ArgumentParser: "and is worth 80us (355 vs 434)", ) p.add_argument("--sdma-queues", type=int, default=8) + p.add_argument( + "--gemm-impl", + choices=("8wave", "preshuffle4w"), + default="8wave", + help="8wave: the pinned aiter kernel, the only one with a fused " + "epilogue. preshuffle4w: CK's shape -- 4 waves, B loaded straight " + "to registers instead of through LDS. gemm-only.", + ) + p.add_argument("--tile-m", type=int, default=64, help="preshuffle4w only") + p.add_argument("--tile-n", type=int, default=256, help="preshuffle4w only") + p.add_argument("--tile-k", type=int, default=128, help="preshuffle4w only") p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) p.add_argument("--eager", action="store_true") diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py new file mode 100644 index 000000000..d330ef5a0 --- /dev/null +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -0,0 +1,1109 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Preshuffle GEMM (layout API): f16/bf16/fp8/int8, ping-pong scf.for loop with scheduler hints. + +Pinned from ``aiter.ops.flydsl.kernels.preshuffle_gemm`` and extended with +``quant="blockscale"`` -- the A 1x128 / B 128x128 fp32 quantisation DSV4 runs, +which upstream does not have (its fp8 scales are per-row x per-col, applied once +in the epilogue). + +Why this kernel and not ``kernels_fused``'s 8-wave one: it is already the shape +CK's fastest blockscale instance uses. 4 waves, 256 threads, and **B never +touches LDS** -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR +double-buffered, the way CK's ``blockscale_b_preshuffle_v1`` pipeline does with +``b_thread_bufs``. Only A is staged, so LDS is ``tile_m * tile_k`` x 2 stages +(16 KB at 64x128) against the 8-wave kernel's 96 KB, which caps that one at one +block per CU. B is already preshuffled to ``N0 K0 KLane NLane KPack``, so every +lane pulls its own MFMA operand with one 128-bit load and there is nothing for +LDS to amortise. + +Measured at ``[16384, 7168] K=2048`` on MI355X, warm cache, and the outcome is +not the one the shape predicts: + +=========================== ========= =========== ========== +kernel no scale blockscale scale cost +=========================== ========= =========== ========== +8-wave 128x256 (kernels_fused) 224.0 370.5 +146.5 +this, 128x128x256 290.9 857.3 +566.4 +this, 64x128x128 (best) -- 532.4 -- +CK 64x256 Intrawave v1 -- 300.4 ~0 +=========================== ========= =========== ========== + +So **the shape is not the problem and neither is B-in-LDS**. Our unscaled GEMM +is already faster than CK's scaled one (224.0 against 300.4); what costs us is +applying the per-K-block scale, and this kernel makes that worse, not better. + +CK gets it for free because it promotes **one MFMA tile at a time**: the loop is +``static_ford`` with k innermost, and the +partial lives in ``c_thread_buf_per_scale``, a single +``GetRegSizePerXdlops()``-wide buffer -- 4 VGPRs. ``fx.gemm`` issues the MFMAs +for a whole fragment at once, so the only two things expressible here are +whole-fragment, and both are expensive: a running rescale needs a divide per +element per K-block (660us; 88us of that recovered with ``v_rcp_f32``), and +whole-fragment promotion needs a duplicate accumulator, which spills -- 2753us +at 128x128x256, the same failure the 8-wave kernel hit at 4052us. + +Closing this means driving MFMAs per accumulator tile rather than per fragment, +which is a deeper change than swapping the pipeline. +""" + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import const_expr, gpu, math, range_constexpr, rocdl +from flydsl.expr.typing import ( + BFloat16, + Float8E4M3FN, + Float8E4M3FNUZ, + Float16, + Float32, + Int8, + Int32, + T, +) +from flydsl.expr.typing import Vector as Vec +from flydsl.runtime.device import get_rocm_arch + +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.mfma_preshuffle_pipeline import xcd_remap_bx_by +from aiter.ops.flydsl.kernels.splitk_epilogue import ( + CPOL_COHERENT, + splitk_reduce_epilogue, +) + +# (dsrd_preload, dvmem_preload) per (tile_m, tile_n, tile_k). +_TILE_PRELOAD_TABLE = { + # ── tile_m = 16 ── + (16, 64, 256): (2, 2), + (16, 64, 512): (4, 4), + (16, 128, 256): (2, 2), + (16, 128, 512): (2, 2), + (16, 192, 256): (2, 2), + (16, 256, 256): (2, 2), + (16, 256, 512): (2, 2), + (16, 512, 256): (2, 2), + # ── tile_m = 32 ── + (32, 64, 128): (6, 6), + (32, 64, 256): (6, 6), + (32, 64, 512): (2, 2), + (32, 128, 128): (6, 6), + (32, 128, 256): (6, 6), + (32, 192, 128): (6, 6), + (32, 192, 256): (6, 6), + (32, 256, 128): (6, 6), + (32, 256, 256): (6, 6), + # ── tile_m = 48 ── + (48, 64, 128): (8, 8), + (48, 64, 256): (2, 2), + (48, 128, 256): (6, 6), + (48, 192, 256): (6, 6), + (48, 256, 256): (6, 6), + # ── tile_m = 64 ── + (64, 64, 128): (4, 4), + (64, 64, 256): (4, 4), + (64, 128, 128): (8, 8), + (64, 128, 256): (8, 8), + (64, 192, 128): (8, 8), + (64, 192, 256): (8, 8), + (64, 256, 64): (8, 8), + (64, 256, 128): (8, 8), + (64, 256, 256): (8, 8), + # ── tile_m = 80 ── + (80, 64, 256): (4, 4), + (80, 128, 256): (8, 8), + (80, 192, 256): (8, 8), + (80, 256, 256): (8, 8), + # ── tile_m = 96 ── + (96, 64, 128): (6, 6), + (96, 64, 256): (6, 6), + (96, 128, 128): (8, 8), + (96, 128, 256): (6, 6), + (96, 192, 128): (8, 8), + (96, 192, 256): (8, 8), + (96, 256, 128): (8, 8), + (96, 256, 256): (8, 8), + # ── tile_m = 112 ── + (112, 64, 256): (8, 8), + (112, 128, 256): (4, 4), + (112, 192, 256): (8, 8), + (112, 256, 256): (8, 8), + # ── tile_m = 128 ── + (128, 64, 128): (6, 6), + (128, 64, 256): (8, 8), + (128, 128, 64): (4, 4), + (128, 128, 128): (8, 8), + (128, 128, 256): (4, 4), + (128, 192, 128): (8, 8), + (128, 192, 256): (8, 8), + (128, 256, 128): (6, 6), + (128, 256, 256): (4, 4), + # ── tile_m = 160 ── + (160, 192, 128): (8, 8), + # ── tile_m = 192 ── + (192, 64, 128): (6, 6), + (192, 128, 128): (6, 6), + # ── tile_m = 224 ── + (224, 64, 128): (4, 4), + (224, 128, 128): (6, 6), + (224, 192, 128): (6, 6), + # ── tile_m = 256 ── + (256, 64, 128): (4, 4), + (256, 128, 128): (6, 6), + (256, 192, 128): (6, 6), + (256, 256, 128): (4, 4), +} + +_TILE_PRELOAD_DEFAULT = (0, 0) + + +def _get_preload(tile_m, tile_n, tile_k): + """Look up (dsrd_preload, dvmem_preload) from the tile table.""" + return _TILE_PRELOAD_TABLE.get( + (int(tile_m), int(tile_n), int(tile_k)), _TILE_PRELOAD_DEFAULT + ) + + +# Row bound of the A/C layout views; gemm_kernels mirrors it to guard on the host. +PRESHUFFLE_M_MAX = 65536 + + +@functools.lru_cache(maxsize=1024) +def compile_preshuffle_gemm( + *, + N: int, + K: int, + tile_m: int, + tile_n: int, + tile_k: int, + in_dtype: str = "fp8", + out_dtype: str = "bf16", + quant: str = "ptpc", # "ptpc" (upstream) or "blockscale" (A 1x128, B 128x128) + epilogue: str = "none", # "none", "bias", "bias_relu", "bias_silu", "bias_gelu" + waves_per_eu: int | None = None, + enable_scheduler: bool = True, + use_async_copy: bool = False, + xcd_swizzle: int = 0, + lds_stage: int = 2, + split_k: int = 1, +): + """Compile preshuffle GEMM (fp8/int8/fp16/bf16). + Signature: fn(C, out, semaphore, A, B, scale_a, scale_b, bias, M, N, stream). + For split_k > 1, C is an fp32 partial workspace and the last arriving split + performs the reduction and final conversion inside the same kernel launch. + """ + if in_dtype not in ("fp8", "int8", "fp16", "bf16"): + raise ValueError(f"in_dtype must be fp8/int8/fp16/bf16, got {in_dtype!r}") + if split_k < 1 or K % split_k != 0: + raise ValueError(f"split_k must divide K; got split_k={split_k}, K={K}") + split_k_extent = K // split_k + if tile_k <= 0 or split_k_extent % tile_k != 0: + raise ValueError( + "tile_k must be a positive divisor of K/split_k; " + f"got tile_k={tile_k}, K={K}, split_k={split_k}" + ) + if epilogue not in ("none", "bias", "bias_relu", "bias_silu", "bias_gelu"): + raise ValueError( + f"epilogue must be none/bias/bias_relu/bias_silu/bias_gelu, got {epilogue!r}" + ) + if lds_stage not in (1, 2): + raise ValueError(f"lds_stage must be 1 or 2, got {lds_stage}") + _has_epilogue = epilogue != "none" + _has_bias = epilogue in ("bias", "bias_relu", "bias_silu", "bias_gelu") + _has_relu = epilogue == "bias_relu" + _has_silu = epilogue == "bias_silu" + _has_gelu = epilogue == "bias_gelu" + if split_k > 1 and _has_epilogue: + raise ValueError("split_k > 1 does not support fused bias or activation") + + if quant not in ("ptpc", "blockscale"): + raise ValueError(f"quant must be ptpc or blockscale, got {quant!r}") + is_blockscale = quant == "blockscale" + if is_blockscale: + # tile_k == 128 makes one main-loop tile exactly one K-scale block, so + # the rescale lands on a loop boundary that already exists. + if in_dtype != "fp8": + raise ValueError("blockscale is fp8 only") + if tile_k % 128: + raise ValueError(f"blockscale needs tile_k a multiple of 128, got {tile_k}") + if K % 128 or N % 128: + raise ValueError(f"blockscale needs K and N multiples of 128, got K={K} N={N}") + if split_k != 1 or _has_epilogue: + raise ValueError("blockscale does not support split_k or a fused epilogue") + # Only the 2-stage pipeline threads the tile index down to mma_kloop, + # which is what the rescale is indexed by. + if lds_stage != 2: + raise ValueError(f"blockscale needs lds_stage=2, got {lds_stage}") + + is_fp8 = in_dtype == "fp8" + is_int8 = in_dtype == "int8" + is_f16 = in_dtype == "fp16" + is_bf16 = in_dtype == "bf16" + is_f16_or_bf16 = is_f16 or is_bf16 + is_8bit = is_fp8 or is_int8 + elem_bytes = 1 if is_8bit else 2 + + # The async gmem->LDS DMA (buffer_load_lds 128b) only lowers for 8-bit inputs. + if use_async_copy and not is_8bit: + raise ValueError("use_async_copy is only supported for 8-bit inputs (fp8/int8)") + + gpu_arch = get_rocm_arch() + is_gfx942 = str(gpu_arch).startswith("gfx942") + is_gfx950 = str(gpu_arch).startswith("gfx950") + use_mfma_scale_128 = is_fp8 and is_gfx950 and (tile_k % 128 == 0) + use_mfma_k32 = is_f16_or_bf16 and is_gfx950 + + if is_f16_or_bf16: + layout_elem = Float16 if is_f16 else BFloat16 + elif is_int8: + layout_elem = Int8 + else: + layout_elem = Float8E4M3FN if is_gfx950 else Float8E4M3FNUZ + final_out_elem_cls = BFloat16 if out_dtype == "bf16" else Float16 + out_elem_cls = Float32 if split_k > 1 else final_out_elem_cls + out_elem_bytes = 4 if split_k > 1 else 2 + + # Tile geometry (tile_K_perm = K-elements grouped per MMA k-step) + tile_K_perm = 128 if use_mfma_scale_128 else (64 if is_8bit else 32) + k_iters = tile_k // tile_K_perm + num_tiles = split_k_extent // tile_k + #: K-scale blocks; with the tile_k=128 blockscale requires, one per tile. + kb_count = K // 128 if is_blockscale else 0 + m_repeat = tile_m // 16 + num_waves = 4 + n_per_wave = tile_n // num_waves + num_acc_n = n_per_wave // 16 + acc_size = m_repeat * num_acc_n * 4 + + total_threads = 256 + a_load_bytes = 16 + bytes_per_thread_a = (tile_m * tile_k * elem_bytes) // total_threads + num_a_loads = bytes_per_thread_a // a_load_bytes + num_b_loads = (tile_n * tile_k * elem_bytes) // total_threads // 16 + num_ds_load = (tile_m * tile_k * elem_bytes) // 64 // 16 # A LDS reads per wave + num_gmem_loads = num_a_loads + num_b_loads + if is_8bit and is_gfx950: + dsrd_preload, dvmem_preload = _get_preload(tile_m, tile_n, tile_k) + else: + dsrd_preload, dvmem_preload = (0, 0) + + a_lds_elems = tile_m * tile_k + + @fx.struct + class SharedStorage: + a0: fx.Array[layout_elem, a_lds_elems, 16] + if lds_stage == 2: + a1: fx.Array[layout_elem, a_lds_elems, 16] + if split_k > 1: + split_flag: fx.Array[Int32, 1, 4] + + # ── Kernel ──────────────────────────────────────────────────────── + @flyc.kernel + def kernel_gemm( + arg_c: fx.Tensor, + arg_out: fx.Tensor, + arg_semaphore: fx.Tensor, + arg_a: fx.Tensor, + arg_b: fx.Tensor, + arg_scale_a: fx.Tensor, + arg_scale_b: fx.Tensor, + arg_bias: fx.Tensor, + i32_m: fx.Int32, + i32_n: fx.Int32, + tiled_mma_arg: fx.TiledMma, + tiled_copy_g2s: fx.TiledCopy, + ): + tid = fx.thread_idx.x + bid_x, bid_y, bid_z = fx.block_idx + k_off = fx.Int32(bid_z) * num_tiles + + if const_expr(xcd_swizzle > 0): + _bx, _by = xcd_remap_bx_by( + gpu.block_id("x"), + gpu.block_id("y"), + fx.Index(i32_m), + tile_m=tile_m, + tile_n=tile_n, + N=N, + xcd_swizzle=xcd_swizzle, + ) + bid_x, bid_y = Int32(_bx), Int32(_by) + + if const_expr(use_mfma_scale_128): + _scale_atom = fx.make_mma_atom( + fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, layout_elem) + ) + tiled_mma = fx.make_tiled_mma( + _scale_atom, + fx.make_layout((1, 4, 1), (0, 1, 0)), + fx.make_tile(None, None, fx.make_layout((32, 4), (1, 32))), + ) + else: + tiled_mma = tiled_mma_arg + + # Bound A (read) and C (store) to the actual M extent so blocks covering + # rows past M (ragged M) drop their OOB loads/stores at the descriptor + # instead of faulting / writing past the allocation. B and scale_b/bias + # are per-N (exact multiple) and stay max_size; scale_a is per-row (M) and + # is bounded the same way below (see its epilogue load). + gA = fx.rocdl.make_buffer_tensor( + arg_a, + max_size=False, + num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes), + ) + gB = fx.rocdl.make_buffer_tensor(arg_b) + c_tensor = arg_c + if const_expr(split_k > 1): + c_split_offset = fx.Int64(bid_z) * fx.Int64(i32_m) * fx.Int64(N) + c_tensor = fx.Tensor( + fx.make_view( + fx.add_offset(fx.get_iter(arg_c), c_split_offset), + fx.make_layout((PRESHUFFLE_M_MAX, N), (N, 1)), + ) + ) + gC = fx.rocdl.make_buffer_tensor( + c_tensor, + max_size=False, + num_records_bytes=fx.Int64(i32_m) * fx.Int64(N) * fx.Int64(out_elem_bytes), + ) + + tA = fx.flat_divide(gA, fx.make_tile(tile_m, tile_k))[None, None, bid_x, None] + tB = fx.flat_divide(gB, fx.make_tile(tile_n, tile_k))[None, None, bid_y, None] + tC = fx.flat_divide(gC, fx.make_tile(tile_m, tile_n))[None, None, bid_x, bid_y] + + buf_copy = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), layout_elem) + uni_copy = fx.make_copy_atom(fx.UniversalCopy128b(), layout_elem) + + # Per-thread slices + thr_mma = tiled_mma.thr_slice(tid) + thr_g2s = tiled_copy_g2s.get_slice(tid) + thr_s2r = fx.make_tiled_copy_A(buf_copy, tiled_mma).get_slice(tid) + thr_g2r_B = fx.make_tiled_copy_B(buf_copy, tiled_mma).get_slice(tid) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + if const_expr(is_8bit): + k_blocks16 = (tile_k * elem_bytes) // 16 + if k_blocks16 <= 0 or (k_blocks16 & (k_blocks16 - 1)) != 0: + raise ValueError( + f"Unsupported tile_k for 8-bit LDS swizzle: tile_k={tile_k}, elem_bytes={elem_bytes} (k_blocks16={k_blocks16}); " + "expected tile_k*elem_bytes to be a positive multiple of 16 with (tile_k*elem_bytes/16) a power of two." + ) + swz_bits = k_blocks16.bit_length() - 1 # log2 + swz = fx.SwizzleType.get(swz_bits, 4, swz_bits) + else: + swz = fx.SwizzleType.get(3, 3, 3) + + def _make_sA(arr): + return fx.make_view( + arr.ptr, + fx.make_composed_layout( + fx.static(swz), + fx.make_ordered_layout((tile_m, tile_k), (1, 0)), + ), + ) + + if const_expr(lds_stage == 1): + sA_stages = [_make_sA(lds.a0)] + else: + sA_stages = [_make_sA(lds.a0), _make_sA(lds.a1)] + + # Partitions + pA_g = thr_g2s.partition_S(tA) + pA_s_stages = [thr_g2s.partition_D(s) for s in sA_stages] + pA_s2r_stages = [thr_s2r.partition_S(s) for s in sA_stages] + pB_g = thr_g2r_B.partition_S(tB) + + # Fragments — 2 separate B fragments (split double buffer for VGPR lifetime) + frag_copy_A = fx.make_fragment_like(pA_s_stages[0][None, None, None]) + frag_A = thr_mma.make_fragment_A(sA_stages[0]) + frag_B_single_layout = thr_mma.partition_B(tB).layout(None, None, None, 0) + frag_B_stages = [ + fx.make_fragment_like(frag_B_single_layout, layout_elem.ir_type) + for _ in range(2) + ] + frag_C = thr_mma.make_fragment_C(tC) + frag_A_retile = thr_s2r.retile(frag_A) + frag_B_retile_stages = [thr_g2r_B.retile(b) for b in frag_B_stages] + out_cpol = CPOL_COHERENT if split_k > 1 else 0 + copy_op = fx.rocdl.BufferCopy32b if split_k > 1 else fx.rocdl.BufferCopy16b + buf_copy_out = fx.make_copy_atom(copy_op(out_cpol), out_elem_cls) + thr_r2g_C = fx.make_tiled_copy_C(buf_copy_out, tiled_mma).get_slice(tid) + pC_g = thr_r2g_C.partition_S(tC) + frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) + frag_C_retile = thr_r2g_C.retile(frag_C_out) + + # ── Async gmem->LDS DMA (buffer_load_lds) for the A tile ── + if const_expr(use_async_copy): + dma_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + # Bound to the real M extent (like the sync gA) so ragged-M blocks DMA-read + # OOB rows as 0 instead of faulting past the allocation. + gA_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor( + fx.make_view( + fx.get_iter(arg_a), fx.make_layout(PRESHUFFLE_M_MAX * K, 1) + ) + ), + max_size=False, + num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes), + ) + gA_div = fx.logical_divide(gA_flat, fx.make_layout(1, 1)) + sA_i8_ptr = [fx.recast_iter(Int8, lds.a0.ptr)] + if const_expr(lds_stage == 2): + sA_i8_ptr.append(fx.recast_iter(Int8, lds.a1.ptr)) + bx_m = bid_x * tile_m + wave_id = tid // 64 + step_bytes = total_threads * a_load_bytes + wave_stride_bytes = 64 * a_load_bytes + k_blocks16_dma = (tile_k * elem_bytes) // 16 + elems_per_16b = 16 // elem_bytes + + def dma_a_to_lds(k_tile_val, stage): + wave_off = rocdl.readfirstlane( + fx.Int32.ir_type, wave_id * wave_stride_bytes + ) + lds_ptr = fx.add_offset(sA_i8_ptr[stage], wave_off) + base_k = (k_off + k_tile_val) * tile_k + for i in range_constexpr(num_a_loads): + if const_expr(i > 0): + lds_ptr = fx.add_offset(lds_ptr, step_bytes) + pos_bytes = i * total_threads * a_load_bytes + tid * a_load_bytes + elem_idx = pos_bytes // elem_bytes + m = elem_idx // tile_k + k = elem_idx % tile_k + k_swz = k ^ ((m % k_blocks16_dma) * elems_per_16b) + gmem_byte = ((bx_m + m) * K + base_k + k_swz) * elem_bytes + dst = fx.make_view(lds_ptr, fx.make_layout(1, 1)) + src = fx.slice(gA_div, (None, fx.Int32(gmem_byte))) + fx.copy(dma_atom, src, dst) + + # ── Scheduling hints (ported from old pipeline) ─────────── + def build_scheduler(numer: int, denom: int): + if const_expr(denom <= 0): + return [] + if const_expr(numer <= 0): + return [0] * denom + out = [] + prev = 0 + for i in range_constexpr(denom): + cur = ((i + 1) * numer + (denom - 1)) // denom + out.append(cur - prev) + prev = cur + return out + + def hot_loop_scheduler(): + mfma_group = num_acc_n + + if const_expr(is_gfx942): + mfma_total = (k_iters * 2) * m_repeat * mfma_group + mfma_per_iter = 2 * mfma_group + sche_iters = 0 if mfma_per_iter == 0 else (mfma_total // mfma_per_iter) + + rocdl.sched_dsrd(2) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + + if const_expr(num_acc_n < 4): + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + + dswr_tail = num_a_loads + dstr_advance = 2 + if const_expr(dswr_tail > sche_iters): + dswr_tail = sche_iters + dswr_start = max(sche_iters - dswr_tail - dstr_advance, 0) + + for sche_i in range_constexpr(sche_iters): + rocdl.sched_vmem(1) + rocdl.sched_mfma(mfma_group) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(mfma_group) + if const_expr(sche_i >= dswr_start - 1): + rocdl.sched_dswr(1) + else: + if const_expr(use_mfma_scale_128): + element_k_per_mfma = 128 + else: + element_k_per_mfma = 32 + num_mfma_per_tile_k = tile_k // element_k_per_mfma + mfma_total = num_mfma_per_tile_k * m_repeat * mfma_group + dswr_tail = num_a_loads + dstr_advance = 2 + if const_expr(dswr_tail > mfma_total): + dswr_tail = mfma_total + dsrd_preload_eff = min(int(dsrd_preload), num_ds_load) + dvmem_preload_eff = min(int(dvmem_preload), num_gmem_loads) + vmem_remaining = num_gmem_loads - dvmem_preload_eff + dsrd_remaining = num_ds_load - dsrd_preload_eff + if const_expr(vmem_remaining > 0 and vmem_remaining < mfma_total): + vmem_schedule = build_scheduler(vmem_remaining, vmem_remaining) + [ + 0 + ] * (mfma_total - vmem_remaining) + else: + vmem_schedule = build_scheduler(vmem_remaining, mfma_total) + dsrd_schedule = build_scheduler(dsrd_remaining, mfma_total) + dswr_start = max(mfma_total - dswr_tail - dstr_advance, 0) + last_dsrd_mfma_idx = -1 + for sched_idx in range_constexpr(mfma_total): + if const_expr(dsrd_schedule[sched_idx]): + last_dsrd_mfma_idx = sched_idx + dswr_start = max(dswr_start, last_dsrd_mfma_idx + 1) + idx_ds_read = dsrd_preload_eff + idx_gmem_load = dvmem_preload_eff + idx_ds_write = 0 + if const_expr(dvmem_preload_eff): + rocdl.sched_vmem(dvmem_preload_eff) + if const_expr(dsrd_preload_eff): + rocdl.sched_dsrd(dsrd_preload_eff) + for mfma_idx in range_constexpr(mfma_total): + rocdl.sched_mfma(1) + n_dsrd = dsrd_schedule[mfma_idx] + if const_expr(n_dsrd and (idx_ds_read < num_ds_load)): + if const_expr(idx_ds_read + n_dsrd > num_ds_load): + n_dsrd = num_ds_load - idx_ds_read + if const_expr(n_dsrd): + rocdl.sched_dsrd(n_dsrd) + idx_ds_read += n_dsrd + n_vmem = vmem_schedule[mfma_idx] + if const_expr(n_vmem and (idx_gmem_load < num_gmem_loads)): + if const_expr(idx_gmem_load + n_vmem > num_gmem_loads): + n_vmem = num_gmem_loads - idx_gmem_load + if const_expr(n_vmem): + rocdl.sched_vmem(n_vmem) + idx_gmem_load += n_vmem + if const_expr( + (not use_async_copy) + and (idx_ds_write < dswr_tail) + and (mfma_idx >= dswr_start) + ): + rocdl.sched_dswr(1) + idx_ds_write += 1 + if const_expr((not use_async_copy) and idx_ds_write < num_a_loads): + rocdl.sched_dswr(num_a_loads - idx_ds_write) + + rocdl.sched_barrier(0) + + # Accumulator addressing, shared by the block-scale rescale and the + # epilogue: element p of frag_C is row bx_m + mi*16 + lane_div_16*4 + ii, + # column by_n + (ni*num_waves + wave_id)*16 + lane_mod_16. + bx_m = bid_x * tile_m + by_n = bid_y * tile_n + wave_id = gpu.thread_id("x") // 64 + lane_id = gpu.thread_id("x") % 64 + lane_div_16 = lane_id // 16 + lane_mod_16 = lane_id % 16 + + # Scale operands go through the layout API rather than aiter's + # buffer_ops: on this FlyDSL, RawPtrBufferLoadOp wants `aux` as an + # operand, and buffer_ops passes None when no cache modifier is set, so + # every buffer_ops.buffer_load in this file raises. That breaks the + # upstream ptpc epilogue too; blockscale simply does not use it. + if const_expr(is_blockscale): + gSA = fx.rocdl.make_buffer_tensor( + arg_scale_a, + max_size=False, + num_records_bytes=fx.Int64(i32_m) * fx.Int64(kb_count) * fx.Int64(4), + ) + gSB = fx.rocdl.make_buffer_tensor(arg_scale_b) + sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) + sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + sc_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + sc_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) + sc_reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + sc_reg_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) + + def _load_sa(mi, kb): + """This lane's four A-scale rows of M-tile ``mi`` at K-block ``kb``. + + ``arg_scale_a`` is physically [K/128, M] (column-major [M, K/128]), + the layout aiter_per1x128_quant(transpose_scale=True) emits, so the + four rows are contiguous and land in one vec4. + """ + idx = fx.Int32(kb) * fx.Int32(i32_m) + bx_m + mi * 16 + lane_div_16 * 4 + fx.copy(sc_atom_4, fx.slice(sa_div, (None, fx.Int32(idx))), sc_reg_4) + return Vec(fx.memref_load_vec(sc_reg_4)) + + def _load_sb(ni, kb): + """The single B scale covering N-tile ``ni`` at K-block ``kb``. + + ``arg_scale_b`` is [N/128, K/128] row-major, and a 16-column group + lies wholly inside one 128-wide block, so this is one scalar. + """ + nblk = (by_n + (ni * num_waves + wave_id) * 16) // 128 + idx = nblk * fx.Int32(kb_count) + fx.Int32(kb) + fx.copy(sc_atom_1, fx.slice(sb_div, (None, fx.Int32(idx))), sc_reg_1) + return Vec(fx.memref_load_vec(sc_reg_1))[0] + + def rebase_to(kb): + """Rebase the accumulator from K-block ``kb-1``'s scales to ``kb``'s. + + The accumulator is kept as ``C_k / s_k``: one multiply per element + per K-block and no second accumulator, since promotion into a + duplicate of a 64-VGPR frag_C is what spilled 81 registers last + time. Keeping the multiply on the live accumulator also leaves a + real data dependency, so the scheduler cannot hoist every MFMA above + every rescale -- the failure `rocdl.sched_barrier(0)` did not stop. + + The M-tiles are walked load -> divide -> apply one at a time rather + than loading all of them up front. Both orders emit the same + arithmetic, but the batched one keeps 2*m_repeat vec4 of scales live + at once -- as many registers as the accumulator itself -- and that + cost 643us at 128x128x256 against 7 here. + + ``kb - 1`` is clamped rather than branched on: at kb=0 the ratio is + 1, and the accumulator is zero there anyway. + """ + kb_i = fx.Int32(kb) + zero_i = fx.Int32(0) + prev_kb = (kb_i > zero_i).select(kb_i - fx.Int32(1), zero_i) + # B first: num_acc_n scalars, and every M-tile needs them. + # v_rcp_f32 rather than a divide: an IEEE f32 divide expands to + # ~10 VALU ops and there are m_repeat*4 + num_acc_n of them per + # K-block, which measured 660us against a 291us kernel. rcp is + # 1 ulp, four orders under the 1.66e-3 fp8 floor these operands + # already carry. + rb = [ + _load_sb(ni, prev_kb) * fx.Float32(fx.rocdl.rcp(T.f32, _load_sb(ni, kb_i))) + for ni in range_constexpr(num_acc_n) + ] + acc = Vec(frag_C.load()) + vals = [None] * acc_size + for mi in range_constexpr(m_repeat): + pa = _load_sa(mi, prev_kb) + ca = _load_sa(mi, kb_i) + ra = [ + pa[e] * fx.Float32(fx.rocdl.rcp(T.f32, ca[e])) + for e in range_constexpr(4) + ] + for ni in range_constexpr(num_acc_n): + for ii in range_constexpr(4): + p = ni * (m_repeat * 4) + mi * 4 + ii + vals[p] = acc[p] * (ra[ii] * rb[ni]) + frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) + + def final_scale(): + """Undo the invariant: multiply by the last K-block's scales.""" + last = fx.Int32(kb_count - 1) + sb = [_load_sb(ni, last) for ni in range_constexpr(num_acc_n)] + acc = Vec(frag_C.load()) + vals = [None] * acc_size + for mi in range_constexpr(m_repeat): + sa = _load_sa(mi, last) + for ni in range_constexpr(num_acc_n): + for ii in range_constexpr(4): + p = ni * (m_repeat * 4) + mi * 4 + ii + vals[p] = acc[p] * (sa[ii] * sb[ni]) + frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) + + # ── Pipeline stage (double-buffered B via split fragments) ─ + def mma_kloop(a_stage, cur_frag_B, tile_idx=None): + for ki in range_constexpr(k_iters): + if const_expr(is_blockscale): + # One MMA k-step is one 128-wide scale block. + rebase_to(fx.Int32(tile_idx) * fx.Int32(k_iters) + fx.Int32(ki)) + fx.copy( + uni_copy, + pA_s2r_stages[a_stage][None, None, ki], + frag_A_retile[None, None, ki], + ) + k_coord = ki if (use_mfma_scale_128 or use_mfma_k32) else (None, ki) + fx.gemm( + tiled_mma, + frag_C, + frag_A[None, None, k_coord], + cur_frag_B[None, None, k_coord], + frag_C, + ) + + def pipeline_2stage(read_stage, next_k_val=None, read_next=True, cur_k_val=None): + write_stage = read_stage ^ 1 + a_read = read_stage + a_write = write_stage + cur_frag_B = frag_B_stages[read_stage] + do_next = read_next and next_k_val is not None + if const_expr(use_async_copy): + if const_expr(do_next): + dma_a_to_lds(next_k_val, a_write) + fx.copy( + buf_copy, + pB_g[None, None, None, k_off + next_k_val], + frag_B_retile_stages[write_stage], + ) + mma_kloop(a_read, cur_frag_B, cur_k_val) + if const_expr(enable_scheduler): + hot_loop_scheduler() + if const_expr(do_next): + rocdl.s_waitcnt(num_b_loads) + gpu.barrier() + return + if const_expr(do_next): + global_k_next = k_off + next_k_val + fx.copy(buf_copy, pA_g[None, None, None, global_k_next], frag_copy_A) + fx.copy( + buf_copy, + pB_g[None, None, None, global_k_next], + frag_B_retile_stages[write_stage], + ) + mma_kloop(a_read, cur_frag_B, cur_k_val) + if const_expr(do_next): + fx.copy(uni_copy, frag_copy_A, pA_s_stages[a_write][None, None, None]) + if const_expr(enable_scheduler): + hot_loop_scheduler() + if const_expr(do_next): + gpu.barrier() + + # ── Prologue ────────────────────────────────────────────── + acc_zero = ( + Vec.filled(acc_size, 0, Int32) + if const_expr(is_int8) + else Vec.filled(acc_size, 0.0, Float32) + ) + if const_expr(use_async_copy): + dma_a_to_lds(fx.Int32(0), 0) + fx.copy(buf_copy, pB_g[None, None, None, k_off], frag_B_retile_stages[0]) + frag_C.store(acc_zero) + rocdl.s_waitcnt(num_b_loads) + gpu.barrier() + else: + fx.copy(buf_copy, pA_g[None, None, None, k_off], frag_copy_A) + fx.copy(buf_copy, pB_g[None, None, None, k_off], frag_B_retile_stages[0]) + frag_C.store(acc_zero) + fx.copy(uni_copy, frag_copy_A, pA_s_stages[0][None, None, None]) + gpu.barrier() + rocdl.sched_barrier(0) + + # ── Main tile loop ──────────────────────────────────────────── + if const_expr(lds_stage == 1 and num_tiles > 1): + frag_Bc = frag_B_stages[0] + frag_Bc_retile = frag_B_retile_stages[0] + for iv, state in range(0, num_tiles - 1, 1, init=[frag_C.load()]): + frag_C.store(state[0]) + k_next = fx.Int32(iv + 1) + mma_kloop(0, frag_Bc) + if const_expr(not use_async_copy): + fx.copy( + buf_copy, pA_g[None, None, None, k_off + k_next], frag_copy_A + ) + fx.copy( + buf_copy, + pB_g[None, None, None, k_off + k_next], + frag_Bc_retile, + ) + gpu.barrier() # single buffer: all reads done before overwrite + if const_expr(use_async_copy): + dma_a_to_lds(k_next, 0) + else: + fx.copy(uni_copy, frag_copy_A, pA_s_stages[0][None, None, None]) + if const_expr(enable_scheduler): + hot_loop_scheduler() + if const_expr(use_async_copy): + rocdl.s_waitcnt(num_b_loads) + gpu.barrier() + results = yield [frag_C.load()] + frag_C.store(results) + elif const_expr(lds_stage == 2): + # 2-tile/iter ping-pong: middle loop runs 2 tiles/iter + is_odd_tiles = (num_tiles % 2) == 1 + tail = 1 if is_odd_tiles else 2 + loop_end = (num_tiles - tail) // 2 + + def two_tiles(k_base): + pipeline_2stage( + read_stage=0, next_k_val=k_base + fx.Int32(1), cur_k_val=k_base + ) + pipeline_2stage( + read_stage=1, + next_k_val=k_base + fx.Int32(2), + cur_k_val=k_base + fx.Int32(1), + ) + + if const_expr(loop_end > 0 and use_async_copy): + for iv in range_constexpr(loop_end): + two_tiles(fx.Int32(iv * 2)) + elif const_expr(loop_end > 0): + for iv, state in range(0, loop_end, 1, init=[frag_C.load()]): + frag_C.store(state[0]) + two_tiles(fx.Int32(iv * 2)) + results = yield [frag_C.load()] + frag_C.store(results) + k_tail0 = num_tiles - tail # first tile handled by the peeled tail + for j in range_constexpr(tail - 1): + pipeline_2stage( + read_stage=(k_tail0 + j) % 2, + next_k_val=fx.Int32(k_tail0 + j + 1), + cur_k_val=fx.Int32(k_tail0 + j), + ) + + # ── Epilogue-operand preloads (scale_a / scale_b / bias) ───────── + def load_epi_operands(): + s_a = s_b = bias = None + if const_expr(is_blockscale): + # Nothing to preload: final_scale() applies the last block's + # scales to frag_C directly, so the epilogue below sees an + # already-scaled accumulator. + pass + elif const_expr(is_8bit): + # Per-row(scale_a) × per-col(scale_b) scaling, applied in the epilogue. + scale_b_rsrc = buffer_ops.create_buffer_resource( + arg_scale_b, max_size=True + ) + s_b = [ + buffer_ops.buffer_load( + scale_b_rsrc, + fx.Int32(by_n + (ni * num_waves + wave_id) * 16 + lane_mod_16), + vec_width=1, + dtype=T.f32, + ) + for ni in range_constexpr(num_acc_n) + ] + scale_a_rsrc = buffer_ops.create_buffer_resource( + arg_scale_a, + max_size=False, + num_records_bytes=fx.Int64(i32_m) * fx.Int64(4), + ) + s_a = [ + Vec( + buffer_ops.buffer_load( + scale_a_rsrc, + fx.Int32(bx_m + mi * 16 + lane_div_16 * 4), + vec_width=4, + dtype=T.f32, + ) + ).bitcast(fx.Float32) + for mi in range_constexpr(m_repeat) + ] + if const_expr(_has_bias): + # Per-column bias (out_dtype), one scalar per N-block, shared across rows. + bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=True) + bias_elem_ty = T.bf16 if out_dtype == "bf16" else T.f16 + bias = [ + fx.Float32( + buffer_ops.buffer_load( + bias_rsrc, + fx.Int32( + by_n + (ni * num_waves + wave_id) * 16 + lane_mod_16 + ), + vec_width=1, + dtype=bias_elem_ty, + ) + ) + for ni in range_constexpr(num_acc_n) + ] + return s_a, s_b, bias + + overlap_epi_load = ( + acc_size <= 64 + ) # small enough accumulator to keep operands live over the MMA + s_a_vals = s_b_vals = bias_vals = None + if const_expr(overlap_epi_load): + s_a_vals, s_b_vals, bias_vals = load_epi_operands() + + # Final MMA stage — overlaps the epilogue-operand loads when issued above. + if const_expr(lds_stage == 1): + mma_kloop(0, frag_B_stages[0]) + else: + pipeline_2stage( + read_stage=(num_tiles - 1) % 2, + read_next=False, + cur_k_val=fx.Int32(num_tiles - 1), + ) + + # ── Epilogue ───────────────────────────────────────────── + if const_expr(is_blockscale): + final_scale() + if const_expr((not is_8bit or is_blockscale) and not _has_epilogue): + frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) + else: + if const_expr(not overlap_epi_load): + s_a_vals, s_b_vals, bias_vals = load_epi_operands() + + def apply_activation(val_s): + # ReLU/SiLU/GeLU : maximumf for relu; exp+rcp for silu; + # tanh-approx gelu expanded through a non-positive exponent (no overflow). + if const_expr(_has_relu): + return fx.Float32(val_s).maximumf(fx.Float32(0.0)) + if const_expr(_has_silu): + exp_neg = math.exp(val_s * fx.Float32(-1.0)) + return val_s * (fx.Float32(1.0) / (fx.Float32(1.0) + exp_neg)) + if const_expr(_has_gelu): + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = val_s * val_s * val_s + y = fx.Float32(0.7978845608) * (val_s + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * val_s * (numerator * (one_f32 / denom)) + return val_s + + acc_vec = Vec(frag_C.load()) + out_elems = [] + for p in range_constexpr(acc_size): + ni = p // (m_repeat * 4) + mi = (p // 4) % m_repeat + ii = p % 4 + val = acc_vec[p] + if const_expr(is_int8): + val = val.to(Float32) + if const_expr(is_8bit and not is_blockscale): + val_s = (val * s_a_vals[mi][ii]) * s_b_vals[ni] + else: + val_s = val + if const_expr(_has_bias): + val_s = val_s + bias_vals[ni] + val_s = apply_activation(val_s) + out_elems.append(val_s.to(out_elem_cls)) + + out_vec = vector.from_elements( + T.vec(acc_size, out_elem_cls.ir_type), out_elems + ) + frag_C_out.store(out_vec) + + fx.copy(buf_copy_out, frag_C_retile, pC_g) + if const_expr(split_k > 1): + # The store carries sc0|sc1, so waiting on it is the whole release. + rocdl.s_waitcnt(0) + gpu.barrier() + splitk_reduce_epilogue( + arg_c, + fx.Tensor( + fx.make_view( + fx.get_iter(arg_out), + fx.make_layout((PRESHUFFLE_M_MAX, N), (N, 1)), + ) + ), + arg_semaphore, + lds.split_flag.ptr, + tile_m, + tile_n, + total_threads, + final_out_elem_cls, + tid, + bid_x, + bid_y, + i32_m, + N, + split_k, + ) + + # ── Host launcher ───────────────────────────────────────────── + @flyc.jit + def launch_gemm( + arg_c: fx.Tensor, + arg_out: fx.Tensor, + arg_semaphore: fx.Tensor, + arg_a: fx.Tensor, + arg_b: fx.Tensor, + arg_scale_a: fx.Tensor, + arg_scale_b: fx.Tensor, + arg_bias: fx.Tensor, + i32_m: fx.Int32, + i32_n: fx.Int32, + stream: fx.Stream, + ): + CompilationContext.get_current() + + # MMA atom — layout_elem carries the dtype (Float16/BFloat16/Float8E4M3FN/etc) + if const_expr(use_mfma_k32): + mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, layout_elem)) + k_perm = fx.make_layout((8, 4), (1, 8)) + elif const_expr(is_f16_or_bf16): + mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 16, layout_elem)) + k_perm = fx.make_layout((4, 4, 2), (1, 8, 4)) + elif const_expr(is_int8): + mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, layout_elem, Int32)) + k_perm = fx.make_layout((8, 4, 2), (1, 16, 8)) + else: + # fp8: narrow atom here; the scale (16x16x128) tiled_mma is rebuilt in-kernel + mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, layout_elem)) + k_perm = fx.make_layout((8, 4, 2), (1, 16, 8)) + + tiled_mma = fx.make_tiled_mma( + mma_atom, + fx.make_layout((1, 4, 1), (0, 1, 0)), + fx.make_tile(None, None, k_perm), + ) + + # G2S tiled copy + val_per_thr = a_load_bytes // elem_bytes + thrs_k = tile_k // val_per_thr + thrs_m = total_threads // thrs_k + tiled_copy_g2s = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy128b(), layout_elem), + fx.make_layout( + ((thrs_k, thrs_m), (1, val_per_thr)), + ((thrs_m * val_per_thr, 1), (1, thrs_m)), + ), + fx.make_tile(thrs_m, tile_k), + ) + + # Preshuffle B layout (2D hierarchical) + kp_bytes = 16 + kp_elems = kp_bytes if elem_bytes == 1 else kp_bytes // elem_bytes + k_bytes_b = K * elem_bytes + n0 = N // 16 + k0 = k_bytes_b // 64 + s_nlane = kp_elems + s_klane = 16 * s_nlane + s_k0 = 4 * s_klane + s_n0 = k0 * s_k0 + preshuffle_B = fx.Tensor( + fx.make_view( + fx.get_iter(arg_b), + fx.make_layout( + ((16, n0), (kp_elems, 4, k0)), ((s_nlane, s_n0), (1, s_klane, s_k0)) + ), + ) + ) + + # Reshape A and C to 2D + M_max = PRESHUFFLE_M_MAX + arg_a_2d = fx.Tensor( + fx.make_view(fx.get_iter(arg_a), fx.make_layout((M_max, K), (K, 1))) + ) + arg_c_2d = fx.Tensor( + fx.make_view(fx.get_iter(arg_c), fx.make_layout((M_max, N), (N, 1))) + ) + + gx = (i32_m + (tile_m - 1)) // tile_m + gy = i32_n // tile_n + + kernel_gemm( + arg_c_2d, + arg_out, + arg_semaphore, + arg_a_2d, + preshuffle_B, + arg_scale_a, + arg_scale_b, + arg_bias, + i32_m, + i32_n, + tiled_mma, + tiled_copy_g2s, + value_attrs={"rocdl.waves_per_eu": waves_per_eu}, + ).launch( + grid=(gx, gy, split_k), + block=(256, 1, 1), + stream=stream, + ) + + if const_expr(is_f16_or_bf16 and num_acc_n <= 2): + launch_gemm.compile_hints["llvm_options"] = {"enable-post-misched": False} + + return launch_gemm From 1d6c5681ce4f1a7d15cc6000755561953223ce4f Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 09:43:38 +0000 Subject: [PATCH 34/79] bench(cco): apply the scale a tile at a time in preshuffle4w, as CK does The 4-wave port applied the block scale to the whole fragment, which is the expensive way round: a running rescale costs a divide per element per K-block (+204.7us over the unscaled kernel) and whole-fragment promotion needs a duplicate accumulator and spills (2753us). CK does neither -- its loop is static_ford with k innermost, and the partial lives in c_thread_buf_per_scale, one GetRegSizePerXdlops() buffer, 4 VGPRs. That is expressible here after all. frag_C's modes are ((4,1), m_repeat, num_acc_n) and frag_A / frag_B's are (elems, m_repeat|num_acc_n, k_iters), so a tile is a plain slice of each. fx.gemm on the bare atom will not take the slice directly -- it leaves a ub.poison the rmem-to-SSA promotion pass rejects -- so the operands are copied into rmem tensors first, the way Mfma16x16x128 ._make_operand_frag does, hoisted to one per k-step. Worth 36% of the scale cost, 64x128x128 at [16384, 7168] K=2048: whole-fragment rescale 532.4us (+204.7 over the 327.7 unscaled) per-tile promotion 459.5us (+131.8) relL2 1.658e-3 throughout, the fp8 floor. Still not competitive: kernels_fused is 370.5 and CK 300.4. The remaining 131.8us is not the algorithm. Per k-step this adds ~150 VALU ops -- 8 partials zeroed, 6 operand materialisations, 64 promote FMAs -- against ~256 cycles of MFMA, and that VALU only co-issues with the matrix pipe if there are waves to hide it behind. There are not: the same promotion at 128x128x256, whose unscaled form is the fastest tile measured at 284.8us, spills outright. Next lever is occupancy, not the scale loop. The operand materialisation is this kernel's tax, not CK's method's. kernels_fused's Mfma16x16x128 already drives MFMA per tile with materialised operands, and starts from a 224.0us unscaled skeleton against this one's 284.8, so that is where the same promotion should go next. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 2 +- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 177 +++++++++--------- 2 files changed, 92 insertions(+), 87 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 79ad79a45..e4137cee3 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -72,7 +72,7 @@ blockscale instance (64x256, Intrawave v1) measured the same way. The gap is the per-K-block scale, not the GEMM -- unscaled we are 224.0us. See ``kernels_preshuffle4w.py``, which ports CK's 4-wave B-out-of-LDS shape and -lands at 532.4us for exactly that reason. +lands at 459.5us for exactly that reason. Fusing over SDMA wins, and only because of ``--chunks``; see the table at its definition in ``run()``. It was pinned to 1 while the aiter GEMM's diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index d330ef5a0..1b1ac5b07 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -21,31 +21,43 @@ Measured at ``[16384, 7168] K=2048`` on MI355X, warm cache, and the outcome is not the one the shape predicts: -=========================== ========= =========== ========== -kernel no scale blockscale scale cost -=========================== ========= =========== ========== -8-wave 128x256 (kernels_fused) 224.0 370.5 +146.5 -this, 128x128x256 290.9 857.3 +566.4 -this, 64x128x128 (best) -- 532.4 -- -CK 64x256 Intrawave v1 -- 300.4 ~0 -=========================== ========= =========== ========== - -So **the shape is not the problem and neither is B-in-LDS**. Our unscaled GEMM -is already faster than CK's scaled one (224.0 against 300.4); what costs us is -applying the per-K-block scale, and this kernel makes that worse, not better. - -CK gets it for free because it promotes **one MFMA tile at a time**: the loop is -``static_ford`` with k innermost, and the -partial lives in ``c_thread_buf_per_scale``, a single -``GetRegSizePerXdlops()``-wide buffer -- 4 VGPRs. ``fx.gemm`` issues the MFMAs -for a whole fragment at once, so the only two things expressible here are -whole-fragment, and both are expensive: a running rescale needs a divide per -element per K-block (660us; 88us of that recovered with ``v_rcp_f32``), and -whole-fragment promotion needs a duplicate accumulator, which spills -- 2753us -at 128x128x256, the same failure the 8-wave kernel hit at 4052us. - -Closing this means driving MFMAs per accumulator tile rather than per fragment, -which is a deeper change than swapping the pipeline. +============================= ========= =========== ========== +kernel no scale blockscale scale cost +============================= ========= =========== ========== +8-wave 128x256 (kernels_fused) 224.0 370.5 +146.5 +this, 64x128x128 327.7 459.5 +131.8 +this, 128x128x256 284.8 2464.6 (spills) +CK 64x256 Intrawave v1 -- 300.4 ~0 +============================= ========= =========== ========== + +Two things that table says. **The shape is not the problem and neither is +B-in-LDS**: our unscaled 8-wave GEMM at 224.0 is already faster than CK's +*scaled* kernel, and this one's best unscaled tile is 284.8. What costs us is +applying the per-K-block scale. + +And **CK's structure is the right one, but it is not sufficient here**. The +scale is applied a tile at a time, mirroring +``static_ford`` with k innermost and the +partial in ``c_thread_buf_per_scale`` -- one ``GetRegSizePerXdlops()`` buffer, +4 VGPRs. That is worth 36%: the whole-fragment alternatives cost far more, a +running rescale +204.7us (a divide per element per K-block, of which 88us came +back with ``v_rcp_f32``) and a whole-fragment promotion 2753us, the same +duplicate-accumulator spill the 8-wave kernel hit at 4052us. But 131.8us +remains against CK's ~0. + +What is left is not algorithm. Per k-step this adds ~150 VALU ops -- zeroing 8 +partials, 6 operand materialisations, 64 promote FMAs -- against ~256 cycles of +MFMA, and on CDNA that VALU co-issues with the matrix pipe only if there are +waves to hide it behind. Register pressure here allows very few: the same +promotion at 128x128x256, whose unscaled form is the fastest tile measured, +spills outright. So the next lever is occupancy, not the scale loop. + +The operand materialisation is an artefact of this kernel, not of CK's method: +``fx.gemm`` on the bare atom rejects a sliced fragment view (it leaves a +``ub.poison`` the rmem-to-SSA promotion pass then fails on), so A and B have to +be copied into rmem tensors first. ``kernels_fused``'s ``Mfma16x16x128`` already +materialises per tile for its own reasons, so porting this promotion there +would not pay that cost -- and it starts from a 224.0us skeleton. """ import functools @@ -331,10 +343,12 @@ def kernel_gemm( ) bid_x, bid_y = Int32(_bx), Int32(_by) + mma_atom = None if const_expr(use_mfma_scale_128): _scale_atom = fx.make_mma_atom( fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, layout_elem) ) + mma_atom = _scale_atom tiled_mma = fx.make_tiled_mma( _scale_atom, fx.make_layout((1, 4, 1), (0, 1, 0)), @@ -646,78 +660,71 @@ def _load_sb(ni, kb): fx.copy(sc_atom_1, fx.slice(sb_div, (None, fx.Int32(idx))), sc_reg_1) return Vec(fx.memref_load_vec(sc_reg_1))[0] - def rebase_to(kb): - """Rebase the accumulator from K-block ``kb-1``'s scales to ``kb``'s. - - The accumulator is kept as ``C_k / s_k``: one multiply per element - per K-block and no second accumulator, since promotion into a - duplicate of a 64-VGPR frag_C is what spilled 81 registers last - time. Keeping the multiply on the live accumulator also leaves a - real data dependency, so the scheduler cannot hoist every MFMA above - every rescale -- the failure `rocdl.sched_barrier(0)` did not stop. - - The M-tiles are walked load -> divide -> apply one at a time rather - than loading all of them up front. Both orders emit the same - arithmetic, but the batched one keeps 2*m_repeat vec4 of scales live - at once -- as many registers as the accumulator itself -- and that - cost 643us at 128x128x256 against 7 here. - - ``kb - 1`` is clamped rather than branched on: at kb=0 the ratio is - 1, and the accumulator is zero there anyway. + def mma_promote(a_stage, cur_frag_B, ki, kb): + """One scale block's MFMAs, promoted a tile at a time. + + CK's shape (blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle + _v1.hpp): the loop is (m0, n0, kscale) with k innermost and the + partial lives in ``c_thread_buf_per_scale``, a single + ``GetRegSizePerXdlops()``-wide buffer. Only one tile's 4 VGPRs are + live at a time, which is why CK pays almost nothing for the block + scale where both whole-fragment forms are expensive here: a running + rescale needs a divide per element per block (+660us) and a + whole-fragment promotion needs a duplicate accumulator and spills + (2753us). + + frag_C's modes are ((4,1), m_repeat, num_acc_n) and frag_A / + frag_B's are (elems, m_repeat|num_acc_n, k_iters), so one tile is a + plain slice of each; ``fx.gemm`` is handed the bare atom rather than + tiled_mma, exactly as Mfma16x16x128._do_mma does. """ - kb_i = fx.Int32(kb) - zero_i = fx.Int32(0) - prev_kb = (kb_i > zero_i).select(kb_i - fx.Int32(1), zero_i) - # B first: num_acc_n scalars, and every M-tile needs them. - # v_rcp_f32 rather than a divide: an IEEE f32 divide expands to - # ~10 VALU ops and there are m_repeat*4 + num_acc_n of them per - # K-block, which measured 660us against a 291us kernel. rcp is - # 1 ulp, four orders under the 1.66e-3 fp8 floor these operands - # already carry. - rb = [ - _load_sb(ni, prev_kb) * fx.Float32(fx.rocdl.rcp(T.f32, _load_sb(ni, kb_i))) - for ni in range_constexpr(num_acc_n) - ] - acc = Vec(frag_C.load()) - vals = [None] * acc_size - for mi in range_constexpr(m_repeat): - pa = _load_sa(mi, prev_kb) - ca = _load_sa(mi, kb_i) - ra = [ - pa[e] * fx.Float32(fx.rocdl.rcp(T.f32, ca[e])) - for e in range_constexpr(4) - ] - for ni in range_constexpr(num_acc_n): - for ii in range_constexpr(4): - p = ni * (m_repeat * 4) + mi * 4 + ii - vals[p] = acc[p] * (ra[ii] * rb[ni]) - frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) - - def final_scale(): - """Undo the invariant: multiply by the last K-block's scales.""" - last = fx.Int32(kb_count - 1) - sb = [_load_sb(ni, last) for ni in range_constexpr(num_acc_n)] + sb = [_load_sb(ni, kb) for ni in range_constexpr(num_acc_n)] + # Operands are materialised into bare rmem tensors, the way + # Mfma16x16x128._make_operand_frag does: fx.gemm on the atom does + # not take a sliced view, it leaves a ub.poison the rmem-to-SSA + # promotion pass then rejects. Each is built once per k-step, not + # once per (mi, ni) pair. + b_ops = [] + for ni in range_constexpr(num_acc_n): + b_op = fx.make_rmem_tensor(8, fx.Int32) + b_op.store(Vec(cur_frag_B[None, ni, ki].load()).bitcast(fx.Int32)) + b_ops.append(b_op) acc = Vec(frag_C.load()) - vals = [None] * acc_size + vals = [acc[p] for p in range_constexpr(acc_size)] for mi in range_constexpr(m_repeat): - sa = _load_sa(mi, last) + sa = _load_sa(mi, kb) + a_op = fx.make_rmem_tensor(8, fx.Int32) + a_op.store(Vec(frag_A[None, mi, ki].load()).bitcast(fx.Int32)) for ni in range_constexpr(num_acc_n): + b_op = b_ops[ni] + part = fx.make_rmem_tensor(4, fx.Float32) + part.store(Vec.filled(4, 0.0, fx.Float32)) + fx.gemm(mma_atom, part, a_op, b_op, part) + pv = Vec(part.load()) for ii in range_constexpr(4): - p = ni * (m_repeat * 4) + mi * 4 + ii - vals[p] = acc[p] * (sa[ii] * sb[ni]) + # frag_C.load() yields logical order: ii fastest, then + # m_repeat, then num_acc_n. + idx = ni * (m_repeat * 4) + mi * 4 + ii + vals[idx] = vals[idx] + pv[ii] * (sa[ii] * sb[ni]) frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) # ── Pipeline stage (double-buffered B via split fragments) ─ def mma_kloop(a_stage, cur_frag_B, tile_idx=None): for ki in range_constexpr(k_iters): - if const_expr(is_blockscale): - # One MMA k-step is one 128-wide scale block. - rebase_to(fx.Int32(tile_idx) * fx.Int32(k_iters) + fx.Int32(ki)) fx.copy( uni_copy, pA_s2r_stages[a_stage][None, None, ki], frag_A_retile[None, None, ki], ) + if const_expr(is_blockscale): + # One MMA k-step is one 128-wide scale block. + mma_promote( + a_stage, + cur_frag_B, + ki, + fx.Int32(tile_idx) * fx.Int32(k_iters) + fx.Int32(ki), + ) + continue k_coord = ki if (use_mfma_scale_128 or use_mfma_k32) else (None, ki) fx.gemm( tiled_mma, @@ -850,8 +857,8 @@ def two_tiles(k_base): def load_epi_operands(): s_a = s_b = bias = None if const_expr(is_blockscale): - # Nothing to preload: final_scale() applies the last block's - # scales to frag_C directly, so the epilogue below sees an + # Nothing to preload: mma_promote() folds every K-block's + # scales into frag_C as it goes, so the epilogue below sees an # already-scaled accumulator. pass elif const_expr(is_8bit): @@ -921,8 +928,6 @@ def load_epi_operands(): ) # ── Epilogue ───────────────────────────────────────────── - if const_expr(is_blockscale): - final_scale() if const_expr((not is_8bit or is_blockscale) and not _has_epilogue): frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) else: From 4f9b998f3aebcd7b926f411e40b06b0946829abf Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 10:02:56 +0000 Subject: [PATCH 35/79] bench(cco): align preshuffle4w to CK at 64x256, and the GEMM lands on it Moving to CK's own tile first was the right order: at 64x256x128 the per-tile promotion was 760.2us, worse than 64x128's 459.5, because the promote spent three ops per accumulator element and num_acc_n=4 made that four times over. Fixing the shape unlocks the tile rather than the other way round. whole-fragment rescale (divide) 760.2us per-tile promotion 459.5 + explicit fma, scale hoisted 391.7 + MFMAs batched per M-row 385.9 scale removed, per-tile path 303.1 CK 64x256 Intrawave v1 300.4 Written as `acc + pv * (sa * sb)` the compiler emits v_pk_mul_f32 + v_pk_add_f32. Hoisting the scale product out of the element loop and calling fx.math.fma contracts the rest, and the loop body then reads CK 128 v_fmac_f32, 8 v_mul_f32 -> VALU/MFMA 5.81 this 128 v_(pk_)fma, 128 v_(pk_)mul -> VALU/MFMA 5.53 i.e. fewer instructions than CK, same 32 MFMA and 20 ds, VGPR 248 against 232 so the same 2 waves/SIMD. Strip the scale and the kernel lands on CK: 303.1 against 300.4. The remaining +82.6us is a layout difference, not arithmetic. CK needs one FMA per element because c_scale_thread is a single scalar for a whole MFMA tile; we need a multiply first, because A's scale is per row (ScaleBlockM=1) and an MFMA 16x16 lane holds four different rows. CK's four values must be four columns of one row instead -- what kernels_fused's --swap-ab arranges. That is what closing the last 28% takes. Also worth recording: driving MFMAs per accumulator tile beats a single fx.gemm over the whole fragment even with no scaling at all, 303.1us against 556.0 at this tile. relL2 1.658e-3 throughout. Default tile is now 64x256x128. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 2 +- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 100 ++++++++++-------- 2 files changed, 57 insertions(+), 45 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index e4137cee3..b4bf5fc78 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -72,7 +72,7 @@ blockscale instance (64x256, Intrawave v1) measured the same way. The gap is the per-K-block scale, not the GEMM -- unscaled we are 224.0us. See ``kernels_preshuffle4w.py``, which ports CK's 4-wave B-out-of-LDS shape and -lands at 459.5us for exactly that reason. +lands at 385.9us: aligned on the GEMM, still paying for the scale. Fusing over SDMA wins, and only because of ``--chunks``; see the table at its definition in ``run()``. It was pinned to 1 while the aiter GEMM's diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index 1b1ac5b07..f5cfcc343 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -18,46 +18,45 @@ lane pulls its own MFMA operand with one 128-bit load and there is nothing for LDS to amortise. -Measured at ``[16384, 7168] K=2048`` on MI355X, warm cache, and the outcome is -not the one the shape predicts: - -============================= ========= =========== ========== -kernel no scale blockscale scale cost -============================= ========= =========== ========== -8-wave 128x256 (kernels_fused) 224.0 370.5 +146.5 -this, 64x128x128 327.7 459.5 +131.8 -this, 128x128x256 284.8 2464.6 (spills) -CK 64x256 Intrawave v1 -- 300.4 ~0 -============================= ========= =========== ========== - -Two things that table says. **The shape is not the problem and neither is -B-in-LDS**: our unscaled 8-wave GEMM at 224.0 is already faster than CK's -*scaled* kernel, and this one's best unscaled tile is 284.8. What costs us is -applying the per-K-block scale. - -And **CK's structure is the right one, but it is not sufficient here**. The -scale is applied a tile at a time, mirroring -``static_ford`` with k innermost and the -partial in ``c_thread_buf_per_scale`` -- one ``GetRegSizePerXdlops()`` buffer, -4 VGPRs. That is worth 36%: the whole-fragment alternatives cost far more, a -running rescale +204.7us (a divide per element per K-block, of which 88us came -back with ``v_rcp_f32``) and a whole-fragment promotion 2753us, the same -duplicate-accumulator spill the 8-wave kernel hit at 4052us. But 131.8us -remains against CK's ~0. - -What is left is not algorithm. Per k-step this adds ~150 VALU ops -- zeroing 8 -partials, 6 operand materialisations, 64 promote FMAs -- against ~256 cycles of -MFMA, and on CDNA that VALU co-issues with the matrix pipe only if there are -waves to hide it behind. Register pressure here allows very few: the same -promotion at 128x128x256, whose unscaled form is the fastest tile measured, -spills outright. So the next lever is occupancy, not the scale loop. - -The operand materialisation is an artefact of this kernel, not of CK's method: -``fx.gemm`` on the bare atom rejects a sliced fragment view (it leaves a -``ub.poison`` the rmem-to-SSA promotion pass then fails on), so A and B have to -be copied into rmem tensors first. ``kernels_fused``'s ``Mfma16x16x128`` already -materialises per tile for its own reasons, so porting this promotion there -would not pay that cost -- and it starts from a 224.0us skeleton. +Measured at ``[16384, 7168] K=2048`` on MI355X, warm cache, tile 64x256x128 -- +CK's own winning shape: + +============================= ======= ======== +stage us vs CK +============================= ======= ======== +whole-fragment rescale (divide) 760.2 +153% +per-tile promotion 459.5 +53% + + explicit fma, scale hoisted 391.7 +30% + + MFMAs batched per M-row 385.9 +28% +this kernel, scale removed 303.1 +1% +CK 64x256 Intrawave v1 300.4 . +============================= ======= ======== + +So the GEMM itself is aligned: strip the scale and this lands on CK. What is +left is the scale, +82.6us against CK's ~0, and the loop body says why. Per +128 accumulator elements: + + CK 128 v_fmac_f32, 8 v_mul_f32 -> VALU/MFMA 5.81 + this 128 v_(pk_)fma, 128 v_(pk_)mul -> VALU/MFMA 5.53 + +Our instruction *count* is already under CK's; the shape is not. CK needs one +FMA per element because ``c_scale_thread`` is a single scalar for a whole MFMA +tile, where we need a multiply per element first: A's scale is per row +(ScaleBlockM=1) and an MFMA 16x16 lane holds four different rows, so four +different A scales. CK's four values must instead be four *columns* of one row, +which is what ``kernels_fused``'s ``--swap-ab`` arranges. Closing the last 28% +means adopting that output layout, not more arithmetic. + +Two structural notes that came out of the port. B never touches LDS here -- +``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's +``blockscale_b_preshuffle_v1`` does with ``b_thread_bufs`` -- and B is already +preshuffled to ``N0 K0 KLane NLane KPack``, the same layout CK's +``preShuffleBuffer`` emits, so every lane pulls its own MFMA operand with one +128-bit load. And driving MFMAs per accumulator tile beats one ``fx.gemm`` over +the whole fragment even before any scaling: 303.1us against 556.0 at this tile. + +For scale, the 8-wave kernel in ``kernels_fused`` is 224.0us unscaled and 370.5 +scaled at its own best tile. """ import functools @@ -695,17 +694,30 @@ def mma_promote(a_stage, cur_frag_B, ki, kb): sa = _load_sa(mi, kb) a_op = fx.make_rmem_tensor(8, fx.Int32) a_op.store(Vec(frag_A[None, mi, ki].load()).bitcast(fx.Int32)) + # All of this M-row's MFMAs are issued before any of their + # results is read. With tile_k=128 a scale block is a single + # MFMA, so promoting each one where it is issued exposes the + # full MFMA latency; CK hides it by keeping k innermost, which + # is not available here. num_acc_n partials cost 4*num_acc_n + # VGPRs. + parts = [] for ni in range_constexpr(num_acc_n): - b_op = b_ops[ni] part = fx.make_rmem_tensor(4, fx.Float32) part.store(Vec.filled(4, 0.0, fx.Float32)) - fx.gemm(mma_atom, part, a_op, b_op, part) - pv = Vec(part.load()) + fx.gemm(mma_atom, part, a_op, b_ops[ni], part) + parts.append(part) + for ni in range_constexpr(num_acc_n): + sab = [sa[ii] * sb[ni] for ii in range_constexpr(4)] + pv = Vec(parts[ni].load()) + # Explicit fma: written as ``acc + pv * (sa * sb)`` the + # compiler emits v_pk_mul_f32 + v_pk_add_f32, three ops per + # element against CK's one v_fmac_f32. Hoisting the scale + # product and contracting the rest gets it to two. for ii in range_constexpr(4): # frag_C.load() yields logical order: ii fastest, then # m_repeat, then num_acc_n. idx = ni * (m_repeat * 4) + mi * 4 + ii - vals[idx] = vals[idx] + pv[ii] * (sa[ii] * sb[ni]) + vals[idx] = fx.math.fma(pv[ii], sab[ii], vals[idx]) frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) # ── Pipeline stage (double-buffered B via split fragments) ─ From fc89289489cc43fd5451df04f24676afcdd18e14 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 10:33:30 +0000 Subject: [PATCH 36/79] bench(cco): migrate swap_ab to preshuffle4w -- it aligns the instructions and still loses Exchanging the MFMA operands so a lane owns four consecutive N of one row instead of four M is the last piece of CK's shape, and it delivers both things it should: A's scale collapses to one scalar per MFMA tile, because the four values are now one row rather than four, so the promote is a single FMA like CK's c_scale_thread -- VALU/MFMA 5.53 -> 4.47, against CK's 5.81. C's four values are four consecutive columns, 8 contiguous bytes in a row-major C -- 64 buffer_store_short -> 16 buffer_store_dwordx2. And it is 31% slower: 385.9us -> 506.9 at 64x256x128, [16384, 7168] K=2048. VGPRs go 248 -> 256 and the main loop starts spilling; scratch_load_dwordx2 x3 and scratch_store_dwordx2 x2 appear inside it. That it is the loop and not the new epilogue is settled by quadrupling K: the gap widens with it, 1.30x at K=2048 against 1.47x at K=8192. Moving the store registers out of kernel scope does not recover it, so the pressure is somewhere else. Kept behind swap_ab=False (--p4w-swap-ab), because the instruction-level result is right and reusable and what blocks it is a register budget, not the idea. Its own flag rather than the 8-wave --swap-ab, which is on by default and which --permlane requires. bench also stops compiling the 8-wave kernel when it is not the one under test. relL2 1.658e-3 on both paths. tests/python/cco/test_gemm_ar.py currently shows 3 failures, and they reproduce on the unmodified tree (verified by stashing): another container is back on the box at 89% VRAM, which starves the subprocess tests and trips the known aiter large-grid non-determinism the bitwise test checks. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 11 +- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 103 ++++++++++++++++-- 2 files changed, 106 insertions(+), 8 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index b4bf5fc78..6424e26ab 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -318,13 +318,14 @@ def run(args) -> int: in_dtype="fp8", out_dtype="bf16", quant=args.quant, + swap_ab=args.p4w_swap_ab, waves_per_eu=args.waves_per_eu, xcd_swizzle=args.xcd_swizzle, ) semaphore = torch.zeros(1, device="cuda", dtype=torch.int32) bias_unused = torch.zeros(1, device="cuda", dtype=torch.bfloat16) - gemm = compile_fused_gemm_scatter( + gemm = None if args.gemm_impl == "preshuffle4w" else compile_fused_gemm_scatter( cfg, rank, K=args.k, @@ -664,6 +665,14 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("--tile-m", type=int, default=64, help="preshuffle4w only") p.add_argument("--tile-n", type=int, default=256, help="preshuffle4w only") p.add_argument("--tile-k", type=int, default=128, help="preshuffle4w only") + p.add_argument( + "--p4w-swap-ab", + action=argparse.BooleanOptionalAction, + default=False, + help="preshuffle4w: exchange the MFMA operands so a lane owns four " + "consecutive N instead of four M. Off by default -- it does what it " + "should to the instructions and still loses; see kernels_preshuffle4w.", + ) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) p.add_argument("--eager", action="store_true") diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index f5cfcc343..b988ecc37 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -47,6 +47,27 @@ which is what ``kernels_fused``'s ``--swap-ab`` arranges. Closing the last 28% means adopting that output layout, not more arithmetic. +``swap_ab`` exchanges the MFMA operands so a lane owns four consecutive N of +one row instead of four M. It does exactly what it is supposed to on both +counts CK benefits from -- A's scale collapses to one scalar per tile, so the +promote is a single FMA like ``c_scale_thread``, and the C store goes from 64 +``buffer_store_short`` to 16 ``buffer_store_dwordx2``: + +============================ ========= ========== ============ + VALU/MFMA C stores total (us) +============================ ========= ========== ============ +off 5.53 64 short 385.9 +on 4.47 16 dwordx2 506.9 +============================ ========= ========== ============ + +and it is 31% slower anyway, because at VGPR 248 -> 256 the main loop starts +spilling: `scratch_load_dwordx2` x3 and `scratch_store_dwordx2` x2 appear +inside it. Confirmed to be the loop and not the epilogue by quadrupling K -- +the gap widens with it (1.30x at K=2048, 1.47x at K=8192). Moving the store +registers out of kernel scope does not recover it, so the pressure is +elsewhere. Left in behind ``swap_ab=False``: the instruction-level result is +right and reusable, and what blocks it is a register budget, not the idea. + Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's ``blockscale_b_preshuffle_v1`` does with ``b_thread_bufs`` -- and B is already @@ -192,6 +213,7 @@ def compile_preshuffle_gemm( in_dtype: str = "fp8", out_dtype: str = "bf16", quant: str = "ptpc", # "ptpc" (upstream) or "blockscale" (A 1x128, B 128x128) + swap_ab: bool = False, # blockscale only; see the module docstring epilogue: str = "none", # "none", "bias", "bias_relu", "bias_silu", "bias_gelu" waves_per_eu: int | None = None, enable_scheduler: bool = True, @@ -232,6 +254,7 @@ def compile_preshuffle_gemm( if quant not in ("ptpc", "blockscale"): raise ValueError(f"quant must be ptpc or blockscale, got {quant!r}") is_blockscale = quant == "blockscale" + swap_ab = swap_ab and is_blockscale if is_blockscale: # tile_k == 128 makes one main-loop tile exactly one K-scale block, so # the rescale lands on a loop boundary that already exists. @@ -632,18 +655,45 @@ def hot_loop_scheduler(): gSB = fx.rocdl.make_buffer_tensor(arg_scale_b) sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + # C write-out for the swapped accumulator. The tiled_mma-derived + # path (thr_r2g_C / pC_g) partitions C for the *unswapped* layout, + # so it would scatter these four values down four rows. They are + # four consecutive columns of one row instead, which in a row-major + # C is 8 contiguous bytes -- one store, where the unswapped layout + # can only ever emit buffer_store_short. + gC_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor( + fx.make_view( + fx.get_iter(c_tensor), + fx.make_layout(PRESHUFFLE_M_MAX * N, 1), + ) + ), + max_size=False, + num_records_bytes=fx.Int64(i32_m) + * fx.Int64(N) + * fx.Int64(out_elem_bytes), + ) + c_div = fx.logical_divide(gC_flat, fx.make_layout(1, 1)) + c_atom_64 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), out_elem_cls) sc_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) sc_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) sc_reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + sa_reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) sc_reg_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) def _load_sa(mi, kb): - """This lane's four A-scale rows of M-tile ``mi`` at K-block ``kb``. + """This lane's A scale for M-tile ``mi`` at K-block ``kb``. ``arg_scale_a`` is physically [K/128, M] (column-major [M, K/128]), - the layout aiter_per1x128_quant(transpose_scale=True) emits, so the - four rows are contiguous and land in one vec4. + the layout aiter_per1x128_quant(transpose_scale=True) emits. With + the operands swapped a lane owns a single row -- ``lane % 16`` -- + so this is one scalar, and combined with B's it gives the one + per-tile ``c_scale_thread`` CK promotes with. """ + if const_expr(swap_ab): + idx = fx.Int32(kb) * fx.Int32(i32_m) + bx_m + mi * 16 + lane_mod_16 + fx.copy(sc_atom_1, fx.slice(sa_div, (None, fx.Int32(idx))), sa_reg_1) + return Vec(fx.memref_load_vec(sa_reg_1))[0] idx = fx.Int32(kb) * fx.Int32(i32_m) + bx_m + mi * 16 + lane_div_16 * 4 fx.copy(sc_atom_4, fx.slice(sa_div, (None, fx.Int32(idx))), sc_reg_4) return Vec(fx.memref_load_vec(sc_reg_4)) @@ -704,10 +754,21 @@ def mma_promote(a_stage, cur_frag_B, ki, kb): for ni in range_constexpr(num_acc_n): part = fx.make_rmem_tensor(4, fx.Float32) part.store(Vec.filled(4, 0.0, fx.Float32)) - fx.gemm(mma_atom, part, a_op, b_ops[ni], part) + # Operands exchanged: fx.gemm(atom, c, b, a, c) computes + # (A B)^T in the accumulator's native layout, so a lane's + # four values become four consecutive *columns* of one row + # instead of four rows. gcnasm's mfma_adaptor_swap_ab. + if const_expr(swap_ab): + fx.gemm(mma_atom, part, b_ops[ni], a_op, part) + else: + fx.gemm(mma_atom, part, a_op, b_ops[ni], part) parts.append(part) for ni in range_constexpr(num_acc_n): - sab = [sa[ii] * sb[ni] for ii in range_constexpr(4)] + sab = ( + [sa * sb[ni]] * 4 + if const_expr(swap_ab) + else [sa[ii] * sb[ni] for ii in range_constexpr(4)] + ) pv = Vec(parts[ni].load()) # Explicit fma: written as ``acc + pv * (sa * sb)`` the # compiler emits v_pk_mul_f32 + v_pk_add_f32, three ops per @@ -940,7 +1001,34 @@ def load_epi_operands(): ) # ── Epilogue ───────────────────────────────────────────── - if const_expr((not is_8bit or is_blockscale) and not _has_epilogue): + if const_expr(swap_ab): + # Allocated here, not at kernel scope: held across the main loop + # these push VGPRs over the 2-waves/SIMD line and the loop spills + # to scratch, which cost more than the wider store saved. + c_reg_4 = [ + fx.make_rmem_tensor(fx.make_layout(4, 1), out_elem_cls) + for _ in range(m_repeat * num_acc_n) + ] + acc_vec = Vec(frag_C.load()) + for mi in range_constexpr(m_repeat): + row = bx_m + mi * 16 + lane_mod_16 + for ni in range_constexpr(num_acc_n): + col = by_n + (ni * num_waves + wave_id) * 16 + lane_div_16 * 4 + creg = c_reg_4[mi * num_acc_n + ni] + creg.store( + vector.from_elements( + T.vec(4, out_elem_cls.ir_type), + [ + acc_vec[ni * (m_repeat * 4) + mi * 4 + ii].to( + out_elem_cls + ) + for ii in range_constexpr(4) + ], + ) + ) + idx = fx.Int32(row) * fx.Int32(N) + fx.Int32(col) + fx.copy(c_atom_64, creg, fx.slice(c_div, (None, idx))) + elif const_expr((not is_8bit or is_blockscale) and not _has_epilogue): frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) else: if const_expr(not overlap_epi_load): @@ -991,7 +1079,8 @@ def apply_activation(val_s): ) frag_C_out.store(out_vec) - fx.copy(buf_copy_out, frag_C_retile, pC_g) + if const_expr(not swap_ab): + fx.copy(buf_copy_out, frag_C_retile, pC_g) if const_expr(split_k > 1): # The store carries sc0|sc1, so waiting on it is the whole release. rocdl.s_waitcnt(0) From ed2e66aa10e65af70b4bcb7c8d7d9c46c4cd9b8d Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 11:02:54 +0000 Subject: [PATCH 37/79] bench(cco): attribute swap_ab's register cost -- it is entirely the C write-out --waves-per-eu 2 caps VGPRs at 256, so the question is which part of the swap crosses it. Compiling the three parts separately (single-GPU probe, no communicator needed, since the ISA is dumped at compile time): variant VGPR scratch off (the default) 256 8 exchange MFMA operands 260 8 scalar A scale 240 0 explicit C store 284 8 operands + scalar scale 248 0 all three 256 8 The scale change saves 16 registers and the operand exchange costs 4. Together they are the only variant that does not spill -- and they spill less than the current default, which turns out to spill too. The entire cost is the hand written C store. Which part of it was not obvious. Not its landing registers: one shared register measures the same as sixteen. Not frag_C_out left allocated for a path that no longer uses it. Not interleaving the bf16 conversion with the stores. It was partly the addressing -- rebuilding 16 runtime indices as one base plus compile-time strides (mi steps 16 rows, ni steps num_waves*16 columns, both static) is worth 284 -> 270 unconstrained and 506.9 -> 485.1us, and that improvement is kept -- but under the cap it still spills. So what is missing is a store for the transposed accumulator that costs no registers: the tiled-copy path rebuilt against a transposed C partitioning, instead of written by hand. With that, the measured 248/0 applies and swap_ab becomes strictly better than a default that spills. Still behind swap_ab=False. relL2 1.658e-3 on both paths; 388.9us off, 488.7 on at 64x256x128, [16384, 7168] K=2048. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 77 ++++++++++++++----- 1 file changed, 59 insertions(+), 18 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index b988ecc37..3606adcf7 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -60,13 +60,34 @@ on 4.47 16 dwordx2 506.9 ============================ ========= ========== ============ -and it is 31% slower anyway, because at VGPR 248 -> 256 the main loop starts -spilling: `scratch_load_dwordx2` x3 and `scratch_store_dwordx2` x2 appear -inside it. Confirmed to be the loop and not the epilogue by quadrupling K -- -the gap widens with it (1.30x at K=2048, 1.47x at K=8192). Moving the store -registers out of kernel scope does not recover it, so the pressure is -elsewhere. Left in behind ``swap_ab=False``: the instruction-level result is -right and reusable, and what blocks it is a register budget, not the idea. +and it is 26% slower anyway, on a register budget. ``--waves-per-eu 2`` caps +VGPRs at 256, and compiling the three parts of the swap separately says where +they go: + +=========================== ====== ========= +variant VGPR scratch +=========================== ====== ========= +off (this kernel's default) 256 8 + exchange MFMA operands 260 8 + scalar A scale 240 0 + explicit C store 284 8 +**operands + scalar scale** **248** **0** +all three (swap_ab=True) 256 8 +=========================== ====== ========= + +The scale change *saves* 16 registers and the operand exchange costs 4; both +together are the only variant here that does not spill, and it spills less than +the default does. **The whole cost is the C write-out.** Not its registers +(one shared landing register measures the same as sixteen), not ``frag_C_out`` +left allocated, not the conversion order, and not the addressing -- rebuilding +the 16 runtime indices as one base plus compile-time strides took it 284 -> 270 +unconstrained and 506.9 -> 485.1us, but under the cap it still spills. + +So the missing piece is a store for the transposed accumulator that costs no +registers, i.e. the tiled-copy path (``thr_r2g_C``/``pC_g``) rebuilt against a +transposed C partitioning, rather than written out by hand. Then the measured +248/0 applies and swap_ab is strictly better than the default, which spills. +Left in behind ``swap_ab=False`` until that exists. Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's @@ -465,10 +486,14 @@ def _make_sA(arr): out_cpol = CPOL_COHERENT if split_k > 1 else 0 copy_op = fx.rocdl.BufferCopy32b if split_k > 1 else fx.rocdl.BufferCopy16b buf_copy_out = fx.make_copy_atom(copy_op(out_cpol), out_elem_cls) - thr_r2g_C = fx.make_tiled_copy_C(buf_copy_out, tiled_mma).get_slice(tid) - pC_g = thr_r2g_C.partition_S(tC) - frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) - frag_C_retile = thr_r2g_C.retile(frag_C_out) + # Only built for the tiled write-out. Under the swapped layout the + # explicit store replaces it, and leaving frag_C_out allocated costs + # acc_size out-elements of register space for nothing. + if const_expr(not swap_ab): + thr_r2g_C = fx.make_tiled_copy_C(buf_copy_out, tiled_mma).get_slice(tid) + pC_g = thr_r2g_C.partition_S(tC) + frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) + frag_C_retile = thr_r2g_C.retile(frag_C_out) # ── Async gmem->LDS DMA (buffer_load_lds) for the A tile ── if const_expr(use_async_copy): @@ -1009,25 +1034,41 @@ def load_epi_operands(): fx.make_rmem_tensor(fx.make_layout(4, 1), out_elem_cls) for _ in range(m_repeat * num_acc_n) ] + # Convert first, store second: interleaving them keeps the f32 + # accumulator alive alongside the growing set of converted values. + # One base address plus compile-time strides, not 16 runtime + # indices: mi steps 16 rows and ni steps num_waves*16 columns, both + # static, so the stores use the buffer's immediate offset instead of + # a VGPR each. + base = ( + (bx_m + lane_mod_16) * fx.Int32(N) + + by_n + + wave_id * 16 + + lane_div_16 * 4 + ) + cview = fx.make_view( + fx.add_offset(fx.get_iter(gC_flat), fx.Int32(base)), + fx.make_layout( + (4, m_repeat, num_acc_n), (1, 16 * N, num_waves * 16) + ), + ) acc_vec = Vec(frag_C.load()) + out_elems = [ + acc_vec[p].to(out_elem_cls) for p in range_constexpr(acc_size) + ] for mi in range_constexpr(m_repeat): - row = bx_m + mi * 16 + lane_mod_16 for ni in range_constexpr(num_acc_n): - col = by_n + (ni * num_waves + wave_id) * 16 + lane_div_16 * 4 creg = c_reg_4[mi * num_acc_n + ni] creg.store( vector.from_elements( T.vec(4, out_elem_cls.ir_type), [ - acc_vec[ni * (m_repeat * 4) + mi * 4 + ii].to( - out_elem_cls - ) + out_elems[ni * (m_repeat * 4) + mi * 4 + ii] for ii in range_constexpr(4) ], ) ) - idx = fx.Int32(row) * fx.Int32(N) + fx.Int32(col) - fx.copy(c_atom_64, creg, fx.slice(c_div, (None, idx))) + fx.copy(c_atom_64, creg, cview[None, mi, ni]) elif const_expr((not is_8bit or is_blockscale) and not _has_epilogue): frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) else: From 63bc1ab51831fa6476c9a77cd735ed107d24e3fb Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 11:23:07 +0000 Subject: [PATCH 38/79] bench(cco): give swap_ab a register-free write-out -- and it changes nothing The missing piece was a store for the transposed accumulator that costs no registers. Two fixes get there: The tiled copy, not a hand-written store. The swapped accumulator maps to C with M and N exchanged, so the ordinary partitioning applied to C^T[N, M] (strides (1, N)) lands on the right addresses: tiled_mma rebuilt with the waves raked along the first mode, partition_S against a transposed C tile, and a 64b copy atom since a tile's four values are now contiguous. The hand-written store this replaces was the entire register cost -- 28 VGPRs, against the scalar A scale's -16 and the operand exchange's +4. And the accumulator written in the swapped slot order from the start. The T-partition wants (v, n_rep, m_rep) where frag_C is (v, m_rep, n_rep); permuting in the epilogue is correct but costs 8 VGPRs because source and destination are live together. mma_promote chooses the order instead, which is free. That lands exactly where the bisect predicted: 248 VGPR, zero scratch, 16 buffer_store_dwordx2, one FMA per accumulator element. And 387.4us against the unswapped path's 387.5 -- parity, at [16384, 7168] K=2048, 64x256x128. So none of it was the bottleneck. Not the promote arithmetic, three ops per element down to one. Not the spill, which the default path had too. Not the store width. The +85us the block scale costs over this kernel's own 303.1us skeleton survives all three, while CK pays ~0 for the same work -- same v_mfma_f32_16x16x128_f8f6f4, same 2 waves/SIMD, and a VALU/MFMA of 5.81 against our 4.47. We issue less and it hides worse, which points at scheduling: Intrawave plus explicit sched_group_barrier placement on CK's side. That is not visible in an instruction mix and is where this should go next. relL2 1.658e-3 on both paths. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 193 +++++++++--------- 1 file changed, 91 insertions(+), 102 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index 3606adcf7..9c6fe5270 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -60,34 +60,37 @@ on 4.47 16 dwordx2 506.9 ============================ ========= ========== ============ -and it is 26% slower anyway, on a register budget. ``--waves-per-eu 2`` caps -VGPRs at 256, and compiling the three parts of the swap separately says where -they go: - -=========================== ====== ========= -variant VGPR scratch -=========================== ====== ========= -off (this kernel's default) 256 8 - exchange MFMA operands 260 8 - scalar A scale 240 0 - explicit C store 284 8 -**operands + scalar scale** **248** **0** -all three (swap_ab=True) 256 8 -=========================== ====== ========= - -The scale change *saves* 16 registers and the operand exchange costs 4; both -together are the only variant here that does not spill, and it spills less than -the default does. **The whole cost is the C write-out.** Not its registers -(one shared landing register measures the same as sixteen), not ``frag_C_out`` -left allocated, not the conversion order, and not the addressing -- rebuilding -the 16 runtime indices as one base plus compile-time strides took it 284 -> 270 -unconstrained and 506.9 -> 485.1us, but under the cap it still spills. - -So the missing piece is a store for the transposed accumulator that costs no -registers, i.e. the tiled-copy path (``thr_r2g_C``/``pC_g``) rebuilt against a -transposed C partitioning, rather than written out by hand. Then the measured -248/0 applies and swap_ab is strictly better than the default, which spills. -Left in behind ``swap_ab=False`` until that exists. +Both are done now, and neither moved the clock. Getting there took two fixes +that are worth recording, because each was found by attribution rather than by +guessing: + +* ``--waves-per-eu 2`` caps VGPRs at 256, and the swap crosses it. Compiling + its three parts separately: the scalar A scale *saves* 16 registers, the + operand exchange costs 4, and the hand-written C store costs 28. The store + was the whole problem -- and not its landing registers (one measures the same + as sixteen), nor ``frag_C_out``, nor the conversion order. Going through the + tiled copy against a transposed C partitioning instead costs nothing: + ``tiled_mma`` rebuilt with the waves raked along the first mode, partitioning + ``C^T[N, M]`` (strides ``(1, N)``), which is the same map with M and N + exchanged. +* That partition wants the fragment in ``(v, n_rep, m_rep)`` order where + ``frag_C`` is ``(v, m_rep, n_rep)``. Permuting in the epilogue is correct but + costs 8 VGPRs, because source and destination are live together. The slot + order is ours to choose, so ``mma_promote`` writes the swapped order from the + start and there is nothing to permute. + +Result: 248 VGPR, no spill, 16 ``buffer_store_dwordx2``, one FMA per element. +And 387.4us against the unswapped path's 387.5 -- **exactly parity**. + +So none of it was the bottleneck: not the promote arithmetic (three ops per +element down to one), not the spill, not the store width. The +85us the block +scale costs over this kernel's own 303.1us skeleton survives all three, while +CK pays ~0 for the same work with the same ``v_mfma_f32_16x16x128_f8f6f4`` and +the same 2 waves/SIMD. Our VALU/MFMA is already *under* CK's, 4.47 against +5.81, so what is left is that CK's VALU hides behind its MFMAs and ours does +not -- scheduling, which for CK is Intrawave plus explicit sched_group_barrier +placement. That is the next thing to look at, and it is not something the +instruction mix shows. Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's @@ -486,13 +489,55 @@ def _make_sA(arr): out_cpol = CPOL_COHERENT if split_k > 1 else 0 copy_op = fx.rocdl.BufferCopy32b if split_k > 1 else fx.rocdl.BufferCopy16b buf_copy_out = fx.make_copy_atom(copy_op(out_cpol), out_elem_cls) - # Only built for the tiled write-out. Under the swapped layout the - # explicit store replaces it, and leaving frag_C_out allocated costs - # acc_size out-elements of register space for nothing. - if const_expr(not swap_ab): + frag_C_out = None + if const_expr(swap_ab): + # The swapped accumulator maps to C with M and N exchanged, so the + # ordinary partitioning applied to C^T lands on the right addresses. + # C is row-major [M, N], so C^T is [N, M] with strides (1, N) -- a + # tile's four values are then contiguous, which is what lets the + # store be one 64-bit instead of four 16-bit ones. Going through + # the tiled copy rather than by hand is the point: a hand-written + # store needs landing registers and costs the 8 VGPRs that push + # this over --waves-per-eu 2 and make the main loop spill. + tiled_mma_T = fx.make_tiled_mma( + mma_atom, + fx.make_layout((4, 1, 1), (1, 0, 0)), + fx.make_tile(None, None, fx.make_layout((32, 4), (1, 32))), + ) + c_tensor_T = fx.Tensor( + fx.make_view( + fx.get_iter(arg_c), + fx.make_layout((N, PRESHUFFLE_M_MAX), (1, N)), + ) + ) + gC_T = fx.rocdl.make_buffer_tensor( + c_tensor_T, + max_size=False, + num_records_bytes=fx.Int64(i32_m) + * fx.Int64(N) + * fx.Int64(out_elem_bytes), + ) + tC_T = fx.flat_divide(gC_T, fx.make_tile(tile_n, tile_m))[ + None, None, bid_y, bid_x + ] + copy_out_c = fx.make_copy_atom( + fx.rocdl.BufferCopy64b(out_cpol), out_elem_cls + ) + # The fragment has to carry the transposed mode order too: the + # T-partition expects (v, n_rep, m_rep) where frag_C is + # (v, m_rep, n_rep), so the fill below permutes the slots. + frag_C_out = fx.make_fragment_like( + tiled_mma_T.thr_slice(tid).make_fragment_C(tC_T), + out_elem_cls.ir_type, + ) + thr_r2g_C = fx.make_tiled_copy_C(copy_out_c, tiled_mma_T).get_slice(tid) + pC_g = thr_r2g_C.partition_S(tC_T) + frag_C_retile = thr_r2g_C.retile(frag_C_out) + else: + frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) + copy_out_c = buf_copy_out thr_r2g_C = fx.make_tiled_copy_C(buf_copy_out, tiled_mma).get_slice(tid) pC_g = thr_r2g_C.partition_S(tC) - frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) frag_C_retile = thr_r2g_C.retile(frag_C_out) # ── Async gmem->LDS DMA (buffer_load_lds) for the A tile ── @@ -680,26 +725,6 @@ def hot_loop_scheduler(): gSB = fx.rocdl.make_buffer_tensor(arg_scale_b) sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) - # C write-out for the swapped accumulator. The tiled_mma-derived - # path (thr_r2g_C / pC_g) partitions C for the *unswapped* layout, - # so it would scatter these four values down four rows. They are - # four consecutive columns of one row instead, which in a row-major - # C is 8 contiguous bytes -- one store, where the unswapped layout - # can only ever emit buffer_store_short. - gC_flat = fx.rocdl.make_buffer_tensor( - fx.Tensor( - fx.make_view( - fx.get_iter(c_tensor), - fx.make_layout(PRESHUFFLE_M_MAX * N, 1), - ) - ), - max_size=False, - num_records_bytes=fx.Int64(i32_m) - * fx.Int64(N) - * fx.Int64(out_elem_bytes), - ) - c_div = fx.logical_divide(gC_flat, fx.make_layout(1, 1)) - c_atom_64 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), out_elem_cls) sc_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) sc_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) sc_reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) @@ -800,9 +825,17 @@ def mma_promote(a_stage, cur_frag_B, ki, kb): # element against CK's one v_fmac_f32. Hoisting the scale # product and contracting the rest gets it to two. for ii in range_constexpr(4): - # frag_C.load() yields logical order: ii fastest, then - # m_repeat, then num_acc_n. - idx = ni * (m_repeat * 4) + mi * 4 + ii + # Slot order is ours to choose, and choosing it once + # here is free: the swapped write-out wants + # (ii, num_acc_n, m_repeat), and producing the + # accumulator in that order avoids a 64-value permute + # in the epilogue, which costs 8 VGPRs because source + # and destination are live together. + idx = ( + mi * (num_acc_n * 4) + ni * 4 + ii + if const_expr(swap_ab) + else ni * (m_repeat * 4) + mi * 4 + ii + ) vals[idx] = fx.math.fma(pv[ii], sab[ii], vals[idx]) frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) @@ -1026,50 +1059,7 @@ def load_epi_operands(): ) # ── Epilogue ───────────────────────────────────────────── - if const_expr(swap_ab): - # Allocated here, not at kernel scope: held across the main loop - # these push VGPRs over the 2-waves/SIMD line and the loop spills - # to scratch, which cost more than the wider store saved. - c_reg_4 = [ - fx.make_rmem_tensor(fx.make_layout(4, 1), out_elem_cls) - for _ in range(m_repeat * num_acc_n) - ] - # Convert first, store second: interleaving them keeps the f32 - # accumulator alive alongside the growing set of converted values. - # One base address plus compile-time strides, not 16 runtime - # indices: mi steps 16 rows and ni steps num_waves*16 columns, both - # static, so the stores use the buffer's immediate offset instead of - # a VGPR each. - base = ( - (bx_m + lane_mod_16) * fx.Int32(N) - + by_n - + wave_id * 16 - + lane_div_16 * 4 - ) - cview = fx.make_view( - fx.add_offset(fx.get_iter(gC_flat), fx.Int32(base)), - fx.make_layout( - (4, m_repeat, num_acc_n), (1, 16 * N, num_waves * 16) - ), - ) - acc_vec = Vec(frag_C.load()) - out_elems = [ - acc_vec[p].to(out_elem_cls) for p in range_constexpr(acc_size) - ] - for mi in range_constexpr(m_repeat): - for ni in range_constexpr(num_acc_n): - creg = c_reg_4[mi * num_acc_n + ni] - creg.store( - vector.from_elements( - T.vec(4, out_elem_cls.ir_type), - [ - out_elems[ni * (m_repeat * 4) + mi * 4 + ii] - for ii in range_constexpr(4) - ], - ) - ) - fx.copy(c_atom_64, creg, cview[None, mi, ni]) - elif const_expr((not is_8bit or is_blockscale) and not _has_epilogue): + if const_expr((not is_8bit or is_blockscale) and not _has_epilogue): frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) else: if const_expr(not overlap_epi_load): @@ -1120,8 +1110,7 @@ def apply_activation(val_s): ) frag_C_out.store(out_vec) - if const_expr(not swap_ab): - fx.copy(buf_copy_out, frag_C_retile, pC_g) + fx.copy(copy_out_c, frag_C_retile, pC_g) if const_expr(split_k > 1): # The store carries sc0|sc1, so waiting on it is the whole release. rocdl.s_waitcnt(0) From c4c536a775f957e30ddb4e3e89b7ea4a15ff850c Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 11:46:41 +0000 Subject: [PATCH 39/79] bench(cco): the scheduling lead is dead, and so is the premise it rested on Chasing the interleaving: CK's loop body alternates MFMA and VALU evenly, a median of 4 VALU between MFMAs with 2 of 31 gaps empty, where ours puts MFMA clusters next to 40-80 instruction VALU stretches with 8 of 31 gaps empty. FlyDSL's sched_* helpers cover MFMA/DS/VMEM and not VALU, so nothing pins the promote -- which looked like the whole explanation. It is not. sched_group_barrier(0x002, n, 0) -- LLVM's VALU mask, which the helpers simply do not wrap -- is worth 0-2us at every group size from 2 to 8. Batching more MFMAs before promoting them is worth under 4us across 1, 2, 4, 8 and 16 tiles. Hoisting the scale addresses' loop invariants is actively worse, 409us against 387, because twelve extra live values cost more than the signed divide-by-128 expansion they remove -- that expansion, v_mul_hi_u32 / ashrrev / xor / sub / max sitting in the middle of the promote's FMAs, was what the longest VALU run turned out to contain. And the premise was wrong. Replacing the promote's FMA with a plain add measures 387.4us against 385.4 with it; feeding constant scales instead of loading them measures 396.8 against 395.6. The block scale is free in this kernel. The 303.1us "skeleton" this whole line was aimed at closing was measured the same way before the swap work and does not reproduce, so the +85us it implied never existed as stated. The gap to CK is therefore the GEMM itself, not the quantisation and not the scale's scheduling: 388.4 (swap off) / 388.8 (swap on) against 300.5, and both paths measure the same with the scale removed. Recorded rather than pursued. No functional change; docstring only. relL2 1.658e-3 both paths. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 31 +++++++++++++------ 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index 9c6fe5270..f13d8077a 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -82,15 +82,28 @@ Result: 248 VGPR, no spill, 16 ``buffer_store_dwordx2``, one FMA per element. And 387.4us against the unswapped path's 387.5 -- **exactly parity**. -So none of it was the bottleneck: not the promote arithmetic (three ops per -element down to one), not the spill, not the store width. The +85us the block -scale costs over this kernel's own 303.1us skeleton survives all three, while -CK pays ~0 for the same work with the same ``v_mfma_f32_16x16x128_f8f6f4`` and -the same 2 waves/SIMD. Our VALU/MFMA is already *under* CK's, 4.47 against -5.81, so what is left is that CK's VALU hides behind its MFMAs and ours does -not -- scheduling, which for CK is Intrawave plus explicit sched_group_barrier -placement. That is the next thing to look at, and it is not something the -instruction mix shows. +So none of it was the bottleneck. Nor, it turns out, is the block scale: +replacing the promote's FMA with a plain add measures 387.4us against 385.4 +with it, and feeding constant scales instead of loading them measures 396.8 +against 395.6. **The scale is free in this kernel.** An earlier 303.1us +reading, taken the same way before the swap work, does not reproduce and should +not be relied on. + +Which moves the question: the gap to CK is the GEMM, not the quantisation. The +scheduling hypothesis this was chasing is dead too. FlyDSL's ``sched_*`` +helpers cover MFMA/DS/VMEM but not VALU, so nothing pins the promote -- but +adding VALU groups by hand (``sched_group_barrier(0x002, n, 0)``, LLVM's VALU +mask) is worth 0-2us at any group size from 2 to 8. So is batching more MFMAs +before promoting them: 1, 2, 4, 8 and 16 tiles all land within 4us of each +other. Hoisting the scale addresses' loop invariants out is actively worse, 409 +against 387, because the twelve live values cost more than the signed +divide-by-128 expansion they remove. + +The one thing the ISA does show, and which nothing above changed, is shape: +CK's loop body interleaves MFMA and VALU evenly -- median 4 VALU between MFMAs, +2 of 31 gaps empty -- where ours alternates MFMA clusters with 40-80 +instruction VALU stretches, 8 of 31 gaps empty. Whatever produces that sits +upstream of the group barriers. Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's From ec1f31b662ea8c9604c88881a80d00918e04fe3a Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 12:36:44 +0000 Subject: [PATCH 40/79] bench(cco): read the loop bodies side by side -- the gap is stalls, not work Our loop body is shorter than CK's: 312 instructions swap-off and 272 swap-on against 322, with identical MFMA (32) and LDS (20) counts and fewer VALU (184 / 145 against 186). Fewer instructions and 29% slower means cycles per instruction, and two things in the ISA fit. We drain the memory queue much harder. CK waits at s_waitcnt vmcnt(15), vmcnt(16), vmcnt(10) -- ten to sixteen loads still in flight, waiting only on the oldest. Ours are five s_waitcnt vmcnt(1) per loop body swap-off, 4/7/9 swap-on: a near-full drain, five times an iteration. And CK spends more s_nop, not fewer: 25 against 12, over nop counts 1 to 11. MFMA hazard slots filled deliberately rather than covered by a wait. What that does not settle is which loads the tight waits are on. Constant scales do not change the time, which points at the B fragment rather than the scale loads. Separating a load from its use is the next thing to try; CK advances its scale copies a window ahead (a_scale_thread_copy plus MoveSrcSliceWindow). The ATT capture that would settle it directly did not work out, and the two failure modes are worth recording. The single-GPU probe is blocked by an HSA teardown assertion (ScratchCache not empty at shutdown) that aborts before rocprofv3 finalises, and exiting hard with os._exit skips the finalisation instead -- no trace either way. At 8 ranks the .att files come out healthy but all eight decodes paired to a null code object, three attempts at [4096, 7168] K=1024, which is the flakiness bench/att/README already documents. Docstring only, no functional change. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 34 ++++++++++++++++--- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index f13d8077a..b8e331b33 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -99,11 +99,35 @@ against 387, because the twelve live values cost more than the signed divide-by-128 expansion they remove. -The one thing the ISA does show, and which nothing above changed, is shape: -CK's loop body interleaves MFMA and VALU evenly -- median 4 VALU between MFMAs, -2 of 31 gaps empty -- where ours alternates MFMA clusters with 40-80 -instruction VALU stretches, 8 of 31 gaps empty. Whatever produces that sits -upstream of the group barriers. +Reading the two loop bodies against each other says the same thing from the +other side. Ours is *shorter* than CK's -- 312 instructions swap-off and 272 +swap-on against 322 -- with identical MFMA (32) and LDS (20) counts and fewer +VALU. Fewer instructions, 29% slower, so the gap is cycles per instruction. + +Two things in the ISA fit that: + +* **We drain the memory queue far harder.** CK's waits are + ``s_waitcnt vmcnt(15)``, ``vmcnt(16)``, ``vmcnt(10)`` -- it keeps ten to + sixteen loads in flight and waits only for the oldest. Ours are five + ``s_waitcnt vmcnt(1)`` per loop body swap-off (4/7/9 swap-on), i.e. a near + full drain, five times per iteration. +* **CK uses more s_nop, not fewer.** 25 against our 12, spread over nop counts + 1 to 11 -- MFMA hazard slots filled deliberately rather than covered by a + wait. + +Not settled: feeding constant scales instead of loading them does not change +the time, which argues the tight waits are on the B fragment rather than on the +scale loads. Separating a load from its use -- CK advances its scale copies a +window ahead, ``a_scale_thread_copy`` plus ``MoveSrcSliceWindow`` -- is the +obvious thing to try next. + +An ATT capture would settle it directly and did not work out. The single-GPU +probe route is blocked by an HSA teardown assertion (``ScratchCache not empty +at shutdown``) that kills the process before rocprofv3 finalises, and exiting +hard with ``os._exit`` skips the finalisation instead. At 8 ranks the trace +files come out healthy but every decode paired to a null code object, over +three attempts at ``[4096, 7168] K=1024`` -- the flakiness ``bench/att/README`` +already documents, where the only lever is to run it again. Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's From 72b2b8cd958754d8da3a7279d7ee2775a20d47d6 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 12:45:20 +0000 Subject: [PATCH 41/79] bench(cco): the tight vmcnt waits are the scale loads, and they cost nothing Following the two leads from the loop-body read, both to a stop. Which loads the near-full drains are on is now settled: every s_waitcnt vmcnt(0..1) is followed straight by the promote's v_pk_mul / v_pk_fma, and every wait sitting next to an MFMA is vmcnt(14..17). So the tight waits are the scale loads, and the B fragment already runs in the same range CK does -- the part of the pipeline that looked like the problem is not one. And the scale waits are not stalls. Issuing the loads a full M-row of MFMAs ahead of their use leaves them at 0/0/0/1/2, i.e. the compiler sinks them back to the point of use, and the clock does not move: 389.9 -> 386.4us swap-off, 386.7 -> 386.5 swap-on. Removing the loads outright does not move it either. The scale arrays are small and L1-resident, so the wait is satisfied as soon as it is reached. A static "tight wait" is not a stall, and this one was a misleading signal. That closes the list this investigation could reach without a working stall profile: promote arithmetic, register spill, store width, VALU scheduling groups, MFMA batch size, the scale loads themselves, address-arithmetic hoisting, and now load-to-use distance -- none of them move 388us. The loop body is shorter than CK's with the same MFMA and LDS counts, so the 29% is cycles the ISA does not account for, and ATT does not decode for this kernel today. Docstring only, no functional change. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .../cco/flydsl/gemm_ar/kernels_preshuffle4w.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index b8e331b33..b3093d32c 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -115,11 +115,15 @@ 1 to 11 -- MFMA hazard slots filled deliberately rather than covered by a wait. -Not settled: feeding constant scales instead of loading them does not change -the time, which argues the tight waits are on the B fragment rather than on the -scale loads. Separating a load from its use -- CK advances its scale copies a -window ahead, ``a_scale_thread_copy`` plus ``MoveSrcSliceWindow`` -- is the -obvious thing to try next. +Both were followed up, and both are dead ends. The tight waits are the scale +loads, not the B fragment: every ``vmcnt(0..1)`` is followed straight by the +promote's ``v_pk_mul``/``v_pk_fma``, while every wait next to an MFMA sits at +``vmcnt(14..17)`` -- so our B pipeline already runs where CK's does. And those +scale waits cost nothing. Issuing the loads a full M-row of MFMAs ahead of +their use does not relax them (they stay at 0/0/0/1/2) and does not move the +clock; neither does removing them outright. The data is small and L1-resident, +so the wait is satisfied as soon as it is reached. A static "tight wait" is not +a stall. An ATT capture would settle it directly and did not work out. The single-GPU probe route is blocked by an HSA teardown assertion (``ScratchCache not empty From 1dead3fce2ef9465ed0636497547aeaae41c1951 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 12:56:39 +0000 Subject: [PATCH 42/79] bench(cco): ATT decodes with --eager, and it puts the gap in VALU and the store The decoder was not flaky here, it was the hipGraph. bench_gemm_ar times through a graph by default and a graph-replayed dispatch pairs to a null code object every time -- 0 for 4 with graph timing, 1 for 1 with --eager, same shape and kernel. The code object is dumped either way (llvm-nm finds the kernel's symbols in one of the att_gfx950_code_object_id_*.out), so the failure was in the pairing. Recorded in bench/att/README, along with the other dead end: a single-process capture of a FlyDSL kernel cannot finish, because normal exit trips an HSA teardown assertion downstream of FlyDSL's leaked signals and os._exit skips rocprofv3's finalisation. Against CK at [4096, 7168] K=1024, as a share of each kernel's own attributed latency: this / stalled CK / stalled VALU 29.3% 46% 18.4% 31% global load 24.7% 89% 18.7% 90% s_waitcnt 19.4% 100% 20.0% 100% MFMA 7.9% 70% 3.8% 55% C store 5.5% 79% 0.3% 64% LDS 3.3% 26% 15.5% 75% s_nop 1.7% 100% 9.9% 100% Both run about three quarters stalled (72% against 75%), so the question is where, and the answer is not memory: global load and s_waitcnt are within a few points of each other, which retires the vmcnt lead for good. Ours is in VALU -- 60% more of the latency at half again the stall rate -- and in the MFMAs feeding it, 70% stalled against 55%. That is the promote's dependency, a tile's MFMA result read by the FMA that follows it, and it survived the batching sweep that was meant to break it. The C store is the other outlier at 5.5% against 0.3%. CK's profile is the mirror image: LDS at 15.5% and 75% stalled against our 3.3%, plus 9.9% s_nop. Its time goes into the A pipeline and into deliberately filled hazard slots, not into waiting on arithmetic. Docstring and notes only. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 37 +++++++++++++++---- 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index b3093d32c..796a922fe 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -125,13 +125,36 @@ so the wait is satisfied as soon as it is reached. A static "tight wait" is not a stall. -An ATT capture would settle it directly and did not work out. The single-GPU -probe route is blocked by an HSA teardown assertion (``ScratchCache not empty -at shutdown``) that kills the process before rocprofv3 finalises, and exiting -hard with ``os._exit`` skips the finalisation instead. At 8 ranks the trace -files come out healthy but every decode paired to a null code object, over -three attempts at ``[4096, 7168] K=1024`` -- the flakiness ``bench/att/README`` -already documents, where the only lever is to run it again. +An ATT capture does decode, once ``--eager`` is passed: a hipGraph-replayed +dispatch pairs to a null code object every time, which is what had made this +look flaky (0 for 4 with graph timing, 1 for 1 without). See +``bench/att/README``. What it shows, against CK on the same +``[4096, 7168] K=1024`` shape, as a share of each kernel's own attributed +latency: + +============= =============== =============== +category this / stalled CK / stalled +============= =============== =============== +VALU 29.3% 46% 18.4% 31% +global load 24.7% 89% 18.7% 90% +s_waitcnt 19.4% 100% 20.0% 100% +MFMA 7.9% 70% 3.8% 55% +C store 5.5% 79% 0.3% 64% +LDS 3.3% 26% 15.5% 75% +s_nop 1.7% 100% 9.9% 100% +============= =============== =============== + +Both run about three quarters stalled overall (72% against 75%), so the +difference is where. Ours is in VALU -- 60% more of the latency and half again +the stall rate -- and in the MFMAs feeding it, 70% stalled against 55%. That is +the promote's dependency: a tile's MFMA result is read by the FMA that follows +it. Which is what the batching sweep was meant to break and did not, so the +chain survives reordering. The C store is the other outlier, 5.5% against 0.3%. + +CK's profile is the mirror image: it spends its time in LDS (15.5% at 75% +stalled, against our 3.3%) and in s_nop (9.9%), i.e. in the A pipeline and in +deliberately filled hazard slots, not waiting on arithmetic. + Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's From 9f9906379e87f3202a8d3749f2c03bcb81aa2bb0 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 13:03:50 +0000 Subject: [PATCH 43/79] bench(cco): the C store outlier resolves, and warns about reading ATT shares It is 64 buffer_store_short at 36504 cycles against CK's 8 buffer_store_dwordx4 at 5064 -- the same 128 bytes a lane, eight times the instructions, seven times the latency. CK earns the wide store with the CShuffle its kernel name advertises: the accumulator goes through LDS to be reassembled into eight contiguous elements a lane (CShuffleBlockTransferScalarPerVector = 8), where we store straight out of registers at the MFMA's native four-element granularity. But swap_ab already cuts ours to 16 buffer_store_dwordx2, four times fewer, and that measured at parity: 387.4us against 387.5. A 5.5% share of attributed latency is not 5.5% of the clock. The epilogue runs once per block and with two waves a SIMD the other wave's MFMAs cover it. The same caveat applies to the VALU share, which is why this is worth writing down rather than just noting the store is fixed. (The swap build's own trace would have shown this directly and would not decode -- eight consecutive pairing failures against one first-try success for the default build, so the flakiness is build-specific, not just luck.) What the trace leaves is the MFMA-to-VALU dependency: our MFMAs 70% stalled against CK's 55%, feeding a VALU 46% stalled against 31%. Docstring only. Co-Authored-By: Claude Opus 5 (1M context) --- .../cco/flydsl/gemm_ar/kernels_preshuffle4w.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index 796a922fe..7380a7911 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -155,6 +155,23 @@ stalled, against our 3.3%) and in s_nop (9.9%), i.e. in the A pipeline and in deliberately filled hazard slots, not waiting on arithmetic. +The C store outlier resolves, and it is a caution about reading these shares. +It is 64 ``buffer_store_short`` at 36504 cycles against CK's 8 +``buffer_store_dwordx4`` at 5064 -- the same 128 bytes a lane, eight times the +instructions, seven times the latency. CK earns the wide store with the +CShuffle its kernel name advertises: the accumulator goes through LDS to be +reassembled into eight contiguous elements a lane +(``CShuffleBlockTransferScalarPerVector = 8``), where we store straight out of +registers at the MFMA's native four-element granularity. ``swap_ab`` already +cuts ours to 16 ``buffer_store_dwordx2`` -- and buys nothing, 387.4us against +387.5. **A 5.5% share of attributed latency is not 5.5% of the clock**: the +epilogue runs once per block, and with two waves a SIMD the other wave's MFMAs +cover it. The same caveat applies to the VALU number above. + +So the trace has retired the store as well, and what it leaves is the +MFMA-to-VALU dependency -- our MFMAs 70% stalled against CK's 55%, feeding a +VALU that is 46% stalled against 31%. + Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's From 54fef9600cb946bc2f06de061b154a4576e6c95b Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 13:48:57 +0000 Subject: [PATCH 44/79] bench(cco): the MFMA/VALU clumping is real, has a cause, and is not the cost The trace shows CK alternating one MFMA with four v_fmac while ours puts MFMAs in clumps followed by 40-80 instruction VALU stretches. That difference is in the source. CK promotes each tile where it is issued and pins it with an empty asm volatile("" : "+v"(value)) on the accumulator, so nothing can be moved across; mma_promote batches num_acc_n MFMAs and promotes them afterwards with nothing pinning the order, so LLVM sinks every promote below every MFMA. Reproducing CK's shape does not pay. One MFMA per promote plus rocdl.sched_barrier(0) after each takes the longest VALU run from 81 to 11 and the median gap from 4 to 6, much closer to CK's MVVVV, and the kernel goes 385.6 -> 393.7us. Slightly worse. Reverted. A K sweep is more useful than any of this. At K=2048 and K=8192 we measure 386.2 and 1390.1 against CK's 298.4 and 995.5, so our fixed cost extrapolates *below* CK's -- 51.6us against 66.0, CK paying for its CShuffle epilogue -- and the entire gap is per-iteration: 20.9us a K-block against 14.5, 44% more, on a loop body with fewer instructions than CK's. Occupancy cannot absorb it: --waves-per-eu 3 and 4 spill outright (2729 and 4719us), so both kernels run at two waves a SIMD. Docstring only, no functional change. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index 7380a7911..ee6f7dc71 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -168,9 +168,27 @@ epilogue runs once per block, and with two waves a SIMD the other wave's MFMAs cover it. The same caveat applies to the VALU number above. -So the trace has retired the store as well, and what it leaves is the -MFMA-to-VALU dependency -- our MFMAs 70% stalled against CK's 55%, feeding a -VALU that is 46% stalled against 31%. +So the trace has retired the store as well, and the interleaving with it. + +The shapes really do differ, and for a reason in the source. CK promotes each +tile where it is issued -- one MFMA then its four ``v_fmac`` -- and pins that +with an empty ``asm volatile("" : "+v"(value))`` on the accumulator, so nothing +can be re-clumped. ``mma_promote`` batches ``num_acc_n`` MFMAs and promotes +them afterwards, with nothing pinning the order, so LLVM sinks every promote +below every MFMA: MFMA in clumps, then a 40-80 instruction VALU stretch. + +Reproducing CK's shape does not pay. One MFMA per promote plus +``rocdl.sched_barrier(0)`` after each takes the longest VALU run from 81 down +to 11 and the median gap from 4 to 6 -- much closer to CK's ``MVVVV`` -- and +the kernel goes 385.6 -> 393.7us, slightly *worse*. The clumping is not what +costs the 29%. + +What a K sweep does say is where to look. Extrapolating 386.2us at K=2048 and +1390.1 at K=8192 against CK's 298.4 and 995.5: our fixed cost is *lower*, 51.6 +against 66.0 (CK pays for the CShuffle), and every bit of the gap is +per-iteration -- 20.9us a K-block against 14.5, 44% more, on a loop body with +fewer instructions. Occupancy cannot absorb it either: --waves-per-eu 3 and 4 +spill outright (2729 and 4719us), so both kernels sit at two waves a SIMD. Two structural notes that came out of the port. B never touches LDS here -- From cc696762e74655c2c9a7e099f7a3f72d547e1d53 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Thu, 10 Sep 2026 14:07:34 +0000 Subject: [PATCH 45/79] bench(cco): hardware counters say the work is identical and the wait is not rocprofv3 --pmc, same shape, against CK: SQ_INSTS_MFMA 7,340,032 7,340,032 1.00 SQ_INSTS_VALU 55,078,912 55,881,728 0.99 SQ_VALU_MFMA_BUSY_CYCLES 234,881,024 234,881,024 1.00 MemUnitStalled 0.11% 0.05% - FetchSize 483.6 MB 520.5 MB 0.93 SQ_WAIT_INST_LDS 2,442,413 33,294,145 0.07 SQ_WAIT_ANY 156,021,602 120,383,045 1.30 SQ_WAVE_CYCLES 386,929,422 333,626,547 1.16 Identical MFMA instruction count, identical MFMA-pipe busy cycles, VALU within 1%. Memory is not the bottleneck for either -- MemUnitStalled is ~0 and we fetch less than CK does -- and it is not LDS, where CK waits 13x more than we do. The whole difference is 36M more cycles of SQ_WAIT_ANY and an MFMA pipe idle 53M cycles longer. SQ_VALU_MFMA_COEXEC_CYCLES looked like the lever and turned out to be a symptom. Ours is 50.8M against CK's 67.1M, but with swap_ab on we issue 44.8M VALU instructions against CK's 55.9M and overlap 57.6M cycles -- a ratio of 1.28 against CK's 1.20. Better overlap per instruction, fewer instructions to overlap, still slower. And it does not respond to anything: promote batch sizes of 1, 2 and 4, with and without rocdl.sched_barrier(0) pinning each group, all measure 50.8-50.9M. So the 44%-per-iteration gap is cycles spent waiting on something none of these counters name. Recorded rather than solved; this is where the investigation stops unless someone has a better instrument than PC-sampled ATT and the SQ counter set. Docstring only. tests/python/cco/test_gemm_ar.py: 14 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py index ee6f7dc71..f09383ca4 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py @@ -190,6 +190,39 @@ fewer instructions. Occupancy cannot absorb it either: --waves-per-eu 3 and 4 spill outright (2729 and 4719us), so both kernels sit at two waves a SIMD. +Hardware counters close the question of what the difference is *not* +(``rocprofv3 --pmc``, same shape, 8 ranks against CK's single process): + +========================= ============== ============== ====== +counter ours CK ratio +========================= ============== ============== ====== +SQ_INSTS_MFMA 7,340,032 7,340,032 1.00 +SQ_INSTS_VALU 55,078,912 55,881,728 0.99 +SQ_VALU_MFMA_BUSY_CYCLES 234,881,024 234,881,024 1.00 +MemUnitStalled 0.11% 0.05% - +FetchSize 483.6 MB 520.5 MB 0.93 +SQ_WAIT_INST_LDS 2,442,413 33,294,145 0.07 +SQ_WAIT_ANY 156,021,602 120,383,045 1.30 +SQ_WAVE_CYCLES 386,929,422 333,626,547 1.16 +========================= ============== ============== ====== + +Identical MFMA instructions, identical MFMA-pipe busy cycles, VALU within 1%. +Memory is not the bottleneck for either -- ``MemUnitStalled`` is ~0 and we +fetch *less* -- and it is not LDS, where we wait 13x less than CK does. What is +left is 36M more cycles of ``SQ_WAIT_ANY``, and an MFMA pipe idle 53M cycles +longer. + +``SQ_VALU_MFMA_COEXEC_CYCLES`` looked like the lever and is not. Ours is 50.8M +against CK's 67.1M, but that is a symptom: with ``swap_ab`` on we issue 44.8M +VALU instructions against CK's 55.9M and overlap 57.6M cycles, a ratio of 1.28 +against CK's 1.20 -- better overlap per instruction, fewer instructions to +overlap, still slower. And nothing moves it: promote batch sizes of 1, 2 and 4, +with and without ``rocdl.sched_barrier(0)`` pinning each group, all measure +50.8-50.9M. + +So the 44%-per-iteration gap is cycles the wave spends waiting on something +none of these counters name. Recorded rather than solved. + Two structural notes that came out of the port. B never touches LDS here -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's From e7225e6eb03f0f15c572e4fe4def6de94cc607ae Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 13:34:06 +0000 Subject: [PATCH 46/79] delete the preshuffle4w negative result, keeping its findings A port of CK's 4-wave B-out-of-LDS shape, carried to answer why our block-scale GEMM is 22% behind CK's at the same shape. It reached CK's instruction mix and never reached its speed, and it goes out because keeping it would mean vendoring ~1000 more lines of aiter (buffer_ops, vector, splitk_epilogue, xcd_remap_bx_by) to serve a kernel nothing calls. What it ruled out is kept, in python/mori/ops/gemm_ar/README.md, so the same ground is not walked twice: ten hypotheses falsified by measurement, SQ_INSTS_MFMA and SQ_VALU_MFMA_BUSY_CYCLES identical to CK's, and a K-sweep that puts the whole difference per-iteration rather than in the fixed cost. The traces are in cc696762, 54fef960, 9f990637 and 1dead3fc. Co-Authored-By: Claude Opus 5 (1M context) --- .../flydsl/gemm_ar/kernels_preshuffle4w.py | 1377 ----------------- 1 file changed, 1377 deletions(-) delete mode 100644 benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py diff --git a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py b/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py deleted file mode 100644 index f09383ca4..000000000 --- a/benchmark/cco/flydsl/gemm_ar/kernels_preshuffle4w.py +++ /dev/null @@ -1,1377 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Preshuffle GEMM (layout API): f16/bf16/fp8/int8, ping-pong scf.for loop with scheduler hints. - -Pinned from ``aiter.ops.flydsl.kernels.preshuffle_gemm`` and extended with -``quant="blockscale"`` -- the A 1x128 / B 128x128 fp32 quantisation DSV4 runs, -which upstream does not have (its fp8 scales are per-row x per-col, applied once -in the epilogue). - -Why this kernel and not ``kernels_fused``'s 8-wave one: it is already the shape -CK's fastest blockscale instance uses. 4 waves, 256 threads, and **B never -touches LDS** -- ``thr_g2r_B``/``frag_B_stages`` load it global->VGPR -double-buffered, the way CK's ``blockscale_b_preshuffle_v1`` pipeline does with -``b_thread_bufs``. Only A is staged, so LDS is ``tile_m * tile_k`` x 2 stages -(16 KB at 64x128) against the 8-wave kernel's 96 KB, which caps that one at one -block per CU. B is already preshuffled to ``N0 K0 KLane NLane KPack``, so every -lane pulls its own MFMA operand with one 128-bit load and there is nothing for -LDS to amortise. - -Measured at ``[16384, 7168] K=2048`` on MI355X, warm cache, tile 64x256x128 -- -CK's own winning shape: - -============================= ======= ======== -stage us vs CK -============================= ======= ======== -whole-fragment rescale (divide) 760.2 +153% -per-tile promotion 459.5 +53% - + explicit fma, scale hoisted 391.7 +30% - + MFMAs batched per M-row 385.9 +28% -this kernel, scale removed 303.1 +1% -CK 64x256 Intrawave v1 300.4 . -============================= ======= ======== - -So the GEMM itself is aligned: strip the scale and this lands on CK. What is -left is the scale, +82.6us against CK's ~0, and the loop body says why. Per -128 accumulator elements: - - CK 128 v_fmac_f32, 8 v_mul_f32 -> VALU/MFMA 5.81 - this 128 v_(pk_)fma, 128 v_(pk_)mul -> VALU/MFMA 5.53 - -Our instruction *count* is already under CK's; the shape is not. CK needs one -FMA per element because ``c_scale_thread`` is a single scalar for a whole MFMA -tile, where we need a multiply per element first: A's scale is per row -(ScaleBlockM=1) and an MFMA 16x16 lane holds four different rows, so four -different A scales. CK's four values must instead be four *columns* of one row, -which is what ``kernels_fused``'s ``--swap-ab`` arranges. Closing the last 28% -means adopting that output layout, not more arithmetic. - -``swap_ab`` exchanges the MFMA operands so a lane owns four consecutive N of -one row instead of four M. It does exactly what it is supposed to on both -counts CK benefits from -- A's scale collapses to one scalar per tile, so the -promote is a single FMA like ``c_scale_thread``, and the C store goes from 64 -``buffer_store_short`` to 16 ``buffer_store_dwordx2``: - -============================ ========= ========== ============ - VALU/MFMA C stores total (us) -============================ ========= ========== ============ -off 5.53 64 short 385.9 -on 4.47 16 dwordx2 506.9 -============================ ========= ========== ============ - -Both are done now, and neither moved the clock. Getting there took two fixes -that are worth recording, because each was found by attribution rather than by -guessing: - -* ``--waves-per-eu 2`` caps VGPRs at 256, and the swap crosses it. Compiling - its three parts separately: the scalar A scale *saves* 16 registers, the - operand exchange costs 4, and the hand-written C store costs 28. The store - was the whole problem -- and not its landing registers (one measures the same - as sixteen), nor ``frag_C_out``, nor the conversion order. Going through the - tiled copy against a transposed C partitioning instead costs nothing: - ``tiled_mma`` rebuilt with the waves raked along the first mode, partitioning - ``C^T[N, M]`` (strides ``(1, N)``), which is the same map with M and N - exchanged. -* That partition wants the fragment in ``(v, n_rep, m_rep)`` order where - ``frag_C`` is ``(v, m_rep, n_rep)``. Permuting in the epilogue is correct but - costs 8 VGPRs, because source and destination are live together. The slot - order is ours to choose, so ``mma_promote`` writes the swapped order from the - start and there is nothing to permute. - -Result: 248 VGPR, no spill, 16 ``buffer_store_dwordx2``, one FMA per element. -And 387.4us against the unswapped path's 387.5 -- **exactly parity**. - -So none of it was the bottleneck. Nor, it turns out, is the block scale: -replacing the promote's FMA with a plain add measures 387.4us against 385.4 -with it, and feeding constant scales instead of loading them measures 396.8 -against 395.6. **The scale is free in this kernel.** An earlier 303.1us -reading, taken the same way before the swap work, does not reproduce and should -not be relied on. - -Which moves the question: the gap to CK is the GEMM, not the quantisation. The -scheduling hypothesis this was chasing is dead too. FlyDSL's ``sched_*`` -helpers cover MFMA/DS/VMEM but not VALU, so nothing pins the promote -- but -adding VALU groups by hand (``sched_group_barrier(0x002, n, 0)``, LLVM's VALU -mask) is worth 0-2us at any group size from 2 to 8. So is batching more MFMAs -before promoting them: 1, 2, 4, 8 and 16 tiles all land within 4us of each -other. Hoisting the scale addresses' loop invariants out is actively worse, 409 -against 387, because the twelve live values cost more than the signed -divide-by-128 expansion they remove. - -Reading the two loop bodies against each other says the same thing from the -other side. Ours is *shorter* than CK's -- 312 instructions swap-off and 272 -swap-on against 322 -- with identical MFMA (32) and LDS (20) counts and fewer -VALU. Fewer instructions, 29% slower, so the gap is cycles per instruction. - -Two things in the ISA fit that: - -* **We drain the memory queue far harder.** CK's waits are - ``s_waitcnt vmcnt(15)``, ``vmcnt(16)``, ``vmcnt(10)`` -- it keeps ten to - sixteen loads in flight and waits only for the oldest. Ours are five - ``s_waitcnt vmcnt(1)`` per loop body swap-off (4/7/9 swap-on), i.e. a near - full drain, five times per iteration. -* **CK uses more s_nop, not fewer.** 25 against our 12, spread over nop counts - 1 to 11 -- MFMA hazard slots filled deliberately rather than covered by a - wait. - -Both were followed up, and both are dead ends. The tight waits are the scale -loads, not the B fragment: every ``vmcnt(0..1)`` is followed straight by the -promote's ``v_pk_mul``/``v_pk_fma``, while every wait next to an MFMA sits at -``vmcnt(14..17)`` -- so our B pipeline already runs where CK's does. And those -scale waits cost nothing. Issuing the loads a full M-row of MFMAs ahead of -their use does not relax them (they stay at 0/0/0/1/2) and does not move the -clock; neither does removing them outright. The data is small and L1-resident, -so the wait is satisfied as soon as it is reached. A static "tight wait" is not -a stall. - -An ATT capture does decode, once ``--eager`` is passed: a hipGraph-replayed -dispatch pairs to a null code object every time, which is what had made this -look flaky (0 for 4 with graph timing, 1 for 1 without). See -``bench/att/README``. What it shows, against CK on the same -``[4096, 7168] K=1024`` shape, as a share of each kernel's own attributed -latency: - -============= =============== =============== -category this / stalled CK / stalled -============= =============== =============== -VALU 29.3% 46% 18.4% 31% -global load 24.7% 89% 18.7% 90% -s_waitcnt 19.4% 100% 20.0% 100% -MFMA 7.9% 70% 3.8% 55% -C store 5.5% 79% 0.3% 64% -LDS 3.3% 26% 15.5% 75% -s_nop 1.7% 100% 9.9% 100% -============= =============== =============== - -Both run about three quarters stalled overall (72% against 75%), so the -difference is where. Ours is in VALU -- 60% more of the latency and half again -the stall rate -- and in the MFMAs feeding it, 70% stalled against 55%. That is -the promote's dependency: a tile's MFMA result is read by the FMA that follows -it. Which is what the batching sweep was meant to break and did not, so the -chain survives reordering. The C store is the other outlier, 5.5% against 0.3%. - -CK's profile is the mirror image: it spends its time in LDS (15.5% at 75% -stalled, against our 3.3%) and in s_nop (9.9%), i.e. in the A pipeline and in -deliberately filled hazard slots, not waiting on arithmetic. - -The C store outlier resolves, and it is a caution about reading these shares. -It is 64 ``buffer_store_short`` at 36504 cycles against CK's 8 -``buffer_store_dwordx4`` at 5064 -- the same 128 bytes a lane, eight times the -instructions, seven times the latency. CK earns the wide store with the -CShuffle its kernel name advertises: the accumulator goes through LDS to be -reassembled into eight contiguous elements a lane -(``CShuffleBlockTransferScalarPerVector = 8``), where we store straight out of -registers at the MFMA's native four-element granularity. ``swap_ab`` already -cuts ours to 16 ``buffer_store_dwordx2`` -- and buys nothing, 387.4us against -387.5. **A 5.5% share of attributed latency is not 5.5% of the clock**: the -epilogue runs once per block, and with two waves a SIMD the other wave's MFMAs -cover it. The same caveat applies to the VALU number above. - -So the trace has retired the store as well, and the interleaving with it. - -The shapes really do differ, and for a reason in the source. CK promotes each -tile where it is issued -- one MFMA then its four ``v_fmac`` -- and pins that -with an empty ``asm volatile("" : "+v"(value))`` on the accumulator, so nothing -can be re-clumped. ``mma_promote`` batches ``num_acc_n`` MFMAs and promotes -them afterwards, with nothing pinning the order, so LLVM sinks every promote -below every MFMA: MFMA in clumps, then a 40-80 instruction VALU stretch. - -Reproducing CK's shape does not pay. One MFMA per promote plus -``rocdl.sched_barrier(0)`` after each takes the longest VALU run from 81 down -to 11 and the median gap from 4 to 6 -- much closer to CK's ``MVVVV`` -- and -the kernel goes 385.6 -> 393.7us, slightly *worse*. The clumping is not what -costs the 29%. - -What a K sweep does say is where to look. Extrapolating 386.2us at K=2048 and -1390.1 at K=8192 against CK's 298.4 and 995.5: our fixed cost is *lower*, 51.6 -against 66.0 (CK pays for the CShuffle), and every bit of the gap is -per-iteration -- 20.9us a K-block against 14.5, 44% more, on a loop body with -fewer instructions. Occupancy cannot absorb it either: --waves-per-eu 3 and 4 -spill outright (2729 and 4719us), so both kernels sit at two waves a SIMD. - -Hardware counters close the question of what the difference is *not* -(``rocprofv3 --pmc``, same shape, 8 ranks against CK's single process): - -========================= ============== ============== ====== -counter ours CK ratio -========================= ============== ============== ====== -SQ_INSTS_MFMA 7,340,032 7,340,032 1.00 -SQ_INSTS_VALU 55,078,912 55,881,728 0.99 -SQ_VALU_MFMA_BUSY_CYCLES 234,881,024 234,881,024 1.00 -MemUnitStalled 0.11% 0.05% - -FetchSize 483.6 MB 520.5 MB 0.93 -SQ_WAIT_INST_LDS 2,442,413 33,294,145 0.07 -SQ_WAIT_ANY 156,021,602 120,383,045 1.30 -SQ_WAVE_CYCLES 386,929,422 333,626,547 1.16 -========================= ============== ============== ====== - -Identical MFMA instructions, identical MFMA-pipe busy cycles, VALU within 1%. -Memory is not the bottleneck for either -- ``MemUnitStalled`` is ~0 and we -fetch *less* -- and it is not LDS, where we wait 13x less than CK does. What is -left is 36M more cycles of ``SQ_WAIT_ANY``, and an MFMA pipe idle 53M cycles -longer. - -``SQ_VALU_MFMA_COEXEC_CYCLES`` looked like the lever and is not. Ours is 50.8M -against CK's 67.1M, but that is a symptom: with ``swap_ab`` on we issue 44.8M -VALU instructions against CK's 55.9M and overlap 57.6M cycles, a ratio of 1.28 -against CK's 1.20 -- better overlap per instruction, fewer instructions to -overlap, still slower. And nothing moves it: promote batch sizes of 1, 2 and 4, -with and without ``rocdl.sched_barrier(0)`` pinning each group, all measure -50.8-50.9M. - -So the 44%-per-iteration gap is cycles the wave spends waiting on something -none of these counters name. Recorded rather than solved. - - -Two structural notes that came out of the port. B never touches LDS here -- -``thr_g2r_B``/``frag_B_stages`` load it global->VGPR double-buffered, as CK's -``blockscale_b_preshuffle_v1`` does with ``b_thread_bufs`` -- and B is already -preshuffled to ``N0 K0 KLane NLane KPack``, the same layout CK's -``preShuffleBuffer`` emits, so every lane pulls its own MFMA operand with one -128-bit load. And driving MFMAs per accumulator tile beats one ``fx.gemm`` over -the whole fragment even before any scaling: 303.1us against 556.0 at this tile. - -For scale, the 8-wave kernel in ``kernels_fused`` is 224.0us unscaled and 370.5 -scaled at its own best tile. -""" - -import functools - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import const_expr, gpu, math, range_constexpr, rocdl -from flydsl.expr.typing import ( - BFloat16, - Float8E4M3FN, - Float8E4M3FNUZ, - Float16, - Float32, - Int8, - Int32, - T, -) -from flydsl.expr.typing import Vector as Vec -from flydsl.runtime.device import get_rocm_arch - -from aiter.ops.flydsl.kernels import buffer_ops, vector -from aiter.ops.flydsl.kernels.mfma_preshuffle_pipeline import xcd_remap_bx_by -from aiter.ops.flydsl.kernels.splitk_epilogue import ( - CPOL_COHERENT, - splitk_reduce_epilogue, -) - -# (dsrd_preload, dvmem_preload) per (tile_m, tile_n, tile_k). -_TILE_PRELOAD_TABLE = { - # ── tile_m = 16 ── - (16, 64, 256): (2, 2), - (16, 64, 512): (4, 4), - (16, 128, 256): (2, 2), - (16, 128, 512): (2, 2), - (16, 192, 256): (2, 2), - (16, 256, 256): (2, 2), - (16, 256, 512): (2, 2), - (16, 512, 256): (2, 2), - # ── tile_m = 32 ── - (32, 64, 128): (6, 6), - (32, 64, 256): (6, 6), - (32, 64, 512): (2, 2), - (32, 128, 128): (6, 6), - (32, 128, 256): (6, 6), - (32, 192, 128): (6, 6), - (32, 192, 256): (6, 6), - (32, 256, 128): (6, 6), - (32, 256, 256): (6, 6), - # ── tile_m = 48 ── - (48, 64, 128): (8, 8), - (48, 64, 256): (2, 2), - (48, 128, 256): (6, 6), - (48, 192, 256): (6, 6), - (48, 256, 256): (6, 6), - # ── tile_m = 64 ── - (64, 64, 128): (4, 4), - (64, 64, 256): (4, 4), - (64, 128, 128): (8, 8), - (64, 128, 256): (8, 8), - (64, 192, 128): (8, 8), - (64, 192, 256): (8, 8), - (64, 256, 64): (8, 8), - (64, 256, 128): (8, 8), - (64, 256, 256): (8, 8), - # ── tile_m = 80 ── - (80, 64, 256): (4, 4), - (80, 128, 256): (8, 8), - (80, 192, 256): (8, 8), - (80, 256, 256): (8, 8), - # ── tile_m = 96 ── - (96, 64, 128): (6, 6), - (96, 64, 256): (6, 6), - (96, 128, 128): (8, 8), - (96, 128, 256): (6, 6), - (96, 192, 128): (8, 8), - (96, 192, 256): (8, 8), - (96, 256, 128): (8, 8), - (96, 256, 256): (8, 8), - # ── tile_m = 112 ── - (112, 64, 256): (8, 8), - (112, 128, 256): (4, 4), - (112, 192, 256): (8, 8), - (112, 256, 256): (8, 8), - # ── tile_m = 128 ── - (128, 64, 128): (6, 6), - (128, 64, 256): (8, 8), - (128, 128, 64): (4, 4), - (128, 128, 128): (8, 8), - (128, 128, 256): (4, 4), - (128, 192, 128): (8, 8), - (128, 192, 256): (8, 8), - (128, 256, 128): (6, 6), - (128, 256, 256): (4, 4), - # ── tile_m = 160 ── - (160, 192, 128): (8, 8), - # ── tile_m = 192 ── - (192, 64, 128): (6, 6), - (192, 128, 128): (6, 6), - # ── tile_m = 224 ── - (224, 64, 128): (4, 4), - (224, 128, 128): (6, 6), - (224, 192, 128): (6, 6), - # ── tile_m = 256 ── - (256, 64, 128): (4, 4), - (256, 128, 128): (6, 6), - (256, 192, 128): (6, 6), - (256, 256, 128): (4, 4), -} - -_TILE_PRELOAD_DEFAULT = (0, 0) - - -def _get_preload(tile_m, tile_n, tile_k): - """Look up (dsrd_preload, dvmem_preload) from the tile table.""" - return _TILE_PRELOAD_TABLE.get( - (int(tile_m), int(tile_n), int(tile_k)), _TILE_PRELOAD_DEFAULT - ) - - -# Row bound of the A/C layout views; gemm_kernels mirrors it to guard on the host. -PRESHUFFLE_M_MAX = 65536 - - -@functools.lru_cache(maxsize=1024) -def compile_preshuffle_gemm( - *, - N: int, - K: int, - tile_m: int, - tile_n: int, - tile_k: int, - in_dtype: str = "fp8", - out_dtype: str = "bf16", - quant: str = "ptpc", # "ptpc" (upstream) or "blockscale" (A 1x128, B 128x128) - swap_ab: bool = False, # blockscale only; see the module docstring - epilogue: str = "none", # "none", "bias", "bias_relu", "bias_silu", "bias_gelu" - waves_per_eu: int | None = None, - enable_scheduler: bool = True, - use_async_copy: bool = False, - xcd_swizzle: int = 0, - lds_stage: int = 2, - split_k: int = 1, -): - """Compile preshuffle GEMM (fp8/int8/fp16/bf16). - Signature: fn(C, out, semaphore, A, B, scale_a, scale_b, bias, M, N, stream). - For split_k > 1, C is an fp32 partial workspace and the last arriving split - performs the reduction and final conversion inside the same kernel launch. - """ - if in_dtype not in ("fp8", "int8", "fp16", "bf16"): - raise ValueError(f"in_dtype must be fp8/int8/fp16/bf16, got {in_dtype!r}") - if split_k < 1 or K % split_k != 0: - raise ValueError(f"split_k must divide K; got split_k={split_k}, K={K}") - split_k_extent = K // split_k - if tile_k <= 0 or split_k_extent % tile_k != 0: - raise ValueError( - "tile_k must be a positive divisor of K/split_k; " - f"got tile_k={tile_k}, K={K}, split_k={split_k}" - ) - if epilogue not in ("none", "bias", "bias_relu", "bias_silu", "bias_gelu"): - raise ValueError( - f"epilogue must be none/bias/bias_relu/bias_silu/bias_gelu, got {epilogue!r}" - ) - if lds_stage not in (1, 2): - raise ValueError(f"lds_stage must be 1 or 2, got {lds_stage}") - _has_epilogue = epilogue != "none" - _has_bias = epilogue in ("bias", "bias_relu", "bias_silu", "bias_gelu") - _has_relu = epilogue == "bias_relu" - _has_silu = epilogue == "bias_silu" - _has_gelu = epilogue == "bias_gelu" - if split_k > 1 and _has_epilogue: - raise ValueError("split_k > 1 does not support fused bias or activation") - - if quant not in ("ptpc", "blockscale"): - raise ValueError(f"quant must be ptpc or blockscale, got {quant!r}") - is_blockscale = quant == "blockscale" - swap_ab = swap_ab and is_blockscale - if is_blockscale: - # tile_k == 128 makes one main-loop tile exactly one K-scale block, so - # the rescale lands on a loop boundary that already exists. - if in_dtype != "fp8": - raise ValueError("blockscale is fp8 only") - if tile_k % 128: - raise ValueError(f"blockscale needs tile_k a multiple of 128, got {tile_k}") - if K % 128 or N % 128: - raise ValueError(f"blockscale needs K and N multiples of 128, got K={K} N={N}") - if split_k != 1 or _has_epilogue: - raise ValueError("blockscale does not support split_k or a fused epilogue") - # Only the 2-stage pipeline threads the tile index down to mma_kloop, - # which is what the rescale is indexed by. - if lds_stage != 2: - raise ValueError(f"blockscale needs lds_stage=2, got {lds_stage}") - - is_fp8 = in_dtype == "fp8" - is_int8 = in_dtype == "int8" - is_f16 = in_dtype == "fp16" - is_bf16 = in_dtype == "bf16" - is_f16_or_bf16 = is_f16 or is_bf16 - is_8bit = is_fp8 or is_int8 - elem_bytes = 1 if is_8bit else 2 - - # The async gmem->LDS DMA (buffer_load_lds 128b) only lowers for 8-bit inputs. - if use_async_copy and not is_8bit: - raise ValueError("use_async_copy is only supported for 8-bit inputs (fp8/int8)") - - gpu_arch = get_rocm_arch() - is_gfx942 = str(gpu_arch).startswith("gfx942") - is_gfx950 = str(gpu_arch).startswith("gfx950") - use_mfma_scale_128 = is_fp8 and is_gfx950 and (tile_k % 128 == 0) - use_mfma_k32 = is_f16_or_bf16 and is_gfx950 - - if is_f16_or_bf16: - layout_elem = Float16 if is_f16 else BFloat16 - elif is_int8: - layout_elem = Int8 - else: - layout_elem = Float8E4M3FN if is_gfx950 else Float8E4M3FNUZ - final_out_elem_cls = BFloat16 if out_dtype == "bf16" else Float16 - out_elem_cls = Float32 if split_k > 1 else final_out_elem_cls - out_elem_bytes = 4 if split_k > 1 else 2 - - # Tile geometry (tile_K_perm = K-elements grouped per MMA k-step) - tile_K_perm = 128 if use_mfma_scale_128 else (64 if is_8bit else 32) - k_iters = tile_k // tile_K_perm - num_tiles = split_k_extent // tile_k - #: K-scale blocks; with the tile_k=128 blockscale requires, one per tile. - kb_count = K // 128 if is_blockscale else 0 - m_repeat = tile_m // 16 - num_waves = 4 - n_per_wave = tile_n // num_waves - num_acc_n = n_per_wave // 16 - acc_size = m_repeat * num_acc_n * 4 - - total_threads = 256 - a_load_bytes = 16 - bytes_per_thread_a = (tile_m * tile_k * elem_bytes) // total_threads - num_a_loads = bytes_per_thread_a // a_load_bytes - num_b_loads = (tile_n * tile_k * elem_bytes) // total_threads // 16 - num_ds_load = (tile_m * tile_k * elem_bytes) // 64 // 16 # A LDS reads per wave - num_gmem_loads = num_a_loads + num_b_loads - if is_8bit and is_gfx950: - dsrd_preload, dvmem_preload = _get_preload(tile_m, tile_n, tile_k) - else: - dsrd_preload, dvmem_preload = (0, 0) - - a_lds_elems = tile_m * tile_k - - @fx.struct - class SharedStorage: - a0: fx.Array[layout_elem, a_lds_elems, 16] - if lds_stage == 2: - a1: fx.Array[layout_elem, a_lds_elems, 16] - if split_k > 1: - split_flag: fx.Array[Int32, 1, 4] - - # ── Kernel ──────────────────────────────────────────────────────── - @flyc.kernel - def kernel_gemm( - arg_c: fx.Tensor, - arg_out: fx.Tensor, - arg_semaphore: fx.Tensor, - arg_a: fx.Tensor, - arg_b: fx.Tensor, - arg_scale_a: fx.Tensor, - arg_scale_b: fx.Tensor, - arg_bias: fx.Tensor, - i32_m: fx.Int32, - i32_n: fx.Int32, - tiled_mma_arg: fx.TiledMma, - tiled_copy_g2s: fx.TiledCopy, - ): - tid = fx.thread_idx.x - bid_x, bid_y, bid_z = fx.block_idx - k_off = fx.Int32(bid_z) * num_tiles - - if const_expr(xcd_swizzle > 0): - _bx, _by = xcd_remap_bx_by( - gpu.block_id("x"), - gpu.block_id("y"), - fx.Index(i32_m), - tile_m=tile_m, - tile_n=tile_n, - N=N, - xcd_swizzle=xcd_swizzle, - ) - bid_x, bid_y = Int32(_bx), Int32(_by) - - mma_atom = None - if const_expr(use_mfma_scale_128): - _scale_atom = fx.make_mma_atom( - fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, layout_elem) - ) - mma_atom = _scale_atom - tiled_mma = fx.make_tiled_mma( - _scale_atom, - fx.make_layout((1, 4, 1), (0, 1, 0)), - fx.make_tile(None, None, fx.make_layout((32, 4), (1, 32))), - ) - else: - tiled_mma = tiled_mma_arg - - # Bound A (read) and C (store) to the actual M extent so blocks covering - # rows past M (ragged M) drop their OOB loads/stores at the descriptor - # instead of faulting / writing past the allocation. B and scale_b/bias - # are per-N (exact multiple) and stay max_size; scale_a is per-row (M) and - # is bounded the same way below (see its epilogue load). - gA = fx.rocdl.make_buffer_tensor( - arg_a, - max_size=False, - num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes), - ) - gB = fx.rocdl.make_buffer_tensor(arg_b) - c_tensor = arg_c - if const_expr(split_k > 1): - c_split_offset = fx.Int64(bid_z) * fx.Int64(i32_m) * fx.Int64(N) - c_tensor = fx.Tensor( - fx.make_view( - fx.add_offset(fx.get_iter(arg_c), c_split_offset), - fx.make_layout((PRESHUFFLE_M_MAX, N), (N, 1)), - ) - ) - gC = fx.rocdl.make_buffer_tensor( - c_tensor, - max_size=False, - num_records_bytes=fx.Int64(i32_m) * fx.Int64(N) * fx.Int64(out_elem_bytes), - ) - - tA = fx.flat_divide(gA, fx.make_tile(tile_m, tile_k))[None, None, bid_x, None] - tB = fx.flat_divide(gB, fx.make_tile(tile_n, tile_k))[None, None, bid_y, None] - tC = fx.flat_divide(gC, fx.make_tile(tile_m, tile_n))[None, None, bid_x, bid_y] - - buf_copy = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), layout_elem) - uni_copy = fx.make_copy_atom(fx.UniversalCopy128b(), layout_elem) - - # Per-thread slices - thr_mma = tiled_mma.thr_slice(tid) - thr_g2s = tiled_copy_g2s.get_slice(tid) - thr_s2r = fx.make_tiled_copy_A(buf_copy, tiled_mma).get_slice(tid) - thr_g2r_B = fx.make_tiled_copy_B(buf_copy, tiled_mma).get_slice(tid) - - lds = fx.SharedAllocator().allocate(SharedStorage).peek() - if const_expr(is_8bit): - k_blocks16 = (tile_k * elem_bytes) // 16 - if k_blocks16 <= 0 or (k_blocks16 & (k_blocks16 - 1)) != 0: - raise ValueError( - f"Unsupported tile_k for 8-bit LDS swizzle: tile_k={tile_k}, elem_bytes={elem_bytes} (k_blocks16={k_blocks16}); " - "expected tile_k*elem_bytes to be a positive multiple of 16 with (tile_k*elem_bytes/16) a power of two." - ) - swz_bits = k_blocks16.bit_length() - 1 # log2 - swz = fx.SwizzleType.get(swz_bits, 4, swz_bits) - else: - swz = fx.SwizzleType.get(3, 3, 3) - - def _make_sA(arr): - return fx.make_view( - arr.ptr, - fx.make_composed_layout( - fx.static(swz), - fx.make_ordered_layout((tile_m, tile_k), (1, 0)), - ), - ) - - if const_expr(lds_stage == 1): - sA_stages = [_make_sA(lds.a0)] - else: - sA_stages = [_make_sA(lds.a0), _make_sA(lds.a1)] - - # Partitions - pA_g = thr_g2s.partition_S(tA) - pA_s_stages = [thr_g2s.partition_D(s) for s in sA_stages] - pA_s2r_stages = [thr_s2r.partition_S(s) for s in sA_stages] - pB_g = thr_g2r_B.partition_S(tB) - - # Fragments — 2 separate B fragments (split double buffer for VGPR lifetime) - frag_copy_A = fx.make_fragment_like(pA_s_stages[0][None, None, None]) - frag_A = thr_mma.make_fragment_A(sA_stages[0]) - frag_B_single_layout = thr_mma.partition_B(tB).layout(None, None, None, 0) - frag_B_stages = [ - fx.make_fragment_like(frag_B_single_layout, layout_elem.ir_type) - for _ in range(2) - ] - frag_C = thr_mma.make_fragment_C(tC) - frag_A_retile = thr_s2r.retile(frag_A) - frag_B_retile_stages = [thr_g2r_B.retile(b) for b in frag_B_stages] - out_cpol = CPOL_COHERENT if split_k > 1 else 0 - copy_op = fx.rocdl.BufferCopy32b if split_k > 1 else fx.rocdl.BufferCopy16b - buf_copy_out = fx.make_copy_atom(copy_op(out_cpol), out_elem_cls) - frag_C_out = None - if const_expr(swap_ab): - # The swapped accumulator maps to C with M and N exchanged, so the - # ordinary partitioning applied to C^T lands on the right addresses. - # C is row-major [M, N], so C^T is [N, M] with strides (1, N) -- a - # tile's four values are then contiguous, which is what lets the - # store be one 64-bit instead of four 16-bit ones. Going through - # the tiled copy rather than by hand is the point: a hand-written - # store needs landing registers and costs the 8 VGPRs that push - # this over --waves-per-eu 2 and make the main loop spill. - tiled_mma_T = fx.make_tiled_mma( - mma_atom, - fx.make_layout((4, 1, 1), (1, 0, 0)), - fx.make_tile(None, None, fx.make_layout((32, 4), (1, 32))), - ) - c_tensor_T = fx.Tensor( - fx.make_view( - fx.get_iter(arg_c), - fx.make_layout((N, PRESHUFFLE_M_MAX), (1, N)), - ) - ) - gC_T = fx.rocdl.make_buffer_tensor( - c_tensor_T, - max_size=False, - num_records_bytes=fx.Int64(i32_m) - * fx.Int64(N) - * fx.Int64(out_elem_bytes), - ) - tC_T = fx.flat_divide(gC_T, fx.make_tile(tile_n, tile_m))[ - None, None, bid_y, bid_x - ] - copy_out_c = fx.make_copy_atom( - fx.rocdl.BufferCopy64b(out_cpol), out_elem_cls - ) - # The fragment has to carry the transposed mode order too: the - # T-partition expects (v, n_rep, m_rep) where frag_C is - # (v, m_rep, n_rep), so the fill below permutes the slots. - frag_C_out = fx.make_fragment_like( - tiled_mma_T.thr_slice(tid).make_fragment_C(tC_T), - out_elem_cls.ir_type, - ) - thr_r2g_C = fx.make_tiled_copy_C(copy_out_c, tiled_mma_T).get_slice(tid) - pC_g = thr_r2g_C.partition_S(tC_T) - frag_C_retile = thr_r2g_C.retile(frag_C_out) - else: - frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) - copy_out_c = buf_copy_out - thr_r2g_C = fx.make_tiled_copy_C(buf_copy_out, tiled_mma).get_slice(tid) - pC_g = thr_r2g_C.partition_S(tC) - frag_C_retile = thr_r2g_C.retile(frag_C_out) - - # ── Async gmem->LDS DMA (buffer_load_lds) for the A tile ── - if const_expr(use_async_copy): - dma_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) - # Bound to the real M extent (like the sync gA) so ragged-M blocks DMA-read - # OOB rows as 0 instead of faulting past the allocation. - gA_flat = fx.rocdl.make_buffer_tensor( - fx.Tensor( - fx.make_view( - fx.get_iter(arg_a), fx.make_layout(PRESHUFFLE_M_MAX * K, 1) - ) - ), - max_size=False, - num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes), - ) - gA_div = fx.logical_divide(gA_flat, fx.make_layout(1, 1)) - sA_i8_ptr = [fx.recast_iter(Int8, lds.a0.ptr)] - if const_expr(lds_stage == 2): - sA_i8_ptr.append(fx.recast_iter(Int8, lds.a1.ptr)) - bx_m = bid_x * tile_m - wave_id = tid // 64 - step_bytes = total_threads * a_load_bytes - wave_stride_bytes = 64 * a_load_bytes - k_blocks16_dma = (tile_k * elem_bytes) // 16 - elems_per_16b = 16 // elem_bytes - - def dma_a_to_lds(k_tile_val, stage): - wave_off = rocdl.readfirstlane( - fx.Int32.ir_type, wave_id * wave_stride_bytes - ) - lds_ptr = fx.add_offset(sA_i8_ptr[stage], wave_off) - base_k = (k_off + k_tile_val) * tile_k - for i in range_constexpr(num_a_loads): - if const_expr(i > 0): - lds_ptr = fx.add_offset(lds_ptr, step_bytes) - pos_bytes = i * total_threads * a_load_bytes + tid * a_load_bytes - elem_idx = pos_bytes // elem_bytes - m = elem_idx // tile_k - k = elem_idx % tile_k - k_swz = k ^ ((m % k_blocks16_dma) * elems_per_16b) - gmem_byte = ((bx_m + m) * K + base_k + k_swz) * elem_bytes - dst = fx.make_view(lds_ptr, fx.make_layout(1, 1)) - src = fx.slice(gA_div, (None, fx.Int32(gmem_byte))) - fx.copy(dma_atom, src, dst) - - # ── Scheduling hints (ported from old pipeline) ─────────── - def build_scheduler(numer: int, denom: int): - if const_expr(denom <= 0): - return [] - if const_expr(numer <= 0): - return [0] * denom - out = [] - prev = 0 - for i in range_constexpr(denom): - cur = ((i + 1) * numer + (denom - 1)) // denom - out.append(cur - prev) - prev = cur - return out - - def hot_loop_scheduler(): - mfma_group = num_acc_n - - if const_expr(is_gfx942): - mfma_total = (k_iters * 2) * m_repeat * mfma_group - mfma_per_iter = 2 * mfma_group - sche_iters = 0 if mfma_per_iter == 0 else (mfma_total // mfma_per_iter) - - rocdl.sched_dsrd(2) - rocdl.sched_mfma(1) - if const_expr(tile_m == 16): - rocdl.sched_vmem(1) - rocdl.sched_mfma(1) - if const_expr(tile_m == 16): - rocdl.sched_vmem(1) - - if const_expr(num_acc_n < 4): - rocdl.sched_dsrd(1) - rocdl.sched_mfma(1) - if const_expr(tile_m == 16): - rocdl.sched_vmem(1) - rocdl.sched_dsrd(1) - rocdl.sched_mfma(1) - if const_expr(tile_m == 16): - rocdl.sched_vmem(1) - rocdl.sched_mfma(1) - - dswr_tail = num_a_loads - dstr_advance = 2 - if const_expr(dswr_tail > sche_iters): - dswr_tail = sche_iters - dswr_start = max(sche_iters - dswr_tail - dstr_advance, 0) - - for sche_i in range_constexpr(sche_iters): - rocdl.sched_vmem(1) - rocdl.sched_mfma(mfma_group) - rocdl.sched_dsrd(1) - rocdl.sched_mfma(mfma_group) - if const_expr(sche_i >= dswr_start - 1): - rocdl.sched_dswr(1) - else: - if const_expr(use_mfma_scale_128): - element_k_per_mfma = 128 - else: - element_k_per_mfma = 32 - num_mfma_per_tile_k = tile_k // element_k_per_mfma - mfma_total = num_mfma_per_tile_k * m_repeat * mfma_group - dswr_tail = num_a_loads - dstr_advance = 2 - if const_expr(dswr_tail > mfma_total): - dswr_tail = mfma_total - dsrd_preload_eff = min(int(dsrd_preload), num_ds_load) - dvmem_preload_eff = min(int(dvmem_preload), num_gmem_loads) - vmem_remaining = num_gmem_loads - dvmem_preload_eff - dsrd_remaining = num_ds_load - dsrd_preload_eff - if const_expr(vmem_remaining > 0 and vmem_remaining < mfma_total): - vmem_schedule = build_scheduler(vmem_remaining, vmem_remaining) + [ - 0 - ] * (mfma_total - vmem_remaining) - else: - vmem_schedule = build_scheduler(vmem_remaining, mfma_total) - dsrd_schedule = build_scheduler(dsrd_remaining, mfma_total) - dswr_start = max(mfma_total - dswr_tail - dstr_advance, 0) - last_dsrd_mfma_idx = -1 - for sched_idx in range_constexpr(mfma_total): - if const_expr(dsrd_schedule[sched_idx]): - last_dsrd_mfma_idx = sched_idx - dswr_start = max(dswr_start, last_dsrd_mfma_idx + 1) - idx_ds_read = dsrd_preload_eff - idx_gmem_load = dvmem_preload_eff - idx_ds_write = 0 - if const_expr(dvmem_preload_eff): - rocdl.sched_vmem(dvmem_preload_eff) - if const_expr(dsrd_preload_eff): - rocdl.sched_dsrd(dsrd_preload_eff) - for mfma_idx in range_constexpr(mfma_total): - rocdl.sched_mfma(1) - n_dsrd = dsrd_schedule[mfma_idx] - if const_expr(n_dsrd and (idx_ds_read < num_ds_load)): - if const_expr(idx_ds_read + n_dsrd > num_ds_load): - n_dsrd = num_ds_load - idx_ds_read - if const_expr(n_dsrd): - rocdl.sched_dsrd(n_dsrd) - idx_ds_read += n_dsrd - n_vmem = vmem_schedule[mfma_idx] - if const_expr(n_vmem and (idx_gmem_load < num_gmem_loads)): - if const_expr(idx_gmem_load + n_vmem > num_gmem_loads): - n_vmem = num_gmem_loads - idx_gmem_load - if const_expr(n_vmem): - rocdl.sched_vmem(n_vmem) - idx_gmem_load += n_vmem - if const_expr( - (not use_async_copy) - and (idx_ds_write < dswr_tail) - and (mfma_idx >= dswr_start) - ): - rocdl.sched_dswr(1) - idx_ds_write += 1 - if const_expr((not use_async_copy) and idx_ds_write < num_a_loads): - rocdl.sched_dswr(num_a_loads - idx_ds_write) - - rocdl.sched_barrier(0) - - # Accumulator addressing, shared by the block-scale rescale and the - # epilogue: element p of frag_C is row bx_m + mi*16 + lane_div_16*4 + ii, - # column by_n + (ni*num_waves + wave_id)*16 + lane_mod_16. - bx_m = bid_x * tile_m - by_n = bid_y * tile_n - wave_id = gpu.thread_id("x") // 64 - lane_id = gpu.thread_id("x") % 64 - lane_div_16 = lane_id // 16 - lane_mod_16 = lane_id % 16 - - # Scale operands go through the layout API rather than aiter's - # buffer_ops: on this FlyDSL, RawPtrBufferLoadOp wants `aux` as an - # operand, and buffer_ops passes None when no cache modifier is set, so - # every buffer_ops.buffer_load in this file raises. That breaks the - # upstream ptpc epilogue too; blockscale simply does not use it. - if const_expr(is_blockscale): - gSA = fx.rocdl.make_buffer_tensor( - arg_scale_a, - max_size=False, - num_records_bytes=fx.Int64(i32_m) * fx.Int64(kb_count) * fx.Int64(4), - ) - gSB = fx.rocdl.make_buffer_tensor(arg_scale_b) - sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) - sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) - sc_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) - sc_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) - sc_reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - sa_reg_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - sc_reg_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) - - def _load_sa(mi, kb): - """This lane's A scale for M-tile ``mi`` at K-block ``kb``. - - ``arg_scale_a`` is physically [K/128, M] (column-major [M, K/128]), - the layout aiter_per1x128_quant(transpose_scale=True) emits. With - the operands swapped a lane owns a single row -- ``lane % 16`` -- - so this is one scalar, and combined with B's it gives the one - per-tile ``c_scale_thread`` CK promotes with. - """ - if const_expr(swap_ab): - idx = fx.Int32(kb) * fx.Int32(i32_m) + bx_m + mi * 16 + lane_mod_16 - fx.copy(sc_atom_1, fx.slice(sa_div, (None, fx.Int32(idx))), sa_reg_1) - return Vec(fx.memref_load_vec(sa_reg_1))[0] - idx = fx.Int32(kb) * fx.Int32(i32_m) + bx_m + mi * 16 + lane_div_16 * 4 - fx.copy(sc_atom_4, fx.slice(sa_div, (None, fx.Int32(idx))), sc_reg_4) - return Vec(fx.memref_load_vec(sc_reg_4)) - - def _load_sb(ni, kb): - """The single B scale covering N-tile ``ni`` at K-block ``kb``. - - ``arg_scale_b`` is [N/128, K/128] row-major, and a 16-column group - lies wholly inside one 128-wide block, so this is one scalar. - """ - nblk = (by_n + (ni * num_waves + wave_id) * 16) // 128 - idx = nblk * fx.Int32(kb_count) + fx.Int32(kb) - fx.copy(sc_atom_1, fx.slice(sb_div, (None, fx.Int32(idx))), sc_reg_1) - return Vec(fx.memref_load_vec(sc_reg_1))[0] - - def mma_promote(a_stage, cur_frag_B, ki, kb): - """One scale block's MFMAs, promoted a tile at a time. - - CK's shape (blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle - _v1.hpp): the loop is (m0, n0, kscale) with k innermost and the - partial lives in ``c_thread_buf_per_scale``, a single - ``GetRegSizePerXdlops()``-wide buffer. Only one tile's 4 VGPRs are - live at a time, which is why CK pays almost nothing for the block - scale where both whole-fragment forms are expensive here: a running - rescale needs a divide per element per block (+660us) and a - whole-fragment promotion needs a duplicate accumulator and spills - (2753us). - - frag_C's modes are ((4,1), m_repeat, num_acc_n) and frag_A / - frag_B's are (elems, m_repeat|num_acc_n, k_iters), so one tile is a - plain slice of each; ``fx.gemm`` is handed the bare atom rather than - tiled_mma, exactly as Mfma16x16x128._do_mma does. - """ - sb = [_load_sb(ni, kb) for ni in range_constexpr(num_acc_n)] - # Operands are materialised into bare rmem tensors, the way - # Mfma16x16x128._make_operand_frag does: fx.gemm on the atom does - # not take a sliced view, it leaves a ub.poison the rmem-to-SSA - # promotion pass then rejects. Each is built once per k-step, not - # once per (mi, ni) pair. - b_ops = [] - for ni in range_constexpr(num_acc_n): - b_op = fx.make_rmem_tensor(8, fx.Int32) - b_op.store(Vec(cur_frag_B[None, ni, ki].load()).bitcast(fx.Int32)) - b_ops.append(b_op) - acc = Vec(frag_C.load()) - vals = [acc[p] for p in range_constexpr(acc_size)] - for mi in range_constexpr(m_repeat): - sa = _load_sa(mi, kb) - a_op = fx.make_rmem_tensor(8, fx.Int32) - a_op.store(Vec(frag_A[None, mi, ki].load()).bitcast(fx.Int32)) - # All of this M-row's MFMAs are issued before any of their - # results is read. With tile_k=128 a scale block is a single - # MFMA, so promoting each one where it is issued exposes the - # full MFMA latency; CK hides it by keeping k innermost, which - # is not available here. num_acc_n partials cost 4*num_acc_n - # VGPRs. - parts = [] - for ni in range_constexpr(num_acc_n): - part = fx.make_rmem_tensor(4, fx.Float32) - part.store(Vec.filled(4, 0.0, fx.Float32)) - # Operands exchanged: fx.gemm(atom, c, b, a, c) computes - # (A B)^T in the accumulator's native layout, so a lane's - # four values become four consecutive *columns* of one row - # instead of four rows. gcnasm's mfma_adaptor_swap_ab. - if const_expr(swap_ab): - fx.gemm(mma_atom, part, b_ops[ni], a_op, part) - else: - fx.gemm(mma_atom, part, a_op, b_ops[ni], part) - parts.append(part) - for ni in range_constexpr(num_acc_n): - sab = ( - [sa * sb[ni]] * 4 - if const_expr(swap_ab) - else [sa[ii] * sb[ni] for ii in range_constexpr(4)] - ) - pv = Vec(parts[ni].load()) - # Explicit fma: written as ``acc + pv * (sa * sb)`` the - # compiler emits v_pk_mul_f32 + v_pk_add_f32, three ops per - # element against CK's one v_fmac_f32. Hoisting the scale - # product and contracting the rest gets it to two. - for ii in range_constexpr(4): - # Slot order is ours to choose, and choosing it once - # here is free: the swapped write-out wants - # (ii, num_acc_n, m_repeat), and producing the - # accumulator in that order avoids a 64-value permute - # in the epilogue, which costs 8 VGPRs because source - # and destination are live together. - idx = ( - mi * (num_acc_n * 4) + ni * 4 + ii - if const_expr(swap_ab) - else ni * (m_repeat * 4) + mi * 4 + ii - ) - vals[idx] = fx.math.fma(pv[ii], sab[ii], vals[idx]) - frag_C.store(vector.from_elements(T.vec(acc_size, T.f32), vals)) - - # ── Pipeline stage (double-buffered B via split fragments) ─ - def mma_kloop(a_stage, cur_frag_B, tile_idx=None): - for ki in range_constexpr(k_iters): - fx.copy( - uni_copy, - pA_s2r_stages[a_stage][None, None, ki], - frag_A_retile[None, None, ki], - ) - if const_expr(is_blockscale): - # One MMA k-step is one 128-wide scale block. - mma_promote( - a_stage, - cur_frag_B, - ki, - fx.Int32(tile_idx) * fx.Int32(k_iters) + fx.Int32(ki), - ) - continue - k_coord = ki if (use_mfma_scale_128 or use_mfma_k32) else (None, ki) - fx.gemm( - tiled_mma, - frag_C, - frag_A[None, None, k_coord], - cur_frag_B[None, None, k_coord], - frag_C, - ) - - def pipeline_2stage(read_stage, next_k_val=None, read_next=True, cur_k_val=None): - write_stage = read_stage ^ 1 - a_read = read_stage - a_write = write_stage - cur_frag_B = frag_B_stages[read_stage] - do_next = read_next and next_k_val is not None - if const_expr(use_async_copy): - if const_expr(do_next): - dma_a_to_lds(next_k_val, a_write) - fx.copy( - buf_copy, - pB_g[None, None, None, k_off + next_k_val], - frag_B_retile_stages[write_stage], - ) - mma_kloop(a_read, cur_frag_B, cur_k_val) - if const_expr(enable_scheduler): - hot_loop_scheduler() - if const_expr(do_next): - rocdl.s_waitcnt(num_b_loads) - gpu.barrier() - return - if const_expr(do_next): - global_k_next = k_off + next_k_val - fx.copy(buf_copy, pA_g[None, None, None, global_k_next], frag_copy_A) - fx.copy( - buf_copy, - pB_g[None, None, None, global_k_next], - frag_B_retile_stages[write_stage], - ) - mma_kloop(a_read, cur_frag_B, cur_k_val) - if const_expr(do_next): - fx.copy(uni_copy, frag_copy_A, pA_s_stages[a_write][None, None, None]) - if const_expr(enable_scheduler): - hot_loop_scheduler() - if const_expr(do_next): - gpu.barrier() - - # ── Prologue ────────────────────────────────────────────── - acc_zero = ( - Vec.filled(acc_size, 0, Int32) - if const_expr(is_int8) - else Vec.filled(acc_size, 0.0, Float32) - ) - if const_expr(use_async_copy): - dma_a_to_lds(fx.Int32(0), 0) - fx.copy(buf_copy, pB_g[None, None, None, k_off], frag_B_retile_stages[0]) - frag_C.store(acc_zero) - rocdl.s_waitcnt(num_b_loads) - gpu.barrier() - else: - fx.copy(buf_copy, pA_g[None, None, None, k_off], frag_copy_A) - fx.copy(buf_copy, pB_g[None, None, None, k_off], frag_B_retile_stages[0]) - frag_C.store(acc_zero) - fx.copy(uni_copy, frag_copy_A, pA_s_stages[0][None, None, None]) - gpu.barrier() - rocdl.sched_barrier(0) - - # ── Main tile loop ──────────────────────────────────────────── - if const_expr(lds_stage == 1 and num_tiles > 1): - frag_Bc = frag_B_stages[0] - frag_Bc_retile = frag_B_retile_stages[0] - for iv, state in range(0, num_tiles - 1, 1, init=[frag_C.load()]): - frag_C.store(state[0]) - k_next = fx.Int32(iv + 1) - mma_kloop(0, frag_Bc) - if const_expr(not use_async_copy): - fx.copy( - buf_copy, pA_g[None, None, None, k_off + k_next], frag_copy_A - ) - fx.copy( - buf_copy, - pB_g[None, None, None, k_off + k_next], - frag_Bc_retile, - ) - gpu.barrier() # single buffer: all reads done before overwrite - if const_expr(use_async_copy): - dma_a_to_lds(k_next, 0) - else: - fx.copy(uni_copy, frag_copy_A, pA_s_stages[0][None, None, None]) - if const_expr(enable_scheduler): - hot_loop_scheduler() - if const_expr(use_async_copy): - rocdl.s_waitcnt(num_b_loads) - gpu.barrier() - results = yield [frag_C.load()] - frag_C.store(results) - elif const_expr(lds_stage == 2): - # 2-tile/iter ping-pong: middle loop runs 2 tiles/iter - is_odd_tiles = (num_tiles % 2) == 1 - tail = 1 if is_odd_tiles else 2 - loop_end = (num_tiles - tail) // 2 - - def two_tiles(k_base): - pipeline_2stage( - read_stage=0, next_k_val=k_base + fx.Int32(1), cur_k_val=k_base - ) - pipeline_2stage( - read_stage=1, - next_k_val=k_base + fx.Int32(2), - cur_k_val=k_base + fx.Int32(1), - ) - - if const_expr(loop_end > 0 and use_async_copy): - for iv in range_constexpr(loop_end): - two_tiles(fx.Int32(iv * 2)) - elif const_expr(loop_end > 0): - for iv, state in range(0, loop_end, 1, init=[frag_C.load()]): - frag_C.store(state[0]) - two_tiles(fx.Int32(iv * 2)) - results = yield [frag_C.load()] - frag_C.store(results) - k_tail0 = num_tiles - tail # first tile handled by the peeled tail - for j in range_constexpr(tail - 1): - pipeline_2stage( - read_stage=(k_tail0 + j) % 2, - next_k_val=fx.Int32(k_tail0 + j + 1), - cur_k_val=fx.Int32(k_tail0 + j), - ) - - # ── Epilogue-operand preloads (scale_a / scale_b / bias) ───────── - def load_epi_operands(): - s_a = s_b = bias = None - if const_expr(is_blockscale): - # Nothing to preload: mma_promote() folds every K-block's - # scales into frag_C as it goes, so the epilogue below sees an - # already-scaled accumulator. - pass - elif const_expr(is_8bit): - # Per-row(scale_a) × per-col(scale_b) scaling, applied in the epilogue. - scale_b_rsrc = buffer_ops.create_buffer_resource( - arg_scale_b, max_size=True - ) - s_b = [ - buffer_ops.buffer_load( - scale_b_rsrc, - fx.Int32(by_n + (ni * num_waves + wave_id) * 16 + lane_mod_16), - vec_width=1, - dtype=T.f32, - ) - for ni in range_constexpr(num_acc_n) - ] - scale_a_rsrc = buffer_ops.create_buffer_resource( - arg_scale_a, - max_size=False, - num_records_bytes=fx.Int64(i32_m) * fx.Int64(4), - ) - s_a = [ - Vec( - buffer_ops.buffer_load( - scale_a_rsrc, - fx.Int32(bx_m + mi * 16 + lane_div_16 * 4), - vec_width=4, - dtype=T.f32, - ) - ).bitcast(fx.Float32) - for mi in range_constexpr(m_repeat) - ] - if const_expr(_has_bias): - # Per-column bias (out_dtype), one scalar per N-block, shared across rows. - bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=True) - bias_elem_ty = T.bf16 if out_dtype == "bf16" else T.f16 - bias = [ - fx.Float32( - buffer_ops.buffer_load( - bias_rsrc, - fx.Int32( - by_n + (ni * num_waves + wave_id) * 16 + lane_mod_16 - ), - vec_width=1, - dtype=bias_elem_ty, - ) - ) - for ni in range_constexpr(num_acc_n) - ] - return s_a, s_b, bias - - overlap_epi_load = ( - acc_size <= 64 - ) # small enough accumulator to keep operands live over the MMA - s_a_vals = s_b_vals = bias_vals = None - if const_expr(overlap_epi_load): - s_a_vals, s_b_vals, bias_vals = load_epi_operands() - - # Final MMA stage — overlaps the epilogue-operand loads when issued above. - if const_expr(lds_stage == 1): - mma_kloop(0, frag_B_stages[0]) - else: - pipeline_2stage( - read_stage=(num_tiles - 1) % 2, - read_next=False, - cur_k_val=fx.Int32(num_tiles - 1), - ) - - # ── Epilogue ───────────────────────────────────────────── - if const_expr((not is_8bit or is_blockscale) and not _has_epilogue): - frag_C_out.store(Vec(frag_C.load()).to(out_elem_cls)) - else: - if const_expr(not overlap_epi_load): - s_a_vals, s_b_vals, bias_vals = load_epi_operands() - - def apply_activation(val_s): - # ReLU/SiLU/GeLU : maximumf for relu; exp+rcp for silu; - # tanh-approx gelu expanded through a non-positive exponent (no overflow). - if const_expr(_has_relu): - return fx.Float32(val_s).maximumf(fx.Float32(0.0)) - if const_expr(_has_silu): - exp_neg = math.exp(val_s * fx.Float32(-1.0)) - return val_s * (fx.Float32(1.0) / (fx.Float32(1.0) + exp_neg)) - if const_expr(_has_gelu): - half_f32 = fx.Float32(0.5) - one_f32 = fx.Float32(1.0) - zero_f32 = fx.Float32(0.0) - two_f32 = fx.Float32(2.0) - x3 = val_s * val_s * val_s - y = fx.Float32(0.7978845608) * (val_s + fx.Float32(0.044715) * x3) - abs_y = fx.Float32(y).maximumf(zero_f32 - y) - e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) - denom = one_f32 + e_neg2abs - numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) - return half_f32 * val_s * (numerator * (one_f32 / denom)) - return val_s - - acc_vec = Vec(frag_C.load()) - out_elems = [] - for p in range_constexpr(acc_size): - ni = p // (m_repeat * 4) - mi = (p // 4) % m_repeat - ii = p % 4 - val = acc_vec[p] - if const_expr(is_int8): - val = val.to(Float32) - if const_expr(is_8bit and not is_blockscale): - val_s = (val * s_a_vals[mi][ii]) * s_b_vals[ni] - else: - val_s = val - if const_expr(_has_bias): - val_s = val_s + bias_vals[ni] - val_s = apply_activation(val_s) - out_elems.append(val_s.to(out_elem_cls)) - - out_vec = vector.from_elements( - T.vec(acc_size, out_elem_cls.ir_type), out_elems - ) - frag_C_out.store(out_vec) - - fx.copy(copy_out_c, frag_C_retile, pC_g) - if const_expr(split_k > 1): - # The store carries sc0|sc1, so waiting on it is the whole release. - rocdl.s_waitcnt(0) - gpu.barrier() - splitk_reduce_epilogue( - arg_c, - fx.Tensor( - fx.make_view( - fx.get_iter(arg_out), - fx.make_layout((PRESHUFFLE_M_MAX, N), (N, 1)), - ) - ), - arg_semaphore, - lds.split_flag.ptr, - tile_m, - tile_n, - total_threads, - final_out_elem_cls, - tid, - bid_x, - bid_y, - i32_m, - N, - split_k, - ) - - # ── Host launcher ───────────────────────────────────────────── - @flyc.jit - def launch_gemm( - arg_c: fx.Tensor, - arg_out: fx.Tensor, - arg_semaphore: fx.Tensor, - arg_a: fx.Tensor, - arg_b: fx.Tensor, - arg_scale_a: fx.Tensor, - arg_scale_b: fx.Tensor, - arg_bias: fx.Tensor, - i32_m: fx.Int32, - i32_n: fx.Int32, - stream: fx.Stream, - ): - CompilationContext.get_current() - - # MMA atom — layout_elem carries the dtype (Float16/BFloat16/Float8E4M3FN/etc) - if const_expr(use_mfma_k32): - mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, layout_elem)) - k_perm = fx.make_layout((8, 4), (1, 8)) - elif const_expr(is_f16_or_bf16): - mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 16, layout_elem)) - k_perm = fx.make_layout((4, 4, 2), (1, 8, 4)) - elif const_expr(is_int8): - mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, layout_elem, Int32)) - k_perm = fx.make_layout((8, 4, 2), (1, 16, 8)) - else: - # fp8: narrow atom here; the scale (16x16x128) tiled_mma is rebuilt in-kernel - mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, layout_elem)) - k_perm = fx.make_layout((8, 4, 2), (1, 16, 8)) - - tiled_mma = fx.make_tiled_mma( - mma_atom, - fx.make_layout((1, 4, 1), (0, 1, 0)), - fx.make_tile(None, None, k_perm), - ) - - # G2S tiled copy - val_per_thr = a_load_bytes // elem_bytes - thrs_k = tile_k // val_per_thr - thrs_m = total_threads // thrs_k - tiled_copy_g2s = fx.make_tiled_copy( - fx.make_copy_atom(fx.UniversalCopy128b(), layout_elem), - fx.make_layout( - ((thrs_k, thrs_m), (1, val_per_thr)), - ((thrs_m * val_per_thr, 1), (1, thrs_m)), - ), - fx.make_tile(thrs_m, tile_k), - ) - - # Preshuffle B layout (2D hierarchical) - kp_bytes = 16 - kp_elems = kp_bytes if elem_bytes == 1 else kp_bytes // elem_bytes - k_bytes_b = K * elem_bytes - n0 = N // 16 - k0 = k_bytes_b // 64 - s_nlane = kp_elems - s_klane = 16 * s_nlane - s_k0 = 4 * s_klane - s_n0 = k0 * s_k0 - preshuffle_B = fx.Tensor( - fx.make_view( - fx.get_iter(arg_b), - fx.make_layout( - ((16, n0), (kp_elems, 4, k0)), ((s_nlane, s_n0), (1, s_klane, s_k0)) - ), - ) - ) - - # Reshape A and C to 2D - M_max = PRESHUFFLE_M_MAX - arg_a_2d = fx.Tensor( - fx.make_view(fx.get_iter(arg_a), fx.make_layout((M_max, K), (K, 1))) - ) - arg_c_2d = fx.Tensor( - fx.make_view(fx.get_iter(arg_c), fx.make_layout((M_max, N), (N, 1))) - ) - - gx = (i32_m + (tile_m - 1)) // tile_m - gy = i32_n // tile_n - - kernel_gemm( - arg_c_2d, - arg_out, - arg_semaphore, - arg_a_2d, - preshuffle_B, - arg_scale_a, - arg_scale_b, - arg_bias, - i32_m, - i32_n, - tiled_mma, - tiled_copy_g2s, - value_attrs={"rocdl.waves_per_eu": waves_per_eu}, - ).launch( - grid=(gx, gy, split_k), - block=(256, 1, 1), - stream=stream, - ) - - if const_expr(is_f16_or_bf16 and num_acc_n <= 2): - launch_gemm.compile_hints["llvm_options"] = {"enable-post-misched": False} - - return launch_gemm From b00acd6e8ec2c8cec48a598a99bc88335e7e9e51 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 13:34:41 +0000 Subject: [PATCH 47/79] move the fused GEMM+AR into mori.ops.gemm_ar, and drop the aiter dependency SGLang calls this kernel in production, so it is library code and not a benchmark. Living under benchmark/ was also the only reason every module carried a sys.path.insert and imported its siblings by bare name; inside a package those all become ordinary relative imports. Nothing under python/mori/ imported aiter before this, and nothing does now. The two pieces that were imported are vendored instead: _gemm_a8w8_8wave.py aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py, verbatim but for split_row_major_2d inlined -- three lines against a 1230-line module. _shuffle.py preshuffle_b, the layout=(16, 16) path of aiter's shuffle_weight, checked bitwise against it on four shapes. kernels_fused.py keeps its deliberate divergence from the vendored pipeline -- the corrected s_waitcnt count, still wrong upstream -- and is marked at that site so a re-sync cannot quietly undo it. GemmAllReduceOp in op.py is the public API, mirroring EpDispatchCombineOp: it takes a Communicator and owns the symmetric window, the per-M kernel cache and the row padding that SGLang was hand-rolling. What stays with the caller is deployment policy -- the profitability threshold, and how m_max is chosen -- because those are shape- and workload-specific, while supports() answers only whether a shape is expressible. counter_chunks keeps the largest-divisor rule: chunks must divide the row bands per destination or the completion counter's modulo test never fires, and a fixed 8 has nothing to divide at m_pad=4096. Taking the largest is also worth 1104.8us against 1434.0us at M=15360. gemm_ar is registered lazily in mori.ops like dispatch_combine_v2, since it needs the optional flydsl extra. The 420-line research docstring that was the head of kernels_fused.py moves to README.md next to it. Measured after the move, 8x MI355X, [16384, 7168] K=2048, blockscale -- unchanged against the numbers this work was reported on: fused-sdma 1152.4us (1146.2 before) split-sdma 1465.9 (1465.5) gemm-only 369.4 (388.5) Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/__init__.py | 8 +- python/mori/ops/gemm_ar/README.md | 464 ++++++++++++++ python/mori/ops/gemm_ar/__init__.py | 74 +++ .../ar => python/mori/ops/gemm_ar}/_compat.py | 0 python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py | 599 ++++++++++++++++++ python/mori/ops/gemm_ar/_shuffle.py | 65 ++ .../mori/ops}/gemm_ar/kernels_fused.py | 473 ++------------ .../mori/ops/gemm_ar}/kernels_lsa.py | 7 +- .../mori/ops/gemm_ar}/kernels_sdma.py | 9 +- .../ar => python/mori/ops/gemm_ar}/layout.py | 0 python/mori/ops/gemm_ar/op.py | 270 ++++++++ 11 files changed, 1529 insertions(+), 440 deletions(-) create mode 100644 python/mori/ops/gemm_ar/README.md create mode 100644 python/mori/ops/gemm_ar/__init__.py rename {benchmark/cco/flydsl/ar => python/mori/ops/gemm_ar}/_compat.py (100%) create mode 100644 python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py create mode 100644 python/mori/ops/gemm_ar/_shuffle.py rename {benchmark/cco/flydsl => python/mori/ops}/gemm_ar/kernels_fused.py (74%) rename {benchmark/cco/flydsl/ar => python/mori/ops/gemm_ar}/kernels_lsa.py (98%) rename {benchmark/cco/flydsl/ar => python/mori/ops/gemm_ar}/kernels_sdma.py (98%) rename {benchmark/cco/flydsl/ar => python/mori/ops/gemm_ar}/layout.py (100%) create mode 100644 python/mori/ops/gemm_ar/op.py 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..80d742fce --- /dev/null +++ b/python/mori/ops/gemm_ar/README.md @@ -0,0 +1,464 @@ +# 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. + +## 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 + +```python +from mori.ops.gemm_ar import GemmAllReduceOp, preshuffle_b + +op = GemmAllReduceOp(comm, n=7168, k=2048, m_max=16384) +b_shuffled = preshuffle_b(b_fp8) # [N, K] -> MFMA preshuffled +out = op(a_fp8, b_shuffled, a_scale, b_scale) # [M, N] bf16, already reduced +``` + +`a_scale` is the A 1x128 block scale, physically `[K/128, M]` fp32 (i.e. `[M, +K/128]` column-major -- what `aiter_per1x128_quant(transpose_scale=True)` +emits). `b_scale` is `[N/128, K/128]` fp32 row-major. + +`M` must be a multiple of `world_size * block_m`; `op.padded_m(m, world_size)` +rounds up and `op.pad_rows` zero-extends to it. A padded row produces a zero +output row, so the caller slices the result back to `m`. + +## 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 -- having the epilogue store straight into +a peer -- would first require coalescing ``StoreC._store_bf16``, which writes one +bf16 at a time through ``BufferCopy16b``; gcnasm measured that lane-scatter at +0.26x when pushed to a peer. See ``LSA_NOTE`` at the bottom. + +## 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..34943ab1f --- /dev/null +++ b/python/mori/ops/gemm_ar/__init__.py @@ -0,0 +1,74 @@ +# 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, 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", + "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/benchmark/cco/flydsl/ar/_compat.py b/python/mori/ops/gemm_ar/_compat.py similarity index 100% rename from benchmark/cco/flydsl/ar/_compat.py rename to python/mori/ops/gemm_ar/_compat.py 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..63fee77c4 --- /dev/null +++ b/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py @@ -0,0 +1,599 @@ +# 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/benchmark/cco/flydsl/gemm_ar/kernels_fused.py b/python/mori/ops/gemm_ar/kernels_fused.py similarity index 74% rename from benchmark/cco/flydsl/gemm_ar/kernels_fused.py rename to python/mori/ops/gemm_ar/kernels_fused.py index f0d2b43aa..b6fa23653 100644 --- a/benchmark/cco/flydsl/gemm_ar/kernels_fused.py +++ b/python/mori/ops/gemm_ar/kernels_fused.py @@ -19,36 +19,27 @@ # 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. -"""aiter's 8-wave fp8 GEMM with the all-reduce's scatter fused into its epilogue. +"""The 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. +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 ``ar.kernels_sdma`` +phases still run as their own kernels, reused verbatim from ``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. - -## 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 -- having the epilogue store straight into -a peer -- would first require coalescing ``StoreC._store_bf16``, which writes one -bf16 at a time through ``BufferCopy16b``; gcnasm measured that lane-scatter at -0.26x when pushed to a peer. See ``LSA_NOTE`` at the bottom. +``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: +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 @@ -56,425 +47,47 @@ 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## 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. - -## Pinned copy +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`` -(aiter ``aiter/ops/flydsl/kernels/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 +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. -``tests/python/cco/test_gemm_ar.py`` asserts this kernel's C output is -bit-identical to aiter's unfused kernel on the same inputs, which is what stops -the copy from rotting silently. - -The copy carries **one deliberate divergence** from aiter, marked at its site in -the main loop: the final ``wait_barrier`` count. aiter's is too permissive and -corrupts the output non-deterministically on large grids; the test above still -passes because it runs shapes small enough that aiter is correct there too. Do -not "restore" it when re-syncing the copy -- see the comment at that line, and -push the fix upstream instead. +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 aiter's -# gemm_a8w8_8wave.py omits it for the same reason: `fx.struct` reads the +# 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 os -import sys - 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 -from aiter.ops.flydsl.kernels.gemm_a8w8_8wave import ( - G2SLoader, - Mfma16x16x128, - S2RLoader, - StoreC, - _xcd_swizzle_any, - ceildiv, - compute_global_swizzle, - make_fp8_buffer_tensor, - wait_barrier, -) -from aiter.ops.flydsl.kernels.mfma_preshuffle_pipeline import split_row_major_2d - import mori.cco.device.flydsl as cco from mori.cco.device.flydsl import _bindings as raw_cco -_AR_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ar") -sys.path.insert(0, _AR_DIR) -from flydsl._mlir import ir # noqa: E402 -from flydsl._mlir.dialects import llvm as _llvm_d # noqa: E402 -from flydsl._mlir.dialects import scf # noqa: E402 - -from _compat import ( # noqa: E402 +from ._compat import ( CM_CACHED, CM_SC0_SC1, atomic_add_u32, @@ -483,11 +96,21 @@ buffer_store, create_buffer_resource_from_addr, i32_type, - wave_uniform_i64, - buffer_store, - create_buffer_resource_from_addr, 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, ) @@ -603,7 +226,7 @@ class _BlockScaleK: (``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 - ``aiter.gemm_a8w8_blockscale_bpreshuffle`` consumes, via sglang's + ``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 @@ -1158,7 +781,7 @@ def compile_fused_gemm_scatter( ``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 aiter's wrapper, which + 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 @@ -1240,7 +863,7 @@ def compile_fused_gemm_scatter( if not swap_ab: raise ValueError( "--quant blockscale needs --swap-ab: the unswapped epilogue " - "applies its scales inside aiter's StoreC, which this copy does " + "applies its scales inside the stock StoreC, which this copy does " "not override" ) if K % 128: @@ -1359,7 +982,7 @@ def kernel_gemm_scatter( dev_comm: Int64, win: Int64, ): - # ---- begin pinned copy of aiter kernel_gemm ---- + # ---- begin pinned copy of _gemm_a8w8_8wave kernel_gemm ---- F8_IR_t = fx.Float8E4M3FN.ir_type n_blocks = ceildiv(c_n, BLOCK_N) diff --git a/benchmark/cco/flydsl/ar/kernels_lsa.py b/python/mori/ops/gemm_ar/kernels_lsa.py similarity index 98% rename from benchmark/cco/flydsl/ar/kernels_lsa.py rename to python/mori/ops/gemm_ar/kernels_lsa.py index 8da12a041..ec07a8e84 100644 --- a/benchmark/cco/flydsl/ar/kernels_lsa.py +++ b/python/mori/ops/gemm_ar/kernels_lsa.py @@ -60,8 +60,6 @@ from __future__ import annotations -import os -import sys import flydsl.compiler as flyc import flydsl.expr as fx @@ -74,8 +72,7 @@ import mori.cco.device.flydsl as cco from mori.cco.device.flydsl import _bindings as raw_cco -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from _compat import ( # noqa: E402 +from ._compat import ( CM_CACHED, buffer_load, buffer_store, @@ -88,7 +85,7 @@ wave_uniform_i64, signal_store_u32, ) -from layout import MAX_WORLD # noqa: E402 +from .layout import MAX_WORLD def _spin_until(rsrc, flag): diff --git a/benchmark/cco/flydsl/ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py similarity index 98% rename from benchmark/cco/flydsl/ar/kernels_sdma.py rename to python/mori/ops/gemm_ar/kernels_sdma.py index 8f36755b8..ae8613c9a 100644 --- a/benchmark/cco/flydsl/ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -75,8 +75,6 @@ from __future__ import annotations -import os -import sys import flydsl.compiler as flyc import flydsl.expr as fx @@ -87,8 +85,7 @@ import mori.cco.device.flydsl as cco from mori.cco.device.flydsl import _bindings as raw_cco -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from _compat import ( # noqa: E402 +from ._compat import ( CM_CACHED, CM_SC1, buffer_load, @@ -101,8 +98,8 @@ signal_store_u32, wave_uniform_i64, ) -from kernels_lsa import _spin_until # noqa: E402 -from layout import MAX_WORLD # noqa: E402 +from .kernels_lsa import _spin_until +from .layout import MAX_WORLD #: 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 diff --git a/benchmark/cco/flydsl/ar/layout.py b/python/mori/ops/gemm_ar/layout.py similarity index 100% rename from benchmark/cco/flydsl/ar/layout.py rename to python/mori/ops/gemm_ar/layout.py diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py new file mode 100644 index 000000000..298ee18bd --- /dev/null +++ b/python/mori/ops/gemm_ar/op.py @@ -0,0 +1,270 @@ +# 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 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 compile_fused_gemm_scatter +from .kernels_sdma import build_sdma_phases +from .layout import 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 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 + + +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) + return max(c for c in range(1, min(MAX_CHUNKS, bands) + 1) if bands % c == 0) + + +def supports( + m: int, + n: int, + k: int, + world_size: int, + *, + block_n: int = DEFAULT_BLOCK_N, +) -> 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``. + """ + return ( + 2 <= world_size <= 8 + and n % block_n == 0 + and n % SCALE_BLOCK_K == 0 + and k % SCALE_BLOCK_K == 0 + and m > 0 + ) + + +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, + ): + self.comm = comm + self.rank = comm.rank + self.world_size = comm.world_size + self.n, self.k = n, k + self.block_m, self.block_n = block_m, block_n + self.sdma_queues = sdma_queues + self.m_max = padded_m(m_max, self.world_size, block_m) + + if not supports(self.m_max, n, k, self.world_size, block_n=block_n): + raise ValueError( + f"unsupported shape: M<={self.m_max} N={n} K={k} " + f"world_size={self.world_size}" + ) + + self.window_bytes = self._make_cfg(self.m_max).window_bytes + 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) + + self._cache: dict[int, tuple] = {} + self._pad_in: Optional[torch.Tensor] = None + + 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 _make_cfg(self, m: int) -> ArConfig: + cfg = ArConfig( + world_size=self.world_size, + m=m, + n=self.n, + recv_slots=self.world_size, + counter_chunks=counter_chunks(m, self.world_size, self.block_m), + ) + cfg.validate() + return cfg + + 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) + c = from_gpu_ptr( + self.mem.ptr + cfg.input_off, (m, self.n), torch.bfloat16 + ) + out = from_gpu_ptr( + self.mem.ptr + cfg.output_off, (m, self.n), torch.bfloat16 + ) + hit = (gemm, parts, c, out) + self._cache[m] = hit + return hit + + 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. + """ + m, k = x.shape + if self._pad_in is None or self._pad_in.shape[0] < m_pad: + self._pad_in = torch.zeros((m_pad, k), dtype=x.dtype, device=x.device) + buf = self._pad_in[: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`. + + Both scale buffers are read linearly, in physical order: ``a_scale`` is + logically ``[M, K/128]`` but column-major, i.e. physically + ``[K/128, M]``; ``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. + """ + 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()" + ) + gemm, parts, c, out = self._compiled(m) + stream = fx.Stream(torch.cuda.current_stream()) + gemm( + a_fp8.contiguous().view(torch.int8).view(-1), + b_preshuffled.contiguous().view(torch.int8).view(-1), + c.view(-1), + a_scale.reshape(-1), + b_scale.reshape(-1), + m, + self.n, + self.dev_comm.ptr, + self.win.handle, + stream=stream, + ) + parts["drain"](self.dev_comm.ptr, self.win.handle, stream=stream) + parts["reduce"](self.dev_comm.ptr, self.win.handle, stream=stream) + parts["gather"](self.dev_comm.ptr, self.win.handle, stream=stream) + return out From e760674677dc9ddbae78f47522c21c9ea84953ca Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 13:35:06 +0000 Subject: [PATCH 48/79] point the benchmarks and tests at the package, and de-aiter the tests The benchmarks keep their place under benchmark/ and now import from mori.ops.gemm_ar; bench_gemm_ar.py loses --gemm-impl and --tile-* along with the preshuffle4w kernel they selected. Three tests were written against aiter as the reference and are rebuilt on references mori owns: test_pinned_copy_matches_aiter_kernel_bitwise is gone. It asserted the copy is bit-identical to aiter's kernel, which stops being the contract once mori owns the copy. Replaced by test_ptpc_matches_an_fp32_reference, which runs the same kernel on the same shapes against fp32 in torch, so the ptpc path keeps a guard. test_blockscale_* drops the cross-check against gemm_a8w8_blockscale_bpreshuffle and keeps the fp32 reference it already computed. The gate is unchanged: relL2 < 3e-3, the fp8 floor. test_swap_ab_is_bitwise_identical now builds its reference from this same kernel with swap_ab=False. That is the property under test -- swapping the MFMA operands must be a pure register relabelling -- and it is a stronger check than comparing to a third implementation. test_gemm_ar_bench.py is new: it drives bench_gemm_ar.py over the modes and parses its RESULT_JSON. Accuracy is a hard gate (validated, rel_l2); absolute timings are printed but not asserted, because a shared runner would fail them on a busy box and pass a real regression on an idle one. The one perf assertion is relative -- fused-sdma beats split-sdma, both measured back to back on the same box, where the recorded margin is 21%. The optional aiter cross-checks that remain (test_flydsl_ar's bitwise comparison against aiter's custom AR, and bench_ar's AR_CHECK_VS_AITER) are importorskip and skip cleanly without it. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_probe_lsa.py | 3 +- benchmark/cco/flydsl/ar/bench_ar.py | 7 +- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 101 +++--------- tests/python/cco/test_flydsl_ar.py | 16 +- tests/python/cco/test_gemm_ar.py | 145 ++++++++---------- tests/python/cco/test_gemm_ar_bench.py | 136 ++++++++++++++++ 6 files changed, 228 insertions(+), 180 deletions(-) create mode 100644 tests/python/cco/test_gemm_ar_bench.py diff --git a/benchmark/cco/flydsl/ar/_probe_lsa.py b/benchmark/cco/flydsl/ar/_probe_lsa.py index df171857d..5e3d42d3d 100644 --- a/benchmark/cco/flydsl/ar/_probe_lsa.py +++ b/benchmark/cco/flydsl/ar/_probe_lsa.py @@ -26,8 +26,7 @@ import mori.cco.device.flydsl as cco from mori.cco import CCODevCommRequirements, Communicator, GDA_CONNECTION_NONE -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from _compat import ( # noqa: E402 +from mori.ops.gemm_ar._compat import ( CM_CACHED, FLYDSL_VERSION, buffer_load, diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py index dc88c5a19..944f3aeb2 100644 --- a/benchmark/cco/flydsl/ar/bench_ar.py +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -77,8 +77,7 @@ ) from mori.tensor_utils import from_gpu_ptr -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from layout import ArConfig # noqa: E402 +from mori.ops.gemm_ar import ArConfig VMM_SLACK = 256 * 1024 * 1024 @@ -189,13 +188,13 @@ def run(args) -> int: dc = comm.create_dev_comm(reqs) if args.backend == "lsa": - from kernels_lsa import build_lsa_ar + 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 kernels_sdma import build_sdma_phases + from mori.ops.gemm_ar import build_sdma_phases parts = build_sdma_phases( cfg, diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 6424e26ab..62889acda 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -22,7 +22,7 @@ # SOFTWARE. """GEMM + all-reduce at the DSV4-Pro ``wo_b`` shape: fused vs split. - MORI_SOCKET_IFNAME=lo MORI_ENABLE_SDMA=1 PYTHONPATH=/path/to/aiter \\ + 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 @@ -55,9 +55,9 @@ 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 aiter -8-wave kernel's native per-token/per-channel form, kept because the two bitwise -tests and every earlier measurement are on it. +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 @@ -69,20 +69,19 @@ 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. The gap is the -per-K-block scale, not the GEMM -- unscaled we are 224.0us. See -``kernels_preshuffle4w.py``, which ports CK's 4-wave B-out-of-LDS shape and -lands at 385.9us: aligned on the GEMM, still paying for the scale. +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 aiter GEMM's -under-counted ``s_waitcnt`` made every chunked run intermittently wrong, and +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 ``kernels_fused.py`` before trusting any fused number from this -benchmark: several of its options are intermittently wrong, and a single +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. """ @@ -107,14 +106,13 @@ ) from mori.tensor_utils import from_gpu_ptr -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, _HERE) -sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "ar")) -from kernels_fused import compile_fused_gemm_scatter # noqa: E402 -from kernels_preshuffle4w import compile_preshuffle_gemm # noqa: E402 -from kernels_lsa import build_lsa_ar # noqa: E402 -from kernels_sdma import build_sdma_phases # noqa: E402 -from layout import ArConfig # noqa: E402 +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") @@ -221,13 +219,7 @@ def _median_us(fn, warmup: int, iters: int, *, graph: bool = True) -> float: def run(args) -> int: - from aiter.ops.shuffle import shuffle_weight - local_rank, rank, world_size, uid = _setup_distributed() - if args.gemm_impl == "preshuffle4w" and args.mode != "gemm-only": - raise SystemExit( - "--gemm-impl preshuffle4w has no fused epilogue; use --mode gemm-only" - ) # 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. @@ -282,7 +274,7 @@ def run(args) -> int: cfg.validate() a, b, sa, sb = make_operands(rank, args.m, args.n, args.k, args.quant) - b_shuf = shuffle_weight(b, layout=(16, 16)) + b_shuf = preshuffle_b(b) vmm = max(4 * cfg.window_bytes + VMM_SLACK, VMM_SLACK) result = None @@ -306,26 +298,7 @@ def run(args) -> int: reqs.sdma_queue_count = args.sdma_queues dc = comm.create_dev_comm(reqs) - if args.gemm_impl == "preshuffle4w": - # CK's shape: 4 waves, B never staged in LDS. gemm-only, so the - # window's input region is just an ordinary output buffer here. - launch4w = compile_preshuffle_gemm( - N=args.n, - K=args.k, - tile_m=args.tile_m, - tile_n=args.tile_n, - tile_k=args.tile_k, - in_dtype="fp8", - out_dtype="bf16", - quant=args.quant, - swap_ab=args.p4w_swap_ab, - waves_per_eu=args.waves_per_eu, - xcd_swizzle=args.xcd_swizzle, - ) - semaphore = torch.zeros(1, device="cuda", dtype=torch.int32) - bias_unused = torch.zeros(1, device="cuda", dtype=torch.bfloat16) - - gemm = None if args.gemm_impl == "preshuffle4w" else compile_fused_gemm_scatter( + gemm = compile_fused_gemm_scatter( cfg, rank, K=args.k, @@ -368,10 +341,6 @@ def run(args) -> int: sa_arg, sb_arg = sa, sb def run_gemm(stream): - if args.gemm_impl == "preshuffle4w": - launch4w(c, c, semaphore, a, b_shuf, sa_arg, sb_arg, bias_unused, - args.m, args.n, stream=stream) - return gemm(a_i8, b_i8, c_flat, sa_arg, sb_arg, args.m, args.n, dc.ptr, win.handle, stream=stream) @@ -485,12 +454,6 @@ def once(): "rel_l2": rel_l2, "validated": validated, "timing": "eager" if args.eager else "graph", - "gemm_impl": args.gemm_impl, - "tile": ( - [args.tile_m, args.tile_n, args.tile_k] - if args.gemm_impl == "preshuffle4w" - else None - ), } print("RESULT_JSON " + json.dumps(result, sort_keys=True), flush=True) print( @@ -599,9 +562,8 @@ def build_parser() -> argparse.ArgumentParser: # 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: there the - # default has to stay "aiter's kernel verbatim", which is what - # ``test_pinned_copy_matches_aiter_kernel_bitwise`` checks. + # ``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, @@ -654,25 +616,6 @@ def build_parser() -> argparse.ArgumentParser: "and is worth 80us (355 vs 434)", ) p.add_argument("--sdma-queues", type=int, default=8) - p.add_argument( - "--gemm-impl", - choices=("8wave", "preshuffle4w"), - default="8wave", - help="8wave: the pinned aiter kernel, the only one with a fused " - "epilogue. preshuffle4w: CK's shape -- 4 waves, B loaded straight " - "to registers instead of through LDS. gemm-only.", - ) - p.add_argument("--tile-m", type=int, default=64, help="preshuffle4w only") - p.add_argument("--tile-n", type=int, default=256, help="preshuffle4w only") - p.add_argument("--tile-k", type=int, default=128, help="preshuffle4w only") - p.add_argument( - "--p4w-swap-ab", - action=argparse.BooleanOptionalAction, - default=False, - help="preshuffle4w: exchange the MFMA operands so a lane owns four " - "consecutive N instead of four M. Off by default -- it does what it " - "should to the instructions and still loses; see kernels_preshuffle4w.", - ) p.add_argument("--warmup", type=int, default=10) p.add_argument("--iters", type=int, default=51) p.add_argument("--eager", action="store_true") diff --git a/tests/python/cco/test_flydsl_ar.py b/tests/python/cco/test_flydsl_ar.py index ae55287a8..22bb9b871 100644 --- a/tests/python/cco/test_flydsl_ar.py +++ b/tests/python/cco/test_flydsl_ar.py @@ -26,7 +26,6 @@ torchrun and parses its ``RESULT_JSON`` line. """ -import importlib.util import json import os from pathlib import Path @@ -36,21 +35,10 @@ 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" -sys.path.insert(0, str(AR_DIR)) - - -def _load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -layout = _load("flydsl_ar_layout", AR_DIR / "layout.py") # -------------------------------------------------------------------------- diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index b05e766c2..ef2c9e96e 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -19,16 +19,14 @@ # 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 GEMM + all-reduce: anti-rot guard for the pinned kernel copy, plus smoke. +"""Numerics for ``mori.ops.gemm_ar``: the GEMM, the epilogue, and the layout math. -``benchmark/cco/flydsl/gemm_ar/kernels_fused.py`` holds a copy of aiter's -``kernel_gemm`` with an epilogue added. The single most valuable test here is not -that the fused path is correct end to end -- ``bench_gemm_ar.py`` checks that -against a host reference -- but that the *copy* still computes what aiter's -original computes. That one runs on a single GPU and needs no cco at all. +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 importlib.util import json import os from pathlib import Path @@ -38,24 +36,10 @@ import pytest import torch -REPO_ROOT = Path(__file__).resolve().parents[3] -GEMM_AR_DIR = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "gemm_ar" -AR_DIR = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "ar" -for _d in (str(AR_DIR), str(GEMM_AR_DIR)): - if _d not in sys.path: - sys.path.insert(0, _d) - - -def _load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - sys.modules[name] = module - spec.loader.exec_module(module) - return module - +from mori.ops.gemm_ar import layout -layout = _load("flydsl_ar_layout_gemm", AR_DIR / "layout.py") +REPO_ROOT = Path(__file__).resolve().parents[3] +BENCH = REPO_ROOT / "benchmark" / "cco" / "flydsl" / "gemm_ar" / "bench_gemm_ar.py" # -------------------------------------------------------------------------- @@ -96,38 +80,35 @@ def test_validate_rejects_zero_chunks(): # -------------------------------------------------------------------------- -# the anti-rot guard: our pinned copy must match aiter's original bit-for-bit +# 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_pinned_copy_matches_aiter_kernel_bitwise(m, n, k): - """The unfused copy must be indistinguishable from aiter's kernel. - - Bitwise, not approximate: the two are meant to be the same instruction - sequence over the same inputs, so any drift in the copied pipeline shows up - here rather than as a mystery in a fused benchmark. +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("aiter", reason="aiter not importable (set PYTHONPATH)") pytest.importorskip("flydsl") - from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 - from aiter.ops.shuffle import shuffle_weight - import flydsl.expr as fx - from kernels_fused import compile_fused_gemm_scatter + 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 = shuffle_weight(b, layout=(16, 16)) + b_shuf = preshuffle_b(b) - ref = torch.zeros(m, n, device="cuda", dtype=torch.bfloat16) - flydsl_8wave_gemm_a8(a, b_shuf, sa, sb, ref, 256, 256) + 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 @@ -150,38 +131,35 @@ def test_pinned_copy_matches_aiter_kernel_bitwise(m, n, k): stream=fx.Stream(torch.cuda.current_stream()), ) torch.cuda.synchronize() - assert torch.equal(got, ref), ( - "the pinned copy of aiter's kernel_gemm has drifted from the original; " - "re-sync benchmark/cco/flydsl/gemm_ar/kernels_fused.py" + + 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_agrees_with_the_kernel_the_model_dispatches(m, n, k): - """``--quant blockscale`` against ``gemm_a8w8_blockscale_bpreshuffle``. +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 - and dispatches that aiter kernel; this benchmark's own GEMM has to agree - with it, or the fused-vs-split numbers are being drawn at the wrong - operating point. Not a bitwise check: the two accumulate in different - orders, so both are compared to an fp32 reference and to each other. + 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 that kernel consumes -- ``sa`` logically - ``[M, K/128]`` but **column-major** (sglang's + 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 against it, - only this one lands at the fp8 floor; the others give 0.17-0.32. + ``[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("aiter", reason="aiter not importable (set PYTHONPATH)") + pytest.importorskip("flydsl") import flydsl.expr as fx - from aiter.ops.gemm_op_a8w8 import gemm_a8w8_blockscale_bpreshuffle - from aiter.ops.shuffle import shuffle_weight - sys.path.insert(0, str(GEMM_AR_DIR)) - sys.path.insert(0, str(AR_DIR)) - from kernels_fused import compile_fused_gemm_scatter + from mori.ops.gemm_ar import compile_fused_gemm_scatter, preshuffle_b BK = 128 g = torch.Generator(device="cuda").manual_seed(11) @@ -194,7 +172,7 @@ def test_blockscale_agrees_with_the_kernel_the_model_dispatches(m, n, k): + 0.01 ) sa = sa.t().contiguous().t() - b_shuf = shuffle_weight(b, layout=(16, 16)) + b_shuf = preshuffle_b(b) af, bf = a.float(), b.float() ref = torch.zeros(m, n, device="cuda", dtype=torch.float32) @@ -207,9 +185,6 @@ def test_blockscale_agrees_with_the_kernel_the_model_dispatches(m, n, k): ) rn = ref.norm() - theirs = torch.empty(m, n, device="cuda", dtype=torch.bfloat16) - gemm_a8w8_blockscale_bpreshuffle(a, b_shuf, sa, sb, out=theirs) - 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, @@ -228,12 +203,8 @@ def test_blockscale_agrees_with_the_kernel_the_model_dispatches(m, n, k): ) torch.cuda.synchronize() - rel_ours = ((ours.float() - ref).norm() / rn).item() - rel_theirs = ((theirs.float() - ref).norm() / rn).item() - rel_pair = ((ours.float() - theirs.float()).norm() / rn).item() - assert rel_ours < 3e-3, f"ours {rel_ours:.3e} vs the fp32 reference" - assert rel_theirs < 3e-3, f"aiter {rel_theirs:.3e} vs the fp32 reference" - assert rel_pair < 3e-3, f"ours vs aiter {rel_pair:.3e}" + 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)]) @@ -249,25 +220,35 @@ def test_swap_ab_is_bitwise_identical(m, n, k): """ if not torch.cuda.is_available(): pytest.skip("requires a GPU") - pytest.importorskip("aiter", reason="aiter not importable (set PYTHONPATH)") - from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 - from aiter.ops.shuffle import shuffle_weight - + pytest.importorskip("flydsl") import flydsl.expr as fx - from kernels_fused import compile_fused_gemm_scatter + 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 = shuffle_weight(b, layout=(16, 16)) + 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) - flydsl_8wave_gemm_a8(a, b_shuf, sa, sb, ref, 256, 256) + 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() - cfg = layout.ArConfig(world_size=2, m=m, n=n) for extra in ( {}, {"permlane": True}, @@ -286,7 +267,9 @@ def test_swap_ab_is_bitwise_identical(m, n, k): stream=fx.Stream(torch.cuda.current_stream()), ) torch.cuda.synchronize() - assert torch.equal(got, ref), f"swap_ab {extra} drifted from aiter" + 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(): @@ -325,7 +308,7 @@ def _run_bench(world_size, mode, m, n, k, extra=()): "torch.distributed.run", "--standalone", f"--nproc_per_node={world_size}", - str(GEMM_AR_DIR / "bench_gemm_ar.py"), + str(BENCH), "--mode", mode, "-m", str(m), "-n", str(n), "-k", str(k), "--warmup", "1", "--iters", "3", @@ -360,7 +343,7 @@ def test_fused_is_stable_across_repeats(): """ if torch.cuda.device_count() < 8: pytest.skip("the fused race only reproduces at world_size=8") - pytest.importorskip("aiter", reason="aiter not importable (set PYTHONPATH)") + pytest.importorskip("flydsl") seen = [] for _ in range(3): record = _run_bench(8, "fused-sdma", 4096, 7168, 1024) @@ -384,7 +367,7 @@ def test_gemm_ar_modes_agree(mode): """ if torch.cuda.device_count() < 2: pytest.skip("requires 2 GPUs") - pytest.importorskip("aiter", reason="aiter not importable (set PYTHONPATH)") + pytest.importorskip("flydsl") record = _run_bench(2, mode, 512, 1024, 256) assert record["validated"] is True assert record["max_rank_time_us"] > 0 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..04c2da398 --- /dev/null +++ b/tests/python/cco/test_gemm_ar_bench.py @@ -0,0 +1,136 @@ +# 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." + ) From 089d43eea59d00e7a3cc22b68673462151021e97 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 13:35:07 +0000 Subject: [PATCH 49/79] gate the fused GEMM+AR on CI: numerics assert, timings record ci_cco.yml did not touch any of this. Two steps, on the runner that already builds WITH_FLYDSL=1, GPU_ARCHS=gfx950 and BUILD_CCO_SDMA=ON with 8 GPUs: numerics test_gemm_ar.py + test_flydsl_ar.py, run from /tmp. From a repo cwd python/mori/ shadows the install, so a subpackage find_packages never picked up would read as green -- the same reason the JIT v2 step already runs from outside the checkout. Preceded by an explicit import of the op. benchmark test_gemm_ar_bench.py with -s, so the timings reach the log for a bisect to read. Accuracy is the gate; the only perf assertion is fused beating split. timeout-minutes 30 -> 60: FlyDSL compiles each shape on first use, and the two steps measure at about 4 minutes together on an idle box, but the job was already close to its limit. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci_cco.yml | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci_cco.yml b/.github/workflows/ci_cco.yml index ebec967f6..d72b94ecf 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 @@ -169,6 +169,28 @@ 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. + $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 1200 python -m pytest -v \ + $GITHUB_WORKSPACE/tests/python/cco/test_gemm_ar.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 "\ From 2182d60909d13f5426a2db87ccd0ec3a3e7a56fe Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 13:40:32 +0000 Subject: [PATCH 50/79] make the gemm_ar sources pass pre-commit These files never passed. ruff reports 8 errors on them at e7225e6e~1, and black would reformat all seven -- so pre-commit CI was already red on this branch before any of the work above. Beyond the mechanical reformat: * five dead locals removed (tiles_per_peer, elems_per_pack, blocks, stride_packs, and one more) -- all leftovers from earlier revisions of the index arithmetic, none of them read; * the `raw = lambda v: ...` unwrapper was duplicated verbatim in _permlane16_swap and _ds_bpermute; hoisted to a module-level `_raw`. The vendored _gemm_a8w8_8wave.py gets mori's MIT header like every other Python file here, since the License inserter hook has no exclude. Its upstream Apache-2.0 notice is kept below that, reworded so it reads as the attribution it is rather than as a competing grant on the file. Numerics re-run after the reformat: 38 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/_probe_lsa.py | 25 ++- benchmark/cco/flydsl/ar/bench_ar.py | 6 +- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 55 ++++- python/mori/ops/gemm_ar/_compat.py | 11 +- python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py | 32 ++- python/mori/ops/gemm_ar/kernels_fused.py | 191 +++++++++++------- python/mori/ops/gemm_ar/kernels_lsa.py | 13 +- python/mori/ops/gemm_ar/kernels_sdma.py | 11 +- python/mori/ops/gemm_ar/layout.py | 5 +- python/mori/ops/gemm_ar/op.py | 12 +- tests/python/cco/test_gemm_ar.py | 84 ++++++-- tests/python/cco/test_gemm_ar_bench.py | 18 +- 12 files changed, 327 insertions(+), 136 deletions(-) diff --git a/benchmark/cco/flydsl/ar/_probe_lsa.py b/benchmark/cco/flydsl/ar/_probe_lsa.py index 5e3d42d3d..790f18a5a 100644 --- a/benchmark/cco/flydsl/ar/_probe_lsa.py +++ b/benchmark/cco/flydsl/ar/_probe_lsa.py @@ -1,5 +1,26 @@ #!/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``. @@ -83,9 +104,7 @@ def probe_ar_kernel(dev_comm: Int64, win: Int64, flag: Int32): # 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 - ) + 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) diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py index 944f3aeb2..af9e96054 100644 --- a/benchmark/cco/flydsl/ar/bench_ar.py +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -248,7 +248,7 @@ def once(): 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()) + d = out.view(-1)[lo:hi].float() - ref.view(-1)[lo:hi].float() if d.norm().item() != 0.0: bad.append(r) print( @@ -258,9 +258,7 @@ def once(): ) comm.barrier() - elapsed_us = _median_us( - once, args.warmup, args.iters, graph=not args.eager - ) + elapsed_us = _median_us(once, args.warmup, args.iters, graph=not args.eager) comm.barrier() stats = torch.tensor([elapsed_us], dtype=torch.float64) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 62889acda..01274f482 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -157,13 +157,19 @@ def make_operands(rank: int, m: int, n: int, k: int, quant: str = "ptpc"): 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 + 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 + 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 @@ -261,7 +267,9 @@ def run(args) -> int: 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) + 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( @@ -341,8 +349,18 @@ def run(args) -> int: 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) + 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( @@ -391,7 +409,10 @@ def once(): 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) + print( + f"[rank {rank}] GEMM VALIDATION FAILED relL2={rel_l2:.3e}", + flush=True, + ) del ref elif ( args.mode != "gemm-only" @@ -485,8 +506,12 @@ def build_parser() -> argparse.ArgumentParser: "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-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) @@ -516,8 +541,16 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument( "--fence", choices=( - "all", "agent", "agent-leader", "nt-agent", "leader", "none", - "writethrough", "wt-agent", "raw-wt", "raw-wt-agent", + "all", + "agent", + "agent-leader", + "nt-agent", + "leader", + "none", + "writethrough", + "wt-agent", + "raw-wt", + "raw-wt-agent", "raw-wt-leader", ), default="none", diff --git a/python/mori/ops/gemm_ar/_compat.py b/python/mori/ops/gemm_ar/_compat.py index 610efb576..ee94964b4 100644 --- a/python/mori/ops/gemm_ar/_compat.py +++ b/python/mori/ops/gemm_ar/_compat.py @@ -161,11 +161,12 @@ def _byte_offset(offset, elem_ty): # *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 + _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.""" diff --git a/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py b/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py index 63fee77c4..c4e09cf93 100644 --- a/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py +++ b/python/mori/ops/gemm_ar/_gemm_a8w8_8wave.py @@ -1,6 +1,32 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors -# Modifications Copyright (C) 2026 Advanced Micro Devices, Inc. +# 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). diff --git a/python/mori/ops/gemm_ar/kernels_fused.py b/python/mori/ops/gemm_ar/kernels_fused.py index b6fa23653..9bfc171e4 100644 --- a/python/mori/ops/gemm_ar/kernels_fused.py +++ b/python/mori/ops/gemm_ar/kernels_fused.py @@ -276,7 +276,9 @@ def a_scales(self, base_row, kb): 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)) + 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. @@ -304,9 +306,9 @@ def ratio(prev, cur): 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): + 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 @@ -363,8 +365,9 @@ class _SwapABStoreC(StoreC): 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): + 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 @@ -432,13 +435,20 @@ def store(self, c_frag, base_row, base_col): 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) + 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. @@ -452,21 +462,23 @@ def _permlane16_swap(x, y): shuffle. """ st = ir.Type.parse("!llvm.struct<(i32, i32)>") - raw = lambda v: v.ir_value() if hasattr(v, "ir_value") else v - res = fx.rocdl.permlane16_swap(st, raw(x), raw(y), False, True) + 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), + 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) - raw = lambda v: v.ir_value() if hasattr(v, "ir_value") else v return fx.Int32( - fx.rocdl.ds_bpermute(i32, raw(fx.Int32(src_lane) * fx.Int32(4)), raw(value)) + fx.rocdl.ds_bpermute(i32, _raw(fx.Int32(src_lane) * fx.Int32(4)), _raw(value)) ) @@ -495,8 +507,13 @@ class _PermlaneStoreC(_SwapABStoreC): _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)) + 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. @@ -511,8 +528,9 @@ def store_all(self, frags, base_row, base_col, row_step, col_step): 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]) + 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 @@ -542,7 +560,9 @@ def _emit(self, c_frag, base_row, base_col, a_scales, b_scales): 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) + self._scaled(vec_f32[k], a_scales[ti], b_scales[tj][k]).to( + fx.BFloat16 + ) for k in range_constexpr(4) ], fx.BFloat16, @@ -600,7 +620,8 @@ def _emit(self, c_frag, base_row, base_col, a_scales, b_scales): else: fx.memref_store_vec(out8, self.reg_bf16_8) fx.copy( - self.out_atom_8, self.reg_bf16_8, + self.out_atom_8, + self.reg_bf16_8, fx.slice(self.c_div, (None, fx.Int32(idx))), ) @@ -638,8 +659,6 @@ class _WideStoreProbeC(_SwapABStoreC): ``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 = [ @@ -656,7 +675,9 @@ def store(self, c_frag, base_row, base_col): 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) + 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 @@ -878,13 +899,9 @@ def compile_fused_gemm_scatter( 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}" - ) + 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" - ) + 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 @@ -924,7 +941,6 @@ def compile_fused_gemm_scatter( # 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 - tiles_per_peer = m_tiles_per_peer * n_blocks_const # 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 @@ -940,8 +956,17 @@ def compile_fused_gemm_scatter( 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", + "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}" @@ -1095,7 +1120,14 @@ def kernel_gemm_scatter( else _LaneTransposeStoreC ) store_c = _cls( - A_scale, B_scale, C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B, + 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, @@ -1105,7 +1137,14 @@ def kernel_gemm_scatter( 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, + 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, @@ -1115,7 +1154,14 @@ def kernel_gemm_scatter( 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, + 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, @@ -1215,15 +1261,15 @@ def kernel_gemm_scatter( 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, - ) + (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) @@ -1275,16 +1321,16 @@ def kernel_gemm_scatter( # 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, - ) - ) + (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() @@ -1314,16 +1360,16 @@ def kernel_gemm_scatter( # 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, - ) - ) + (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) @@ -1348,8 +1394,10 @@ def kernel_gemm_scatter( # 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, + (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) @@ -1381,9 +1429,16 @@ def kernel_gemm_scatter( 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, + [ + (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) diff --git a/python/mori/ops/gemm_ar/kernels_lsa.py b/python/mori/ops/gemm_ar/kernels_lsa.py index ec07a8e84..faa8564cc 100644 --- a/python/mori/ops/gemm_ar/kernels_lsa.py +++ b/python/mori/ops/gemm_ar/kernels_lsa.py @@ -137,7 +137,6 @@ def build_lsa_ar( threads = cfg.threads blocks = cfg.blocks - elems_per_pack = cfg.elems_per_pack elem_dtype = fx.BFloat16 if cfg.elem_bytes == 2 else fx.Float32 num_packs = cfg.num_packs stride_packs = blocks * threads @@ -216,7 +215,9 @@ def ar_1stage(dev_comm: Int64, win: Int64): 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))) + 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): @@ -246,7 +247,9 @@ def ar_2stage(dev_comm: Int64, win: Int64): ) for i in range(ws) ] - tmp = create_buffer_resource_from_addr(wave_uniform_i64(w.lsa_ptr(rank, tmp_off))) + 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): @@ -294,7 +297,9 @@ def ar_2stage(dev_comm: Int64, win: Int64): # 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))) + 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)) diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index ae8613c9a..a12367e76 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -99,7 +99,6 @@ wave_uniform_i64, ) from .kernels_lsa import _spin_until -from .layout import MAX_WORLD #: 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 @@ -155,7 +154,7 @@ def build_sdma_phases( if queues < 1: raise ValueError(f"queues must be >= 1, got {queues}") - threads, blocks = cfg.threads, cfg.blocks + 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 @@ -164,7 +163,6 @@ def build_sdma_phases( 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 - stride_packs = blocks * threads part = cfg.packs_per_rank slice_bytes = cfg.slice_bytes @@ -244,7 +242,8 @@ def push(dev_comm: Int64, win: Int64): 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), + src_off_expr=lambda tid: fx.Int64(in_off) + + fx.Int64(tid) * fx.Int64(slice_bytes), ) @flyc.kernel(known_block_size=[threads, 1, 1]) @@ -263,9 +262,7 @@ def sdma_reduce(dev_comm: Int64, win: Int64): # 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 + cfg.recv_slot_off(rank) if reduce_self_from_recv else in_off + my_slice_off ) srcs = [ create_buffer_resource_from_addr( diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index 52c3d78c9..9ada9e0c0 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -338,7 +338,10 @@ def validate(self) -> None: "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: + 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" diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index 298ee18bd..a62181fad 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -57,9 +57,7 @@ def padded_m(m: int, world_size: int, block_m: int = DEFAULT_BLOCK_M) -> int: return (m + granule - 1) // granule * granule -def counter_chunks( - m_pad: int, world_size: int, block_m: int = DEFAULT_BLOCK_M -) -> int: +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 @@ -193,12 +191,8 @@ def _compiled(self, m: int): lane_transpose=True, ) parts = build_sdma_phases(cfg, self.rank, queues=self.sdma_queues) - c = from_gpu_ptr( - self.mem.ptr + cfg.input_off, (m, self.n), torch.bfloat16 - ) - out = from_gpu_ptr( - self.mem.ptr + cfg.output_off, (m, self.n), torch.bfloat16 - ) + c = from_gpu_ptr(self.mem.ptr + cfg.input_off, (m, self.n), torch.bfloat16) + out = from_gpu_ptr(self.mem.ptr + cfg.output_off, (m, self.n), torch.bfloat16) hit = (gemm, parts, c, out) self._cache[m] = hit return hit diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index ef2c9e96e..32170decf 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -165,7 +165,10 @@ def test_blockscale_lands_at_the_fp8_floor(m, n, k): 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 + 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 @@ -187,8 +190,16 @@ def test_blockscale_lands_at_the_fp8_floor(m, n, k): 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, + 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) @@ -198,7 +209,10 @@ def test_blockscale_lands_at_the_fp8_floor(m, n, k): ours.view(-1), sa.t().reshape(-1).contiguous(), sb.reshape(-1).contiguous(), - m, n, 0, 0, + m, + n, + 0, + 0, stream=fx.Stream(torch.cuda.current_stream()), ) torch.cuda.synchronize() @@ -237,14 +251,26 @@ def test_swap_ab_is_bitwise_identical(m, n, k): # 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, + 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, + ref.view(-1), + sa, + sb, + m, + n, + 0, + 0, stream=fx.Stream(torch.cuda.current_stream()), ) torch.cuda.synchronize() @@ -256,20 +282,33 @@ def test_swap_ab_is_bitwise_identical(m, n, k): {"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, + 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, + 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" - ) + 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(): @@ -290,7 +329,9 @@ def test_rotated_tile_order_is_a_permutation_of_the_linear_one(): 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)} + assert seen == { + (bm, bn) for bm in range(m // block_m) for bn in range(n_blocks) + } # -------------------------------------------------------------------------- @@ -309,9 +350,18 @@ def _run_bench(world_size, mode, m, n, k, extra=()): "--standalone", f"--nproc_per_node={world_size}", str(BENCH), - "--mode", mode, - "-m", str(m), "-n", str(n), "-k", str(k), - "--warmup", "1", "--iters", "3", + "--mode", + mode, + "-m", + str(m), + "-n", + str(n), + "-k", + str(k), + "--warmup", + "1", + "--iters", + "3", *extra, ] result = subprocess.run( diff --git a/tests/python/cco/test_gemm_ar_bench.py b/tests/python/cco/test_gemm_ar_bench.py index 04c2da398..b0407468a 100644 --- a/tests/python/cco/test_gemm_ar_bench.py +++ b/tests/python/cco/test_gemm_ar_bench.py @@ -68,10 +68,20 @@ def _run(world_size, mode, m, n, k, quant="blockscale", iters=11): "--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), + "--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 From 854bcbe0859100d8764114c887d8f2f62cc67783 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 14:06:19 +0000 Subject: [PATCH 51/79] gemm_ar: add window_bytes_for, and take nranks off the Communicator Two things the SGLang integration found when it was rebuilt on this API. window_bytes_for is a sizing helper the op cannot do without: the window is allocated out of the communicator's VMM reservation, so a caller has to know how big it will be *before* Communicator.init, which is before the op exists. Same arithmetic the constructor runs. Communicator exposes nranks; DevCommHandle calls the same quantity world_size. The op read world_size off the Communicator, which raised at construction time on every rank. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/op.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index a62181fad..7bab7320d 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -117,7 +117,9 @@ def __init__( ): self.comm = comm self.rank = comm.rank - self.world_size = comm.world_size + # 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 @@ -152,6 +154,27 @@ def __init__( self._cache: dict[int, tuple] = {} self._pad_in: Optional[torch.Tensor] = None + @staticmethod + def window_bytes_for( + world_size: int, *, m_max: int, n: int, block_m: int = DEFAULT_BLOCK_M + ) -> 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. + """ + m_pad = padded_m(m_max, world_size, block_m) + 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), + ) + cfg.validate() + return cfg.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) From f5d82cfa450f4328515ba03aff80d806f295ae78 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 15:00:24 +0000 Subject: [PATCH 52/79] gemm_ar: give ArConfig an fp8 wire format, regions only First step of moving the two transfer legs to fp8. This lays out where the quantised payload and its scales live; no kernel reads them yet, and `comm_dtype` defaults to "bf16", so nothing changes behaviour. Why it is worth doing, measured in-server at M=16384 N=7168 on 8 ranks: gemm 437.7us 39.8% drain 180.2 16.4% 196 MiB *1140 GB/s* reduce 43.4 3.9% gather 437.1 39.8% 196 MiB 470 GB/s gather is at the physical ceiling -- 470 GB/s is 7 xGMI links flat out -- so halving its bytes halves its time. drain's 1140 GB/s is not a bandwidth; 7 links cannot do that. It is the tell that the pushes already went out from the GEMM epilogue and drain is only waiting for the tail, which is why the scatter leg has much less to give than its byte count suggests. Scale granularity differs per leg, and not by preference: gather one scale per row. The reduce kernel owns whole rows of its slice, so an amax over all 7168 columns is free. [M] fp32 = 64 KiB. scatter one scale per (row, N-tile). 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 at all -- getting one would need a cross-block reduction, which serialises precisely what the fusion exists to overlap. Per-(row, N-tile) is what a block can compute, and at 28 scales per row instead of 1 it is finer than per-row rather than coarser. [M, N/256] fp32 = 1.75 MiB, +1.6%. `output` stays bf16 throughout: it is the consumer's tensor, not a wire format, and the reduce still accumulates in fp32. On the bf16 wire `gout` aliases `output` and every scale region is zero-sized, so the shipped layout is byte-identical -- asserted, at the 734009088 the running server logs. At M=16384 the fp8 window is 591.6 MiB against bf16's 700.0, and one put carries 14.0 MiB instead of 28.0. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/layout.py | 161 +++++++++++++++++++++++++++++- tests/python/cco/test_gemm_ar.py | 100 +++++++++++++++++++ 2 files changed, 256 insertions(+), 5 deletions(-) diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index 9ada9e0c0..f04fabe29 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -123,6 +123,16 @@ class ArConfig: #: 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 + #: What travels on the wire. ``"bf16"`` sends the payload as-is; + #: ``"fp8"`` quantises both transfer legs to e4m3 with fp32 scales, halving + #: the bytes. The reduce still accumulates in fp32 and ``output`` is still + #: bf16 -- only the wire changes. See ``scale_*`` below for the granularity, + #: which differs per leg because of where each leg's data lives. + comm_dtype: str = "bf16" + #: 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. @@ -155,6 +165,86 @@ def num_packs(self) -> int: 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_wire(self) -> bool: + return self.comm_dtype == "fp8" + + @property + def wire_elem_bytes(self) -> int: + """Bytes per element *on the wire*, as opposed to in ``output``.""" + return 1 if self.fp8_wire else self.elem_bytes + + @property + def wire_nbytes(self) -> int: + return self.num_elems * self.wire_elem_bytes + + @property + def wire_slice_bytes(self) -> int: + """One peer slice on the wire. The unit of a single SDMA put.""" + return self.slice_rows * self.n * self.wire_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_wire: + 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_wire: + 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_wire: + return 0 + return _align_up(self.m * 4, SIGNAL_ALIGN) + + @property + def gather_scale_slice_bytes(self) -> int: + if not self.fp8_wire: + return 0 + return self.slice_rows * 4 + # --- reduce-scatter partition (2-stage), sharded along m --- @property @@ -273,15 +363,58 @@ def signal_slot(self, block: int, peer: int) -> int: @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.""" + return _align_up(self.input_off + self.wire_nbytes, SIGNAL_ALIGN) + + def input_scale_slice_off(self, peer: int) -> int: + """Where peer ``peer``'s row band's scales start within ``input_scale``.""" + return ( + self.input_scale_off + + peer * self.slice_rows * self.scatter_tiles_per_row * 4 + ) + @property def output_off(self) -> int: - return self.input_off + self.nbytes + """The consumer's tensor. Always bf16 -- only the wire changes dtype.""" + return self.input_scale_off + self.scatter_scale_bytes + + @property + def output_end(self) -> int: + return _align_up(self.output_off + self.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. + """ + return self.output_end if self.fp8_wire else self.output_off + + @property + def gout_bytes(self) -> int: + return self.wire_nbytes if self.fp8_wire else 0 + + @property + def gout_scale_off(self) -> int: + return _align_up(self.gout_off + self.gout_bytes, SIGNAL_ALIGN) + + def gout_scale_slice_off(self, peer: int) -> int: + return self.gout_scale_off + peer * self.slice_rows * 4 @property def tmp_off(self) -> int: - return self.output_off + self.nbytes + # 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. + return max(self.output_end, self.gout_scale_off + self.gather_scale_bytes) @property def tmp_bytes(self) -> int: @@ -294,7 +427,7 @@ def recv_off(self) -> int: @property def recv_bytes(self) -> int: - return self.recv_slots * self.slice_bytes + return self.recv_slots * self.wire_slice_bytes def recv_slot_off(self, peer: int) -> int: """Where peer ``peer``'s contribution to *my* slice lands.""" @@ -304,11 +437,29 @@ def recv_slot_off(self, peer: int) -> int: raise IndexError( f"peer {peer} has no landing slot; recv_slots={self.recv_slots}" ) - return self.recv_off + peer * self.slice_bytes + return self.recv_off + peer * self.wire_slice_bytes + + @property + def recv_scale_off(self) -> int: + return _align_up(self.recv_off + self.recv_bytes, SIGNAL_ALIGN) + + @property + def recv_scale_bytes(self) -> int: + if not self.fp8_wire: + return 0 + return _align_up(self.recv_slots * self.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}" + ) + return self.recv_scale_off + peer * self.scatter_scale_slice_bytes @property def window_bytes(self) -> int: - return self.recv_off + self.recv_bytes + return self.recv_scale_off + self.recv_scale_bytes # --- traffic model, for reporting alongside measured time --- diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 32170decf..8e029606a 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -79,6 +79,106 @@ def test_validate_rejects_zero_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.comm_dtype == "bf16" + # 700.0056 MiB -- the size the running server logs, before any of the + # fp8 regions existed. + assert c.window_bytes == 734009088 + 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.wire_slice_bytes == c.slice_bytes + + +def test_fp8_wire_halves_the_payload_and_every_region_is_disjoint(): + c = layout.ArConfig( + world_size=8, + m=16384, + n=7168, + recv_slots=8, + counter_chunks=8, + comm_dtype="fp8", + ) + c.validate() + assert c.wire_slice_bytes * 2 == c.slice_bytes + assert c.wire_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.wire_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.ArConfig( + 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_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.ArConfig( + 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.wire_nbytes * 0.02 + + # -------------------------------------------------------------------------- # the anti-rot guard: the pinned copy still computes what the pipeline does # -------------------------------------------------------------------------- From 4f38b8be4df945d8e8a2af48e4091db8fe1e2e2e Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 15:51:56 +0000 Subject: [PATCH 53/79] gemm_ar: fp8 on the gather leg The gather is where the time is: 437us of a 1098us fused layer, moving 196 MiB at 470 GB/s, which is 7 xGMI links at their ceiling. Halving its bytes halves its time. It is also the leg where fp8 is cheap in accuracy terms: it carries *finished* values, so each element is rounded exactly once. (The scatter leg carries partial sums that then get added across every rank; it is still bf16 and is refused with a clear error rather than silently reading an fp8-sized region full of bf16.) Four pieces: quantize_gather bf16 -> e4m3 + one fp32 scale per row, on my own slice, after the reduce. gather pushes `gout` instead of `output`, at half the bytes, with the row scales riding the same queue so they cannot overtake the payload (8 KiB against 14 MiB -- about 2us). dequantize_gather the 7 slices that travelled, back to bf16. My own slice never travels and is left exactly as the reduce wrote it. parts["order"] the phase order now ships with the phases. A caller holding its own ("scatter", "reduce", "gather") list skips the two new kernels and all-reduces into zeros -- which is exactly what happened, and reads as relL2 sqrt(7/8) = 0.9354. One wave owns one row, so the per-row amax is a butterfly over ds_bpermute with no LDS and no barrier, and every lane ends up with the scale it needs. Conversion is `v_cvt_pk_fp8_f32` / `v_cvt_pk_f32_fp8`: the one-shot pk8 form is gfx1250, and `arith.truncf` to fp8 is not an option at all -- it has no LLVM lowering and fails in the translation pass rather than at trace time. Measured, standalone SDMA all-reduce, [M, 7168] on 8x MI355X: M bf16 fp8 gather 4096 295.5us 302.5us +2.4% 16384 1045.9us 928.5us -11.2% It loses at small M and wins at large, which is the fixed cost of the two conversions against a transfer that shrinks with M. Per-phase at M=16384: quantize 12.4us, dequantize 88.2us, against ~218us saved on the push. The accuracy cost is real and is worth stating plainly: **relL2 2.1e-2**, against a bf16 wire that is bitwise exact and an end-to-end 3.7-4.0e-3 for the fused wo_b in the model today. It is also a floor, not a tuning problem -- on a normal payload, per-row scaling costs 2.65e-2 and per-32 costs 2.40e-2, 9% better for 200x the scale bytes. The error is e4m3's 3-bit mantissa, and no granularity avoids it. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/ar/bench_ar.py | 44 ++- python/mori/ops/gemm_ar/__init__.py | 3 +- python/mori/ops/gemm_ar/kernels_sdma.py | 339 +++++++++++++++++++++++- python/mori/ops/gemm_ar/layout.py | 91 +++++-- tests/python/cco/test_flydsl_ar.py | 33 ++- tests/python/cco/test_gemm_ar.py | 47 +++- 6 files changed, 505 insertions(+), 52 deletions(-) diff --git a/benchmark/cco/flydsl/ar/bench_ar.py b/benchmark/cco/flydsl/ar/bench_ar.py index af9e96054..f3a1cc6f6 100644 --- a/benchmark/cco/flydsl/ar/bench_ar.py +++ b/benchmark/cco/flydsl/ar/bench_ar.py @@ -79,6 +79,19 @@ 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 @@ -162,6 +175,8 @@ def run(args) -> int: # 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() @@ -204,8 +219,11 @@ def run(args) -> int: reduce_blocks=args.reduce_blocks, ) - def launch(dc_ptr, win_h, stream=None): - for k in ("scatter", "reduce", "gather"): + # 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 @@ -237,7 +255,12 @@ def once(): diff = (out.float() - ref.float()).norm().item() denom = ref.float().norm().item() rel_l2 = diff / denom if denom else diff - validated = rel_l2 == 0.0 + # 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 @@ -249,7 +272,7 @@ def once(): 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() != 0.0: + if d.norm().item() > tol * denom: bad.append(r) print( f"[rank {rank}] VALIDATION FAILED relL2={rel_l2:.3e} " @@ -314,6 +337,19 @@ def once(): 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( diff --git a/python/mori/ops/gemm_ar/__init__.py b/python/mori/ops/gemm_ar/__init__.py index 34943ab1f..ea81b00a8 100644 --- a/python/mori/ops/gemm_ar/__init__.py +++ b/python/mori/ops/gemm_ar/__init__.py @@ -29,7 +29,7 @@ # touching the op or the kernels raises. import importlib -from .layout import ArConfig, MAX_WORLD, select_stage +from .layout import ArConfig, MAX_WORLD, ar_config, select_stage from ._shuffle import preshuffle_b _LAZY = { @@ -49,6 +49,7 @@ __all__ = [ "ArConfig", + "ar_config", "GemmAllReduceOp", "MAX_WORLD", "build_lsa_ar", diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index a12367e76..c107aa304 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -85,6 +85,8 @@ 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_SC1, @@ -100,6 +102,90 @@ ) from .kernels_lsa import _spin_until +#: 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 + +#: Wave width, and the width of the amax butterfly below (2**6). +WAVE = 64 +WAVE_LOG2 = 6 + +#: Elements one lane converts per chunk: 8 bf16 is one 16B load, and the 8 fp8 +#: it becomes is one 8B store. Both are contiguous across the wave, so both are +#: perfectly coalesced. +CHUNK_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. @@ -153,6 +239,16 @@ def build_sdma_phases( ) 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 @@ -188,15 +284,28 @@ def _signal_and_wait(w, arr_off, flag, tid): 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): + 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): @@ -217,11 +326,23 @@ def push(dev_comm: Int64, win: Int64): fx.Int64(dst_off_of_peer), win, src_off_expr(tid), - fx.Int64(slice_bytes), + 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 @@ -309,12 +430,187 @@ def sdma_reduce(dev_comm: Int64, win: Int64): ) buffer_store(packed, out, i32_off, cache_modifier=CM_CACHED) - # Gather. My reduced slice goes to the same offset in every peer's output. - 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), - ) + # --- 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)) + + 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 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.Window(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): + i32_off = (row * cfg.n + (c * WAVE) * CHUNK_ELEMS) // 2 + lane * 4 + v = ( + fx.Vector( + buffer_load( + src, + i32_off, + vec_width=4, + dtype=i32_type(), + cache_modifier=CM_CACHED, + ) + ) + .bitcast(fx.BFloat16) + .to(fx.Float32) + ) + for e in range_constexpr(CHUNK_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): + base = row * cfg.n + (c * WAVE) * CHUNK_ELEMS + v = ( + fx.Vector( + buffer_load( + src, + base // 2 + lane * 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(CHUNK_ELEMS)]) + buffer_store( + q, + dst, + (base + lane * CHUNK_ELEMS) // 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.Window(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): + base = row * cfg.n + (c * WAVE) * CHUNK_ELEMS + packed = fx.Vector( + buffer_load( + src, + (base + lane * CHUNK_ELEMS) // 4, + vec_width=2, + dtype=i32_type(), + cache_modifier=CM_SC1, + ) + ) + wide = [e * scale for e in _unpack8_fp8(packed)] + v = fx.Vector.from_elements(wide, fx.Float32) + buffer_store( + v.to(fx.BFloat16).bitcast(fx.Int32), + dst, + base // 2 + lane * 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), + ) # 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 @@ -330,12 +626,27 @@ def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): return go - return { + 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 cfg.fp8_gather: + # Bracket the gather. `quantize` has to follow the reduce and precede the + # push; `dequantize` has to follow the barrier the push ends with. + phases["quantize"] = _phase(quantize_gather, quant_blocks, QUANT_THREADS) + 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. + phases["order"] = ( + ("scatter", "reduce", "quantize", "gather", "dequantize") + if cfg.fp8_gather + else ("scatter", "reduce", "gather") + ) + return phases def build_sdma_ar(cfg, rank: int, *, queues: int = 8, signal: bool = False): @@ -346,12 +657,14 @@ def build_sdma_ar(cfg, rank: int, *, queues: int = 8, signal: bool = False): payload to every peer. """ parts = build_sdma_phases(cfg, rank, queues=queues, signal=signal) - scatter, reduce_, gather = parts["scatter"], parts["reduce"], parts["gather"] + # 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)): - scatter(dev_comm, win, stream=stream) - reduce_(dev_comm, win, stream=stream) - gather(dev_comm, win, stream=stream) + for phase in seq: + phase(dev_comm, win, stream=stream) return run, 2 diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index f04fabe29..410e8d234 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -94,6 +94,22 @@ def select_stage(world_size: int, nbytes: int) -> int: 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. @@ -123,12 +139,18 @@ class ArConfig: #: 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 - #: What travels on the wire. ``"bf16"`` sends the payload as-is; - #: ``"fp8"`` quantises both transfer legs to e4m3 with fp32 scales, halving - #: the bytes. The reduce still accumulates in fp32 and ``output`` is still - #: bf16 -- only the wire changes. See ``scale_*`` below for the granularity, - #: which differs per leg because of where each leg's data lives. - comm_dtype: str = "bf16" + #: 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" #: 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. @@ -170,23 +192,44 @@ def stage(self) -> int: # 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 fp8_wire(self) -> bool: - return self.comm_dtype == "fp8" + """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 wire_elem_bytes(self) -> int: - """Bytes per element *on the wire*, as opposed to in ``output``.""" - return 1 if self.fp8_wire else self.elem_bytes + def gather_nbytes(self) -> int: + return self.num_elems * self.gather_elem_bytes @property - def wire_nbytes(self) -> int: - return self.num_elems * self.wire_elem_bytes + 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 wire_slice_bytes(self) -> int: - """One peer slice on the wire. The unit of a single SDMA put.""" - return self.slice_rows * self.n * self.wire_elem_bytes + def gather_slice_bytes(self) -> int: + return self.slice_rows * self.n * self.gather_elem_bytes @property def slice_rows(self) -> int: @@ -221,27 +264,27 @@ def scatter_tiles_per_row(self) -> int: @property def scatter_scale_bytes(self) -> int: """fp32 scales for the whole ``[m, n]`` payload, scatter granularity.""" - if not self.fp8_wire: + 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_wire: + 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_wire: + 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_wire: + if not self.fp8_gather: return 0 return self.slice_rows * 4 @@ -369,7 +412,7 @@ def input_off(self) -> int: @property def input_scale_off(self) -> int: """Scatter-leg scales for ``input``. Zero-sized on the bf16 wire.""" - return _align_up(self.input_off + self.wire_nbytes, SIGNAL_ALIGN) + return _align_up(self.input_off + self.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``.""" @@ -396,11 +439,11 @@ def gout_off(self) -> int: 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. """ - return self.output_end if self.fp8_wire else self.output_off + return self.output_end if self.fp8_gather else self.output_off @property def gout_bytes(self) -> int: - return self.wire_nbytes if self.fp8_wire else 0 + return self.gather_nbytes if self.fp8_gather else 0 @property def gout_scale_off(self) -> int: @@ -427,7 +470,7 @@ def recv_off(self) -> int: @property def recv_bytes(self) -> int: - return self.recv_slots * self.wire_slice_bytes + return self.recv_slots * self.scatter_slice_bytes def recv_slot_off(self, peer: int) -> int: """Where peer ``peer``'s contribution to *my* slice lands.""" @@ -437,7 +480,7 @@ def recv_slot_off(self, peer: int) -> int: raise IndexError( f"peer {peer} has no landing slot; recv_slots={self.recv_slots}" ) - return self.recv_off + peer * self.wire_slice_bytes + return self.recv_off + peer * self.scatter_slice_bytes @property def recv_scale_off(self) -> int: diff --git a/tests/python/cco/test_flydsl_ar.py b/tests/python/cco/test_flydsl_ar.py index 22bb9b871..bd1d262cf 100644 --- a/tests/python/cco/test_flydsl_ar.py +++ b/tests/python/cco/test_flydsl_ar.py @@ -191,7 +191,7 @@ def test_validate_rejects_unvectorizable_payload(): # -------------------------------------------------------------------------- -def _run_bench(world_size, backend, m, n, extra_env=None): +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") @@ -214,6 +214,7 @@ def _run_bench(world_size, backend, m, n, extra_env=None): "1", "--iters", "3", + *extra, ] result = subprocess.run( command, cwd=REPO_ROOT, env=env, capture_output=True, text=True, timeout=600 @@ -240,6 +241,36 @@ def test_flydsl_ar_gpu_smoke(backend, world_size, m, n): 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 diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 8e029606a..6426ef163 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -88,7 +88,7 @@ def test_bf16_wire_layout_is_unchanged_by_the_fp8_option(): """ c = layout.ArConfig(world_size=8, m=16384, n=7168, recv_slots=8, counter_chunks=8) c.validate() - assert c.comm_dtype == "bf16" + assert c.scatter_dtype == "bf16" and c.gather_dtype == "bf16" # 700.0056 MiB -- the size the running server logs, before any of the # fp8 regions existed. assert c.window_bytes == 734009088 @@ -97,11 +97,12 @@ def test_bf16_wire_layout_is_unchanged_by_the_fp8_option(): assert c.scatter_scale_bytes == 0 assert c.gather_scale_bytes == 0 assert c.recv_scale_bytes == 0 - assert c.wire_slice_bytes == c.slice_bytes + 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.ArConfig( + c = layout.ar_config( world_size=8, m=16384, n=7168, @@ -110,13 +111,14 @@ def test_fp8_wire_halves_the_payload_and_every_region_is_disjoint(): comm_dtype="fp8", ) c.validate() - assert c.wire_slice_bytes * 2 == c.slice_bytes - assert c.wire_nbytes * 2 == c.nbytes + 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.wire_nbytes), + ("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), @@ -133,7 +135,7 @@ def test_fp8_wire_halves_the_payload_and_every_region_is_disjoint(): def test_fp8_scale_slots_are_peer_major_and_stay_inside_their_region(): - c = layout.ArConfig( + c = layout.ar_config( world_size=8, m=16384, n=7168, @@ -156,6 +158,33 @@ def test_fp8_scale_slots_are_peer_major_and_stay_inside_their_region(): ) +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. @@ -165,7 +194,7 @@ def test_scatter_scale_granularity_is_finer_than_per_row(): rather than 1 -- finer, so it cannot be less accurate than what was asked for. """ - c = layout.ArConfig( + c = layout.ar_config( world_size=8, m=16384, n=7168, @@ -176,7 +205,7 @@ def test_scatter_scale_granularity_is_finer_than_per_row(): 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.wire_nbytes * 0.02 + assert c.scatter_scale_bytes < c.scatter_nbytes * 0.02 # -------------------------------------------------------------------------- From a4ad014d67a9e81dca1cbd37a3320443d008e1da Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Fri, 11 Sep 2026 15:58:53 +0000 Subject: [PATCH 54/79] gemm_ar: expose the fp8 gather on the op and the fused benchmark `GemmAllReduceOp(gather_dtype="fp8")` and `bench_gemm_ar.py --gather-dtype fp8`. Both callers were chaining the phases from a literal ("drain", "reduce", "gather"), which on the fp8 wire skips the quantise and the dequantise and all-reduces into zeros, so the order now ships with the phases as parts["order"] / parts["fused_order"] and every caller follows it. Measured on the path this was built for -- fused-sdma, [16384, 7168] K=2048, 8x MI355X, median of 11: gather wire us relL2 bf16 1151.5 2.35e-3 fp8 1033.9 2.49e-2 -10.2% For scale, split-sdma on the same shape is 1465.9us, so fused with the fp8 gather is -29.5% against not fusing at all. The benchmark's fp8 gate is two-sided on purpose: below 4e-2 says the kernel is right, and above 5e-3 says the fp8 wire was actually taken. A one-sided gate would pass a silent fallback to bf16 as a success. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 30 +++++++++++++++--- python/mori/ops/gemm_ar/README.md | 31 +++++++++++++++++++ python/mori/ops/gemm_ar/kernels_sdma.py | 10 ++++-- python/mori/ops/gemm_ar/op.py | 22 ++++++++++--- 4 files changed, 82 insertions(+), 11 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 01274f482..d616d9a2e 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -278,6 +278,7 @@ def run(args) -> int: n=args.n, recv_slots=world_size if needs_sdma else 0, counter_chunks=chunks, + gather_dtype=args.gather_dtype, ) cfg.validate() @@ -385,11 +386,15 @@ def once(): # 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["drain" if fused else "scatter"](dc.ptr, win.handle, stream=stream) + # 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 - parts["reduce"](dc.ptr, win.handle, stream=stream) - parts["gather"](dc.ptr, win.handle, stream=stream) + for phase in order[1:]: + parts[phase](dc.ptr, win.handle, stream=stream) comm.barrier() once() @@ -433,7 +438,16 @@ def once(): # 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. - validated = rel_l2 < 3e-3 + # + # --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 @@ -649,6 +663,14 @@ def build_parser() -> argparse.ArgumentParser: "and is worth 80us (355 vs 434)", ) p.add_argument("--sdma-queues", type=int, default=8) + 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") diff --git a/python/mori/ops/gemm_ar/README.md b/python/mori/ops/gemm_ar/README.md index 80d742fce..5e5ce6625 100644 --- a/python/mori/ops/gemm_ar/README.md +++ b/python/mori/ops/gemm_ar/README.md @@ -47,6 +47,37 @@ emits). `b_scale` is `[N/128, K/128]` fp32 row-major. rounds up and `op.pad_rows` zero-extends to it. A padded row produces a zero output row, so the caller slices the result back to `m`. +## 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** | + +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 diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index c107aa304..d60eb431a 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -641,11 +641,15 @@ def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): # 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. - phases["order"] = ( - ("scatter", "reduce", "quantize", "gather", "dequantize") + tail = ( + ("reduce", "quantize", "gather", "dequantize") if cfg.fp8_gather - else ("scatter", "reduce", "gather") + else ("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 diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index 7bab7320d..94d3d5404 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -114,6 +114,7 @@ def __init__( block_m: int = DEFAULT_BLOCK_M, block_n: int = DEFAULT_BLOCK_N, sdma_queues: int = 1, + gather_dtype: str = "bf16", ): self.comm = comm self.rank = comm.rank @@ -123,6 +124,10 @@ def __init__( 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 self.m_max = padded_m(m_max, self.world_size, block_m) if not supports(self.m_max, n, k, self.world_size, block_n=block_n): @@ -156,7 +161,12 @@ def __init__( @staticmethod def window_bytes_for( - world_size: int, *, m_max: int, n: int, block_m: int = DEFAULT_BLOCK_M + world_size: int, + *, + m_max: int, + n: int, + block_m: int = DEFAULT_BLOCK_M, + gather_dtype: str = "bf16", ) -> int: """Symmetric-window bytes an op for this shape will allocate. @@ -171,6 +181,7 @@ def window_bytes_for( n=n, recv_slots=world_size, counter_chunks=counter_chunks(m_pad, world_size, block_m), + gather_dtype=gather_dtype, ) cfg.validate() return cfg.window_bytes @@ -186,6 +197,7 @@ def _make_cfg(self, m: int) -> ArConfig: n=self.n, recv_slots=self.world_size, counter_chunks=counter_chunks(m, self.world_size, self.block_m), + gather_dtype=self.gather_dtype, ) cfg.validate() return cfg @@ -281,7 +293,9 @@ def __call__( self.win.handle, stream=stream, ) - parts["drain"](self.dev_comm.ptr, self.win.handle, stream=stream) - parts["reduce"](self.dev_comm.ptr, self.win.handle, stream=stream) - parts["gather"](self.dev_comm.ptr, self.win.handle, stream=stream) + # parts["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. + for phase in parts["fused_order"]: + parts[phase](self.dev_comm.ptr, self.win.handle, stream=stream) return out From 67bffcfc7bb9f37660e611c30994d90361458817 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Sat, 12 Sep 2026 04:39:56 +0000 Subject: [PATCH 55/79] gemm_ar: widen the fp8 conversions to a full 16B access per lane The dequantise read 8 bytes per lane (one dwordx2 of fp8) while writing 16 (one dwordx4 of bf16), so the read side was half width. The kernel ran at 3.4 TB/s against roughly 8 TB/s of HBM. CHUNK_ELEMS 8 -> 16 makes the fp8 side a full dwordx4 on both kernels; the bf16 side becomes 32 bytes, issued as two. The conversion helpers still work 8 at a time -- that is `v_cvt_pk_fp8_f32`'s grouping here -- so a chunk is simply two of them. In-server, per layer, against the same bf16 baseline: phase bf16 fp8 8B fp8 16B gather 437.8 256.3 256.4 quantize - 12.7 11.8 dequantize - 85.9 61.0 wo_b layer 1101.6 994.9 980.8 -9.7% -11.0% dequantise -29%, to 4.8 TB/s. relL2 is unchanged at 2.49e-2 -- this moves bytes, not arithmetic. Note the narrowed shape constraint: n must now be a multiple of CHUNK_ELEMS*WAVE = 1024, where it was 512. 7168 satisfies it; the build raises with the requirement spelled out rather than failing later. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/kernels_sdma.py | 110 ++++++++++++++---------- 1 file changed, 64 insertions(+), 46 deletions(-) diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index d60eb431a..e75427f63 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -110,10 +110,16 @@ WAVE = 64 WAVE_LOG2 = 6 -#: Elements one lane converts per chunk: 8 bf16 is one 16B load, and the 8 fp8 -#: it becomes is one 8B store. Both are contiguous across the wave, so both are -#: perfectly coalesced. -CHUNK_ELEMS = 8 +#: 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): @@ -481,23 +487,24 @@ def quantize_gather(dev_comm: Int64, win: Int64): # pass 1: this row's amax, over the whole wave local = fx.Float32(0.0) for c in range_constexpr(chunks_per_row): - i32_off = (row * cfg.n + (c * WAVE) * CHUNK_ELEMS) // 2 + lane * 4 - v = ( - fx.Vector( - buffer_load( - src, - i32_off, - vec_width=4, - dtype=i32_type(), - cache_modifier=CM_CACHED, + 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) ) - .bitcast(fx.BFloat16) - .to(fx.Float32) - ) - for e in range_constexpr(CHUNK_ELEMS): - a = v[e] - local = local.maximumf((-a).maximumf(a)) + 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. @@ -511,26 +518,30 @@ def quantize_gather(dev_comm: Int64, win: Int64): # pass 2: scale and narrow for c in range_constexpr(chunks_per_row): - base = row * cfg.n + (c * WAVE) * CHUNK_ELEMS - v = ( - fx.Vector( - buffer_load( - src, - base // 2 + lane * 4, - vec_width=4, - dtype=i32_type(), - cache_modifier=CM_CACHED, + 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) ) - .bitcast(fx.BFloat16) - .to(fx.Float32) - ) - scaled = v * inv - q = _pack8_fp8([scaled[e] for e in range_constexpr(CHUNK_ELEMS)]) + 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( - q, + fx.Vector.from_elements(words, fx.Int32), dst, - (base + lane * CHUNK_ELEMS) // 4, + off // 4, cache_modifier=CM_CACHED, ) @@ -568,24 +579,31 @@ def dequantize_gather(dev_comm: Int64, win: Int64): ) ) for c in range_constexpr(chunks_per_row): - base = row * cfg.n + (c * WAVE) * CHUNK_ELEMS + 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, - (base + lane * CHUNK_ELEMS) // 4, - vec_width=2, + off // 4, + vec_width=4, dtype=i32_type(), cache_modifier=CM_SC1, ) ) - wide = [e * scale for e in _unpack8_fp8(packed)] - v = fx.Vector.from_elements(wide, fx.Float32) - buffer_store( - v.to(fx.BFloat16).bitcast(fx.Int32), - dst, - base // 2 + lane * 4, - cache_modifier=CM_CACHED, - ) + 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, + ) # 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; From 2579649cded7a1194b5198287ac94aa9604b47d1 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Sat, 12 Sep 2026 05:52:13 +0000 Subject: [PATCH 56/79] gemm_ar: pull the fp8 gather over LSA and widen it in the same kernel 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 -- 98 MiB a layer at the model's shape. A CU pull has those bytes in registers already, so that read simply does not happen. At [16384, 7168] K=2048 on 8 ranks, whole fused layer: gather=bf16 / sdma 1150.7us gather=fp8 / sdma 1018.9 -11.4% gather=fp8 / lsa 957.3 -16.8% so -6.1% against the SDMA gather, and the whole fp8 option is now -16.8% rather than -11.4%. relL2 is unchanged at 2.49e-2. The grid is the entire story and it is not obvious. These are xGMI reads, so the grid throttles outstanding remote requests instead of 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%. Flat from 48 to 80, steep either side. Note this is also not LSA_BLOCK_CAP's 24: that cap is for a kernel moving bf16 with no arithmetic, while this moves half the bytes and dequantises them, so it needs more waves in flight to keep the links fed. kernels_lsa.ar_2stage's two documented traps both apply and are both handled: the acquire fence after the barrier (which orders but does not invalidate, so without it we read our own stale L2 for a peer's gout), and issuing every peer's load before any store so all 7 links stream at once rather than draining one link at a time. `gather_transport="sdma"` keeps the push path, so the two are one flag apart. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 18 ++ python/mori/ops/gemm_ar/kernels_sdma.py | 281 +++++++++++++++++- python/mori/ops/gemm_ar/layout.py | 14 + python/mori/ops/gemm_ar/op.py | 7 + 4 files changed, 312 insertions(+), 8 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index d616d9a2e..3fbef600f 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -279,6 +279,7 @@ def run(args) -> int: recv_slots=world_size if needs_sdma else 0, counter_chunks=chunks, gather_dtype=args.gather_dtype, + gather_transport=args.gather_transport, ) cfg.validate() @@ -370,6 +371,7 @@ def run_gemm(stream): queues=args.sdma_queues, reduce_self_from_recv=direct_lsa, recv_uncached=direct_lsa, + fuse_quantize=args.fuse_quantize, ) if args.mode == "split-lsa": lsa_ar, _ = build_lsa_ar(cfg, rank) @@ -663,6 +665,22 @@ def build_parser() -> argparse.ArgumentParser: "and is worth 80us (355 vs 434)", ) p.add_argument("--sdma-queues", type=int, default=8) + p.add_argument( + "--fuse-quantize", + action=argparse.BooleanOptionalAction, + default=True, + help="fold the fp8 narrowing into the reduce instead of running it as " + "its own kernel. Same numbers either way; it saves reading the bf16 " + "back out of memory.", + ) + 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"), diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index e75427f63..00d4c54de 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -75,6 +75,8 @@ from __future__ import annotations +import os + import flydsl.compiler as flyc import flydsl.expr as fx @@ -106,6 +108,11 @@ #: 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 @@ -207,6 +214,7 @@ def build_sdma_phases( reduce_blocks=None, reduce_self_from_recv: bool = False, recv_uncached: bool = False, + fuse_quantize: bool = False, ): """Compile the phases separately, so the fused GEMM can reuse the tail. @@ -456,12 +464,157 @@ def sdma_reduce(dev_comm: Int64, win: Int64): 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.Window(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. @@ -605,6 +758,96 @@ def dequantize_gather(dev_comm: Int64, win: Int64): 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.Window(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 @@ -630,6 +873,11 @@ def dequantize_gather(dev_comm: Int64, win: Int64): 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. @@ -651,19 +899,36 @@ def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): "gather": _phase(gather, 1, PUSH_THREADS), } if cfg.fp8_gather: - # Bracket the gather. `quantize` has to follow the reduce and precede the - # push; `dequantize` has to follow the barrier the push ends with. + # `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["dequantize"] = _phase(dequantize_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. - tail = ( - ("reduce", "quantize", "gather", "dequantize") - if cfg.fp8_gather - else ("reduce", "gather") - ) + 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") + ) + 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. diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index 410e8d234..e25be6e9f 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -151,6 +151,15 @@ class ArConfig: #: below. ``ArConfig(comm_dtype=...)`` sets both at once. scatter_dtype: str = "bf16" gather_dtype: str = "bf16" + #: 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. @@ -200,6 +209,11 @@ def fp8_scatter(self) -> bool: 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.""" diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index 94d3d5404..25711842c 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -115,6 +115,7 @@ def __init__( block_n: int = DEFAULT_BLOCK_N, sdma_queues: int = 1, gather_dtype: str = "bf16", + gather_transport: str = "lsa", ): self.comm = comm self.rank = comm.rank @@ -128,6 +129,11 @@ def __init__( #: 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.m_max = padded_m(m_max, self.world_size, block_m) if not supports(self.m_max, n, k, self.world_size, block_n=block_n): @@ -198,6 +204,7 @@ def _make_cfg(self, m: int) -> ArConfig: recv_slots=self.world_size, counter_chunks=counter_chunks(m, self.world_size, self.block_m), gather_dtype=self.gather_dtype, + gather_transport=self.gather_transport, ) cfg.validate() return cfg From 3eaaa3f6ad286ca7bf0fc8ca7d4c5ff66abfd0dc Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Sat, 12 Sep 2026 05:53:07 +0000 Subject: [PATCH 57/79] gemm_ar: record the LSA-pull sweep and the fused-reduce loss in the README Both are results a reader needs and neither is visible from the source: the pull grid wants ~64 blocks where the local kernels want 512, and folding the narrowing into the reduce loses to the thread map it forces. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/README.md | 43 +++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/python/mori/ops/gemm_ar/README.md b/python/mori/ops/gemm_ar/README.md index 5e5ce6625..268209147 100644 --- a/python/mori/ops/gemm_ar/README.md +++ b/python/mori/ops/gemm_ar/README.md @@ -60,6 +60,49 @@ Measured, fused-sdma at `[16384, 7168]` K=2048 on 8x MI355X: | 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 From 53b418a4b4fb52cf054c8879cdf17ec8dbd51755 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Sat, 12 Sep 2026 06:36:58 +0000 Subject: [PATCH 58/79] docs: add MORI-GEMM-AR-BENCHMARK Consolidates the measurements for mori.ops.gemm_ar into the slot the repo already has for this -- alongside MORI-EP-BENCHMARK and MORI-IO-BENCHMARK, paired with the design README next to the code the way the GUIDE docs are. Everything in it is reproducible from the commands it gives, and every number is stated with the conditions it was taken under, including the 2% run-to-run spread at this shape so readers do not over-read small differences. Covers the headline (fusing -22% at M=16384, the fp8 gather a further -15%), the per-phase breakdown captured in the model rather than the benchmark, the pull grid sweep that is the single most important tuning parameter and is invisible in the source, and two sections that are easy to leave out and should not be: * what fp8 costs numerically, including the granularity table that shows ~2.5e-2 is e4m3's floor rather than something to tune; * the model-level evaluation, which is the part that decides whether the fp8 option is usable -- and which reports that it is *not detectable* at the model's output while being 6.7x worse at the layer, with the bf16-vs-bf16 control that makes that claim mean anything, and an explicit statement of what the evidence does not cover. Also records the two negative results with their causes, and the capture protocol, whose absence once produced a reported +8.3% regression that did not exist. Co-Authored-By: Claude Opus 5 (1M context) --- docs/MORI-GEMM-AR-BENCHMARK.md | 328 ++++++++++++++++++++++++++++++ docs/index.rst | 1 + python/mori/ops/gemm_ar/README.md | 4 + 3 files changed, 333 insertions(+) create mode 100644 docs/MORI-GEMM-AR-BENCHMARK.md diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md new file mode 100644 index 000000000..3b25ee3ff --- /dev/null +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -0,0 +1,328 @@ +# 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%**. + +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 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. + +**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 | vs unfused | +|---|---:|---:| +| unfused (GEMM + NCCL) | 1101.9 ms | — | +| fused, bf16 wire | 1077.2 | -2.2% | +| fused, fp8 / sdma | 1050.4 | -4.7% | +| fused, fp8 / lsa | **1048.2** | **-4.9%** | + +The three fused rows are one capture session; the unfused baseline is a separate +one, where the bf16 fused configuration measured 1071.7 ms rather than 1077.2. +Read the `-2.2%` as approximate and the differences between the fused rows as +the reliable part. + +The layer-level win is much 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/ops/gemm_ar/README.md b/python/mori/ops/gemm_ar/README.md index 268209147..3722a8102 100644 --- a/python/mori/ops/gemm_ar/README.md +++ b/python/mori/ops/gemm_ar/README.md @@ -10,6 +10,10 @@ The public API is `GemmAllReduceOp` in `op.py`. `kernels_fused.py`, 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. From 693fd5451280bc0a8664b67847fb981284e07839 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Sat, 12 Sep 2026 07:47:28 +0000 Subject: [PATCH 59/79] gemm_ar: run on the flydsl the CI image actually pins Every GPU test in the CCO job failed with AttributeError: module 'flydsl.expr' has no attribute 'to_llvm_ptr' and `NameError: name '_rocdl' is not defined` behind it. Both are the same mistake: `_compat.py` exists to span flydsl versions, and two things in it were written against 0.3.x only. The CI image pins **`flydsl<0.3`** (docker/Dockerfile.dev), because 0.3 dropped `flydsl.expr.buffer_ops` and the cco Python examples 04/05/07/08 still import it. Locally every venv here has 0.3.x, so nothing caught this -- the failures only appear at *trace* time, and `build_sdma_phases` compiles lazily, so even building every kernel under 0.2.4 comes back clean. It has to be run. * `signal_ptr` used `fx.to_llvm_ptr`, which does not exist before 0.3. It now calls `llvm.inttoptr` directly, which is byte-identical in 0.2.4 and 0.3.x -- same signature, same dialect module -- so this is one code path rather than a version probe. `to_llvm_ptr` only resolved the address space from the backend, and for Global on ROCm that is 1; the constant is named and commented rather than bare. * `_rocdl` was imported inside the 0.3 arm of the buffer_ops try/except but used unconditionally by `wave_uniform_i64` and the fp8 conversions. Hoisted to module scope; it exists in both. Nothing else needed changing. Checked directly rather than assumed: every `flydsl.expr`, `expr.rocdl` and `Vector` name this package uses is present in 0.2.4. Verified by installing 0.2.4 alongside and running the real kernels, not just building them: flydsl 0.2.4 44 passed, 1 skipped + bench 4 passed flydsl 0.3.0 44 passed, 1 skipped Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/_compat.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/python/mori/ops/gemm_ar/_compat.py b/python/mori/ops/gemm_ar/_compat.py index ee94964b4..ff60b4632 100644 --- a/python/mori/ops/gemm_ar/_compat.py +++ b/python/mori/ops/gemm_ar/_compat.py @@ -88,6 +88,12 @@ _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. @@ -106,7 +112,6 @@ def _dtype(name): 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 import rocdl as _rocdl from flydsl.expr.typing import T as _T HAS_BUFFER_OPS = False @@ -218,9 +223,6 @@ def buffer_store(data, rsrc, offset, cache_modifier=CM_CACHED): # 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 -from flydsl.expr.typing import AddressSpace as _AS # noqa: E402 -from flydsl.expr.typing import PointerType as _PT # noqa: E402 -from flydsl.expr.typing import inttoptr as _inttoptr # noqa: E402 _SYSTEM_SCOPE = "" # LLVM AMDGPU: the empty syncscope is system-wide _AGENT_SCOPE = "agent" @@ -245,15 +247,24 @@ def wave_uniform_i64(addr): 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. - ``inttoptr`` yields a ``!fly.ptr``; the llvm dialect ops need the lowered - ``!llvm.ptr``, hence the extra hop (same shape as aiter's - ``kernels/mxfp4_gemm_common.py::_global_base_ptr1``). + 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 fx.to_llvm_ptr( - _inttoptr(_PT.get(_dtype("i32"), _AS.Global), fx.Int64(addr_i64)) + return _llvm_d.inttoptr( + ir.Type.parse(_LLVM_GLOBAL_PTR), fx.Int64(addr_i64).ir_value() ) From 2673ea8cbbde6441db33227a09177f8434cd2850 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 06:36:19 +0000 Subject: [PATCH 60/79] gemm_ar: fire the gather's puts from the reduce -- measured, and it loses The third fusion candidate, and on paper the most promising of them: the gather push sends this rank's *own* reduced slice, so unlike the LSA pull it has no cross-rank dependency, and SDMA is a copy engine so it spends no CU time. It is the same shape as the GEMM's fused scatter, and it reuses that machinery -- the submit lock, and a monotonic counter per band. The traversal goes band-major so bands can complete early; under the flat grid stride every block touches every band and nothing is publishable until the kernel is nearly over. The inner map, and so the coalescing, is unchanged -- which is what made this look better than the quantize fusion, whose row-major map cost it 20us. It loses anyway, at every band count, against the unfused 1148.0us: bands 1 4 8 16 32 us 1227.3 1380.1 1630.3 2137.2 3026.7 Linear, ~61us per band, and behind already at one band. The cause is neither the traversal nor the puts. Publishing a range to a copy engine needs a system-scope release from *every block that wrote it*, so the bill is 256 blocks x bands fences at about 0.24us each. The whole reduce is 42us; the overlap cannot repay that. Isolated by running the two halves separately at 8 bands: fence, no vmcnt drain 1649.8us vmcnt drain, no fence 1164.5us so the drain is free and the fence is the entire cost. Dropping the fence is not the fix it looks like. Without it the kernel reaches 1157.9us at 4 bands -- still short of 1148.0 -- and it is racy: at 32 bands it produced relL2 1.8e-2 against the 2.35e-3 floor, differing per rank (4.5e-3, 8.1e-3, 1.6e-2). The same 32 bands *with* the fence validate exactly, which is what rules out an indexing bug and pins it on the missing release. No unsafe knob is left behind. 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 only when what is hidden is much larger than the cost of publishing it -- which is the transferable result, and the reason the two remaining fusion ideas in this area are not worth trying either. Off by default; the bf16 and fp8 paths are untouched (1147.6us unfused, 44 tests pass). ArConfig gains a 128-byte band-counter region -- the scatter's counters cannot be borrowed, since the GEMM bumps its own rank's slots too. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 10 + docs/MORI-GEMM-AR-BENCHMARK.md | 28 +++ python/mori/ops/gemm_ar/kernels_sdma.py | 182 ++++++++++++++++++ python/mori/ops/gemm_ar/layout.py | 25 ++- tests/python/cco/test_gemm_ar.py | 9 +- 5 files changed, 250 insertions(+), 4 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 3fbef600f..66ad37553 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -280,6 +280,7 @@ def run(args) -> int: counter_chunks=chunks, gather_dtype=args.gather_dtype, gather_transport=args.gather_transport, + gather_bands=args.gather_bands, ) cfg.validate() @@ -372,6 +373,7 @@ def run_gemm(stream): reduce_self_from_recv=direct_lsa, recv_uncached=direct_lsa, fuse_quantize=args.fuse_quantize, + fuse_reduce_push=args.fuse_reduce_push, ) if args.mode == "split-lsa": lsa_ar, _ = build_lsa_ar(cfg, rank) @@ -665,6 +667,14 @@ def build_parser() -> argparse.ArgumentParser: "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( + "--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, diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md index 3b25ee3ff..aa69bb6f6 100644 --- a/docs/MORI-GEMM-AR-BENCHMARK.md +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -267,6 +267,34 @@ It is the thread map: a per-row amax cannot be taken by a block holding only par 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. +**Firing the gather's puts from inside the reduce** (`fuse_reduce_push=True`) +looked like the safest fusion 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 is also exactly the shape of the GEMM's fused +scatter. It loses anyway, at every band count, against the unfused 1148.0us: + +| bands | 1 | 4 | 8 | 16 | 32 | +|---|---:|---:|---:|---:|---:| +| us | 1227.3 | 1380.1 | 1630.3 | 2137.2 | 3026.7 | + +Linear, about 61us per band. The cause is neither the traversal nor the puts: +publishing a range to a copy engine needs a **system-scope release from every +block that wrote it**, so the bill is 256 blocks x bands fences at ~0.24us each. +The whole reduce is 42us, so the overlap can never repay it. Isolated by running +the halves separately at 8 bands: fence but no `vmcnt` drain 1649.8us, drain but +no fence 1164.5us — the drain is free and the fence is all of it. + +Dropping the fence is not the fix it appears to be. Without it the kernel gets to +1157.9us at 4 bands — still short of 1148.0 — and it is racy: at 32 bands it +produced relL2 1.8e-2 against the 2.35e-3 floor, differing per rank. The same 32 +bands *with* the fence validate exactly, which rules out an indexing bug and pins +it on the missing release. + +The contrast with the GEMM's fused scatter is the useful 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 only pays when what is +hidden is much larger than the cost of publishing it. + **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` diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index 00d4c54de..563364596 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -92,6 +92,8 @@ from ._compat import ( CM_CACHED, CM_SC1, + atomic_add_u32, + atomic_store_u32, buffer_load, buffer_store, create_buffer_resource_from_addr, @@ -103,6 +105,9 @@ 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. @@ -215,6 +220,7 @@ def build_sdma_phases( reduce_self_from_recv: bool = False, recv_uncached: bool = False, fuse_quantize: bool = False, + fuse_reduce_push: bool = False, ): """Compile the phases separately, so the fused GEMM can reuse the tail. @@ -277,6 +283,7 @@ def build_sdma_phases( 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. @@ -444,6 +451,174 @@ def sdma_reduce(dev_comm: Int64, win: Int64): ) 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 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. + + **This loses, and is off by default.** At [16384, 7168] K=2048 on 8 + ranks, whole fused layer, against the unfused 1148.0us: + + bands 1 4 8 16 32 + us 1227.3 1380.1 1630.3 2137.2 3026.7 + + Linear in the band count, about 61us each, and losing already at one + band. The cause is not the traversal and not the puts: it is that + publishing a range to a copy engine needs a *system-scope release from + every block that wrote it*, so the cost is 256 blocks x bands fences, + around 0.24us apiece. The whole reduce is only 42us, so the overlap can + never repay it. + + Isolated by running the two halves separately at 8 bands: with the + fence but no vmcnt drain, 1649.8us; with the drain but no fence, + 1164.5us. The drain is free, the fence is all of it. + + Dropping the fence is not a fix even though it looks like one. Without + it the kernel reaches 1157.9us at 4 bands -- still short of 1148.0 -- + and it is racy: at 32 bands it produced relL2 1.8e-2 against the 2.35e-3 + floor, differing per rank (4.5e-3, 8.1e-3, 1.6e-2, ...). The same 32 + bands *with* the fence validate exactly, which is what rules out an + indexing bug and pins it on the missing release. + + 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, and here it is amortised against 42us of + bandwidth-saturated reduce. The mechanism is only worth it when what is + being hidden is much larger than the cost of publishing it. + """ + tid = fx.thread_idx.x + bid = fx.block_idx.x + w = cco.Window(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)) + ) + + 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_CACHED) + + # Retire this block's stores, then release them where a copy engine + # can see them. `wait_barrier(0)` is s_waitcnt vmcnt(0) + s_barrier: + # a bare s_barrier would only prove the stores were *issued*, which + # is not enough to hand the range to SDMA. + # Publishing a band to a copy engine takes both of these, and the + # second is where all the cost is -- see the class docstring. + # `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. The system fence then releases them somewhere the + # engine can read. Dropping the fence is *not* an option: it is + # nearly free in wall time and it is what makes the result correct. + wait_barrier(0) + raw_cco.cco_system_fence(fx.Int32(0)) + if tid == fx.Int32(0): + ctr = signal_ptr( + fx.Int64(w.lsa_ptr(rank, gather_counter_off)) + + 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( + fx.Int64(w.lsa_ptr(rank, lock_off)) + + 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 @@ -898,6 +1073,11 @@ def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): "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 @@ -927,6 +1107,8 @@ def go(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): 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 diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index e25be6e9f..11f1c305c 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -151,6 +151,11 @@ class ArConfig: #: 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 @@ -354,7 +359,9 @@ def blocks(self) -> int: @property def signal_bytes(self) -> int: - return self.lock_off + self.lock_bytes + # Everything before `input`: the barrier arrays, the scatter's tile + # counters, the submit locks, and the gather's band counters. + return self.gather_counter_off + self.gather_counter_bytes @property def start_off(self) -> int: @@ -403,6 +410,22 @@ def lock_off(self) -> int: def lock_bytes(self) -> int: return _align_up(self.world_size * 4, SIGNAL_ALIGN) + @property + def gather_counter_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_bytes(self) -> int: + return _align_up(self.gather_bands * 4, 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})") diff --git a/tests/python/cco/test_gemm_ar.py b/tests/python/cco/test_gemm_ar.py index 6426ef163..882d7c0fc 100644 --- a/tests/python/cco/test_gemm_ar.py +++ b/tests/python/cco/test_gemm_ar.py @@ -89,9 +89,12 @@ def test_bf16_wire_layout_is_unchanged_by_the_fp8_option(): 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" - # 700.0056 MiB -- the size the running server logs, before any of the - # fp8 regions existed. - assert c.window_bytes == 734009088 + # 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 From 92a379c88857dde4cae401d254b37ecc7494c14d Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 07:04:32 +0000 Subject: [PATCH 61/79] gemm_ar: pay the publish release in the store policy, not in a fence Correcting the previous commit, which blamed the whole cost on "the system fence" and concluded the mechanism could not work. The fence was the expensive *way* to pay for the release, not the only one. Handing a range to a copy engine does need a release -- the engine reads over the fabric, not through a CU's cache, so `s_waitcnt vmcnt(0)` alone only retires the stores as far as this XCD's L2. But it can be paid two ways, and they differ by an order of magnitude: publish="fence" leave the bytes in L2, release to system scope after. L2 writeback work, charged once per block per band: 256 x bands, ~61us per band. publish="writethrough" store with sc0+sc1 so they never stop in L1 or L2, and the waitcnt above *is* the release. At [16384, 7168] K=2048 on 8 ranks, whole fused layer: 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 470us apart at 8 bands. The write-through store is free in its own right, measured separately: applying that policy to the plain unfused reduce moves it 1153.4 -> 1148.7us. This output is written once and nothing local reads it again before the gather, so holding it in L2 was buying nothing to begin with. It still does not win. Against 1148.7us unfused the best is 1157.2 at 4 bands, and the two remaining effects nearly cancel: bands=1 costs 13us for the mechanism with no overlap to show for it (per-band 256-block wait_barrier, counter atomic, the elected block's locked puts), and four bands buy about 5us back before the sync cost takes over. So: parity, not a win, and it stays off. Also corrected: the earlier claim that dropping the release "is not an option" stands, but for a sharper reason than given. Cached stores with no fence reach 1157.9us and are racy -- relL2 1.8e-2 at 32 bands, varying per rank -- while the same 32 bands are *exact* under either correct publish mode. That is what rules out an indexing bug and pins it on the missing release rather than on the fence specifically. Default is publish="writethrough"; no unsafe mode is reachable. bf16 and fp8 paths untouched, 44 tests pass. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 4 + docs/MORI-GEMM-AR-BENCHMARK.md | 57 +++++---- python/mori/ops/gemm_ar/kernels_sdma.py | 116 ++++++++++++------ 3 files changed, 117 insertions(+), 60 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 66ad37553..33b7afa72 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -374,6 +374,7 @@ def run_gemm(stream): 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) @@ -668,6 +669,9 @@ def build_parser() -> argparse.ArgumentParser: ) 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, diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md index aa69bb6f6..be9ff9167 100644 --- a/docs/MORI-GEMM-AR-BENCHMARK.md +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -268,31 +268,44 @@ of a row, so fusing forces one-wave-per-row, where `sdma_reduce` walks packs wit a flat grid stride and streams a block through all 8 source slices at once. **Firing the gather's puts from inside the reduce** (`fuse_reduce_push=True`) -looked like the safest fusion 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 is also exactly the shape of the GEMM's fused -scatter. It loses anyway, at every band count, against the unfused 1148.0us: +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 | |---|---:|---:|---:|---:|---:| -| us | 1227.3 | 1380.1 | 1630.3 | 2137.2 | 3026.7 | - -Linear, about 61us per band. The cause is neither the traversal nor the puts: -publishing a range to a copy engine needs a **system-scope release from every -block that wrote it**, so the bill is 256 blocks x bands fences at ~0.24us each. -The whole reduce is 42us, so the overlap can never repay it. Isolated by running -the halves separately at 8 bands: fence but no `vmcnt` drain 1649.8us, drain but -no fence 1164.5us — the drain is free and the fence is all of it. - -Dropping the fence is not the fix it appears to be. Without it the kernel gets to -1157.9us at 4 bands — still short of 1148.0 — and it is racy: at 32 bands it -produced relL2 1.8e-2 against the 2.35e-3 floor, differing per rank. The same 32 -bands *with* the fence validate exactly, which rules out an indexing bug and pins -it on the missing release. - -The contrast with the GEMM's fused scatter is the useful 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 only pays when what is +| `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. **A CK-shaped 4-wave GEMM**, chasing a 22% gap against CK's block-scale kernel diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index 563364596..f4dfe4338 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -91,6 +91,7 @@ from ._compat import ( CM_CACHED, + CM_SC0_SC1, CM_SC1, atomic_add_u32, atomic_store_u32, @@ -221,6 +222,7 @@ def build_sdma_phases( 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. @@ -470,6 +472,8 @@ def sdma_reduce(dev_comm: Int64, win: Int64): # 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 " @@ -493,35 +497,54 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): difference is when the bytes start moving. The following `gather` phase becomes drain-only. - **This loses, and is off by default.** At [16384, 7168] K=2048 on 8 - ranks, whole fused layer, against the unfused 1148.0us: - - bands 1 4 8 16 32 - us 1227.3 1380.1 1630.3 2137.2 3026.7 - - Linear in the band count, about 61us each, and losing already at one - band. The cause is not the traversal and not the puts: it is that - publishing a range to a copy engine needs a *system-scope release from - every block that wrote it*, so the cost is 256 blocks x bands fences, - around 0.24us apiece. The whole reduce is only 42us, so the overlap can - never repay it. - - Isolated by running the two halves separately at 8 bands: with the - fence but no vmcnt drain, 1649.8us; with the drain but no fence, - 1164.5us. The drain is free, the fence is all of it. - - Dropping the fence is not a fix even though it looks like one. Without - it the kernel reaches 1157.9us at 4 bands -- still short of 1148.0 -- - and it is racy: at 32 bands it produced relL2 1.8e-2 against the 2.35e-3 - floor, differing per rank (4.5e-3, 8.1e-3, 1.6e-2, ...). The same 32 - bands *with* the fence validate exactly, which is what rules out an - indexing bug and pins it on the missing release. + **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, and here it is amortised against 42us of - bandwidth-saturated reduce. The mechanism is only worth it when what is - being hidden is much larger than the cost of publishing it. + 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 @@ -572,21 +595,38 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): if elem_dtype is fx.Float32 else acc.to(elem_dtype).bitcast(fx.Int32) ) - buffer_store(packed, out, i32_off, cache_modifier=CM_CACHED) - - # Retire this block's stores, then release them where a copy engine - # can see them. `wait_barrier(0)` is s_waitcnt vmcnt(0) + s_barrier: - # a bare s_barrier would only prove the stores were *issued*, which - # is not enough to hand the range to SDMA. - # Publishing a band to a copy engine takes both of these, and the - # second is where all the cost is -- see the class docstring. + 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. The system fence then releases them somewhere the - # engine can read. Dropping the fence is *not* an option: it is - # nearly free in wall time and it is what makes the result correct. + # 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) - raw_cco.cco_system_fence(fx.Int32(0)) + if const_expr(publish == "fence"): + raw_cco.cco_system_fence(fx.Int32(0)) if tid == fx.Int32(0): ctr = signal_ptr( fx.Int64(w.lsa_ptr(rank, gather_counter_off)) From e5f712ed38efe7bbcf2ce344dc30eb40e700cd8f Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 07:22:44 +0000 Subject: [PATCH 62/79] gemm_ar: hoist the loop-invariant lsa_ptr calls out of the band loop `Window.lsa_ptr` is not free and not foldable. It is an opaque extern call that loads two fields out of the window descriptor and combines them: cco_lsa_ptr(window, peer, offset): stride = w->stride4G << 32 // load return w->winBase + peer*stride + offset // load Because the compiler cannot see through the call it cannot CSE or hoist these itself, so every call re-reads both fields. The counter and lock base pointers in `sdma_reduce_push` are loop-invariant and are now taken once instead of once per (block, band). **This does not measure.** bands=4 goes 1157.2 -> 1154.8/1152.3/1157.0 over three repeats, i.e. inside the run-to-run band. The call count at these sites is simply too small to matter: the counter pointer is thread 0 only, once per band, and the lock pointer only on the elected block. It is committed because loop-invariant work does not belong in a loop and the reason the compiler cannot move it is worth writing down, not because it bought anything. Checked the rest of the package while here: every call in a *hot* path is already hoisted, deliberately. `sdma_reduce`, `lsa_pull_gather` and `kernels_lsa.ar_2stage` all build their buffer descriptors at kernel scope before the pack loop -- which, together with `wave_uniform_i64`, is the established idiom here. The only other in-loop-looking sites, `kernels_fused` 1503/1522, are per-block epilogue code with the same negligible call count. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/kernels_sdma.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index f4dfe4338..db772b30c 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -568,6 +568,16 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): wave_uniform_i64(w.lsa_ptr(rank, out_off + my_slice_off)) ) + # `lsa_ptr` is an opaque extern call that loads two fields out of the + # window descriptor (`winBase`, `stride4G`) and combines them. The + # compiler cannot see through it, so it cannot hoist or CSE the calls + # itself -- every call re-reads both. Both of these bases are + # loop-invariant, so they are taken once here rather than once per + # (block, band). This is the same reason the source/destination buffer + # descriptors above are built at kernel scope and not in the pack loop. + 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 @@ -628,10 +638,7 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): if const_expr(publish == "fence"): raw_cco.cco_system_fence(fx.Int32(0)) if tid == fx.Int32(0): - ctr = signal_ptr( - fx.Int64(w.lsa_ptr(rank, gather_counter_off)) - + fx.Int64(band) * fx.Int64(4) - ) + 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 @@ -641,10 +648,7 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): 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( - fx.Int64(w.lsa_ptr(rank, lock_off)) - + fx.Int64(dest) * fx.Int64(4) - ) + lock = signal_ptr(lock_base + fx.Int64(dest) * fx.Int64(4)) _acquire_peer_lock(lock) sdma.put( dest, From 2bc532be6bc9047057650b081418bd97064ffdd6 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 07:51:31 +0000 Subject: [PATCH 63/79] cco: add Window.lsa_geometry / lsa_ptr_at, and measure that it changes nothing `cco_lsa_ptr` is `winBase + peer*stride + offset` and reads both fields out of the window descriptor on every call. It is opaque to FlyDSL's tracer, and although `CCO_DEV` 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. `lsa_geometry()` takes base and stride once; `lsa_ptr_at()` then computes addresses as ordinary DSL arithmetic the compiler can fold, which for the common case of a Python-int peer and offset collapses to one add. No C++ change is needed for this: peer=1, offset=0 is exactly `base + stride`, so two existing calls yield both values. **It does not measure, and the reason is worth keeping.** The LSA pull gather builds 15 descriptors and was the best target in the package; converting it moves a 230us kernel by nothing -- 973.2us before, then 972.5 / 975.8 / 976.1. Every `lsa_ptr` call in `gemm_ar` is already at kernel scope, because building descriptors before the hot loop is the established idiom there (the same idiom `wave_uniform_i64` belongs to), so what this removes is on the order of 26 loads per launch. Committed anyway, with the measurement in the docstring, because the primitive is the right one to reach for when an address really is needed per iteration -- and because the next person to notice that `lsa_ptr` looks expensive should find the answer rather than re-run the experiment. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/cco/device/flydsl/handles.py | 41 ++++++++++++++++++++++++ python/mori/ops/gemm_ar/kernels_sdma.py | 36 ++++++++++++++------- 2 files changed, 65 insertions(+), 12 deletions(-) diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index 30dce38d1..a5c6911d1 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -110,6 +110,47 @@ def lsa_ptr(self, peer_lsa_rank, offset=0): """Peer's LSA-accessible VA inside this window (uint64), for direct load/store.""" return raw.cco_lsa_ptr(self.handle, peer_lsa_rank, offset) + def lsa_geometry(self): + """Read the window's base and peer stride once, for reuse. + + ``cco_lsa_ptr`` is ``winBase + peer * stride + offset``, and it loads + *both* fields out of the window descriptor on every call. FlyDSL emits + it as an extern call so the tracer cannot fold it, and although it is + ``always_inline`` -- so LLVM does see the two loads -- a kernel that + stores through addresses derived from that same base gives LLVM no way + to prove the loads are not clobbered, so it must reload them. + + Taking the geometry once turns every later address into ordinary DSL + arithmetic over SSA values, which the compiler *can* fold. For the + common case here, where the peer and the offset are both Python ints, + it collapses to a single add against a value already in a register. + + Costs one extra call here (peer=1, offset=0 is exactly + ``base + stride``) to remove one from every use site. Returns + ``(base, stride)`` for :meth:`lsa_ptr_at`. + + **It does not measure on mori's own kernels, and that is expected.** + Swapping the 15 descriptors in ``gemm_ar``'s LSA pull gather over to it + moves a 230us kernel by nothing: 973.2us before, 972.5/975.8/976.1 + after. Every ``lsa_ptr`` call in that package is already at kernel + scope -- descriptors get built once before the hot loop, which is the + established idiom there -- so what this removes is on the order of 26 + loads per launch. + + Where it *would* pay is a kernel that needs an address inside a loop, + which is exactly the case the idiom above exists to avoid. Use it when + the call count is large or per-iteration; do not expect it to show up + otherwise. + """ + base = fx.Int64(raw.cco_lsa_ptr(self.handle, 0, 0)) + stride = fx.Int64(raw.cco_lsa_ptr(self.handle, 1, 0)) - base + return base, stride + + @staticmethod + def lsa_ptr_at(base, stride, peer_lsa_rank, offset=0): + """:meth:`lsa_ptr` from a cached :meth:`lsa_geometry` -- no extern call.""" + return base + fx.Int64(peer_lsa_rank) * stride + fx.Int64(offset) + @cco_struct class Gda: diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index db772b30c..02dfb42de 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -706,15 +706,22 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): ), ) - 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)) - ) + # These build one descriptor per (peer, region), and `lsa_pull_gather` + # wants 15 of them. `lsa_ptr` reads the window's base and stride on + # every call, so the geometry is taken once and the rest is arithmetic + # -- see Window.lsa_geometry. + def _geom(w): + return w.lsa_geometry() + + def _row_addr_at(w, peer, byte_off, geom=None): + if geom is None: + addr = w.lsa_ptr(peer, byte_off) + else: + addr = w.lsa_ptr_at(geom[0], geom[1], peer, byte_off) + return create_buffer_resource_from_addr(wave_uniform_i64(addr)) + + def _row_addr(w, byte_off, geom=None): + return _row_addr_at(w, rank, byte_off, geom) @flyc.kernel(known_block_size=[QUANT_THREADS, 1, 1]) def sdma_reduce_quant(dev_comm: Int64, win: Int64): @@ -1012,17 +1019,22 @@ def lsa_pull_gather(dev_comm: Int64, win: Int64): raw_cco.cco_system_fence(fx.Int32(0)) + # 15 descriptors below, so the window geometry is read once here + # rather than once per descriptor. + geom = _geom(w) # 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) + _row_addr_at(w, p, cfg.gout_off + p * cfg.slice_rows * cfg.n, geom) for p in peers ] - scas = [_row_addr_at(w, p, cfg.gout_scale_slice_off(p)) for p in peers] + scas = [ + _row_addr_at(w, p, cfg.gout_scale_slice_off(p), geom) 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) + out = _row_addr(w, out_off, geom) for row in range( bid * QUANT_WAVES + wave, slice_rows, pull_blocks * QUANT_WAVES From 046391f6363e3e5a2d81665284f6c3b936ae2dad Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 08:36:40 +0000 Subject: [PATCH 64/79] cco: global_load the window geometry, and convert every lsa_ptr in kernels_sdma Does what was asked, and the mechanism is confirmed at the IR level. It still does not move any measurement, which is now the third independent test of this idea. C++: `cco_lsa_win_base` / `cco_lsa_stride` read through an `address_space(1)` pointer. `cco_lsa_ptr` casts its uint64 to a *generic* pointer, and a generic access has to be a `flat_load` -- the compiler cannot rule out LDS -- so it is counted against `lgkmcnt` as well as `vmcnt`, giving any following `s_waitcnt` one more counter to wait on. The disassembled bitcode shows the difference plainly: cco_lsa_win_base: inttoptr i64 %0 to ptr addrspace(1) load ptr, ptr addrspace(1) %2 <- global, 1 load cco_lsa_ptr: inttoptr i64 %0 to ptr load i32, ptr %5 / load ptr, ptr %4 <- flat, 2 loads FlyDSL: `Window.lsa_geometry()` now goes through those two instead of deriving the stride from `lsa_ptr(1,0) - lsa_ptr(0,0)`. kernels_sdma.py: 21 `lsa_ptr` call sites -> 0. Every kernel takes the geometry once and computes addresses with `lsa_ptr_at`, including `_signal_and_wait`, which is the one site whose peer is a *runtime* lane index and so the only one where the address genuinely varies per lane. Measured, [16384, 7168] K=2048 on 8 ranks, all three wire configs, each against where it sat before: bf16/sdma 1148.0 -> 1142.5 fp8/sdma 1031 -> 1031.4 fp8/lsa 973.2 -> 971.9 i.e. nothing, inside the ~2% band. Consistent with the two earlier attempts: hoisting the band-loop calls, and the pure-Python geometry. The reason has been the same every time -- every call is at kernel scope and there are a handful per launch, so this removes on the order of tens of loads from kernels that run for hundreds of microseconds. kernels_fused.py and kernels_lsa.py are deliberately left alone. Same shape of call site (per-block epilogue and per-block barrier), so the same zero, and kernels_fused is the file with the documented races -- not worth touching for a change measured at zero three times over. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/cco/device/flydsl/_bindings.py | 3 + python/mori/cco/device/flydsl/handles.py | 12 +- python/mori/cco/device/ops.py | 7 ++ python/mori/ops/gemm_ar/kernels_sdma.py | 126 +++++++++++---------- src/cco/device/cco_device_wrapper.cpp | 25 ++++ 5 files changed, 110 insertions(+), 63 deletions(-) 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 a5c6911d1..271bfd969 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -125,8 +125,12 @@ def lsa_geometry(self): common case here, where the peer and the offset are both Python ints, it collapses to a single add against a value already in a register. - Costs one extra call here (peer=1, offset=0 is exactly - ``base + stride``) to remove one from every use site. Returns + The two reads 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 is counted against ``lgkmcnt`` as well as ``vmcnt``, and any + following ``s_waitcnt`` has one more counter to wait on. Returns ``(base, stride)`` for :meth:`lsa_ptr_at`. **It does not measure on mori's own kernels, and that is expected.** @@ -142,8 +146,8 @@ def lsa_geometry(self): the call count is large or per-iteration; do not expect it to show up otherwise. """ - base = fx.Int64(raw.cco_lsa_ptr(self.handle, 0, 0)) - stride = fx.Int64(raw.cco_lsa_ptr(self.handle, 1, 0)) - base + base = fx.Int64(raw.cco_lsa_win_base(self.handle)) + stride = fx.Int64(raw.cco_lsa_stride(self.handle)) return base, stride @staticmethod 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/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index 02dfb42de..258944a0c 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -294,17 +294,25 @@ def build_sdma_phases( I32_PER_PACK = 4 # a 16B pack always moves as 4 x i32, whatever the payload - def _next_flag(w, bid): + def _next_flag(w, bid, geom): """``_flag[block] + 1`` -- monotonic, never reset (graph-replay safe).""" - base = fx.Int64(w.lsa_ptr(rank, flag_off)) + fx.Int64(bid) * fx.Int64(4) + base = w.lsa_ptr_at(geom[0], geom[1], 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)) + def _signal_and_wait(w, arr_off, flag, tid, geom): + """Row-0 barrier: publish to peer ``tid``, then wait on ``tid``'s slot. + + The peer here is a *runtime* lane index, so this is the one site where + the address genuinely varies per lane -- and therefore the one where + computing it from a cached geometry, rather than calling `lsa_ptr`, + replaces a descriptor load pair with a multiply-add. + """ + peer_arr = w.lsa_ptr_at(geom[0], geom[1], 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)) + self_arr = w.lsa_ptr_at(geom[0], geom[1], rank, arr_off) _spin_until(signal_ptr(self_arr + fx.Int64(tid) * fx.Int64(4)), flag) def _push_kernel( @@ -336,7 +344,8 @@ def push(dev_comm: Int64, win: Int64): w = cco.Window(win) sdma = cco.DevComm(dev_comm).sdma() - flag, flag_rsrc = _next_flag(w, 0) + geom = w.lsa_geometry() + flag, flag_rsrc = _next_flag(w, 0, geom) if tid < ws: if tid != rank: if const_expr(pushes): @@ -373,7 +382,7 @@ def push(dev_comm: Int64, win: Int64): # 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) + _signal_and_wait(w, arr_off, flag, tid, geom) fgpu.barrier() if tid == 0: local_store_u32(flag_rsrc, flag) @@ -405,22 +414,21 @@ def sdma_reduce(dev_comm: Int64, win: Int64): # 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. + # Nine descriptors from one window: take the geometry once. + geom = w.lsa_geometry() + + def _rsrc(byte_off): + return create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr_at(geom[0], geom[1], rank, byte_off)) + ) + 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) + srcs = [_rsrc(self_off)] + [ + _rsrc(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)) - ) + out = _rsrc(out_off + my_slice_off) gtid = bid * threads + tid for pk in range(gtid, part, red_stride): @@ -551,22 +559,21 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): w = cco.Window(win) sdma = cco.DevComm(dev_comm).sdma() + # Nine descriptors from one window: take the geometry once. + geom = w.lsa_geometry() + + def _rsrc(byte_off): + return create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr_at(geom[0], geom[1], rank, byte_off)) + ) + 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) + srcs = [_rsrc(self_off)] + [ + _rsrc(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)) - ) + out = _rsrc(out_off + my_slice_off) # `lsa_ptr` is an opaque extern call that loads two fields out of the # window descriptor (`winBase`, `stride4G`) and combines them. The @@ -575,8 +582,8 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): # loop-invariant, so they are taken once here rather than once per # (block, band). This is the same reason the source/destination buffer # descriptors above are built at kernel scope and not in the pack loop. - ctr_base = fx.Int64(w.lsa_ptr(rank, gather_counter_off)) - lock_base = fx.Int64(w.lsa_ptr(rank, lock_off)) + ctr_base = w.lsa_ptr_at(geom[0], geom[1], rank, gather_counter_off) + lock_base = w.lsa_ptr_at(geom[0], geom[1], rank, lock_off) gtid = bid * threads + tid for band in range_constexpr(cfg.gather_bands): @@ -708,19 +715,14 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): # These build one descriptor per (peer, region), and `lsa_pull_gather` # wants 15 of them. `lsa_ptr` reads the window's base and stride on - # every call, so the geometry is taken once and the rest is arithmetic - # -- see Window.lsa_geometry. - def _geom(w): - return w.lsa_geometry() - - def _row_addr_at(w, peer, byte_off, geom=None): - if geom is None: - addr = w.lsa_ptr(peer, byte_off) - else: - addr = w.lsa_ptr_at(geom[0], geom[1], peer, byte_off) - return create_buffer_resource_from_addr(wave_uniform_i64(addr)) - - def _row_addr(w, byte_off, geom=None): + # every call, so callers take the geometry once and the rest is + # arithmetic -- see Window.lsa_geometry. + def _row_addr_at(w, peer, byte_off, geom): + return create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr_at(geom[0], geom[1], peer, byte_off)) + ) + + def _row_addr(w, byte_off, geom): return _row_addr_at(w, rank, byte_off, geom) @flyc.kernel(known_block_size=[QUANT_THREADS, 1, 1]) @@ -765,18 +767,20 @@ def sdma_reduce_quant(dev_comm: Int64, win: Int64): w = cco.Window(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) + geom = w.lsa_geometry() 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) + srcs = [_row_addr(w, self_off, geom)] + [ + _row_addr(w, cfg.recv_slot_off((rank + j) % ws), geom) + 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)) + out = _row_addr(w, out_off + my_slice_off, geom) + gdst = _row_addr(w, cfg.gout_off + rank * cfg.slice_rows * cfg.n, geom) + gsca = _row_addr(w, cfg.gout_scale_slice_off(rank), geom) for row in range( bid * QUANT_WAVES + wave, slice_rows, quant_blocks * QUANT_WAVES @@ -855,10 +859,11 @@ def quantize_gather(dev_comm: Int64, win: Int64): w = cco.Window(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) + geom = w.lsa_geometry() - 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)) + src = _row_addr(w, out_off + my_slice_off, geom) + dst = _row_addr(w, cfg.gout_off + rank * cfg.slice_rows * cfg.n, geom) + sca = _row_addr(w, cfg.gout_scale_slice_off(rank), geom) for row in range( bid * QUANT_WAVES + wave, slice_rows, quant_blocks * QUANT_WAVES @@ -935,12 +940,15 @@ def dequantize_gather(dev_comm: Int64, win: Int64): w = cco.Window(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) + geom = w.lsa_geometry() 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) + # All three are in *my own* window: the SDMA push already + # delivered peer `peer`'s slice and scales here. + src = _row_addr(w, cfg.gout_off + peer * cfg.slice_rows * cfg.n, geom) + sca = _row_addr(w, cfg.gout_scale_slice_off(peer), geom) + dst = _row_addr(w, out_off + peer * cfg.slice_rows * cfg.n * 2, geom) for row in range( bid * QUANT_WAVES + wave, slice_rows, quant_blocks * QUANT_WAVES @@ -1021,7 +1029,7 @@ def lsa_pull_gather(dev_comm: Int64, win: Int64): # 15 descriptors below, so the window geometry is read once here # rather than once per descriptor. - geom = _geom(w) + geom = w.lsa_geometry() # 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 = [ 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 From 81fe2fea9d08e1e5d6d49ec1f5587c84319ab7b8 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 09:53:51 +0000 Subject: [PATCH 65/79] cco: put the window-geometry caching in a class, not in every call site 046391f6 threaded a `geom` tuple through every helper in kernels_sdma -- 21 call sites and three changed signatures -- to hoist the two descriptor loads out of `lsa_ptr`. Same optimisation, wrong seam. `CachedWindow` reads the geometry in its constructor and computes `lsa_ptr` from it, so the kernels change by one line each: w = cco.Window(win) -> w = cco.CachedWindow(win) and nothing else. kernels_sdma.py goes back to the pre-046391f6 text; the 14 lines that remain are the 7 swaps. `Window` keeps `lsa_ptr` unchanged, and `lsa_geometry`/`lsa_ptr_at` are gone -- the class is the API now. The C++ side from 046391f6 stands: `cco_lsa_win_base` / `cco_lsa_stride` take an `address_space(1)` pointer, so each is one `global_load` rather than the `flat_load` that `cco_lsa_ptr`'s generic cast forces, and a `flat_load` counts against `lgkmcnt` as well as `vmcnt`. It still does not measure. A/B/A' in one session, m=16384 n=7168 k=2048, 8x MI355X, max_rank_time_us: wire Window CachedWindow Window again bf16/sdma 1114.5 1111.4 1115.4 fp8/sdma 983.3 983.6 983.6 fp8/lsa 926.6 924.3 926.0 B-A is smaller than A'-A. This is the fourth independent null for the idea, after the band-loop hoist, the Python geometry API and the global_load conversion. Note for anyone re-measuring: this box drifts ~4% between mornings, across every configuration at once. Today's absolute numbers are all ~4% under 046391f6's on identical code, which is why the A/B is the only line worth reading. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/cco/device/flydsl/__init__.py | 12 +- python/mori/cco/device/flydsl/handles.py | 102 ++++++++------ python/mori/ops/gemm_ar/kernels_sdma.py | 158 +++++++++------------- 3 files changed, 136 insertions(+), 136 deletions(-) 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/handles.py b/python/mori/cco/device/flydsl/handles.py index 271bfd969..ade444678 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -110,50 +110,64 @@ def lsa_ptr(self, peer_lsa_rank, offset=0): """Peer's LSA-accessible VA inside this window (uint64), for direct load/store.""" return raw.cco_lsa_ptr(self.handle, peer_lsa_rank, offset) - def lsa_geometry(self): - """Read the window's base and peer stride once, for reuse. - - ``cco_lsa_ptr`` is ``winBase + peer * stride + offset``, and it loads - *both* fields out of the window descriptor on every call. FlyDSL emits - it as an extern call so the tracer cannot fold it, and although it is - ``always_inline`` -- so LLVM does see the two loads -- a kernel that - stores through addresses derived from that same base gives LLVM no way - to prove the loads are not clobbered, so it must reload them. - - Taking the geometry once turns every later address into ordinary DSL - arithmetic over SSA values, which the compiler *can* fold. For the - common case here, where the peer and the offset are both Python ints, - it collapses to a single add against a value already in a register. - - The two reads 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 is counted against ``lgkmcnt`` as well as ``vmcnt``, and any - following ``s_waitcnt`` has one more counter to wait on. Returns - ``(base, stride)`` for :meth:`lsa_ptr_at`. - - **It does not measure on mori's own kernels, and that is expected.** - Swapping the 15 descriptors in ``gemm_ar``'s LSA pull gather over to it - moves a 230us kernel by nothing: 973.2us before, 972.5/975.8/976.1 - after. Every ``lsa_ptr`` call in that package is already at kernel - scope -- descriptors get built once before the hot loop, which is the - established idiom there -- so what this removes is on the order of 26 - loads per launch. - - Where it *would* pay is a kernel that needs an address inside a loop, - which is exactly the case the idiom above exists to avoid. Use it when - the call count is large or per-iteration; do not expect it to show up - otherwise. - """ - base = fx.Int64(raw.cco_lsa_win_base(self.handle)) - stride = fx.Int64(raw.cco_lsa_stride(self.handle)) - return base, stride - - @staticmethod - def lsa_ptr_at(base, stride, peer_lsa_rank, offset=0): - """:meth:`lsa_ptr` from a cached :meth:`lsa_geometry` -- no extern call.""" - return base + fx.Int64(peer_lsa_rank) * stride + fx.Int64(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. + + **It does not measure on mori's own kernels, and that is expected.** + A/B/A' at ``m=16384 n=7168 k=2048`` on 8x MI355X, us, lower is better:: + + wire Window CachedWindow Window again + bf16/sdma 1114.5 1111.4 1115.4 + fp8/sdma 983.3 983.6 983.6 + fp8/lsa 926.6 924.3 926.0 + + The difference is smaller than the drift between the two Window runs. Every + ``lsa_ptr`` call in ``gemm_ar`` is already at kernel scope -- descriptors get + built once before the hot loop, which is the established idiom -- so what + this removes is on the order of tens of loads per launch against a kernel + that runs for a millisecond. Reach for it when a kernel needs addresses *per + iteration*; do not expect it to show up otherwise. + + Measure A/B in one session if you revisit this. Absolute numbers on this box + move ~4% between mornings, across every configuration at once, which is + several times the effect being looked for. + + Unlike :class:`Window` this is a plain Python object, not a ``cco_struct``, + so it cannot be carried across an ``scf.if``/``scf.for`` boundary. Build it + inside whichever region needs it, or use ``Window`` there. + """ + + __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 diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index 258944a0c..4d2c42820 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -294,25 +294,17 @@ def build_sdma_phases( I32_PER_PACK = 4 # a 16B pack always moves as 4 x i32, whatever the payload - def _next_flag(w, bid, geom): + def _next_flag(w, bid): """``_flag[block] + 1`` -- monotonic, never reset (graph-replay safe).""" - base = w.lsa_ptr_at(geom[0], geom[1], rank, flag_off) + fx.Int64( - bid - ) * fx.Int64(4) + 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, geom): - """Row-0 barrier: publish to peer ``tid``, then wait on ``tid``'s slot. - - The peer here is a *runtime* lane index, so this is the one site where - the address genuinely varies per lane -- and therefore the one where - computing it from a cached geometry, rather than calling `lsa_ptr`, - replaces a descriptor load pair with a multiply-add. - """ - peer_arr = w.lsa_ptr_at(geom[0], geom[1], tid, arr_off) + 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 = w.lsa_ptr_at(geom[0], geom[1], rank, arr_off) + 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( @@ -341,11 +333,10 @@ def _push_kernel( @flyc.kernel(known_block_size=[PUSH_THREADS, 1, 1]) def push(dev_comm: Int64, win: Int64): tid = fx.thread_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) sdma = cco.DevComm(dev_comm).sdma() - geom = w.lsa_geometry() - flag, flag_rsrc = _next_flag(w, 0, geom) + flag, flag_rsrc = _next_flag(w, 0) if tid < ws: if tid != rank: if const_expr(pushes): @@ -382,7 +373,7 @@ def push(dev_comm: Int64, win: Int64): # 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, geom) + _signal_and_wait(w, arr_off, flag, tid) fgpu.barrier() if tid == 0: local_store_u32(flag_rsrc, flag) @@ -409,26 +400,27 @@ def sdma_reduce(dev_comm: Int64, win: Int64): """ tid = fx.thread_idx.x bid = fx.block_idx.x - w = cco.Window(win) + 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. - # Nine descriptors from one window: take the geometry once. - geom = w.lsa_geometry() - - def _rsrc(byte_off): - return create_buffer_resource_from_addr( - wave_uniform_i64(w.lsa_ptr_at(geom[0], geom[1], rank, byte_off)) - ) - self_off = ( cfg.recv_slot_off(rank) if reduce_self_from_recv else in_off + my_slice_off ) - srcs = [_rsrc(self_off)] + [ - _rsrc(cfg.recv_slot_off((rank + j) % ws)) for j in range(1, ws) + 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 = _rsrc(out_off + my_slice_off) + 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): @@ -556,34 +548,25 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): """ tid = fx.thread_idx.x bid = fx.block_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) sdma = cco.DevComm(dev_comm).sdma() - # Nine descriptors from one window: take the geometry once. - geom = w.lsa_geometry() - - def _rsrc(byte_off): - return create_buffer_resource_from_addr( - wave_uniform_i64(w.lsa_ptr_at(geom[0], geom[1], rank, byte_off)) - ) - self_off = ( cfg.recv_slot_off(rank) if reduce_self_from_recv else in_off + my_slice_off ) - srcs = [_rsrc(self_off)] + [ - _rsrc(cfg.recv_slot_off((rank + j) % ws)) for j in range(1, ws) + 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 = _rsrc(out_off + my_slice_off) - - # `lsa_ptr` is an opaque extern call that loads two fields out of the - # window descriptor (`winBase`, `stride4G`) and combines them. The - # compiler cannot see through it, so it cannot hoist or CSE the calls - # itself -- every call re-reads both. Both of these bases are - # loop-invariant, so they are taken once here rather than once per - # (block, band). This is the same reason the source/destination buffer - # descriptors above are built at kernel scope and not in the pack loop. - ctr_base = w.lsa_ptr_at(geom[0], geom[1], rank, gather_counter_off) - lock_base = w.lsa_ptr_at(geom[0], geom[1], rank, lock_off) + out = create_buffer_resource_from_addr( + wave_uniform_i64(w.lsa_ptr(rank, out_off + my_slice_off)) + ) gtid = bid * threads + tid for band in range_constexpr(cfg.gather_bands): @@ -645,7 +628,10 @@ def _rsrc(byte_off): 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)) + ctr = signal_ptr( + fx.Int64(w.lsa_ptr(rank, gather_counter_off)) + + 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 @@ -655,7 +641,10 @@ def _rsrc(byte_off): 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)) + lock = signal_ptr( + fx.Int64(w.lsa_ptr(rank, lock_off)) + + fx.Int64(dest) * fx.Int64(4) + ) _acquire_peer_lock(lock) sdma.put( dest, @@ -713,17 +702,15 @@ def _rsrc(byte_off): ), ) - # These build one descriptor per (peer, region), and `lsa_pull_gather` - # wants 15 of them. `lsa_ptr` reads the window's base and stride on - # every call, so callers take the geometry once and the rest is - # arithmetic -- see Window.lsa_geometry. - def _row_addr_at(w, peer, byte_off, geom): + def _row_addr_at(w, peer, byte_off): return create_buffer_resource_from_addr( - wave_uniform_i64(w.lsa_ptr_at(geom[0], geom[1], peer, byte_off)) + wave_uniform_i64(w.lsa_ptr(peer, byte_off)) ) - def _row_addr(w, byte_off, geom): - return _row_addr_at(w, rank, byte_off, geom) + 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): @@ -764,23 +751,21 @@ def sdma_reduce_quant(dev_comm: Int64, win: Int64): """ tid = fx.thread_idx.x bid = fx.block_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) - geom = w.lsa_geometry() self_off = ( cfg.recv_slot_off(rank) if reduce_self_from_recv else in_off + my_slice_off ) - srcs = [_row_addr(w, self_off, geom)] + [ - _row_addr(w, cfg.recv_slot_off((rank + j) % ws), geom) - for j in range(1, ws) + 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, geom) - gdst = _row_addr(w, cfg.gout_off + rank * cfg.slice_rows * cfg.n, geom) - gsca = _row_addr(w, cfg.gout_scale_slice_off(rank), geom) + 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 @@ -856,14 +841,13 @@ def quantize_gather(dev_comm: Int64, win: Int64): """ tid = fx.thread_idx.x bid = fx.block_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) - geom = w.lsa_geometry() - src = _row_addr(w, out_off + my_slice_off, geom) - dst = _row_addr(w, cfg.gout_off + rank * cfg.slice_rows * cfg.n, geom) - sca = _row_addr(w, cfg.gout_scale_slice_off(rank), geom) + 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 @@ -937,18 +921,15 @@ def dequantize_gather(dev_comm: Int64, win: Int64): """ tid = fx.thread_idx.x bid = fx.block_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) - geom = w.lsa_geometry() for j in range_constexpr(1, ws): peer = (rank + j) % ws - # All three are in *my own* window: the SDMA push already - # delivered peer `peer`'s slice and scales here. - src = _row_addr(w, cfg.gout_off + peer * cfg.slice_rows * cfg.n, geom) - sca = _row_addr(w, cfg.gout_scale_slice_off(peer), geom) - dst = _row_addr(w, out_off + peer * cfg.slice_rows * cfg.n * 2, geom) + 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 @@ -1021,28 +1002,23 @@ def lsa_pull_gather(dev_comm: Int64, win: Int64): """ tid = fx.thread_idx.x bid = fx.block_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) lane = tid % fx.Int32(WAVE) wave = tid // fx.Int32(WAVE) raw_cco.cco_system_fence(fx.Int32(0)) - # 15 descriptors below, so the window geometry is read once here - # rather than once per descriptor. - geom = w.lsa_geometry() # 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, geom) + _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), geom) 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, geom) + out = _row_addr(w, out_off) for row in range( bid * QUANT_WAVES + wave, slice_rows, pull_blocks * QUANT_WAVES From 536eba6f2ded04a6dab2d95fb6a432277f29619e Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 10:42:24 +0000 Subject: [PATCH 66/79] cco: CachedWindow in kernels_lsa and kernels_fused -- and the first real win Four more sites: both kernels in kernels_lsa.py, and kernels_fused.py's w_pre and its one inline lsa_ptr. kernels_fused.py:1488 stays on Window, see below. split-lsa is the first configuration where this measures. A/B, three alternating repeats, m=16384 n=7168 k=2048, 8x MI355X, max_rank_time_us: 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 -0.85% against a Window spread of 0.6us across its three runs. The reason is the shape of the kernel, not the number of call sites. ar_1stage/ar_2stage build nine peer addresses in *every block* of a short kernel, so two flat_loads apiece is a real fraction of it. Every other site converted so far -- 21 in kernels_sdma, both here in kernels_fused -- runs once per launch against a body lasting a millisecond, and all of them measured nothing. Count address constructions per launch, not `grep -c lsa_ptr`. kernels_fused.py:1488 cannot be converted. That `w` is read inside a dynamic `if`, and FlyDSL captures every variable an if body reads as state, requiring each to be a single MLIR value: TypeError: state variable 'w' is CachedWindow, not an MLIR Value Making CachedWindow a cco_struct does not help -- three fields extract to three values: TypeError: state variable 'w' is list, not an MLIR Value Window passes only because it has exactly one field. Reverted that experiment; CachedWindow stays a plain object, and the limitation is now documented with the error text so nobody retries it. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/cco/device/flydsl/handles.py | 35 ++++++++++++++---------- python/mori/ops/gemm_ar/kernels_fused.py | 4 +-- python/mori/ops/gemm_ar/kernels_lsa.py | 4 +-- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index ade444678..2eec48c52 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -134,28 +134,35 @@ class CachedWindow: counts against ``lgkmcnt`` as well as ``vmcnt``, giving any following ``s_waitcnt`` one more counter to wait on. - **It does not measure on mori's own kernels, and that is expected.** - A/B/A' at ``m=16384 n=7168 k=2048`` on 8x MI355X, us, lower is better:: + **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:: - wire Window CachedWindow Window again - bf16/sdma 1114.5 1111.4 1115.4 - fp8/sdma 983.3 983.6 983.6 - fp8/lsa 926.6 924.3 926.0 + 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 - The difference is smaller than the drift between the two Window runs. Every - ``lsa_ptr`` call in ``gemm_ar`` is already at kernel scope -- descriptors get - built once before the hot loop, which is the established idiom -- so what - this removes is on the order of tens of loads per launch against a kernel - that runs for a millisecond. Reach for it when a kernel needs addresses *per - iteration*; do not expect it to show up otherwise. + ``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. Absolute numbers on this box move ~4% between mornings, across every configuration at once, which is several times the effect being looked for. Unlike :class:`Window` this is a plain Python object, not a ``cco_struct``, - so it cannot be carried across an ``scf.if``/``scf.for`` boundary. Build it - inside whichever region needs it, or use ``Window`` there. + 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"). Build it inside + whichever region needs it, or use ``Window`` there -- + ``kernels_fused.py:1488`` does the latter. """ __slots__ = ("handle", "base", "stride") diff --git a/python/mori/ops/gemm_ar/kernels_fused.py b/python/mori/ops/gemm_ar/kernels_fused.py index 9bfc171e4..54eb428bc 100644 --- a/python/mori/ops/gemm_ar/kernels_fused.py +++ b/python/mori/ops/gemm_ar/kernels_fused.py @@ -1086,7 +1086,7 @@ def kernel_gemm_scatter( mfma = _SwappedMfma(mfma_raw) else: mfma = Mfma16x16x128(N_TILES_A, N_TILES_B) - w_pre = cco.Window(win) + 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) @@ -1202,7 +1202,7 @@ def kernel_gemm_scatter( N_TILES_A, N_TILES_B, c_rsrc=create_buffer_resource_from_addr( - fx.Int64(cco.Window(win).lsa_ptr(rank, in_off)), + fx.Int64(cco.CachedWindow(win).lsa_ptr(rank, in_off)), num_records_bytes=cfg.nbytes, ), ) diff --git a/python/mori/ops/gemm_ar/kernels_lsa.py b/python/mori/ops/gemm_ar/kernels_lsa.py index faa8564cc..de169f09a 100644 --- a/python/mori/ops/gemm_ar/kernels_lsa.py +++ b/python/mori/ops/gemm_ar/kernels_lsa.py @@ -200,7 +200,7 @@ def _acc_peers(rsrcs, pack_idx): def ar_1stage(dev_comm: Int64, win: Int64): bid = fx.block_idx.x tid = fx.thread_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) flag, flag_rsrc = _next_flag(w, bid) if tid < ws: @@ -228,7 +228,7 @@ def ar_1stage(dev_comm: Int64, win: Int64): def ar_2stage(dev_comm: Int64, win: Int64): bid = fx.block_idx.x tid = fx.thread_idx.x - w = cco.Window(win) + w = cco.CachedWindow(win) flag_a, flag_rsrc = _next_flag(w, bid) if tid < ws: From d5e26faeaebaaca65f360fcda42fd6ef30338e45 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 13:00:40 +0000 Subject: [PATCH 67/79] gemm_ar: reuse w_pre instead of rebuilding the window twice more kernel_gemm_scatter built the window three times: w_pre at the top, then a throwaway `cco.Window(win)` inline for c_rsrc, and a third one for the scatter epilogue. All three are in the same function with no nested def between them, so w_pre was in scope for all of them -- the rebuilds were left over, not needed. The epilogue one also needed the two addresses moved out of the dynamic `if`, which is what let it use w_pre at all: FlyDSL captures every variable an if body reads as state and requires single MLIR values, so a CachedWindow read inside fails. Both bases are rank-local with constant offsets, identical across threads, so hoisting them is free and leaves only plain Int64 in the captured state. That retires the "kernels_fused.py:1488 stays on Window" exception from 536eba6f -- the file now builds the window exactly once. No timing change, which is what the per-launch rule predicts: these addresses are built once per launch, not per block. fused-sdma bf16 1111.3us, fp8/lsa 930.0, fused-lsa 1360.3, all inside the spreads in 536eba6f. relL2 unchanged at 2.35e-3 bf16 / 2.49e-2 fp8, validated on all three. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/kernels_fused.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/python/mori/ops/gemm_ar/kernels_fused.py b/python/mori/ops/gemm_ar/kernels_fused.py index 54eb428bc..f36f376e7 100644 --- a/python/mori/ops/gemm_ar/kernels_fused.py +++ b/python/mori/ops/gemm_ar/kernels_fused.py @@ -1202,7 +1202,7 @@ def kernel_gemm_scatter( N_TILES_A, N_TILES_B, c_rsrc=create_buffer_resource_from_addr( - fx.Int64(cco.CachedWindow(win).lsa_ptr(rank, in_off)), + fx.Int64(w_pre.lsa_ptr(rank, in_off)), num_records_bytes=cfg.nbytes, ), ) @@ -1485,7 +1485,12 @@ def kernel_gemm_scatter( # "writethrough": nothing to do -- the stores already went to memory, # and wait_barrier(0) above retired them. - w = cco.Window(win) + # 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 @@ -1499,10 +1504,7 @@ def kernel_gemm_scatter( m_tiles_per_chunk ) slot = dest * fx.Int32(chunks) + chunk - ctr = signal_ptr( - fx.Int64(w.lsa_ptr(rank, counter_off)) - + fx.Int64(slot) * fx.Int64(4) - ) + 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) @@ -1518,10 +1520,7 @@ def kernel_gemm_scatter( # 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( - fx.Int64(w.lsa_ptr(rank, lock_off)) - + fx.Int64(dest) * fx.Int64(4) - ) + lock = signal_ptr(lock_base + fx.Int64(dest) * fx.Int64(4)) if const_expr(chunks > 1): _acquire_peer_lock(lock) sdma.put( From 7c8c9c71677085d59e625212827a9ceab89e8141 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 13:01:14 +0000 Subject: [PATCH 68/79] cco: the scf-boundary note now points at hoisting, not at a line that changed d5e26fae removed the kernels_fused.py:1488 Window the docstring cited as the worked example. Hoisting the addresses out of the branch is the general answer anyway, and it does not rot when a line number moves. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/cco/device/flydsl/handles.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index 2eec48c52..cfbe48c75 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -160,9 +160,11 @@ class CachedWindow: ``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"). Build it inside - whichever region needs it, or use ``Window`` there -- - ``kernels_fused.py:1488`` does the latter. + 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") From 96bd9522eda096ad455bde09e2251516a936b24b Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 13:50:17 +0000 Subject: [PATCH 69/79] gemm_ar: fix fuse-reduce-push, broken by the CachedWindow conversion 81fe2fea converted kernels_sdma's 21 lsa_ptr sites and I measured three wire configurations, all green. --fuse-reduce-push is off by default, so neither the test suite nor any sweep ever traced sdma_reduce_push, and it had been dead since: TypeError: state variable 'w' is CachedWindow, not an MLIR Value Same cause and same fix as kernels_fused: the two bases are rank-local at a constant offset, so they are invariant across bands and identical on every thread. Hoisting them above the band loop keeps the window out of the if's captured state and computes them once instead of once per band. Smoke-tested the other off-default flags for the same fault -- publish=fence, fuse-quantize both ways, direct-fence=all, gather-bands=4, sdma-queues=1 -- all validate. Answering the original question: fusing the push into the reduce still does not pay. M=16384, three alternating repeats: off 1111.4 1111.5 1112.9 ON 1127.2 1129.4 1128.8 +16.4us, +1.5% Worse than the +8.5us recorded before CachedWindow, and the spreads are 1.5 and 2.2us. Staying off by default. Also: --fuse-quantize defaulted to True while build_sdma_phases defaults it to False, and its help claimed "same numbers either way". It costs 23.4us on fp8/lsa (926.6 vs 903.2, three repeats each) and 18.5us on fp8/sdma. Flipped the bench default to match the library, so the bench stops reporting the op ~23us slower than it runs. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 6 +++--- python/mori/ops/gemm_ar/kernels_sdma.py | 17 +++++++++-------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 33b7afa72..643219d8c 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -682,10 +682,10 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument( "--fuse-quantize", action=argparse.BooleanOptionalAction, - default=True, + default=False, help="fold the fp8 narrowing into the reduce instead of running it as " - "its own kernel. Same numbers either way; it saves reading the bf16 " - "back out of memory.", + "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", diff --git a/python/mori/ops/gemm_ar/kernels_sdma.py b/python/mori/ops/gemm_ar/kernels_sdma.py index 4d2c42820..db375446f 100644 --- a/python/mori/ops/gemm_ar/kernels_sdma.py +++ b/python/mori/ops/gemm_ar/kernels_sdma.py @@ -568,6 +568,13 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): 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 @@ -628,10 +635,7 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): if const_expr(publish == "fence"): raw_cco.cco_system_fence(fx.Int32(0)) if tid == fx.Int32(0): - ctr = signal_ptr( - fx.Int64(w.lsa_ptr(rank, gather_counter_off)) - + fx.Int64(band) * fx.Int64(4) - ) + 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 @@ -641,10 +645,7 @@ def sdma_reduce_push(dev_comm: Int64, win: Int64): 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( - fx.Int64(w.lsa_ptr(rank, lock_off)) - + fx.Int64(dest) * fx.Int64(4) - ) + lock = signal_ptr(lock_base + fx.Int64(dest) * fx.Int64(4)) _acquire_peer_lock(lock) sdma.put( dest, From 83046d455a01092b0b3be96ae650f353fc5cd5e1 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 15:25:00 +0000 Subject: [PATCH 70/79] docs: the window-geometry result, and retract the "machine drift" claim Adds the lsa_ptr / CachedWindow line to the negative results: four ways of hoisting the window geometry, all null on kernels_sdma and kernels_fused, one win of -10.7us on split-lsa, and the rule that separates them -- address constructions per launch, not call sites. Also re-measures the two existing negative results with repeats: fuse_reduce_push +16.4us, fuse_quantize +23.4us. And retracts something I got wrong and committed. I read today's numbers against this page and found everything ~4% faster, and wrote that up -- in 81fe2fea's message and in CachedWindow's docstring -- as the box drifting between mornings. It does not. The benchmark defaults to --quant ptpc; every number on this page is --quant blockscale. Measured back to back: mode doc today blockscale today ptpc gemm-only 369.4 367.9 221.0 split-sdma 1472.9 1462.8 1259.0 fused bf16 1148.8 1144.9 1112.1 The page reproduces within ~1%. There is no drift; there was a flag I did not pin. The A/B pairs are unaffected -- both arms always ran the same flags -- but the explanation attached to them was wrong, so it is corrected at both sites. Also confirms the PR comparison in blockscale, the mode the headline is quoted in: split-lsa 1463.6 -> 1455.6us, everything else flat. Co-Authored-By: Claude Opus 5 (1M context) --- docs/MORI-GEMM-AR-BENCHMARK.md | 52 ++++++++++++++++++++++++ python/mori/cco/device/flydsl/handles.py | 7 ++-- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md index be9ff9167..d8667adc3 100644 --- a/docs/MORI-GEMM-AR-BENCHMARK.md +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -267,6 +267,13 @@ It is the thread map: a per-row amax cannot be taken by a block holding only par 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 @@ -308,6 +315,51 @@ 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` diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index cfbe48c75..70b34d3b0 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -151,9 +151,10 @@ class CachedWindow: So: count address constructions per launch, not ``grep -c lsa_ptr``. - Measure A/B in one session if you revisit this. Absolute numbers on this box - move ~4% between mornings, across every configuration at once, which is - several times the effect being looked for. + 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 From e3514e25236c36ff0867ef5f976f3020b7ffbfb7 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Mon, 14 Sep 2026 16:44:55 +0000 Subject: [PATCH 71/79] gemm_ar: fix the review's P1/P2 findings in the public op, and test that layer The op layer had no tests at all -- test_gemm_ar.py goes through bench_gemm_ar.py, which calls the low-level builders -- so three correctness bugs sat behind a green suite. Adds tests/python/cco/test_gemm_ar_op.py (27 cases) and a worker that drives GemmAllReduceOp itself. P1, all reproduced before fixing: * The A scale was flattened with reshape(-1). The kernel reads (row, kb) at kb*M + row, so for the column-major [M, K/128] tensor the quantiser returns that pairs every scale with the wrong K block: relL2 0.26 instead of 2.4e-3. _flatten_a_scale now distinguishes [M, K/128], [K/128, M] and an already-flat physical-order buffer, and the three agree bitwise. * counter_chunks changes with M, which moved lock_off and input_off, so a second M read the previous shape's payload as control state. The counter region is now sized by capacity (MAX_CHUNKS) and carries one counter set per shape -- the election test is `seq % tiles_per_chunk == 0` on a counter that is never reset, and tiles_per_chunk depends on M, so sharing a set means electing on the other shape's residue. * supports() admitted K=128, which returns finite garbage (relL2 1.41): the mainloop prefetches block 1 and runs two tail steps, so K_ITERS must be >= 2. It also disagreed with ArConfig.validate() on world sizes 3/5/6/7. It now builds the config and validates it rather than keeping a second copy of the rules. P2: * ArConfig.validate() rejects misspelled mode strings ("FP8" selected bf16, "sla" selected SDMA), non-positive sizes, and the fp8 N multiple -- before anything is allocated. * __call__ validates operand dtype/shape/device. A/B go in as int8, so a bf16 tensor used to reach the kernel and come back finite and wrong. * close() / context manager. It releases the window and memory, not the dev-comm: destroying that and then letting the communicator tear down aborts the process. The window is the point -- ~700 MiB at the model shape. * pad_rows keys its buffer on (device, dtype, K), not just capacity, and checks m_pad >= m. A bf16 buffer reused for an fp32 input silently rounded it. * RESULT_JSON records the effective wire, transport, quantize fusion, queues and resolved tile order. An fp8/lsa and an fp8/sdma run both said only mode="fused-sdma". * The cached plan is a NamedTuple with the phase tail already resolved. Also: removed LSA_NOTE, which said the LSA transport is not implemented in the file that implements it; corrected the --fence help, which named 'agent' as the default when the default is 'none' and called 'none' incorrect. Measured, twice each at m=16384 blockscale on 8 ranks: 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, so the explicit fences cost 650-790us and did not move accuracy. README gains a runnable end-to-end example; op.padded_m(m, world_size) in the old one raised TypeError. One defect found by the new tests and NOT fixed: after a call at some M, the first call at a *larger* M is wrong about half the time, always with the same value, with the repeat correct. 8 runs each on 2 ranks: 512->1024 fails 4/8, 1024->512 fails 0/8, and 2048->4096 fails too, so it is the growth and not the chunk count. Five calls at one M with different operands each time are all correct, which rules out a stale read at a fixed shape. The op now raises instead of returning the wrong answer, and the test pins that. 73 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py | 31 +- python/mori/ops/gemm_ar/README.md | 75 ++- python/mori/ops/gemm_ar/kernels_fused.py | 22 +- python/mori/ops/gemm_ar/layout.py | 110 ++++- python/mori/ops/gemm_ar/op.py | 460 +++++++++++++++--- tests/python/cco/gemm_ar_op_worker.py | 247 ++++++++++ tests/python/cco/test_gemm_ar_op.py | 349 +++++++++++++ 7 files changed, 1179 insertions(+), 115 deletions(-) create mode 100644 tests/python/cco/gemm_ar_op_worker.py create mode 100644 tests/python/cco/test_gemm_ar_op.py diff --git a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py index 643219d8c..806a6178f 100644 --- a/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py +++ b/benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py @@ -488,6 +488,25 @@ def once(): "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, @@ -573,10 +592,14 @@ def build_parser() -> argparse.ArgumentParser: "raw-wt-leader", ), default="none", - help="release before the epilogue push. Correct: 'agent' (default, " - "cheapest) and 'all' (system scope, adds a buffer_inv that costs 20us). " - "Incorrect, and present only to price the fence: 'leader', 'none', " - "'writethrough'. 'wt-agent' is correct but no faster than 'agent'", + 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", diff --git a/python/mori/ops/gemm_ar/README.md b/python/mori/ops/gemm_ar/README.md index 3722a8102..31037cdb0 100644 --- a/python/mori/ops/gemm_ar/README.md +++ b/python/mori/ops/gemm_ar/README.md @@ -35,21 +35,68 @@ 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 from mori.ops.gemm_ar import GemmAllReduceOp, preshuffle_b -op = GemmAllReduceOp(comm, n=7168, k=2048, m_max=16384) -b_shuffled = preshuffle_b(b_fp8) # [N, K] -> MFMA preshuffled -out = op(a_fp8, b_shuffled, a_scale, b_scale) # [M, N] bf16, already reduced +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) + +with Communicator.init(world, rank, 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)) + out = op(a_fp8, b_shuffled, a_scale, b_scale)[: x.shape[0]] ``` -`a_scale` is the A 1x128 block scale, physically `[K/128, M]` fp32 (i.e. `[M, -K/128]` column-major -- what `aiter_per1x128_quant(transpose_scale=True)` -emits). `b_scale` is `[N/128, K/128]` fp32 row-major. +**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. Pass it as `[M, K/128]` +in either memory order, as `[K/128, M]`, or already flat in physical order -- +the op distinguishes these. What it cannot distinguish is a `[M, K/128]` tensor +you have already flattened yourself with `reshape(-1)`: for the column-major +tensor `aiter_per1x128_quant(transpose_scale=True)` returns, that walks the +logical rows and pairs every scale with the wrong K block, which validates, +runs, and returns relL2 0.26. `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` (8). Warm each M up once before capturing a graph. -`M` must be a multiple of `world_size * block_m`; `op.padded_m(m, world_size)` -rounds up and `op.pad_rows` zero-extends to it. A padded row produces a zero -output row, so the caller slices the result back to `m`. +**Cleanup.** `close()` (or the `with` above) releases the window, memory and +dev-comm. The communicator holds a strong reference to each, so dropping the op +is not enough -- at the model shape that is about 700 MiB per rank. Synchronise +first if anything may still be in flight. ## fp8 on the wire @@ -168,10 +215,12 @@ traces. ``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 -- having the epilogue store straight into -a peer -- would first require coalescing ``StoreC._store_bf16``, which writes one -bf16 at a time through ``BufferCopy16b``; gcnasm measured that lane-scatter at -0.26x when pushed to a peer. See ``LSA_NOTE`` at the bottom. +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 diff --git a/python/mori/ops/gemm_ar/kernels_fused.py b/python/mori/ops/gemm_ar/kernels_fused.py index f36f376e7..08fb2874d 100644 --- a/python/mori/ops/gemm_ar/kernels_fused.py +++ b/python/mori/ops/gemm_ar/kernels_fused.py @@ -1579,24 +1579,4 @@ def launch( return launch -LSA_NOTE = """\ -The LSA transport is not implemented, and the reason is a measurement rather than -a preference. - -Fusing LSA means the epilogue stores into a peer instead of into local C. But -`StoreC._store_bf16` writes **one bf16 at a time** through `BufferCopy16b`, at -`c_index = (row + i) * c_cols + col` with `row = base_row + ti*16 + lane//16*4 + i` -and `col = base_col + tj*16 + lane%16` -- a lane-scatter. gcnasm measured that -pattern at 0.26x when the destination is a peer, so the variant is worthless -until `StoreC` is coalesced (their winning form was a wave-local `ds_bpermute` -pair-coalesced store; staging C through LDS regressed 1-3%). - -That rewrite is the expensive half of this step, and the ceiling it is competing -for is small: at [4096, 7168] the GEMM is ~40us against an all-reduce of ~290us -(LSA) or ~340us (SDMA), so even perfect overlap removes at most the GEMM, ~12%. -Fusing the cheap transport first establishes whether the completion-counter -structure works at all; if the answer is no, the `StoreC` rewrite is wasted -regardless of transport. -""" - -__all__ = ["compile_fused_gemm_scatter", "LSA_NOTE", "BLOCK_K"] +__all__ = ["compile_fused_gemm_scatter", "BLOCK_K"] diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index 11f1c305c..8598d2c1a 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -74,6 +74,15 @@ 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 @@ -139,6 +148,21 @@ class ArConfig: #: 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 #: 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 @@ -361,7 +385,9 @@ def blocks(self) -> int: 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. - return self.gather_counter_off + self.gather_counter_bytes + # 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: @@ -383,18 +409,37 @@ def reduce_blocks(self) -> int: return max(1, min(SDMA_REDUCE_BLOCK_CAP, need)) @property - def counter_off(self) -> int: + 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_chunks`` and grows with M, so borrowing the 16 - spare ``_flag`` slots would silently stop fitting at M >= 8192. + 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.world_size * self.counter_chunks * 4, SIGNAL_ALIGN) + return _align_up( + self.counter_set_bytes * self.counter_shape_slots, SIGNAL_ALIGN + ) @property def lock_off(self) -> int: @@ -403,15 +448,18 @@ def lock_off(self) -> int: 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_off + self.counter_bytes + 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_off(self) -> int: + 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 @@ -422,9 +470,17 @@ def gather_counter_off(self) -> int: """ 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, SIGNAL_ALIGN) + 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: @@ -579,6 +635,41 @@ def validate(self) -> None: ) 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}], " @@ -590,6 +681,9 @@ def validate(self) -> None: __all__ = [ "ArConfig", + "WIRE_DTYPES", + "GATHER_TRANSPORTS", + "FP8_GATHER_N_MULTIPLE", "select_stage", "K_MAX_BLOCKS", "LSA_BLOCK_CAP", diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index 25711842c..4bcbbc319 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -23,7 +23,7 @@ from __future__ import annotations -from typing import Optional +from typing import NamedTuple, Optional import torch @@ -31,9 +31,9 @@ from mori.cco import CCODevCommRequirements, GDA_CONNECTION_NONE from mori.tensor_utils import from_gpu_ptr -from .kernels_fused import compile_fused_gemm_scatter +from .kernels_fused import BLOCK_K, compile_fused_gemm_scatter from .kernels_sdma import build_sdma_phases -from .layout import ArConfig +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 @@ -50,6 +50,17 @@ # 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. Each needs its own counter set -- +# see ArConfig.counter_shape_slots -- and each also costs a FlyDSL compile, so +# this is a small number by nature rather than a tuning knob. +DEFAULT_MAX_SHAPES = 8 + 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.""" @@ -72,28 +83,156 @@ def counter_chunks(m_pad: int, world_size: int, block_m: int = DEFAULT_BLOCK_M) return max(c for c in range(1, min(MAX_CHUNKS, bands) + 1) if bands % c == 0) +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, +) -> 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, + 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. """ - return ( - 2 <= world_size <= 8 - and n % block_n == 0 - and n % SCALE_BLOCK_K == 0 - and k % SCALE_BLOCK_K == 0 - and m > 0 + 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, + shape_index=0, + ) + except ValueError: + return False + return True + + +#: fp8 encodings the kernel's byte reinterpretation is valid for. Both are 1 +#: byte and the GEMM only ever sees the bytes, but forwarding a bf16 tensor here +#: reinterprets two rows as one and produces finite nonsense. +FP8_DTYPES = (torch.float8_e4m3fnuz, 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`` -- the layout + ``gemm_a8w8_blockscale_bpreshuffle`` consumes, and what sglang's + ``materialize_bpreshuffle_fp8_scale`` produces. Three spellings are accepted + and they mean different things: + + * ``[M, K/128]`` -- the logical shape. Transposed here, which is free for the + column-major tensor the quantiser returns (``.t()`` is then contiguous) and + a copy for a row-major one. **Passing this through ``reshape(-1)`` instead + is the bug this function exists to prevent**: for the column-major case + that walks the logical rows, pairing every scale with the wrong K block, + which still validates and still runs (relL2 0.26 against 2.4e-3). + * ``[K/128, M]`` -- the physical shape, taken as-is. + * 1-D -- already flattened in physical order, taken as-is. + + ``m == kb`` would be ambiguous; the logical reading wins, and K would have to + be ``128 * M`` for it to arise. + """ + 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 [M, K/128], [K/128, M] or flat, got shape " + f"{tuple(a_scale.shape)}" + ) + if tuple(a_scale.shape) == (m, kb): + return a_scale.t().reshape(-1) + if tuple(a_scale.shape) == (kb, m): + return a_scale.reshape(-1) + raise ValueError( + f"a_scale shape {tuple(a_scale.shape)} matches neither [M, K/128] = " + f"{(m, kb)} nor its 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 + + class GemmAllReduceOp: """Fused fp8 block-scale GEMM + all-reduce over cco SDMA. @@ -116,7 +255,32 @@ def __init__( sdma_queues: int = 1, gather_dtype: str = "bf16", gather_transport: str = "lsa", + max_shapes: int = DEFAULT_MAX_SHAPES, ): + self._closed = True # so a failed constructor leaves close() a no-op + 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" + ) + if block_m < 1 or m_max < 1 or sdma_queues < 1 or max_shapes < 1: + raise ValueError( + f"block_m, m_max, sdma_queues and max_shapes must all be >= 1; got " + f"block_m={block_m} m_max={m_max} sdma_queues={sdma_queues} " + f"max_shapes={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 @@ -134,36 +298,46 @@ def __init__( #: 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] = {} - if not supports(self.m_max, n, k, self.world_size, block_n=block_n): - raise ValueError( - f"unsupported shape: M<={self.m_max} N={n} K={k} " - f"world_size={self.world_size}" - ) - + self.mem = self.win = self.dev_comm = None self.window_bytes = self._make_cfg(self.m_max).window_bytes - 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) - - self._cache: dict[int, tuple] = {} + 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 + + self._cache: dict[int, _Plan] = {} self._pad_in: Optional[torch.Tensor] = None + #: M values that have actually run, for :meth:`_check_growth`. + self._m_seen: set[int] = set() @staticmethod def window_bytes_for( @@ -173,6 +347,8 @@ def window_bytes_for( n: int, block_m: int = DEFAULT_BLOCK_M, gather_dtype: str = "bf16", + gather_transport: str = "lsa", + max_shapes: int = DEFAULT_MAX_SHAPES, ) -> int: """Symmetric-window bytes an op for this shape will allocate. @@ -180,34 +356,75 @@ def window_bytes_for( VMM reservation, so ``Communicator.init(per_rank_vmm=...)`` has to be sized first. Same arithmetic the constructor uses. """ - m_pad = padded_m(m_max, world_size, block_m) - cfg = ArConfig( + return _build_cfg( world_size=world_size, - m=m_pad, + m_pad=padded_m(m_max, world_size, block_m), n=n, - recv_slots=world_size, - counter_chunks=counter_chunks(m_pad, world_size, block_m), + block_m=block_m, gather_dtype=gather_dtype, - ) - cfg.validate() - return cfg.window_bytes + gather_transport=gather_transport, + counter_capacity=MAX_CHUNKS, + shape_slots=max_shapes, + 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 _check_growth(self, m: int) -> None: + """Refuse the first call at an M larger than one already run. + + Known-bad and not yet root-caused: after a call at some M, the first call + at a *larger* M comes back wrong about half the time, always with the + same value, while the repeat is correct. Measured on 2 ranks, 8 runs + each: 512 then 1024 fails 4/8; 1024 then 512 fails 0/8; 2048 then 4096 + (same chunk count) also fails, so it is the growth and not the chunk + count. Zeroing the window and synchronising between the calls avoids it, + which points at a read of a region the smaller shape left dirty rather + than at the counters -- those are already per-shape. + + Until that is found, this is an error rather than a wrong answer. Serve + the largest M first (padding to it is usually what a server does anyway), + or give each M its own op. + """ + if self._m_seen and m > max(self._m_seen): + raise ValueError( + f"M={m} is larger than every M this instance has run " + f"{sorted(self._m_seen)}. Growing M within one op is a known " + f"defect -- the first call at the larger M is intermittently " + f"wrong. Call the largest M first, or build a second op." + ) + + 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: - cfg = ArConfig( + return _build_cfg( world_size=self.world_size, - m=m, + m_pad=m, n=self.n, - recv_slots=self.world_size, - counter_chunks=counter_chunks(m, self.world_size, self.block_m), + 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), ) - cfg.validate() - return cfg def _compiled(self, m: int): """Kernel plus SDMA phases for this M, compiled on first sight.""" @@ -233,12 +450,59 @@ def _compiled(self, m: int): lane_transpose=True, ) parts = build_sdma_phases(cfg, self.rank, queues=self.sdma_queues) - c = from_gpu_ptr(self.mem.ptr + cfg.input_off, (m, self.n), torch.bfloat16) - out = from_gpu_ptr(self.mem.ptr + cfg.output_off, (m, self.n), torch.bfloat16) - hit = (gemm, parts, c, out) + # 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, + 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 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. @@ -246,11 +510,30 @@ def pad_rows(self, x: torch.Tensor, m_pad: int) -> torch.Tensor: 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 self._pad_in is None or self._pad_in.shape[0] < m_pad: - self._pad_in = torch.zeros((m_pad, k), dtype=x.dtype, device=x.device) - buf = self._pad_in[:m_pad] + 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 @@ -268,12 +551,15 @@ def __call__( block_m`` (see :meth:`padded_m`), ``b_preshuffled`` is ``[N, K]`` through :func:`preshuffle_b`. - Both scale buffers are read linearly, in physical order: ``a_scale`` is - logically ``[M, K/128]`` but column-major, i.e. physically - ``[K/128, M]``; ``b_scale`` is ``[N/128, K/128]`` row-major. + ``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. + 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: @@ -286,13 +572,16 @@ def __call__( f"M={m} must be a multiple of world_size * block_m = " f"{self.world_size * self.block_m}; use padded_m()/pad_rows()" ) - gemm, parts, c, out = self._compiled(m) + self._check_operands(a_fp8, b_preshuffled, a_scale, b_scale, m) + self._check_growth(m) + plan = self._compiled(m) + self._m_seen.add(m) stream = fx.Stream(torch.cuda.current_stream()) - gemm( + plan.gemm( a_fp8.contiguous().view(torch.int8).view(-1), b_preshuffled.contiguous().view(torch.int8).view(-1), - c.view(-1), - a_scale.reshape(-1), + plan.input.view(-1), + _flatten_a_scale(a_scale, m, self.k // SCALE_BLOCK_K), b_scale.reshape(-1), m, self.n, @@ -300,9 +589,42 @@ def __call__( self.win.handle, stream=stream, ) - # parts["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. - for phase in parts["fused_order"]: - parts[phase](self.dev_comm.ptr, self.win.handle, stream=stream) - return out + 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/tests/python/cco/gemm_ar_op_worker.py b/tests/python/cco/gemm_ar_op_worker.py new file mode 100644 index 000000000..4227e5b15 --- /dev/null +++ b/tests/python/cco/gemm_ar_op_worker.py @@ -0,0 +1,247 @@ +#!/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 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() + results[label] = _rel_l2(got, _reference(world, m, n, k)) + _emit(rank, case="alternating_m", **results) + + +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) + + +def case_growing_m(op, rank, world, m_small, m_large, n, k): + """The op must refuse a larger M after a smaller one, not answer wrongly.""" + a, b, sa, sb = _operands(rank, m_small, n, k) + op(a, preshuffle_b(b), sa, sb) + refused = False + try: + a, b, sa, sb = _operands(rank, m_large, n, k) + op(a, preshuffle_b(b), sa, sb) + except ValueError: + refused = True + _emit(rank, case="growing_m", refused=refused) + + +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() + 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 == "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 == "changing_data": + case_changing_data(op, rank, world, args.m, args.n, args.k, args.calls) + elif args.case == "growing_m": + case_growing_m( + op, rank, world, args.m_small, args.m_large, args.n, args.k + ) + 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_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py new file mode 100644 index 000000000..aea4e7479 --- /dev/null +++ b/tests/python/cco/test_gemm_ar_op.py @@ -0,0 +1,349 @@ +# 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 + 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, whatever spelling came in.""" + 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], physically [kb, M] + assert col_major.stride() == (1, m) + + assert torch.equal(_flatten_a_scale(col_major, m, kb), want) + assert torch.equal(_flatten_a_scale(logical, 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 exact regression: 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)) + + +@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 + + +# --- 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 + + +def _run_worker(world_size: int, case: str, *extra: str, timeout: int = 900): + 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 + assert result.returncode == 0, output + records = [ + json.loads(line.removeprefix("RESULT_JSON ")) + for line in output.splitlines() + if line.startswith("RESULT_JSON ") + ] + assert records, output + return records[0] + + +requires_two_gpus = pytest.mark.skipif( + torch.cuda.device_count() < 2, reason="needs 2 GPUs" +) + + +@requires_two_gpus +def test_public_op_takes_the_column_major_scale(): + """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 = _run_worker(2, "scale_order", "-m", "512", "-n", "1024", "-k", "512") + assert r["rel_l2"] < FP8_FLOOR, r + assert r["flat_matches"], r + assert r["physical_matches"], r + + +@requires_two_gpus +def test_one_instance_serves_a_smaller_m_after_a_larger_one(): + """M=4096 then M=8192 used to hang: the control region moved under it. + + Runs both twice, alternating, so a counter that carried a residue from the + other shape would show up on the repeat as well as on the first switch. + Shrinking only -- see :func:`test_growing_m_is_refused`. + """ + r = _run_worker( + 2, + "alternating_m", + "--m-small", + "1024", + "--m-large", + "512", + "-n", + "1024", + "-k", + "512", + ) + 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(): + """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 = _run_worker( + 2, "changing_data", "-m", "512", "-n", "1024", "-k", "512", "--calls", "5" + ) + for key in (f"call{i}" for i in range(5)): + assert r[key] < FP8_FLOOR, (key, r) + + +@requires_two_gpus +def test_growing_m_is_refused(): + """Growing M within one instance is a known defect, so it must not run. + + After a call at some M, the first call at a *larger* M comes back wrong + about half the time -- always the same value, with the repeat correct. + 8 runs each on 2 ranks: + + 512 -> 1024 4 of 8 wrong + 1024 -> 512 0 of 8 wrong + 2048 -> 4096 also wrong, so it is the growth, not the chunk count + 512 x5, different operands each call all correct + + The last line rules out a stale read of a peer's slice at a fixed shape, and + the per-shape counter sets tested above rule out a counter residue. Root + cause still open; until then the op raises instead of returning the wrong + answer, and this test pins that behaviour. + """ + r = _run_worker( + 2, + "growing_m", + "--m-small", + "512", + "--m-large", + "1024", + "-n", + "1024", + "-k", + "512", + ) + assert r["refused"], r + + +@requires_two_gpus +def test_close_releases_and_is_idempotent(): + """The communicator holds the handles, so dropping the op is not enough.""" + r = _run_worker(2, "close", "-n", "1024", "-k", "512", "-m", "512") + assert r["refused_after_close"], r From 5576f5008ae82cff5e0b7abc1f8c8c6b811f5e02 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 15 Sep 2026 02:36:44 +0000 Subject: [PATCH 72/79] gemm_ar: pin the payload map to m_max, so two shapes stop aliasing each other Root-causes and fixes the defect e3514e25 could only guard against: after a call at some M, the first call at a larger M came back wrong about half the time. The row map was the same on every failing run -- at world=2 n=1024, m=1024 had rows 512..767 wrong, contiguous, and nothing else. Those rows are bytes 3152128..3676416, which is exactly where m=512's recv slot for peer 1 sat, and their contents matched peer 1's *unreduced* GEMM output at m=512 to 9.5e-07, i.e. that peer's scatter payload from the previous call, still in place. 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. The overlaps at that shape: m=1024 input x m=512 output 1048576 B m=1024 output x m=512 recv0/1 524288 B each ArConfig.capacity_m lays the payload out for the largest m an instance will serve, while sizes stay m-based -- the map moves, not the bytes moved. The op passes m_max for every shape, so a region only ever aliases itself. The 15 offset properties now read their sizes from `self._cap`, a copy with m replaced; with capacity_m unset it is `self`, so nothing else changes and the benchmark, which builds its own configs, is untouched. Verified: 512 -> 1024 is 8/8 correct where it was 4/8 wrong; 2048 -> 4096 6/6. The growth guard and its test are gone -- growing M is supported now, and test_one_instance_serves_two_m_values_in_any_order covers both directions. A host-only test pins the map being M-independent, and asserts it is *not* without capacity_m, which is the property that broke. No performance change, as expected since the map is identical when m == m_max: blockscale at the model shape, split-sdma 1460.0us, fused bf16 1149.4, fp8/lsa 949.0, split-lsa 1453.4, all within noise of the same run before. 74 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/layout.py | 73 ++++++++++++----- python/mori/ops/gemm_ar/op.py | 34 ++------ tests/python/cco/gemm_ar_op_worker.py | 74 +++++++++++++----- tests/python/cco/test_gemm_ar_op.py | 108 ++++++++++++++++---------- 4 files changed, 183 insertions(+), 106 deletions(-) diff --git a/python/mori/ops/gemm_ar/layout.py b/python/mori/ops/gemm_ar/layout.py index 8598d2c1a..2744f7831 100644 --- a/python/mori/ops/gemm_ar/layout.py +++ b/python/mori/ops/gemm_ar/layout.py @@ -42,7 +42,7 @@ from __future__ import annotations -from dataclasses import dataclass +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 @@ -163,6 +163,17 @@ class ArConfig: #: 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 @@ -497,6 +508,18 @@ def signal_slot(self, block: int, peer: int) -> int: 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.""" @@ -505,23 +528,24 @@ def input_off(self) -> int: @property def input_scale_off(self) -> int: """Scatter-leg scales for ``input``. Zero-sized on the bf16 wire.""" - return _align_up(self.input_off + self.scatter_nbytes, SIGNAL_ALIGN) + 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``.""" - return ( - self.input_scale_off - + peer * self.slice_rows * self.scatter_tiles_per_row * 4 - ) + 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.""" - return self.input_scale_off + self.scatter_scale_bytes + c = self._cap + return c.input_scale_off + c.scatter_scale_bytes @property def output_end(self) -> int: - return _align_up(self.output_off + self.nbytes, SIGNAL_ALIGN) + c = self._cap + return _align_up(c.output_off + c.nbytes, SIGNAL_ALIGN) @property def gout_off(self) -> int: @@ -532,7 +556,8 @@ def gout_off(self) -> int: 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. """ - return self.output_end if self.fp8_gather else self.output_off + c = self._cap + return c.output_end if c.fp8_gather else c.output_off @property def gout_bytes(self) -> int: @@ -540,17 +565,20 @@ def gout_bytes(self) -> int: @property def gout_scale_off(self) -> int: - return _align_up(self.gout_off + self.gout_bytes, SIGNAL_ALIGN) + c = self._cap + return _align_up(c.gout_off + c.gout_bytes, SIGNAL_ALIGN) def gout_scale_slice_off(self, peer: int) -> int: - return self.gout_scale_off + peer * self.slice_rows * 4 + 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. - return max(self.output_end, self.gout_scale_off + self.gather_scale_bytes) + c = self._cap + return max(c.output_end, c.gout_scale_off + c.gather_scale_bytes) @property def tmp_bytes(self) -> int: @@ -559,11 +587,13 @@ def tmp_bytes(self) -> int: @property def recv_off(self) -> int: - return self.tmp_off + self.tmp_bytes + c = self._cap + return c.tmp_off + c.tmp_bytes @property def recv_bytes(self) -> int: - return self.recv_slots * self.scatter_slice_bytes + 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.""" @@ -573,17 +603,20 @@ def recv_slot_off(self, peer: int) -> int: raise IndexError( f"peer {peer} has no landing slot; recv_slots={self.recv_slots}" ) - return self.recv_off + peer * self.scatter_slice_bytes + c = self._cap + return c.recv_off + peer * c.scatter_slice_bytes @property def recv_scale_off(self) -> int: - return _align_up(self.recv_off + self.recv_bytes, SIGNAL_ALIGN) + 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 - return _align_up(self.recv_slots * self.scatter_scale_slice_bytes, SIGNAL_ALIGN) + 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.""" @@ -591,11 +624,13 @@ def recv_scale_slot_off(self, peer: int) -> int: raise IndexError( f"peer {peer} has no landing slot; recv_slots={self.recv_slots}" ) - return self.recv_scale_off + peer * self.scatter_scale_slice_bytes + c = self._cap + return c.recv_scale_off + peer * c.scatter_scale_slice_bytes @property def window_bytes(self) -> int: - return self.recv_scale_off + self.recv_scale_bytes + c = self._cap + return c.recv_scale_off + c.recv_scale_bytes # --- traffic model, for reporting alongside measured time --- diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index 4bcbbc319..c367e53de 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -110,6 +110,7 @@ def _build_cfg( 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. @@ -126,6 +127,7 @@ def _build_cfg( 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, ) @@ -336,8 +338,6 @@ def __init__( self._cache: dict[int, _Plan] = {} self._pad_in: Optional[torch.Tensor] = None - #: M values that have actually run, for :meth:`_check_growth`. - self._m_seen: set[int] = set() @staticmethod def window_bytes_for( @@ -372,30 +372,6 @@ 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 _check_growth(self, m: int) -> None: - """Refuse the first call at an M larger than one already run. - - Known-bad and not yet root-caused: after a call at some M, the first call - at a *larger* M comes back wrong about half the time, always with the - same value, while the repeat is correct. Measured on 2 ranks, 8 runs - each: 512 then 1024 fails 4/8; 1024 then 512 fails 0/8; 2048 then 4096 - (same chunk count) also fails, so it is the growth and not the chunk - count. Zeroing the window and synchronising between the calls avoids it, - which points at a read of a region the smaller shape left dirty rather - than at the counters -- those are already per-shape. - - Until that is found, this is an error rather than a wrong answer. Serve - the largest M first (padding to it is usually what a server does anyway), - or give each M its own op. - """ - if self._m_seen and m > max(self._m_seen): - raise ValueError( - f"M={m} is larger than every M this instance has run " - f"{sorted(self._m_seen)}. Growing M within one op is a known " - f"defect -- the first call at the larger M is intermittently " - f"wrong. Call the largest M first, or build a second op." - ) - 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) @@ -424,6 +400,10 @@ def _make_cfg(self, m: int) -> ArConfig: 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): @@ -573,9 +553,7 @@ def __call__( f"{self.world_size * self.block_m}; use padded_m()/pad_rows()" ) self._check_operands(a_fp8, b_preshuffled, a_scale, b_scale, m) - self._check_growth(m) plan = self._compiled(m) - self._m_seen.add(m) stream = fx.Stream(torch.cuda.current_stream()) plan.gemm( a_fp8.contiguous().view(torch.int8).view(-1), diff --git a/tests/python/cco/gemm_ar_op_worker.py b/tests/python/cco/gemm_ar_op_worker.py index 4227e5b15..926ed475a 100644 --- a/tests/python/cco/gemm_ar_op_worker.py +++ b/tests/python/cco/gemm_ar_op_worker.py @@ -133,6 +133,47 @@ def case_scale_order(op, rank, world, m, n, k): ) +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. @@ -149,7 +190,21 @@ def case_alternating_m(op, rank, world, m_small, m_large, n, k): ): a, b, sa, sb = _operands(rank, m, n, k) got = op(a, preshuffle_b(b), sa, sb).clone() - results[label] = _rel_l2(got, _reference(world, m, n, k)) + 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="alternating_m", **results) @@ -191,19 +246,6 @@ def case_changing_data(op, rank, world, m, n, k, calls): _emit(rank, case="changing_data", **out) -def case_growing_m(op, rank, world, m_small, m_large, n, k): - """The op must refuse a larger M after a smaller one, not answer wrongly.""" - a, b, sa, sb = _operands(rank, m_small, n, k) - op(a, preshuffle_b(b), sa, sb) - refused = False - try: - a, b, sa, sb = _operands(rank, m_large, n, k) - op(a, preshuffle_b(b), sa, sb) - except ValueError: - refused = True - _emit(rank, case="growing_m", refused=refused) - - def main() -> int: p = argparse.ArgumentParser() p.add_argument("--case", required=True) @@ -229,10 +271,6 @@ def main() -> int: case_scale_order(op, rank, world, args.m, 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 == "growing_m": - case_growing_m( - op, rank, world, args.m_small, args.m_large, args.n, args.k - ) elif args.case == "alternating_m": case_alternating_m( op, rank, world, args.m_small, args.m_large, args.n, args.k diff --git a/tests/python/cco/test_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py index aea4e7479..3a31dc885 100644 --- a/tests/python/cco/test_gemm_ar_op.py +++ b/tests/python/cco/test_gemm_ar_op.py @@ -270,20 +270,30 @@ def test_public_op_takes_the_column_major_scale(): @requires_two_gpus -def test_one_instance_serves_a_smaller_m_after_a_larger_one(): - """M=4096 then M=8192 used to hang: the control region moved under it. - - Runs both twice, alternating, so a counter that carried a residue from the - other shape would show up on the repeat as well as on the first switch. - Shrinking only -- see :func:`test_growing_m_is_refused`. +@pytest.mark.parametrize("m_small,m_large", [(512, 1024), (2048, 4096)]) +def test_one_instance_serves_two_m_values_in_any_order(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 = _run_worker( 2, "alternating_m", "--m-small", - "1024", + str(m_small), "--m-large", - "512", + str(m_large), "-n", "1024", "-k", @@ -293,6 +303,55 @@ def test_one_instance_serves_a_smaller_m_after_a_larger_one(): assert r[key] < FP8_FLOOR, (key, r) +def test_the_payload_map_is_the_same_for_every_m(): + """What the test above checks on hardware, as arithmetic. + + Regression guard for the aliasing itself: without capacity_m every region + below the control area sits at a different byte for every M. + """ + + def cfg(m): + c = ArConfig( + world_size=2, + m=m, + n=1024, + recv_slots=2, + counter_chunks=counter_chunks(m, 2), + counter_capacity=MAX_CHUNKS, + counter_shape_slots=8, + capacity_m=4096, + ) + c.validate() + return c + + maps = { + m: ( + cfg(m).input_off, + cfg(m).output_off, + cfg(m).tmp_off, + cfg(m).recv_off, + cfg(m).recv_slot_off(1), + cfg(m).window_bytes, + ) + for m in (512, 1024, 2048, 4096) + } + assert len(set(maps.values())) == 1, maps + + # And without it they all differ, which is the bug this guards. + loose = { + m: ArConfig( + world_size=2, + m=m, + n=1024, + recv_slots=2, + counter_chunks=counter_chunks(m, 2), + counter_capacity=MAX_CHUNKS, + ).recv_slot_off(1) + for m in (512, 1024, 2048, 4096) + } + assert len(set(loose.values())) == 4, loose + + @requires_two_gpus def test_changing_operands_between_calls_stays_correct(): """Five calls at one M, different data each time. @@ -309,39 +368,6 @@ def test_changing_operands_between_calls_stays_correct(): assert r[key] < FP8_FLOOR, (key, r) -@requires_two_gpus -def test_growing_m_is_refused(): - """Growing M within one instance is a known defect, so it must not run. - - After a call at some M, the first call at a *larger* M comes back wrong - about half the time -- always the same value, with the repeat correct. - 8 runs each on 2 ranks: - - 512 -> 1024 4 of 8 wrong - 1024 -> 512 0 of 8 wrong - 2048 -> 4096 also wrong, so it is the growth, not the chunk count - 512 x5, different operands each call all correct - - The last line rules out a stale read of a peer's slice at a fixed shape, and - the per-shape counter sets tested above rule out a counter residue. Root - cause still open; until then the op raises instead of returning the wrong - answer, and this test pins that behaviour. - """ - r = _run_worker( - 2, - "growing_m", - "--m-small", - "512", - "--m-large", - "1024", - "-n", - "1024", - "-k", - "512", - ) - assert r["refused"], r - - @requires_two_gpus def test_close_releases_and_is_idempotent(): """The communicator holds the handles, so dropping the op is not enough.""" From 465547babfb4dec13d67ffc3f168095aa54c3ef1 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 15 Sep 2026 03:37:22 +0000 Subject: [PATCH 73/79] docs: the fp8 gather does not pay in the server, and the pull costs End-to-end regression after the op-layer fixes, same protocol as the page describes, two or three runs each. GPU busy: unfused 1095.5ms (1100.2 1095.1 1091.2) fused, bf16 1021.0 (1018.7 1025.1 1019.2) fused, fp8/sdma 1026.5 (1025.9 1027.1) fused, fp8/lsa 1049.3 (1048.3 1050.2) This page's layer table has fp8/lsa at 949us against bf16's 1149 at the same shape, so fp8/lsa should have been the best of the three fused rows and is the worst. The likely reason is rank skew -- a pull needs the peer's slice finished at the moment it reads, a push lets each producer send as soon as its own band is done -- and only the benchmark has every rank in lockstep on an idle box. SGLang now defaults the transport to sdma; GemmAllReduceOp keeps lsa, which is right for the benchmark it was measured in. Fusing itself is worth more than the old capture said: -6.8% against -2.2%. The unfused row reproduces, the bf16 fused row does not and is 56ms better. That measurement predates the move to mori.ops.gemm_ar and several fixes here; the cause was not isolated, and the page says so rather than picking one. Pointers added at the headline and at the transport table, because a reader who stops at either would take away the opposite conclusion. Co-Authored-By: Claude Opus 5 (1M context) --- docs/MORI-GEMM-AR-BENCHMARK.md | 54 ++++++++++++++++++++++++++-------- 1 file changed, 42 insertions(+), 12 deletions(-) diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md index d8667adc3..6ee6f110b 100644 --- a/docs/MORI-GEMM-AR-BENCHMARK.md +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -75,6 +75,11 @@ Each run prints a `RESULT_JSON` line with `max_rank_time_us`, `rel_l2` and Fusing is worth **-22%** at M=16384; the fp8 gather a further **-15%**. +> The fp8 row does **not** survive to the server: end to end it is no better +> than the bf16 wire and the LSA pull is 28ms worse. See +> [the end-to-end numbers](#reproducing-the-end-to-end-numbers). A layer +> measured in isolation is measured without the rank skew it will meet. + 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**. @@ -131,6 +136,12 @@ layer. A CU pull has those bytes in registers already. | fp8 / sdma | 1018.9 | | **fp8 / lsa** | **957.3** | +This ordering reverses in the server: 1021.0ms of GPU busy for bf16 against +1026.5 for fp8/sdma and 1049.3 for fp8/lsa. A pull needs the peer's slice +finished at the moment it reads; a push lets each producer send as soon as its +own band is done, and only the benchmark has every rank in lockstep. SGLang +defaults the transport to `sdma` for that reason. + 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 @@ -405,17 +416,36 @@ gen(main); post("/stop_profile") GPU busy time over that capture: +GPU busy, means over two or three runs of the protocol above, each row's +individual runs in brackets: + | | GPU busy | vs unfused | |---|---:|---:| -| unfused (GEMM + NCCL) | 1101.9 ms | — | -| fused, bf16 wire | 1077.2 | -2.2% | -| fused, fp8 / sdma | 1050.4 | -4.7% | -| fused, fp8 / lsa | **1048.2** | **-4.9%** | - -The three fused rows are one capture session; the unfused baseline is a separate -one, where the bf16 fused configuration measured 1071.7 ms rather than 1077.2. -Read the `-2.2%` as approximate and the differences between the fused rows as -the reliable part. - -The layer-level win is much larger than the end-to-end one because `wo_b` is -about 12% of the profile. +| unfused (GEMM + NCCL) | 1095.5 ms (1100.2 1095.1 1091.2) | — | +| **fused, bf16 wire** | **1021.0** (1018.7 1025.1 1019.2) | **-6.8%** | +| fused, fp8 / sdma | 1026.5 (1025.9 1027.1) | -6.3% | +| fused, fp8 / lsa | 1049.3 (1048.3 1050.2) | -4.2% | + +The request's own wall time agrees: 1.1854s unfused against 1.0975 fused. + +**The fp8 gather does not pay end to end, and the pull costs.** That inverts this +page's own layer table, where fp8/lsa is 949us against bf16's 1149 at the model +shape -- fp8/lsa should have been the best row and is the worst of the three +fused ones. The likely reason is rank skew: a pull needs the peer's slice +finished at the moment it reads, while a push lets each producer send as soon as +its own band is done. The layer benchmark runs every rank in lockstep on an idle +box; a prefill does not. SGLang therefore defaults the transport to `sdma`, +while `GemmAllReduceOp` keeps `lsa` -- that default is right for the benchmark +it was measured in, and wrong for the server. + +Treat it as the general caution: every number above this section is a layer in +isolation, and a collective measured that way is measured without the skew it +will actually meet. + +An earlier capture of the same three configurations read 1101.9 / 1077.2 / +1050.4 / 1048.2. The unfused row reproduces; the bf16 fused row does not, and is +56ms better here. That measurement predates the move to `mori.ops.gemm_ar` and +several fixes in this branch, and the cause was not isolated. + +The layer-level win is larger than the end-to-end one because `wo_b` is about +12% of the profile. From 3bde626cf5d0937ca0be17db238d6b249827e0cd Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 15 Sep 2026 06:41:28 +0000 Subject: [PATCH 74/79] gemm_ar: stop guessing what an [M, K/128] scale means, and retract the doc's e2e claims Two corrections, one of them to a bug I introduced yesterday. **_flatten_a_scale must not infer the layout from the shape.** It treated a 2-D [M, K/128] tensor as logical and transposed it. That is right for the column-major one the review reported, and wrong for the one aiter_per1x128_quant(transpose_scale=True) returns, which writes K/128-major data and keeps the (M, K/128) shape -- shape and stride are identical in both cases and only the producer knows which. Guessing took SGLang from relL2 3.9e-3 to 0.36. It now accepts 1-D, [K/128, M], and [M, K/128] column-major, and rejects [M, K/128] row-major with a message naming both readings; SGLang passes .reshape(-1) explicitly. Verified on the server's own dumped operands at [11264, 7168] K=2048 on 8 ranks: flat 3.868e-3, ten repeats identical; the transposed spelling 0.365. **The end-to-end section is restored.** Yesterday's rewrite -- "fusing is worth -6.8%", "the fp8 gather does not pay", "the LSA pull costs 28ms" -- was measured against a mori built with BUILD_CCO_SDMA=OFF, where every put silently does nothing. The all-reduce then returns mostly the local slice, the model still answers fluently, every mori kernel still shows up in the profile, and the fused path measures faster than it is because it moves no data. fp8/lsa looked worst of the three for the same reason: the pull is the one leg that does not use SDMA, so it was the only one actually working. Re-measured on a build that has it: unfused 1096.0ms wall 1.1969s fused, bf16 1070.3 1.1728 -2.3% fused, fp8/sdma 1052.1 1.1455 -4.0% fused, fp8/lsa 1041.2 1.1385 -5.0% which is the page's original 1101.9 / 1077.2 / 1050.4 / 1048.2 to about half a percent, in the same order. The page now leads with how to notice the failure: perplexity, 862511 against 3.26, because generated text and the kernel trace both look fine, and a short prompt never enters the path at all (fusing needs M >= 4096). 75 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- docs/MORI-GEMM-AR-BENCHMARK.md | 62 +++++++++++------------------ python/mori/ops/gemm_ar/op.py | 60 +++++++++++++++++----------- tests/python/cco/test_gemm_ar_op.py | 30 ++++++++++++-- 3 files changed, 87 insertions(+), 65 deletions(-) diff --git a/docs/MORI-GEMM-AR-BENCHMARK.md b/docs/MORI-GEMM-AR-BENCHMARK.md index 6ee6f110b..a5601cea4 100644 --- a/docs/MORI-GEMM-AR-BENCHMARK.md +++ b/docs/MORI-GEMM-AR-BENCHMARK.md @@ -75,10 +75,8 @@ Each run prints a `RESULT_JSON` line with `max_rank_time_us`, `rel_l2` and Fusing is worth **-22%** at M=16384; the fp8 gather a further **-15%**. -> The fp8 row does **not** survive to the server: end to end it is no better -> than the bf16 wire and the LSA pull is 28ms worse. See -> [the end-to-end numbers](#reproducing-the-end-to-end-numbers). A layer -> measured in isolation is measured without the rank skew it will meet. +> 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 @@ -136,11 +134,8 @@ layer. A CU pull has those bytes in registers already. | fp8 / sdma | 1018.9 | | **fp8 / lsa** | **957.3** | -This ordering reverses in the server: 1021.0ms of GPU busy for bf16 against -1026.5 for fp8/sdma and 1049.3 for fp8/lsa. A pull needs the peer's slice -finished at the moment it reads; a push lets each producer send as soon as its -own band is done, and only the benchmark has every rank in lockstep. SGLang -defaults the transport to `sdma` for that reason. +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 @@ -416,36 +411,27 @@ gen(main); post("/stop_profile") GPU busy time over that capture: -GPU busy, means over two or three runs of the protocol above, each row's -individual runs in brackets: +GPU busy over that capture: -| | GPU busy | vs unfused | -|---|---:|---:| -| unfused (GEMM + NCCL) | 1095.5 ms (1100.2 1095.1 1091.2) | — | -| **fused, bf16 wire** | **1021.0** (1018.7 1025.1 1019.2) | **-6.8%** | -| fused, fp8 / sdma | 1026.5 (1025.9 1027.1) | -6.3% | -| fused, fp8 / lsa | 1049.3 (1048.3 1050.2) | -4.2% | - -The request's own wall time agrees: 1.1854s unfused against 1.0975 fused. - -**The fp8 gather does not pay end to end, and the pull costs.** That inverts this -page's own layer table, where fp8/lsa is 949us against bf16's 1149 at the model -shape -- fp8/lsa should have been the best row and is the worst of the three -fused ones. The likely reason is rank skew: a pull needs the peer's slice -finished at the moment it reads, while a push lets each producer send as soon as -its own band is done. The layer benchmark runs every rank in lockstep on an idle -box; a prefill does not. SGLang therefore defaults the transport to `sdma`, -while `GemmAllReduceOp` keeps `lsa` -- that default is right for the benchmark -it was measured in, and wrong for the server. - -Treat it as the general caution: every number above this section is a layer in -isolation, and a collective measured that way is measured without the skew it -will actually meet. - -An earlier capture of the same three configurations read 1101.9 / 1077.2 / -1050.4 / 1048.2. The unfused row reproduces; the bf16 fused row does not, and is -56ms better here. That measurement predates the move to `mori.ops.gemm_ar` and -several fixes in this branch, and the cause was not isolated. +| | 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/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index c367e53de..b0ea99966 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -183,22 +183,27 @@ def supports( 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`` -- the layout - ``gemm_a8w8_blockscale_bpreshuffle`` consumes, and what sglang's - ``materialize_bpreshuffle_fp8_scale`` produces. Three spellings are accepted - and they mean different things: - - * ``[M, K/128]`` -- the logical shape. Transposed here, which is free for the - column-major tensor the quantiser returns (``.t()`` is then contiguous) and - a copy for a row-major one. **Passing this through ``reshape(-1)`` instead - is the bug this function exists to prevent**: for the column-major case - that walks the logical rows, pairing every scale with the wrong K block, - which still validates and still runs (relL2 0.26 against 2.4e-3). - * ``[K/128, M]`` -- the physical shape, taken as-is. - * 1-D -- already flattened in physical order, taken as-is. - - ``m == kb`` would be ambiguous; the logical reading wins, and K would have to - be ``128 * M`` for it to arise. + 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: @@ -208,16 +213,25 @@ def _flatten_a_scale(a_scale: torch.Tensor, m: int, kb: int) -> torch.Tensor: return a_scale if a_scale.dim() != 2: raise ValueError( - f"a_scale must be [M, K/128], [K/128, M] or flat, got shape " - f"{tuple(a_scale.shape)}" + f"a_scale must be [K/128, M], [M, K/128] column-major, or flat; got " + f"shape {tuple(a_scale.shape)}" ) - if tuple(a_scale.shape) == (m, kb): - return a_scale.t().reshape(-1) - if tuple(a_scale.shape) == (kb, m): + 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 {tuple(a_scale.shape)} matches neither [M, K/128] = " - f"{(m, kb)} nor its transpose" + f"a_scale shape {shape} matches neither [M, K/128] = {(m, kb)} nor its " + f"transpose" ) diff --git a/tests/python/cco/test_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py index 3a31dc885..4b6940f70 100644 --- a/tests/python/cco/test_gemm_ar_op.py +++ b/tests/python/cco/test_gemm_ar_op.py @@ -172,22 +172,21 @@ def test_validate_rejects_misspelled_modes(field, value): def test_flatten_a_scale_preserves_physical_order(): - """The kernel reads (row, kb) at kb*M + row, whatever spelling came in.""" + """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], physically [kb, M] + 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, 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 exact regression: reshape(-1) on the column-major tensor.""" + """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() @@ -195,6 +194,29 @@ def test_flatten_a_scale_is_not_a_reshape(): 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): From f03469fcfa40352b6e4e7e65906e26fe501f1fbc Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 15 Sep 2026 07:30:10 +0000 Subject: [PATCH 75/79] ci: teach the Triton ABI tests about the two new symbols, and run the op tests cco_lsa_win_base and cco_lsa_stride went into CCO_DEVICE_FUNCTIONS without the two tests that pin that surface: the count tripwire (68 -> 70) and the flydsl-vs-triton ABI comparison, whose extern list had to name them. Both are doing their job -- the surface changed and they said so; I had only run the gemm_ar files locally. Also wires tests/python/cco/test_gemm_ar_op.py into the numerics step. It was added in e3514e25 and never run by CI, which is the same gap it exists to close: the other two files drive the benchmark, which calls the low-level builders, so nothing covered the public op. Raises that step's timeout to 1800s -- it spawns several 2-rank workers with an 20s settle between them. Ran the CCO steps that work here: triton bindings/lsa/sdma 11 passed, rocm import order 2 passed, sdma api 12 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci_cco.yml | 8 +++++++- tests/python/cco/test_triton_bindings.py | 4 +++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci_cco.yml b/.github/workflows/ci_cco.yml index 65f5284b5..e78f9d184 100644 --- a/.github/workflows/ci_cco.yml +++ b/.github/workflows/ci_cco.yml @@ -184,10 +184,16 @@ jobs: # 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 1200 python -m pytest -v \ + 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) 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, From c3177cba9b22cbbf5fde378a379606ea52bdeb01 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Tue, 15 Sep 2026 12:52:55 +0000 Subject: [PATCH 76/79] gemm_ar: give every legal M a counter slot, instead of rationing eight max_shapes defaulted to 8. Each distinct M needs its own counter set -- the election modulus depends on M, so two shapes sharing a set elect on each other's residue -- and the ninth M raised. A server's M changes with the batch, so under real traffic it raised within seconds: mori fused wo_b failed and is disabled for this process; falling back to the split path: this op already serves 8 distinct M values [4096, 5120, 6144, ... and SGLang turns any exception from this path into a permanent fallback, so the whole optimisation switched itself off. Two thirds of a 128k-prefix benchmark matrix measured the unfused path without saying so. The default is now every M the window can legally take: a padded M is a multiple of world_size*block_m and at most m_max, so m_max/(world_size*block_m) slots can never be exhausted by a legal call -- 16 at the model shape, covering the 13 values from 4096 to 16384. A slot 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 and cost the feature. 75 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/op.py | 38 +++++++++++++++++++++++++++-------- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index b0ea99966..c8669cef6 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -56,10 +56,23 @@ # (relL2 1.41 against 2.4e-3 at K=256). MIN_K = 2 * BLOCK_K -# Distinct M values one instance can serve. Each needs its own counter set -- -# see ArConfig.counter_shape_slots -- and each also costs a FlyDSL compile, so -# this is a small number by nature rather than a tuning knob. -DEFAULT_MAX_SHAPES = 8 + +# 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: @@ -166,7 +179,9 @@ def supports( gather_dtype=gather_dtype, gather_transport=gather_transport, counter_capacity=MAX_CHUNKS, - shape_slots=DEFAULT_MAX_SHAPES, + shape_slots=default_max_shapes( + padded_m(m, world_size, block_m), world_size, block_m + ), shape_index=0, ) except ValueError: @@ -271,7 +286,7 @@ def __init__( sdma_queues: int = 1, gather_dtype: str = "bf16", gather_transport: str = "lsa", - max_shapes: int = DEFAULT_MAX_SHAPES, + max_shapes: Optional[int] = None, ): self._closed = True # so a failed constructor leaves close() a no-op if gather_dtype not in WIRE_DTYPES: @@ -288,6 +303,10 @@ def __init__( f"block_n must be {DEFAULT_BLOCK_N}: the op always compiles with " f"permlane, whose lane transpose is written for that width" ) + if max_shapes is None: + max_shapes = default_max_shapes( + padded_m(m_max, comm.nranks, block_m), comm.nranks, block_m + ) if block_m < 1 or m_max < 1 or sdma_queues < 1 or max_shapes < 1: raise ValueError( f"block_m, m_max, sdma_queues and max_shapes must all be >= 1; got " @@ -362,7 +381,7 @@ def window_bytes_for( block_m: int = DEFAULT_BLOCK_M, gather_dtype: str = "bf16", gather_transport: str = "lsa", - max_shapes: int = DEFAULT_MAX_SHAPES, + max_shapes: Optional[int] = None, ) -> int: """Symmetric-window bytes an op for this shape will allocate. @@ -378,7 +397,10 @@ def window_bytes_for( gather_dtype=gather_dtype, gather_transport=gather_transport, counter_capacity=MAX_CHUNKS, - shape_slots=max_shapes, + 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 From 1531cfe8ddd23abed2897182b7cc73ccf73b11c3 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 16 Sep 2026 03:34:40 +0000 Subject: [PATCH 77/79] gemm_ar: prove the transport moves bytes before trusting anything it measures A mori built without BUILD_CCO_SDMA=ON -- the default, since the flag only flips with BUILD_BENCHMARK -- compiles the SDMA puts out. Every symbol is still there, every kernel still launches, every put still returns, and nothing moves. The all-reduce then yields 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 data. An end-to-end campaign read -6.8% instead of -2.3% that way, and the only thing that caught it was perplexity, 862511 against 3.26. GemmAllReduceOp.self_test() writes rank+1 over the input region, runs the split order (scatter -> reduce -> gather; fused_order opens with drain, which assumes a GEMM epilogue already pushed), and checks every element came back as sum(1..world_size). With the puts inert each rank sees only its own share and the check fails. It runs at m_max, compiling the shape the first real call needs anyway, and reads mori.cco.device._build_flags.BUILD_CCO_SDMA only to make the error message name the cause. The test asserts both halves, because a check that cannot fail is what let this happen. Zeroing the landing slots before the negative half is not a detail: the op has served other cases by then, so the slots still hold a previous call's peer data, reducing that produces a plausible sum, and the teeth test passes while testing nothing -- the same stale-read effect that made the dead transport look correct. Also in this commit, because the tests needed it: * All six GPU cases now come from one 2-rank spawn behind a module fixture. Six back-to-back launches outran the SDMA queue reclaim (hsaKmtCreateQueueExt, anvil.cpp:237) and the casualties were as often the neighbouring file's tests as this file's, since the cost falls on whoever launches next. The file also drops from 305s to 177s. * A spawn that fails with that signature is retried once, with a longer settle. Nothing else is retried. * counter_chunks() raises with the shape it was given instead of "max() iterable argument is empty" when m_pad is below one row band. 75 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/op.py | 77 +++++++++++++ tests/python/cco/gemm_ar_op_worker.py | 96 +++++++++++++++- tests/python/cco/test_gemm_ar_op.py | 151 +++++++++++++------------- 3 files changed, 249 insertions(+), 75 deletions(-) diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index c8669cef6..8e28c12e1 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -93,6 +93,12 @@ def counter_chunks(m_pad: int, world_size: int, block_m: int = DEFAULT_BLOCK_M) 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) @@ -262,6 +268,9 @@ class _Plan(NamedTuple): 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: @@ -473,6 +482,7 @@ def _compiled(self, m: int): hit = _Plan( gemm=gemm, tail=tail, + parts=parts, input=from_gpu_ptr( self.mem.ptr + cfg.input_off, (m, self.n), torch.bfloat16 ), @@ -483,6 +493,73 @@ def _compiled(self, m: int): 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. diff --git a/tests/python/cco/gemm_ar_op_worker.py b/tests/python/cco/gemm_ar_op_worker.py index 926ed475a..a758cfcff 100644 --- a/tests/python/cco/gemm_ar_op_worker.py +++ b/tests/python/cco/gemm_ar_op_worker.py @@ -205,7 +205,53 @@ def case_alternating_m(op, rank, world, m_small, m_large, n, k): (seen.float() - stale.float()).abs().max() ) results[f"{label}_rows"] = rep - _emit(rank, case="alternating_m", **results) + _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): @@ -246,6 +292,43 @@ def case_changing_data(op, rank, world, m, n, k, calls): _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) @@ -258,17 +341,26 @@ def main() -> int: args = p.parse_args() rank, world, uid = _setup() - m_max = max(args.m, args.m_large, args.m_small) + 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": diff --git a/tests/python/cco/test_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py index 4b6940f70..f836c7b74 100644 --- a/tests/python/cco/test_gemm_ar_op.py +++ b/tests/python/cco/test_gemm_ar_op.py @@ -239,7 +239,25 @@ def test_padded_m_rounds_to_whole_bands(): _SETTLE_SECONDS = 20 -def _run_worker(world_size: int, case: str, *extra: str, timeout: int = 900): +#: 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") @@ -263,14 +281,22 @@ def _run_worker(world_size: int, case: str, *extra: str, timeout: int = 900): timeout=timeout, ) output = result.stdout + result.stderr - assert result.returncode == 0, output + 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[0] + 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( @@ -278,14 +304,34 @@ def _run_worker(world_size: int, case: str, *extra: str, timeout: int = 900): ) +@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(): +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 = _run_worker(2, "scale_order", "-m", "512", "-n", "1024", "-k", "512") + r = _case(worker_results, "scale_order") assert r["rel_l2"] < FP8_FLOOR, r assert r["flat_matches"], r assert r["physical_matches"], r @@ -293,7 +339,9 @@ def test_public_op_takes_the_column_major_scale(): @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(m_small, m_large): +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 @@ -309,73 +357,13 @@ def test_one_instance_serves_two_m_values_in_any_order(m_small, m_large): payload to the instance's m_max. Every shape now has the same map, so a region only ever aliases itself. """ - r = _run_worker( - 2, - "alternating_m", - "--m-small", - str(m_small), - "--m-large", - str(m_large), - "-n", - "1024", - "-k", - "512", - ) + 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) -def test_the_payload_map_is_the_same_for_every_m(): - """What the test above checks on hardware, as arithmetic. - - Regression guard for the aliasing itself: without capacity_m every region - below the control area sits at a different byte for every M. - """ - - def cfg(m): - c = ArConfig( - world_size=2, - m=m, - n=1024, - recv_slots=2, - counter_chunks=counter_chunks(m, 2), - counter_capacity=MAX_CHUNKS, - counter_shape_slots=8, - capacity_m=4096, - ) - c.validate() - return c - - maps = { - m: ( - cfg(m).input_off, - cfg(m).output_off, - cfg(m).tmp_off, - cfg(m).recv_off, - cfg(m).recv_slot_off(1), - cfg(m).window_bytes, - ) - for m in (512, 1024, 2048, 4096) - } - assert len(set(maps.values())) == 1, maps - - # And without it they all differ, which is the bug this guards. - loose = { - m: ArConfig( - world_size=2, - m=m, - n=1024, - recv_slots=2, - counter_chunks=counter_chunks(m, 2), - counter_capacity=MAX_CHUNKS, - ).recv_slot_off(1) - for m in (512, 1024, 2048, 4096) - } - assert len(set(loose.values())) == 4, loose - - @requires_two_gpus -def test_changing_operands_between_calls_stays_correct(): +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 @@ -383,15 +371,32 @@ def test_changing_operands_between_calls_stays_correct(): previous call's bytes -- bit-identical to the right answer. This is the only thing in the suite that would notice. """ - r = _run_worker( - 2, "changing_data", "-m", "512", "-n", "1024", "-k", "512", "--calls", "5" - ) + 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_close_releases_and_is_idempotent(): +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 = _run_worker(2, "close", "-n", "1024", "-k", "512", "-m", "512") + r = _case(worker_results, "close") assert r["refused_after_close"], r From f4caced0d66db2defc9da8721b8432e6a6ce49e1 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 16 Sep 2026 10:06:09 +0000 Subject: [PATCH 78/79] gemm_ar: reject FNUZ and illegal tiles at the boundary, and fix the rollback Four findings from the second review, all reproduced. **FNUZ is not the encoding this kernel implements.** FP8_DTYPES admitted both e4m3 variants on the theory that the launch only sees bytes and both are one byte wide. The width is the same; the exponent bias is not -- 8 against OCP's 7 -- and the MMA atom is fixed to fx.Float8E4M3FN. Every A and every B element therefore reads at twice its value and the product comes back 4x too large: finite, plausible, and wrong. The reviewer measured exactly that, 4096 where 1024 was expected, relL2 ~3. sglang reaches this on gfx942, where normalize_e4m3fn_to_e4m3fnuz converts weights to FNUZ, so the accepting-and-corrupting version was reachable from a real caller. Accept only float8_e4m3fn. The comment justifying the old tuple was the actual defect -- it argued from element size, which is evidence of nothing. **The rollback could not run.** __init__'s except block calls close(), whose first statement is self._cache.clear(), but _cache was assigned after the try block. Any acquisition failure died in cleanup with AttributeError, demoting the real error to __context__ and skipping the handle-release loop the rollback exists for. Fault injection at create_dev_comm showed 2 resources acquired and 0 closed. Initialise all cleanup state before the first thing that can raise. **Tile constraints were not shared.** The kernel asserts BLOCK_M >= 128 and BLOCK_M % 128 == 0 (bare, so an empty AssertionError), while the op only checked block_m >= 1. block_m=64 passed supports(), constructed, acquired all three resources, and failed on the first real call. block_m=0 was worse: default_max_shapes and padded_m both divide by world_size * block_m, and they ran before the positivity check, so it raised ZeroDivisionError from inside the padding arithmetic. _tile_constraints() now states the granules once and runs first in both supports() and __init__, before anything divides. supports() also answers False for world_size < 1 for the same reason. **README.** Communicator.init takes the UniqueId wrapper, not the bytes that survive a broadcast, so the quickstart raised TypeError as written. Three other places had drifted from the code: the a_scale contract (row- major [M,K/128] is rejected now, and the reason is that the two 2-D spellings are indistinguishable by shape yet need opposite treatment), the max_shapes default (every legal M, not 8), and close() (the dev-comm is deliberately left to the communicator). The example now clones the result before the context frees the window behind it. Tests: the tile boundary at 64/1/192/0/-128 and block_n 128/384/0, that 256x512 is still accepted, world_size 0 and -1, that FNUZ is not in FP8_DTYPES, and a fault injection at each of the three acquisition steps asserting the original error propagates and everything acquired before it was released. The injection stubs from_gpu_ptr: the fake handles carry a null pointer and the window-zeroing write faults the context otherwise, which then surfaces in whatever runs next. 44 pass in test_gemm_ar_op.py, including the 2-rank GPU cases. test_flydsl_ar.py::test_flydsl_ar_matches_aiter_bitwise fails here, but it fails identically on the unmodified tree. Co-Authored-By: Claude Opus 5 (1M context) --- python/mori/ops/gemm_ar/README.md | 53 +++++++++---- python/mori/ops/gemm_ar/op.py | 83 +++++++++++++++++---- tests/python/cco/test_gemm_ar_op.py | 112 ++++++++++++++++++++++++++++ 3 files changed, 219 insertions(+), 29 deletions(-) diff --git a/python/mori/ops/gemm_ar/README.md b/python/mori/ops/gemm_ar/README.md index 31037cdb0..eaeddf60e 100644 --- a/python/mori/ops/gemm_ar/README.md +++ b/python/mori/ops/gemm_ar/README.md @@ -42,7 +42,7 @@ all-reduce quietly produces zeros. ```python import torch, torch.distributed as dist -from mori.cco import Communicator +from mori.cco import Communicator, UniqueId from mori.ops.gemm_ar import GemmAllReduceOp, preshuffle_b rank, world = dist.get_rank(), dist.get_world_size() @@ -56,26 +56,40 @@ 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) -with Communicator.init(world, rank, uid[0], per_rank_vmm=vmm) as comm: +# 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. Pass it as `[M, K/128]` -in either memory order, as `[K/128, M]`, or already flat in physical order -- -the op distinguishes these. What it cannot distinguish is a `[M, K/128]` tensor -you have already flattened yourself with `reshape(-1)`: for the column-major -tensor `aiter_per1x128_quant(transpose_scale=True)` returns, that walks the -logical rows and pairs every scale with the wrong K block, which validates, -runs, and returns relL2 0.26. `b_scale` is `[N/128, K/128]` fp32 row-major. +**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 @@ -91,12 +105,21 @@ 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` (8). Warm each M up once before capturing a graph. - -**Cleanup.** `close()` (or the `with` above) releases the window, memory and -dev-comm. The communicator holds a strong reference to each, so dropping the op -is not enough -- at the model shape that is about 700 MiB per rank. Synchronise -first if anything may still be in flight. +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 diff --git a/python/mori/ops/gemm_ar/op.py b/python/mori/ops/gemm_ar/op.py index 8e28c12e1..134aab8a5 100644 --- a/python/mori/ops/gemm_ar/op.py +++ b/python/mori/ops/gemm_ar/op.py @@ -42,6 +42,16 @@ 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 @@ -102,6 +112,28 @@ def counter_chunks(m_pad: int, world_size: int, block_m: int = DEFAULT_BLOCK_M) 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: @@ -174,6 +206,9 @@ def supports( 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: @@ -195,10 +230,20 @@ def supports( return True -#: fp8 encodings the kernel's byte reinterpretation is valid for. Both are 1 -#: byte and the GEMM only ever sees the bytes, but forwarding a bf16 tensor here -#: reinterprets two rows as one and produces finite nonsense. -FP8_DTYPES = (torch.float8_e4m3fnuz, torch.float8_e4m3fn) +#: 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: @@ -297,7 +342,15 @@ def __init__( 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)}" @@ -312,16 +365,22 @@ def __init__( 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 block_m < 1 or m_max < 1 or sdma_queues < 1 or max_shapes < 1: - raise ValueError( - f"block_m, m_max, sdma_queues and max_shapes must all be >= 1; got " - f"block_m={block_m} m_max={m_max} sdma_queues={sdma_queues} " - f"max_shapes={max_shapes}" - ) + 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}") @@ -348,7 +407,6 @@ def __init__( #: are disjoint, so two M values never elect on each other's residue. self._shape_slot: dict[int, int] = {} - self.mem = self.win = self.dev_comm = None self.window_bytes = self._make_cfg(self.m_max).window_bytes try: self.mem = comm.alloc_mem(self.window_bytes) @@ -378,9 +436,6 @@ def __init__( raise self._closed = False - self._cache: dict[int, _Plan] = {} - self._pad_in: Optional[torch.Tensor] = None - @staticmethod def window_bytes_for( world_size: int, diff --git a/tests/python/cco/test_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py index f836c7b74..9325e5e85 100644 --- a/tests/python/cco/test_gemm_ar_op.py +++ b/tests/python/cco/test_gemm_ar_op.py @@ -44,6 +44,7 @@ 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, @@ -229,6 +230,117 @@ def test_padded_m_rounds_to_whole_bands(): 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 -------------------------------------------- From dc8ba7fbb6dd4f8d2dea523de76237c49a51e533 Mon Sep 17 00:00:00 2001 From: "Xiang.Chen" Date: Wed, 16 Sep 2026 10:17:38 +0000 Subject: [PATCH 79/79] gemm_ar: format the new assert the way black 25.1.0 wants it pre-commit's only complaint on f4caced. The repo pins black 25.1.0, which still wraps the condition rather than the message on a multi-line assert. Co-Authored-By: Claude Opus 5 (1M context) --- tests/python/cco/test_gemm_ar_op.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/python/cco/test_gemm_ar_op.py b/tests/python/cco/test_gemm_ar_op.py index 9325e5e85..08d7e5599 100644 --- a/tests/python/cco/test_gemm_ar_op.py +++ b/tests/python/cco/test_gemm_ar_op.py @@ -336,9 +336,9 @@ def create_dev_comm(self, reqs): 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}" - ) + assert ( + set(acquired[:-1]) == closed + ), f"failing at {fail_at}: acquired {acquired}, closed {closed}" # --- multi-rank: the op itself --------------------------------------------