diff --git a/CHANGELOG.md b/CHANGELOG.md index cb74e1d..f6ff726 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Changed +- `q8_0` `quantized_matmul` at one activation row takes less time on every GPU. + `KQ_QMV_SPLITK=0` restores the previous kernels. +- The iq2_s fused MoE gate/up and down gathers at one activation row take less + time on every GPU. `KQ_MOE_C32=0` restores the previous kernels. + ## [0.5.1] ### Added diff --git a/README.md b/README.md index 48e9f28..a8b2206 100644 --- a/README.md +++ b/README.md @@ -303,6 +303,11 @@ All optional; the defaults are right for normal use. shared-expert slots stay on the float path). Read per call. Default off: the outputs differ from the float kernels at half rounding level, and an activation past 65504 in magnitude turns to inf. A/B lever. +- `KQ_MOE_C32` - on Metal, the 32-weight decode step of the iq2_s fused MoE gate/up and down + gathers. By default it runs at one activation row on every GPU. `0` restores the 16-weight + kernels, and `1` also runs it at wider rows where the dedupe kernels do not apply. `KQ_MOE_NX=16` + selects its 16-lane form, and `KQ_MOE_NX=32` keeps the 16-weight kernels. Read per call. The + outputs differ from the 16-weight kernels only in summation order. A/B lever. - `KQ_SDPA_IDX_NAX=0` - run `sdpa_fa_indexed` on its simdgroup kernel even where the NAX tile kernel is available. Read once per process. The simdgroup kernel matches `sdpa_fa_verify` over diff --git a/docs/kernels.md b/docs/kernels.md index 881341c..f139eee 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -207,6 +207,15 @@ and one FMA. The row width is a Metal function constant, so each width gets its the superblock count and row stride folded. Under sustained load that ran about 1% faster than reading the width at run time. +A `q8_0` `quantized_matmul` with one activation row and no batch dimensions runs a split-K kernel +too, when N is even and 256 divides K. The eight simdgroups of a threadgroup share two output rows. +They stride the 256-weight steps between them and add their partial sums in threadgroup memory. A +lane keeps the per-row kernels' share of a step, eight weights of one block. It reads them as four +16-bit words, since a quarter block sits 2 bytes past a 2-byte boundary. The split puts more +simdgroups in flight at layer-size N. On M3 Max, back to back over ten shapes, it took 1.6% to 6.6% +less time than the per-row kernels: 2.7% less at K 8192 and N 4096, and 3.2% less at K 4096 and N +8192. In a synced decode step of eight Muse Spark UD-IQ2_S layers, it saved 1.1% per layer. + Under sustained decode the GPU runs below its top clock, where the per-row kernels turn ALU-bound. Measured on M5 Max over 90 seconds of back-to-back mat-vecs at the gemma-4 31B MLP shapes, the split-K kernel read 450-465 GB/s at the lowest clock, against 385-395 GB/s for the per-row kernels @@ -339,10 +348,10 @@ Tuning levers (defaults are right for normal use): - `KQ_QMV_FINE` - tiling of the per-row mat-vec. `1` forces two output rows per threadgroup, `0` forces the coarse tiling, and unset takes the fine tiling at one activation row up to the per-codec N ceiling that `kq.qmv_fine_max_n(codec)` reports. Read live per call. -- `KQ_QMV_SPLITK` - the `q6_k` split-K kernel for one activation row. `0` restores the per-row - kernels, as does a set `KQ_QMV_FINE`, so a tiling A/B keeps measuring the tiling it names. A - lever that forces another route at one row, such as `KQ_VERIFY_NAX=1`, takes precedence. Read - live per call. +- `KQ_QMV_SPLITK` - the `q6_k` and `q8_0` split-K kernels for one activation row. `0` restores + the per-row kernels, as does a set `KQ_QMV_FINE`, so a tiling A/B keeps measuring the tiling it + names. A lever that forces another route at one row, such as `KQ_VERIFY_NAX=1`, takes + precedence. Read live per call. - `KQ_VERIFY_EXT` - the mat-vec route on the M 2-12 band. `1` forces `mv_ext` for every codec with the kernel, `0` forces `verify_qmv` where the codec has it and per-row qmv elsewhere, and unset takes the per-codec default. On NAX GPUs a set value also keeps per-row qmv and NAX split-K off @@ -1029,6 +1038,26 @@ The GLU activation is selected per model: plain SwiGLU/GELU, the clamped `silu_l `swiglu_clamp` (gpt-oss clamped SwiGLU: biases added, sigmoid slope `alpha`, and a `(u + 1)` linear term; requires the expert biases and is instantiated for `mxfp4`/`nvfp4` only). +At one activation row the iq2_s gate/up and down gathers decode 32 weights per lane step, one +32-weight group of a block, where the other codecs decode 16. The block scale, the high-bit byte +and the scale byte then load once per 32 weights, and the grid indices and sign bytes load as +16-bit pairs. Each lane covers one 32-weight group per step at 8 K-lanes per output row, or 16 +under `KQ_MOE_NX=16`. Outputs differ from the 16-weight kernels only in summation order. +`KQ_MOE_C32=0` restores the 16-weight kernels. `KQ_MOE_C32=1` also runs the 32-weight kernels at +wider rows where the dedupe kernels do not apply. A forced `KQ_MOE_NX=32` keeps the 16-weight +kernels. + +On M3 Max at E=256, N=K=4096 and top 8, back to back, the gate/up gather took 0.303 ms against +0.340 ms (-11.1%), and the down gather took 0.166 ms against 0.191 ms (-12.5%). In a synced decode +step of eight Muse Spark UD-IQ2_S layers, the MoE blocks took 8.3% less time. With the `q8_0` +split-K kernel, the step went from 1.414 to 1.335 ms per layer (-5.6%). + +The grid lookup is the next lever. With the grid read removed (wrong output), the 32-weight +gathers took about 25% less time, and 37% less than the 16-weight gathers. That bounds the gain of +a cheaper lookup. Two other grid sources gave no gain. A 64-bit grid in a device buffer matched the +constant table. A 2-bit code table, in threadgroup memory or in a device buffer, was about 20% +slower. + ## Attention Scaled-dot-product variants for shapes stock MLX's fused allowlist excludes, plus the sparse diff --git a/metal/kq_moe_glu_kq.metal b/metal/kq_moe_glu_kq.metal index 8f5215e..f74b23e 100644 --- a/metal/kq_moe_glu_kq.metal +++ b/metal/kq_moe_glu_kq.metal @@ -314,6 +314,29 @@ instantiate_kq_ext_sx2_lowbit(iq1_m, KqIq1_mExt) instantiate_kq_ext_sx_mix(q2_0, KqQ2_0Ext, iq4_nl, KqIq4_nlExt, bfloat16_t) instantiate_kq_ext_sx_mix(q2_0, KqQ2_0Ext, iq4_nl, KqIq4_nlExt, float16_t) +// 32-weight decode step (KqExtC32) for the uniform gathers at decode. +#define instantiate_kq_ext_c32_nx(codec, traits, type, nx, sfx) \ + instantiate_kernel( \ + "kq_" #codec "_moe_glu_gather_silu_c32" sfx "_" #type, \ + kq_ext_moe_glu_gather, type, traits, KQ_GLU_ACT_SILU, nx, true) \ + instantiate_kernel( \ + "kq_" #codec "_moe_glu_gather_gelu_c32" sfx "_" #type, \ + kq_ext_moe_glu_gather, type, traits, KQ_GLU_ACT_GELU, nx, true) \ + instantiate_kernel( \ + "kq_" #codec "_moe_glu_gather_silu_limit_c32" sfx "_" #type, \ + kq_ext_moe_glu_gather, type, traits, KQ_GLU_ACT_SILU_LIMIT, nx, true) \ + instantiate_kernel( \ + "kq_" #codec "_gather_qmv_c32" sfx "_" #type, \ + kq_ext_gather_qmv, type, traits, nx, true) + +#define instantiate_kq_ext_c32(codec, traits) \ + instantiate_kq_ext_c32_nx(codec, traits, bfloat16_t, 8, "") \ + instantiate_kq_ext_c32_nx(codec, traits, bfloat16_t, 16, "_nx16") \ + instantiate_kq_ext_c32_nx(codec, traits, float16_t, 8, "") \ + instantiate_kq_ext_c32_nx(codec, traits, float16_t, 16, "_nx16") + +instantiate_kq_ext_c32(iq2_s, KqIq2_sExt) + // q6_k/q8_0 mix_ns (generic; no tuned ns kernels exist). instantiate_kq_ext_mix_ns(q6_k, KqQ6_KExt) instantiate_kq_ext_mix_ns(q8_0, KqQ8_0Ext) diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index f15d390..5e992b6 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -158,6 +158,7 @@ instantiate_kquant_batched(qmv_fast, type, 32, 8, 0, q8_0) \ instantiate_kquant_batched(qmv_fast, type, 32, 8, 1, q8_0) \ instantiate_kquant_batched(qmv_fast_fine, type, 32, 8, 0, q8_0) \ + instantiate_kquant_batched(qmv_splitk, type, 32, 8, 0, q8_0) \ instantiate_kquant_batched(qmv_fine, type, 32, 8, 0, q8_0) \ instantiate_kquant_batched(qmv, type, 32, 8, 0, q8_0) \ instantiate_kquant_batched(qmv, type, 32, 8, 1, q8_0) \ diff --git a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h index daf5bd9..4c63950 100644 --- a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h +++ b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h @@ -780,11 +780,32 @@ template #define KQ_EXT_NXPSG 8 +// Dot of a 32-weight group (l0, l1 unscaled, sc their two scales) with +// x[0..31]. +template +METAL_FUNC float kq_ext_dot32( + const thread float4x4& l0, + const thread float4x4& l1, + float2 sc, + const device T* xp) { + return sc.x * + (dot(l0[0], float4(*(const device vec*)(xp + 0))) + + dot(l0[1], float4(*(const device vec*)(xp + 4))) + + dot(l0[2], float4(*(const device vec*)(xp + 8))) + + dot(l0[3], float4(*(const device vec*)(xp + 12)))) + + sc.y * + (dot(l1[0], float4(*(const device vec*)(xp + 16))) + + dot(l1[1], float4(*(const device vec*)(xp + 20))) + + dot(l1[2], float4(*(const device vec*)(xp + 24))) + + dot(l1[3], float4(*(const device vec*)(xp + 28)))); +} + // Partial dot of x (one activation row) against wire-byte weight row `row`, // this thread's K-stripe only (chunks tx, tx + NX, ...). `luts` is the // codec's threadgroup-staged table block (KqTgLuts; unused for table-free -// codecs). -template +// codecs). With C32 a lane steps one 32-weight group at a time through +// KqExtC32 (the "_c32" kernels). +template METAL_FUNC float kq_ext_row_partial( const device uint8_t* w, const device T* x, @@ -796,6 +817,19 @@ METAL_FUNC float kq_ext_row_partial( const int nb = K / Codec::superblock; const device uint8_t* w_row = w + row * (int64_t)nb * Codec::block_bytes; float acc = 0.0f; + if (C32 && KqExtC32::ok) { + constexpr short gpb = Codec::superblock / 32; + for (int ig = tx; 32 * ig < K; ig += NX) { + float4x4 l0, l1; + const float2 sc = KqExtC32::deq_group32s( + w_row + (int64_t)(ig / gpb) * Codec::block_bytes, + short(ig % gpb), + l0, + l1); + acc += kq_ext_dot32(l0, l1, sc, x + ig * 32); + } + return acc; + } for (int ich = tx; 16 * ich < K; ich += NX) { const device uint8_t* block = w_row + (int64_t)(ich / chpb) * Codec::block_bytes; @@ -813,7 +847,7 @@ METAL_FUNC float kq_ext_row_partial( } // GLU pair variant: gate and up rows share each activation chunk load. -template +template METAL_FUNC float2 kq_ext_glu_row_partial( const device uint8_t* gw, const device uint8_t* uw, @@ -828,6 +862,35 @@ METAL_FUNC float2 kq_ext_glu_row_partial( const device uint8_t* g_row = gw + row_off; const device uint8_t* u_row = uw + row_off; float2 acc = float2(0.0f); + if (C32 && KqExtC32::ok) { + constexpr short gpb = Codec::superblock / 32; + for (int ig = tx; 32 * ig < K; ig += NX) { + const int64_t boff = (int64_t)(ig / gpb) * Codec::block_bytes; + const short g8 = short(ig % gpb); + const device T* xp = x + ig * 32; + float4 a[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + a[i] = float4(*(const device vec*)(xp + 4 * i)); + } + float4x4 l0, l1; + float2 sc = KqExtC32::deq_group32s(g_row + boff, g8, l0, l1); + acc.x += sc.x * + (dot(l0[0], a[0]) + dot(l0[1], a[1]) + dot(l0[2], a[2]) + + dot(l0[3], a[3])) + + sc.y * + (dot(l1[0], a[4]) + dot(l1[1], a[5]) + dot(l1[2], a[6]) + + dot(l1[3], a[7])); + sc = KqExtC32::deq_group32s(u_row + boff, g8, l0, l1); + acc.y += sc.x * + (dot(l0[0], a[0]) + dot(l0[1], a[1]) + dot(l0[2], a[2]) + + dot(l0[3], a[3])) + + sc.y * + (dot(l1[0], a[4]) + dot(l1[1], a[5]) + dot(l1[2], a[6]) + + dot(l1[3], a[7])); + } + return acc; + } for (int ich = tx; 16 * ich < K; ich += NX) { const int64_t boff = (int64_t)(ich / chpb) * Codec::block_bytes; const short cch = short(ich % chpb); @@ -873,7 +936,12 @@ METAL_FUNC float kq_ext_reduce(float v) { return v; } -template +template < + typename T, + typename Codec, + int ACT, + int NX = KQ_EXT_NXPSG, + bool C32 = false> [[kernel]] void kq_ext_moe_glu_gather( const device uint8_t* gw [[buffer(0)]], const device uint8_t* uw [[buffer(1)]], @@ -899,7 +967,7 @@ template out += ((int64_t)tid.z * R + tid.y) * N; KQ_EXT_STAGE_LUTS(Codec, kq_luts) - const float2 gu = kq_ext_glu_row_partial( + const float2 gu = kq_ext_glu_row_partial( gw, uw, x, (int64_t)expert * N + out_row, K, tx, kq_luts); const float g = kq_ext_reduce(gu.x); const float u = kq_ext_reduce(gu.y); @@ -997,7 +1065,7 @@ template } } -template +template [[kernel]] void kq_ext_gather_qmv( const device uint8_t* w [[buffer(0)]], const device T* x [[buffer(1)]], @@ -1022,7 +1090,7 @@ template out += row_idx * N; KQ_EXT_STAGE_LUTS(Codec, kq_luts) - const float r = kq_ext_reduce(kq_ext_row_partial( + const float r = kq_ext_reduce(kq_ext_row_partial( w, x, (int64_t)expert * N + out_row, K, tx, kq_luts)); if (tx == 0) { out[out_row] = static_cast(r); diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index 2528f03..4d3c85d 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -1711,6 +1711,104 @@ template w, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); } +// Split-K M=1 mat-vec. The nsg simdgroups of a threadgroup share rps rows +// and stride the 256-weight steps by nsg, so the grid keeps N / rps +// threadgroups while each simdgroup walks only K / nsg of a row. Lane l +// owns 8 weights (block l / 4, quarter l % 4) of every 8-block step, the +// lane mapping of kq_q8_0_qmv_fast_impl. The quarter sits 2 bytes past a +// 2-byte boundary, so it loads as four ushorts. Partial sums reduce through +// `red` (nsg * rps floats). Needs N % rps == 0 and K % 256 == 0. +template +METAL_FUNC void kq_q8_0_qmv_splitk_impl( + const device uint8_t* w, + const device T* x, + device T* y, + const int in_vec_size, + threadgroup float* red, + uint3 tid, + uint simd_gid, + uint simd_lid) { + typedef float U; + constexpr int step = 8 * KQ_Q8_0_GROUP; + thread U result[rps] = {0}; + + const int row_bytes = in_vec_size * KQ_Q8_0_BLOCK_BYTES / KQ_Q8_0_GROUP; + const int out_row = tid.y * rps; + const int lane_k = simd_lid * 8; + const int quarter = simd_lid % 4; + const device uint8_t* w0 = w + static_cast(out_row) * row_bytes + + (simd_lid / 4) * KQ_Q8_0_BLOCK_BYTES; + + for (int k = int(simd_gid) * step; k < in_vec_size; k += nsg * step) { + const vec xa = *(const device vec*)(x + k + lane_k); + const vec xb = *(const device vec*)(x + k + lane_k + 4); + const U x0 = U(xa[0]), x1 = U(xa[1]), x2 = U(xa[2]), x3 = U(xa[3]); + const U x4 = U(xb[0]), x5 = U(xb[1]), x6 = U(xb[2]), x7 = U(xb[3]); + const int boff = (k / KQ_Q8_0_GROUP) * KQ_Q8_0_BLOCK_BYTES; +#pragma unroll + for (int row = 0; row < rps; row++) { + const device uint8_t* b = w0 + row * row_bytes + boff; + const U d = U(*(const device half*)b); + const device ushort* q = + (const device ushort*)(b + KQ_Q8_0_Q_OFFSET + quarter * 8); + const float2 q0 = float2(as_type(q[0])); + const float2 q1 = float2(as_type(q[1])); + const float2 q2 = float2(as_type(q[2])); + const float2 q3 = float2(as_type(q[3])); + const U p = q0.x * x0 + q0.y * x1 + q1.x * x2 + q1.y * x3 + q2.x * x4 + + q2.y * x5 + q3.x * x6 + q3.y * x7; + result[row] += d * p; + } + } + +#pragma unroll + for (int row = 0; row < rps; row++) { + const U s = simd_sum(result[row]); + if (simd_lid == 0) { + red[simd_gid * rps + row] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_gid == 0 && simd_lid < uint(rps)) { + U t = U(0); + for (int g = 0; g < nsg; g++) { + t += red[g * rps + simd_lid]; + } + y[out_row + simd_lid] = static_cast(t); + } +} + +#define KQ_Q8_0_DEFINE_QMV_SPLITK(name, RPS, NSG) \ + template \ + [[kernel, max_total_threads_per_threadgroup(32 * NSG)]] void name( \ + const device uint8_t* w, \ + const device uint8_t* /* scales */, \ + const device T* x, \ + device T* y, \ + const constant int& in_vec_size, \ + const constant int& /* out_vec_size */, \ + const constant int& /* x_batch_ndims */, \ + const constant int* /* x_shape */, \ + const constant int64_t* /* x_strides */, \ + const constant int& /* w_batch_ndims */, \ + const constant int* /* w_shape */, \ + const constant int64_t* /* w_strides */, \ + const constant int64_t* /* s_strides */, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + static_assert( \ + group_size == KQ_Q8_0_GROUP, "Q8_0 kernel requires group_size=32"); \ + static_assert(bits == 8, "Q8_0 kernel requires bits=8"); \ + threadgroup float red[NSG * RPS]; \ + kq_q8_0_qmv_splitk_impl( \ + w, x, y, in_vec_size, red, tid, simd_gid, simd_lid); \ + } + +// Split-K M=1 q8_0 mat-vec. The host copy of the tile is +// kq_q8_0_splitk_rps/nsg in src/kquant_metal_internal.h. +KQ_Q8_0_DEFINE_QMV_SPLITK(kq_q8_0_qmv_splitk, 2, 8) + // Bias-fused variant of kq_q8_0_qmv_fast: decode-only (B=1) fast path for a // KQuantLinear whose GGUF weight carries a real linear bias (e.g. gpt-oss // attention QKVO). Fuses the post-matmul add into this kernel's own output diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index 8a7ba2b..38b8fdc 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -5257,6 +5257,55 @@ struct KqTgLuts { } }; +// 32-weight decode step for the Ext gathers (the "_c32" kernels): one call +// decodes a whole 32-weight group, so the block scale, the qh byte and the +// scale byte load once per 32 weights and qs and signs load as ushort +// pairs. r0 holds weights 0..15 and r1 weights 16..31, sign-applied and +// unscaled; the return value holds their two scales. Codecs without a +// 32-weight decode keep ok = false and never reach these kernels. +template +struct KqExtC32 { + MLX_MTL_CONST bool ok = false; + static METAL_FUNC float2 deq_group32s( + const device uint8_t*, + short, + thread float4x4&, + thread float4x4&) { + return float2(0.0f); + } +}; + +template <> +struct KqExtC32 { + MLX_MTL_CONST bool ok = true; + static METAL_FUNC float2 deq_group32s( + const device uint8_t* block, + short ib32, + thread float4x4& r0, + thread float4x4& r1) { + const float d = float(*(const device half*)block); + const uint sc = block[KQ_IQ2_S_SCALES_OFFSET + ib32]; + const uint qhb = block[KQ_IQ2_S_QH_OFFSET + ib32]; + const device ushort* qp = + (const device ushort*)(block + KQ_IQ2_S_QS_OFFSET + ib32 * 4); + const device ushort* sp = + (const device ushort*)(block + KQ_IQ2_S_SIGNS_OFFSET + ib32 * 4); + const uint q01 = qp[0]; + const uint q23 = qp[1]; + const uint s01 = sp[0]; + const uint s23 = sp[1]; + kq_grid_u64_signed( + iq2s_grid[(q01 & 0xFF) | ((qhb << 8) & 0x300)], s01, r0[0], r0[1]); + kq_grid_u64_signed( + iq2s_grid[(q01 >> 8) | ((qhb << 6) & 0x300)], s01 >> 8, r0[2], r0[3]); + kq_grid_u64_signed( + iq2s_grid[(q23 & 0xFF) | ((qhb << 4) & 0x300)], s23, r1[0], r1[1]); + kq_grid_u64_signed( + iq2s_grid[(q23 >> 8) | ((qhb << 2) & 0x300)], s23 >> 8, r1[2], r1[3]); + return d * (0.5f + float2(float(sc & 0xf), float(sc >> 4))) * 0.25f; + } +}; + // --------------------------------------------------------------------------- // Half-dot decode tables (KQ_MOE_HALF): the grid codecs' 8-bit magnitude // entries pre-expanded to half4 pairs in threadgroup memory, so the chunk diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 4007153..1e866a9 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -1816,11 +1816,13 @@ void qvm( ce.dispatch_threadgroups(grid_dims, group_dims); } -// Dispatches the split-K M=1 q6_k mat-vec, kq_q6_k_qmv_splitk. The -// simdgroups of a threadgroup split each row along K, which keeps the kernel -// on memory bandwidth at the lower GPU clocks of sustained decode, where the -// per-row kernels turn ALU-bound. Returns false, dispatching nothing, when -// the shape or pipeline does not fit. +// Dispatches the split-K M=1 mat-vec of q6_k (kq_q6_k_qmv_splitk) or q8_0 +// (kq_q8_0_qmv_splitk). The simdgroups of a threadgroup split each row along +// K. For q6_k this keeps the kernel on memory bandwidth at the lower GPU +// clocks of sustained decode, where the per-row kernels turn ALU-bound. For +// q8_0 it puts more simdgroups in flight on layer-size N, where the per-row +// grid leaves the GPU short of loads in flight. Returns false, dispatching +// nothing, when the shape or pipeline does not fit. static bool qmv_splitk( const array& x, const array& w, @@ -1833,9 +1835,13 @@ static bool qmv_splitk( Device& d, const Stream& s, const std::string& kquant_type) { - constexpr int rps = kq_q6_k_splitk_rps; - constexpr int nsg = kq_q6_k_splitk_nsg; - if (kquant_type != "q6_k" || N % rps != 0 || K % qmv_fast_k_align() != 0) { + const bool q6 = kquant_type == "q6_k"; + if (!q6 && kquant_type != "q8_0") { + return false; + } + const int rps = q6 ? kq_q6_k_splitk_rps : kq_q8_0_splitk_rps; + const int nsg = q6 ? kq_q6_k_splitk_nsg : kq_q8_0_splitk_nsg; + if (N % rps != 0 || K % qmv_fast_k_align() != 0) { return false; } std::string kname; @@ -1849,12 +1855,17 @@ static bool qmv_splitk( "_b_", bits, "_batch_0"); - // K is function constant 340, so the pipeline is keyed by row width too. - int k_const = K; - mx::metal::MTLFCList func_consts = { - {&k_const, MTL::DataType::DataTypeInt, 340}}; - auto kernel = - kq_get_kernel(d, kname, kname + "_k" + std::to_string(K), func_consts); + MTL::ComputePipelineState* kernel = nullptr; + if (q6) { + // K is function constant 340, so the pipeline is keyed by row width too. + int k_const = K; + mx::metal::MTLFCList func_consts = { + {&k_const, MTL::DataType::DataTypeInt, 340}}; + kernel = + kq_get_kernel(d, kname, kname + "_k" + std::to_string(K), func_consts); + } else { + kernel = kq_get_kernel(d, kname); + } if (int(kernel->maxTotalThreadsPerThreadgroup()) < 32 * nsg) { return false; } diff --git a/src/kquant_metal_internal.h b/src/kquant_metal_internal.h index 26c031c..e516875 100644 --- a/src/kquant_metal_internal.h +++ b/src/kquant_metal_internal.h @@ -81,6 +81,11 @@ inline int kquant_qmv_fast_bn(const std::string& kquant_type) { constexpr int kq_q6_k_splitk_rps = 4; constexpr int kq_q6_k_splitk_nsg = 8; +// Output rows and simdgroups per threadgroup of the split-K M=1 q8_0 qmv, +// kq_q8_0_qmv_splitk in kq_quantized.h. +constexpr int kq_q8_0_splitk_rps = 2; +constexpr int kq_q8_0_splitk_nsg = 8; + // quantized.cpp:63-65 (kquant branch). KQuant blocks are 32 or 256 weights. inline int qmv_fast_k_align() { return 256; diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index bc5e3dd..3a5c779 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -58,9 +58,10 @@ inline std::string kq_gather_stem(const std::string& t, int K) { // geometry (E=512 top-10, K=2560 N=640, iq3_xxs/iq4_xs + q8_0 shexp) nx16 // beats nx8 by 3.6-9.8% at every t in 2..8, still positive at the 7040-tg // t=8 point, so the cap sits at 8192 tgs; true prefill grids stay NX = 8. -inline int kq_moe_pick_nx(int64_t rows, int K, bool wide_ok) { - // Re-read per call only when the variable exists at all (interleaved A/B - // flips it in-process); unset costs one static check. +// The width KQ_MOE_NX forces, or 0. Re-read per call only when the variable +// exists at all (interleaved A/B flips it in-process); unset costs one +// static check. +inline int kq_moe_nx_env() { static const bool has_env = std::getenv("KQ_MOE_NX") != nullptr; if (has_env) { const char* e = std::getenv("KQ_MOE_NX"); @@ -69,6 +70,13 @@ inline int kq_moe_pick_nx(int64_t rows, int K, bool wide_ok) { return v; } } + return 0; +} + +inline int kq_moe_pick_nx(int64_t rows, int K, bool wide_ok) { + if (const int v = kq_moe_nx_env()) { + return v; + } if (wide_ok && K / 16 >= 32 && (rows * 8) / 64 < 8192) { return 16; } @@ -337,6 +345,23 @@ inline int kq_moe_pick_sg_h(int N, int nx) { return sg; } +// 32-weight decode step ("_c32" kernels, KqExtC32) for the iq2_s uniform +// gathers: by default at t = 1 only. KQ_MOE_C32=0 forces the 16-weight +// kernels and =1 the c32 kernels at every t. Read per call. The c32 kernels +// have NX 8 and 16 forms, so a forced KQ_MOE_NX=32 keeps the 16-weight +// kernels. +inline bool kq_moe_c32(const std::string& codec, int T) { + if (codec != "iq2_s" || kq_moe_nx_env() == 32) { + return false; + } + const int v = kq_env_flag("KQ_MOE_C32", -1); + return v == -1 ? T == 1 : v != 0; +} + +inline int kq_moe_c32_nx() { + return kq_moe_nx_env() == 16 ? 16 : 8; +} + inline bool kq_gather_fine(int64_t coarse_tgs) { static const bool has_env = std::getenv("KQ_GATHER_FINE") != nullptr; const int v = has_env ? kq_env_flag("KQ_GATHER_FINE", -1) : -1; @@ -543,14 +568,18 @@ void KQuantMoEGLUKQ::eval_gpu( // staged activation row and half grid (nx 16 measured 10% slower). const bool use_half = !biased && kq_moe_half(kquant_type_, K); const bool dd = !biased && !use_half && kq_moe_dedup(T, KQ_DD_MAX_T); - const int nx = use_half ? 8 : kq_moe_pick_nx((int64_t)N * R * T, K, true); + const bool c32 = !biased && !use_half && !dd && kq_moe_c32(kquant_type_, T); + const int nx = use_half ? 8 + : c32 ? kq_moe_c32_nx() + : kq_moe_pick_nx((int64_t)N * R * T, K, true); std::string stem = kq_gather_stem_nx(kquant_type_, K, nx); if (dd && (stem == "q6_k" || stem == "q8_0")) { stem += "_ext"; // the dedupe kernels are Ext-only } std::string kname = "kq_" + stem + "_moe_glu_gather_" + (biased ? "bias_" : "") + (use_half ? "h_" : "") + (dd ? "dd_" : "") + - act_ + kq_nx_suffix(nx) + "_" + kq_type_string(x.dtype()); + act_ + (c32 ? "_c32" : "") + kq_nx_suffix(nx) + "_" + + kq_type_string(x.dtype()); auto kernel = kq_get_kernel(d, kname); kq_moe_log_kname(kname, int(kernel->maxTotalThreadsPerThreadgroup())); auto& ce = mx::metal::get_command_encoder(s); @@ -614,7 +643,10 @@ void KQuantGatherQMVKQ::eval_gpu( int K = x.shape(-1); const bool use_half = !biased && kq_moe_half(kquant_type_, K); - const int nx = use_half ? 8 : kq_moe_pick_nx((int64_t)N * R * T, K, false); + const bool c32 = !biased && !use_half && kq_moe_c32(kquant_type_, T); + const int nx = use_half ? 8 + : c32 ? kq_moe_c32_nx() + : kq_moe_pick_nx((int64_t)N * R * T, K, false); const std::string stem = kq_gather_stem_nx(kquant_type_, K, nx); // Fine tiling is instantiated only on the tuned q6_k/q8_0 kernels; the Ext // equivalents (2x and 4x threadgroup variants) measured E2E-neutral and @@ -623,8 +655,8 @@ void KQuantGatherQMVKQ::eval_gpu( const bool fine = tuned && !biased && nx == 8 && kq_gather_fine((int64_t)(N / 8) * R * T); std::string kname = "kq_" + stem + "_gather_qmv" + (biased ? "_bias" : "") + - (use_half ? "_h" : "") + (fine ? "_fine" : kq_nx_suffix(nx)) + "_" + - kq_type_string(x.dtype()); + (use_half ? "_h" : "") + (c32 ? "_c32" : "") + + (fine ? "_fine" : kq_nx_suffix(nx)) + "_" + kq_type_string(x.dtype()); auto kernel = kq_get_kernel(d, kname); kq_moe_log_kname(kname, int(kernel->maxTotalThreadsPerThreadgroup())); auto& ce = mx::metal::get_command_encoder(s); diff --git a/tests/test_matmul_synth.py b/tests/test_matmul_synth.py index ab3e578..fb922b2 100644 --- a/tests/test_matmul_synth.py +++ b/tests/test_matmul_synth.py @@ -135,14 +135,16 @@ def test_matmul_synth_tail(codec): ) -# q6_k at M 1 runs the split-K mat-vec when 4 divides N, and the per-row -# kernels otherwise or under KQ_QMV_SPLITK=0. K 256 leaves seven of the -# eight simdgroups without a superblock, and K 5376 splits 21 superblocks -# unevenly. Both arms meet the oracle in every dtype and agree bit for bit -# where N rules the split out. The op promotes float32 x to bfloat16. The -# tolerance scales with the output because the CPU path quantizes x to 8 -# bits, which costs it about 1% of the largest output at these K. -SPLITK_SHAPES = ((256, 256), (260, 5376), (1024, 4096), (258, 512)) +# q6_k and q8_0 at M 1 run the split-K mat-vec when the tile's row count +# divides N (4 for q6_k, 2 for q8_0), and the per-row kernels otherwise or +# under KQ_QMV_SPLITK=0. K 256 leaves most simdgroups without a step, K 5376 +# splits 21 steps unevenly, and N 255 rules the split out. Both arms meet the +# oracle in every dtype and agree bit for bit where N rules the split out. +# The op promotes float32 x to bfloat16. The tolerance scales with the +# output because the CPU path quantizes x to 8 bits, which costs it about 1% +# of the largest output at these K. +SPLITK_SHAPES = ((256, 256), (260, 5376), (1024, 4096), (258, 512), (255, 512)) +SPLITK_RPS = {"q6_k": 4, "q8_0": 2} # Levers that move an M=1 call off the qmv route or change its tiling. # KQ_VERIFY_EXT and KQ_QMM_SPLITK are read once per process, so a set one # skips the tests. @@ -171,30 +173,31 @@ def _clear_qmv_levers(monkeypatch): @pytest.mark.parametrize("dtype", [mx.float32, mx.float16, mx.bfloat16]) @pytest.mark.parametrize("n, k", SPLITK_SHAPES) -def test_q6_k_qmv_splitk(n, k, dtype, monkeypatch): +@pytest.mark.parametrize("codec", list(SPLITK_RPS)) +def test_qmv_splitk(codec, n, k, dtype, monkeypatch): _clear_qmv_levers(monkeypatch) rng = np.random.default_rng(1) w_np = (rng.standard_normal((n, k)) * 0.1).astype(np.float32) - wq, _ = kq.quantize(mx.array(w_np), "q6_k") + wq, _ = kq.quantize(mx.array(w_np), codec) mx.eval(wq) packed = np.ascontiguousarray(np.array(wq).astype(np.uint8)) w = mx.array(packed) scales = mx.zeros((1,), dtype=mx.uint8) - deq = quants.dequantize(packed, GT.Q6_K).astype(np.float32) + deq = quants.dequantize(packed, CODECS[codec]).astype(np.float32) x_np = (rng.standard_normal((1, k)) * 0.1).astype(np.float32) x = mx.array(x_np).astype(dtype) r = np.array(x.astype(mx.float32)) @ deq.T outs = {} for lever in ("1", "0"): monkeypatch.setenv("KQ_QMV_SPLITK", lever) - got = kq.quantized_matmul(x, w, scales, "q6_k", transpose=True) + got = kq.quantized_matmul(x, w, scales, codec, transpose=True) mx.eval(got) outs[lever] = got err = float(np.abs(np.array(got.astype(mx.float32)) - r).max()) assert err < 2e-2 * float(np.abs(r).max()), ( - f"N={n} K={k} KQ_QMV_SPLITK={lever}: max_abs={err:.3e}" + f"{codec} N={n} K={k} KQ_QMV_SPLITK={lever}: max_abs={err:.3e}" ) - if n % 4: + if n % SPLITK_RPS[codec]: assert bool(mx.array_equal(outs["1"], outs["0"]).item()) # The two kernels differ only in float32 summation order, so their # outputs agree to one unit in the last place. @@ -254,3 +257,48 @@ def run(n, lever=None, fine=None): assert run(4, fine="1") == [256.0] * 4 assert run(4, fine="0") == [256.0] * 4 assert run(6) == [256.0] * 6 + + +def _q8_0_block(q, d): + d16 = np.array([d], np.float16).view(np.uint8) + return np.concatenate([d16, np.asarray(q, np.int8).view(np.uint8)]) + + +# The reduction order shows which q8_0 kernel ran. With x all ones, lane 0 +# of the first 256-weight step contributes +B (B = 32768 * 1016, exact in +# float32), lane 1 of the second step -B, and lane 0 of the third step 0.5. +# The per-row kernels add 0.5 to +B inside lane 0, where it is below half a +# unit of B, and return 0. The split-K kernel gives each step its own +# simdgroup and returns 0.5. +@pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")) or mx.default_device() != mx.gpu, + reason="Metal kernel routing", +) +@pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16]) +def test_q8_0_qmv_splitk_routing(dtype, monkeypatch): + _clear_qmv_levers(monkeypatch) + zero = _q8_0_block(np.zeros(32), 0.0) + blocks = [zero] * 24 + blocks[0] = _q8_0_block([127] * 8 + [0] * 24, 32768.0) + blocks[8] = _q8_0_block([0] * 8 + [-127] * 8 + [0] * 16, 32768.0) + blocks[16] = _q8_0_block([1] + [0] * 31, 0.5) + row = np.concatenate(blocks) + scales = mx.zeros((1,), dtype=mx.uint8) + x = mx.ones((1, 768), dtype=dtype) + + def run(n, lever=None, fine=None): + for name, v in (("KQ_QMV_SPLITK", lever), ("KQ_QMV_FINE", fine)): + if v is None: + monkeypatch.delenv(name, raising=False) + else: + monkeypatch.setenv(name, v) + w = mx.array(np.tile(row, (n, 1))) + y = kq.quantized_matmul(x, w, scales, "q8_0", transpose=True) + return np.array(y.astype(mx.float32)).ravel().tolist() + + assert run(8) == [0.5] * 8 + assert run(8, lever="1") == [0.5] * 8 + assert run(8, lever="0") == [0.0] * 8 + assert run(8, fine="1") == [0.0] * 8 + assert run(8, fine="0") == [0.0] * 8 + assert run(3) == [0.0] * 3 diff --git a/tests/test_moe_glu_variants.py b/tests/test_moe_glu_variants.py index 039be01..ad7f17e 100644 --- a/tests/test_moe_glu_variants.py +++ b/tests/test_moe_glu_variants.py @@ -514,3 +514,137 @@ def test_shexp_mix_slot_parallel_matches_dispatch(): stem = codec if scodec == codec else f"{codec}_sx_{scodec}" name = f"kq_{stem}_gather_qmv_mix_sp_" assert (name in r.stderr) == (listed == "1"), (codec, scodec, listed) + + +_C32_SNIPPET = r""" +import sys +import numpy as np +import mlx.core as mx +import mlx_kquant as kq + +T, dtype, out_path = int(sys.argv[1]), getattr(mx, sys.argv[2]), sys.argv[3] +rng = np.random.default_rng(31) +E, N, K, S = 8, 256, 1024, 4 +def wire(e, n, k): + w = mx.array((rng.standard_normal((e, n, k)) * 0.1).astype(np.float32)) + wq, _ = kq.quantize(w, "iq2_s", mx.ones((k,), dtype=mx.float32)) + mx.eval(wq) + return wq +gw, uw, dw = wire(E, N, K), wire(E, N, K), wire(E, K, N) +x = mx.array((rng.standard_normal((T, K)) * 0.05).astype(np.float32)).astype(dtype) +h = mx.array((rng.standard_normal((T, S, N)) * 0.05).astype(np.float32)).astype(dtype) +inds = mx.array(rng.integers(0, E, size=(T, S)).astype(np.uint32)) +glu = kq.moe_glu_gather_kq(x, gw, uw, "iq2_s", inds, act="silu") +qmv = kq.gather_qmv_kq(h, dw, "iq2_s", inds) +mx.eval(glu, qmv) +np.savez(out_path, glu=np.array(glu.astype(mx.float32)), + qmv=np.array(qmv.astype(mx.float32))) +""" + + +def _c32_run(T, dtype, env, tmp_path, arm): + f = tmp_path / f"{arm}.npz" + r = _run_child( + [sys.executable, "-c", _C32_SNIPPET, str(T), dtype, str(f)], + {**env, "KQ_MOE_NX_LOG": "1"}, + f"c32 {arm} T={T} {dtype}", + ) + return np.load(f), r.stderr + + +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +def test_c32_dispatch(dtype, tmp_path): + """iq2_s gate/up and down gathers take the 32-weight-step kernels at one + row by default, at NX = 8 or a forced 16. KQ_MOE_C32=0 and a forced + KQ_MOE_NX=32 keep the 16-weight kernels. Two rows keep them too unless + KQ_MOE_C32=1. The forms differ only in summation order.""" + want_glu = f"kq_iq2_s_moe_glu_gather_silu_c32_{dtype}_t" + want_qmv = f"kq_iq2_s_gather_qmv_c32_{dtype}_t" + new, names = _c32_run(1, dtype, {}, tmp_path, "default") + assert want_glu in names and want_qmv in names, names + old, names = _c32_run(1, dtype, {"KQ_MOE_C32": "0"}, tmp_path, "old") + assert "_c32" not in names, names + wide, names = _c32_run(1, dtype, {"KQ_MOE_NX": "16"}, tmp_path, "nx16") + assert "kq_iq2_s_moe_glu_gather_silu_c32_nx16_" in names, names + assert "kq_iq2_s_gather_qmv_c32_nx16_" in names, names + for key in ("glu", "qmv"): + scale = np.abs(old[key]).max() + for got in (new, wide): + rel = np.abs(got[key] - old[key]).max() / scale + assert rel < 2e-3, (key, rel) + _, names = _c32_run(1, dtype, {"KQ_MOE_NX": "32"}, tmp_path, "nx32") + assert "_c32" not in names, names + assert "kq_iq2_s_moe_glu_gather_silu_nx32_" in names, names + assert "kq_iq2_s_gather_qmv_nx32_" in names, names + _, names = _c32_run(2, dtype, {}, tmp_path, "t2") + assert "_c32" not in names, names + _, names = _c32_run( + 2, dtype, {"KQ_MOE_C32": "1", "KQ_MOE_DEDUP": "0"}, tmp_path, "t2c32" + ) + assert want_glu in names and want_qmv in names, names + + +_C32_LIVE_SNIPPET = r""" +import os +import numpy as np +import mlx.core as mx +import mlx_kquant as kq + +rng = np.random.default_rng(5) +w = mx.array((rng.standard_normal((4, 256, 256)) * 0.1).astype(np.float32)) +wq, _ = kq.quantize(w, "iq2_s", mx.ones((256,), dtype=mx.float32)) +x = mx.ones((1, 1, 256), dtype=mx.bfloat16) +inds = mx.array(np.array([[1]], np.uint32)) +os.environ.pop("KQ_MOE_C32", None) +mx.eval(kq.gather_qmv_kq(x, wq, "iq2_s", inds)) +os.environ["KQ_MOE_C32"] = "0" +mx.eval(kq.gather_qmv_kq(x, wq, "iq2_s", inds)) +""" + + +def test_c32_reads_live(): + """KQ_MOE_C32 takes effect on the next call, also when the process + started without it.""" + r = _run_child( + [sys.executable, "-c", _C32_LIVE_SNIPPET], + {"KQ_MOE_NX_LOG": "1"}, + "c32 live read", + ) + assert "kq_iq2_s_gather_qmv_bfloat16_t" in r.stderr, r.stderr + assert "kq_iq2_s_gather_qmv_c32_bfloat16_t" in r.stderr, r.stderr + + +_C32_REF_SNIPPET = r""" +import json +import sys +sys.path.insert(0, sys.argv[1]) +import mlx.core as mx +import test_moe_glu as tm + +res = tm._check_codec("iq2_s", dtype=getattr(mx, sys.argv[2])) +print("RESULT " + json.dumps( + None if res is None else [(n, float(r), bool(o)) for n, r, o in res])) +""" + + +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +def test_c32_accuracy(dtype): + """The 32-weight-step kernels, forced at the reference test's three rows, + stay inside the gather bound of the f32 dequantized reference.""" + tests_dir = os.path.dirname(os.path.abspath(__file__)) + r = _run_child( + [sys.executable, "-c", _C32_REF_SNIPPET, tests_dir, dtype], + {"KQ_MOE_C32": "1", "KQ_MOE_DEDUP": "0", "KQ_MOE_NX_LOG": "1"}, + f"c32 reference {dtype}", + ) + assert "kq_iq2_s_moe_glu_gather_silu_c32_" in r.stderr, r.stderr + assert "kq_iq2_s_gather_qmv_c32_" in r.stderr, r.stderr + line = [ln for ln in r.stdout.splitlines() if ln.startswith("RESULT ")] + assert line, r.stdout + r.stderr + import json + + res = json.loads(line[-1][len("RESULT ") :]) + if res is None: + pytest.skip("iq2_s reference unavailable") + for name, rel, ok in res: + assert ok, (name, rel)