perf(metal): faster moe prefill gather below 64 rows per expert

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.

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)

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<KqIq2_xxsExt>::deq_chunk16s(src, short(2 * sb), r0, s0);
+    KqExtDeq<KqIq2_xxsExt>::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]);
     }
   }
 
@@ -3508,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<AccumType, TM, TN> Dtile;
     Dtile.clear();
 
@@ -3534,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<T, TM, TK> Atile;
-            NAXTile<T, BR, BC> Btile;
+          if (sg_active) {
+            STEEL_PRAGMA_NO_UNROLL
+            for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+              NAXTile<T, TM, TK> Atile;
+              NAXTile<T, BR, BC> 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<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
-            } else {
-              Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * BN_padded);
-            }
+              if constexpr (transpose) {
+                Btile.template load<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
+              } else {
+                Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * BN_padded);
+              }
 
-            tile_matmad_nax(
-                Dtile,
-                Atile,
-                metal::bool_constant<false>{},
-                Btile,
-                metal::bool_constant<transpose>{});
+              tile_matmad_nax(
+                  Dtile,
+                  Atile,
+                  metal::bool_constant<false>{},
+                  Btile,
+                  metal::bool_constant<transpose>{});
 
-            (void)compiler_barrier;
+              (void)compiler_barrier;
+            }
           }
 
           xn += BK;
@@ -3574,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<T, TM, TK> Atile;
-            NAXTile<T, BR, BC> Btile;
+          if (sg_active) {
+            STEEL_PRAGMA_NO_UNROLL
+            for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+              NAXTile<T, TM, TK> Atile;
+              NAXTile<T, BR, BC> 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<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
-            } else {
-              Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * BN_padded);
-            }
+              if constexpr (transpose) {
+                Btile.template load<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
+              } else {
+                Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * BN_padded);
+              }
 
-            tile_matmad_nax(
-                Dtile,
-                Atile,
-                metal::bool_constant<false>{},
-                Btile,
-                metal::bool_constant<transpose>{});
+              tile_matmad_nax(
+                  Dtile,
+                  Atile,
+                  metal::bool_constant<false>{},
+                  Btile,
+                  metal::bool_constant<transpose>{});
 
-            (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));
         }
       });
     });

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<int>(x.size() / K),
           N,
           K,
+          E,
           d,
           s,
           kquant_type_);