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_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(), 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