# Copyright 2026 Jakevin # SPDX-License-Identifier: Apache-2.0 """Ternary GEMV and GEMM for MLX. Codes {-1,0,1} are stored as 2-bit values {0,1,2}, four per byte, low bits first. One scale per 64 weights (or 128, as an option) and no zero point: value = (code - 1) * scale. The inner loop is the MLX 2-bit mask dot: pre-scale x by 1, 1/4, 1/16, 1/64 and AND the byte. Small batches use a simdgroup GEMV. K multiple of 512 follows the MLX qmv_fast walk (two rows per simdgroup, 16 weights per lane). Other K, including 640, uses a 4-wide dot. Batch >= 32 with the row count and K divisible by 32 uses a 32x32x32 simdgroup-matrix GEMM; if the batch is not a multiple of 32, the leftover rows go through the GEMV. The weight tile is half, the accumulator is float, and a tile never crosses a scale group. Batches up to 8 run through a small LRU of compiled graphs keyed by shape. Weights are passed as arguments, so the cache holds no weight arrays. Larger batches dispatch directly. """ from __future__ import annotations import mlx.core as mx from .pack import GROUP, GROUPS def set_cache_limit(gb: float = 2) -> int: """Opt-in: cap the MLX cache at ``gb`` gigabytes. Returns the previous limit. Importing the package does not change the MLX cache limit. """ return mx.set_cache_limit(int(gb * 1024**3)) def configure() -> int: """Cap the MLX cache at 2GB. Same as ``set_cache_limit(2)``.""" return set_cache_limit(2) _KERNELS: dict = {} _META: dict = {} QMV_ROWS = 2 QMV_NSG = 1 GEMM_BM = 32 GEMM_BN = 32 GEMM_BK = 32 GEMM_MIN_BATCH = 32 _HEADER = r""" #include #include #include using namespace metal; // code*x for four 2-bit lanes packed in one byte. x1/x2/x3 are already // divided by 4, 16 and 64 so the masked field does not need a shift. #define TERN_DOT4(b, x0, x1, x2, x3) \ (((x0) * float((b) & 3u) + (x1) * float((b) & 0x0cu)) \ + ((x2) * float((b) & 0x30u) + (x3) * float((b) & 0xc0u))) #define TERN_DOT4H(b, x0, x1, x2, x3) \ (((x0) * half((b) & 3u) + (x1) * half((b) & 0x0cu)) \ + ((x2) * half((b) & 0x30u) + (x3) * half((b) & 0xc0u))) """ def metal(name: str, source: str, header: str = _HEADER): # MLX 0.22 metal_kernel has no compile_options. The cache key does not # include the source, so a body change must use a new name. key = (name, header[:40]) if key not in _KERNELS: _KERNELS[key] = mx.fast.metal_kernel( name=name, input_names=["x", "codes", "scales", "meta"], output_names=["y"], header=header, source=source, ensure_row_contiguous=False, ) return _KERNELS[key] def meta_u32(values: list[int]) -> mx.array: key = tuple(int(v) for v in values) buf = _META.get(key) if buf is None: buf = mx.array(list(key), dtype=mx.uint32) mx.eval(buf) _META[key] = buf return buf def _prepare(x, codes, scales, k, group): if group not in GROUPS: raise ValueError(f"group {group} not in {GROUPS}") if k % group != 0: raise ValueError(f"K {k} is not a multiple of {group}") if x.ndim == 1: x = x.reshape(1, -1) if not isinstance(x, mx.array): x = mx.array(x) if x.dtype != mx.float32: x = x.astype(mx.float32) if not isinstance(codes, mx.array): codes = mx.array(codes) if codes.dtype != mx.uint8: codes = codes.astype(mx.uint8) if not isinstance(scales, mx.array): scales = mx.array(scales) if int(x.shape[-1]) != k: raise ValueError(f"x last dim {x.shape[-1]} != K {k}") # The kernels index raw row-major buffers (ensure_row_contiguous=False), # so transposed or sliced views must be copied first. No-op if contiguous. x = mx.contiguous(x) codes = mx.contiguous(codes) scales = mx.contiguous(scales) return x, codes, scales, int(x.shape[0]), int(codes.shape[0]) # Kernel sources are written for group 64. Group 128 only changes how a # weight index maps to its scale: k >> 7 instead of k >> 6, and one more # lane bit per scale. Group 64 sources are left byte-identical. _REGROUP_128 = ( ("k >> 6", "k >> 7"), ("t >> 6", "t >> 7"), ("(lane >> 2)", "(lane >> 3)"), ("(lane >> 4)", "(lane >> 5)"), ("((k0 + wc) >> 6)", "((k0 + wc) >> 7)"), ) def _regroup(source: str, group: int) -> str: if group == 64: return source if group != 128: raise ValueError(f"group {group} not in {GROUPS}") for old, new in _REGROUP_128: source = source.replace(old, new) if ">> 6" in source: raise AssertionError("group 128 source still indexes scales by 64") return source def _gname(group: int) -> str: return "" if group == 64 else f"_gs{group}" def _dot16_expr(w: str, vecdot: bool = False) -> str: if not vecdot: return ( f"TERN_DOT4({w} & 0xffu, xv0, xv1, xv2, xv3)" f" + TERN_DOT4(({w} >> 8u) & 0xffu, xv4, xv5, xv6, xv7)" f" + TERN_DOT4(({w} >> 16u) & 0xffu, xv8, xv9, xv10, xv11)" f" + TERN_DOT4(({w} >> 24u) & 0xffu, xv12, xv13, xv14, xv15)" ) # Packed dot so the four masked lanes stay in one float4. def one(byte: str, xs: str) -> str: return ( f"dot(float4(float({byte} & 3u), float({byte} & 0x0cu), " f"float({byte} & 0x30u), float({byte} & 0xc0u)), float4({xs}))" ) return ( f"{one(f'({w} & 0xffu)', 'xv0, xv1, xv2, xv3')}" f" + {one(f'(({w} >> 8u) & 0xffu)', 'xv4, xv5, xv6, xv7')}" f" + {one(f'(({w} >> 16u) & 0xffu)', 'xv8, xv9, xv10, xv11')}" f" + {one(f'(({w} >> 24u) & 0xffu)', 'xv12, xv13, xv14, xv15')}" ) def _qmv_source(nsg: int, rows: int, guarded: bool, blocks: int | None = None, vecdot: bool = False, narrow: bool = False, half_dot: bool = False, chunk128: bool = False) -> str: if rows not in (2, 4, 8) or nsg not in (1, 2, 4): raise ValueError(f"qmv rows={rows} nsg={nsg}") if blocks is not None and blocks < 1: raise ValueError(f"qmv blocks={blocks}") acc_init = "\n".join(f" float acc{r} = 0.0f;" for r in range(rows)) def w16_loads() -> str: lines = [] for r in range(rows): addr = ( f"*reinterpret_cast(codes + (row0 + {r}u) * row_bytes + byte0 + (lane << 2))" ) if guarded: lines.append(f"uint w{r} = 0u; if (row0 + {r}u < N) w{r} = {addr};") else: lines.append(f"uint w{r} = {addr};") return "\n ".join(lines) def w4_loads() -> str: lines = [] for r in range(rows): addr = f"uint(codes[(row0 + {r}u) * row_bytes + byte0 + lane])" if guarded: lines.append(f"uint c{r} = 0u; if (row0 + {r}u < N) c{r} = {addr};") else: lines.append(f"uint c{r} = {addr};") return "\n ".join(lines) def apply(scale_shift: int, src: str) -> str: lines = [] for r in range(rows): fma = ( f"acc{r} = fma(float(scales[(row0 + {r}u) * n_groups + g0 + (lane >> {scale_shift})]), " f"d{r} - sumx, acc{r});" ) if guarded: lines.append(f"if (row0 + {r}u < N) {{ {fma} }}") else: lines.append(fma) return "\n ".join(lines) def tail_body() -> str: lines = [ "uint sh = (t & 3u) * 2u;", "uint byte = t >> 2;", "uint g = t >> 6;", ] for r in range(rows): body = ( f"uint code = (uint(codes[(row0 + {r}u) * row_bytes + byte]) >> sh) & 3u;" f"acc{r} = fma((float(code) - 1.0f) * float(scales[(row0 + {r}u) * n_groups + g]), xv, acc{r});" ) if guarded: lines.append(f"if (row0 + {r}u < N) {{ {body} }}") else: lines.append("{ " + body + " }") return "\n ".join(lines) red = "\n ".join(f"float red{r} = simd_sum(acc{r});" for r in range(rows)) writes = [] for r in range(rows): cond = f"lane == 0u && row0 + {r}u < N" if guarded else "lane == 0u" writes.append(f"if ({cond}) y[b * N + row0 + {r}u] = red{r};") store = red + "\n " + "\n ".join(writes) dots16 = "\n ".join( f"float d{r} = {_dot16_expr(f'w{r}', vecdot)};" for r in range(rows) ) dots4 = "\n ".join( f"float d{r} = TERN_DOT4(c{r}, xv0, xv1, xv2, xv3);" for r in range(rows) ) step512 = f""" const device float* xp = x + b * K + k + (lane << 4); float4 a0 = *reinterpret_cast(xp); float4 a1 = *reinterpret_cast(xp + 4); float4 a2 = *reinterpret_cast(xp + 8); float4 a3 = *reinterpret_cast(xp + 12); float sumx = ((a0.x + a0.y) + (a0.z + a0.w)) + ((a1.x + a1.y) + (a1.z + a1.w)); sumx += ((a2.x + a2.y) + (a2.z + a2.w)) + ((a3.x + a3.y) + (a3.z + a3.w)); float xv0 = a0.x; float xv1 = a0.y * 0.25f; float xv2 = a0.z * 0.0625f; float xv3 = a0.w * 0.015625f; float xv4 = a1.x; float xv5 = a1.y * 0.25f; float xv6 = a1.z * 0.0625f; float xv7 = a1.w * 0.015625f; float xv8 = a2.x; float xv9 = a2.y * 0.25f; float xv10 = a2.z * 0.0625f; float xv11 = a2.w * 0.015625f; float xv12 = a3.x; float xv13 = a3.y * 0.25f; float xv14 = a3.z * 0.0625f; float xv15 = a3.w * 0.015625f; uint byte0 = k >> 2; uint g0 = k >> 6; {w16_loads()} {dots16} {apply(2, "w")}""" if narrow: # One float4 of x is live at a time so the large-N grid can stay resident. shifts = (0, 8, 16, 24) nlines = [ "const device float* xp = x + b * K + k + (lane << 4);", "uint byte0 = k >> 2;", "uint g0 = k >> 6;", w16_loads(), ] for r in range(rows): nlines.append(f"float part{r} = 0.0f;") for bi, sh in enumerate(shifts): nlines.append("{") nlines.append( f"float4 a = *reinterpret_cast(xp + {bi * 4});" ) nlines.append("float sum4 = (a.x + a.y) + (a.z + a.w);") nlines.append( "float u0 = a.x; float u1 = a.y * 0.25f; " "float u2 = a.z * 0.0625f; float u3 = a.w * 0.015625f;" ) for r in range(rows): expr = ( f"part{r} += TERN_DOT4((w{r} >> {sh}u) & 0xffu, u0, u1, u2, u3) - sum4;" ) if guarded: nlines.append(f"if (row0 + {r}u < N) {{ {expr} }}") else: nlines.append(expr) nlines.append("}") for r in range(rows): fma = ( f"acc{r} = fma(float(scales[(row0 + {r}u) * n_groups + g0 + (lane >> 2)]), " f"part{r}, acc{r});" ) if guarded: nlines.append(f"if (row0 + {r}u < N) {{ {fma} }}") else: nlines.append(fma) step512 = "\n ".join(nlines) if half_dot: # Half products, float partial. One float4 of x is live per byte. shifts = (0, 8, 16, 24) hlines = [ "const device float* xp = x + b * K + k + (lane << 4);", "uint byte0 = k >> 2;", "uint g0 = k >> 6;", w16_loads(), ] for r in range(rows): hlines.append(f"float part{r} = 0.0f;") for bi, sh in enumerate(shifts): hlines.append("{") hlines.append( f"float4 a = *reinterpret_cast(xp + {bi * 4});" ) hlines.append( "half hx = half(a.x); half hy = half(a.y); " "half hz = half(a.z); half hw = half(a.w);" ) hlines.append("half sum4 = (hx + hy) + (hz + hw);") hlines.append( "half u0 = hx; half u1 = hy * 0.25h; " "half u2 = hz * 0.0625h; half u3 = hw * 0.015625h;" ) for r in range(rows): expr = ( f"part{r} += float(TERN_DOT4H((w{r} >> {sh}u) & 0xffu, u0, u1, u2, u3)) " f"- float(sum4);" ) if guarded: hlines.append(f"if (row0 + {r}u < N) {{ {expr} }}") else: hlines.append(expr) hlines.append("}") for r in range(rows): fma = ( f"acc{r} = fma(float(scales[(row0 + {r}u) * n_groups + g0 + (lane >> 2)]), " f"part{r}, acc{r});" ) if guarded: hlines.append(f"if (row0 + {r}u < N) {{ {fma} }}") else: hlines.append(fma) step512 = "\n ".join(hlines) if chunk128: # Four values per lane. Fewer live registers than the 16-wide body. k_loops = ( " uint k = 0u;\n" " for (; k + 128u <= K; k += 128u) {\n" " const device float* xp = x + b * K + k + (lane << 2);\n" " float4 a0 = *reinterpret_cast(xp);\n" " float sumx = (a0.x + a0.y) + (a0.z + a0.w);\n" " float xv0 = a0.x; float xv1 = a0.y * 0.25f; float xv2 = a0.z * 0.0625f; float xv3 = a0.w * 0.015625f;\n" " uint byte0 = k >> 2;\n" " uint g0 = k >> 6;\n" f" {w4_loads()}\n" f" {dots4}\n" f" {apply(4, 'c')}\n" " }\n" " for (uint t = k + lane; t < K; t += 32u) {\n" " float xv = x[b * K + t];\n" f" {tail_body()}\n" " }\n" ) elif blocks is None: # K is not a multiple of 512 (expert down is 640). Keep the short tail. k_loops = ( " uint k = 0u;\n" " for (; k + 512u <= K; k += 512u) {\n" f"{step512}\n" " }\n" " for (; k + 128u <= K; k += 128u) {\n" " const device float* xp = x + b * K + k + (lane << 2);\n" " float4 a0 = *reinterpret_cast(xp);\n" " float sumx = (a0.x + a0.y) + (a0.z + a0.w);\n" " float xv0 = a0.x; float xv1 = a0.y * 0.25f; float xv2 = a0.z * 0.0625f; float xv3 = a0.w * 0.015625f;\n" " uint byte0 = k >> 2;\n" " uint g0 = k >> 6;\n" f" {w4_loads()}\n" f" {dots4}\n" f" {apply(4, 'c')}\n" " }\n" " for (uint t = k + lane; t < K; t += 32u) {\n" " float xv = x[b * K + t];\n" f" {tail_body()}\n" " }\n" ) else: # Bake the trip count and drop the tail so the mask body can unroll. k_loops = ( " #pragma unroll\n" f" for (uint blk = 0u; blk < {int(blocks)}u; ++blk) {{\n" " uint k = blk * 512u;\n" f"{step512}\n" " }\n" ) return f""" constexpr uint NSG = {nsg}u; constexpr uint ROWS = {rows}u; uint lane = thread_index_in_simdgroup; uint sg = simdgroup_index_in_threadgroup; uint tile = threadgroup_position_in_grid.x; uint b = threadgroup_position_in_grid.y; uint N = meta[0]; uint K = meta[1]; uint row_bytes = meta[2]; uint n_groups = meta[3]; uint row0 = (tile * NSG + sg) * ROWS; {acc_init} {k_loops} {store} """ def _gemm_source(bm: int, bn: int, bk: int, half_w: bool) -> str: if bm % 16 or bn % 16 or bk % 32 or 64 % bk: raise ValueError(f"bad tile {bm} {bn} {bk}") nsg_n = bn // 16 nsg_m = bm // 16 tg = nsg_n * nsg_m * 32 per_w = (bn * bk) // tg per_x = (bm * bk) // tg if per_w % 4 or per_x % 4 or bk % per_w or bk % per_x or per_w > bk or per_x > bk: raise ValueError(f"tile {bm}x{bn}x{bk} does not split across {tg} threads") lda_x = bk + 4 lda_w = bk + (8 if half_w else 4) wtype = "half" if half_w else "float" nbytes = per_w // 4 if nbytes % 4 == 0: w_load = f""" uint e = tid * {per_w}u; uint wr = e / BK; uint wc = e - wr * BK; uint grow = n0 + wr; const device uint* wp = reinterpret_cast( codes + grow * row_bytes + ((k0 + wc) >> 2)); float s = float(scales[grow * n_groups + ((k0 + wc) >> 6)]); float s1 = s * 0.25f; float s2 = s * 0.0625f; float s3 = s * 0.015625f; #pragma unroll for (uint j = 0u; j < {nbytes // 4}u; ++j) {{ uint w32 = wp[j]; #pragma unroll for (uint i = 0u; i < 4u; ++i) {{ uint bcode = (w32 >> (i * 8u)) & 0xffu; float d0 = s * float(bcode & 3u) - s; float d1 = s1 * float(bcode & 0x0cu) - s; float d2 = s2 * float(bcode & 0x30u) - s; float d3 = s3 * float(bcode & 0xc0u) - s; uint kc = wc + (j * 4u + i) * 4u; *reinterpret_cast(&WBUF[wr][kc]) = {wtype}4({wtype}(d0), {wtype}(d1), {wtype}(d2), {wtype}(d3)); }} }} """ else: w_load = f""" uint e = tid * {per_w}u; uint wr = e / BK; uint wc = e - wr * BK; uint grow = n0 + wr; const device uint8_t* wp = codes + grow * row_bytes + ((k0 + wc) >> 2); float s = float(scales[grow * n_groups + ((k0 + wc) >> 6)]); float s1 = s * 0.25f; float s2 = s * 0.0625f; float s3 = s * 0.015625f; #pragma unroll for (uint i = 0u; i < {nbytes}u; ++i) {{ uint bcode = uint(wp[i]); float d0 = s * float(bcode & 3u) - s; float d1 = s1 * float(bcode & 0x0cu) - s; float d2 = s2 * float(bcode & 0x30u) - s; float d3 = s3 * float(bcode & 0xc0u) - s; uint kc = wc + i * 4u; *reinterpret_cast(&WBUF[wr][kc]) = {wtype}4({wtype}(d0), {wtype}(d1), {wtype}(d2), {wtype}(d3)); }} """ n4 = per_x // 4 x_load = f""" uint ex = tid * {per_x}u; uint xr = ex / BK; uint xc = ex - xr * BK; const device float* xp = x + (b0 + xr) * K + (k0 + xc); #pragma unroll for (uint i = 0u; i < {n4}u; ++i) {{ float4 v = *reinterpret_cast(xp + i * 4u); *reinterpret_cast(&XBUF[xr][xc + i * 4u]) = v; }} """ return f""" constexpr uint BM = {bm}u; constexpr uint BN = {bn}u; constexpr uint BK = {bk}u; constexpr uint LDAX = {lda_x}u; constexpr uint LDAW = {lda_w}u; uint tid = thread_position_in_threadgroup.x; uint sg = simdgroup_index_in_threadgroup; uint N = meta[0]; uint K = meta[1]; uint row_bytes = meta[2]; uint n_groups = meta[3]; uint B = meta[4]; uint n0 = threadgroup_position_in_grid.x * BN; uint b0 = threadgroup_position_in_grid.y * BM; threadgroup float XBUF[BM][LDAX]; threadgroup {wtype} WBUF[BN][LDAW]; uint sg_n = sg % {nsg_n}u; uint sg_m = sg / {nsg_n}u; uint mbase = sg_m * 16u; uint nbase = sg_n * 16u; ushort lane = thread_index_in_simdgroup; short qid = short(lane) >> 2; short fm = (qid & 4) + ((short(lane) >> 1) & 3); short fn = ((qid & 2) << 1) + ((short(lane) & 1) << 1); uint ufm = uint(fm); uint ufn = uint(fn); simdgroup_matrix c00(0.0f); simdgroup_matrix c01(0.0f); simdgroup_matrix c10(0.0f); simdgroup_matrix c11(0.0f); for (uint k0 = 0u; k0 < K; k0 += BK) {{ {x_load} {w_load} threadgroup_barrier(mem_flags::mem_threadgroup); #pragma unroll for (uint kk = 0u; kk < BK; kk += 8u) {{ simdgroup_matrix a0(0.0f); simdgroup_matrix a1(0.0f); simdgroup_matrix b0m(0.0f); simdgroup_matrix b1m(0.0f); float2 av0 = *reinterpret_cast(&XBUF[mbase + ufm][kk + ufn]); float2 av1 = *reinterpret_cast(&XBUF[mbase + 8u + ufm][kk + ufn]); a0.thread_elements()[0] = av0.x; a0.thread_elements()[1] = av0.y; a1.thread_elements()[0] = av1.x; a1.thread_elements()[1] = av1.y; b0m.thread_elements()[0] = float(WBUF[nbase + ufn][kk + ufm]); b0m.thread_elements()[1] = float(WBUF[nbase + ufn + 1u][kk + ufm]); b1m.thread_elements()[0] = float(WBUF[nbase + 8u + ufn][kk + ufm]); b1m.thread_elements()[1] = float(WBUF[nbase + 8u + ufn + 1u][kk + ufm]); simdgroup_multiply_accumulate(c00, a0, b0m, c00); simdgroup_multiply_accumulate(c01, a0, b1m, c01); simdgroup_multiply_accumulate(c10, a1, b0m, c10); simdgroup_multiply_accumulate(c11, a1, b1m, c11); }} threadgroup_barrier(mem_flags::mem_threadgroup); }} uint gm = b0 + mbase + ufm; uint gn = n0 + nbase + ufn; if (gm < B) {{ if (gn < N) y[gm * N + gn] = c00.thread_elements()[0]; if (gn + 1u < N) y[gm * N + gn + 1u] = c00.thread_elements()[1]; if (gn + 8u < N) y[gm * N + gn + 8u] = c01.thread_elements()[0]; if (gn + 9u < N) y[gm * N + gn + 9u] = c01.thread_elements()[1]; }} uint gm8 = gm + 8u; if (gm8 < B) {{ if (gn < N) y[gm8 * N + gn] = c10.thread_elements()[0]; if (gn + 1u < N) y[gm8 * N + gn + 1u] = c10.thread_elements()[1]; if (gn + 8u < N) y[gm8 * N + gn + 8u] = c11.thread_elements()[0]; if (gn + 9u < N) y[gm8 * N + gn + 9u] = c11.thread_elements()[1]; }} """ def _qmv_mlx_source(nsg: int, rows: int, guarded: bool) -> str: """MLX qmv_fast layout: 16 values per lane, one K loop, byte loads.""" if rows not in (2, 4, 8) or nsg not in (1, 2, 4): raise ValueError(f"qmv mlx rows={rows} nsg={nsg}") init = "\n ".join(f"float result{r} = 0.0f;" for r in range(rows)) dots = [] for r in range(rows): body = ( f"const device uint8_t* wl = codes + (out_row + {r}u) * row_bytes + (lane << 2) + (k >> 2);" f"float s = float(scales[(out_row + {r}u) * n_groups + (k >> 6) + (lane >> 2)]);" "float accum = 0.0f;" "accum += xt0 * float(wl[0] & 3u) + xt1 * float(wl[0] & 0x0cu);" "accum += xt2 * float(wl[0] & 0x30u) + xt3 * float(wl[0] & 0xc0u);" "accum += xt4 * float(wl[1] & 3u) + xt5 * float(wl[1] & 0x0cu);" "accum += xt6 * float(wl[1] & 0x30u) + xt7 * float(wl[1] & 0xc0u);" "accum += xt8 * float(wl[2] & 3u) + xt9 * float(wl[2] & 0x0cu);" "accum += xt10 * float(wl[2] & 0x30u) + xt11 * float(wl[2] & 0xc0u);" "accum += xt12 * float(wl[3] & 3u) + xt13 * float(wl[3] & 0x0cu);" "accum += xt14 * float(wl[3] & 0x30u) + xt15 * float(wl[3] & 0xc0u);" f"result{r} = fma(s, accum - sum, result{r});" ) if guarded: dots.append(f"if (out_row + {r}u < N) {{ {body} }}") else: dots.append("{ " + body + " }") row_gate = "\n ".join(dots) if guarded: store = "\n ".join( f"float red{r} = simd_sum(result{r}); if (lane == 0u && out_row + {r}u < N) y[b * N + out_row + {r}u] = red{r};" for r in range(rows) ) else: store = "\n ".join( f"float red{r} = simd_sum(result{r}); if (lane == 0u) y[b * N + out_row + {r}u] = red{r};" for r in range(rows) ) return f""" uint lane = thread_index_in_simdgroup; uint sg = simdgroup_index_in_threadgroup; uint tile = threadgroup_position_in_grid.x; uint b = threadgroup_position_in_grid.y; uint N = meta[0]; uint K = meta[1]; uint row_bytes = meta[2]; uint n_groups = meta[3]; uint out_row = (tile * {nsg}u + sg) * {rows}u; {init} for (uint k = 0u; k < K; k += 512u) {{ const device float* xp = x + b * K + k + (lane << 4); float4 a0 = *reinterpret_cast(xp); float4 a1 = *reinterpret_cast(xp + 4); float4 a2 = *reinterpret_cast(xp + 8); float4 a3 = *reinterpret_cast(xp + 12); float sum = ((a0.x + a0.y) + (a0.z + a0.w)) + ((a1.x + a1.y) + (a1.z + a1.w)); sum += ((a2.x + a2.y) + (a2.z + a2.w)) + ((a3.x + a3.y) + (a3.z + a3.w)); float xt0 = a0.x; float xt1 = a0.y * 0.25f; float xt2 = a0.z * 0.0625f; float xt3 = a0.w * 0.015625f; float xt4 = a1.x; float xt5 = a1.y * 0.25f; float xt6 = a1.z * 0.0625f; float xt7 = a1.w * 0.015625f; float xt8 = a2.x; float xt9 = a2.y * 0.25f; float xt10 = a2.z * 0.0625f; float xt11 = a2.w * 0.015625f; float xt12 = a3.x; float xt13 = a3.y * 0.25f; float xt14 = a3.z * 0.0625f; float xt15 = a3.w * 0.015625f; {row_gate} }} {store} """ def run_qmv_mlx(x, codes, scales, k, rows: int = 4, nsg: int = 2, group: int = GROUP): x, codes, scales, b, n = _prepare(x, codes, scales, k, group) if k % 512: raise ValueError(f"qmv mlx needs K multiple of 512, got {k}") rows_per_tg = nsg * rows guarded = (n % rows_per_tg) != 0 tg = nsg * 32 n_tiles = (n + rows_per_tg - 1) // rows_per_tg name = f"mlxtern3_qmvmlx_r{rows}_s{nsg}" + ("_g" if guarded else "_a") + _gname(group) ker = metal(name, _regroup(_qmv_mlx_source(nsg, rows, guarded), group)) (y,) = ker( inputs=[x, codes, scales, meta_u32([n, k, k // 4, k // group, b])], template=[("T", mx.float32)], grid=(n_tiles * tg, b, 1), threadgroup=(tg, 1, 1), output_shapes=[(b, n)], output_dtypes=[mx.float32], ) return y def run_qmv(x, codes, scales, k, rows: int = QMV_ROWS, nsg: int = QMV_NSG, group: int = GROUP, vecdot: bool = False, specialize: bool = True, narrow: bool = False, half_dot: bool = False, chunk128: bool = False): x, codes, scales, b, n = _prepare(x, codes, scales, k, group) if rows not in (2, 4, 8) or nsg not in (1, 2, 4): raise ValueError(f"qmv rows={rows} nsg={nsg}") rows_per_tg = nsg * rows guarded = (n % rows_per_tg) != 0 tg = nsg * 32 n_tiles = (n + rows_per_tg - 1) // rows_per_tg blocks = None if chunk128 else ((k // 512) if (specialize and k % 512 == 0) else None) name = f"mlxtern3_u2q2_r{rows}_s{nsg}" if chunk128: name += "_c128" elif blocks is not None: name += f"_k{k}" if vecdot: name += "_v" if narrow: name += "_n" if half_dot: name += "_h" name += "_g" if guarded else "_a" name += "_fast" + _gname(group) ker = metal(name, _regroup(_qmv_source(nsg, rows, guarded, blocks, vecdot, narrow, half_dot, chunk128), group)) (y,) = ker( inputs=[x, codes, scales, meta_u32([n, k, k // 4, k // group, b])], template=[("T", mx.float32)], grid=(n_tiles * tg, b, 1), threadgroup=(tg, 1, 1), output_shapes=[(b, n)], output_dtypes=[mx.float32], ) return y def run_gemm(x, codes, scales, k, bm: int = GEMM_BM, bn: int = GEMM_BN, bk: int = GEMM_BK, group: int = GROUP, half_w: bool = True): x, codes, scales, b, n = _prepare(x, codes, scales, k, group) if k % bk or n % bn or b % bm: raise ValueError(f"gemm tile {bm}x{bn}x{bk} does not cover N={n} K={k} B={b}") nsg = (bm // 16) * (bn // 16) tg = nsg * 32 name = f"mlxtern3_u2g2_{bm}x{bn}x{bk}" + ("" if half_w else "_fw") + "_fast" + _gname(group) ker = metal(name, _regroup(_gemm_source(bm, bn, bk, half_w), group)) (y,) = ker( inputs=[x, codes, scales, meta_u32([n, k, k // 4, k // group, b])], template=[("T", mx.float32)], grid=((n // bn) * tg, b // bm, 1), threadgroup=(tg, 1, 1), output_shapes=[(b, n)], output_dtypes=[mx.float32], ) return y from collections import OrderedDict COMPILE_MAX_BATCH = 8 COMPILE_CACHE_SIZE = 64 _COMPILED: OrderedDict = OrderedDict() def _compiled(k: int, n: int, b: int, group: int, x_dtype, c_dtype, s_dtype): key = (k, n, b, group, x_dtype, c_dtype, s_dtype) fn = _COMPILED.get(key) if fn is not None: _COMPILED.move_to_end(key) return fn def _run(z, c, s): return _dispatch(z, c, s, k, group) fn = mx.compile(_run) _COMPILED[key] = fn while len(_COMPILED) > COMPILE_CACHE_SIZE: _COMPILED.popitem(last=False) return fn def _dispatch(x, codes, scales, k: int, group: int): b = 1 if x.ndim == 1 else int(x.shape[0]) n = int(codes.shape[0]) tile = _gemm_launch(b, n, k) if tile is not None: bm, bn, bk = tile return run_gemm(x, codes, scales, k, bm, bn, bk, group=group) head = _gemm_head(b, n, k) if head: # Batch >= 32 but not a multiple of 32: the leading multiple of 32 # rows go through the GEMM, the remainder through the GEMV. top = run_gemm(x[:head], codes, scales, k, GEMM_BM, GEMM_BN, GEMM_BK, group=group) rest = _gemv(x[head:], codes, scales, k, group) return mx.concatenate([top, rest], axis=0) return _gemv(x, codes, scales, k, group) def _gemv(x, codes, scales, k: int, group: int): # K multiple of 512 uses the MLX qmv_fast walk (2 rows, 1 simdgroup). # Shorter K, including the 640-wide expert down projection, keeps the # 128-wide body, which is faster there. if k % 512 == 0: return run_qmv_mlx(x, codes, scales, k, rows=2, nsg=1, group=group) return run_qmv(x, codes, scales, k, group=group) def _gemm_head(b: int, n: int, k: int) -> int: """Rows of a ragged batch that the GEMM can take; 0 if none.""" if b < GEMM_MIN_BATCH or n % 32 or k % 32: return 0 return (b // 32) * 32 def _gemm_launch(b: int, n: int, k: int) -> tuple[int, int, int] | None: if b < GEMM_MIN_BATCH or b % 32 or n % 32 or k % 32: return None # 32x32x32 keeps the shared tile small enough to stay resident on the M4 Pro. # Larger tiles were slower on the expert and are re-checked on the wide d3 shapes. return 32, 32, 32 def tern_matmul(x, codes, scales, k: int, group: int = GROUP) -> mx.array: """x [B, K] or [K], codes uint8 [N, K/4], scales [N, K/group] -> [B, N]. ``group`` is 64 (default) or 128 weights per scale. Batches up to ``COMPILE_MAX_BATCH`` reuse a compiled graph from a bounded LRU keyed by shape and dtype; weights are arguments, not captured, so the cache never keeps them alive. Larger batches dispatch the kernel directly. A vector input comes back as shape [1, N]. """ if not isinstance(x, mx.array): x = mx.array(x) if x.ndim == 0 or x.ndim > 2: raise ValueError(f"x must have shape [K] or [B, K], got ndim {x.ndim}") k = int(k) group = int(group) if group not in GROUPS: raise ValueError(f"group {group} not in {GROUPS}") if k <= 0 or k % group: raise ValueError(f"K {k} is not a positive multiple of {group}") if int(x.shape[-1]) != k: raise ValueError(f"x last dim {x.shape[-1]} != K {k}") if codes.ndim != 2 or int(codes.shape[1]) != k // 4: raise ValueError(f"codes shape {tuple(codes.shape)} != [N, K/4] with K={k}") n = int(codes.shape[0]) if tuple(scales.shape) != (n, k // group): raise ValueError(f"scales shape {tuple(scales.shape)} != ({n}, {k // group})") b = 1 if x.ndim == 1 else int(x.shape[0]) if b == 0 or n == 0: # An empty launch would hand Metal a zero-size buffer. return mx.zeros((b, n), dtype=mx.float32) if b > COMPILE_MAX_BATCH: return _dispatch(x, codes, scales, k, int(group)) fn = _compiled(k, n, b, int(group), x.dtype, codes.dtype, scales.dtype) return fn(x, codes, scales)