feat(kernel): skinny_matmul, dense small-N matmul for speculative verify widths

CHANGELOG.md

@@ -6,6 +6,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Added
+- `skinny_matmul`: x @ w.T at token widths 1..16 against small-N large-K
+  nn.Linear-layout weights, 4-8x faster than the stock GEMM at widths
+  2..16 (router gates, indexer projections at speculative verify widths).
+
 ## [0.3.9]
 
 ### Changed

CMakeLists.txt

@@ -128,6 +128,7 @@ target_sources(
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_moe_glu.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_norm_fused.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_route_shed.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_skinny_mv.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_gather.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_encode.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_arena.cpp
@@ -209,6 +210,7 @@ if(MLX_BUILD_METAL)
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu_kq.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_norm_fused.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_route_shed.metal
+    ${CMAKE_CURRENT_LIST_DIR}/metal/kq_skinny_mv.metal
     INCLUDE_DIRS
     ${CMAKE_CURRENT_LIST_DIR}/metal
     ${MLX_INCLUDE_DIRS}

bindings.cpp

@@ -1151,6 +1151,28 @@ NB_MODULE(_ext, m) {
             array: same shape and dtype as a.
       )");
 
+  m.def(
+      "skinny_matmul",
+      &mlx_kquant::skinny_matmul,
+      "x"_a,
+      "w"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        y = x @ w.T at token widths 1..16 against a small-N, large-K
+        weight in nn.Linear layout, f32 accumulate. Fills the GEMV-to-GEMM
+        cliff MLX hits at M >= 2 on these shapes.
+
+        Args:
+            x (array): [..., M, K], 1 <= M <= 16, K a multiple of 4;
+                float16/bfloat16/float32.
+            w (array): [N, K] weight, dtype matching x or float32.
+
+        Returns:
+            array: [..., M, N]; float32 when either operand is, else the
+            x dtype.
+      )");
+
   m.def(
       "route_shed",
       &mlx_kquant::route_shed,

metal/kq_skinny_mv.metal

@@ -0,0 +1,11 @@
+// clang-format off
+// Skinny matmul kernel instantiations; see kq_skinny_mv.h.
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/kq_skinny_mv.h"
+
+instantiate_kernel("kq_skinny_mv_float16_t_float16_t", kq_skinny_mv, float16_t, float16_t, float16_t)
+instantiate_kernel("kq_skinny_mv_bfloat16_t_bfloat16_t", kq_skinny_mv, bfloat16_t, bfloat16_t, bfloat16_t)
+instantiate_kernel("kq_skinny_mv_float_float", kq_skinny_mv, float, float, float)
+instantiate_kernel("kq_skinny_mv_float16_t_float", kq_skinny_mv, float16_t, float, float)
+instantiate_kernel("kq_skinny_mv_bfloat16_t_float", kq_skinny_mv, bfloat16_t, float, float)
+    // clang-format on

metal/mlx/backend/metal/kernels/kq_skinny_mv.h

@@ -0,0 +1,52 @@
+// Skinny matmul: y = x @ w.T for token widths 1..16 against a small-N,
+// large-K weight (w is [N, K], nn.Linear layout). MLX's steel GEMM leaves
+// the GEMV fast path at M >= 2 and runs these shapes 3-10x slower than
+// their bytes (router gates, indexer weight projections, hyper-connection
+// mixes at speculative verify widths); this kernel keeps the GEMV shape:
+// one simdgroup per output column streams the w row coalesced (float4 per
+// lane) while holding all M row accumulators in registers, f32 accumulate,
+// one round to OT at the write.
+//
+// Grid: (ceil(N / KQ_SKINNY_NSG), T, 1) threadgroups of 32 * KQ_SKINNY_NSG
+// threads; T collapses the batch dims ahead of the [M, K] tail. Requires
+// K % 4 == 0 and M <= KQ_SKINNY_MMAX.
+
+#define KQ_SKINNY_NSG 8
+#define KQ_SKINNY_MMAX 16
+
+template <typename XT, typename WT, typename OT>
+[[kernel]] void kq_skinny_mv(
+    const device XT* x [[buffer(0)]],
+    const device WT* w [[buffer(1)]],
+    device OT* out [[buffer(2)]],
+    const constant int& M [[buffer(3)]],
+    const constant int& N [[buffer(4)]],
+    const constant int& K [[buffer(5)]],
+    uint2 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  const int n = tid.x * KQ_SKINNY_NSG + simd_gid;
+  if (n >= N) {
+    return;
+  }
+  using XT4 = vec<XT, 4>;
+  using WT4 = vec<WT, 4>;
+  const int K4 = K / 4;
+  const device WT4* wrow = (const device WT4*)(w + (int64_t)n * K);
+  const device XT4* xrows = (const device XT4*)(x + (int64_t)tid.y * M * K);
+
+  float acc[KQ_SKINNY_MMAX] = {0};
+  for (int k4 = simd_lid; k4 < K4; k4 += 32) {
+    const float4 wv = float4(wrow[k4]);
+    for (int m = 0; m < M; m++) {
+      acc[m] += metal::dot(float4(xrows[m * K4 + k4]), wv);
+    }
+  }
+  device OT* orow = out + (int64_t)tid.y * M * N;
+  for (int m = 0; m < M; m++) {
+    const float v = simd_sum(acc[m]);
+    if (simd_lid == 0) {
+      orow[m * N + n] = static_cast<OT>(v);
+    }
+  }
+}

mlx_kquant/__init__.py

@@ -76,6 +76,7 @@
     shared_event_read,
     shared_event_set,
     shared_event_wait,
+    skinny_matmul,
     verify_zero_copy_views,
     zero_copy_view_count,
 )
@@ -129,6 +130,7 @@
     "sdpa_decode_gqa_paged",
     "sdpa_fa_verify",
     "sdpa_vector",
+    "skinny_matmul",
     "residency_commit",
     "residency_erase",
     "residency_insert",

src/kquant.h

@@ -611,6 +611,15 @@ mx::array rmsnorm2_add(
     float eps,
     mx::StreamOrDevice s = {});
 
+// y = x @ w.T for token widths 1..16 against a small-N, large-K weight in
+// nn.Linear layout ([N, K]). x is [..., M, K] with 1 <= M <= 16 and
+// K % 4 == 0; x float16/bfloat16/float32, w matching x or float32; output
+// float32 when either operand is, else x dtype. f32 accumulate. Fills the
+// GEMV-to-GEMM cliff MLX's steel path hits at M >= 2 on these shapes
+// (router gates, indexer weight projections, hyper-connection mixes at
+// speculative verify widths).
+mx::array skinny_matmul(mx::array x, mx::array w, mx::StreamOrDevice s = {});
+
 // GPU-side routed-expert slot remap + residency shed for streamed MoE decode
 // (the gpu-dispatch autonomous-token front end; see kq_route_shed.h).
 // slot_table maps expert id -> arena slot, negative = non-resident. Every
@@ -1540,6 +1549,28 @@ class KQuantRMSNorm2Add : public mx::Primitive {
   float eps_;
 };
 
+// Skinny matmul y = x @ w.T at token widths 1..16 (see skinny_matmul).
+// Inference-only.
+class KQuantSkinnyMV : public mx::Primitive {
+ public:
+  explicit KQuantSkinnyMV(mx::Stream stream) : mx::Primitive(stream) {}
+
+  void eval_cpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+  void eval_gpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+
+  std::vector<mx::Shape> output_shapes(
+      const std::vector<mx::array>& inputs) override;
+
+  const char* name() const override {
+    return "KQuantSkinnyMV";
+  }
+  bool is_equivalent(const mx::Primitive& other) const override;
+};
+
 // Routed-expert slot remap + residency shed (see route_shed).
 // Inference-only.
 class KQuantRouteShed : public mx::Primitive {

src/kquant_norm_fused.cpp

@@ -27,15 +27,18 @@ namespace mlx_kquant {
 
 namespace {
 
-// Row-contiguous activation with a supported dtype, or throw.
+// Row-contiguous activation with a supported dtype, or throw. The
+// Contiguous node is unconditional: flags of an unevaluated array are not
+// meaningful at op-build time (see gather_qmm), and Contiguous::eval
+// aliases zero-copy when the input turns out packed.
 mx::array
 prep_act(const mx::array& x, const char* op, const char* what, mx::Stream s) {
   auto dt = x.dtype();
   if (dt != mx::float16 && dt != mx::bfloat16) {
     throw std::invalid_argument(
         std::string(op) + " " + what + " must be float16 or bfloat16.");
   }
-  return x.flags().row_contiguous ? x : mx::contiguous(x, false, s);
+  return mx::contiguous(x, false, s);
 }
 
 // 1-D [D] weight matching the activation dtype, or throw.
@@ -55,7 +58,7 @@ mx::array prep_norm_weight(
     throw std::invalid_argument(
         std::string(op) + " " + what + " dtype must match the activations.");
   }
-  return w.flags().row_contiguous ? w : mx::contiguous(w, false, s);
+  return mx::contiguous(w, false, s);
 }
 
 } // namespace
@@ -423,8 +426,7 @@ mx::array add_rmsnorm(
       throw std::invalid_argument(
           std::string(op) + " scale dtype must match the activations.");
     }
-    inputs.push_back(
-        sc.flags().row_contiguous ? sc : mx::contiguous(sc, false, s));
+    inputs.push_back(mx::contiguous(sc, false, s));
   }
   return mx::array(
       h.shape(),