diff --git a/CHANGELOG.md b/CHANGELOG.md index 41c4921..2c81404 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +## [0.4.18] + +### Fixed +- On GPUs of two dies such as M5 Ultra, `hc_front_expand_collapse` could read stale values from + other threadgroups and give garbage GLM-5.3 decode; its handoff is now device coherent. +- Hyper-connection kernels capped below 64 threads per threadgroup raise an error instead of hanging. + ## [0.4.17] ### Added diff --git a/metal/mlx/backend/metal/kernels/kq_hc_glue.h b/metal/mlx/backend/metal/kernels/kq_hc_glue.h index 460da5b..25d8911 100644 --- a/metal/mlx/backend/metal/kernels/kq_hc_glue.h +++ b/metal/mlx/backend/metal/kernels/kq_hc_glue.h @@ -88,6 +88,52 @@ inline float kq_hc_tg_sum( acc = fma(a3.z, b3.z, acc); \ acc = fma(a3.w, b3.w, acc); +// Pointer types of the buffers the fused kernel hands from one threadgroup +// to another (h, the mix dots, the sum of squares). By default device +// memory has threadgroup coherence: a plain store is visible only to the +// threads of its own threadgroup. A coherent(device) pointer makes its +// loads and stores visible across the threadgroups of a dispatch once +// they are synchronized (MSL spec 4.8, Memory Coherency). The split +// kernels hand off at a dispatch boundary and keep plain pointers, so the +// bodies take the qualifier as a parameter (C) and every access in them +// goes through these types. A vector of four 16-bit values (T4) crosses as +// its raw bits in a uint2, with put on the store and get on the load: the +// paravirtual GPU of the macos-26 CI runners lowers coherent accesses of +// bfloat vectors wrongly and those of 32-bit elements right. +template +struct kq_hc_ptr { + typedef device T* type; + typedef const device T* ctype; + static inline T put(T v) { + return v; + } + static inline T get(T v) { + return v; + } +}; +template +struct kq_hc_ptr { + typedef device coherent(device) T* type; + typedef const device coherent(device) T* ctype; + static inline T put(T v) { + return v; + } + static inline T get(T v) { + return v; + } +}; +template +struct kq_hc_ptr { + typedef device coherent(device) uint2* type; + typedef const device coherent(device) uint2* ctype; + static inline uint2 put(T v) { + return as_type(v); + } + static inline T get(uint2 v) { + return as_type(v); + } +}; + template [[kernel]] void kq_hc_front_reduce( const device T* x [[buffer(0)]], @@ -145,17 +191,18 @@ template } // Shared body of kq_hc_front_expand_reduce, also run by every threadgroup -// of the fused kernel. partial is threadgroup scratch [KQ_HC_NSG_MAX]. -template +// of the fused kernel, with coherent(device) outputs there (C). +// partial is threadgroup scratch [KQ_HC_NSG_MAX]. +template inline void kq_hc_front_expand_reduce_body( const device T* x_sub, const device T* resid, const device float* post, const device float* comb, const device float* fn, - device T* h_out, - device float* mixes_raw, - device float* sumsq, + typename kq_hc_ptr::type h_out, + typename kq_hc_ptr::type mixes_raw, + typename kq_hc_ptr::type sumsq, const int D, uint row, uint m, @@ -169,7 +216,7 @@ inline void kq_hc_front_expand_reduce_body( const device T* xs = x_sub + (int64_t)row * D; const device T* rr = resid + (int64_t)row * KTOT; - device T* hout = h_out + (int64_t)row * KTOT; + typename kq_hc_ptr::type hout = h_out + (int64_t)row * KTOT; const uint pb = row * 4, cb = row * 16; const float p0 = post[pb + 0], p1 = post[pb + 1]; @@ -183,7 +230,7 @@ inline void kq_hc_front_expand_reduce_body( using T4 = vec; const device T4* xs4 = (const device T4*)xs; const device T4* rr4 = (const device T4*)rr; - device T4* h4 = (device T4*)hout; + typename kq_hc_ptr::type h4 = (typename kq_hc_ptr::type)hout; const device float4* f4 = (const device float4*)(fn + (int64_t)m * KTOT); float acc = 0.0f; @@ -231,10 +278,10 @@ inline void kq_hc_front_expand_reduce_body( float4 f3 = f4[3 * D4q + d4]; KQ_HC_DOT16(acc, h0, h1, h2, h3, f0, f1, f2, f3) } else { - h4[0 * D4q + d4] = hv0; - h4[1 * D4q + d4] = hv1; - h4[2 * D4q + d4] = hv2; - h4[3 * D4q + d4] = hv3; + h4[0 * D4q + d4] = kq_hc_ptr::put(hv0); + h4[1 * D4q + d4] = kq_hc_ptr::put(hv1); + h4[2 * D4q + d4] = kq_hc_ptr::put(hv2); + h4[3 * D4q + d4] = kq_hc_ptr::put(hv3); KQ_HC_DOT16(acc, h0, h1, h2, h3, h0, h1, h2, h3) } } @@ -264,7 +311,7 @@ template uint tid [[thread_position_in_threadgroup]], uint nt [[threads_per_threadgroup]]) { threadgroup float partial[KQ_HC_NSG_MAX]; - kq_hc_front_expand_reduce_body( + kq_hc_front_expand_reduce_body( x_sub, resid, post, @@ -286,12 +333,13 @@ template // the streams with the pre gates, which every thread derives on its own // from the mixes, while the last simdgroup runs the sinkhorn (post and // comb). ssq_shared is [KQ_HC_NSG_MAX] and inv_shared [1] of threadgroup -// scratch. -template +// scratch. The fused kernel reads x, mixes_raw and sumsq through +// coherent(device) pointers (C). +template inline void kq_hc_sinkhorn_collapse_body( - const device T* x, - const device float* mixes_raw, - const device float* sumsq, + typename kq_hc_ptr::ctype x, + typename kq_hc_ptr::ctype mixes_raw, + typename kq_hc_ptr::ctype sumsq, const device float* scale, const device float* base, const device T* w, @@ -317,20 +365,21 @@ inline void kq_hc_sinkhorn_collapse_body( const bool worker = tid < NC; const bool spare = (sg == NSG - 1); - const device float* mix = mixes_raw + row * KQ_HC_MIX; + typename kq_hc_ptr::ctype mix = mixes_raw + row * KQ_HC_MIX; device float* post_out = post + row * KQ_HC; device float* comb_out = comb + row * KQ_HC * KQ_HC; const float factor = metal::rsqrt(sumsq[row] / (float)(KQ_HC * D) + NEPS); - const device T* x_row = x + (int64_t)row * (KQ_HC * D); + typename kq_hc_ptr::ctype x_row = x + (int64_t)row * (KQ_HC * D); device T* out_row = collapsed + (int64_t)row * D; using T4 = vec; - const device T4* x_row0 = (const device T4*)(x_row + 0 * D); - const device T4* x_row1 = (const device T4*)(x_row + 1 * D); - const device T4* x_row2 = (const device T4*)(x_row + 2 * D); - const device T4* x_row3 = (const device T4*)(x_row + 3 * D); + using X4 = typename kq_hc_ptr::ctype; + X4 x_row0 = (X4)(x_row + 0 * D); + X4 x_row1 = (X4)(x_row + 1 * D); + X4 x_row2 = (X4)(x_row + 2 * D); + X4 x_row3 = (X4)(x_row + 3 * D); device T4* out4 = (device T4*)out_row; const uint D4 = (uint)D / 4; @@ -359,9 +408,10 @@ inline void kq_hc_sinkhorn_collapse_body( post_out[lane] = post_v; } - float4 v = - (*(const device float4*)(mix + BASE_OFF + llane * KQ_HC) * comb_scale + - *(const device float4*)(base + BASE_OFF + llane * KQ_HC)) * + float4 v = (*(typename kq_hc_ptr::ctype)( + mix + BASE_OFF + llane * KQ_HC) * + comb_scale + + *(const device float4*)(base + BASE_OFF + llane * KQ_HC)) * active; float row_max = metal::max(metal::max(v.x, v.y), metal::max(v.z, v.w)); @@ -397,10 +447,10 @@ inline void kq_hc_sinkhorn_collapse_body( for (uint c = 0; c < chunks; ++c) { uint d4 = c * NC + tid; if (worker && d4 < D4) { - float4 x0 = float4(x_row0[d4]); - float4 x1 = float4(x_row1[d4]); - float4 x2 = float4(x_row2[d4]); - float4 x3 = float4(x_row3[d4]); + float4 x0 = float4(kq_hc_ptr::get(x_row0[d4])); + float4 x1 = float4(kq_hc_ptr::get(x_row1[d4])); + float4 x2 = float4(kq_hc_ptr::get(x_row2[d4])); + float4 x3 = float4(kq_hc_ptr::get(x_row3[d4])); float4 result = fma(float4(p0), x0, @@ -420,10 +470,10 @@ inline void kq_hc_sinkhorn_collapse_body( for (uint c = 0; c < chunks; ++c) { uint d4 = c * NC + tid; if (worker && d4 < D4) { - float4 x0 = float4(x_row0[d4]); - float4 x1 = float4(x_row1[d4]); - float4 x2 = float4(x_row2[d4]); - float4 x3 = float4(x_row3[d4]); + float4 x0 = float4(kq_hc_ptr::get(x_row0[d4])); + float4 x1 = float4(kq_hc_ptr::get(x_row1[d4])); + float4 x2 = float4(kq_hc_ptr::get(x_row2[d4])); + float4 x3 = float4(kq_hc_ptr::get(x_row3[d4])); float4 result = fma(float4(p0), x0, @@ -456,7 +506,7 @@ template uint nt [[threads_per_threadgroup]]) { threadgroup float ssq_shared[KQ_HC_NSG_MAX]; threadgroup float inv_shared[1]; - kq_hc_sinkhorn_collapse_body( + kq_hc_sinkhorn_collapse_body( x, mixes_raw, sumsq, @@ -484,9 +534,11 @@ template // over the mixes and h the others published. Nothing waits, so the grid // needs no co-residency and any row count works; the counter is zero on // entry (the host clears a fresh one per dispatch) and is left at -// KQ_HC_MIX + 1. The collapsing threadgroup reads h and the mixes after -// a device-scope fence on the acquiring side that matches the release -// fence each publisher issues before its increment. +// KQ_HC_MIX + 1. h, the mixes and the sum of squares pass between +// threadgroups through coherent(device) pointers (MSL spec 4.8): a plain +// device store is visible only within its own threadgroup, which on a GPU +// of two dies can leave the collapse reading stale values. Every thread +// fences at device scope after it publishes and before it reads. template [[kernel]] void kq_hc_front_expand_collapse( const device T* x_sub [[buffer(0)]], @@ -497,9 +549,9 @@ template const device float* scale [[buffer(5)]], const device float* base [[buffer(6)]], const device T* w [[buffer(7)]], - device T* h_out [[buffer(8)]], - device float* mixes_raw [[buffer(9)]], - device float* sumsq [[buffer(10)]], + device coherent(device) T* h_out [[buffer(8)]], + device coherent(device) float* mixes_raw [[buffer(9)]], + device coherent(device) float* sumsq [[buffer(10)]], device atomic_uint* arrive [[buffer(11)]], device T* collapsed [[buffer(12)]], device float* post_out [[buffer(13)]], @@ -514,12 +566,14 @@ template threadgroup float partial[KQ_HC_NSG_MAX]; threadgroup float ssq_shared[KQ_HC_NSG_MAX]; threadgroup float inv_shared[1]; - threadgroup uint last_flag; + // An array, which tid 0 writes before the barrier that precedes every + // read: as a scalar the compiler sees the paths that skip the write. + threadgroup uint last_flag[1]; const uint row = tg / (KQ_HC_MIX + 1); const uint m = tg % (KQ_HC_MIX + 1); - kq_hc_front_expand_reduce_body( + kq_hc_front_expand_reduce_body( x_sub, resid, post_in, @@ -535,24 +589,32 @@ template nt, partial); - // publish this threadgroup's mix dot (or h and sumsq), then arrive + // Publish this threadgroup's mix dot (or h and sumsq), then arrive. + // A seq_cst fence is a release and an acquire fence (MSL spec 6.16.3); + // the acquire and release orders on a fence need Metal 4.1 (6.16.1). + // The barrier orders every thread's stores before tid 0, which fences + // before its fetch_add (release). The tid 0 of the last group fences + // after the fetch_add (acquire), and the barrier carries that to the + // other threads before they read. threadgroup_barrier(mem_flags::mem_device); if (tid == 0) { - atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst); + atomic_thread_fence( + mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); uint prev = atomic_fetch_add_explicit(arrive + row, 1u, memory_order_relaxed); - last_flag = (prev == (uint)KQ_HC_MIX) ? 1u : 0u; - if (last_flag != 0u) { - atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst); + const uint last = (prev == (uint)KQ_HC_MIX) ? 1u : 0u; + if (last != 0u) { + atomic_thread_fence( + mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); } + last_flag[0] = last; } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (last_flag == 0u) { + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + if (last_flag[0] == 0u) { return; } - threadgroup_barrier(mem_flags::mem_device); - kq_hc_sinkhorn_collapse_body( + kq_hc_sinkhorn_collapse_body( h_out, mixes_raw, sumsq, diff --git a/mlx_kquant/_version.py b/mlx_kquant/_version.py index ac15521..2dd5536 100644 --- a/mlx_kquant/_version.py +++ b/mlx_kquant/_version.py @@ -1 +1 @@ -__version__ = "0.4.17" +__version__ = "0.4.18" diff --git a/src/kquant_hc_glue.cpp b/src/kquant_hc_glue.cpp index 60091f7..b4458ff 100644 --- a/src/kquant_hc_glue.cpp +++ b/src/kquant_hc_glue.cpp @@ -94,10 +94,12 @@ namespace { // simdgroup. Register pressure puts the pipeline limit below the preferred // count on some GPUs, and an oversized dispatch is silent garbage rather // than an error (the same trap kq_sdpa_vector guards), so the kernels read -// their own [[threads_per_threadgroup]] and stride by it. KQ_HC_NT_LOG=1 -// prints the limit and the choice once per kernel. +// their own [[threads_per_threadgroup]] and stride by it. least is the +// smallest width the kernel runs at. KQ_HC_NT_LOG=1 prints the limit and +// the choice once per kernel. template -inline int kq_hc_threads(K kernel, const std::string& kname, int want) { +inline int +kq_hc_threads(K kernel, const std::string& kname, int want, int least) { int cap = int(kernel->maxTotalThreadsPerThreadgroup()); static const int forced = []() { const char* e = std::getenv("KQ_HC_NT"); @@ -107,10 +109,11 @@ inline int kq_hc_threads(K kernel, const std::string& kname, int want) { cap = std::min(cap, forced); } const int nt = (std::min(want, cap) / 32) * 32; - if (nt < 32) { + if (nt < least) { throw std::runtime_error( "[mlx_kquant.hc] " + kname + " supports only " + std::to_string(cap) + - " threads per threadgroup, below the one simdgroup it needs."); + " threads per threadgroup, below the " + std::to_string(least) + + " it needs."); } static const bool log = std::getenv("KQ_HC_NT_LOG") != nullptr; if (log) { @@ -136,7 +139,9 @@ inline int kq_hc_threads(K kernel, const std::string& kname, int want) { // hc_front_expand_collapse claims against the split pair. Register // pressure differs per kernel, so the caps can differ on one GPU. Take // the smallest cap in the family and give it to all four. hc_expand -// reduces nothing and keeps its own width. +// reduces nothing and keeps its own width. The collapse gives one +// simdgroup to the sinkhorn and the others to the columns, so the family +// needs two simdgroups. inline int kq_hc_family_threads(mx::metal::Device& d, const std::string& t) { static std::mutex mu; static std::map cached; @@ -153,8 +158,8 @@ inline int kq_hc_family_threads(mx::metal::Device& d, const std::string& t) { int nt = kHcThreads; for (const char* base : family) { const std::string kname = base + t; - nt = - std::min(nt, kq_hc_threads(kq_get_kernel(d, kname), kname, kHcThreads)); + nt = std::min( + nt, kq_hc_threads(kq_get_kernel(d, kname), kname, kHcThreads, 64)); } cached.emplace(t, nt); return nt; @@ -333,7 +338,7 @@ void KQuantHcExpand::eval_gpu( ce.set_input_array(inputs[3], 3); ce.set_output_array(out, 4); ce.set_bytes(D, 5); - const int nt = kq_hc_threads(kernel, kname, kHcExpandThreads); + const int nt = kq_hc_threads(kernel, kname, kHcExpandThreads, 32); ce.dispatch_threadgroups(MTL::Size(rows * 2, 1, 1), MTL::Size(nt, 1, 1)); } diff --git a/tests/test_hc_glue.py b/tests/test_hc_glue.py index 16b6329..03fd225 100644 --- a/tests/test_hc_glue.py +++ b/tests/test_hc_glue.py @@ -13,6 +13,8 @@ from __future__ import annotations import os +import subprocess +import sys import mlx.core as mx import numpy as np @@ -237,6 +239,117 @@ def test_front_expand_collapse_multi_row(rows, D): assert np.array_equal(np.array(cb), np.array(cb_ref)) +# GLM-5.3's hyper-connection decode: hidden 4096, 20 sinkhorn iterations, +# norm eps 1e-5, and 92 cycles per token. +GLM_D, GLM_NORM_EPS, GLM_CYCLES = 4096, 1e-5, 92 + + +def _glm_inputs(rows, dtype, seed): + ks = mx.random.split(mx.random.key(seed), 8) + x_sub = (mx.random.normal((1, rows, GLM_D), key=ks[0]) * 0.5).astype(dtype) + resid = (mx.random.normal((1, rows, HC, GLM_D), key=ks[1]) * 2.0).astype(dtype) + post = 2 * mx.sigmoid(mx.random.normal((1, rows, HC), key=ks[2])) + c = mx.softmax(mx.random.normal((1, rows, HC, HC), key=ks[3]), axis=-1) + comb = c / c.sum(axis=-2, keepdims=True) + fn = mx.random.normal((MIX, HC * GLM_D), key=ks[4]) * 0.02 + scale = mx.array([0.3, 0.2, 0.25], dtype=mx.float32) + base = mx.random.normal((MIX,), key=ks[5]) * 0.5 + w = (1.0 + 0.1 * mx.random.normal((GLM_D,), key=ks[6])).astype(dtype) + mx.eval(x_sub, resid, post, comb, fn, scale, base, w) + return x_sub, resid, post, comb, fn, scale, base, w + + +def _glm_fused(x_sub, resid, post, comb, fn, scale, base, w): + return kq.hc_front_expand_collapse( + x_sub, + resid, + post, + comb, + fn, + scale, + base, + w, + iters=ITERS, + hc_eps=HC_EPS, + norm_eps=GLM_NORM_EPS, + ) + + +def _glm_pair(x_sub, resid, post, comb, fn, scale, base, w): + h, mixes, ssq = kq.hc_front_expand_reduce(x_sub, resid, post, comb, fn) + x, p, cb = kq.hc_sinkhorn_collapse( + h, + mixes, + ssq, + scale, + base, + w, + iters=ITERS, + hc_eps=HC_EPS, + norm_eps=GLM_NORM_EPS, + ) + return h, x, p, cb + + +def _glm_chain(fun, args): + # The cycles of one token in one evaluation, each fed by the last; the + # collapsed output stands in for the sublayer's. + x_sub, resid, post, comb, fn, scale, base, w = args + outs = [] + for _ in range(GLM_CYCLES): + h, x, p, cb = fun(x_sub, resid, post, comb, fn, scale, base, w) + outs.append((h, x, p, cb)) + x_sub, resid, post, comb = x, h, p, cb + mx.eval(outs) + return outs + + +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) +def test_front_expand_collapse_glm53_chain(rows, dtype): + # At GLM-5.3's decode shapes the fused cycle equals the split pair bit + # for bit, every cycle of a token's chain. + args = _glm_inputs(rows, dtype, 211 + rows) + fused = _glm_chain(_glm_fused, args) + pair = _glm_chain(_glm_pair, args) + for i, (f, p) in enumerate(zip(fused, pair, strict=True)): + for name, a, b in zip(("h", "x", "post", "comb"), f, p, strict=True): + assert mx.array_equal(a, b).item(), (i, name) + + +_NT_CHILD = r""" +import mlx.core as mx +import mlx_kquant as kq +x = mx.zeros((1, 1, 4, 1024), dtype=mx.bfloat16) +fn = mx.zeros((24, 4096), dtype=mx.float32) +try: + mx.eval(kq.hc_front_reduce(x, fn)) + print("no error") +except Exception as e: + print(f"{type(e).__name__}: {e}") +""" + + +@pytest.mark.parametrize("nt", ["32", "63"]) +def test_family_width_floor(nt): + # The collapse gives one simdgroup to the sinkhorn and the rest to the + # columns, so a launch of one simdgroup would have no column workers. + # Capped below two simdgroups, the family raises before any launch. + env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + env["KQ_HC_NT"] = nt + r = subprocess.run( + [sys.executable, "-c", _NT_CHILD], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert r.returncode == 0, r.stderr[-4000:] + out = r.stdout.strip().splitlines()[-1] + assert out.startswith("RuntimeError: [mlx_kquant.hc] kq_hc_front_reduce_"), out + assert "below the 64 it needs" in out, out + + def test_input_validation(): x, fn, scale, base, w = _mk(17, 2048) with pytest.raises(ValueError): diff --git a/tests/test_hc_glue_ir.py b/tests/test_hc_glue_ir.py new file mode 100644 index 0000000..c1a1e3b --- /dev/null +++ b/tests/test_hc_glue_ir.py @@ -0,0 +1,115 @@ +"""The handoff buffers of the fused hyper-connection cycle in the compiled IR. + +kq_hc_front_expand_collapse hands h, the mix dots and the sum of squares +from every threadgroup of a row to the last one to arrive, inside one +dispatch. By default device memory has threadgroup coherence (MSL spec +4.8, Memory Coherency), so those accesses must be coherent(device), which +the compiler lowers to air.store.device_coherent and +air.load.device_coherent. h crosses as raw bits in uint2 vectors: the +paravirtual GPU of the macos-26 CI runners lowers the coherent forms of +bfloat vectors wrongly. The split kernels hand off at a dispatch +boundary and keep plain accesses. Compiles metal/kq_hc_glue.metal to LLVM +IR with xcrun metal, as the build does; skips where that is missing. + +Usage: test_hc_glue_ir.py +""" + +from __future__ import annotations + +import os +import re +import shutil +import subprocess + +import mlx.core as mx +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SOURCE = os.path.join(ROOT, "metal", "kq_hc_glue.metal") +MLX_INCLUDE = os.path.join(os.path.dirname(mx.__file__), "include") + +_DEFINE = re.compile(r"^define [^@]*@([\w.$]+)\(") +_CALL = re.compile(r"@([\w.$]+)\(") +_COHERENT_BF16 = re.compile(r"@air\.(?:load|store)\.device_coherent\.[\w.]*bf16[\w.]*") + + +def _metal(): + if shutil.which("xcrun") is None or not os.path.exists(SOURCE): + return None + r = subprocess.run( + ["xcrun", "-sdk", "macosx", "--find", "metal"], + capture_output=True, + text=True, + ) + return r.stdout.strip() if r.returncode == 0 else None + + +@pytest.fixture(scope="module") +def defines(tmp_path_factory): + """Function name -> body lines of the compiled kq_hc_glue.metal.""" + metal = _metal() + if metal is None: + pytest.skip("needs xcrun metal and the metal/ sources") + out = tmp_path_factory.mktemp("ir") / "kq_hc_glue.ll" + # The build's Metal deployment target (setup.py). + env = dict(os.environ, MACOSX_DEPLOYMENT_TARGET="26.2") + cmd = [metal, "-S", "-emit-llvm", "-I", os.path.join(ROOT, "metal")] + cmd += ["-I", MLX_INCLUDE, "-I", os.path.join(MLX_INCLUDE, "metal_cpp")] + cmd += ["-Wno-c++17-extensions", SOURCE, "-o", str(out)] + r = subprocess.run(cmd, capture_output=True, text=True, env=env) + assert r.returncode == 0, r.stderr[-4000:] + funcs, name = {}, None + for line in out.read_text().splitlines(): + m = _DEFINE.match(line) + if m: + name = m.group(1) + funcs[name] = [] + elif line == "}": + name = None + elif name is not None: + funcs[name].append(line) + return funcs + + +def _reach(funcs, kernel): + # The kernel's body and the bodies of the defined functions it calls. + seen, todo, lines = set(), [kernel], [] + while todo: + f = todo.pop() + if f in seen or f not in funcs: + continue + seen.add(f) + lines += funcs[f] + todo += [c for line in funcs[f] for c in _CALL.findall(line)] + return "\n".join(lines) + + +def _kernels(funcs, prefix): + names = [f for f in funcs if re.fullmatch(rf"{prefix}_(bfloat16|float16)_t", f)] + assert len(names) == 2, (prefix, names) + return names + + +def test_fused_handoff_is_device_coherent(defines): + # Both the stores that publish and the loads of the collapse. + for name in _kernels(defines, "kq_hc_front_expand_collapse"): + body = _reach(defines, name) + assert "@air.store.device_coherent." in body, name + assert "@air.load.device_coherent." in body, name + + +def test_fused_handoff_has_no_bfloat_coherent_access(defines): + # Every coherent access of h is a bitcast to uint2, for every T. + for name in _kernels(defines, "kq_hc_front_expand_collapse"): + body = _reach(defines, name) + bad = set(_COHERENT_BF16.findall(body)) + assert not bad, (name, sorted(bad)) + + +@pytest.mark.parametrize( + "prefix", + ["kq_hc_front_reduce", "kq_hc_front_expand_reduce", "kq_hc_sinkhorn_collapse"], +) +def test_split_kernels_stay_plain(defines, prefix): + for name in _kernels(defines, prefix): + assert "device_coherent" not in _reach(defines, name), name