From f6577ab637e61bff472a3535787c5185b04f5a51 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:04:51 -0700 Subject: [PATCH 1/3] (metal): iq2_xxs nax gather loader dequants via vectorized chunk16s, bit-identical --- .../backend/metal/kernels/kq_quantized_nax.h | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index 0988bd3..7597735 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h @@ -2837,21 +2837,18 @@ struct KqNaxIq2_xxsBlockLoader { void load_unsafe() const { const short sb = (reduction_dim == 1 ? kt_base : fixed_kt_base) + bj / k_tile_size; - const float d = float(*(const device half*)src); - const device uint8_t* qs = src + KQ_IQ2_XXS_QS_OFFSET + sb * 8; - const uint signbits = uint(qs[4]) | (uint(qs[5]) << 8) | - (uint(qs[6]) << 16) | (uint(qs[7]) << 24); - const float db = d * (0.5f + float(signbits >> 28)) * 0.25f; + // The 32-weight k-tile is one ib32 group = chunk16s chunks 2*sb and + // 2*sb+1 (shared scale word). Sign-select on the exact grid bytes then + // one scale multiply is bit-identical to the scalar per-weight form: + // (db*gb)*(+-1) == db*(+-gb) in IEEE. + float4x4 r0, r1; + float s0, s1; + KqExtDeq::deq_chunk16s(src, short(2 * sb), r0, s0); + KqExtDeq::deq_chunk16s(src, short(2 * sb + 1), r1, s1); #pragma unroll - for (short i = 0; i < n_reads; i++) { - const int p = i; - const int l = p / 8; - const int j = p % 8; - const uint8_t signs = ksigns_iq2xs[(signbits >> (7 * l)) & 127]; - const uint64_t g = iq2xxs_grid[qs[l]]; - const float gb = float((g >> (8 * j)) & 0xff); - const float sgn = (signs & kmask_iq2xs[j]) ? -1.f : 1.f; - dst[i] = T(db * gb * sgn); + for (short i = 0; i < 16; i++) { + dst[i] = T(s0 * r0[i / 4][i % 4]); + dst[16 + i] = T(s1 * r1[i / 4][i % 4]); } } From b8ee030fd4a6d248d51737fe5896174a12489921 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:20:21 -0700 Subject: [PATCH 2/3] perf(metal): gather rhs nax simdgroups skip mma work for expert segments outside their row band --- .../backend/metal/kernels/kq_quantized_nax.h | 132 ++++++++++-------- 1 file changed, 74 insertions(+), 58 deletions(-) diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index 7597735..2ce294a 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h @@ -3505,6 +3505,16 @@ METAL_FUNC void kq_gather_qmm_rhs_nax_tgp_impl( } threadgroup_barrier(mem_flags::mem_none); + // Rows this simdgroup owns within the current expert segment. When the + // segment misses the simdgroup's row band the store below is an empty + // slice, so its A loads and matmads are dead work: skip them. The + // threadgroup barriers and the cooperative Ws dequant stay TG-uniform + // (every segment intersects at least one band, and loop trip counts are + // identical across the TG). + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_lo_lim < m_hi_lim; + NAXTile Dtile; Dtile.clear(); @@ -3531,35 +3541,37 @@ METAL_FUNC void kq_gather_qmm_rhs_nax_tgp_impl( threadgroup_barrier(mem_flags::mem_threadgroup); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + if (sg_active) { + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; - // Prevents the Metal compiler from reordering loads across - // iterations. - volatile int compiler_barrier; + // Prevents the Metal compiler from reordering loads across + // iterations. + volatile int compiler_barrier; - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - (void)compiler_barrier; + (void)compiler_barrier; + } } xn += BK; @@ -3571,53 +3583,57 @@ METAL_FUNC void kq_gather_qmm_rhs_nax_tgp_impl( loader_w.load_safe(tile_w); threadgroup_barrier(mem_flags::mem_threadgroup); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + if (sg_active) { + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; - // Prevents the Metal compiler from reordering loads across - // iterations. - volatile int compiler_barrier; + // Prevents the Metal compiler from reordering loads across + // iterations. + volatile int compiler_barrier; - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - (void)compiler_barrier; + (void)compiler_barrier; + } } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); + if (sg_active) { + if constexpr (kAlignedN.value) { + if (m_lo_lim == 0 && m_hi_lim == SM) { + Dtile.store(y + tm * N + tn, N); + } else { + Dtile.store_slice( + y + tm * N + tn, + N, + short2(0, m_lo_lim), + short2(SN, m_hi_lim)); + } } else { Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); + y + tm * N + tn, + N, + short2(0, m_lo_lim), + short2(sgp_sn, m_hi_lim)); } - } else { - Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); } }); }); From 25164daa310f6a3718363f3b5080162a2e4fc4e4 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:20:21 -0700 Subject: [PATCH 3/3] perf(metal): bm32 gather rhs nax tile dispatched below 64 rows per expert --- CHANGELOG.md | 4 ++++ metal/kq_quantized_nax.metal | 35 +++++++++++++++++++++++++++++++++++ src/kquant_gather.cpp | 18 +++++++++++++++++- 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a5c764c..81b326a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,10 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). fine buffers during deep prefill. ### Changed +- MoE prefill gather (sorted-rhs NAX path) is 12-28% faster per call + below ~64 rows per expert, biggest at 128-529-token chunks on top-8 + 256-expert shapes (bit-identical; KQ_GATHER_RHS_NAX_BM forces the + tile height). - iq2_xxs / iq2_xs / iq2_s / iq3_s MoE gather decode is 9-12% faster per call (hoisted block scale, byte-indexed grids); the ext mat-vec at verify widths 2..8 gains 7-10% on the same codecs. diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index 3e3432e..e5be078 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -207,6 +207,41 @@ instantiate_kquant_nax_bm128(iq2_s, 256, 2) instantiate_kquant_nax_bm128(iq1_s, 256, 1) instantiate_kquant_nax_bm128(iq1_m, 256, 1) +// Small-BM gather_qmm_rhs tile for the few-rows-per-expert prefill regime +// (MoE top-k at chat chunk sizes: rows/E ~ 4-32). Every expert segment in +// a tile re-runs the full-tile MMA K-loop, so segments/tile ~ BM/(rows/E)+1 +// sets the MMA redundancy and BM=32 roughly halves it, for a modest extra +// Ws dequant (expert tiles straddle more row tiles). Same wm2 wn2 -> SM=16, +// TM=1, TN=2 paired-N fragment shape as the proven qmm_t smallbm tile +// (BM=16 is the silently-empty-kernel trap, see above). nt only: the +// sorted-rhs route is transpose-only and the nn TM=1 pairing with tb=false +// is unproven. No float x variant, same reachability argument as the db64 +// macro. +#define instantiate_kquant_nax_gather_rhs_bm32(codec, gs, bits) \ + instantiate_kquant_nax_gather_qmm_rhs( \ + float16_t, gs, bits, true, nt, 32, 64, 2, 2, codec) \ + instantiate_kquant_nax_gather_qmm_rhs( \ + bfloat16_t, gs, bits, true, nt, 32, 64, 2, 2, codec) +instantiate_kquant_nax_gather_rhs_bm32(q6_k, 256, 6) +instantiate_kquant_nax_gather_rhs_bm32(q8_0, 32, 8) +instantiate_kquant_nax_gather_rhs_bm32(q4_k, 256, 4) +instantiate_kquant_nax_gather_rhs_bm32(q5_k, 256, 5) +instantiate_kquant_nax_gather_rhs_bm32(q3_k, 256, 3) +instantiate_kquant_nax_gather_rhs_bm32(q2_k, 256, 2) +instantiate_kquant_nax_gather_rhs_bm32(q5_1, 32, 5) +instantiate_kquant_nax_gather_rhs_bm32(q4_0, 32, 4) +instantiate_kquant_nax_gather_rhs_bm32(q4_1, 32, 4) +instantiate_kquant_nax_gather_rhs_bm32(q5_0, 32, 5) +instantiate_kquant_nax_gather_rhs_bm32(iq4_nl, 32, 4) +instantiate_kquant_nax_gather_rhs_bm32(iq4_xs, 256, 4) +instantiate_kquant_nax_gather_rhs_bm32(iq3_xxs, 256, 3) +instantiate_kquant_nax_gather_rhs_bm32(iq3_s, 256, 3) +instantiate_kquant_nax_gather_rhs_bm32(iq2_xxs, 256, 2) +instantiate_kquant_nax_gather_rhs_bm32(iq2_xs, 256, 2) +instantiate_kquant_nax_gather_rhs_bm32(iq2_s, 256, 2) +instantiate_kquant_nax_gather_rhs_bm32(iq1_s, 256, 1) +instantiate_kquant_nax_gather_rhs_bm32(iq1_m, 256, 1) + instantiate_kquant_nax_smallbm(q6_k, 256, 6) instantiate_kquant_nax_smallbm(q8_0, 32, 8) instantiate_kquant_nax_smallbm(q4_k, 256, 4) diff --git a/src/kquant_gather.cpp b/src/kquant_gather.cpp index 3e5f213..8dfac06 100644 --- a/src/kquant_gather.cpp +++ b/src/kquant_gather.cpp @@ -271,10 +271,25 @@ void gather_qmm_rhs_nax( int M, int N, int K, + int E, Device& d, const Stream& s, const std::string& kquant_type) { - int bm = 64, bn = 64, bk = 64, wm = 2, wn = 2; + int bn = 64, bk = 64, wm = 2, wn = 2; + // Each expert segment in a row tile pays a full-tile MMA K-loop (the + // in-kernel simdgroup skip trims bands the segment misses), so at few + // rows per expert the smaller tile roughly halves the redundant MMA for + // a modest extra Ws dequant. Crossover measured on DSV4-Flash shapes. + const int rows_per_expert = M / std::max(E, 1); + int bm = rows_per_expert < 64 ? 32 : 64; + // Tuning lever: force the NAX rhs tile height (32/64). Read live - only + // reached on the sorted prefill path, so the getenv cost is negligible. + if (const char* e = std::getenv("KQ_GATHER_RHS_NAX_BM")) { + int v = std::atoi(e); + if (v == 32 || v == 64) { + bm = v; + } + } const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -852,6 +867,7 @@ void KQuantGatherQMM::eval_gpu( /*M=*/static_cast(x.size() / K), N, K, + E, d, s, kquant_type_);