From 221ea949dc2ce19c61796122b0c751bfaeca20e7 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 11:21:27 +0100 Subject: [PATCH 01/16] profile(): warm-up and best-of-3 timing; MLXR_BENCH_NOPROFILE The first sweep ran cold (first touch of the loaded buffers, GPU clock ramp) and read 2-3x slower than the same matvec inside a decode step. Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Transformer.swift | 16 +++++++++++----- Sources/mlxr/main.swift | 3 ++- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index eac569b..9c3182a 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -1017,12 +1017,18 @@ extension Transformer { /// Developer micro-benchmarks: matvec bandwidth and per-dispatch overhead. public func profile() throws -> String { let q = GPU.shared.queue + // Best of 3 command buffers after one discarded warm-up: the first touch of freshly + // loaded buffers and the GPU clock ramp otherwise make the first sweep look 2-3x slower. func time(_ n: Int, _ body: (MTLComputeCommandEncoder) -> Void) -> Double { - let cb = q.makeCommandBuffer()! - let enc = cb.makeComputeCommandEncoder(dispatchType: .serial)! - for _ in 0.. 0 { best = min(best, (cb.gpuEndTime - cb.gpuStartTime) / Double(n)) } + } + return best } var out = "" let H = config.hiddenSize diff --git a/Sources/mlxr/main.swift b/Sources/mlxr/main.swift index 2a258f6..c5789d8 100644 --- a/Sources/mlxr/main.swift +++ b/Sources/mlxr/main.swift @@ -144,7 +144,8 @@ case "bench": let t = try Transformer(directory: dir, contextLength: 4096) print("weights: \(formatBytes(t.weightBytes))") let quick = ProcessInfo.processInfo.environment["MLXR_BENCH_DECODE"] != nil - if !quick { print(try t.profile()) } + let noProfile = quick || ProcessInfo.processInfo.environment["MLXR_BENCH_NOPROFILE"] != nil + if !noProfile { print(try t.profile()) } _ = try t.forward([1], offset: 0) for n in (quick ? [] : [128, 512, 2048]) { let toks = (0.. Date: Mon, 2 Nov 2026 12:47:38 +0100 Subject: [PATCH 02/16] Prefill: 64x64 quantized GEMM tiles and a faster flash_prefill (bit-identical) qgemm is now one template over the per-simdgroup accumulator grid: the 32x32 tile (2x2 accumulators) replaces the old qgemm, and a 64x64 tile (4x4 accumulators, 2 MMAs per fragment load instead of 1) runs long batches. Both keep each output's chain (fp32 dequant, k ascending in 8-steps), so results are bit-identical to the old kernel and to each other. Full unrolling is what makes the 4x4 tile work (4x slower without it); padded rows and register prefetch measured slower on M2 Pro and are not used. Linear.useGemm64 picks the 64x64 tile when it pads at most 1.2x more rows and yields >= 64 threadgroups. flash_prefill stages K/V as half (exact: the cache is half) in two buffers with 16-byte loads of the next block during the current one's MMAs (one barrier per block), rescales O element-wise instead of by a diagonal MMA, and schedules the latest query blocks first. Same key blocking and softmax order, bit-identical output; ~1.9x faster at D=128, ~1.4-1.6x at D=256. Tests: both tiles against the CPU reference (odd O, K % 64 != 0, T up to 517), tiles agree bitwise across batch sizes, and prefill vs incremental decode for a 600-token prompt (two batches, the second at offset 512). Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Kernels.swift | 304 ++++++++++++++++--------- Sources/MLXEngine/Transformer.swift | 28 ++- Tests/MLXEngineTests/KernelTests.swift | 38 ++++ Tests/MLXEngineTests/ModelTests.swift | 24 ++ 4 files changed, 282 insertions(+), 112 deletions(-) diff --git a/Sources/MLXEngine/Kernels.swift b/Sources/MLXEngine/Kernels.swift index c767161..036681d 100644 --- a/Sources/MLXEngine/Kernels.swift +++ b/Sources/MLXEngine/Kernels.swift @@ -391,10 +391,10 @@ kernel void qmvb4(device const uchar* wA [[buffer(0)]], } // --------------------------------------------------------------------------- -// Prefill path: tiled GEMM on simdgroup matrices. Y[T, O] = X[T, K] · W[O, K]ᵀ +// Prefill path for dense weights (dgemm): tiled GEMM on simdgroup matrices. Y[T, O] = X[T, K] · W[O, K]ᵀ // Threadgroup tile: 32 tokens × 32 output rows, K-step 32, 4 simdgroups (2×2), // each owning a 16×16 block of 8×8 accumulators. Weights are dequantized into -// threadgroup memory once per tile and shared by all 32 tokens. +// threadgroup memory once per tile and shared by all 32 tokens. (Quantized weights: qgemm below.) // --------------------------------------------------------------------------- #define GBM 32 #define GBN 32 @@ -438,7 +438,34 @@ inline void gemm_store(threadgroup float* S, thread simdgroup_matrix +inline void gemm_dequant8(device const uchar* wrow, device const T* srow, device const T* brow, + int k, int gs, bool ok, thread float4& qa, thread float4& qb) { + qa = float4(0); qb = float4(0); + if (ok) { + unpack8(wrow + ulong(k) * BITS / 8, qa, qb); + const int g = k / gs; + const float s = float(srow[g]), b = float(brow[g]); + qa = qa * s + b; qb = qb * s + b; + } +} + +template kernel void qgemm(device const uchar* w [[buffer(0)]], device const T* scales [[buffer(1)]], device const T* biases [[buffer(2)]], @@ -448,33 +475,62 @@ kernel void qgemm(device const uchar* w [[buffer(0)]], uint2 tgid [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], uint sgid [[simdgroup_index_in_threadgroup]]) { - threadgroup float Xs[GBM * GBK]; - threadgroup float Ws[GBN * GBK]; - const int oBase = int(tgid.x) * GBN, tBase = int(tgid.y) * GBM; - const uint sm = sgid / 2, sn = sgid % 2; - simdgroup_matrix C[2][2]; - for (int i = 0; i < 2; i++) for (int j = 0; j < 2; j++) C[i][j] = simdgroup_matrix(0); + constexpr int BM = 16 * NI, BN = 16 * NJ, BK = GBK; + constexpr int XV = BM / 16; // float4s of x per thread per K-step + constexpr int WV = BN / 32; // 8-weight chunks per thread per K-step + constexpr int SH = (BM + BN) * BK > BM * BN ? (BM + BN) * BK : BM * BN; + threadgroup float sh[SH]; // Xs then Ws; reused to stage edge tiles + threadgroup float* Xs = sh; + threadgroup float* Ws = sh + BM * BK; + const int oBase = int(tgid.x) * BN, tBase = int(tgid.y) * BM; + const int sm = int(sgid) / 2, sn = int(sgid) % 2; + simdgroup_matrix C[NI][NJ]; + UNROLL for (int i = 0; i < NI; i++) UNROLL for (int j = 0; j < NJ; j++) C[i][j] = simdgroup_matrix(0); const ulong rowBytes = ulong(p.K) * BITS / 8; const int gpr = p.K / p.gs; - // each thread dequantizes one 8-value chunk: row = tid / 4, chunk = tid % 4 + // Weights: thread → rows tid / 4 + 32 v, 8-value chunk tid % 4. const int wr = int(tid) / 4, wc = int(tid) % 4; - const int orow = oBase + wr; - for (int k0 = 0; k0 < p.K; k0 += GBK) { - gemm_load_x(Xs, x, p, tBase, k0, tid); - float4 qa = float4(0), qb = float4(0); - if (orow < p.O) { - unpack8(w + ulong(orow) * rowBytes + ulong(k0 + wc * 8) * BITS / 8, qa, qb); - const int g = (k0 + wc * 8) / p.gs; - const float s = float(scales[orow * gpr + g]), b = float(biases[orow * gpr + g]); - qa = qa * s + b; qb = qb * s + b; + ulong wo[WV]; + UNROLL for (int v = 0; v < WV; v++) wo[v] = ulong(min(oBase + wr + 32 * v, p.O - 1)); + // Activations: thread → token rows tid / 8 + 16 v, columns (tid % 8) * 4 ..< + 4. + const int xr = int(tid) / 8, xc = (int(tid) % 8) * 4; + int xo[XV]; + UNROLL for (int v = 0; v < XV; v++) xo[v] = min(tBase + xr + 16 * v, p.T - 1) * p.xStride + xc; + for (int k0 = 0; k0 < p.K; k0 += BK) { + UNROLL for (int v = 0; v < WV; v++) { + float4 qa, qb; + gemm_dequant8(w + wo[v] * rowBytes, scales + wo[v] * gpr, biases + wo[v] * gpr, k0 + wc * 8, p.gs, true, qa, qb); + *(threadgroup float4*)(Ws + (wr + 32 * v) * BK + wc * 8) = qa; + *(threadgroup float4*)(Ws + (wr + 32 * v) * BK + wc * 8 + 4) = qb; } - *(threadgroup float4*)(Ws + wr * GBK + wc * 8) = qa; - *(threadgroup float4*)(Ws + wr * GBK + wc * 8 + 4) = qb; + UNROLL for (int v = 0; v < XV; v++) + *(threadgroup float4*)(Xs + (xr + 16 * v) * BK + xc) = *(device const float4*)(x + xo[v] + k0); threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_tile_mma(Xs, Ws, C, sm, sn); + UNROLL for (int kk = 0; kk < BK; kk += 8) { + simdgroup_matrix A[NI], B[NJ]; + UNROLL for (int i = 0; i < NI; i++) simdgroup_load(A[i], Xs + (sm * 8 * NI + i * 8) * BK + kk, BK); + UNROLL for (int j = 0; j < NJ; j++) simdgroup_load(B[j], Ws + (sn * 8 * NJ + j * 8) * BK + kk, BK, ulong2(0, 0), true); + UNROLL for (int i = 0; i < NI; i++) + UNROLL for (int j = 0; j < NJ; j++) + simdgroup_multiply_accumulate(C[i][j], A[i], B[j], C[i][j]); + } threadgroup_barrier(mem_flags::mem_threadgroup); } - gemm_store(Xs, C, y, p, tBase, oBase, sm, sn, tid); + if (tBase + BM <= p.T && oBase + BN <= p.O) { + device float* yb = y + (tBase + sm * 8 * NI) * p.yStride + oBase + sn * 8 * NJ; + UNROLL for (int i = 0; i < NI; i++) + UNROLL for (int j = 0; j < NJ; j++) + simdgroup_store(C[i][j], yb + (i * 8) * p.yStride + j * 8, p.yStride); + return; + } + UNROLL for (int i = 0; i < NI; i++) + UNROLL for (int j = 0; j < NJ; j++) + simdgroup_store(C[i][j], sh + (sm * 8 * NI + i * 8) * BN + sn * 8 * NJ + j * 8, BN); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = int(tid); idx < BM * BN; idx += 128) { + const int r = idx / BN, c = idx % BN; + if (tBase + r < p.T && oBase + c < p.O) y[(tBase + r) * p.yStride + oBase + c] = sh[idx]; + } } // --------------------------------------------------------------------------- @@ -489,18 +545,6 @@ kernel void qgemm(device const uchar* w [[buffer(0)]], // --------------------------------------------------------------------------- #define SBK 64 -template -inline void gemm_dequant8(device const uchar* wrow, device const T* srow, device const T* brow, - int k, int gs, bool ok, thread float4& qa, thread float4& qb) { - qa = float4(0); qb = float4(0); - if (ok) { - unpack8(wrow + ulong(k) * BITS / 8, qa, qb); - const int g = k / gs; - const float s = float(srow[g]), b = float(brow[g]); - qa = qa * s + b; qb = qb * s + b; - } -} - template kernel void qgemms(device const uchar* w [[buffer(0)]], device const T* scales [[buffer(1)]], @@ -1080,9 +1124,50 @@ kernel void attention_prefill(device const float* qkv [[buffer(0)]], // Flash attention for prefill on simdgroup matrices. // Threadgroup = (head, 32 queries); 4 simdgroups × 8 query rows; key blocks of BK -// (16, or 8 for head_dim 256 so the K/V block fits threadgroup memory). -// S = Q·Kᵀ and O += P·V use 8×8 MMA; the online-softmax rescale of O is a -// multiply by a diagonal matrix. DD = head_dim / 8. +// (16, or 8 for head_dim 256). K/V blocks are staged as half (the cache type) in two +// threadgroup buffers: the next block is read into registers with 16-byte loads while +// the current one is multiplied, so each block costs one barrier. S = Q·Kᵀ and +// O += P·V use 8×8 MMA on fragments converted exactly from half; the online-softmax +// rescale of O is applied to each lane's fragment elements. Query blocks are +// scheduled latest (heaviest) first. DD = head_dim / 8. +template +inline void flash_load(device const half* kb, device const half* vb, int j0, int L, int maxSeq, uint tid, + thread uint4 (&kq)[CPT], thread uint4 (&vq)[CPT]) { + constexpr int RC = D / 8, CH = BK * RC; // 8-half chunks per row / per block + for (int v = 0; v < CPT; v++) { + const int idx = int(tid) + 128 * v; + if (CH % 128 != 0 && idx >= CH) break; + const int r = idx / RC, j = j0 + r; + if (j < L) { + const ulong src = ulong(kv_slot(j, maxSeq)) * D + (idx - r * RC) * 8; + kq[v] = *(device const uint4*)(kb + src); + vq[v] = *(device const uint4*)(vb + src); + } else { + kq[v] = uint4(0); vq[v] = uint4(0); + } + } +} + +template +inline void flash_stage(threadgroup half* Ks, uint tid, thread uint4 (&kq)[CPT], thread uint4 (&vq)[CPT]) { + constexpr int CH = BK * D / 8; + for (int v = 0; v < CPT; v++) { + const int idx = int(tid) + 128 * v; + if (CH % 128 != 0 && idx >= CH) break; + *(threadgroup uint4*)(Ks + idx * 8) = kq[v]; + *(threadgroup uint4*)(Ks + BK * D + idx * 8) = vq[v]; + } +} + +inline simdgroup_matrix flash_frag(threadgroup const half* src, int ld, bool transpose) { + simdgroup_matrix hm; + simdgroup_load(hm, src, ld, ulong2(0, 0), transpose); + simdgroup_matrix fm; + fm.thread_elements()[0] = float(hm.thread_elements()[0]); + fm.thread_elements()[1] = float(hm.thread_elements()[1]); + return fm; +} + template kernel void flash_prefill(device const float* qkv [[buffer(0)]], device const half* kc [[buffer(1)]], @@ -1096,13 +1181,14 @@ kernel void flash_prefill(device const float* qkv [[buffer(0)]], constexpr int D = DD * 8; constexpr int NB = BK / 8; // 8×8 score tiles per key block constexpr int CPL = BK / 4; // score columns per lane (4 lanes per row) - threadgroup float KV[2 * BK * D]; // K block then V block; reused for the output + constexpr int CPT = (BK * D / 8 + 127) / 128; // 16-byte K (and V) chunks per thread per block + threadgroup half KV[2][2 * BK * D]; // two stages of [K block | V block] threadgroup float Ssh[4][8 * BK]; - threadgroup float Dg[4][64]; - threadgroup float* Ks = KV; - threadgroup float* Vs = KV + BK * D; + // Per-row factors handed from the softmax lanes to the fragment lanes (simd_shuffle by + // a computed lane gave wrong values here). + threadgroup float Rf[4][8]; - const int h = int(tg.x), t0 = int(tg.y) * 32; + const int h = int(tg.x), t0 = ((p.T + 31) / 32 - 1 - int(tg.y)) * 32; const int kvh = h / (p.nH / p.nKV); const int qRow0 = t0 + int(sg) * 8; // first query (batch index) of this simdgroup const int lastQ = min(t0 + 32, p.T) - 1; @@ -1113,87 +1199,91 @@ kernel void flash_prefill(device const float* qkv [[buffer(0)]], simdgroup_load(Q[d], qkv + qRow0 * p.qStride + h * D + d * 8, p.qStride); O[d] = simdgroup_matrix(0); } - for (int i = int(lane); i < 64; i += 32) Dg[sg][i] = 0.0f; const int r = int(lane) / 4, c0 = (int(lane) % 4) * CPL; // softmax ownership: row r, CPL columns const int qAbs = p.offset + qRow0 + r; // absolute position of row r const bool rowValid = qRow0 + r < p.T; float m = -1e30f, l = 0.0f; + // Fragment ownership: this lane holds row fr (columns fc, fc + 1) of every 8×8 fragment. + const int fr = (int(lane) / 4 & 4) + (int(lane) / 2) % 4; device const half* kb = kc + ulong(kvh) * ulong(p.maxSeq) * ulong(D); device const half* vb = vc + ulong(kvh) * ulong(p.maxSeq) * ulong(D); const int qFirst = attn_first_key(p, p.offset + t0); // first key of the block's first query + uint4 kq[CPT], vq[CPT]; + int stage = 0; + flash_load(kb, vb, qFirst / BK * BK, L, p.maxSeq, tid, kq, vq); + flash_stage(KV[0], tid, kq, vq); + threadgroup_barrier(mem_flags::mem_threadgroup); for (int j0 = qFirst / BK * BK; j0 < L; j0 += BK) { - threadgroup_barrier(mem_flags::mem_threadgroup); - // D is a power of two: row / column from i are shifts; the ring slot costs a modulo only - // for positions past the cache length. - for (int i = int(tid); i < BK * D; i += 128) { - const int r = i / D, j = j0 + r; - const bool ok = j < L; - const ulong src = ulong(kv_slot(j, p.maxSeq)) * D + (i - r * D); - Ks[i] = ok ? float(kb[src]) : 0.0f; - Vs[i] = ok ? float(vb[src]) : 0.0f; - } - threadgroup_barrier(mem_flags::mem_threadgroup); + const int jn = j0 + BK; + if (jn < L) flash_load(kb, vb, jn, L, p.maxSeq, tid, kq, vq); + threadgroup const half* Ks = KV[stage]; + threadgroup const half* Vs = KV[stage] + BK * D; // Skip key blocks that are entirely in the future for all 8 rows of this simdgroup. - if (j0 > p.offset + qRow0 + 7) continue; - - simdgroup_matrix S[NB], Kt; - for (int b = 0; b < NB; b++) S[b] = simdgroup_matrix(0); - for (int d = 0; d < DD; d++) { - for (int b = 0; b < NB; b++) { - simdgroup_load(Kt, Ks + b * 8 * D + d * 8, D, ulong2(0, 0), true); - simdgroup_multiply_accumulate(S[b], Q[d], Kt, S[b]); + if (j0 <= p.offset + qRow0 + 7) { + simdgroup_matrix S[NB]; + for (int b = 0; b < NB; b++) S[b] = simdgroup_matrix(0); + for (int d = 0; d < DD; d++) { + for (int b = 0; b < NB; b++) { + simdgroup_matrix Kt = flash_frag(Ks + b * 8 * D + d * 8, D, true); + simdgroup_multiply_accumulate(S[b], Q[d], Kt, S[b]); + } } - } - for (int b = 0; b < NB; b++) simdgroup_store(S[b], &Ssh[sg][b * 8], BK); - simdgroup_barrier(mem_flags::mem_threadgroup); - - float s[CPL]; - float bmax = -1e30f; - const int jMin = attn_first_key(p, qAbs); - for (int c = 0; c < CPL; c++) { - const int j = j0 + c0 + c; - s[c] = (rowValid && j <= qAbs && j >= jMin && j < L) ? attn_cap(Ssh[sg][r * BK + c0 + c] * p.scale, p.softcap) : -1e30f; - bmax = max(bmax, s[c]); - } - bmax = max(bmax, simd_shuffle_xor(bmax, 1)); - bmax = max(bmax, simd_shuffle_xor(bmax, 2)); - const float mn = max(m, bmax); - float bsum = 0.0f; - for (int c = 0; c < CPL; c++) { - const float pe = s[c] <= -1e29f ? 0.0f : exp(s[c] - mn); - Ssh[sg][r * BK + c0 + c] = pe; - bsum += pe; - } - bsum += simd_shuffle_xor(bsum, 1); - bsum += simd_shuffle_xor(bsum, 2); - const float corr = exp(m - mn); - l = l * corr + bsum; - m = mn; - if (c0 == 0) Dg[sg][r * 8 + r] = corr; - simdgroup_barrier(mem_flags::mem_threadgroup); - - simdgroup_matrix Dm, P[NB], Vb; - simdgroup_load(Dm, &Dg[sg][0], 8); - for (int b = 0; b < NB; b++) simdgroup_load(P[b], &Ssh[sg][b * 8], BK); - for (int d = 0; d < DD; d++) { - simdgroup_multiply(O[d], Dm, O[d]); - for (int b = 0; b < NB; b++) { - simdgroup_load(Vb, Vs + b * 8 * D + d * 8, D); - simdgroup_multiply_accumulate(O[d], P[b], Vb, O[d]); + for (int b = 0; b < NB; b++) simdgroup_store(S[b], &Ssh[sg][b * 8], BK); + simdgroup_barrier(mem_flags::mem_threadgroup); + + float s[CPL]; + float bmax = -1e30f; + const int jMin = attn_first_key(p, qAbs); + for (int c = 0; c < CPL; c++) { + const int j = j0 + c0 + c; + s[c] = (rowValid && j <= qAbs && j >= jMin && j < L) ? attn_cap(Ssh[sg][r * BK + c0 + c] * p.scale, p.softcap) : -1e30f; + bmax = max(bmax, s[c]); + } + bmax = max(bmax, simd_shuffle_xor(bmax, 1)); + bmax = max(bmax, simd_shuffle_xor(bmax, 2)); + const float mn = max(m, bmax); + float bsum = 0.0f; + for (int c = 0; c < CPL; c++) { + const float pe = s[c] <= -1e29f ? 0.0f : exp(s[c] - mn); + Ssh[sg][r * BK + c0 + c] = pe; + bsum += pe; + } + bsum += simd_shuffle_xor(bsum, 1); + bsum += simd_shuffle_xor(bsum, 2); + const float corr = exp(m - mn); + l = l * corr + bsum; + m = mn; + if (c0 == 0) Rf[sg][r] = corr; + simdgroup_barrier(mem_flags::mem_threadgroup); + const float fcorr = Rf[sg][fr]; // the factor of fragment row fr + + simdgroup_matrix P[NB]; + for (int b = 0; b < NB; b++) simdgroup_load(P[b], &Ssh[sg][b * 8], BK); + for (int d = 0; d < DD; d++) { + O[d].thread_elements()[0] *= fcorr; + O[d].thread_elements()[1] *= fcorr; + for (int b = 0; b < NB; b++) { + simdgroup_matrix Vb = flash_frag(Vs + b * 8 * D + d * 8, D, false); + simdgroup_multiply_accumulate(O[d], P[b], Vb, O[d]); + } } } + if (jn < L) flash_stage(KV[stage ^ 1], tid, kq, vq); + threadgroup_barrier(mem_flags::mem_threadgroup); + stage ^= 1; } - // Write out: stage each simdgroup's 8×D block in KV, normalise rows by l. KV holds - // 2·BK·D floats: all four blocks at once for BK = 16, two simdgroups at a time for BK = 8. - constexpr int STAGE = (2 * BK * D) / (8 * D); // simdgroups staged per round (4 or 2) + // Write out: stage each simdgroup's 8×D block in the K/V stages (2·BK·D floats: all four blocks + // at once for BK = 16, two at a time for BK = 8) and normalise rows by l. (Scaling the fragment + // elements in registers instead was not bit-identical to this.) const float inv = l > 0 ? 1.0f / l : 0.0f; + constexpr int STAGE = BK / 4; // simdgroups staged per round (4 or 2) for (int round = 0; round < 4 / STAGE; round++) { - threadgroup_barrier(mem_flags::mem_threadgroup); + if (round > 0) threadgroup_barrier(mem_flags::mem_threadgroup); if (int(sg) / STAGE != round) continue; - threadgroup float* Osh = KV + (int(sg) % STAGE) * 8 * D; + threadgroup float* Osh = (threadgroup float*)&KV[0][0] + (int(sg) % STAGE) * 8 * D; for (int d = 0; d < DD; d++) simdgroup_store(O[d], Osh + d * 8, D); simdgroup_barrier(mem_flags::mem_threadgroup); if (rowValid) { @@ -1509,7 +1599,7 @@ kernel void sample_token(device float* logits [[buffer(0)]], #define INST_QNEW(B, T, TN) \ template [[host_name("qmv_" #B "_" #TN)]] kernel void qmv(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, device const float*, uint2, uint, uint); \ template [[host_name("qmv_gate_" #B "_" #TN)]] kernel void qmv(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, device const float*, uint2, uint, uint); \ - template [[host_name("qgemm_" #B "_" #TN)]] kernel void qgemm(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); \ + template [[host_name("qgemm_" #B "_" #TN)]] kernel void qgemm(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); \ template [[host_name("qgemms_" #B "_" #TN "_8")]] kernel void qgemms(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); \ template [[host_name("qgemms_" #B "_" #TN "_16")]] kernel void qgemms(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); // 4-bit remainder variants (TAIL 4 / 8 / 16), plain and gated. @@ -1521,6 +1611,10 @@ kernel void sample_token(device float* logits [[buffer(0)]], template [[host_name("qmv_gate_t8_4_" #TN)]] kernel void qmv<4, T, 1, 8>(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, device const float*, uint2, uint, uint); \ template [[host_name("qmv_gate_t16_4_" #TN)]] kernel void qmv<4, T, 1, 16>(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, device const float*, uint2, uint, uint); INST_QMV_TAIL(half, f16) INST_QMV_TAIL(bfloat, bf16) INST_QMV_TAIL(float, f32) +#define INST_QGEMM64(B, T, TN) \ + template [[host_name("qgemm64_" #B "_" #TN)]] kernel void qgemm(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); +#define INST_QGEMM64_ALL(B) INST_QGEMM64(B, half, f16) INST_QGEMM64(B, bfloat, bf16) INST_QGEMM64(B, float, f32) +INST_QGEMM64_ALL(2) INST_QGEMM64_ALL(3) INST_QGEMM64_ALL(4) INST_QGEMM64_ALL(5) INST_QGEMM64_ALL(6) INST_QGEMM64_ALL(8) #define INST_QMM_ALL(B) \ INST_QNEW(B, half, f16) INST_QNEW(B, bfloat, bf16) INST_QNEW(B, float, f32) \ INST_QMM(B, half, f16, 1) INST_QMM(B, half, f16, 8) \ diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index 9c3182a..7664ffc 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -24,6 +24,7 @@ final class Linear { private var gemmPSO: MTLComputePipelineState? // prefill tiles (simdgroup matrices) private var gemm8PSO: MTLComputePipelineState? // 5...8 tokens: qgemm on an 8-token tile private var gemm16PSO: MTLComputePipelineState? // 9...16 tokens: qgemm on a 16-token tile + private var gemm64PSO: MTLComputePipelineState? // long prompts: 64 × 64 tiles private(set) var qmvPSO: MTLComputePipelineState? // fused decode kernel (bias folded in) private(set) var qmvGatePSO: MTLComputePipelineState? // decode with y = act(y) · e private(set) var qmvbPSO: MTLComputePipelineState? // 4-bit, 2...4 tokens (speculative verification) @@ -75,7 +76,10 @@ final class Linear { qmvGatePSO = try gpu.pipeline("qmv_gate_\(tail)\(bits)_\(sfx)") if bits == 4 { qmvbPSO = try gpu.pipeline("qmvb4_\(sfx)_4") } } - if inDim % 32 == 0 { gemmPSO = try gpu.pipeline("qgemm_\(bits)_\(sfx)") } + if inDim % 32 == 0 { + gemmPSO = try gpu.pipeline("qgemm_\(bits)_\(sfx)") + gemm64PSO = try gpu.pipeline("qgemm64_\(bits)_\(sfx)") + } if inDim % 64 == 0 { // the short-prompt tiles step K by 64 gemm8PSO = try gpu.pipeline("qgemms_\(bits)_\(sfx)_8") gemm16PSO = try gpu.pipeline("qgemms_\(bits)_\(sfx)_16") @@ -154,6 +158,14 @@ final class Linear { threadsPerThreadgroup: MTLSize(width: qmvThreads, height: 1, depth: 1)) } + /// The 64 × 64 GEMM tile runs ~1.2× faster per row than the 32 × 32 one, but pads the batch to a + /// multiple of 64 tokens and makes a quarter as many threadgroups: use it when it pads at most + /// 1.2× more rows and still yields enough threadgroups to fill the GPU (M2 Pro measurements). + static func useGemm64(T: Int, outDim: Int) -> Bool { + let rows64 = (T + 63) / 64 * 64, rows32 = (T + 31) / 32 * 32 + return rows64 * 5 <= rows32 * 6 && (T + 63) / 64 * ((outDim + 63) / 64) >= 64 + } + /// y[t, :outDim] (row stride yStride, starting at yOffset bytes) = x[t] · Wᵀ (+ bias) func encode(_ enc: MTLComputeCommandEncoder, x: MTLBuffer, xOffset: Int = 0, xStride: Int, y: MTLBuffer, yOffset: Int = 0, yStride: Int, T: Int) { @@ -162,12 +174,14 @@ final class Linear { return } // 5+ tokens: tiled GEMM beats the small-batch kernel. A full tile is 32 tokens; short prompts - // use an 8- or 16-token tile instead of paying for 32 (same arithmetic per output element). - var gemm: (pso: MTLComputePipelineState, rows: Int)? + // use an 8- or 16-token tile instead of paying for 32, long ones a 64 × 64 tile (same + // arithmetic per output element in all of them). + var gemm: (pso: MTLComputePipelineState, rows: Int, cols: Int)? if T > Transformer.maxVerifyTokens { - if T <= 8, let p = gemm8PSO { gemm = (p, 8) } - else if T <= 16, let p = gemm16PSO { gemm = (p, 16) } - else if let p = gemmPSO { gemm = (p, 32) } + if T <= 8, let p = gemm8PSO { gemm = (p, 8, 32) } + else if T <= 16, let p = gemm16PSO { gemm = (p, 16, 32) } + else if Linear.useGemm64(T: T, outDim: outDim), let p = gemm64PSO { gemm = (p, 64, 64) } + else if let p = gemmPSO { gemm = (p, 32, 32) } } let tt = T == 1 ? 1 : prefillTT enc.setComputePipelineState(gemm?.pso ?? (T == 1 ? pso1 : psoN)) @@ -181,7 +195,7 @@ final class Linear { enc.params(MMParams(K: Int32(inDim), O: Int32(outDim), T: Int32(T), gs: Int32(groupSize), xStride: Int32(xStride), yStride: Int32(yStride)), 5) if let gemm { - enc.dispatchThreadgroups(MTLSize(width: (outDim + 31) / 32, height: (T + gemm.rows - 1) / gemm.rows, depth: 1), + enc.dispatchThreadgroups(MTLSize(width: (outDim + gemm.cols - 1) / gemm.cols, height: (T + gemm.rows - 1) / gemm.rows, depth: 1), threadsPerThreadgroup: MTLSize(width: 128, height: 1, depth: 1)) } else { enc.dispatchThreadgroups(MTLSize(width: (outDim + rowsPerThreadgroup - 1) / rowsPerThreadgroup, diff --git a/Tests/MLXEngineTests/KernelTests.swift b/Tests/MLXEngineTests/KernelTests.swift index fc62152..46c2caa 100644 --- a/Tests/MLXEngineTests/KernelTests.swift +++ b/Tests/MLXEngineTests/KernelTests.swift @@ -145,6 +145,44 @@ struct MatmulKernelTests { } } + /// Prefill GEMM tiles on long batches: the 32 × 32 tile (few output rows, or T just past a + /// multiple of 64) and the 64 × 64 tile (many rows), with odd O and K % 64 != 0. + @Test(arguments: [3, 4, 8]) + func prefillTilesMatchReference(bits: Int) throws { + // (T, O, K, uses the 64 × 64 tile) + let shapes = [(64, 41, 1056, false), (128, 41, 544, false), (517, 41, 544, false), (96, 2050, 544, false), + (128, 2050, 544, true), (517, 520, 544, true)] + for (T, O, K, big) in shapes { + #expect(Linear.useGemm64(T: T, outDim: O) == big, "tile choice T \(T) O \(O)") + let fx = try QuantFixture(O: O, K: K, bits: bits, gs: 32, seed: UInt64(T + O)) + var rng = RNG(s: UInt64(T * 7 + bits)) + let x = (0..<(T * K)).map { _ in rng.float() } + let yb = try GPU.shared.buffer(T * O * 4) + runGPU { fx.linear.encode($0, x: makeBuffer(x), xStride: K, y: yb, yStride: O, T: T) } + assertClose(readFloats(yb, T * O), fx.reference(x, T: T, K: K), "bits \(bits) T \(T) O \(O) K \(K)") + } + } + + /// Both prefill tiles (and the short-prompt tiles) compute every output with the same chain of + /// operations, so a token's output does not depend on the batch it was computed in. + @Test func prefillTilesAgreeBitwise() throws { + let O = 2050, K = 1024 + let fx = try QuantFixture(O: O, K: K, bits: 4, gs: 64, seed: 3) + var rng = RNG(s: 11) + let x = (0..<(256 * K)).map { _ in rng.float() } + let xb = makeBuffer(x) + func run(_ T: Int) throws -> [UInt32] { + let yb = try GPU.shared.buffer(T * O * 4) + runGPU { fx.linear.encode($0, x: xb, xStride: K, y: yb, yStride: O, T: T) } + return readFloats(yb, T * O).map(\.bitPattern) + } + #expect(Linear.useGemm64(T: 256, outDim: O) && !Linear.useGemm64(T: 96, outDim: O)) + let full = try run(256) + for T in [9, 16, 33, 96] { + #expect(try run(T) == Array(full[..<(T * O)]), "T \(T) vs T 256") + } + } + /// K values that are not a multiple of the 4-bit decode kernel's 1024-wide pass: the 128/256/512 /// remainders are spread over all lanes, others fall back to a partial pass (Gemma, Qwen2.5). @Test(arguments: [1152, 1536, 2304, 2560, 896, 1664, 6912]) diff --git a/Tests/MLXEngineTests/ModelTests.swift b/Tests/MLXEngineTests/ModelTests.swift index ebd1a24..c773fcf 100644 --- a/Tests/MLXEngineTests/ModelTests.swift +++ b/Tests/MLXEngineTests/ModelTests.swift @@ -55,6 +55,30 @@ struct ModelTests { #expect(maxDiff < 0.1, "max logit difference \(maxDiff)") } + /// A prompt longer than one prefill batch (512 tokens): the 64 × 64 GEMM tiles and flash + /// attention over many key blocks, then a second batch at offset 512, against token-by-token decode. + @Test(.enabled(if: installedModel(qwen3) != nil), arguments: [qwen3, llama, phi, mistral, "mlx-community/Qwen3-8B-4bit"]) + func longPrefillMatchesIncrementalDecode(repo: String) throws { + guard let dir = installedModel(repo) else { return } + let model = try Transformer(directory: dir, contextLength: 1024) + let tok = try Tokenizer(directory: dir) + let text = (0..<40).map { "Item \($0): the \(["red", "green", "blue", "amber"][$0 % 4]) fox jumps over \($0 * 7) lazy dogs." }.joined(separator: " ") + let ids = Array(tok.encode(text).map { Int32($0) }.prefix(600)) + #expect(ids.count > ModelConfig.maxBatch + 32) + let V = model.vocabSize + let pre = Array(UnsafeBufferPointer(start: try model.forward(ids, offset: 0), count: V)) + var inc: [Float] = [] + for (i, t) in ids.enumerated() { + let l = try model.forward([t], offset: i) + if i == ids.count - 1 { inc = Array(UnsafeBufferPointer(start: l, count: V)) } + } + let argPre = pre.indices.max { pre[$0] < pre[$1] }! + let argInc = inc.indices.max { inc[$0] < inc[$1] }! + #expect(argPre == argInc) + let maxDiff = zip(pre, inc).map { abs($0 - $1) }.max()! + #expect(maxDiff < 0.1, "max logit difference \(maxDiff)") + } + /// The long-context split decode attention (attention_split + attention_combine) must give /// bit-identical logits to the single-kernel decode attention, at long and short contexts. @Test(.enabled(if: installedModel(qwen3) != nil), arguments: [qwen3, llama, phi, mistral]) From e4390d3322ba4ca77f67bc0306f1c8ea9c330e55 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 11:44:27 +0100 Subject: [PATCH 03/16] WIP: qmvx4 small-batch 4-bit matmul on simdgroup matrices (T = 2..8) Replaces qmvb4; maxVerifyTokens 4 -> 8. Correct (kernel tests pass), not yet faster. Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Kernels.swift | 151 ++++++++++++++++--------- Sources/MLXEngine/Transformer.swift | 34 ++++-- Sources/mlxr/main.swift | 8 +- Tests/MLXEngineTests/KernelTests.swift | 42 +++++++ 4 files changed, 170 insertions(+), 65 deletions(-) diff --git a/Sources/MLXEngine/Kernels.swift b/Sources/MLXEngine/Kernels.swift index 036681d..039cd7f 100644 --- a/Sources/MLXEngine/Kernels.swift +++ b/Sources/MLXEngine/Kernels.swift @@ -317,11 +317,25 @@ kernel void qmv(device const uchar* wA [[buffer(0)]], } } -// Small-batch 4-bit matvec for speculative verification (2..TT tokens): each -// weight chunk is loaded once and applied to every token, so verifying a few -// draft tokens costs little more than decoding one. Same fused [A|B|C] layout. -template -kernel void qmvb4(device const uchar* wA [[buffer(0)]], +// Small-batch 4-bit matmul (T = 2..8 tokens: speculative verification, short prompts) on +// simdgroup matrices. Each weight is loaded and dequantized once, independent of T; the +// per-token work is one 8×8×8 multiply-accumulate per 64 weights, the tokens padded to 8. +// A simdgroup owns 8 output rows: C[row][token] += W[row][k] · X[token][k], fp32 throughout. +// No threadgroup strip for the weights: every lane dequantizes exactly the two elements of +// the W fragment it holds (thread_elements: lane → row fm, columns fn, fn + 1). The sum +// over k is order-free, so fragment column c of step kk in a 128-wide block stands for +// k = 32 (c / 2) + 2 kk + c % 2: a lane's 16 steps then read one contiguous 16-byte load +// (row fm, bytes 16 (fn / 2) ..+16 of the block). x is staged once per threadgroup in that +// order, 8 tokens × QXK values, so a lane's X fragment elements are one float2 per step. +// Same fused [A|B|C] layout and arguments as `qmv`; K % 32 == 0 and gs % 32 == 0 +// (a 32-value lane run stays in one group), K need not be a multiple of 128. +#define QXNSG 4 // simdgroups per threadgroup +#define QXRF 2 // 8-row fragments per simdgroup +#define QXK 256 // x values staged per chunk +#define QXACC 2 // independent accumulator chains per fragment + +template +kernel void qmvx4(device const uchar* wA [[buffer(0)]], device const T* sA [[buffer(1)]], device const T* bA [[buffer(2)]], device const float* x [[buffer(3)]], @@ -337,55 +351,91 @@ kernel void qmvb4(device const uchar* wA [[buffer(0)]], device const T* biasB [[buffer(13)]], device const T* biasC [[buffer(14)]], uint2 tgid [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], uint sgid [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { - const int gRow = (int(tgid.x) * QVNSG + int(sgid)) * R; - if (gRow >= p.O) return; - device const uchar* w; device const T* sc; device const T* bi; device const T* rb; - int rowBase; - if (gRow < p.o1) { w = wA; sc = sA; bi = bA; rb = biasA; rowBase = gRow; } - else if (gRow < p.o1 + p.o2) { w = wB; sc = sB; bi = bB; rb = biasB; rowBase = gRow - p.o1; } - else { w = wC; sc = sC; bi = bC; rb = biasC; rowBase = gRow - p.o1 - p.o2; } - const int nr = min(R, p.O - gRow); - const int nt = min(TT, p.T); + constexpr int NB = QXK / 128; // 128-wide blocks per staged chunk + threadgroup float Xs[QXK * 8]; // [block][step kk][fragment row c][token] + const int sgRow = (int(tgid.x) * QXNSG + int(sgid)) * (8 * QXRF); + const int nt = min(8, p.T); + const uint qid = lane / 4; + const int fm = int((qid & 4) + ((lane / 2) % 4)); // fragment row: weight row / k slot of X + const int fn = int((qid & 2) * 2 + (lane % 2) * 2); // fragment columns fn, fn + 1 const ulong rowBytes = ulong(p.K) / 2; const int gpr = p.K / p.gs; - float acc[R][TT]; - for (int r = 0; r < R; r++) for (int t = 0; t < TT; t++) acc[r][t] = 0.0f; - device const uchar* wb = w + ulong(rowBase) * rowBytes; - device const T* sb = sc + rowBase * gpr; - device const T* bb = bi + rowBase * gpr; - const float4 sLo = float4(1.0f, 1.0f / 16.0f, 1.0f / 256.0f, 1.0f / 4096.0f); - const float4 sHi = sLo * (1.0f / 65536.0f); - const uint4 mLo = uint4(0xFu, 0xF0u, 0xF00u, 0xF000u); - const uint4 mHi = mLo << 16; - // 8 weights (one packed uint) per lane per step: the activations of every token - // fit in registers, and each weight is dequantised once and reused by all tokens. - for (int k = int(lane) * 8; k < p.K; k += 256) { - float4 xl[TT], xh[TT]; - float xs[TT]; - for (int t = 0; t < TT; t++) { - device const float4* xp = (device const float4*)(x + min(t, nt - 1) * p.xStride + k); - const float4 a = xp[0], b = xp[1]; - xs[t] = dot(a + b, float4(1.0f)); - xl[t] = a * sLo; xh[t] = b * sHi; + const int jOff = (fn / 2) * 32; // this lane's 32-value run in a block + // Per fragment: the slot ([A|B|C]; o1, o2 are multiples of 8) and this lane's row in it. + // Fragments past the last row compute nothing, but every thread stages x and reaches the barriers. + device const uchar* wrow[QXRF]; device const T* srow[QXRF]; device const T* brow[QXRF]; + float bias[QXRF]; + bool live[QXRF]; + int nf = 0; // live fragments (a prefix) + for (int r = 0; r < QXRF; r++) { + const int g = sgRow + 8 * r; + live[r] = g < p.O; + nf += live[r] ? 1 : 0; + device const uchar* w; device const T* sc; device const T* bi; device const T* rb; + int rowBase; + if (g < p.o1) { w = wA; sc = sA; bi = bA; rb = biasA; rowBase = g; } + else if (g < p.o1 + p.o2) { w = wB; sc = sB; bi = bB; rb = biasB; rowBase = g - p.o1; } + else { w = wC; sc = sC; bi = bC; rb = biasC; rowBase = g - p.o1 - p.o2; } + const int row = live[r] ? rowBase + min(fm, p.O - g - 1) : 0; + wrow[r] = w + ulong(row) * rowBytes; + srow[r] = sc + row * gpr; + brow[r] = bi + row * gpr; + bias[r] = live[r] && p.hasBias ? float(rb[row]) : 0.0f; + } + simdgroup_matrix C[QXRF][QXACC]; + for (int r = 0; r < QXRF; r++) for (int a = 0; a < QXACC; a++) C[r][a] = simdgroup_matrix(0); + for (int k0 = 0; k0 < p.K; k0 += QXK) { + // This chunk's weights first: the loads do not wait for the x staging. + uint4 wq[NB][QXRF]; + float s[NB][QXRF], b[NB][QXRF]; + for (int i = 0; i < NB; i++) { + const int k = k0 + i * 128 + jOff; + for (int r = 0; r < QXRF; r++) { + if (live[r] && k < p.K) { + wq[i][r] = *(device const uint4*)(wrow[r] + (ulong(k) >> 1)); + s[i][r] = float(srow[r][k / p.gs]); b[i][r] = float(brow[r][k / p.gs]); + } else { wq[i][r] = uint4(0); s[i][r] = 0.0f; b[i][r] = 0.0f; } + } } - const int g = k / p.gs; - const ulong off = ulong(k) >> 1; - for (int r = 0; r < R; r++) { - const int rr = min(r, nr - 1); - const uint wv = *(device const uint*)(wb + ulong(rr) * rowBytes + off); - const float4 ql = float4(uint4(wv) & mLo), qh = float4(uint4(wv) & mHi); - const float s = float(sb[rr * gpr + g]), b = float(bb[rr * gpr + g]); - for (int t = 0; t < TT; t++) - acc[r][t] += s * (dot(ql, xl[t]) + dot(qh, xh[t])) + b * xs[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); // the previous chunk's reads are done + for (int idx = int(tid); idx < 8 * QXK / 4; idx += QXNSG * 32) { + const int t = idx / (QXK / 4), kl = (idx % (QXK / 4)) * 4; + float4 v = float4(0); + if (t < nt && k0 + kl < p.K) v = *(device const float4*)(x + t * p.xStride + k0 + kl); + const int blk = kl / 128, c = ((kl % 128) / 32) * 2, kk = (kl % 32) / 2; + threadgroup float* d = Xs + ((blk * 16 + kk) * 8 + c) * 8 + t; + d[0] = v.x; d[8] = v.y; d[64] = v.z; d[72] = v.w; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (nf == 0) continue; + for (int i = 0; i < NB; i++) { + if (k0 + i * 128 >= p.K) break; + threadgroup const float* xb = Xs + i * 1024 + fm * 8 + fn; + for (int kk = 0; kk < 16; kk++) { + simdgroup_matrix B; + const float2 xv = *(threadgroup const float2*)(xb + kk * 64); + B.thread_elements()[0] = xv.x; + B.thread_elements()[1] = xv.y; + for (int r = 0; r < QXRF; r++) { + const uint u = wq[i][r][kk / 4] >> (8 * (kk % 4)); + simdgroup_matrix A; + A.thread_elements()[0] = float(u & 0xFu) * s[i][r] + b[i][r]; + A.thread_elements()[1] = float((u >> 4) & 0xFu) * s[i][r] + b[i][r]; + simdgroup_multiply_accumulate(C[r][kk % QXACC], A, B, C[r][kk % QXACC]); + } + } } } - for (int r = 0; r < R; r++) { - for (int t = 0; t < TT; t++) { - const float v = simd_sum(acc[r][t]); - if (lane == 0 && r < nr && t < nt) - y[t * p.yStride + gRow + r] = v + (p.hasBias ? float(rb[rowBase + r]) : 0.0f); + for (int r = 0; r < QXRF; r++) { + const int g = sgRow + 8 * r; + if (!live[r] || fm >= p.O - g) continue; + for (int i = 0; i < 2; i++) { + float v = bias[r]; + for (int a = 0; a < QXACC; a++) v += C[r][a].thread_elements()[i]; + if (fn + i < nt) y[(fn + i) * p.yStride + g + fm] = v; } } } @@ -1631,10 +1681,9 @@ INST_QMM_ALL(5) INST_QMM_ALL(6) INST_QMM_ALL(8) -#define QVBR 4 -#define INST_QMVB(T, TN, TT) \ - template [[host_name("qmvb4_" #TN "_" #TT)]] kernel void qmvb4(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, uint2, uint, uint); -INST_QMVB(half, f16, 4) INST_QMVB(bfloat, bf16, 4) INST_QMVB(float, f32, 4) +#define INST_QMVX(T, TN) \ + template [[host_name("qmvx4_" #TN)]] kernel void qmvx4(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, uint2, uint, uint, uint); +INST_QMVX(half, f16) INST_QMVX(bfloat, bf16) INST_QMVX(float, f32) #define INST_DENSE(T, TN) \ template [[host_name("dmm_" #TN "_1")]] kernel void dmm(device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); \ diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index 7664ffc..42354b6 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -7,7 +7,8 @@ private let rowsPerThreadgroup = 8 // QROWS * QNSG private let qmvSimdgroups = 2 // QVNSG private let qmvRowsPerThreadgroup = 4 * qmvSimdgroups private let qmvThreads = 32 * qmvSimdgroups -private let qmvbRowsPerSimdgroup = 4 // QVBR +private let qmvxSimdgroups = 4 // QXNSG +private let qmvxRowsPerThreadgroup = 8 * 2 * qmvxSimdgroups // 8 * QXRF * QXNSG /// A linear layer, either MLX-quantized or dense. final class Linear { @@ -27,7 +28,7 @@ final class Linear { private var gemm64PSO: MTLComputePipelineState? // long prompts: 64 × 64 tiles private(set) var qmvPSO: MTLComputePipelineState? // fused decode kernel (bias folded in) private(set) var qmvGatePSO: MTLComputePipelineState? // decode with y = act(y) · e - private(set) var qmvbPSO: MTLComputePipelineState? // 4-bit, 2...4 tokens (speculative verification) + private(set) var qmvxPSO: MTLComputePipelineState? // 4-bit, 2...8 tokens (speculative verification, short prompts) private var biasPSO: MTLComputePipelineState? let description: String @@ -74,7 +75,7 @@ final class Linear { if bits == 4, [128, 256, 512].contains(inDim % 1024) { tail = "t\(inDim % 1024 / 32)_" } qmvPSO = try gpu.pipeline("qmv_\(tail)\(bits)_\(sfx)") qmvGatePSO = try gpu.pipeline("qmv_gate_\(tail)\(bits)_\(sfx)") - if bits == 4 { qmvbPSO = try gpu.pipeline("qmvb4_\(sfx)_4") } + if bits == 4 { qmvxPSO = try gpu.pipeline("qmvx4_\(sfx)") } } if inDim % 32 == 0 { gemmPSO = try gpu.pipeline("qgemm_\(bits)_\(sfx)") @@ -113,14 +114,14 @@ final class Linear { /// The fused matvec kernels handle T tokens (1, or up to 8 for 4-bit weights). static func canFuse(_ parts: [Linear], T: Int) -> Bool { - T == 1 ? parts.allSatisfy { $0.qmvPSO != nil } : (T <= Transformer.maxVerifyTokens && parts.allSatisfy { $0.qmvbPSO != nil }) + T == 1 ? parts.allSatisfy { $0.qmvPSO != nil } : (T <= Transformer.maxVerifyTokens && parts.allSatisfy { $0.qmvxPSO != nil }) } /// One dispatch computing [A | B | C] rows for T tokens (T == 1, or up to maxVerifyTokens with 4-bit weights). static func encodeFusedDecode(_ enc: MTLComputeCommandEncoder, _ parts: [Linear], x: MTLBuffer, xOffset: Int = 0, xStride: Int = 0, y: MTLBuffer, yOffset: Int = 0, yStride: Int = 0, T: Int = 1) { let f = parts[0] - enc.setComputePipelineState(T == 1 ? f.qmvPSO! : f.qmvbPSO!) + enc.setComputePipelineState(T == 1 ? f.qmvPSO! : f.qmvxPSO!) for slot in 0..<3 { let p = parts[min(slot, parts.count - 1)] let base = slot == 0 ? 0 : 3 + slot * 3 // 0, 6, 9 @@ -133,12 +134,12 @@ final class Linear { enc.set(y, yOffset, 4) enc.set(x, xOffset, 15) let total = parts.reduce(0) { $0 + $1.outDim } - let rowsPerTG = T == 1 ? qmvRowsPerThreadgroup : qmvbRowsPerSimdgroup * qmvSimdgroups + let rowsPerTG = T == 1 ? qmvRowsPerThreadgroup : qmvxRowsPerThreadgroup enc.params(MMParams(K: Int32(f.inDim), O: Int32(total), T: Int32(T), gs: Int32(f.groupSize), xStride: Int32(xStride), yStride: Int32(yStride), o1: Int32(parts[0].outDim), o2: Int32(parts.count > 1 ? parts[1].outDim : 0), hasBias: f.bias != nil ? 1 : 0), 5) enc.dispatchThreadgroups(MTLSize(width: (total + rowsPerTG - 1) / rowsPerTG, height: 1, depth: 1), - threadsPerThreadgroup: MTLSize(width: qmvThreads, height: 1, depth: 1)) + threadsPerThreadgroup: MTLSize(width: T == 1 ? qmvThreads : 32 * qmvxSimdgroups, height: 1, depth: 1)) } /// Decode-only gated output: y[:outDim] = act(x · Wᵀ + bias) · e[:outDim] in one dispatch. @@ -173,11 +174,11 @@ final class Linear { Linear.encodeFusedDecode(enc, [self], x: x, xOffset: xOffset, xStride: xStride, y: y, yOffset: yOffset, yStride: yStride, T: T) return } - // 5+ tokens: tiled GEMM beats the small-batch kernel. A full tile is 32 tokens; short prompts - // use an 8- or 16-token tile instead of paying for 32, long ones a 64 × 64 tile (same - // arithmetic per output element in all of them). + // 5+ tokens (not handled by qmvx4 above): tiled GEMM beats the small-batch kernel. A full + // tile is 32 tokens; short prompts use an 8- or 16-token tile instead of paying for 32, + // long ones a 64 × 64 tile (same arithmetic per output element in all of them). var gemm: (pso: MTLComputePipelineState, rows: Int, cols: Int)? - if T > Transformer.maxVerifyTokens { + if T > 4 { if T <= 8, let p = gemm8PSO { gemm = (p, 8, 32) } else if T <= 16, let p = gemm16PSO { gemm = (p, 16, 32) } else if Linear.useGemm64(T: T, outDim: outDim), let p = gemm64PSO { gemm = (p, 64, 64) } @@ -474,7 +475,7 @@ public final class Transformer { let logits: MTLBuffer /// Logits for every row of a small batch (speculative verification). let logitsAll: MTLBuffer - public static let maxVerifyTokens = 4 + public static let maxVerifyTokens = 8 /// Inverse RoPE frequencies per layer (layers with the same RoPE share a buffer). private let ropeFreqs: [MTLBuffer] @@ -1085,6 +1086,15 @@ extension Transformer { out += String(format: " %@ -%.3f", cls as NSString, (full - stepTime([cls])) * 1000) } out += String(format: "; matvecs only %.3f ms\n", stepTime(["post", "attn", "norm", "silu"]) * 1000) + // Small-batch step (speculative verification, qmvx4) at T = 4 and 8. + for vT in [4, Self.maxVerifyTokens] { + let fullV = stepTime([], T: vT) + out += String(format: "verify step T=%d: %.3f ms; without", vT, fullV * 1000) + for cls in ["qkv", "o", "gateup", "down", "lmhead", "post", "attn", "norm", "silu"] { + out += String(format: " %@ -%.3f", cls as NSString, (fullV - stepTime([cls], T: vT)) * 1000) + } + out += "\n" + } // Prefill step breakdown at 512 tokens (GEMMs, attention and the elementwise kernels). let fullP = stepTime([], T: 512) out += String(format: "prefill step @512: %.1f ms; without", fullP * 1000) diff --git a/Sources/mlxr/main.swift b/Sources/mlxr/main.swift index c5789d8..c387c0a 100644 --- a/Sources/mlxr/main.swift +++ b/Sources/mlxr/main.swift @@ -188,9 +188,13 @@ case "bench": for T in 1...Transformer.maxVerifyTokens { let steps = 32 let s = Date() - for i in 0.. Date: Mon, 2 Nov 2026 12:26:57 +0100 Subject: [PATCH 04/16] WIP: scalar 2-token kernel, per-token decode attention for T <= 8, long-context verify bench Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Kernels.swift | 126 ++++++++++++++++++++++++---- Sources/MLXEngine/Transformer.swift | 43 ++++++---- Sources/mlxr/main.swift | 26 +++--- 3 files changed, 156 insertions(+), 39 deletions(-) diff --git a/Sources/MLXEngine/Kernels.swift b/Sources/MLXEngine/Kernels.swift index 039cd7f..446308a 100644 --- a/Sources/MLXEngine/Kernels.swift +++ b/Sources/MLXEngine/Kernels.swift @@ -440,6 +440,92 @@ kernel void qmvx4(device const uchar* wA [[buffer(0)]], } } +// Small-batch 4-bit matvec, scalar (TT tokens; instantiated for 2): qmv's algebra (x pre-scaled by +// 16^-j, so a nibble is extracted by a mask only) with every token's x in registers, so each weight +// is loaded and unpacked once per TT tokens. KL values per lane per row, R rows per simdgroup share +// one x read. Token rows past p.T repeat the last token (not stored). Same fused [A|B|C] layout as +// `qmv` (R divides 8, so a simdgroup's rows stay in one slot). +#define QTNSG 2 +template +kernel void qmvt4(device const uchar* wA [[buffer(0)]], + device const T* sA [[buffer(1)]], + device const T* bA [[buffer(2)]], + device const float* x [[buffer(3)]], + device float* y [[buffer(4)]], + constant MMParams& p [[buffer(5)]], + device const uchar* wB [[buffer(6)]], + device const T* sB [[buffer(7)]], + device const T* bB [[buffer(8)]], + device const uchar* wC [[buffer(9)]], + device const T* sC [[buffer(10)]], + device const T* bC [[buffer(11)]], + device const T* biasA [[buffer(12)]], + device const T* biasB [[buffer(13)]], + device const T* biasC [[buffer(14)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const int gRow = (int(tgid.x) * QTNSG + int(sgid)) * R; + if (gRow >= p.O) return; + device const uchar* w; device const T* sc; device const T* bi; device const T* rb; + int rowBase; + if (gRow < p.o1) { w = wA; sc = sA; bi = bA; rb = biasA; rowBase = gRow; } + else if (gRow < p.o1 + p.o2) { w = wB; sc = sB; bi = bB; rb = biasB; rowBase = gRow - p.o1; } + else { w = wC; sc = sC; bi = bC; rb = biasC; rowBase = gRow - p.o1 - p.o2; } + const int nr = min(R, p.O - gRow); + const int nt = min(TT, p.T); + const ulong rowBytes = ulong(p.K) / 2; + const int gpr = p.K / p.gs; + float acc[R][TT]; + for (int r = 0; r < R; r++) for (int t = 0; t < TT; t++) acc[r][t] = 0.0f; + device const uchar* wb = w + ulong(rowBase) * rowBytes; + device const T* sb = sc + rowBase * gpr; + device const T* bb = bi + rowBase * gpr; + const float4 sLo = float4(1.0f, 1.0f / 16.0f, 1.0f / 256.0f, 1.0f / 4096.0f); + const float4 sHi = sLo * (1.0f / 65536.0f); + const uint4 mLo = uint4(0xFu, 0xF0u, 0xF00u, 0xF000u); + const uint4 mHi = mLo << 16; + constexpr int NW = KL / 8; // packed uints per lane per row + for (int k = int(lane) * KL; k < p.K; k += 32 * KL) { + float4 xv[TT][2 * NW]; + float xs[TT]; + for (int t = 0; t < TT; t++) { + device const float4* xp = (device const float4*)(x + min(t, nt - 1) * p.xStride + k); + float4 sum = float4(0); + for (int i = 0; i < 2 * NW; i++) { + const float4 v = xp[i]; + sum += v; + xv[t][i] = v * ((i & 1) ? sHi : sLo); + } + xs[t] = dot(sum, float4(1.0f)); + } + const int g = k / p.gs; + const ulong off = ulong(k) >> 1; + for (int r = 0; r < R; r++) { + const int rr = min(r, nr - 1); + uint wv[NW]; + if (NW == 2) { const uint2 u = *(device const uint2*)(wb + ulong(rr) * rowBytes + off); wv[0] = u.x; wv[NW - 1] = u.y; } + else wv[0] = *(device const uint*)(wb + ulong(rr) * rowBytes + off); + float4 q[2 * NW]; + for (int i = 0; i < NW; i++) { q[2 * i] = float4(uint4(wv[i]) & mLo); q[2 * i + 1] = float4(uint4(wv[i]) & mHi); } + const float s = float(sb[rr * gpr + g]), b = float(bb[rr * gpr + g]); + for (int t = 0; t < TT; t++) { + float d = 0.0f; + for (int i = 0; i < 2 * NW; i++) d += dot(q[i], xv[t][i]); + acc[r][t] += s * d + b * xs[t]; + } + } + } + for (int r = 0; r < R; r++) { + for (int t = 0; t < TT; t++) { + const float v = simd_sum(acc[r][t]); + if (lane == 0 && r < nr && t < nt) + y[t * p.yStride + gRow + r] = v + (p.hasBias ? float(rb[rowBase + r]) : 0.0f); + } + } +} + // --------------------------------------------------------------------------- // Prefill path for dense weights (dgemm): tiled GEMM on simdgroup matrices. Y[T, O] = X[T, K] · W[O, K]ᵀ // Threadgroup tile: 32 tokens × 32 output rows, K-step 32, 4 simdgroups (2×2), @@ -1032,15 +1118,18 @@ kernel void attention(device const float* qkv [[buffer(0)]], (threadgroup const float*)ta, lane, out + t * p.outStride + h * p.D); } -// Long-context decode (T = 1), split across threadgroups. The work is the +// Long-context decode, split across threadgroups. The work is the // `attention` kernel's, re-arranged without changing any result: // virtual simdgroup v = tg.y * ASPLIT_NSG + sg takes keys v, v+ASPLIT_V, ... -// for G query heads of KV head tg.x (each K/V row is read once per GQA +// for G query heads of one KV head (each K/V row is read once per GQA // group), and writes its partial (max, sum, acc) per head. attention_combine // then merges the ASPLIT_V partials per head in the same order. Large heads // (256/512, Gemma) take their query heads in sub-groups of G: tg.z selects the // sub-group, so the registers stay bounded and each K/V row is read once per sub-group. -// Partials layout (floats): pm[nH*V] | pl[nH*V] | pa[nH*V*D], row = h*V + v. +// A small batch (T = 2..8, speculative verification) runs every query token t as its own +// decode step at position offset + t: tg.x = kvh · T + t (the tokens of one KV head are +// neighbours in the grid, so their K/V reads meet in the cache). +// Partials layout (floats), per token: pm[nH*V] | pl[nH*V] | pa[nH*V*D], row = h*V + v. #define ASPLIT_V 32 #define ASPLIT_NSG 4 template @@ -1052,16 +1141,16 @@ kernel void attention_split(device const float* qkv [[buffer(0)]], uint3 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { - const int kvh = int(tg.x); + const int t = int(tg.x) % p.T, kvh = int(tg.x) / p.T; const int v = int(tg.y) * ASPLIT_NSG + int(sg); const int h0 = kvh * (p.nH / p.nKV) + int(tg.z) * G; float qr[G][E], acc[G][E], m[G], l[G]; for (int g = 0; g < G; g++) { - device const float* q = qkv + (h0 + g) * p.D; + device const float* q = qkv + t * p.qStride + (h0 + g) * p.D; for (int e = 0; e < E; e++) { qr[g][e] = q[e * 32 + lane] * p.scale; acc[g][e] = 0.0f; } m[g] = -1e30f; l[g] = 0.0f; } - const int L = p.offset + 1; + const int L = p.offset + t + 1; device const half* kb = kc + ulong(kvh) * ulong(p.maxSeq) * ulong(p.D); device const half* vb = vc + ulong(kvh) * ulong(p.maxSeq) * ulong(p.D); for (int j0 = attn_first_key(p, L - 1) + v; j0 < L; j0 += ASPLIT_V * AU) { @@ -1078,25 +1167,28 @@ kernel void attention_split(device const float* qkv [[buffer(0)]], } // Virtual simdgroups with no keys still write (-1e30, 0, 0), exactly like the fused kernel. const int rows = p.nH * ASPLIT_V; + device float* pt = part + ulong(t) * ulong(rows) * ulong(2 + p.D); for (int g = 0; g < G; g++) { const int r = (h0 + g) * ASPLIT_V + v; - if (lane == 0) { part[r] = m[g]; part[rows + r] = l[g]; } - device float* pa = part + 2 * rows + r * p.D; + if (lane == 0) { pt[r] = m[g]; pt[rows + r] = l[g]; } + device float* pa = pt + 2 * rows + r * p.D; for (int e = 0; e < E; e++) pa[e * 32 + lane] = acc[g][e]; } } -// One simdgroup per head: merge the ASPLIT_V partials of attention_split. +// One simdgroup per (head, token): merge the ASPLIT_V partials of attention_split. template kernel void attention_combine(device const float* part [[buffer(0)]], device float* out [[buffer(1)]], constant AttnParams& p [[buffer(2)]], - uint h [[threadgroup_position_in_grid]], + uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + const int h = int(tg.x), t = int(tg.y); const int rows = p.nH * ASPLIT_V; - const int r = int(h) * ASPLIT_V; - attn_merge((device const float*)(part + r), (device const float*)(part + rows + r), - (device const float*)(part + 2 * rows + r * p.D), lane, out + int(h) * p.D); + const int r = h * ASPLIT_V; + device const float* pt = part + ulong(t) * ulong(rows) * ulong(2 + p.D); + attn_merge((device const float*)(pt + r), (device const float*)(pt + rows + r), + (device const float*)(pt + 2 * rows + r * p.D), lane, out + t * p.outStride + h * p.D); } // Prefill attention: one threadgroup per (head, block of QB queries). Each K/V @@ -1684,6 +1776,12 @@ INST_QMM_ALL(8) #define INST_QMVX(T, TN) \ template [[host_name("qmvx4_" #TN)]] kernel void qmvx4(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, uint2, uint, uint, uint); INST_QMVX(half, f16) INST_QMVX(bfloat, bf16) INST_QMVX(float, f32) +// Two tokens, 4 rows per simdgroup, 8 values per lane per row (measured on Qwen3-8B: 16 values +// per lane or 8 rows per simdgroup are slower; with 4 tokens the registers spill or the per-token +// FMAs make it no faster than qmvx4). +#define INST_QMVT(T, TN) \ + template [[host_name("qmvt4_" #TN)]] kernel void qmvt4(device const uchar*, device const T*, device const T*, device const float*, device float*, constant MMParams&, device const uchar*, device const T*, device const T*, device const uchar*, device const T*, device const T*, device const T*, device const T*, device const T*, uint2, uint, uint, uint); +INST_QMVT(half, f16) INST_QMVT(bfloat, bf16) INST_QMVT(float, f32) #define INST_DENSE(T, TN) \ template [[host_name("dmm_" #TN "_1")]] kernel void dmm(device const T*, device const float*, device float*, constant MMParams&, uint2, uint, uint); \ @@ -1728,7 +1826,7 @@ INST_ASPLIT(8, 1) INST_ASPLIT(8, 2) INST_ASPLIT(8, 4) INST_ASPLIT(16, 1) INST_ASPLIT(16, 2) INST_ASPLIT_C(8, 1, true, "_cap") INST_ASPLIT_C(8, 2, true, "_cap") INST_ASPLIT_C(8, 4, true, "_cap") #define INST_ACOMB(E) \ - template [[host_name("attention_combine_" #E)]] kernel void attention_combine(device const float*, device float*, constant AttnParams&, uint, uint); + template [[host_name("attention_combine_" #E)]] kernel void attention_combine(device const float*, device float*, constant AttnParams&, uint2, uint); INST_ACOMB(2) INST_ACOMB(3) INST_ACOMB(4) INST_ACOMB(8) INST_ACOMB(16) #define INST_ATTNP_C(E, QB, CAP, SFX) \ template [[host_name("attention_prefill_" #E SFX)]] kernel void attention_prefill(device const float*, device const half*, device const half*, device float*, constant AttnParams&, uint2, uint, uint, uint); diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index 42354b6..1561858 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -9,6 +9,7 @@ private let qmvRowsPerThreadgroup = 4 * qmvSimdgroups private let qmvThreads = 32 * qmvSimdgroups private let qmvxSimdgroups = 4 // QXNSG private let qmvxRowsPerThreadgroup = 8 * 2 * qmvxSimdgroups // 8 * QXRF * QXNSG +private let qmvtRowsPerThreadgroup = 4 * 2 // R * QTNSG (64 threads) /// A linear layer, either MLX-quantized or dense. final class Linear { @@ -28,7 +29,8 @@ final class Linear { private var gemm64PSO: MTLComputePipelineState? // long prompts: 64 × 64 tiles private(set) var qmvPSO: MTLComputePipelineState? // fused decode kernel (bias folded in) private(set) var qmvGatePSO: MTLComputePipelineState? // decode with y = act(y) · e - private(set) var qmvxPSO: MTLComputePipelineState? // 4-bit, 2...8 tokens (speculative verification, short prompts) + private(set) var qmvtPSO: MTLComputePipelineState? // 4-bit, 2 tokens (speculative verification) + private(set) var qmvxPSO: MTLComputePipelineState? // 4-bit, 3...8 tokens (speculative verification, short prompts) private var biasPSO: MTLComputePipelineState? let description: String @@ -75,7 +77,10 @@ final class Linear { if bits == 4, [128, 256, 512].contains(inDim % 1024) { tail = "t\(inDim % 1024 / 32)_" } qmvPSO = try gpu.pipeline("qmv_\(tail)\(bits)_\(sfx)") qmvGatePSO = try gpu.pipeline("qmv_gate_\(tail)\(bits)_\(sfx)") - if bits == 4 { qmvxPSO = try gpu.pipeline("qmvx4_\(sfx)") } + if bits == 4 { + qmvxPSO = try gpu.pipeline("qmvx4_\(sfx)") + qmvtPSO = try gpu.pipeline("qmvt4_\(sfx)") + } } if inDim % 32 == 0 { gemmPSO = try gpu.pipeline("qgemm_\(bits)_\(sfx)") @@ -121,7 +126,10 @@ final class Linear { static func encodeFusedDecode(_ enc: MTLComputeCommandEncoder, _ parts: [Linear], x: MTLBuffer, xOffset: Int = 0, xStride: Int = 0, y: MTLBuffer, yOffset: Int = 0, yStride: Int = 0, T: Int = 1) { let f = parts[0] - enc.setComputePipelineState(T == 1 ? f.qmvPSO! : f.qmvxPSO!) + // 2 tokens: scalar qmvt4 (per-token FMAs on weights unpacked once); 3...8: qmvx4 (8×8 MMA, + // tokens padded to 8), whose cost hardly depends on T. Both are ALU-bound on an M2 Pro. + let scalar = T == 2 + enc.setComputePipelineState(T == 1 ? f.qmvPSO! : scalar ? f.qmvtPSO! : f.qmvxPSO!) for slot in 0..<3 { let p = parts[min(slot, parts.count - 1)] let base = slot == 0 ? 0 : 3 + slot * 3 // 0, 6, 9 @@ -134,12 +142,13 @@ final class Linear { enc.set(y, yOffset, 4) enc.set(x, xOffset, 15) let total = parts.reduce(0) { $0 + $1.outDim } - let rowsPerTG = T == 1 ? qmvRowsPerThreadgroup : qmvxRowsPerThreadgroup + let rowsPerTG = T == 1 ? qmvRowsPerThreadgroup : scalar ? qmvtRowsPerThreadgroup : qmvxRowsPerThreadgroup + let threads = T == 1 || scalar ? qmvThreads : 32 * qmvxSimdgroups enc.params(MMParams(K: Int32(f.inDim), O: Int32(total), T: Int32(T), gs: Int32(f.groupSize), xStride: Int32(xStride), yStride: Int32(yStride), o1: Int32(parts[0].outDim), o2: Int32(parts.count > 1 ? parts[1].outDim : 0), hasBias: f.bias != nil ? 1 : 0), 5) enc.dispatchThreadgroups(MTLSize(width: (total + rowsPerTG - 1) / rowsPerTG, height: 1, depth: 1), - threadsPerThreadgroup: MTLSize(width: T == 1 ? qmvThreads : 32 * qmvxSimdgroups, height: 1, depth: 1)) + threadsPerThreadgroup: MTLSize(width: threads, height: 1, depth: 1)) } /// Decode-only gated output: y[:outDim] = act(x · Wᵀ + bias) · e[:outDim] in one dispatch. @@ -412,7 +421,7 @@ extension Transformer { let maxD = c.layers.map(\.headDim).max() ?? 0 var bytes = B * 4 * (3 * H + qkv + attn + 3 * I) // h, xn, tmp, qkv, attn, gate|up, mlp bytes += (1 + maxVerifyTokens) * max(c.vocabSize, 1) * 4 // logits, logitsAll - bytes += c.numHeads * attnSplitVirtualSG * (maxD + 2) * 4 // split-attention partials + bytes += maxVerifyTokens * c.numHeads * attnSplitVirtualSG * (maxD + 2) * 4 // split-attention partials if c.perLayerInputSize > 0 { let LP = c.perLayerInputSize * c.numLayers bytes += B * 4 * (2 * LP + c.perLayerInputSize) // per-layer inputs: raw, inputs, gate @@ -647,7 +656,7 @@ public final class Transformer { } kernels[D] = k } - attnPartials = partialBytes > 0 ? try buffer(partialBytes) : nil + attnPartials = partialBytes > 0 ? try buffer(partialBytes * Self.maxVerifyTokens) : nil // per query token attnKernels = kernels if c.perLayerInputSize > 0 { @@ -850,6 +859,10 @@ public final class Transformer { if !skipNorm { layers[0].inputNorm.encode(enc, x: h, y: xn, dim: H, rows: T, rowStride: H) } let decode = T == 1 + // Up to maxVerifyTokens queries run as T decode steps (decode / split attention per token): + // the prefill kernels give one threadgroup per head the whole context, so a verify pass + // grew with the context length far faster than a decode step. + let smallT = T <= Self.maxVerifyTokens for (li, layer) in layers.enumerated() { let spec = layer.spec let D = spec.headDim, nKV = spec.numKVHeads, qkvCols = layer.qkvCols, I = layer.intermediate @@ -892,7 +905,7 @@ public final class Transformer { outStride: Int32(nH * D), T: Int32(T), window: Int32(spec.sliding ? c.slidingWindow : 0), softcap: c.attnLogitSoftcap) if skipAttn { - } else if decode, offset + 1 >= attnSplitMinContext, let (split, combine) = kernels.split, let partials = attnPartials { + } else if smallT, offset + 1 >= attnSplitMinContext, let (split, combine) = kernels.split, let partials = attnPartials { // Partials are reused by every layer: the serial encoder orders the dispatches. enc.setComputePipelineState(split) enc.set(qkvBuf, 0, 0) @@ -900,23 +913,23 @@ public final class Transformer { enc.set(kv.vCache!, 0, 2) enc.set(partials, 0, 3) enc.params(attnParams, 4) - enc.dispatchThreadgroups(MTLSize(width: nKV, height: Self.attnSplitVirtualSG / Self.attnSplitSG, depth: kernels.splitGroups), + enc.dispatchThreadgroups(MTLSize(width: nKV * T, height: Self.attnSplitVirtualSG / Self.attnSplitSG, depth: kernels.splitGroups), threadsPerThreadgroup: MTLSize(width: Self.attnSplitSG * 32, height: 1, depth: 1)) enc.setComputePipelineState(combine) enc.set(partials, 0, 0) enc.set(attnBuf, 0, 1) enc.params(attnParams, 2) - enc.dispatchThreadgroups(MTLSize(width: nH, height: 1, depth: 1), + enc.dispatchThreadgroups(MTLSize(width: nH, height: T, depth: 1), threadsPerThreadgroup: MTLSize(width: 32, height: 1, depth: 1)) } else { - let (attn, attnThreads) = decode ? kernels.decode : kernels.prefill + let (attn, attnThreads) = smallT ? kernels.decode : kernels.prefill enc.setComputePipelineState(attn) enc.set(qkvBuf, 0, 0) enc.set(kv.kCache!, 0, 1) enc.set(kv.vCache!, 0, 2) enc.set(attnBuf, 0, 3) enc.params(attnParams, 4) - enc.dispatchThreadgroups(MTLSize(width: nH, height: decode ? 1 : (T + kernels.prefillQB - 1) / kernels.prefillQB, depth: 1), + enc.dispatchThreadgroups(MTLSize(width: nH, height: smallT ? T : (T + kernels.prefillQB - 1) / kernels.prefillQB, depth: 1), threadsPerThreadgroup: MTLSize(width: attnThreads, height: 1, depth: 1)) } @@ -1064,8 +1077,10 @@ extension Transformer { sweep("o_proj", { [$0.o] }, x: attnBuf) sweep("gate_up", { $0.gateUp }, x: xn) sweep("down", { [$0.down] }, x: mlpBuf) - sweep("gate_up", { $0.gateUp }, x: xn, T: 4) - sweep("down", { [$0.down] }, x: mlpBuf, T: 4) + for vT in [2, 4, Self.maxVerifyTokens] { + sweep("gate_up", { $0.gateUp }, x: xn, T: vT) + sweep("down", { [$0.down] }, x: mlpBuf, T: vT) + } let tLM = time(20) { lmHead.encode($0, x: xn, xStride: H, y: logits, yStride: lmHead.outDim, T: 1) } out += String(format: "lm_head matvec: %.3f ms (%.0f GB/s, cache-warm)\n", tLM * 1000, bytes(lmHead) / tLM / 1e9) // Decode step breakdown: the step without each kernel class (its cost includes dispatch gaps). diff --git a/Sources/mlxr/main.swift b/Sources/mlxr/main.swift index c387c0a..eb053c0 100644 --- a/Sources/mlxr/main.swift +++ b/Sources/mlxr/main.swift @@ -145,8 +145,7 @@ case "bench": print("weights: \(formatBytes(t.weightBytes))") let quick = ProcessInfo.processInfo.environment["MLXR_BENCH_DECODE"] != nil let noProfile = quick || ProcessInfo.processInfo.environment["MLXR_BENCH_NOPROFILE"] != nil - if !noProfile { print(try t.profile()) } - _ = try t.forward([1], offset: 0) + if !noProfile { print(try t.profile()) } _ = try t.forward([1], offset: 0) for n in (quick ? [] : [128, 512, 2048]) { let toks = (0.. 64 { _ = try t.forward((0.. 64 ? "verify @\(ctx)" : "verify" + print(String(format: "%@ T=%d: %6.2f ms per pass (%5.2f ms per token, gpu %.2f ms)", label as NSString, T, dt * 1000, dt * 1000 / Double(T), gpu / Double(steps) * 1000)) } - let dt = Date().timeIntervalSince(s) / Double(steps) - print(String(format: "verify T=%d: %6.2f ms per pass (%5.2f ms per token, gpu %.2f ms)", T, dt * 1000, dt * 1000 / Double(T), gpu / Double(steps) * 1000)) } for T in [5, 8, 12, 15] { let s = Date() From 48e073d3c92941d406b034d6bb9e0b703dff21e6 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 12:28:04 +0100 Subject: [PATCH 05/16] fix bench main.swift --- Sources/mlxr/main.swift | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Sources/mlxr/main.swift b/Sources/mlxr/main.swift index eb053c0..092c376 100644 --- a/Sources/mlxr/main.swift +++ b/Sources/mlxr/main.swift @@ -145,7 +145,8 @@ case "bench": print("weights: \(formatBytes(t.weightBytes))") let quick = ProcessInfo.processInfo.environment["MLXR_BENCH_DECODE"] != nil let noProfile = quick || ProcessInfo.processInfo.environment["MLXR_BENCH_NOPROFILE"] != nil - if !noProfile { print(try t.profile()) } _ = try t.forward([1], offset: 0) + if !noProfile { print(try t.profile()) } + _ = try t.forward([1], offset: 0) for n in (quick ? [] : [128, 512, 2048]) { let toks = (0.. Date: Mon, 2 Nov 2026 12:42:40 +0100 Subject: [PATCH 06/16] qmvx4: document the measured ALU bound Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Kernels.swift | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/Sources/MLXEngine/Kernels.swift b/Sources/MLXEngine/Kernels.swift index 446308a..2c1f69c 100644 --- a/Sources/MLXEngine/Kernels.swift +++ b/Sources/MLXEngine/Kernels.swift @@ -317,10 +317,13 @@ kernel void qmv(device const uchar* wA [[buffer(0)]], } } -// Small-batch 4-bit matmul (T = 2..8 tokens: speculative verification, short prompts) on +// Small-batch 4-bit matmul (T = 3..8 tokens: speculative verification, short prompts) on // simdgroup matrices. Each weight is loaded and dequantized once, independent of T; the // per-token work is one 8×8×8 multiply-accumulate per 64 weights, the tokens padded to 8. -// A simdgroup owns 8 output rows: C[row][token] += W[row][k] · X[token][k], fp32 throughout. +// An 8-row fragment: C[row][token] += W[row][k] · X[token][k], fp32 throughout. +// On an M2 Pro it is ALU-bound at about 3x the T = 1 matvec (with the weight loads replaced by +// constants it costs the same; half-precision fragments are no faster), so it gains over the +// per-token kernels only from 3 tokens on (see qmvt4 for 2). // No threadgroup strip for the weights: every lane dequantizes exactly the two elements of // the W fragment it holds (thread_elements: lane → row fm, columns fn, fn + 1). The sum // over k is order-free, so fragment column c of step kk in a 128-wide block stands for From 8eaf3f3e2780fbd08f0bb7a054f87fa802db88f0 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 13:22:10 +0100 Subject: [PATCH 07/16] Test: a verify batch matches incremental decode (T = 2...8, plain and split attention) Co-Authored-By: Claude Fable 5.1 --- Tests/MLXEngineTests/ModelTests.swift | 34 +++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/Tests/MLXEngineTests/ModelTests.swift b/Tests/MLXEngineTests/ModelTests.swift index c773fcf..a108501 100644 --- a/Tests/MLXEngineTests/ModelTests.swift +++ b/Tests/MLXEngineTests/ModelTests.swift @@ -101,6 +101,40 @@ struct ModelTests { model.attnSplitMinContext = Transformer.defaultSplitMinContext } + /// A speculative verification batch (T = 2...8: qmvt4 / qmvx4 and per-token decode or split + /// attention) gives every row the next-token distribution of decoding the tokens one at a time. + @Test(.enabled(if: installedModel(qwen3) != nil), + arguments: [qwen3, llama, phi, mistral, "mlx-community/Qwen3-8B-4bit", GemmaModelTests.gemma3, GemmaModelTests.gemma4]) + func verifyBatchMatchesIncrementalDecode(repo: String) throws { + guard let dir = installedModel(repo) else { return } + let model = try Transformer(directory: dir, contextLength: 1024) + let tok = try Tokenizer(directory: dir) + let ids = tok.encode(String(repeating: "The quick brown fox jumps over the lazy dog, then it ran into the forest. ", count: 4)).map { Int32($0) } + let V = model.vocabSize, o = 24 + #expect(ids.count >= o + 8) + var worst: Float = 0 + for split in [false, true] { + model.attnSplitMinContext = split ? 1 : Int.max + for T in [2, 3, 5, 8] { + _ = try model.forward(Array(ids.prefix(o)), offset: 0) + let batch = Array(ids[o..<(o + T)]) + let all = Array(UnsafeBufferPointer(start: try model.forward(batch, offset: o, allLogits: true), count: T * V)) + for (i, t) in batch.enumerated() { + let inc = Array(UnsafeBufferPointer(start: try model.forward([t], offset: o + i), count: V)) + let row = all[(i * V)..<((i + 1) * V)] + let argRow = row.indices.max { row[$0] < row[$1] }! - i * V + let argInc = inc.indices.max { inc[$0] < inc[$1] }! + #expect(argRow == argInc, "\(repo) split \(split) T \(T) row \(i)") + let d = zip(row, inc).map { abs($0 - $1) }.max()! + worst = max(worst, d) + #expect(d < 0.1, "\(repo) split \(split) T \(T) row \(i): max logit difference \(d)") + } + } + } + print("verifyBatchMatchesIncrementalDecode \(repo): max logit difference \(worst)") + model.attnSplitMinContext = Transformer.defaultSplitMinContext + } + @Test(.enabled(if: installedModel(qwen3) != nil)) func chatAnswersAndReusesPromptCache() async throws { let engine = InferenceEngine() From 269f75611dbec274773e12dd43b70c1e2cec2eb7 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 11:40:43 +0100 Subject: [PATCH 08/16] Draft-model speculative decoding (engine, CLI --draft, server usage) A small model of the same family drafts greedy tokens (its own KV cache, catch-up and k-1 GPU-sampled decode steps queued back to back, only the ids read back); the target verifies them in one batched pass with the exact acceptance rule. Prompt lookup stays as a free hybrid source. Load checks tokenizer compatibility and the memory of both models; usage reports draft_tokens / accepted_draft_tokens. Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Engine.swift | 256 ++++++++++++++++++++++++--- Sources/MLXEngine/OpenAIServer.swift | 15 +- Sources/MLXEngine/Transformer.swift | 27 +++ Sources/mlxr/main.swift | 32 +++- 4 files changed, 294 insertions(+), 36 deletions(-) diff --git a/Sources/MLXEngine/Engine.swift b/Sources/MLXEngine/Engine.swift index f93aea2..bdd615a 100644 --- a/Sources/MLXEngine/Engine.swift +++ b/Sources/MLXEngine/Engine.swift @@ -26,6 +26,9 @@ public struct GenerationStats: Sendable, Hashable, Codable { /// Speculative decoding: drafted tokens and how many the model accepted. public var draftTokens: Int? = nil public var acceptedDraftTokens: Int? = nil + /// The part of the drafts above that came from the draft model (the rest is prompt lookup). + public var modelDraftTokens: Int? = nil + public var acceptedModelDraftTokens: Int? = nil public init() {} } @@ -49,6 +52,8 @@ public struct LoadedModelInfo: Sendable, Hashable { public let loadSeconds: Double /// The chat template defines native tool calling. public var supportsTools: Bool = false + /// The draft model used for speculative decoding (nil = none). + public var draftName: String? = nil } /// Minimal xorshift RNG so sampling is seedable. @@ -148,6 +153,24 @@ struct Sampler { guard total > 0 else { return (true, draft) } return (false, draw(rest.map { ($0.id, $0.p / total) })) } + + /// Verifies a deterministic draft against one batched pass. `dist(i)` is the target's processed + /// distribution for row i (the token after draft[i-1], row 0 after the pending token); it is + /// called only once every earlier token has been emitted, so repetition penalties see them. + /// Drafts are accepted in order (`verify`); the first rejection emits its correction, and when + /// every draft is accepted the last row yields a bonus token. `emit(token, acceptedDraft)` + /// returns false to stop. Returns the accepted draft count and whether generation continues. + mutating func verifyDraft(_ draft: [Int], dist: (Int) throws -> [(id: Int, p: Float)], + emit: (Int, Bool) -> Bool) rethrows -> (accepted: Int, running: Bool) { + var accepted = 0 + for (i, d) in draft.enumerated() { + let (ok, t) = verify(try dist(i), draft: d) + if ok { accepted += 1 } + guard emit(t, ok) else { return (accepted, false) } + if !ok { return (accepted, true) } + } + return (accepted, emit(draw(try dist(draft.count)), false)) + } } /// Prompt-lookup drafting: if the last few generated tokens already occurred in the @@ -186,6 +209,10 @@ public final class InferenceEngine: @unchecked Sendable { private var template: ChatTemplate? private var stopIds: Set = [] private var cachedTokens: [Int] = [] + /// Speculative decoding with a small model of the same family: its own KV cache, and the + /// token ids by position whose K/V rows that cache holds. + private var draftModel: Transformer? + private var draftCached: [Int] = [] private let cancelLock = NSLock() private var currentToken: CancelToken? /// Sticky trim cut per conversation (key: caller-supplied id, else the first turn). @@ -205,20 +232,40 @@ public final class InferenceEngine: @unchecked Sendable { public var isLoaded: Bool { queue.sync { model != nil } } - public func load(directory: URL, contextLength: Int = 8192) async throws -> LoadedModelInfo { + /// Loads a model. `draftDirectory`: a small model with the same tokenizer that drafts tokens for + /// speculative decoding (see `checkDraftCompatible`); the output is unchanged. + public func load(directory: URL, draftDirectory: URL? = nil, contextLength: Int = 8192) async throws -> LoadedModelInfo { try await withCheckedThrowingContinuation { cont in queue.async { do { let start = Date() self.model = nil + self.draftModel = nil self.tokenizer = nil self.info = nil self.cachedTokens = [] + self.draftCached = [] self.trimState = [:] self.tokenTexts = nil let tok = try Tokenizer(directory: directory) - try Self.checkMemory(directory: directory, contextLength: contextLength) + if let draftDirectory { + try Self.checkDraftCompatible(target: directory, targetTokenizer: tok, draft: draftDirectory) + } + try Self.checkMemory(directories: [directory] + (draftDirectory.map { [$0] } ?? []), contextLength: contextLength) let model = try Transformer(directory: directory, contextLength: contextLength) + var draft: Transformer? + if let draftDirectory { + guard model.supportsBatchedVerify else { + throw EngineError.unsupported("\(directory.lastPathComponent) has no batched verify kernels (4-bit weights are required), so a draft model cannot help") + } + let d = try Transformer(directory: draftDirectory, contextLength: model.maxSeq) + guard d.vocabSize >= min(model.vocabSize, tok.vocabularySize) else { + throw EngineError.unsupported("draft model \(draftDirectory.lastPathComponent) has a smaller vocabulary than \(directory.lastPathComponent)") + } + d.seedSampler(1) + _ = try d.forward([Int32(tok.tokenId(tok.bosToken ?? "") ?? 0)], offset: 0) + draft = d + } let tpl = ChatTemplate(template: tok.chatTemplate, modelType: model.config.modelType, bosToken: tok.bosToken, toolTemplate: tok.toolUseTemplate) var stops = Set(model.config.eosTokenIds) @@ -227,16 +274,18 @@ public final class InferenceEngine: @unchecked Sendable { // Warm-up pass compiles pipelines and faults in weights. _ = try model.forward([Int32(tok.tokenId(tok.bosToken ?? "") ?? 0)], offset: 0) self.model = model + self.draftModel = draft self.tokenizer = tok self.template = tpl self.stopIds = stops - let info = LoadedModelInfo( + var info = LoadedModelInfo( name: directory.lastPathComponent, directory: directory, modelType: model.config.modelType, quantization: model.quantDescription, parameterBytes: model.weightBytes, contextLength: model.maxSeq, layers: model.config.numLayers, hiddenSize: model.config.hiddenSize, chatFormat: tpl.format.rawValue, supportsThinking: tpl.supportsThinkingToggle, loadSeconds: Date().timeIntervalSince(start), supportsTools: tpl.supportsTools) + info.draftName = draft != nil ? draftDirectory?.lastPathComponent : nil self.info = info cont.resume(returning: info) } catch { @@ -250,10 +299,12 @@ public final class InferenceEngine: @unchecked Sendable { cancel() queue.async { self.model = nil + self.draftModel = nil self.tokenizer = nil self.template = nil self.info = nil self.cachedTokens = [] + self.draftCached = [] self.trimState = [:] } } @@ -409,15 +460,30 @@ public final class InferenceEngine: @unchecked Sendable { var nameText = "" var nameOpen = !(nameChoices ?? []).isEmpty let texts = nameOpen ? tokenTextIndex(tok) : nil - // Speculative decoding (prompt lookup): verifying k drafted tokens costs one batched - // pass, so every accepted draft is a token generated without its own weight read. - let speculative = (parameters.speculative ?? false) && model.supportsBatchedVerify && !nameOpen + // Speculative decoding: verifying k drafted tokens costs one batched pass, so every + // accepted draft is a token generated without its own weight read. Drafts come from prompt + // lookup (free) and, when one is loaded, from a small draft model of the same family + // (on unless `speculative` is explicitly false). + let canVerify = model.supportsBatchedVerify && !nameOpen + let draftModel = canVerify && parameters.speculative != false ? self.draftModel : nil + let promptLookup = canVerify && ((parameters.speculative ?? false) || (draftModel != nil && !Self.draftWithoutPromptLookup)) + let speculative = promptLookup || draftModel != nil // Plain decoding selects tokens on the GPU, so each step is queued before the previous one // is read back and the GPU never waits for sampling or detokenisation. let gpuSampling = !speculative && !nameOpen let gpuParams = SampleParams(parameters, vocab: count) cachedTokens = Array(tokens[0..? // CPU sampling only var firstToken: Int? @@ -429,6 +495,14 @@ public final class InferenceEngine: @unchecked Sendable { } else { logits = try model.forward(prompt, offset: common, isCancelled: { token.isCancelled }) } + if let draftModel { + // The draft is advisory: if it fails, generation simply continues without it. + do { + let n = draftCached.count + _ = try draftModel.forward(tokens[n...].map { Int32($0) }, offset: n, isCancelled: { token.isCancelled }) + draftCached = tokens + } catch EngineError.cancelled { throw EngineError.cancelled } catch { draftOK = false } + } } catch EngineError.cancelled { stats.stopReason = "cancelled" emit(.finished(stats)) @@ -442,6 +516,7 @@ public final class InferenceEngine: @unchecked Sendable { emit(.promptProcessed(stats)) var sampler = Sampler(params: parameters, rng: XorShift(seed: seed)) + var shaper = sampler // builds distributions (no randomness), so verifyDraft can call it var decoder = StreamingDecoder() let tGen = Date() stats.stopReason = "length" @@ -449,6 +524,60 @@ public final class InferenceEngine: @unchecked Sendable { var atStart = true var drafted = 0, acceptedDrafts = 0 var draftLen = 4 + // Draft model state: greedy proposals (deterministic, so `Sampler.verify` stays exact) under + // the same repetition penalty; draftCached[.. [Int] { + guard let d = draftModel, draftOK else { return [] } + let P = history.count - 1 + // Rollback: rows of rejected drafts are stale; keep the prefix that matches the history. + var keep = min(draftChecked, draftCached.count) + let lim = min(draftCached.count, P) + while keep < lim && draftCached[keep] == history[keep] { keep += 1 } + draftCached.removeSubrange(keep...) + let k = min(k, d.maxSeq - 1 - P) + guard k > 0 else { return [] } + let start = Date() + do { + let catchUp = history[keep...].map { Int32($0) } + var out: [Int] = [] + var inflight: Transformer.DecodeStep? + if catchUp.count > 16 { // after long prompt-lookup runs only + out.append(try d.prefill(catchUp, offset: keep, sample: draftParams)) + inflight = k > 1 ? try d.beginDecodeStep(offset: P + 1, sample: draftParams) : nil + } else { + inflight = try d.beginSteps(catchUp, offset: keep, sample: draftParams) + } + draftCached = history + draftChecked = P + 1 + while let cur = inflight { + inflight = nil + if cur.offset + 1 < P + k { + do { inflight = try d.beginDecodeStep(offset: cur.offset + 1, sample: draftParams) } + catch { _ = try? d.finish(cur); throw error } + } + do { out.append(try d.finish(cur)) } catch { if let f = inflight { _ = try? d.finish(f) }; throw error } + } + // The cache now holds rows up to d[k-2]; the last draft was never processed. + draftCached += out.dropLast() + draftTime += Date().timeIntervalSince(start) + return out + } catch { + draftCached = Array(history[0.. Bool { @@ -476,7 +605,7 @@ public final class InferenceEngine: @unchecked Sendable { for (i, v) in keep { row[i] = v } } } - guard let d = sampler.distribution(row, count: count, recent: history.suffix(parameters.repetitionWindow)) else { + guard let d = shaper.distribution(row, count: count, recent: history.suffix(parameters.repetitionWindow)) else { throw EngineError.invalidOutput("non-finite logits") } return d @@ -515,35 +644,51 @@ public final class InferenceEngine: @unchecked Sendable { if token.isCancelled { stats.stopReason = "cancelled"; break } let pending = history[history.count - 1] let room = min(maxNew - stats.generatedTokens, model.maxSeq - history.count) - 1 - let draft = speculative ? NGramDraft.propose(history, maxTokens: min(draftLen, room, Transformer.maxVerifyTokens - 1)) : [] + let cap = speculative ? min(room, Transformer.maxVerifyTokens - 1) : 0 + var draft = promptLookup ? NGramDraft.propose(history, maxTokens: min(draftLen, cap)) : [] + let fromModel = draft.isEmpty && cap > 0 && draftOK + if fromModel { draft = proposeFromModel(min(modelDraftLen, cap)) } if draft.isEmpty { let row = try model.forward([Int32(pending)], offset: history.count - 1) cachedTokens = history running = accept(sampler.draw(try dist(row))) continue } + let tVerify = Date() let rows = try model.forward(([pending] + draft).map { Int32($0) }, offset: history.count - 1, allLogits: true) cachedTokens = history drafted += draft.count - var accepted = 0 - var corrected = false - for (i, d) in draft.enumerated() { - let (ok, t) = sampler.verify(try dist(rows + i * vocab), draft: d) - running = accept(t) - if ok { accepted += 1; if running { cachedTokens = history } } - else { corrected = true } - if !ok || !running { break } + // Bonus token: when every draft matches, the last row predicts one more for free. + let (accepted, cont) = try sampler.verifyDraft(draft, dist: { try dist(rows + $0 * vocab) }) { t, ok in + let r = accept(t) + if ok && r { cachedTokens = history } + return r } + running = cont acceptedDrafts += accepted - // Bonus token: every draft matched, so the last row predicts one more for free. - if running && !corrected { running = accept(sampler.draw(try dist(rows + draft.count * vocab))) } - draftLen = accepted == draft.count ? min(draftLen + 1, Transformer.maxVerifyTokens - 1) : max(1, accepted) + if fromModel { + verifyTime += Date().timeIntervalSince(tVerify) + modelCycles += 1 + modelDrafted += draft.count + modelAccepted += accepted + modelDraftLen = accepted == draft.count ? min(modelDraftLen + 1, Transformer.maxVerifyTokens - 1) + : accepted == 0 ? max(1, modelDraftLen - 1) : modelDraftLen + } else { + draftLen = accepted == draft.count ? min(draftLen + 1, Transformer.maxVerifyTokens - 1) : max(1, accepted) + } } let rest = decoder.flush() if !rest.isEmpty { emit(.text(rest)) } let dt = Date().timeIntervalSince(tGen) stats.tokensPerSecond = Double(stats.generatedTokens) / max(dt, 1e-6) if drafted > 0 { stats.draftTokens = drafted; stats.acceptedDraftTokens = acceptedDrafts } + if draftModel != nil { stats.modelDraftTokens = modelDrafted; stats.acceptedModelDraftTokens = modelAccepted } + if Self.traceSpeculation, modelCycles > 0 { + let c = Double(modelCycles) + FileHandle.standardError.write(String(format: "[draft model: %d cycles, accepted %d/%d (%.1f%%), %.2f tokens/cycle, draft %.2f ms + verify %.2f ms per cycle]\n", + modelCycles, modelAccepted, modelDrafted, Double(modelAccepted) / Double(max(modelDrafted, 1)) * 100, + Double(modelAccepted) / c + 1, draftTime / c * 1000, verifyTime / c * 1000).data(using: .utf8)!) + } emit(.finished(stats)) } } @@ -583,8 +728,8 @@ extension InferenceEngine { return allowed } - /// Refuses to load a model whose weights plus KV cache would not fit in GPU-addressable memory. - static func checkMemory(directory: URL, contextLength: Int) throws { + /// Weights + KV cache + scratch of one model at a context length (and the context it gets). + static func memoryNeed(directory: URL, contextLength: Int) throws -> (bytes: Int, context: Int) { let fm = FileManager.default var weights = 0 for f in (try? WeightStore.shardFiles(in: directory)) ?? [] { @@ -593,14 +738,73 @@ extension InferenceEngine { } let cfg = try ModelConfig(directory: directory) let ctx = min(contextLength, max(cfg.maxPositionEmbeddings, 256)) - let kv = cfg.kvCacheBytes(context: ctx) - let scratch = Transformer.scratchBytes(config: cfg) - let need = UInt64(weights + kv + scratch) - let budget = GPU.shared.device.recommendedMaxWorkingSetSize + return (weights + cfg.kvCacheBytes(context: ctx) + Transformer.scratchBytes(config: cfg), ctx) + } + + /// Refuses to load models (a target and its draft) whose weights plus KV caches would not fit + /// in GPU-addressable memory together. + static func checkMemory(directories: [URL], contextLength: Int, budget: UInt64 = GPU.shared.device.recommendedMaxWorkingSetSize) throws { + var need: UInt64 = 0 + var ctx = contextLength + for (i, d) in directories.enumerated() { + let n = try memoryNeed(directory: d, contextLength: contextLength) + need += UInt64(n.bytes) + if i == 0 { ctx = n.context } + } guard need < budget else { let f = ByteCountFormatter() - throw EngineError.unsupported("needs \(f.string(fromByteCount: Int64(need))) (weights + \(ctx)-token KV cache) but the GPU can use about \(f.string(fromByteCount: Int64(budget))). Try a smaller quantization or a shorter context length.") + let what = directories.count > 1 ? "weights of both models + their \(ctx)-token KV caches" : "weights + \(ctx)-token KV cache" + throw EngineError.unsupported("needs \(f.string(fromByteCount: Int64(need))) (\(what)) but the GPU can use about \(f.string(fromByteCount: Int64(budget))). Try \(directories.count > 1 ? "no draft model, " : "")a smaller quantization or a shorter context length.") + } + } + + // MARK: Draft models + + /// Skip prompt lookup when a draft model is loaded (measurement only: `MLXR_DRAFT_NO_LOOKUP`). + static let draftWithoutPromptLookup = ProcessInfo.processInfo.environment["MLXR_DRAFT_NO_LOOKUP"] != nil + /// Print per-request draft-model statistics to stderr (`MLXR_SPEC_TRACE`). + static let traceSpeculation = ProcessInfo.processInfo.environment["MLXR_SPEC_TRACE"] != nil + + /// SHA-256 of a model's tokenizer.json (nil when unreadable). + public static func tokenizerFingerprint(_ directory: URL) -> String? { + let url = directory.appendingPathComponent("tokenizer.json").resolvingSymlinksInPath() + return try? FileHasher.sha256(of: url) + } + + /// A draft model must tokenize exactly like the target: the same tokenizer.json, or else the same + /// vocabulary (every id decodes to the same bytes) and the same special tokens. + static func checkDraftCompatible(target: URL, targetTokenizer tok: Tokenizer, draft: URL) throws { + let a = tokenizerFingerprint(target), b = tokenizerFingerprint(draft) + if let a, a == b { return } + let name = draft.lastPathComponent + let other: Tokenizer + do { other = try Tokenizer(directory: draft) } catch { + throw EngineError.unsupported("draft model \(name): \(error.localizedDescription)") + } + var same = other.vocabularySize == tok.vocabularySize && other.bosToken == tok.bosToken && other.eosToken == tok.eosToken + if same { + for id in 0.. LocalModel? { + let t = target.standardizedFileURL + guard let me = models.first(where: { $0.directory.standardizedFileURL == t }) + ?? ModelStore.describe(target, repo: target.lastPathComponent, source: .library), + me.sizeBytes > 0, let type = me.modelType, let fp = tokenizerFingerprint(target) else { return nil } + return models + .filter { $0.isSupported && $0.modelType == type && $0.sizeBytes > 0 && $0.sizeBytes * 4 <= me.sizeBytes + && $0.directory.standardizedFileURL != t } + .sorted { $0.sizeBytes < $1.sizeBytes } + .first { tokenizerFingerprint($0.directory) == fp } } } diff --git a/Sources/MLXEngine/OpenAIServer.swift b/Sources/MLXEngine/OpenAIServer.swift index eea5ec6..2daf1cf 100644 --- a/Sources/MLXEngine/OpenAIServer.swift +++ b/Sources/MLXEngine/OpenAIServer.swift @@ -1258,18 +1258,29 @@ public final class OpenAIServer: @unchecked Sendable { static func chatUsage(_ o: Outcome) -> [String: Any] { let s = o.stats - return ["prompt_tokens": s.promptTokens, "completion_tokens": s.generatedTokens, + let u: [String: Any] = ["prompt_tokens": s.promptTokens, "completion_tokens": s.generatedTokens, "total_tokens": s.promptTokens + s.generatedTokens, "prompt_tokens_details": ["cached_tokens": s.cachedPromptTokens, "audio_tokens": 0], "completion_tokens_details": ["reasoning_tokens": reasoningTokens(o), "audio_tokens": 0, "accepted_prediction_tokens": 0, "rejected_prediction_tokens": 0]] + return withDraftUsage(u, s) } static func responsesUsage(_ o: Outcome) -> [String: Any] { let s = o.stats - return ["input_tokens": s.promptTokens, "input_tokens_details": ["cached_tokens": s.cachedPromptTokens], + let u: [String: Any] = ["input_tokens": s.promptTokens, "input_tokens_details": ["cached_tokens": s.cachedPromptTokens], "output_tokens": s.generatedTokens, "output_tokens_details": ["reasoning_tokens": reasoningTokens(o)], "total_tokens": s.promptTokens + s.generatedTokens] + return withDraftUsage(u, s) + } + + /// Speculative decoding (extension fields): drafted tokens and how many the model accepted. + static func withDraftUsage(_ u: [String: Any], _ s: GenerationStats) -> [String: Any] { + guard let d = s.draftTokens else { return u } + var u = u + u["draft_tokens"] = d + u["accepted_draft_tokens"] = s.acceptedDraftTokens ?? 0 + return u } /// Stable per-model fingerprint (FNV-1a of name + quantization + context length). diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index 1561858..64f1334 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -813,6 +813,33 @@ public final class Transformer { return DecodeStep(cb: cb, offset: offset) } + /// Writes a few tokens to positions offset.. DecodeStep { + guard !tokens.isEmpty, tokens.count <= maxBatch, offset >= 0, offset + tokens.count < maxSeq else { throw EngineError.contextFull } + guard tokens.allSatisfy({ $0 >= 0 && Int($0) < embed.outDim }) else { + throw EngineError.invalidFile("token id outside the embedding table (\(embed.outDim) rows)") + } + prepared = nil + let n = tokens.count + let single = n <= 4 + writeHistory(tokens, at: offset) // read when the buffers execute + var cb: MTLCommandBuffer? + for i in single ? 0.. Int { step.cb.waitUntilCompleted() diff --git a/Sources/mlxr/main.swift b/Sources/mlxr/main.swift index 092c376..774b65b 100644 --- a/Sources/mlxr/main.swift +++ b/Sources/mlxr/main.swift @@ -5,9 +5,9 @@ import MLXEngine // mlxr search // mlxr download // mlxr list -// mlxr run [-p prompt] [--raw] [--max N] [--temp T] [--no-think] [--ctx N] +// mlxr run [-p prompt] [--raw] [--max N] [--temp T] [--no-think] [--ctx N] [--draft ] // mlxr tokenize -// mlxr serve [--port 1234] [--ctx N] [--lan] +// mlxr serve [--port 1234] [--ctx N] [--lan] [--draft ] setvbuf(stdout, nil, _IONBF, 0) var args = Array(CommandLine.arguments.dropFirst()) @@ -36,6 +36,17 @@ func resolve(_ s: String) -> URL { fail("model not found: \(s) (try `mlxr download \(s)`)") } +/// `--draft `: the draft model for speculative decoding. +func draftDirectory(_ s: String?, target: URL) -> URL? { + guard let s else { return nil } + guard s == "auto" else { return resolve(s) } + guard let m = InferenceEngine.autoDraft(for: target) else { + print("no installed draft model pairs with \(target.lastPathComponent); running without one") + return nil + } + return m.directory +} + func formatBytes(_ b: Int) -> String { ByteCountFormatter.string(fromByteCount: Int64(b), countStyle: .file) } @@ -95,10 +106,12 @@ case "run": let raw = flag("--raw") let noThink = flag("--no-think") let showIds = flag("--ids") + let draftArg = option("--draft") guard let model = args.first else { fail("usage: mlxr run [-p prompt]") } let engine = InferenceEngine.shared - let info = try await engine.load(directory: resolve(model), contextLength: ctx) - print("Loaded \(info.name) [\(info.modelType), \(info.quantization), \(info.layers) layers, \(info.chatFormat)] in \(String(format: "%.2f", info.loadSeconds))s") + let dir = resolve(model) + let info = try await engine.load(directory: dir, draftDirectory: draftDirectory(draftArg, target: dir), contextLength: ctx) + print("Loaded \(info.name) [\(info.modelType), \(info.quantization), \(info.layers) layers, \(info.chatFormat)]\(info.draftName.map { " + draft \($0)" } ?? "") in \(String(format: "%.2f", info.loadSeconds))s") var params = SamplingParameters() params.temperature = temp params.maxTokens = maxTokens @@ -117,7 +130,8 @@ case "run": print(String(format: "[prompt: %d tokens, %.0f tok/s, ttft %.3fs]", s.promptTokens, s.promptTokensPerSecond, s.timeToFirstToken)) case .text(let t): print(t, terminator: "") case .finished(let s): - print(String(format: "\n[generated %d tokens at %.1f tok/s, stop: %@, drafts accepted %d/%d]", s.generatedTokens, s.tokensPerSecond, s.stopReason, s.acceptedDraftTokens ?? 0, s.draftTokens ?? 0)) + let model = s.modelDraftTokens.map { String(format: " (draft model %d/%d)", s.acceptedModelDraftTokens ?? 0, $0) } ?? "" + print(String(format: "\n[generated %d tokens at %.1f tok/s, stop: %@, drafts accepted %d/%d%@]", s.generatedTokens, s.tokensPerSecond, s.stopReason, s.acceptedDraftTokens ?? 0, s.draftTokens ?? 0, model)) } } @@ -125,12 +139,14 @@ case "serve": let port = UInt16(option("--port") ?? "1234") ?? 1234 let ctx = Int(option("--ctx") ?? "8192") ?? 8192 let lan = flag("--lan") - guard let model = args.first else { fail("usage: mlxr serve [--port 1234] [--ctx N] [--lan] [--api-key KEY]") } + let draftArg = option("--draft") + guard let model = args.first else { fail("usage: mlxr serve [--port 1234] [--ctx N] [--lan] [--api-key KEY] [--draft ]") } // LAN mode always needs a key: use --api-key / MLXR_API_KEY, or generate one. var key = option("--api-key") ?? ProcessInfo.processInfo.environment["MLXR_API_KEY"] if lan && key == nil { key = OpenAIServer.generateAPIKey() } - let info = try await InferenceEngine.shared.load(directory: resolve(model), contextLength: ctx) - print("Loaded \(info.name) in \(String(format: "%.2f", info.loadSeconds))s") + let dir = resolve(model) + let info = try await InferenceEngine.shared.load(directory: dir, draftDirectory: draftDirectory(draftArg, target: dir), contextLength: ctx) + print("Loaded \(info.name)\(info.draftName.map { " + draft \($0)" } ?? "") in \(String(format: "%.2f", info.loadSeconds))s") let server = OpenAIServer(port: port, allowLAN: lan, apiKey: key) server.onLog = { print($0) } do { try server.start() } catch { fail("cannot start server: \(error.localizedDescription)") } From a2983ee723043e17a047f78523823fecfceb5d76 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 11:44:29 +0100 Subject: [PATCH 09/16] Speculative decoding tests: exact acceptance on synthetic logits, self-draft greedy identity Co-Authored-By: Claude Fable 5.1 --- Tests/MLXEngineTests/SpeculativeTests.swift | 169 ++++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 Tests/MLXEngineTests/SpeculativeTests.swift diff --git a/Tests/MLXEngineTests/SpeculativeTests.swift b/Tests/MLXEngineTests/SpeculativeTests.swift new file mode 100644 index 0000000..bc61587 --- /dev/null +++ b/Tests/MLXEngineTests/SpeculativeTests.swift @@ -0,0 +1,169 @@ +import Foundation +import Testing +@testable import MLXEngine + +/// Speculative verification on synthetic logits: the emitted tokens must be exactly what the +/// target alone would produce (argmax when greedy, the target's distribution when sampling). +@Suite("Speculative decoding") +struct SpeculativeTests { + /// Runs `verifyDraft` over `rows` (one logits row per draft token plus the bonus row). + static func verify(_ rows: [[Float]], draft: [Int], params: SamplingParameters, sampler: inout Sampler, + stopAfter: Int = .max) -> (emitted: [Int], accepted: Int, running: Bool) { + var shaper = Sampler(params: params, rng: XorShift(seed: 1)) + var emitted: [Int] = [] + let (accepted, running) = sampler.verifyDraft(draft, dist: { i in + var row = rows[i] + return row.withUnsafeMutableBufferPointer { shaper.distribution($0.baseAddress!, count: $0.count, recent: [])! } + }) { t, _ in + emitted.append(t) + return emitted.count < stopAfter + } + return (emitted, accepted, running) + } + + static var greedy: SamplingParameters { + var p = SamplingParameters() + p.temperature = 0 + p.repetitionPenalty = 1 + return p + } + + /// Temperature 1 with every filter off: the processed distribution is exactly softmax(logits). + static var plain: SamplingParameters { + var p = SamplingParameters() + p.temperature = 1; p.topK = 0; p.topP = 1; p.minP = 0; p.repetitionPenalty = 1 + return p + } + + static func onehot(_ i: Int, _ V: Int = 8) -> [Float] { (0.. [Double] { let m = r.max()!; let e = r.map { Double(exp($0 - m)) }; let z = e.reduce(0, +); return e.map { $0 / z } } + let p0 = softmax(row0), p1 = softmax(row1) + var s = Sampler(params: Self.plain, rng: XorShift(seed: 777)) + let n = 200_000 + var pair = [Int](repeating: 0, count: 4) // second token, given the first is the draft (1) + var firstIsDraft = 0 + for _ in 0.. (String, GenerationStats) { + var text = "", stats = GenerationStats() + for try await ev in e.chat(msgs, parameters: p, enableThinking: false) { + if case .text(let t) = ev { text += t } + if case .finished(let s) = ev { stats = s } + } + return (text, stats) + } + let engine = InferenceEngine() + _ = try await engine.load(directory: dir, contextLength: 1024) + let (plain, plainStats) = try await run(engine) + let info = try await engine.load(directory: dir, draftDirectory: dir, contextLength: 1024) + #expect(info.draftName != nil) + let (spec, s) = try await run(engine) + #expect(spec == plain) + #expect(s.generatedTokens == plainStats.generatedTokens) + let drafted = s.modelDraftTokens ?? 0, accepted = s.acceptedModelDraftTokens ?? 0 + #expect(drafted > 50) + #expect(Double(accepted) >= 0.95 * Double(drafted), "accepted \(accepted)/\(drafted)") + // A second turn reuses both caches (the draft's prefix bookkeeping must survive it). + let (again, _) = try await run(engine) + #expect(again == plain) + } +} From dfc1edb1e18ecf437e363b7491fde58a3601c77e Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 11:51:57 +0100 Subject: [PATCH 10/16] App: Draft model picker (Off / Auto / compatible installed models); README Co-Authored-By: Claude Fable 5.1 --- README.md | 3 ++- Sources/MLXRunner/AppState.swift | 13 ++++++++++++- Sources/MLXRunner/ChatView.swift | 19 ++++++++++++++++++- 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 45c9ecd..5082497 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ Each checkpoint above passes end-to-end generation and the prefill-versus-decode - **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a ring of window + 512 K/V rows, so a chat whose shared prefix with the previous request ends more than 512 tokens before that request's end is processed from the start instead of reusing the cache. - **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3, Gemma, Gemma 4) rather than executed as Jinja, so a model with an unusual template may get a near-match format. - **Prefill speed:** about 165 tok/s on a 7B model on an M2 Pro, because the prefill GEMM reaches about 35% of fp32 peak. -- **Speculative decoding is opt-in.** Prompt-lookup drafting with exact verification is implemented and tested (greedy output is identical), but checking 4 drafted tokens on a 7B model still costs about 3 decode steps. So it only helps when most drafts are accepted (copying or editing text), and it is off by default (`SamplingParameters.speculative`, or `mlxr run --spec`). +- **Speculative decoding is opt-in.** Prompt-lookup drafting with exact verification is implemented and tested (greedy output is identical), but checking 4 drafted tokens on a 7B model still costs about 3 decode steps. So it only helps when most drafts are accepted (copying or editing text), and it is off by default (`SamplingParameters.speculative`, or `mlxr run --spec`). A draft model of the same family (`--draft mlx-community/Qwen3-0.6B-4bit`, `--draft auto`, or "Draft model" in the inspector) drafts when prompt lookup has nothing; with Qwen3-8B it gets about half of its drafts accepted, which does not pay off until verifying costs much less than today. - **Glass effects:** set `MLXRUNNER_NO_GLASS=1` to use plain materials instead of Liquid Glass. ## Performance @@ -148,6 +148,7 @@ mlxr download mlx-community/Qwen3-4B-4bit # download into the app's library mlxr list # installed models mlxr run mlx-community/Qwen3-4B-4bit -p "Explain RoPE" --max 300 [--temp 0.7] [--no-think] mlxr serve mlx-community/Qwen3-4B-4bit --port 1234 # OpenAI-compatible server (--lan prints an API key) +mlxr run mlx-community/Qwen3-8B-4bit --draft mlx-community/Qwen3-0.6B-4bit # speculative decoding (or --draft auto) mlxr bench mlx-community/Qwen3-4B-4bit # prefill/decode throughput + kernel micro-benchmarks mlxr tokenize "some text" ``` diff --git a/Sources/MLXRunner/AppState.swift b/Sources/MLXRunner/AppState.swift index 381a005..6cea8e0 100644 --- a/Sources/MLXRunner/AppState.swift +++ b/Sources/MLXRunner/AppState.swift @@ -127,6 +127,8 @@ final class AppState { var sampling = SamplingParameters() { didSet { saveSettings() } } var systemPrompt: String = "" { didSet { saveSettings() } } var contextLength: Int = 8192 { didSet { saveSettings() } } + /// Speculative decoding draft model: "" off, "auto", or a model directory path. + var draftModel: String = "" { didSet { saveSettings() } } var enableThinking: Bool = true { didSet { saveSettings() } } var lastModelPath: String? { didSet { saveSettings() } } @@ -188,9 +190,16 @@ final class AppState { stopGeneration() loadState = .loading(model.name) let ctx = contextLength + let draftSetting = draftModel + let models = localModels Task { do { - let info = try await engine.load(directory: model.directory, contextLength: ctx) + let draft: URL? = switch draftSetting { + case "": nil + case "auto": InferenceEngine.autoDraft(for: model.directory, among: models)?.directory + default: models.first { $0.directory.path == draftSetting && $0.directory != model.directory }?.directory + } + let info = try await engine.load(directory: model.directory, draftDirectory: draft, contextLength: ctx) loadState = .loaded(info) lastModelPath = model.directory.path } catch { @@ -494,6 +503,7 @@ final class AppState { if let data = d.data(forKey: "sampling"), let s = try? JSONDecoder().decode(SamplingParameters.self, from: data) { sampling = s } systemPrompt = d.string(forKey: "systemPrompt") ?? "" contextLength = d.object(forKey: "contextLength") as? Int ?? 8192 + draftModel = d.string(forKey: "draftModel") ?? "" enableThinking = d.object(forKey: "enableThinking") as? Bool ?? true lastModelPath = d.string(forKey: "lastModelPath") } @@ -504,6 +514,7 @@ final class AppState { if let data = try? JSONEncoder().encode(sampling) { d.set(data, forKey: "sampling") } d.set(systemPrompt, forKey: "systemPrompt") d.set(contextLength, forKey: "contextLength") + d.set(draftModel, forKey: "draftModel") d.set(enableThinking, forKey: "enableThinking") d.set(lastModelPath, forKey: "lastModelPath") } diff --git a/Sources/MLXRunner/ChatView.swift b/Sources/MLXRunner/ChatView.swift index e5002a4..7747393 100644 --- a/Sources/MLXRunner/ChatView.swift +++ b/Sources/MLXRunner/ChatView.swift @@ -604,7 +604,12 @@ struct InspectorView: View { Picker("Context length", selection: $state.contextLength) { ForEach([2048, 4096, 8192, 16384, 32768, 65536], id: \.self) { Text("\($0)").tag($0) } } - Text("Context length applies the next time a model is loaded.") + Picker("Draft model", selection: $state.draftModel) { + Text("Off").tag("") + Text("Auto").tag("auto") + ForEach(draftCandidates) { Text($0.name).tag($0.directory.path) } + } + Text("Context length and draft model apply the next time a model is loaded. A small model of the same family drafts tokens for the loaded one to verify: the output is unchanged.") .font(.caption).foregroundStyle(.secondary) } if let info = state.loadedInfo { @@ -614,6 +619,7 @@ struct InspectorView: View { LabeledContent("Layers", value: "\(info.layers)") LabeledContent("Hidden size", value: "\(info.hiddenSize)") LabeledContent("Context", value: "\(info.contextLength)") + if let d = info.draftName { LabeledContent("Draft model", value: d) } LabeledContent("Prompt format", value: info.chatFormat) LabeledContent("Load time", value: String(format: "%.2fs", info.loadSeconds)) } @@ -621,6 +627,17 @@ struct InspectorView: View { } .formStyle(.grouped) } + + /// Installed models that can draft for the loaded one (same family, at most a quarter of its size). + private var draftCandidates: [LocalModel] { + let loaded = state.localModels.first { $0.directory == state.loadedInfo?.directory } + return state.localModels.filter { m in + guard m.isSupported else { return false } + if m.directory.path == state.draftModel { return true } + guard let loaded else { return false } + return m.modelType == loaded.modelType && m.sizeBytes * 4 <= loaded.sizeBytes + } + } } struct SliderRow: View { From 11ca590a5a3fda45fcbf2f1550849190aa0af8ee Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 12:03:07 +0100 Subject: [PATCH 11/16] beginSteps: document its buffers against the in-flight limit Co-Authored-By: Claude Fable 5.1 --- Sources/MLXEngine/Transformer.swift | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Sources/MLXEngine/Transformer.swift b/Sources/MLXEngine/Transformer.swift index 64f1334..7bf77be 100644 --- a/Sources/MLXEngine/Transformer.swift +++ b/Sources/MLXEngine/Transformer.swift @@ -817,7 +817,8 @@ public final class Transformer { /// selects the next token on the GPU, without waiting (draft models: catching up on the tokens /// the target accepted, followed by `beginDecodeStep`s queued back to back). Nothing may be in /// flight. The returned step's `offset` is the last token's position. A few tokens run as single-token - /// steps (the earlier ones without logits), which the matvec kernels do faster than one batch. + /// steps (the earlier ones without logits), which the matvec kernels do faster than one batch; their + /// buffers count as one step for the in-flight limit (they share no per-position sampler state). public func beginSteps(_ tokens: [Int32], offset: Int, sample: SampleParams) throws -> DecodeStep { guard !tokens.isEmpty, tokens.count <= maxBatch, offset >= 0, offset + tokens.count < maxSeq else { throw EngineError.contextFull } guard tokens.allSatisfy({ $0 >= 0 && Int($0) < embed.outDim }) else { From 52da4cf3c70c13aa739344fa8d35932fb08c50e8 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 13:54:21 +0100 Subject: [PATCH 12/16] 1.3.0: README performance table, CHANGELOG, version Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 28 ++++++++++++++++++++++++++++ README.md | 18 +++++++++--------- Sources/MLXEngine/Version.swift | 2 +- 3 files changed, 38 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 090a9a4..f041992 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,33 @@ # Changelog +## 1.3.0 + +Prefill, speculative decoding with a draft model, and a faster 2-token verification. Greedy output is token-identical to 1.2.2 on every supported architecture (Qwen3, Qwen2.5, Llama 3.2, Phi-3.5, Gemma 3, Mistral 7B, Qwen3 8B; short and 1,772-token prompts). + +- **Prefill is 1.5× faster.** The quantized prefill GEMM is one template over the per-simdgroup accumulator grid: the 32×32 tile stays for short batches, and a 64×64 tile (4×4 accumulators, two multiply-accumulates per fragment load instead of one, fully unrolled) runs long ones. Both keep every output's arithmetic chain, so results are bit-identical to the old kernel. Padded threadgroup rows and register prefetch measured slower on an M2 Pro and are not used. `flash_prefill` stages K/V as half (exact: the cache is half) with 16-byte loads of the next block during the current one's multiplies, rescales element-wise instead of through a diagonal matrix multiply, and schedules the longest query blocks first: about 1.9× faster at head size 128 and 1.4–1.6× at 256. Qwen3 0.6B prefills 512 tokens at 3,000 tok/s instead of 2,040; a 1,772-token prompt on Llama 3.2 1B goes from 1,081 to 1,596 tok/s. Decode is unchanged. +- **Draft-model speculative decoding** (`mlxr run --draft `, `mlxr serve --draft`, a "Draft model" picker in the app's inspector, `draft_tokens` / `accepted_draft_tokens` in the server's `usage`). A smaller model of the same family drafts up to `maxVerifyTokens − 1` tokens with the GPU sampler, queued back to back; the target verifies them in one batched pass with the existing exact acceptance rule, so greedy output is identical and sampled output keeps the target's distribution. The draft keeps its own KV cache and prompt-cache mirror, rolls back by position, and is refused at load unless the tokenizers match (same `tokenizer.json`, or the same vocabulary, specials and bytes); the memory check sums both models. `auto` pairs the smallest installed model with the same type and tokenizer at most a quarter the size. Prompt lookup still runs first when it has a match. Qwen3 0.6B drafting for Qwen3 8B and 14B produced identical text to the plain run on short and 1,890-token prompts. **It is off by default and currently slower** (Qwen3 8B: 35.7 → 27.9 tok/s): a verification pass of 3–8 tokens still costs about three decode steps, and the 0.6B's drafts are accepted only about 45% of the time on chat text. It is shipped for measurement and for the day the verify kernel catches up. +- **Verifying 2 tokens costs half of what it did** (Qwen3 8B: 78.6 → 40.7 ms; Qwen3 0.6B: 10.9 → 5.4 ms): a scalar 2-token matvec reads each weight once for both tokens. Batches of 3–8 tokens use a new simdgroup-matrix kernel (`qmvx4`) that dequantizes each weight once, independent of the token count; on an M2 Pro it is ALU-bound at about 3× a single-token step (replacing the weight loads with constants costs the same), so it is no faster than before for 3–4 tokens and extends the batched path to 8. Batches of 2–8 tokens now use the per-token decode attention kernels, so verification at a 1,900-token context grows with the context like a decode step does (8B: 28.8 ms at T=1, 44 ms at T=2) instead of running the prefill attention over the whole context. `maxVerifyTokens` is 8. `mlxr bench` reports verification at a short and a 1,900-token context. +- `mlxr bench`'s profile sweeps warm up first and take the best of three: the first sweep after loading read 2–3× slower than the same matvec inside a decode step (first touch of the buffers, GPU clock ramp), which made the small matrices look latency-bound. `MLXR_BENCH_NOPROFILE=1` skips the profile. +- Tests: the GEMM tiles against the CPU reference up to 517 tokens with odd sizes, tiles agreeing bitwise across batch sizes, prefill vs incremental decode for a 600-token prompt (two batches, the second at offset 512), a verification batch of 2–8 tokens matching one-at-a-time decode with plain and split attention on every installed architecture, exact speculative acceptance on synthetic logits, and self-draft greedy identity. + +Not shipped from the same pass, measured and rejected: folding qkv_post into the decode attention, SiLU into the down projection and the residual add + RMSNorm into the next matvec's prologue (bit-identical, but the step got slower: the norm prologue alone cost +2 ms on Qwen3 0.6B, and widening the matvec kernel's template slowed the plain single-token path by 30%). + +Against mlx-lm 0.32.2 (mlx 0.32.4) on the same weight files, M2 Pro 16 GB: decode is a 128-step single-token loop at a 16-token context (mlx-lm's `stream_generate` rate, which overlaps its CPU work, in brackets), prefill is a timed raw forward of 512 tokens from an empty cache. + +| Model (4-bit) | Decode: MLX Runner | mlx-lm | Prefill 512: MLX Runner | mlx-lm | +|---|---|---|---|---| +| Qwen3 0.6B | 298 tok/s | 239 (260) | 3,000 tok/s | 2,340 | +| Qwen2.5 0.5B | 339 | 244 (296) | 4,240 | 4,030 | +| Llama 3.2 1B | 202 | 199 (205) | 1,720 | 1,720 | +| Gemma 3 1B | 188 | 144 (181) | 2,080 | 2,010 | +| Gemma 2 2B | 94 | 83 (113) | 782 | 791 | +| Gemma 1.1 2B | 78 | 77 (81) | 763 | 868 | +| Gemma 4 E2B | 73 | 69 (76) | 770 | 639 | +| Phi-3.5 mini | 70 | 69 (66) | 466 | 572 | +| Mistral 7B | 41.0 | 42.0 (41.2) | 227 | 305 | + +Decode is level or ahead everywhere (both reach the memory-bandwidth ceiling on the 7B). Prefill is ahead on Qwen3 and Gemma 4, level on Qwen2.5, Llama, Gemma 3 and Gemma 2, and behind on Gemma 1.1 (−12%), Phi-3.5 (−18%) and Mistral 7B (−26%); at 128 tokens, where the 32×32 tile runs, it is 10–35% behind on every model. + ## 1.2.2 - Fixes the chat window freezing during a long reply (any model) and then showing the rest at once, and the transcript going blank when a reply finished. Sampled on an M2 Pro, the main thread was 100% busy in SwiftUI/AppKit layout while the engine (which streams at 6–7 ms per token) had already finished the reply into the stream buffer. Five causes, all in the chat view: diff --git a/README.md b/README.md index 5082497..a360387 100644 --- a/README.md +++ b/README.md @@ -103,8 +103,8 @@ Each checkpoint above passes end-to-end generation and the prefill-versus-decode - **Unsupported model types:** mixture-of-experts (Mixtral, Qwen3-MoE, Gemma 4 26B-A4B), Gemma 3n, and the vision and audio parts of multimodal models (Gemma 3 and Gemma 4 exports run as text-only). Discover marks unsupported types as not runnable. - **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a ring of window + 512 K/V rows, so a chat whose shared prefix with the previous request ends more than 512 tokens before that request's end is processed from the start instead of reusing the cache. - **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3, Gemma, Gemma 4) rather than executed as Jinja, so a model with an unusual template may get a near-match format. -- **Prefill speed:** about 165 tok/s on a 7B model on an M2 Pro, because the prefill GEMM reaches about 35% of fp32 peak. -- **Speculative decoding is opt-in.** Prompt-lookup drafting with exact verification is implemented and tested (greedy output is identical), but checking 4 drafted tokens on a 7B model still costs about 3 decode steps. So it only helps when most drafts are accepted (copying or editing text), and it is off by default (`SamplingParameters.speculative`, or `mlxr run --spec`). A draft model of the same family (`--draft mlx-community/Qwen3-0.6B-4bit`, `--draft auto`, or "Draft model" in the inspector) drafts when prompt lookup has nothing; with Qwen3-8B it gets about half of its drafts accepted, which does not pay off until verifying costs much less than today. +- **Prefill speed:** about 227 tok/s on a 7B model on an M2 Pro, against about 305 for mlx-lm on the same files. The 64×64 prefill GEMM tile reaches roughly 55–60% of fp32 peak on long batches; prompts shorter than 128 tokens use the 32×32 tile and prefill 10–35% slower than mlx-lm. +- **Speculative decoding is opt-in and, for now, slower.** Prompt lookup (`mlxr run --spec`) and a draft model (`mlxr run --draft `, or the app's "Draft model" setting) both verify exactly, so greedy output is identical to a plain run. But verifying 3–8 drafted tokens costs about three decode steps on this GPU (the batched kernel is ALU-bound, see CHANGELOG 1.3.0), verifying 2 costs about 1.5, and Qwen3 0.6B's drafts for Qwen3 8B are accepted about 45% of the time on chat text: Qwen3 8B generates at 28 tok/s with the draft instead of 36. Both stay off by default until the verify kernel pays for them. - **Glass effects:** set `MLXRUNNER_NO_GLASS=1` to use plain materials instead of Liquid Glass. ## Performance @@ -113,13 +113,13 @@ Measured with `mlxr bench` on an **M2 Pro (16 GB)**. Decode is measured with gre | Model (4-bit) | Weights | Prefill (512 tok) | Decode @16 | Decode @1K ctx | |---|---|---|---|---| -| Qwen2.5 0.5B Instruct | 278 MB | 2,960 tok/s | 358 tok/s | 295 tok/s | -| Qwen3 0.6B | 336 MB | 2,040 tok/s | 289 tok/s | 239 tok/s | -| Llama 3.2 1B Instruct | 695 MB | 1,180 tok/s | 203 tok/s | 188 tok/s | -| Phi-3.5 mini | 2.15 GB | 320 tok/s | 70 tok/s | 61 tok/s | -| Mistral 7B v0.3 | 4.08 GB | 166 tok/s | 41.0 tok/s | 39.5 tok/s | -| Gemma 3 1B (QAT) | 732 MB | 1,450 tok/s | 195 tok/s | 178 tok/s | -| Gemma 3 4B | 2.56 GB | 345 tok/s | 67 tok/s | 62 tok/s | +| Qwen2.5 0.5B Instruct | 278 MB | 4,240 tok/s | 358 tok/s | 295 tok/s | +| Qwen3 0.6B | 336 MB | 3,000 tok/s | 289 tok/s | 239 tok/s | +| Llama 3.2 1B Instruct | 695 MB | 1,720 tok/s | 203 tok/s | 188 tok/s | +| Phi-3.5 mini | 2.15 GB | 466 tok/s | 70 tok/s | 61 tok/s | +| Mistral 7B v0.3 | 4.08 GB | 227 tok/s | 41.0 tok/s | 39.5 tok/s | +| Gemma 3 1B (QAT) | 732 MB | 2,080 tok/s | 195 tok/s | 178 tok/s | +| Gemma 3 4B | 2.56 GB | 494 tok/s | 67 tok/s | 62 tok/s | **The 7B ceiling.** Decoding one token reads every weight once, so speed is capped by memory bandwidth. This M2 Pro measures 197 GB/s with a plain read kernel, which caps a 4.08 GB model at about 48 tok/s. MLX Runner reaches 41.0 tok/s, about 85% of that cap. Reaching 50+ tok/s on a 4-bit 7B model therefore needs one of: a Mac with more bandwidth (M-series Max chips have about 400 GB/s), a 3-bit model (which would raise the ceiling to about 60 tok/s, but only 4-bit has an optimized decode kernel so far), or speculative decoding (see Limitations). diff --git a/Sources/MLXEngine/Version.swift b/Sources/MLXEngine/Version.swift index 51f5fe4..2e6c0f2 100644 --- a/Sources/MLXEngine/Version.swift +++ b/Sources/MLXEngine/Version.swift @@ -1,7 +1,7 @@ /// The MLX Runner release version. The release workflow refuses a tag that doesn't match it, /// and scripts/bundle.sh writes it into the app's Info.plist. public enum MLXRunnerVersion { - public static let current = "1.2.2" + public static let current = "1.3.0" public static let repository = "BackendStack21/mlx-runner" public static let homepage = "https://mlx-runner.21no.de" } From 9f9d023b520c01dbd41f0f486726d498e91eeed5 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 16:58:22 +0100 Subject: [PATCH 13/16] Tokenizer: compare vocabulary keys by exact text String equality is canonical equivalence, so Gemma's ; and U+037E collided: which id a semicolon got varied per process, and the losing id decoded to nothing. Token ids now match transformers on the tested Gemma and Qwen3 tokenizers. Co-Authored-By: Claude Opus 5.5 --- Sources/MLXEngine/Tokenizer.swift | 74 +++++++++++-------- .../MLXEngineTests/TokenizerVocabTests.swift | 53 +++++++++++++ 2 files changed, 97 insertions(+), 30 deletions(-) create mode 100644 Tests/MLXEngineTests/TokenizerVocabTests.swift diff --git a/Sources/MLXEngine/Tokenizer.swift b/Sources/MLXEngine/Tokenizer.swift index 6e1b440..f818591 100644 --- a/Sources/MLXEngine/Tokenizer.swift +++ b/Sources/MLXEngine/Tokenizer.swift @@ -1,5 +1,18 @@ import Foundation +/// A token's exact text. `String` equality is canonical equivalence, but tokenizer.json keys are +/// exact: Gemma has both ";" and U+037E, which would otherwise collide (and which of the two ids +/// won depended on the process's hash seed). +struct ExactKey: Hashable, Sendable { + let text: String + init(_ text: String) { self.text = text } + static func == (a: ExactKey, b: ExactKey) -> Bool { a.text.utf8.elementsEqual(b.text.utf8) } + func hash(into h: inout Hasher) { + var s = text + s.withUTF8 { h.combine(bytes: UnsafeRawBufferPointer($0)) } + } +} + /// BPE tokenizer driven by a Hugging Face `tokenizer.json`. /// Supports byte-level BPE (GPT-2, Llama 3, Qwen, Phi) and /// SentencePiece-style BPE with byte fallback (Llama 2, Mistral). @@ -7,11 +20,11 @@ public final class Tokenizer: @unchecked Sendable { enum Mode { case byteLevel, sentencePiece } private(set) var mode: Mode = .byteLevel - private var vocab: [String: Int] = [:] + private var vocab: [ExactKey: Int] = [:] private var idToToken: [String] = [] - private var mergeRanks: [String: Int] = [:] + private var mergeRanks: [ExactKey: Int] = [:] private var specialIds: Set = [] - private var addedTokens: [String: Int] = [:] + private var addedTokens: [ExactKey: Int] = [:] /// Leftmost-longest matcher over every added token (a 6,000-way regex alternation made /// Gemma 3 encoding 100× slower). private var addedMatcher: AddedTokenMatcher? @@ -24,7 +37,7 @@ public final class Tokenizer: @unchecked Sendable { private var spPrependEverySegment = false private var spPrependUnconditional = false /// Added tokens that swallow adjacent whitespace (`lstrip` / `rstrip` in tokenizer.json). - private var addedStrip: [String: (left: Bool, right: Bool)] = [:] + private var addedStrip: [ExactKey: (left: Bool, right: Bool)] = [:] /// The model expects a BOS token at the start of raw prompts. public private(set) var addBosToken = false /// Decoding drops the dummy-prefix space a SentencePiece model puts in front of the first @@ -34,7 +47,7 @@ public final class Tokenizer: @unchecked Sendable { public var usesNFC: Bool { nfc } private var nfc = false private var unkId: Int? - private var cache: [String: [Int]] = [:] + private var cache: [ExactKey: [Int]] = [:] private let cacheLock = NSLock() private static let byteEncoder: [Character] = { @@ -70,49 +83,50 @@ public final class Tokenizer: @unchecked Sendable { let model = json["model"] as? [String: Any] else { throw EngineError.unsupported("tokenizer.json is missing or unreadable") } - guard (model["type"] as? String ?? "BPE") == "BPE", let v = model["vocab"] as? [String: NSNumber] else { + guard (model["type"] as? String ?? "BPE") == "BPE", let v = model["vocab"] as? NSDictionary else { throw EngineError.unsupported("only BPE tokenizers are supported") } var maxId = 0 vocab.reserveCapacity(v.count) // Ids index arrays and GPU tables: reject anything out of range (untrusted file). let idRange = 0..<(1 << 21) - for (k, id) in v { - let i = id.intValue + // Read as an NSDictionary: bridging to [String: _] would itself merge equivalent keys. + for (key, id) in v { + guard let k = key as? String, let i = (id as? NSNumber)?.intValue else { throw EngineError.invalidFile("tokenizer.json: malformed vocabulary entry") } guard idRange.contains(i) else { throw EngineError.invalidFile("tokenizer.json: token id \(i) is out of range") } - vocab[k] = i; maxId = max(maxId, i) + vocab[ExactKey(k)] = i; maxId = max(maxId, i) } if let merges = model["merges"] as? [Any] { mergeRanks.reserveCapacity(merges.count) for (rank, m) in merges.enumerated() { if let s = m as? String { - if let sp = s.firstIndex(of: " ") { - mergeRanks[String(s[.. Int? { vocab[token] } + public func tokenId(_ token: String) -> Int? { vocab[ExactKey(token)] } public func token(_ id: Int) -> String? { id >= 0 && id < idToToken.count ? idToToken[id] : nil } public func isSpecial(_ id: Int) -> Bool { specialIds.contains(id) } @@ -230,19 +244,19 @@ public final class Tokenizer: @unchecked Sendable { // Message content must not forge control tokens; non-special added tokens (Gemma's // HTML tags, Qwen's ) are matched there as the reference tokenizer does. matches.removeAll { m in - specialIds.contains(addedTokens[ns.substring(with: m.range)] ?? -1) + specialIds.contains(addedTokens[ExactKey(ns.substring(with: m.range))] ?? -1) && plainRanges.contains { NSIntersectionRange($0, m.range).length > 0 } } } for m in matches { - let strip = addedStrip[ns.substring(with: m.range)] + let strip = addedStrip[ExactKey(ns.substring(with: m.range))] var segEnd = m.range.location if strip?.left == true { while segEnd > cursor && Self.isSpace(ns.character(at: segEnd - 1)) { segEnd -= 1 } } if segEnd > cursor { encodeOrdinary(ns.substring(with: NSRange(location: cursor, length: segEnd - cursor)), first: firstSegment, into: &out) } - if let id = addedTokens[ns.substring(with: m.range)] { out.append(id) } + if let id = addedTokens[ExactKey(ns.substring(with: m.range))] { out.append(id) } cursor = m.range.location + m.range.length if strip?.right == true { while cursor < ns.length && Self.isSpace(ns.character(at: cursor)) { cursor += 1 } } firstSegment = false @@ -298,25 +312,25 @@ public final class Tokenizer: @unchecked Sendable { private func bpeCached(_ word: String) -> [Int] { cacheLock.lock() - if let c = cache[word] { cacheLock.unlock(); return c } + if let c = cache[ExactKey(word)] { cacheLock.unlock(); return c } cacheLock.unlock() let r = bpe(word) cacheLock.lock() if cache.count > 100_000 { cache.removeAll(keepingCapacity: true) } - cache[word] = r + cache[ExactKey(word)] = r cacheLock.unlock() return r } private func bpe(_ word: String) -> [Int] { - if ignoreMerges, let id = vocab[word] { return [id] } + if ignoreMerges, let id = vocab[ExactKey(word)] { return [id] } var symbols = word.unicodeScalars.map { String($0) } if symbols.count > 1 { while true { var bestRank = Int.max var bestPair = ("", "") for i in 0..<(symbols.count - 1) { - if let r = mergeRanks[symbols[i] + "\u{0}" + symbols[i + 1]], r < bestRank { + if let r = mergeRanks[ExactKey(symbols[i] + "\u{0}" + symbols[i + 1])], r < bestRank { bestRank = r bestPair = (symbols[i], symbols[i + 1]) } @@ -326,7 +340,7 @@ public final class Tokenizer: @unchecked Sendable { merged.reserveCapacity(symbols.count) var i = 0 while i < symbols.count { - if i < symbols.count - 1 && symbols[i] == bestPair.0 && symbols[i + 1] == bestPair.1 { + if i < symbols.count - 1 && symbols[i].utf8.elementsEqual(bestPair.0.utf8) && symbols[i + 1].utf8.elementsEqual(bestPair.1.utf8) { merged.append(bestPair.0 + bestPair.1) i += 2 } else { @@ -341,11 +355,11 @@ public final class Tokenizer: @unchecked Sendable { var ids: [Int] = [] ids.reserveCapacity(symbols.count) for s in symbols { - if let id = vocab[s] { + if let id = vocab[ExactKey(s)] { ids.append(id) } else if byteFallback { for b in s.utf8 { - if let id = vocab[String(format: "<0x%02X>", b)] { ids.append(id) } + if let id = vocab[ExactKey(String(format: "<0x%02X>", b))] { ids.append(id) } } } else if let u = unkId { ids.append(u) @@ -362,7 +376,7 @@ public final class Tokenizer: @unchecked Sendable { if specialIds.contains(id) { return includeSpecial ? Array(tok.utf8) : [] } // Non-special added tokens still go through the decoder in SentencePiece mode (Gemma // lists runs of ▁ as added tokens). - if addedTokens[tok] != nil, mode == .byteLevel { return Array(tok.utf8) } + if addedTokens[ExactKey(tok)] != nil, mode == .byteLevel { return Array(tok.utf8) } switch mode { case .byteLevel: var out: [UInt8] = [] diff --git a/Tests/MLXEngineTests/TokenizerVocabTests.swift b/Tests/MLXEngineTests/TokenizerVocabTests.swift new file mode 100644 index 0000000..8adf1fd --- /dev/null +++ b/Tests/MLXEngineTests/TokenizerVocabTests.swift @@ -0,0 +1,53 @@ +import Foundation +import Testing +@testable import MLXEngine + +/// Swift's `String` equality is canonical equivalence, but a vocabulary is keyed by exact text: +/// Gemma has both ";" and U+037E (its canonical equivalent), each with its own id. +@Suite("Tokenizer vocabulary keys") +struct TokenizerVocabTests { + /// Raw JSON: a Swift dictionary literal would itself collapse the equivalent keys. + private func tokenizer() throws -> Tokenizer { + let dir = FileManager.default.temporaryDirectory.appendingPathComponent("tok-\(UUID().uuidString)") + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + let json = #""" + {"model": {"type": "BPE", "unk_token": "", + "vocab": {"": 0, ";": 1, ";": 2, "Å": 3, "Å": 4, "Ω": 5, "Ω": 6, + "é": 7, "e": 8, "́": 9, "é": 10, "a": 11}, + "merges": [["e", "́"]]}, + "pre_tokenizer": {"type": "Metaspace", "replacement": "▁", "prepend_scheme": "never"}, + "decoder": {"type": "Metaspace", "replacement": "▁", "prepend_scheme": "never"}} + """# + try json.data(using: .utf8)!.write(to: dir.appendingPathComponent("tokenizer.json")) + return try Tokenizer(directory: dir) + } + + @Test func canonicallyEquivalentTokensKeepTheirOwnIds() throws { + let t = try tokenizer() + let cases: [(String, Int)] = [(";", 1), ("\u{37E}", 2), ("\u{C5}", 3), ("\u{212B}", 4), ("\u{3A9}", 5), ("\u{2126}", 6), ("\u{E9}", 7)] + for (text, id) in cases { + #expect(t.encode(text) == [id], "U+\(text.unicodeScalars.map { String($0.value, radix: 16) })") + #expect(t.bytes(for: id) == Array(text.utf8), "id \(id) must decode to its own bytes") + #expect(t.tokenId(text) == id) + } + } + + @Test func mergesAndTheWordCacheAreExact() throws { + let t = try tokenizer() + // Decomposed é merges to its own token; precomposed é is another one, even once cached. + #expect(t.encode("e\u{301}") == [10]) + #expect(t.encode("\u{E9}") == [7]) + #expect(t.encode("e\u{301}") == [10]) + } + + @Test(.enabled(if: installedModel(GemmaModelTests.gemma3) != nil)) + func gemmaSemicolonIsTheAsciiToken() throws { + let t = try Tokenizer(directory: installedModel(GemmaModelTests.gemma3)!) + #expect(t.encode(";").count == 1) + let semi = try #require(t.tokenId(";")) + let greek = try #require(t.tokenId("\u{37E}")) + #expect(semi != greek) + #expect(t.encode("a; b").contains(semi) && !t.encode("a; b").contains(greek)) + #expect(t.bytes(for: semi) == Array(";".utf8) && t.bytes(for: greek) == Array("\u{37E}".utf8)) + } +} From 437b29ae1a62d3f57fc94422f469ff9339ded99f Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 16:58:22 +0100 Subject: [PATCH 14/16] Draft model: review fixes The draft catches up on the prompt after the first token, a non-4-bit target is refused before loading, a saved pick that does not fit is dropped instead of failing the load, tokenizer hashing leaves the main actor, and draft statistics move from env-gated stderr into GenerationStats. Co-Authored-By: Claude Opus 5.5 --- Sources/MLXEngine/Engine.swift | 36 ++++++++++++-------------------- Sources/MLXRunner/AppState.swift | 25 ++++++++++++++++------ Sources/mlxr/main.swift | 4 ++++ 3 files changed, 36 insertions(+), 29 deletions(-) diff --git a/Sources/MLXEngine/Engine.swift b/Sources/MLXEngine/Engine.swift index bdd615a..b4e9362 100644 --- a/Sources/MLXEngine/Engine.swift +++ b/Sources/MLXEngine/Engine.swift @@ -29,6 +29,10 @@ public struct GenerationStats: Sendable, Hashable, Codable { /// The part of the drafts above that came from the draft model (the rest is prompt lookup). public var modelDraftTokens: Int? = nil public var acceptedModelDraftTokens: Int? = nil + /// Draft-model cycles in this request, and the time they spent drafting and verifying. + public var draftCycles: Int? = nil + public var draftSeconds: Double? = nil + public var verifySeconds: Double? = nil public init() {} } @@ -251,6 +255,9 @@ public final class InferenceEngine: @unchecked Sendable { if let draftDirectory { try Self.checkDraftCompatible(target: directory, targetTokenizer: tok, draft: draftDirectory) } + if draftDirectory != nil, let bits = try ModelConfig(directory: directory).bits, bits != 4 { + throw EngineError.unsupported("\(directory.lastPathComponent) is \(bits)-bit: batched verification needs 4-bit weights, so a draft model cannot help") + } try Self.checkMemory(directories: [directory] + (draftDirectory.map { [$0] } ?? []), contextLength: contextLength) let model = try Transformer(directory: directory, contextLength: contextLength) var draft: Transformer? @@ -278,14 +285,13 @@ public final class InferenceEngine: @unchecked Sendable { self.tokenizer = tok self.template = tpl self.stopIds = stops - var info = LoadedModelInfo( + let info = LoadedModelInfo( name: directory.lastPathComponent, directory: directory, modelType: model.config.modelType, quantization: model.quantDescription, parameterBytes: model.weightBytes, contextLength: model.maxSeq, layers: model.config.numLayers, hiddenSize: model.config.hiddenSize, chatFormat: tpl.format.rawValue, supportsThinking: tpl.supportsThinkingToggle, - loadSeconds: Date().timeIntervalSince(start), supportsTools: tpl.supportsTools) - info.draftName = draft != nil ? draftDirectory?.lastPathComponent : nil + loadSeconds: Date().timeIntervalSince(start), supportsTools: tpl.supportsTools, draftName: draftDirectory?.lastPathComponent) self.info = info cont.resume(returning: info) } catch { @@ -466,7 +472,7 @@ public final class InferenceEngine: @unchecked Sendable { // (on unless `speculative` is explicitly false). let canVerify = model.supportsBatchedVerify && !nameOpen let draftModel = canVerify && parameters.speculative != false ? self.draftModel : nil - let promptLookup = canVerify && ((parameters.speculative ?? false) || (draftModel != nil && !Self.draftWithoutPromptLookup)) + let promptLookup = canVerify && ((parameters.speculative ?? false) || draftModel != nil) let speculative = promptLookup || draftModel != nil // Plain decoding selects tokens on the GPU, so each step is queued before the previous one // is read back and the GPU never waits for sampling or detokenisation. @@ -495,14 +501,8 @@ public final class InferenceEngine: @unchecked Sendable { } else { logits = try model.forward(prompt, offset: common, isCancelled: { token.isCancelled }) } - if let draftModel { - // The draft is advisory: if it fails, generation simply continues without it. - do { - let n = draftCached.count - _ = try draftModel.forward(tokens[n...].map { Int32($0) }, offset: n, isCancelled: { token.isCancelled }) - draftCached = tokens - } catch EngineError.cancelled { throw EngineError.cancelled } catch { draftOK = false } - } + // The draft model catches up on the prompt lazily (`proposeFromModel`, after the first + // token is out), so the time to first token stays the target's own. } catch EngineError.cancelled { stats.stopReason = "cancelled" emit(.finished(stats)) @@ -683,12 +683,7 @@ public final class InferenceEngine: @unchecked Sendable { stats.tokensPerSecond = Double(stats.generatedTokens) / max(dt, 1e-6) if drafted > 0 { stats.draftTokens = drafted; stats.acceptedDraftTokens = acceptedDrafts } if draftModel != nil { stats.modelDraftTokens = modelDrafted; stats.acceptedModelDraftTokens = modelAccepted } - if Self.traceSpeculation, modelCycles > 0 { - let c = Double(modelCycles) - FileHandle.standardError.write(String(format: "[draft model: %d cycles, accepted %d/%d (%.1f%%), %.2f tokens/cycle, draft %.2f ms + verify %.2f ms per cycle]\n", - modelCycles, modelAccepted, modelDrafted, Double(modelAccepted) / Double(max(modelDrafted, 1)) * 100, - Double(modelAccepted) / c + 1, draftTime / c * 1000, verifyTime / c * 1000).data(using: .utf8)!) - } + if modelCycles > 0 { stats.draftCycles = modelCycles; stats.draftSeconds = draftTime; stats.verifySeconds = verifyTime } emit(.finished(stats)) } } @@ -760,11 +755,6 @@ extension InferenceEngine { // MARK: Draft models - /// Skip prompt lookup when a draft model is loaded (measurement only: `MLXR_DRAFT_NO_LOOKUP`). - static let draftWithoutPromptLookup = ProcessInfo.processInfo.environment["MLXR_DRAFT_NO_LOOKUP"] != nil - /// Print per-request draft-model statistics to stderr (`MLXR_SPEC_TRACE`). - static let traceSpeculation = ProcessInfo.processInfo.environment["MLXR_SPEC_TRACE"] != nil - /// SHA-256 of a model's tokenizer.json (nil when unreadable). public static func tokenizerFingerprint(_ directory: URL) -> String? { let url = directory.appendingPathComponent("tokenizer.json").resolvingSymlinksInPath() diff --git a/Sources/MLXRunner/AppState.swift b/Sources/MLXRunner/AppState.swift index 6cea8e0..8999be7 100644 --- a/Sources/MLXRunner/AppState.swift +++ b/Sources/MLXRunner/AppState.swift @@ -178,7 +178,11 @@ final class AppState { func refreshModels() { Task.detached(priority: .userInitiated) { let models = ModelStore.scan() - await MainActor.run { self.localModels = models } + await MainActor.run { + self.localModels = models + // A draft pick whose model was removed falls back to none (the picker has no such row). + if !["", "auto"].contains(self.draftModel), !models.contains(where: { $0.directory.path == self.draftModel }) { self.draftModel = "" } + } } } @@ -194,11 +198,20 @@ final class AppState { let models = localModels Task { do { - let draft: URL? = switch draftSetting { - case "": nil - case "auto": InferenceEngine.autoDraft(for: model.directory, among: models)?.directory - default: models.first { $0.directory.path == draftSetting && $0.directory != model.directory }?.directory - } + // Resolved off the main actor (tokenizer hashing). A saved pick that no longer fits the + // model being loaded (another family, or removed) is dropped rather than failing the load. + let draft: URL? = await Task.detached(priority: .userInitiated) { + switch draftSetting { + case "": return nil + case "auto": return InferenceEngine.autoDraft(for: model.directory, among: models)?.directory + default: + guard let d = models.first(where: { $0.directory.path == draftSetting }), d.isSupported, d.directory != model.directory, + d.modelType == model.modelType, + let fp = InferenceEngine.tokenizerFingerprint(model.directory), InferenceEngine.tokenizerFingerprint(d.directory) == fp + else { return nil } + return d.directory + } + }.value let info = try await engine.load(directory: model.directory, draftDirectory: draft, contextLength: ctx) loadState = .loaded(info) lastModelPath = model.directory.path diff --git a/Sources/mlxr/main.swift b/Sources/mlxr/main.swift index 774b65b..8b62aad 100644 --- a/Sources/mlxr/main.swift +++ b/Sources/mlxr/main.swift @@ -132,6 +132,10 @@ case "run": case .finished(let s): let model = s.modelDraftTokens.map { String(format: " (draft model %d/%d)", s.acceptedModelDraftTokens ?? 0, $0) } ?? "" print(String(format: "\n[generated %d tokens at %.1f tok/s, stop: %@, drafts accepted %d/%d%@]", s.generatedTokens, s.tokensPerSecond, s.stopReason, s.acceptedDraftTokens ?? 0, s.draftTokens ?? 0, model)) + if let c = s.draftCycles, c > 0, let dt = s.draftSeconds, let vt = s.verifySeconds { + print(String(format: "[draft model: %d cycles, %.2f tokens/cycle, draft %.1f ms + verify %.1f ms per cycle]", c, + Double(s.acceptedModelDraftTokens ?? 0) / Double(c) + 1, dt / Double(c) * 1000, vt / Double(c) * 1000)) + } } } From f8aeef967a74de307b8b68bfcd7db4e43b77eb21 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 16:58:22 +0100 Subject: [PATCH 15/16] Chat: keep message action buttons visible while the pointer is on them Co-Authored-By: Claude Opus 5.5 --- Sources/MLXRunner/ChatView.swift | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/Sources/MLXRunner/ChatView.swift b/Sources/MLXRunner/ChatView.swift index 7747393..ffafb68 100644 --- a/Sources/MLXRunner/ChatView.swift +++ b/Sources/MLXRunner/ChatView.swift @@ -253,6 +253,9 @@ struct MessageRow: View { let conversationID: UUID let isLast: Bool @State private var hovering = false + /// The action buttons also track the pointer themselves, so they stay visible while it is on + /// them even if the row reports that its hover ended. + @State private var hoveringButtons = false @State private var editing = false @State private var editText = "" @@ -291,7 +294,10 @@ struct MessageRow: View { IconButton("pencil", "Edit") { editText = message.content; editing = true } IconButton("doc.on.doc", "Copy") { copy(message.content) } } - .opacity(hovering ? 1 : 0) + .onHover { hoveringButtons = $0 } + .onDisappear { hoveringButtons = false } + .onDisappear { hoveringButtons = false } + .opacity(hovering || hoveringButtons ? 1 : 0) } } } @@ -364,7 +370,8 @@ struct MessageRow: View { IconButton("doc.on.doc", "Copy") { copy(message.content) } if isLast { IconButton("arrow.clockwise", "Regenerate") { state.regenerate(in: conversationID) } } } - .opacity(hovering ? 1 : 0) + .onHover { hoveringButtons = $0 } + .opacity(hovering || hoveringButtons ? 1 : 0) } } From 6b19940502f57e60c0395e9841d10594facefcc3 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Mon, 2 Nov 2026 16:58:22 +0100 Subject: [PATCH 16/16] Docs: 1.3.0 numbers from one measured set; corrected mlx-lm comparison Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 36 +++++++++++------------------------- README.md | 34 +++++++++++++++++++++++++--------- site/index.html | 44 ++++++++++++++++++++++---------------------- site/site.js | 6 +++--- 4 files changed, 61 insertions(+), 59 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f041992..b117213 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,31 +2,17 @@ ## 1.3.0 -Prefill, speculative decoding with a draft model, and a faster 2-token verification. Greedy output is token-identical to 1.2.2 on every supported architecture (Qwen3, Qwen2.5, Llama 3.2, Phi-3.5, Gemma 3, Mistral 7B, Qwen3 8B; short and 1,772-token prompts). - -- **Prefill is 1.5× faster.** The quantized prefill GEMM is one template over the per-simdgroup accumulator grid: the 32×32 tile stays for short batches, and a 64×64 tile (4×4 accumulators, two multiply-accumulates per fragment load instead of one, fully unrolled) runs long ones. Both keep every output's arithmetic chain, so results are bit-identical to the old kernel. Padded threadgroup rows and register prefetch measured slower on an M2 Pro and are not used. `flash_prefill` stages K/V as half (exact: the cache is half) with 16-byte loads of the next block during the current one's multiplies, rescales element-wise instead of through a diagonal matrix multiply, and schedules the longest query blocks first: about 1.9× faster at head size 128 and 1.4–1.6× at 256. Qwen3 0.6B prefills 512 tokens at 3,000 tok/s instead of 2,040; a 1,772-token prompt on Llama 3.2 1B goes from 1,081 to 1,596 tok/s. Decode is unchanged. -- **Draft-model speculative decoding** (`mlxr run --draft `, `mlxr serve --draft`, a "Draft model" picker in the app's inspector, `draft_tokens` / `accepted_draft_tokens` in the server's `usage`). A smaller model of the same family drafts up to `maxVerifyTokens − 1` tokens with the GPU sampler, queued back to back; the target verifies them in one batched pass with the existing exact acceptance rule, so greedy output is identical and sampled output keeps the target's distribution. The draft keeps its own KV cache and prompt-cache mirror, rolls back by position, and is refused at load unless the tokenizers match (same `tokenizer.json`, or the same vocabulary, specials and bytes); the memory check sums both models. `auto` pairs the smallest installed model with the same type and tokenizer at most a quarter the size. Prompt lookup still runs first when it has a match. Qwen3 0.6B drafting for Qwen3 8B and 14B produced identical text to the plain run on short and 1,890-token prompts. **It is off by default and currently slower** (Qwen3 8B: 35.7 → 27.9 tok/s): a verification pass of 3–8 tokens still costs about three decode steps, and the 0.6B's drafts are accepted only about 45% of the time on chat text. It is shipped for measurement and for the day the verify kernel catches up. -- **Verifying 2 tokens costs half of what it did** (Qwen3 8B: 78.6 → 40.7 ms; Qwen3 0.6B: 10.9 → 5.4 ms): a scalar 2-token matvec reads each weight once for both tokens. Batches of 3–8 tokens use a new simdgroup-matrix kernel (`qmvx4`) that dequantizes each weight once, independent of the token count; on an M2 Pro it is ALU-bound at about 3× a single-token step (replacing the weight loads with constants costs the same), so it is no faster than before for 3–4 tokens and extends the batched path to 8. Batches of 2–8 tokens now use the per-token decode attention kernels, so verification at a 1,900-token context grows with the context like a decode step does (8B: 28.8 ms at T=1, 44 ms at T=2) instead of running the prefill attention over the whole context. `maxVerifyTokens` is 8. `mlxr bench` reports verification at a short and a 1,900-token context. -- `mlxr bench`'s profile sweeps warm up first and take the best of three: the first sweep after loading read 2–3× slower than the same matvec inside a decode step (first touch of the buffers, GPU clock ramp), which made the small matrices look latency-bound. `MLXR_BENCH_NOPROFILE=1` skips the profile. -- Tests: the GEMM tiles against the CPU reference up to 517 tokens with odd sizes, tiles agreeing bitwise across batch sizes, prefill vs incremental decode for a 600-token prompt (two batches, the second at offset 512), a verification batch of 2–8 tokens matching one-at-a-time decode with plain and split attention on every installed architecture, exact speculative acceptance on synthetic logits, and self-draft greedy identity. - -Not shipped from the same pass, measured and rejected: folding qkv_post into the decode attention, SiLU into the down projection and the residual add + RMSNorm into the next matvec's prologue (bit-identical, but the step got slower: the norm prologue alone cost +2 ms on Qwen3 0.6B, and widening the matvec kernel's template slowed the plain single-token path by 30%). - -Against mlx-lm 0.32.2 (mlx 0.32.4) on the same weight files, M2 Pro 16 GB: decode is a 128-step single-token loop at a 16-token context (mlx-lm's `stream_generate` rate, which overlaps its CPU work, in brackets), prefill is a timed raw forward of 512 tokens from an empty cache. - -| Model (4-bit) | Decode: MLX Runner | mlx-lm | Prefill 512: MLX Runner | mlx-lm | -|---|---|---|---|---| -| Qwen3 0.6B | 298 tok/s | 239 (260) | 3,000 tok/s | 2,340 | -| Qwen2.5 0.5B | 339 | 244 (296) | 4,240 | 4,030 | -| Llama 3.2 1B | 202 | 199 (205) | 1,720 | 1,720 | -| Gemma 3 1B | 188 | 144 (181) | 2,080 | 2,010 | -| Gemma 2 2B | 94 | 83 (113) | 782 | 791 | -| Gemma 1.1 2B | 78 | 77 (81) | 763 | 868 | -| Gemma 4 E2B | 73 | 69 (76) | 770 | 639 | -| Phi-3.5 mini | 70 | 69 (66) | 466 | 572 | -| Mistral 7B | 41.0 | 42.0 (41.2) | 227 | 305 | - -Decode is level or ahead everywhere (both reach the memory-bandwidth ceiling on the 7B). Prefill is ahead on Qwen3 and Gemma 4, level on Qwen2.5, Llama, Gemma 3 and Gemma 2, and behind on Gemma 1.1 (−12%), Phi-3.5 (−18%) and Mistral 7B (−26%); at 128 tokens, where the 32×32 tile runs, it is 10–35% behind on every model. +Faster prefill, a tokenizer fix for Gemma, speculative decoding with a draft model, and cheaper 2-token verification. + +- **Prefill is about 1.45× faster.** The quantized prefill GEMM is one template over the per-simdgroup accumulator grid: the 32×32 tile stays for short batches, and a 64×64 tile (4×4 accumulators, two multiply-accumulates per fragment load instead of one, fully unrolled) runs long ones, reaching about 3.5 TFLOPS instead of 2.4. Both keep every output's arithmetic chain. Padded threadgroup rows and register prefetch measured slower on an M2 Pro and are not used. `flash_prefill` stages K/V as half (exact: the cache is half) with 16-byte loads of the next block during the current one's multiplies, rescales element-wise instead of through a diagonal matrix multiply, and schedules the longest query blocks first. Qwen3 0.6B prefills 512 tokens at 3,020 tok/s (was 2,040) and Mistral 7B at 247 (was 166). Decode is unchanged. +- **Fixes wrong and unstable tokens on tokenizers with canonically equivalent entries.** Swift compares strings by canonical equivalence, but `tokenizer.json` keys are exact. Gemma lists both `;` and the Greek question mark U+037E, so they collided: which id a `;` got changed from run to run, and the losing id decoded to nothing. Vocabulary, merge, added-token and word-cache lookups now compare exact text. +- **Draft-model speculative decoding** (`mlxr run --draft `, `mlxr serve --draft`, a "Draft model" picker in the app's inspector, `draft_tokens` / `accepted_draft_tokens` in the server's `usage`). A smaller model of the same family drafts tokens with the GPU sampler, queued back to back, and the target verifies them in one batched pass with the existing exact acceptance rule. Greedy output stays identical and sampled output keeps the target's distribution. The draft has its own KV cache, rolls back by position, and catches up on the prompt after the first token, so it does not delay the first token. Loading refuses a draft unless the tokenizers match, and the memory check sums both models. `auto` pairs the smallest installed model of the same type and tokenizer at most a quarter of the size. The app drops a saved pick that does not fit the model being loaded. **It is off by default and currently slower:** Qwen3 8B with a Qwen3 0.6B draft goes from 38.1 to 27.9 tok/s on a short prompt and from 35.6 to 21.3 at a 1,890-token context, because verifying 3–8 tokens costs about three decode steps and about 45% of drafts are accepted on the tested prompts. +- **Verifying 2 tokens costs half of what it did** (Qwen3 8B: 78.6 → 40.7 ms; Qwen3 0.6B: 10.9 → 5.4 ms), through a 2-token matvec that reads each weight once. Batches of 3–8 tokens use a new simdgroup-matrix kernel (`qmvx4`) that dequantizes each weight once. On an M2 Pro it is ALU-bound at about 3× a single-token step (replacing the weight loads with constants costs the same), so 3–4 tokens are no faster than before, but the batched path now reaches 8. Batches of 2–8 tokens use the per-token decode attention instead of the prefill attention. +- Keeps a chat message's Copy, Edit and Regenerate buttons visible while the pointer is on them, for a report of them hiding on hover and becoming unclickable. +- `mlxr bench`: profile sweeps warm up first and take the best of three (the first sweep after loading read 2–3× slow), verification is reported at a short and a 1,900-token context, and `MLXR_BENCH_NOPROFILE=1` skips the profile. +- Tests: GEMM tiles against the CPU reference up to 517 tokens, tiles agreeing bitwise across batch sizes, prefill vs incremental decode for a 600-token prompt, verification batches of 2–8 tokens vs one-at-a-time decode, exact speculative acceptance on synthetic logits, self-draft greedy identity, and tokenizer entries that are canonically equivalent. + +Tried in the same pass and left out: folding `qkv_post` into the decode attention, SiLU into the down projection and the residual add + RMSNorm into the next matvec (same output, but the step got slower on Qwen3 0.6B: 3.5 → 6.6 ms). ## 1.2.2 diff --git a/README.md b/README.md index a360387..5615745 100644 --- a/README.md +++ b/README.md @@ -103,8 +103,8 @@ Each checkpoint above passes end-to-end generation and the prefill-versus-decode - **Unsupported model types:** mixture-of-experts (Mixtral, Qwen3-MoE, Gemma 4 26B-A4B), Gemma 3n, and the vision and audio parts of multimodal models (Gemma 3 and Gemma 4 exports run as text-only). Discover marks unsupported types as not runnable. - **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a ring of window + 512 K/V rows, so a chat whose shared prefix with the previous request ends more than 512 tokens before that request's end is processed from the start instead of reusing the cache. - **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3, Gemma, Gemma 4) rather than executed as Jinja, so a model with an unusual template may get a near-match format. -- **Prefill speed:** about 227 tok/s on a 7B model on an M2 Pro, against about 305 for mlx-lm on the same files. The 64×64 prefill GEMM tile reaches roughly 55–60% of fp32 peak on long batches; prompts shorter than 128 tokens use the 32×32 tile and prefill 10–35% slower than mlx-lm. -- **Speculative decoding is opt-in and, for now, slower.** Prompt lookup (`mlxr run --spec`) and a draft model (`mlxr run --draft `, or the app's "Draft model" setting) both verify exactly, so greedy output is identical to a plain run. But verifying 3–8 drafted tokens costs about three decode steps on this GPU (the batched kernel is ALU-bound, see CHANGELOG 1.3.0), verifying 2 costs about 1.5, and Qwen3 0.6B's drafts for Qwen3 8B are accepted about 45% of the time on chat text: Qwen3 8B generates at 28 tok/s with the draft instead of 36. Both stay off by default until the verify kernel pays for them. +- **Prefill speed:** about 247 tok/s on a 7B model on an M2 Pro, against about 305 for mlx-lm on the same files. The 64×64 GEMM tile measures about 3.5 TFLOPS. Prompts shorter than about 128 tokens use the 32×32 tile and are further behind mlx-lm. +- **Speculative decoding is opt-in and, for now, slower.** Prompt lookup (`mlxr run --spec`) and a draft model (`mlxr run --draft `, or the app's "Draft model" setting) both verify exactly, so greedy output is identical to a plain run. But verifying 3–8 drafted tokens costs about three decode steps on this GPU (the batched kernel is ALU-bound, see CHANGELOG 1.3.0), verifying 2 costs about 1.5, and Qwen3 0.6B's drafts for Qwen3 8B are accepted about 45% of the time on chat text: Qwen3 8B generates at 28 tok/s with the draft instead of 38. Both stay off by default until the verify kernel pays for them. - **Glass effects:** set `MLXRUNNER_NO_GLASS=1` to use plain materials instead of Liquid Glass. ## Performance @@ -113,13 +113,29 @@ Measured with `mlxr bench` on an **M2 Pro (16 GB)**. Decode is measured with gre | Model (4-bit) | Weights | Prefill (512 tok) | Decode @16 | Decode @1K ctx | |---|---|---|---|---| -| Qwen2.5 0.5B Instruct | 278 MB | 4,240 tok/s | 358 tok/s | 295 tok/s | -| Qwen3 0.6B | 336 MB | 3,000 tok/s | 289 tok/s | 239 tok/s | -| Llama 3.2 1B Instruct | 695 MB | 1,720 tok/s | 203 tok/s | 188 tok/s | +| Qwen2.5 0.5B Instruct | 278 MB | 4,270 tok/s | 343 tok/s | 278 tok/s | +| Qwen3 0.6B | 336 MB | 3,020 tok/s | 286 tok/s | 234 tok/s | +| Llama 3.2 1B Instruct | 695 MB | 1,720 tok/s | 202 tok/s | 185 tok/s | | Phi-3.5 mini | 2.15 GB | 466 tok/s | 70 tok/s | 61 tok/s | -| Mistral 7B v0.3 | 4.08 GB | 227 tok/s | 41.0 tok/s | 39.5 tok/s | -| Gemma 3 1B (QAT) | 732 MB | 2,080 tok/s | 195 tok/s | 178 tok/s | -| Gemma 3 4B | 2.56 GB | 494 tok/s | 67 tok/s | 62 tok/s | +| Mistral 7B v0.3 | 4.08 GB | 247 tok/s | 41.0 tok/s | 39.3 tok/s | +| Gemma 3 1B (QAT) | 732 MB | 2,070 tok/s | 189 tok/s | 173 tok/s | +| Gemma 3 4B | 2.56 GB | 498 tok/s | 67 tok/s | 62 tok/s | + +**Compared with mlx-lm.** Same weight files on the same M2 Pro, mlx-lm 0.32.2 (mlx 0.32.4), one run each. Decode is 128 tokens at a short context for MLX Runner and `stream_generate` for mlx-lm, with greedy sampling. Prefill is a timed raw forward of 512 tokens from an empty cache. + +| Model (4-bit) | Decode: MLX Runner | mlx-lm | Prefill 512: MLX Runner | mlx-lm | +|---|---|---|---|---| +| Qwen3 0.6B | 286 tok/s | 260 | 3,020 tok/s | 2,340 | +| Qwen2.5 0.5B | 343 | 296 | 4,270 | 4,030 | +| Llama 3.2 1B | 202 | 205 | 1,720 | 1,720 | +| Gemma 3 1B | 189 | 181 | 2,070 | 2,010 | +| Gemma 2 2B | 94 | 113 | 782 | 791 | +| Gemma 1.1 2B | 78 | 81 | 763 | 868 | +| Gemma 4 E2B | 73 | 76 | 770 | 639 | +| Phi-3.5 mini | 70 | 66 | 466 | 572 | +| Mistral 7B | 41.0 | 41.2 | 247 | 305 | + +MLX Runner decodes faster on the smallest Qwen models and Phi-3.5, is close on most others, and is well behind on Gemma 2. It prefills 512 tokens faster on Qwen3 and Gemma 4 and is behind on Gemma 1.1, Phi-3.5 and Mistral 7B. At 128 tokens, where the 32×32 tile runs, mlx-lm prefills 4–36% faster on every model except Gemma 4, which is level. **The 7B ceiling.** Decoding one token reads every weight once, so speed is capped by memory bandwidth. This M2 Pro measures 197 GB/s with a plain read kernel, which caps a 4.08 GB model at about 48 tok/s. MLX Runner reaches 41.0 tok/s, about 85% of that cap. Reaching 50+ tok/s on a 4-bit 7B model therefore needs one of: a Mac with more bandwidth (M-series Max chips have about 400 GB/s), a 3-bit model (which would raise the ceiling to about 60 tok/s, but only 4-bit has an optimized decode kernel so far), or speculative decoding (see Limitations). @@ -149,7 +165,7 @@ mlxr list # installed models mlxr run mlx-community/Qwen3-4B-4bit -p "Explain RoPE" --max 300 [--temp 0.7] [--no-think] mlxr serve mlx-community/Qwen3-4B-4bit --port 1234 # OpenAI-compatible server (--lan prints an API key) mlxr run mlx-community/Qwen3-8B-4bit --draft mlx-community/Qwen3-0.6B-4bit # speculative decoding (or --draft auto) -mlxr bench mlx-community/Qwen3-4B-4bit # prefill/decode throughput + kernel micro-benchmarks +mlxr bench mlx-community/Qwen3-4B-4bit # prefill/decode throughput, verification cost at short and 1,900-token contexts, kernel profile (MLXR_BENCH_NOPROFILE=1 skips it) mlxr tokenize "some text" ``` diff --git a/site/index.html b/site/index.html index ae9314a..77b1061 100644 --- a/site/index.html +++ b/site/index.html @@ -4,7 +4,7 @@ MLX Runner: run MLX models on your Mac, natively and fast - + @@ -24,13 +24,13 @@ - + - + @@ -90,7 +90,7 @@

LLMs at the speed
of unified memory.

-
decode
358 tok/s
+
decode
343 tok/s
tokens
0
elapsed
0.00 s
@@ -100,8 +100,8 @@

LLMs at the speed
of unified memory.

    -
  • 358tok/s decode, Qwen2.5 0.5B
    • -
  • 2,960tok/s prefill, Qwen2.5 0.5B
    • +
  • 343tok/s decode, Qwen2.5 0.5B
    • +
  • 4,270tok/s prefill, Qwen2.5 0.5B
  • ~85%of memory bandwidth, Mistral 7B
  • 0third-party dependencies
@@ -116,28 +116,28 @@

Same prompt. Five models. Real speeds.

    -
  1. Qwen2.5 0.5B 278 MB0.89 s
    2. -
  2. Qwen3 0.6B 336 MB1.14 s
    3. -
  3. Llama 3.2 1B 695 MB1.69 s
    4. -
  4. Phi-3.5 mini 2.15 GB5.26 s
    5. -
  5. Mistral 7B v0.3 4.08 GB9.33 s
    6. -
  6. Gemma 3 1B 732 MB1.67 s
    7. -
  7. Gemma 3 4B 2.56 GB5.30 s
    8. +
  8. Qwen2.5 0.5B 278 MB0.87 s
    9. +
  9. Qwen3 0.6B 336 MB1.06 s
    10. +
  10. Llama 3.2 1B 695 MB1.57 s
    11. +
  11. Phi-3.5 mini 2.15 GB4.76 s
    12. +
  12. Mistral 7B v0.3 4.08 GB8.32 s
    13. +
  13. Gemma 3 1B 732 MB1.60 s
    14. +
  14. Gemma 3 4B 2.56 GB4.85 s
- + - - - - - - - + + + + + + +
mlxr bench, MLX Runner 1.2.1, M2 Pro (16 GB), 4-bit models.mlxr bench, MLX Runner 1.3.0, M2 Pro (16 GB), 4-bit models.
ModelWeightsPrefill (512 tok)Decode @16Decode @1K ctx
Qwen2.5 0.5B Instruct278 MB2,960 tok/s358 tok/s295 tok/s
Qwen3 0.6B336 MB2,040 tok/s289 tok/s239 tok/s
Llama 3.2 1B Instruct695 MB1,180 tok/s203 tok/s188 tok/s
Phi-3.5 mini2.15 GB320 tok/s70 tok/s61 tok/s
Mistral 7B v0.34.08 GB166 tok/s41.0 tok/s39.5 tok/s
Gemma 3 1B (QAT)732 MB1,450 tok/s195 tok/s178 tok/s
Gemma 3 4B2.56 GB345 tok/s67 tok/s62 tok/s
Qwen2.5 0.5B Instruct278 MB4,270 tok/s343 tok/s278 tok/s
Qwen3 0.6B336 MB3,020 tok/s286 tok/s234 tok/s
Llama 3.2 1B Instruct695 MB1,720 tok/s202 tok/s185 tok/s
Phi-3.5 mini2.15 GB466 tok/s70 tok/s61 tok/s
Mistral 7B v0.34.08 GB247 tok/s41.0 tok/s39.3 tok/s
Gemma 3 1B (QAT)732 MB2,070 tok/s189 tok/s173 tok/s
Gemma 3 4B2.56 GB498 tok/s67 tok/s62 tok/s
@@ -148,7 +148,7 @@

Same prompt. Five models. Real speeds.

The 7B ceiling

Every token reads every weight. Once.

Decoding one token streams the whole model through the GPU, so memory bandwidth sets the limit, not compute. This M2 Pro reads 197 GB/s, which caps a 4.08 GB model at about 48 tok/s.

-

MLX Runner decodes Mistral 7B at 41.0 tok/s, about 85% of that cap. Faster means more bandwidth (an M-series Max: about 400 GB/s), fewer bits, or speculative decoding (mlxr run --spec).

+

MLX Runner decodes Mistral 7B at 41.0 tok/s, about 85% of that cap. Faster means more bandwidth (an M-series Max: about 400 GB/s), fewer bits, or speculative decoding (exact, but on this GPU it still costs more than it saves).

diff --git a/site/site.js b/site/site.js index 582e2cd..f2362de 100644 --- a/site/site.js +++ b/site/site.js @@ -39,9 +39,9 @@ // ── 1. hero: stream an answer at each model's measured decode speed ── const MODELS = [ - { name: "Qwen2.5 0.5B Instruct · 4-bit", rate: 358 }, - { name: "Qwen3 0.6B · 4-bit", rate: 289 }, - { name: "Llama 3.2 1B Instruct · 4-bit", rate: 203 }, + { name: "Qwen2.5 0.5B Instruct · 4-bit", rate: 343 }, + { name: "Qwen3 0.6B · 4-bit", rate: 286 }, + { name: "Llama 3.2 1B Instruct · 4-bit", rate: 202 }, { name: "Phi-3.5 mini · 4-bit", rate: 70 }, { name: "Mistral 7B v0.3 · 4-bit", rate: 41.0 }, ];