From 43b23b0e2084ce3233fde7cbae8d107c9739fe16 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:52:52 -0700 Subject: [PATCH 1/3] perf(moe): c32 gathers for the iq1, iq2 and iq3 grid codecs --- metal/kq_moe_glu_kq.metal | 6 + .../mlx/backend/metal/kernels/kq_moe_glu_kq.h | 7 +- .../backend/metal/kernels/kq_quantized_iq.h | 259 +++++++++++++++++- src/kquant_moe_glu.cpp | 19 +- 4 files changed, 278 insertions(+), 13 deletions(-) diff --git a/metal/kq_moe_glu_kq.metal b/metal/kq_moe_glu_kq.metal index f74b23e..1f2a538 100644 --- a/metal/kq_moe_glu_kq.metal +++ b/metal/kq_moe_glu_kq.metal @@ -336,6 +336,12 @@ instantiate_kq_ext_sx_mix(q2_0, KqQ2_0Ext, iq4_nl, KqIq4_nlExt, float16_t) instantiate_kq_ext_c32_nx(codec, traits, float16_t, 16, "_nx16") instantiate_kq_ext_c32(iq2_s, KqIq2_sExt) +instantiate_kq_ext_c32(iq2_xxs, KqIq2_xxsExt) +instantiate_kq_ext_c32(iq2_xs, KqIq2_xsExt) +instantiate_kq_ext_c32(iq3_xxs, KqIq3_xxsExt) +instantiate_kq_ext_c32(iq3_s, KqIq3_sExt) +instantiate_kq_ext_c32(iq1_s, KqIq1_sExt) +instantiate_kq_ext_c32(iq1_m, KqIq1_mExt) // q6_k/q8_0 mix_ns (generic; no tuned ns kernels exist). instantiate_kq_ext_mix_ns(q6_k, KqQ6_KExt) 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 4c63950..d741634 100644 --- a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h +++ b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h @@ -825,7 +825,8 @@ METAL_FUNC float kq_ext_row_partial( w_row + (int64_t)(ig / gpb) * Codec::block_bytes, short(ig % gpb), l0, - l1); + l1, + luts); acc += kq_ext_dot32(l0, l1, sc, x + ig * 32); } return acc; @@ -874,14 +875,14 @@ METAL_FUNC float2 kq_ext_glu_row_partial( a[i] = float4(*(const device vec*)(xp + 4 * i)); } float4x4 l0, l1; - float2 sc = KqExtC32::deq_group32s(g_row + boff, g8, l0, l1); + float2 sc = KqExtC32::deq_group32s(g_row + boff, g8, l0, l1, luts); 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); + sc = KqExtC32::deq_group32s(u_row + boff, g8, l0, l1, luts); acc.y += sc.x * (dot(l0[0], a[0]) + dot(l0[1], a[1]) + dot(l0[2], a[2]) + dot(l0[3], a[3])) + diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index 38b8fdc..6c475ff 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -5261,8 +5261,40 @@ struct KqTgLuts { // 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. +// unscaled; the return value holds their two scales. `luts` is the codec's +// KqTgLuts block. Codecs without a 32-weight decode keep ok = false and +// never reach these kernels. +// Two u32 grid entries (8 weights) from a ushort of two 8-bit indices. +METAL_FUNC void kq_c32_u32_pair( + const threadgroup uint32_t* grid, + uint q, + uint s, + thread float4& lo, + thread float4& hi) { + lo = kq_grid_sign4(kq_grid_u8x4(grid[q & 0xFF]), s); + hi = kq_grid_sign4(kq_grid_u8x4(grid[(q >> 8) & 0xFF]), s >> 4); +} +// Two u32 grid entries (8 weights) at indices i0 and i1. +METAL_FUNC void kq_c32_u32_idx( + const threadgroup uint32_t* grid, + uint i0, + uint i1, + uint s, + thread float4& lo, + thread float4& hi) { + lo = kq_grid_sign4(kq_grid_u8x4(grid[i0]), s); + hi = kq_grid_sign4(kq_grid_u8x4(grid[i1]), s >> 4); +} +// One signed u64 grid entry (8 weights) plus delta, unscaled. +METAL_FUNC void kq_c32_i8_delta( + uint64_t ge, + float delta, + thread float4& lo, + thread float4& hi) { + lo = kq_grid_i8x4(uint32_t(ge)) + delta; + hi = kq_grid_i8x4(uint32_t(ge >> 32)) + delta; +} + template struct KqExtC32 { MLX_MTL_CONST bool ok = false; @@ -5270,7 +5302,8 @@ struct KqExtC32 { const device uint8_t*, short, thread float4x4&, - thread float4x4&) { + thread float4x4&, + const threadgroup uint8_t*) { return float2(0.0f); } }; @@ -5282,7 +5315,8 @@ struct KqExtC32 { const device uint8_t* block, short ib32, thread float4x4& r0, - thread float4x4& r1) { + thread float4x4& r1, + const threadgroup uint8_t*) { 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]; @@ -5306,6 +5340,223 @@ struct KqExtC32 { } }; +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 threadgroup uint8_t* luts) { + const threadgroup uint64_t* grid = + reinterpret_cast(luts); + const threadgroup uint8_t* ksigns = luts + 2048; + const float d = float(*(const device half*)block); + const device ushort* qp = + (const device ushort*)(block + KQ_IQ2_XXS_QS_OFFSET + ib32 * 8); + const uint q01 = qp[0]; + const uint q23 = qp[1]; + const uint sb = uint(qp[2]) | (uint(qp[3]) << 16); + kq_grid_u64_signed(grid[q01 & 0xFF], ksigns[sb & 127], r0[0], r0[1]); + kq_grid_u64_signed(grid[q01 >> 8], ksigns[(sb >> 7) & 127], r0[2], r0[3]); + kq_grid_u64_signed( + grid[q23 & 0xFF], ksigns[(sb >> 14) & 127], r1[0], r1[1]); + kq_grid_u64_signed(grid[q23 >> 8], ksigns[(sb >> 21) & 127], r1[2], r1[3]); + return float2(d * (0.5f + float(sb >> 28)) * 0.25f); + } +}; + +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 threadgroup uint8_t* luts) { + const threadgroup uint64_t* grid = + reinterpret_cast(luts); + const threadgroup uint8_t* ksigns = luts + 4096; + const float d = float(*(const device half*)block); + const uint sc = block[KQ_IQ2_XS_SCALES_OFFSET + ib32]; + const device ushort* qp = + (const device ushort*)(block + KQ_IQ2_XS_QS_OFFSET + ib32 * 8); + const uint q0 = qp[0]; + const uint q1 = qp[1]; + const uint q2 = qp[2]; + const uint q3 = qp[3]; + kq_grid_u64_signed(grid[q0 & 511], ksigns[q0 >> 9], r0[0], r0[1]); + kq_grid_u64_signed(grid[q1 & 511], ksigns[q1 >> 9], r0[2], r0[3]); + kq_grid_u64_signed(grid[q2 & 511], ksigns[q2 >> 9], r1[0], r1[1]); + kq_grid_u64_signed(grid[q3 & 511], ksigns[q3 >> 9], r1[2], r1[3]); + return d * (0.5f + float2(float(sc & 0xf), float(sc >> 4))) * 0.25f; + } +}; + +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 threadgroup uint8_t* luts) { + const threadgroup uint32_t* grid = + reinterpret_cast(luts); + const threadgroup uint8_t* ksigns = luts + 1024; + const float d = float(*(const device half*)block); + const device ushort* qp = + (const device ushort*)(block + KQ_IQ3_XXS_QS_OFFSET + ib32 * 8); + const device ushort* gas = + (const device ushort*)(block + KQ_IQ3_XXS_GAS_OFFSET + ib32 * 4); + const uint aux = uint(gas[0]) | (uint(gas[1]) << 16); + kq_c32_u32_pair(grid, qp[0], ksigns[aux & 127], r0[0], r0[1]); + kq_c32_u32_pair(grid, qp[1], ksigns[(aux >> 7) & 127], r0[2], r0[3]); + kq_c32_u32_pair(grid, qp[2], ksigns[(aux >> 14) & 127], r1[0], r1[1]); + kq_c32_u32_pair(grid, qp[3], ksigns[(aux >> 21) & 127], r1[2], r1[3]); + return float2(d * (0.5f + float(aux >> 28)) * 0.5f); + } +}; + +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 threadgroup uint8_t* luts) { + const threadgroup uint32_t* grid = + reinterpret_cast(luts); + const float d = float(*(const device half*)block); + const device ushort* qp = + (const device ushort*)(block + KQ_IQ3_S_QS_OFFSET + ib32 * 8); + const device ushort* sp = + (const device ushort*)(block + KQ_IQ3_S_SIGNS_OFFSET + ib32 * 4); + const uint qhb = block[KQ_IQ3_S_QH_OFFSET + ib32]; + const uint sc = + (block[KQ_IQ3_S_SCALES_OFFSET + ib32 / 2] >> (4 * (ib32 & 1))) & 0xf; + const uint s01 = sp[0]; + const uint s23 = sp[1]; + const uint q01 = qp[0]; + const uint q23 = qp[1]; + const uint q45 = qp[2]; + const uint q67 = qp[3]; + // bit i of qh is bit 8 of grid index i + kq_c32_u32_idx( + grid, + (q01 & 0xFF) | ((qhb << 8) & 256), + (q01 >> 8) | ((qhb << 7) & 256), + s01, + r0[0], + r0[1]); + kq_c32_u32_idx( + grid, + (q23 & 0xFF) | ((qhb << 6) & 256), + (q23 >> 8) | ((qhb << 5) & 256), + s01 >> 8, + r0[2], + r0[3]); + kq_c32_u32_idx( + grid, + (q45 & 0xFF) | ((qhb << 4) & 256), + (q45 >> 8) | ((qhb << 3) & 256), + s23, + r1[0], + r1[1]); + kq_c32_u32_idx( + grid, + (q67 & 0xFF) | ((qhb << 2) & 256), + (q67 >> 8) | ((qhb << 1) & 256), + s23 >> 8, + r1[2], + r1[3]); + return float2(d * float(1 + 2 * sc)); + } +}; + +// iq1_s and iq1_m return grid + delta, unscaled. +template <> +struct KqExtC32 { + MLX_MTL_CONST bool ok = true; + static METAL_FUNC float2 deq_group32s( + const device uint8_t* block, + short ib, + thread float4x4& r0, + thread float4x4& r1, + const threadgroup uint8_t*) { + const float d = float(*(const device half*)block); + const uint qh = + *(const device ushort*)(block + KQ_IQ1_S_QH_OFFSET + ib * 2); + const device ushort* qp = + (const device ushort*)(block + KQ_IQ1_S_QS_OFFSET + ib * 4); + const uint q01 = qp[0]; + const uint q23 = qp[1]; + const float delta = (qh & 0x8000) ? -0.125f : 0.125f; + kq_c32_i8_delta( + iq1s_grid[(q01 & 0xFF) | ((qh << 8) & 0x700)], delta, r0[0], r0[1]); + kq_c32_i8_delta( + iq1s_grid[(q01 >> 8) | ((qh << 5) & 0x700)], delta, r0[2], r0[3]); + kq_c32_i8_delta( + iq1s_grid[(q23 & 0xFF) | ((qh << 2) & 0x700)], delta, r1[0], r1[1]); + kq_c32_i8_delta( + iq1s_grid[(q23 >> 8) | ((qh >> 1) & 0x700)], delta, r1[2], r1[3]); + return float2(d * float(2 * int((qh >> 12) & 7) + 1)); + } +}; + +template <> +struct KqExtC32 { + MLX_MTL_CONST bool ok = true; + static METAL_FUNC float2 deq_group32s( + const device uint8_t* block, + short ib, + thread float4x4& r0, + thread float4x4& r1, + const threadgroup uint8_t*) { + const device ushort* scw = + (const device ushort*)(block + KQ_IQ1_M_SCALES_OFFSET); + const ushort su = (scw[0] >> 12) | ((scw[1] >> 8) & 0x00f0) | + ((scw[2] >> 4) & 0x0f00) | (scw[3] & 0xf000); + const float d = float(as_type(su)); + const uint w = scw[ib / 2]; + const int sh = 6 * (ib & 1); + const device ushort* qp = + (const device ushort*)(block + KQ_IQ1_M_QS_OFFSET + ib * 4); + const uint q01 = qp[0]; + const uint q23 = qp[1]; + const uint qh = + *(const device ushort*)(block + KQ_IQ1_M_QH_OFFSET + ib * 2); + kq_c32_i8_delta( + iq1s_grid[(q01 & 0xFF) | ((qh << 8) & 0x700)], + (qh & 0x08) ? -0.125f : 0.125f, + r0[0], + r0[1]); + kq_c32_i8_delta( + iq1s_grid[(q01 >> 8) | ((qh << 4) & 0x700)], + (qh & 0x80) ? -0.125f : 0.125f, + r0[2], + r0[3]); + kq_c32_i8_delta( + iq1s_grid[(q23 & 0xFF) | (qh & 0x700)], + (qh & 0x0800) ? -0.125f : 0.125f, + r1[0], + r1[1]); + kq_c32_i8_delta( + iq1s_grid[(q23 >> 8) | ((qh >> 4) & 0x700)], + (qh & 0x8000) ? -0.125f : 0.125f, + r1[2], + r1[3]); + return d * + float2(float(2 * int((w >> sh) & 7) + 1), + float(2 * int((w >> (sh + 3)) & 7) + 1)); + } +}; + // --------------------------------------------------------------------------- // 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_moe_glu.cpp b/src/kquant_moe_glu.cpp index 253ed72..71dbc01 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -346,13 +346,20 @@ 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. +// Codecs with a 32-weight decode step (KqExtC32): the grid codecs. +inline bool kq_moe_has_c32(const std::string& codec) { + return codec == "iq2_s" || codec == "iq2_xxs" || codec == "iq2_xs" || + codec == "iq3_xxs" || codec == "iq3_s" || codec == "iq1_s" || + codec == "iq1_m"; +} + +// 32-weight decode step ("_c32" kernels, KqExtC32) for the grid-codec +// 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) { + if (!kq_moe_has_c32(codec) || kq_moe_nx_env() == 32) { return false; } const int v = kq_env_flag("KQ_MOE_C32", -1); From dbb833d6175af284692f0236d2ccf334624fa263 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:15:07 -0700 Subject: [PATCH 2/3] perf(moe): c32 for every grid-codec gather family and the one-row grid-codec mat-vec on m5 --- CHANGELOG.md | 6 + README.md | 16 +- bindings.cpp | 18 +- docs/kernels.md | 54 ++++- metal/kq_moe_glu_kq.metal | 50 ++++- .../mlx/backend/metal/kernels/kq_moe_glu_kq.h | 63 ++++-- src/kquant.h | 14 +- src/kquant_matmul.cpp | 51 +++++ src/kquant_moe_glu.cpp | 135 +++++++++--- tests/test_moe_c32_class.py | 34 ++- tests/test_moe_glu_variants.py | 199 ++++++++++++------ 11 files changed, 506 insertions(+), 134 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cee7e8..a802678 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). time on every GPU. `KQ_MOE_C32=0` restores the previous kernels. - On M5 GPUs the iq2_s MoE gathers at one activation row take a few percent less time again. `KQ_MOE_NX=8` restores the earlier width. +- On M5 GPUs the iq1, iq2 and iq3 fused MoE gathers at one activation row take + up to 19% less time, shared-expert and score-mixed gathers included. + `KQ_MOE_C32=0` restores the previous kernels. +- On M5 GPUs `quantized_matmul` at one activation row takes up to 26% less + time for the iq1, iq2 and iq3 codecs. `KQ_QMV_C32=0` restores the previous + kernels. ## [0.5.1] diff --git a/README.md b/README.md index 4c4a44d..e74a3f8 100644 --- a/README.md +++ b/README.md @@ -303,12 +303,16 @@ 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. It runs 16 - K-lanes per row on M5 and later GPUs and 8 on earlier ones. `KQ_MOE_NX=8` or `16` forces a - width, 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_MOE_C32` - on Metal, the 32-weight decode step of the grid-codec (iq1, iq2 and iq3) fused + MoE gathers at one activation row. By default it runs for every grid codec on M5 and later GPUs + and for iq2_s on earlier ones. `0` restores the 16-weight kernels, and `1` runs it for every grid + codec on any GPU, also at wider rows where the dedupe kernels do not apply. The gate/up gathers + run 16 K-lanes per row on M5 and later GPUs and 8 on earlier ones, and the down gathers run 8. + `KQ_MOE_NX=8` or `16` forces a width, 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_QMV_C32` - on Metal, the same 32-weight step for the one-row `quantized_matmul` of a grid + codec, on by default on M5 and later GPUs. `0` restores the per-row kernels, `1` runs it on any + GPU, and `8` or `16` sets the K-lanes per row. Read per call. 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/bindings.cpp b/bindings.cpp index 96a0940..544edc4 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -131,9 +131,25 @@ NB_MODULE(_ext, m) { "_moe_c32_nx", &mlx_kquant::moe_c32_nx, "arch"_a, - "The K-lanes per output row (8 or 16) of the iq2_s c32 MoE gathers " + "The K-lanes per output row (8 or 16) of the c32 MoE GLU gathers " "for an architecture string, unless KQ_MOE_NX forces one."); + m.def( + "_moe_c32_on", + &mlx_kquant::moe_c32_on, + "arch"_a, + "codec"_a, + "Whether the MoE gathers of a codec take the c32 kernels at t = 1 by " + "default for an architecture string, unless KQ_MOE_C32 forces it."); + + m.def( + "_qmv_c32_on", + &mlx_kquant::qmv_c32_on, + "arch"_a, + "codec"_a, + "Whether the one-row quantized_matmul of a codec takes the c32 kernel " + "by default for an architecture string, unless KQ_QMV_C32 forces it."); + m.def( "_dequantize_route_takes", &mlx_kquant::dequantize_route_takes, diff --git a/docs/kernels.md b/docs/kernels.md index f68c0da..16bfa20 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -348,6 +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_C32` - the c32 mat-vec of the grid codecs at one activation row (see the fused MoE + section). `0` restores the per-row kernels and `1` runs c32 on any GPU. `8` or `16` runs it at + that many K-lanes per row. Unset takes it on M5 and later GPUs. A set `KQ_QMV_FINE` keeps the + per-row kernels. 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 @@ -1038,23 +1042,53 @@ 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, with 16 K-lanes per output row on -M5 and later GPUs and 8 on earlier ones. `KQ_MOE_NX=8` or `16` forces a width. 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 +At one activation row the grid codecs (iq1_s, iq1_m, iq2_xxs, iq2_xs, iq2_s, iq3_xxs and iq3_s) +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 bytes and the scale bytes then load once per 32 weights, and the grid +indices and sign bytes load as 16-bit pairs. Every gather takes this step: gate/up, the shared-expert +GLU, the plain down and both score-mixed downs (`gather_qmv_mix_kq` with the shared expert in the +routed codec or in q5_k, q6_k or q8_0, and `gather_qmv_mix_ns_kq`). The dedupe and half-dot forms +keep the 16-weight step. Outputs differ from the 16-weight kernels only in summation order. + +| GPU | Codecs on by default | GLU K-lanes per row | Down K-lanes per row | +|---|---|---|---| +| M5 and later | all seven | 16 | 8 | +| Earlier | iq2_s | 8 | 8 | + +`kq._ext._moe_c32_on(arch, codec)` and `kq._ext._moe_c32_nx(arch)` give the defaults for an +architecture string. `KQ_MOE_C32=0` restores the 16-weight kernels, and `KQ_MOE_C32=1` runs the +32-weight kernels for every grid codec, on any GPU and at wider rows where the dedupe kernels do not +apply. `KQ_MOE_NX=8` or `16` forces a width, and a forced `KQ_MOE_NX=32` keeps the 16-weight kernels. +The one-row `quantized_matmul` of a grid codec runs the same row step (`kq__qmv_c32`), on +M5 and later GPUs by default, at the GLU width. `KQ_QMV_C32` controls it (see the mat-vec levers). + 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%). -On M5 Max the 16-lane form is faster. Back to back at the same shapes, a MoE block replica with -both gathers at 16 K-lanes took 2.3% to 5.2% less time than at 8, in both arm orders. +On M5 Max the GLU gathers are faster at 16 lanes. Back to back at the same shapes, a MoE block +replica with both gathers at 16 K-lanes took 2.3% to 5.2% less time than at 8, in both arm orders. Alone, the +down gathers took up to 9% less time at 8 lanes than at 16. + +On M5 Max at E=32 and 64, top 8, N=4096 and 7168, expert widths 1536 to 4096, the 16-weight gathers +took longer than the 32-weight ones, in every family: + +| Codec | Longer with the 16-weight step | +|---|---| +| iq1_s | 15 to 23% | +| iq1_m | 12 to 19% | +| iq3_s | 10 to 19% | +| iq2_s | 6 to 13% | +| iq2_xxs | 2 to 8% | +| iq2_xs | 0 to 4% | +| iq3_xxs | 0 to 13% (gate/up even) | + +The one-row `quantized_matmul` per-row kernels took 2 to 14% longer than `qmv_c32` at K=2048 and +N=512, 9 to 31% longer at K=4096 and N=14336, and 15 to 35% longer at K=14336 and N=4096. + 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 diff --git a/metal/kq_moe_glu_kq.metal b/metal/kq_moe_glu_kq.metal index 1f2a538..f99122e 100644 --- a/metal/kq_moe_glu_kq.metal +++ b/metal/kq_moe_glu_kq.metal @@ -329,11 +329,53 @@ instantiate_kq_ext_sx_mix(q2_0, KqQ2_0Ext, iq4_nl, KqIq4_nlExt, float16_t) "kq_" #codec "_gather_qmv_c32" sfx "_" #type, \ kq_ext_gather_qmv, type, traits, nx, true) +// Shared-expert GLU and shared-expert mix with the shexp in scodec (the +// codec itself or an upcast q5_k, q6_k or q8_0); stem is "codec" or +// "codec_sx_scodec". +#define instantiate_kq_ext_c32_sx_nx(stem, traits, straits, type, nx, sfx) \ + instantiate_kernel( \ + "kq_" stem "_moe_glu_gather_shexp_silu_c32" sfx "_" #type, \ + kq_ext_moe_glu_gather_shexp, type, traits, straits, KQ_GLU_ACT_SILU, nx, \ + true) \ + instantiate_kernel( \ + "kq_" stem "_moe_glu_gather_shexp_gelu_c32" sfx "_" #type, \ + kq_ext_moe_glu_gather_shexp, type, traits, straits, KQ_GLU_ACT_GELU, nx, \ + true) \ + instantiate_kernel( \ + "kq_" stem "_gather_qmv_mix_c32" sfx "_" #type, \ + kq_ext_gather_qmv_mix, type, traits, straits, nx, true) + +#define instantiate_kq_ext_c32_sx(stem, traits, straits, type) \ + instantiate_kq_ext_c32_sx_nx(stem, traits, straits, type, 8, "") \ + instantiate_kq_ext_c32_sx_nx(stem, traits, straits, type, 16, "_nx16") \ + instantiate_kernel( \ + "kq_" stem "_gather_qmv_mix_sp_c32_" #type, \ + kq_ext_gather_qmv_mix_sp, type, traits, straits, 8, true) + +#define instantiate_kq_ext_c32_type(codec, traits, type) \ + instantiate_kq_ext_c32_nx(codec, traits, type, 8, "") \ + instantiate_kq_ext_c32_nx(codec, traits, type, 16, "_nx16") \ + instantiate_kernel( \ + "kq_" #codec "_qmv_c32_" #type, kq_ext_qmv_c32, type, traits, 8) \ + instantiate_kernel( \ + "kq_" #codec "_qmv_c32_nx16_" #type, kq_ext_qmv_c32, type, traits, 16) \ + instantiate_kernel( \ + "kq_" #codec "_gather_qmv_mix_ns_c32_" #type, \ + kq_ext_gather_qmv_mix_ns, type, traits, 8, true) \ + instantiate_kernel( \ + "kq_" #codec "_gather_qmv_mix_ns_c32_nx16_" #type, \ + kq_ext_gather_qmv_mix_ns, type, traits, 16, true) \ + instantiate_kernel( \ + "kq_" #codec "_gather_qmv_mix_ns_sp_c32_" #type, \ + kq_ext_gather_qmv_mix_ns_sp, type, traits, 8, true) \ + instantiate_kq_ext_c32_sx(#codec, traits, traits, type) \ + instantiate_kq_ext_c32_sx(#codec "_sx_q5_k", traits, KqQ5_KExt, type) \ + instantiate_kq_ext_c32_sx(#codec "_sx_q6_k", traits, KqQ6_KExt, type) \ + instantiate_kq_ext_c32_sx(#codec "_sx_q8_0", traits, KqQ8_0Ext, type) + #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_type(codec, traits, bfloat16_t) \ + instantiate_kq_ext_c32_type(codec, traits, float16_t) instantiate_kq_ext_c32(iq2_s, KqIq2_sExt) instantiate_kq_ext_c32(iq2_xxs, KqIq2_xxsExt) 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 d741634..2d036cd 100644 --- a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h +++ b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h @@ -1098,12 +1098,39 @@ template } } +// Dense one-row mat-vec through the c32 row step: the M = 1 quantized_matmul +// of the grid codecs. w is (N, row bytes), x one row of K. Dispatch: group +// (32, sg, 1), grid (N / (sg * 32 / NX), 1, 1). +template +[[kernel]] void kq_ext_qmv_c32( + const device uint8_t* w [[buffer(0)]], + const device T* x [[buffer(1)]], + device T* out [[buffer(2)]], + const constant int& K [[buffer(3)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tptg [[threads_per_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int RPS = 32 / NX; + const short tx = short(simd_lid % NX); + const short ty = short(simd_lid / NX); + const int out_row = tid.x * int(tptg.y) * RPS + int(simd_gid) * RPS + ty; + + KQ_EXT_STAGE_LUTS(Codec, kq_luts) + const float r = kq_ext_reduce(kq_ext_row_partial( + w, x, (int64_t)out_row, K, tx, kq_luts)); + if (tx == 0) { + out[out_row] = static_cast(r); + } +} + template < typename T, typename Codec, typename SCodec, int ACT, - int NX = KQ_EXT_NXPSG> + int NX = KQ_EXT_NXPSG, + bool C32 = false> [[kernel]] void kq_ext_moe_glu_gather_shexp( const device uint8_t* gw [[buffer(0)]], const device uint8_t* uw [[buffer(1)]], @@ -1134,11 +1161,11 @@ template < KQ_EXT_STAGE_LUTS(SCodec, kq_sluts) float2 gu; if (shared_slot) { - gu = kq_ext_glu_row_partial( + gu = kq_ext_glu_row_partial( sgw, suw, x, (int64_t)out_row, K, tx, kq_sluts); } else { const int expert = int(indices[tid.z * n_route + tid.y]); - gu = kq_ext_glu_row_partial( + 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); @@ -1208,7 +1235,7 @@ template < // No-shared-expert mix (gemma-style MoE: plain score-weighted sum over the // routed slots): indices and scores are both [T, S], every slot gathers from // the expert stack. -template +template [[kernel]] void kq_ext_gather_qmv_mix_ns( const device uint8_t* w [[buffer(0)]], const device T* h [[buffer(1)]], @@ -1233,7 +1260,7 @@ template const int expert = int(indices[tid.z * S + slot]); const device T* xs = h + ((int64_t)tid.z * S + slot) * K; result += scores[tid.z * S + slot] * - kq_ext_row_partial( + kq_ext_row_partial( w, xs, (int64_t)expert * N + out_row, K, tx, kq_luts); } result = kq_ext_reduce(result); @@ -1258,7 +1285,7 @@ template // S <= KQ_MOE_SP_MAX_S. #define KQ_MOE_SP_MAX_S 16 -template +template [[kernel]] void kq_ext_gather_qmv_mix_ns_sp( const device uint8_t* w [[buffer(0)]], const device T* h [[buffer(1)]], @@ -1292,7 +1319,7 @@ template const int expert = int(indices[tid.z * S + slot]); const device T* xs = h + ((int64_t)tid.z * S + slot) * K; - parts[lrow][tx][slot] = kq_ext_row_partial( + parts[lrow][tx][slot] = kq_ext_row_partial( w, xs, (int64_t)expert * N + out_row, K, tx, kq_luts); threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1308,7 +1335,12 @@ template } } -template +template < + typename T, + typename Codec, + typename SCodec, + int NX = KQ_EXT_NXPSG, + bool C32 = false> [[kernel]] void kq_ext_gather_qmv_mix( const device uint8_t* w [[buffer(0)]], const device uint8_t* sw [[buffer(1)]], @@ -1336,14 +1368,14 @@ template const int expert = int(indices[tid.z * (S - 1) + slot]); const device T* xs = h + ((int64_t)tid.z * S + slot) * K; result += scores[tid.z * SC + slot] * - kq_ext_row_partial( + kq_ext_row_partial( w, xs, (int64_t)expert * N + out_row, K, tx, kq_luts); } { // SC == S - 1: the shared slot's weight is an implicit 1 const device T* xs = h + ((int64_t)tid.z * S + (S - 1)) * K; result += (SC == S ? scores[tid.z * SC + (S - 1)] : 1.0f) * - kq_ext_row_partial( + kq_ext_row_partial( sw, xs, (int64_t)out_row, K, tx, kq_sluts); } result = kq_ext_reduce(result); @@ -1358,7 +1390,12 @@ template // score chain, so outputs are bit-identical to it. Dispatch: group // (32, 2 * S, 1), grid (N / (2 * RPS), 1, T); host gates // S <= KQ_MOE_SP_MAX_S. -template +template < + typename T, + typename Codec, + typename SCodec, + int NX = KQ_EXT_NXPSG, + bool C32 = false> [[kernel]] void kq_ext_gather_qmv_mix_sp( const device uint8_t* w [[buffer(0)]], const device uint8_t* sw [[buffer(1)]], @@ -1389,10 +1426,10 @@ template const device T* xs = h + ((int64_t)tid.z * S + slot) * K; if (slot < S - 1) { const int expert = int(indices[tid.z * (S - 1) + slot]); - parts[lrow][tx][slot] = kq_ext_row_partial( + parts[lrow][tx][slot] = kq_ext_row_partial( w, xs, (int64_t)expert * N + out_row, K, tx, kq_luts); } else { - parts[lrow][tx][slot] = kq_ext_row_partial( + parts[lrow][tx][slot] = kq_ext_row_partial( sw, xs, (int64_t)out_row, K, tx, kq_sluts); } threadgroup_barrier(mem_flags::mem_threadgroup); diff --git a/src/kquant.h b/src/kquant.h index 930ccdd..7400551 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -84,11 +84,23 @@ bool dequantize_route_takes( int N, int K); -// The K-lanes per output row (8 or 16) of the iq2_s c32 MoE gathers on a GPU +// The K-lanes per output row (8 or 16) of the c32 MoE GLU gathers on a GPU // with this architecture string (as get_architecture gives it), when // KQ_MOE_NX does not force one. int moe_c32_nx(const std::string& arch); +// Whether this codec has the 32-weight decode step of the c32 kernels (the +// grid codecs). +bool codec_has_c32(const std::string& codec); + +// Whether the MoE gathers of this codec take the c32 kernels at t = 1 by +// default on a GPU with this architecture string. +bool moe_c32_on(const std::string& arch, const std::string& codec); + +// Whether the one-row quantized_matmul of this codec takes the c32 kernel by +// default on a GPU with this architecture string. +bool qmv_c32_on(const std::string& arch, const std::string& codec); + // Every kernel name the split-K dispatchers can build (qmm_splitk, // qmm_nax_splitk, verify_mma, verify_nax and their fold), for the kernel // name tests; empty on non-Metal builds. diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 1e866a9..3933df8 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -1884,6 +1884,53 @@ static bool qmv_splitk( return true; } +// Dispatches the M = 1 mat-vec of a grid codec through the c32 row step of +// the MoE gathers (kq__qmv_c32). KQ_QMV_C32, read per call: 0 keeps +// the qmv kernels, 1 takes c32 at this GPU's width, 8 or 16 at that width; +// unset follows qmv_c32_on. Returns false, dispatching nothing, when the +// codec, shape or pipeline does not fit. +static bool qmv_c32( + const array& x, + const array& w, + array& out, + int N, + int K, + Device& d, + const Stream& s, + const std::string& kquant_type) { + const int env = kq_env_int("KQ_QMV_C32", -1); + static const std::string arch = d.get_architecture(); + if (env == 0 || (env == -1 && !qmv_c32_on(arch, kquant_type)) || + !codec_has_c32(kquant_type)) { + return false; + } + const int nx = env == 8 || env == 16 ? env : moe_c32_nx(arch); + const int rps = 32 / nx; + int sg = kquant_type == "iq2_xs" ? 8 : 2; + while (sg > 1 && N % (sg * rps) != 0) { + sg /= 2; + } + if (N % (sg * rps) != 0 || K % 256 != 0) { + return false; + } + const std::string kname = "kq_" + kquant_type + "_qmv_c32" + + (nx == 16 ? "_nx16_" : "_") + kq_type_string(x.dtype()); + auto kernel = kq_get_kernel(d, kname); + if (int(kernel->maxTotalThreadsPerThreadgroup()) < 32 * sg) { + return false; + } + const MTL::Size grid_dims(N / (sg * rps), 1, 1); + kq_log_qmm(kname, 1, N, K, grid_dims, x); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(w, 0); + ce.set_input_array(x, 1); + ce.set_output_array(out, 2); + ce.set_bytes(K, 3); + ce.dispatch_threadgroups(grid_dims, MTL::Size(32, sg, 1)); + return true; +} + // Matrix-times-vector quantized kernel dispatch (no biases). void qmv( const array& x, @@ -1920,6 +1967,10 @@ void qmv( x, w, scales, out, group_size, bits, N, K, d, s, kquant_type)) { return; } + if (M == 1 && B == 1 && qmv_fine_env == -1 && + qmv_c32(x, w, out, N, K, d, s, kquant_type)) { + return; + } const bool fine_ok = B == 1 && codec_has_qmv_fine(kquant_type); const bool use_fine = fine_ok && (qmv_fine_env == 1 || diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index 71dbc01..7be6e71 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -354,19 +354,39 @@ inline bool kq_moe_has_c32(const std::string& codec) { } // 32-weight decode step ("_c32" kernels, KqExtC32) for the grid-codec -// 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. +// gathers: by default at t = 1, where this GPU's class takes it for the +// codec (moe_c32_on). 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. The dedupe +// and half-dot forms have no c32 kernels and keep the 16-weight step. inline bool kq_moe_c32(const std::string& codec, int T) { if (!kq_moe_has_c32(codec) || kq_moe_nx_env() == 32) { return false; } const int v = kq_env_flag("KQ_MOE_C32", -1); - return v == -1 ? T == 1 : v != 0; + if (v != -1) { + return v != 0; + } +#ifdef _METAL_ + static const std::string arch = + mx::metal::device(mx::Device::gpu).get_architecture(); + return T == 1 && moe_c32_on(arch, codec); +#else + return false; +#endif } -// A forced KQ_MOE_NX of 8 or 16, else this GPU's class width (moe_c32_nx). +// A shared expert in the routed codec or an upcast q5_k, q6_k or q8_0: the +// pairs with c32 shexp and mix kernels. +inline bool kq_moe_c32_shexp_ok( + const std::string& codec, + const std::string& shexp) { + return shexp == codec || shexp == "q5_k" || shexp == "q6_k" || + shexp == "q8_0"; +} + +// K-lanes per row of the c32 GLU gathers: a forced KQ_MOE_NX of 8 or 16, +// else this GPU's class width (moe_c32_nx). inline int kq_moe_c32_nx() { const int v = kq_moe_nx_env(); if (v == 8 || v == 16) { @@ -381,6 +401,15 @@ inline int kq_moe_c32_nx() { #endif } +// K-lanes per row of the c32 down gathers (plain, mix and mix_ns): a forced +// KQ_MOE_NX of 8 or 16, else 8. On an M5 Max the c32 down gathers of the +// grid codecs took 0 to 9% less time at 8 than at 16 (E=32..64, top 8, +// K=1536 and 4096). +inline int kq_moe_c32_down_nx() { + const int v = kq_moe_nx_env(); + return v == 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; @@ -664,7 +693,7 @@ void KQuantGatherQMVKQ::eval_gpu( const bool use_half = !biased && kq_moe_half(kquant_type_, K); const bool c32 = !biased && !use_half && kq_moe_c32(kquant_type_, T); const int nx = use_half ? 8 - : c32 ? kq_moe_c32_nx() + : c32 ? kq_moe_c32_down_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 @@ -786,8 +815,10 @@ void KQuantMoEGLUShexpKQ::eval_gpu( (shexp_type_ == kquant_type_ || shexp_type_ == "q5_k" || shexp_type_ == "q6_k" || shexp_type_ == "q8_0"); const bool dd = !use_half && kq_moe_dedup(T, KQ_DD_MAX_T); - const int nx = - use_half ? 8 : kq_moe_pick_nx((int64_t)N * (R + 1) * T, K, true); + const bool c32 = !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 + 1) * T, K, true); std::string stem = shexp_type_ == kquant_type_ ? kq_gather_stem_nx(kquant_type_, K, nx) : kquant_type_ + "_sx_" + shexp_type_; @@ -798,8 +829,8 @@ void KQuantMoEGLUShexpKQ::eval_gpu( // time when the limit is positive (kq_glu_epilogue). const std::string act = act_ == "silu_limit" ? "silu" : act_; std::string kname = "kq_" + stem + "_moe_glu_gather_shexp_" + - (use_half ? "h_" : "") + (dd ? "dd_" : "") + act + kq_nx_suffix(nx) + - "_" + kq_type_string(x.dtype()); + (use_half ? "h_" : "") + (dd ? "dd_" : "") + 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); @@ -861,14 +892,20 @@ void KQuantGatherQMVMixKQ::eval_gpu( const std::string sp_stem = shexp_type_ == kquant_type_ ? kq_gather_stem(kquant_type_, K) : kquant_type_ + "_sx_" + shexp_type_; + // c32 serves the loop and slot-parallel forms; the dedupe form keeps the + // 16-weight step. + const bool c32 = !use_half && kq_moe_c32(kquant_type_, T) && + kq_moe_c32_shexp_ok(kquant_type_, shexp_type_); + const int nx_one = c32 + ? kq_moe_c32_down_nx() + : kq_moe_pick_nx( + (int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, T, K)); const bool sp_ok = !use_half && sp_stem != "q6_k" && sp_stem != "q8_0" && - kq_moe_pick_nx( - (int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, T, K)) == 8 && - kq_moe_sp((int64_t)T * (N / 8), S); + nx_one == 8 && kq_moe_sp((int64_t)T * (N / 8), S); const int form = kq_moe_mix_form(sp_ok, dd_ok); if (form == KQ_MIX_SP) { - const std::string kname = - "kq_" + sp_stem + "_gather_qmv_mix_sp_" + kq_type_string(x.dtype()); + const std::string kname = "kq_" + sp_stem + "_gather_qmv_mix_sp_" + + (c32 ? "c32_" : "") + kq_type_string(x.dtype()); auto kernel = kq_get_kernel(d, kname); // 64 * S threads must fit the pipeline; the loop kernel serves the // call when they do not. @@ -924,8 +961,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( kq_moe_pick_sg(kquant_type_, N, nx)); return; } - int nx = - kq_moe_pick_nx((int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, T, K)); + int nx = nx_one; if (use_half && nx > 16) { nx = 16; } @@ -936,8 +972,8 @@ void KQuantGatherQMVMixKQ::eval_gpu( const bool tuned = stem == "q6_k" || stem == "q8_0"; const bool fine = tuned && nx == 8 && kq_gather_fine((int64_t)(N / 8) * T); std::string kname = "kq_" + stem + "_gather_qmv_mix" + - (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); @@ -999,13 +1035,21 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( // dropped. Decode-scale launches route to the slot-parallel variant // unless the codec widens (kq_moe_mix_ns_wide). const bool dd_ok = S <= 16 && kq_moe_dedup(T, KQ_DD_MAX_T); + const bool c32_ok = kq_moe_c32(kquant_type_, T); const bool sp_ok = - kq_moe_pick_nx( - (int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, T, K)) == 8 && + (c32_ok + ? kq_moe_c32_down_nx() + : kq_moe_pick_nx( + (int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, T, K))) == + 8 && kq_moe_sp((int64_t)T * (N / 8), S); const bool dd = kq_moe_mix_form(sp_ok, dd_ok) == KQ_MIX_DD; - const int nx = kq_moe_pick_nx( - (int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, dd ? 2 : T, K)); + // c32 serves the loop and slot-parallel forms, not the dedupe form. + const bool c32 = c32_ok && !dd; + const int nx = c32 + ? kq_moe_c32_down_nx() + : kq_moe_pick_nx( + (int64_t)N * T, K, kq_moe_mix_ns_wide(kquant_type_, dd ? 2 : T, K)); if (dd) { kq_gather_qmv_pair_gpu( d, @@ -1029,7 +1073,8 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( kq_moe_mix_form(true, false) == KQ_MIX_SP; auto ns_name = [&](bool slot_parallel) { return "kq_" + kquant_type_ + "_gather_qmv_mix_ns" + - (slot_parallel ? "_sp" : kq_nx_suffix(nx)) + "_" + + (slot_parallel ? "_sp" : "") + (c32 ? "_c32" : "") + + (slot_parallel ? "" : kq_nx_suffix(nx)) + "_" + kq_type_string(x.dtype()); }; std::string kname = ns_name(sp); @@ -1336,13 +1381,12 @@ std::vector KQuantGatherQMVMixBias::output_shapes( return {{inputs[4].shape(0), inputs[0].shape(1)}}; } -// applegpu_g, two digit gen, a lower case suffix other than p: -// NX 16 from gen 17 (M5), where it beat NX 8 by 2 to 6% on a Muse Spark MoE -// block (M5 Max, E=256, top 8, N=K=4096). NX 8 elsewhere. -int moe_c32_nx(const std::string& arch) { +// The generation of a Mac GPU, from applegpu_g with a two digit +// gen and a lower case suffix other than p; 0 for a phone or anything else. +static int moe_mac_gpu_gen(const std::string& arch) { const std::string head = "applegpu_g"; if (arch.size() != head.size() + 3 || arch.compare(0, head.size(), head)) { - return 8; + return 0; } const char d0 = arch[head.size()]; const char d1 = arch[head.size() + 1]; @@ -1350,10 +1394,35 @@ int moe_c32_nx(const std::string& arch) { if (!std::isdigit(static_cast(d0)) || !std::isdigit(static_cast(d1)) || !std::islower(static_cast(suffix)) || suffix == 'p') { - return 8; + return 0; } - const int gen = (d0 - '0') * 10 + (d1 - '0'); - return gen >= 17 ? 16 : 8; + return (d0 - '0') * 10 + (d1 - '0'); +} + +// K-lanes per row of the c32 GLU gathers: 16 from gen 17 (M5), where it +// beat 8 by 2 to 6% on a Muse Spark MoE block (M5 Max, E=256, top 8, +// N=K=4096). 8 elsewhere. +int moe_c32_nx(const std::string& arch) { + return moe_mac_gpu_gen(arch) >= 17 ? 16 : 8; +} + +// iq2_s on every GPU (M3 Max and M5 Max measured). The other grid codecs +// from gen 17 (M5), measured on an M5 Max only: at E=32, top 8, N=K=4096 +// the 16-weight gathers took 1 to 23% longer (iq1_s most, iq3_xxs least). +bool moe_c32_on(const std::string& arch, const std::string& codec) { + return codec == "iq2_s" || + (kq_moe_has_c32(codec) && moe_mac_gpu_gen(arch) >= 17); +} + +bool codec_has_c32(const std::string& codec) { + return kq_moe_has_c32(codec); +} + +// The M = 1 quantized_matmul of a grid codec through the c32 row step: from +// gen 17 (M5), measured on an M5 Max only. At K=4096, N=14336 the qmv +// kernels took 2 to 21% longer (iq2_s and iq3_s most, iq2_xxs least). +bool qmv_c32_on(const std::string& arch, const std::string& codec) { + return moe_mac_gpu_gen(arch) >= 17 && kq_moe_has_c32(codec); } // Codecs with the fused kq GLU/gather kernels wired (kq_moe_glu_kq.h): diff --git a/tests/test_moe_c32_class.py b/tests/test_moe_c32_class.py index d4d2baa..7074782 100644 --- a/tests/test_moe_c32_class.py +++ b/tests/test_moe_c32_class.py @@ -1,11 +1,12 @@ -"""The K-lanes per row of the iq2_s c32 MoE gathers by GPU class. +"""The c32 kernel defaults by GPU class: which codecs take them and the +K-lanes per row of the GLU gathers. -``_moe_c32_nx`` answers for any architecture string, so these cases run on -every machine, CPU builds included. +The class functions answer for any architecture string, so these cases run +on every machine, CPU builds included. """ import pytest -from mlx_kquant._ext import _moe_c32_nx +from mlx_kquant._ext import _moe_c32_nx, _moe_c32_on, _qmv_c32_on WIDE = [f"applegpu_g{g}{c}" for g in (17, 18, 19) for c in "gsdc"] NARROW = [f"applegpu_g{g}{c}" for g in (13, 14, 15, 16) for c in "gsdc"] + [ @@ -19,6 +20,8 @@ "gpu_g17s", "", ] +GRID = ["iq2_xxs", "iq2_xs", "iq3_xxs", "iq3_s", "iq1_s", "iq1_m"] +OTHER = ["q4_k", "q6_k", "q8_0", "iq4_xs", "iq4_nl", "mxfp4", "q2_0"] @pytest.mark.parametrize("arch", WIDE) @@ -29,3 +32,26 @@ def test_m5_and_later_take_nx16(arch): @pytest.mark.parametrize("arch", NARROW) def test_earlier_phones_and_malformed_keep_nx8(arch): assert _moe_c32_nx(arch) == 8 + + +@pytest.mark.parametrize("arch", WIDE) +def test_m5_and_later_take_c32_for_every_grid_codec(arch): + for codec in ["iq2_s"] + GRID: + assert _moe_c32_on(arch, codec), codec + assert _qmv_c32_on(arch, codec), codec + + +@pytest.mark.parametrize("arch", NARROW) +def test_earlier_keep_c32_for_iq2_s_gathers_only(arch): + assert _moe_c32_on(arch, "iq2_s") + for codec in GRID: + assert not _moe_c32_on(arch, codec), codec + for codec in ["iq2_s"] + GRID: + assert not _qmv_c32_on(arch, codec), codec + + +@pytest.mark.parametrize("arch", WIDE + NARROW) +def test_other_codecs_never_take_c32(arch): + for codec in OTHER: + assert not _moe_c32_on(arch, codec), codec + assert not _qmv_c32_on(arch, codec), codec diff --git a/tests/test_moe_glu_variants.py b/tests/test_moe_glu_variants.py index c1b5dc1..878b166 100644 --- a/tests/test_moe_glu_variants.py +++ b/tests/test_moe_glu_variants.py @@ -516,85 +516,154 @@ def test_shexp_mix_slot_parallel_matches_dispatch(): assert (name in r.stderr) == (listed == "1"), (codec, scodec, listed) +C32_CODECS = ["iq2_s", "iq2_xxs", "iq2_xs", "iq3_xxs", "iq3_s", "iq1_s", "iq1_m"] + _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] +codec, T, dtype = sys.argv[1], int(sys.argv[2]), getattr(mx, sys.argv[3]) +out_path = sys.argv[4] 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)) +def wire(shape, c=codec): + w = mx.array((rng.standard_normal(shape) * 0.1).astype(np.float32)) + wq, _ = kq.quantize(w, c, mx.ones((shape[-1],), 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) +gw, uw, dw = wire((E, N, K)), wire((E, N, K)), wire((E, K, N)) +sg, su, sd = wire((N, K)), wire((N, K)), wire((K, N)) +qg, qu, qd = wire((N, K), "q6_k"), wire((N, K), "q6_k"), wire((K, N), "q6_k") +def act(shape): + a = (rng.standard_normal(shape) * 0.05).astype(np.float32) + return mx.array(a).astype(dtype) +x, h, hs = act((T, K)), act((T, S, N)), act((T, S + 1, N)) 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))) +sc = mx.array(rng.uniform(0.05, 0.9, size=(T, S)).astype(np.float32)) +sx = {"shexp_kquant_type": "q6_k"} +out = { + "glu": kq.moe_glu_gather_kq(x, gw, uw, codec, inds, act="silu"), + "qmv": kq.gather_qmv_kq(h, dw, codec, inds), + "mix_ns": kq.gather_qmv_mix_ns_kq(h, dw, codec, inds, sc), + "shexp": kq.moe_glu_gather_shexp_kq(x, gw, uw, sg, su, codec, inds), + "shexp_sx": kq.moe_glu_gather_shexp_kq(x, gw, uw, qg, qu, codec, inds, **sx), + "mix": kq.gather_qmv_mix_kq(hs, dw, sd, codec, inds, sc), + "mix_sx": kq.gather_qmv_mix_kq(hs, dw, qd, codec, inds, sc, **sx), + "dense": kq.quantized_matmul( + x[:1], gw[0], mx.zeros((1,), mx.uint8), codec, transpose=True), +} +mx.eval(out) +np.savez(out_path, **{k: np.array(v.astype(mx.float32)) for k, v in out.items()}) """ -def _c32_run(T, dtype, env, tmp_path, arm): - f = tmp_path / f"{arm}.npz" +def _c32_run(codec, T, dtype, env, tmp_path, arm): + f = tmp_path / f"{codec}-{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}", + [sys.executable, "-c", _C32_SNIPPET, codec, str(T), dtype, str(f)], + {**env, "KQ_MOE_NX_LOG": "1", "KQ_QMM_LOG": "1"}, + f"c32 {codec} {arm} T={T} {dtype}", ) return np.load(f), r.stderr -def _c32_class_suffix(): +def _c32_class(codec): + """This GPU's c32 defaults for codec: (MoE on, dense on, GLU suffix).""" import mlx.core as mx - from mlx_kquant._ext import _moe_c32_nx + from mlx_kquant._ext import _moe_c32_nx, _moe_c32_on, _qmv_c32_on + + arch = mx.device_info().get("architecture", "") + nx = _moe_c32_nx(arch) + return ( + _moe_c32_on(arch, codec), + _qmv_c32_on(arch, codec), + "" if nx == 8 else f"_nx{nx}", + ) + + +def _c32_names(codec, dtype, glu, down): + """The c32 kernel names of the snippet's seven gathers at one row, with + GLU suffix glu and down suffix down ("" or "_nx16"); the score-mixed + downs take the slot-parallel form at 8 lanes.""" + t = f"_{dtype}_t" + sx = f"{codec}_sx_q6_k" + mix = "_sp_c32" + t if down == "" else "_c32" + down + t + return [ + f"kq_{codec}_moe_glu_gather_silu_c32{glu}{t}", + f"kq_{codec}_gather_qmv_c32{down}{t}", + f"kq_{codec}_gather_qmv_mix_ns{mix}", + f"kq_{codec}_moe_glu_gather_shexp_silu_c32{glu}{t}", + f"kq_{sx}_moe_glu_gather_shexp_silu_c32{glu}{t}", + f"kq_{codec}_gather_qmv_mix{mix}", + f"kq_{sx}_gather_qmv_mix{mix}", + ] - nx = _moe_c32_nx(mx.device_info().get("architecture", "")) - return "" if nx == 8 else f"_nx{nx}" +def _assert_close(got, ref): + for key in ref.files: + rel = np.abs(got[key] - ref[key]).max() / np.abs(ref[key]).max() + assert rel < 2e-3, (key, rel) + +@pytest.mark.parametrize("codec", C32_CODECS) @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 this GPU class's NX (8, or 16 from M5), or a forced 8 - or 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.""" - sfx = _c32_class_suffix() - want_glu = f"kq_iq2_s_moe_glu_gather_silu_c32{sfx}_{dtype}_t" - want_qmv = f"kq_iq2_s_gather_qmv_c32{sfx}_{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 - narrow, names = _c32_run(1, dtype, {"KQ_MOE_NX": "8"}, tmp_path, "nx8") - assert f"kq_iq2_s_moe_glu_gather_silu_c32_{dtype}_t" in names, names - assert f"kq_iq2_s_gather_qmv_c32_{dtype}_t" 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, narrow, 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") +def test_c32_dispatch(codec, dtype, tmp_path): + """At one row the grid codecs' gathers (gate/up, down, both score-mixed + downs, the shared-expert GLU) and the one-row quantized_matmul take the + c32 kernels where this GPU's class does (_moe_c32_on, _qmv_c32_on): the + GLU gathers at the class width, the downs at 8 lanes, the dense mat-vec + at the class width. KQ_MOE_C32 and KQ_QMV_C32 force either way, and + KQ_MOE_NX or KQ_QMV_C32 of 8 or 16 force a width. The forms differ only + in summation order.""" + moe_on, qmv_on, gsfx = _c32_class(codec) + t = f"_{dtype}_t" + dense = f"kq_{codec}_qmv_c32" + old, names = _c32_run( + codec, 1, dtype, {"KQ_MOE_C32": "0", "KQ_QMV_C32": "0"}, tmp_path, "old" + ) assert "_c32" not in names, names + + new, names = _c32_run(codec, 1, dtype, {}, tmp_path, "default") + for name in _c32_names(codec, dtype, gsfx, ""): + assert (name in names) == moe_on, (name, names) + assert (f"{dense}{gsfx}{t}" in names) == qmv_on, names + _assert_close(new, old) + + on = {"KQ_MOE_C32": "1", "KQ_QMV_C32": "1"} + forced, names = _c32_run(codec, 1, dtype, on, tmp_path, "forced") + for name in _c32_names(codec, dtype, gsfx, "") + [f"{dense}{gsfx}{t}"]: + assert name in names, (name, names) + _assert_close(forced, old) + + for nx, sfx in (("8", ""), ("16", "_nx16")): + env = {"KQ_MOE_C32": "1", "KQ_MOE_NX": nx, "KQ_QMV_C32": nx} + got, names = _c32_run(codec, 1, dtype, env, tmp_path, f"nx{nx}") + for name in _c32_names(codec, dtype, sfx, sfx) + [f"{dense}{sfx}{t}"]: + assert name in names, (name, names) + _assert_close(got, old) + + +@pytest.mark.parametrize("codec", C32_CODECS) +def test_c32_stands_down(codec, tmp_path): + """A forced KQ_MOE_NX=32 keeps the 16-weight gathers, as do two rows + unless KQ_MOE_C32=1 forces c32 there (the dedupe forms, which have no + c32 kernels, then stand down).""" + on = {"KQ_MOE_C32": "1"} _, names = _c32_run( - 2, dtype, {"KQ_MOE_C32": "1", "KQ_MOE_DEDUP": "0"}, tmp_path, "t2c32" + codec, 1, "bfloat16", {**on, "KQ_MOE_NX": "32"}, tmp_path, "nx32" ) - assert want_glu in names and want_qmv in names, names + assert "gather_qmv_c32" not in names and "gather_silu_c32" not in names + assert f"kq_{codec}_moe_glu_gather_silu_nx32_" in names, names + _, names = _c32_run(codec, 2, "bfloat16", {}, tmp_path, "t2") + assert "gather_silu_c32" not in names and "gather_qmv_c32" not in names + _, names = _c32_run( + codec, 2, "bfloat16", {**on, "KQ_MOE_DEDUP": "0"}, tmp_path, "t2c32" + ) + _, _, gsfx = _c32_class(codec) + for name in _c32_names(codec, "bfloat16", gsfx, "")[:2]: + assert name in names, (name, names) _C32_LIVE_SNIPPET = r""" @@ -623,9 +692,8 @@ def test_c32_reads_live(): {"KQ_MOE_NX_LOG": "1"}, "c32 live read", ) - sfx = _c32_class_suffix() assert "kq_iq2_s_gather_qmv_bfloat16_t" in r.stderr, r.stderr - assert f"kq_iq2_s_gather_qmv_c32{sfx}_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""" @@ -635,30 +703,37 @@ def test_c32_reads_live(): import mlx.core as mx import test_moe_glu as tm -res = tm._check_codec("iq2_s", dtype=getattr(mx, sys.argv[2])) +res = tm._check_codec(sys.argv[3], 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("codec", C32_CODECS) @pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) -def test_c32_accuracy(dtype): +def test_c32_accuracy(codec, dtype): """The 32-weight-step kernels, forced at the reference test's three rows, - stay inside the gather bound of the f32 dequantized reference.""" + stay inside the gather bound of the f32 dequantized reference for the + gate/up, down, shared-expert GLU and shared-expert mix gathers.""" tests_dir = os.path.dirname(os.path.abspath(__file__)) r = _run_child( - [sys.executable, "-c", _C32_REF_SNIPPET, tests_dir, dtype], + [sys.executable, "-c", _C32_REF_SNIPPET, tests_dir, dtype, codec], {"KQ_MOE_C32": "1", "KQ_MOE_DEDUP": "0", "KQ_MOE_NX_LOG": "1"}, - f"c32 reference {dtype}", + f"c32 reference {codec} {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 + for op in ( + "moe_glu_gather_silu_c32", + "gather_qmv_c32", + "moe_glu_gather_shexp_silu_c32", + "gather_qmv_mix_sp_c32", + ): + assert f"kq_{codec}_{op}" in r.stderr, (op, 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") + pytest.skip(f"{codec} reference unavailable") for name, rel, ok in res: assert ok, (name, rel) From 493a356fb0b96c2d674872bfca35b5c666411e8a Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:32:45 -0700 Subject: [PATCH 3/3] test(moe): iq3_xxs mix_ns at t=1 expects the c32 slot-parallel kernel where c32 is on --- src/kquant_moe_glu.cpp | 4 +++- tests/test_moe_glu_variants.py | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index 7be6e71..96d2a22 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -94,7 +94,9 @@ inline int kq_moe_pick_nx(int64_t rows, int K, bool wide_ok) { // q2_k / q3_k (scale unpack heavy) lose 8-14%, q6_k / q8_0 are flat, the // iq2 grid codecs and q5_k lose in the chained regime, and every codec is // flat-to-negative at t >= 4 or K < 2048, so those keep the slot-parallel -// form. KQ_MOE_NX still forces a width. +// form. Where c32 is on, iq3_xxs at t = 1 takes the c32 slot-parallel +// kernel instead (M5 Max: 7-9% less time than this pick). KQ_MOE_NX still +// forces a width. inline bool kq_moe_mix_ns_wide(const std::string& codec, int T, int K) { if (T > 2 || K < 2048) { return false; diff --git a/tests/test_moe_glu_variants.py b/tests/test_moe_glu_variants.py index 878b166..052878f 100644 --- a/tests/test_moe_glu_variants.py +++ b/tests/test_moe_glu_variants.py @@ -134,8 +134,10 @@ def test_sp_bit_identical(codec, tmp_path): def test_mix_ns_wide_pick(codec, T, tmp_path): """The codec-keyed nx16 mix_ns pick (T <= 2, K >= 2048) dispatches the _nx16 kernel for the listed codecs, the slot-parallel kernel otherwise, - and matches the forced-nx8 launch to summation-order noise.""" - wide = codec != "q6_k" and T <= 2 + and matches the forced-nx8 launch to summation-order noise. A c32 codec + at T = 1 takes the c32 slot-parallel kernel instead where c32 is on.""" + c32 = T == 1 and _c32_class(codec)[0] + wide = codec != "q6_k" and T <= 2 and not c32 outs, names = {}, {} for arm, env in (("auto", {}), ("nx8", {"KQ_MOE_NX": "8"})): f = tmp_path / f"{arm}.npy" @@ -148,6 +150,7 @@ def test_mix_ns_wide_pick(codec, T, tmp_path): outs[arm] = np.load(f) names[arm] = r.stdout + r.stderr assert ("_nx16_" in names["auto"]) == wide, names["auto"] + assert ("_mix_ns_sp_c32_" in names["auto"]) == c32, names["auto"] assert "_sp_" in names["nx8"], names["nx8"] a, b = outs["auto"], outs["nx8"] scale = np.abs(b).max() + 1e-6