From 254675880ce155c0f6183d203bc9a494ee0d2369 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:52:34 -0700 Subject: [PATCH 1/2] feat(kernel): skinny_matmul, GEMV-shaped x @ w.T for token widths 1..16 MLX's steel GEMM leaves the GEMV fast path at M >= 2 and runs small-N large-K shapes far below their bytes (router gates, indexer weight projections, hyper-connection mixes at speculative verify widths: 19-42us per call where the bytes cost 3-8us). One simdgroup per output column streams the w row coalesced holding all M row accumulators in registers, f32 accumulate, one round at the write. 4-8x per call at M in 2..16, parity with the stock GEMV at M=1. CPU eval mirrors the kernel semantics. --- CHANGELOG.md | 5 + CMakeLists.txt | 2 + bindings.cpp | 22 ++ metal/kq_skinny_mv.metal | 11 + .../mlx/backend/metal/kernels/kq_skinny_mv.h | 52 +++++ mlx_kquant/__init__.py | 2 + src/kquant.h | 31 +++ src/kquant_skinny_mv.cpp | 201 ++++++++++++++++++ tests/test_skinny_matmul.py | 79 +++++++ 9 files changed, 405 insertions(+) create mode 100644 metal/kq_skinny_mv.metal create mode 100644 metal/mlx/backend/metal/kernels/kq_skinny_mv.h create mode 100644 src/kquant_skinny_mv.cpp create mode 100644 tests/test_skinny_matmul.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 98e0e8e..9206562 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a9402e..6c82681 100644 --- a/CMakeLists.txt +++ b/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} diff --git a/bindings.cpp b/bindings.cpp index 0b3a2f5..87d6d27 100644 --- a/bindings.cpp +++ b/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, diff --git a/metal/kq_skinny_mv.metal b/metal/kq_skinny_mv.metal new file mode 100644 index 0000000..6c3c385 --- /dev/null +++ b/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 diff --git a/metal/mlx/backend/metal/kernels/kq_skinny_mv.h b/metal/mlx/backend/metal/kernels/kq_skinny_mv.h new file mode 100644 index 0000000..e1c3b93 --- /dev/null +++ b/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 +[[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; + using WT4 = vec; + 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(v); + } + } +} diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index e1c8d5f..133276d 100644 --- a/mlx_kquant/__init__.py +++ b/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", diff --git a/src/kquant.h b/src/kquant.h index 3e223be..42fea2f 100644 --- a/src/kquant.h +++ b/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& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& 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 { diff --git a/src/kquant_skinny_mv.cpp b/src/kquant_skinny_mv.cpp new file mode 100644 index 0000000..ec6f523 --- /dev/null +++ b/src/kquant_skinny_mv.cpp @@ -0,0 +1,201 @@ +// Skinny matmul: y = x @ w.T for token widths 1..16 against small-N, +// large-K nn.Linear-layout weights. MLX's steel GEMM leaves the GEMV fast +// path at M >= 2 and runs these shapes far below their bytes (router +// gates, indexer weight projections, hyper-connection mixes at +// speculative verify widths); the kernel keeps the GEMV shape with all M +// row accumulators in registers (see kq_skinny_mv.h). The CPU eval +// mirrors the kernel's f32-accumulate / one-round-at-write semantics. +#include +#include +#include + +#include "kquant.h" +#include "kquant_internal.h" // kq_type_string + +#include "mlx/backend/cpu/encoder.h" +#include "mlx/ops.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" // kq_get_kernel +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +#ifdef _METAL_ + +void KQuantSkinnyMV::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + auto& out = outputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + + const auto& x = inputs[0]; + const auto& w = inputs[1]; + int K = x.shape(-1); + int M = x.shape(-2); + int N = w.shape(0); + int T = int(x.size() / (int64_t(M) * K)); + + std::string kname = "kq_skinny_mv_" + kq_type_string(x.dtype()) + "_" + + kq_type_string(w.dtype()); + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(x, 0); + ce.set_input_array(w, 1); + ce.set_output_array(out, 2); + ce.set_bytes(M, 3); + ce.set_bytes(N, 4); + ce.set_bytes(K, 5); + MTL::Size group_dims(32 * 8, 1, 1); + MTL::Size grid_dims((N + 7) / 8, T, 1); + ce.dispatch_threadgroups(grid_dims, group_dims); +} + +#else // !_METAL_ + +void KQuantSkinnyMV::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant.skinny_matmul] requires Metal."); +} + +#endif + +namespace { + +// Dispatch a functor over the admitted (x dtype, w dtype) pairs. +template +void skinny_cpu_dispatch(mx::Dtype xt, mx::Dtype wt, F&& run) { + auto with_x = [&](auto* xtag) { + using XT = std::remove_pointer_t; + if (wt == mx::float32) { + run(static_cast(nullptr), static_cast(nullptr)); + } else { + run(static_cast(nullptr), static_cast(nullptr)); + } + }; + if (xt == mx::float16) { + with_x(static_cast(nullptr)); + } else if (xt == mx::bfloat16) { + with_x(static_cast(nullptr)); + } else { + with_x(static_cast(nullptr)); + } +} + +} // namespace + +void KQuantSkinnyMV::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + auto& out = outputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + + const auto& x = inputs[0]; + const auto& w = inputs[1]; + + auto& encoder = mx::cpu::get_command_encoder(stream()); + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_output_array(out); + encoder.dispatch([x = mx::array::unsafe_weak_copy(x), + w = mx::array::unsafe_weak_copy(w), + out = mx::array::unsafe_weak_copy(out)]() mutable { + const int K = x.shape(-1); + const int M = x.shape(-2); + const int N = w.shape(0); + const int64_t T = x.size() / (int64_t(M) * K); + skinny_cpu_dispatch(x.dtype(), w.dtype(), [&](auto* xtag, auto* wtag) { + using XT = std::remove_pointer_t; + using WT = std::remove_pointer_t; + const bool out_f32 = + std::is_same_v || std::is_same_v; + const XT* xp = x.data(); + const WT* wp = w.data(); + for (int64_t t = 0; t < T; t++) { + for (int m = 0; m < M; m++) { + const XT* xrow = xp + (t * M + m) * K; + for (int n = 0; n < N; n++) { + const WT* wrow = wp + int64_t(n) * K; + float acc = 0; + for (int k = 0; k < K; k++) { + acc += static_cast(xrow[k]) * static_cast(wrow[k]); + } + const int64_t o = (t * M + m) * N + n; + if (out_f32) { + out.data()[o] = acc; + } else { + out.data()[o] = static_cast(acc); + } + } + } + } + }); + }); +} + +bool KQuantSkinnyMV::is_equivalent(const mx::Primitive&) const { + return true; +} + +std::vector KQuantSkinnyMV::output_shapes( + const std::vector& inputs) { + auto shape = inputs[0].shape(); + shape.back() = inputs[1].shape(0); + return {shape}; +} + +mx::array skinny_matmul(mx::array x, mx::array w, mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.skinny_matmul]"; + if (x.ndim() < 2) { + throw std::invalid_argument( + std::string(op) + " x must have at least 2 axes."); + } + if (w.ndim() != 2) { + throw std::invalid_argument(std::string(op) + " w must be 2-D [N, K]."); + } + int K = x.shape(-1); + int M = x.shape(-2); + if (M < 1 || M > 16) { + throw std::invalid_argument( + std::string(op) + " x rows (second-to-last axis) must be 1..16."); + } + if (K % 4 != 0 || w.shape(1) != K) { + throw std::invalid_argument( + std::string(op) + " K must match w and be a multiple of 4."); + } + auto xt = x.dtype(); + auto wt = w.dtype(); + if (xt != mx::float16 && xt != mx::bfloat16 && xt != mx::float32) { + throw std::invalid_argument( + std::string(op) + " x must be float16, bfloat16, or float32."); + } + if (wt != xt && wt != mx::float32) { + throw std::invalid_argument( + std::string(op) + " w dtype must match x or be float32."); + } + auto out_dtype = (xt == mx::float32 || wt == mx::float32) ? mx::float32 : xt; + // Unconditional: flags() on a lazy array are not yet valid, so a + // build-time row_contiguous check can skip a needed copy. Contiguous is + // a zero-copy passthrough at eval when the input already is. + auto x_c = mx::contiguous(x, false, s); + auto w_c = mx::contiguous(w, false, s); + + auto shape = x.shape(); + shape.back() = w.shape(0); + return mx::array( + std::move(shape), + out_dtype, + std::make_shared(s), + {std::move(x_c), std::move(w_c)}); +} + +} // namespace mlx_kquant diff --git a/tests/test_skinny_matmul.py b/tests/test_skinny_matmul.py new file mode 100644 index 0000000..22b9b03 --- /dev/null +++ b/tests/test_skinny_matmul.py @@ -0,0 +1,79 @@ +"""skinny_matmul vs an f32 mx.matmul reference. + +The kernel accumulates in f32 and rounds once at the write. The reference +runs on the CPU stream: the GPU f32 GEMM is TF32-by-default on M5-class +devices and would put ~1e-3 of error in the reference itself. +""" + +import mlx.core as mx +import pytest + +import mlx_kquant as kq + +SHAPES = [ + # (M, K, N) + (1, 4096, 256), # router gate, decode width + (4, 4096, 256), # router gate, verify width + (3, 4096, 64), # indexer weights_proj, rows=2 verify + (4, 16384, 24), # hyper-connection mixes + (16, 4096, 4), # width cap, tiny N + (2, 100, 7), # K not a multiple of 128, odd N +] + +COMBOS = [ + (mx.float16, mx.float16), + (mx.bfloat16, mx.bfloat16), + (mx.float32, mx.float32), + (mx.float16, mx.float32), + (mx.bfloat16, mx.float32), +] + + +def _rel(got, ref): + denom = mx.abs(ref).max().item() + 1e-6 + return mx.abs(got.astype(mx.float32) - ref).max().item() / denom + + +@pytest.mark.parametrize("xt,wt", COMBOS) +@pytest.mark.parametrize("M,K,N", SHAPES) +def test_skinny_matmul_parity(xt, wt, M, K, N): + mx.random.seed(3) + x = (mx.random.normal((2, M, K)) * 0.5).astype(xt) + w = (mx.random.normal((N, K)) * 0.5).astype(wt) + got = kq.skinny_matmul(x, w) + with mx.stream(mx.cpu): + ref = x.astype(mx.float32) @ w.astype(mx.float32).T + mx.eval(got, ref) + + expect_dtype = mx.float32 if mx.float32 in (xt, wt) else xt + assert got.dtype == expect_dtype + assert got.shape == (2, M, N) + tol = 3e-5 if expect_dtype == mx.float32 else 5e-3 + assert _rel(got, ref) < tol + + +def test_skinny_matmul_noncontiguous_inputs(): + x = mx.random.normal((4, 4096)).astype(mx.float16).T.swapaxes(0, 1) + w = mx.random.normal((4096, 64)).astype(mx.float16).T + got = kq.skinny_matmul(x[None], w) + with mx.stream(mx.cpu): + ref = x[None].astype(mx.float32) @ w.astype(mx.float32).T + mx.eval(got, ref) + assert _rel(got, ref) < 5e-3 + + +def test_skinny_matmul_rejects(): + x = mx.zeros((17, 4096), dtype=mx.float16) + w = mx.zeros((8, 4096), dtype=mx.float16) + with pytest.raises(ValueError): + kq.skinny_matmul(x, w) # M > 16 + with pytest.raises(ValueError): + kq.skinny_matmul( + mx.zeros((2, 4098), dtype=mx.float16), + mx.zeros((8, 4098), dtype=mx.float16), + ) # K % 4 + with pytest.raises(ValueError): + kq.skinny_matmul( + mx.zeros((2, 4096), dtype=mx.float32), + mx.zeros((8, 4096), dtype=mx.float16), + ) # f32 x with f16 w From 15d62c1e06b0528cca141e7c0dcb4a64cebf5d88 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:53:34 -0700 Subject: [PATCH 2/2] fix(norm): lazy non-contiguous inputs to the fused norm ops computed garbage The factories short-circuited the Contiguous node on a build-time row_contiguous flags check, but flags of an unevaluated array are not meaningful yet (a transposed lazy view reports contiguous), so the kernels walked the wrong strides. Unconditional Contiguous node, which aliases zero-copy at eval when the input is already packed. Same class as the gather_qmm fix; moe_glu/route_shed sites still use the ternary on activation inputs that are contiguous products today. --- src/kquant_norm_fused.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/kquant_norm_fused.cpp b/src/kquant_norm_fused.cpp index 820373c..3689ac5 100644 --- a/src/kquant_norm_fused.cpp +++ b/src/kquant_norm_fused.cpp @@ -27,7 +27,10 @@ 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(); @@ -35,7 +38,7 @@ prep_act(const mx::array& x, const char* op, const char* what, mx::Stream s) { 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(),