From d96d5362b5267efaf1f9fabaff1e4b2fb20b275e Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 02:08:27 -0700 Subject: [PATCH 01/10] feat(matmul): quantized_matmul_multi runs weights that share x in one verify_nax dispatch --- bindings.cpp | 33 +++ metal/kq_quantized_nax.metal | 37 +++ .../mlx/backend/metal/kernels/kq_verify_nax.h | 69 ++++- mlx_kquant/__init__.py | 2 + src/kquant.h | 61 +++++ src/kquant_matmul.cpp | 205 ++++++++++++++ src/kquant_matmul_fused.cpp | 177 ++++++++++++ tests/test_matmul_multi.py | 258 ++++++++++++++++++ 8 files changed, 838 insertions(+), 4 deletions(-) create mode 100644 tests/test_matmul_multi.py diff --git a/bindings.cpp b/bindings.cpp index 260a009..edf8af9 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -666,6 +666,39 @@ NB_MODULE(_ext, m) { the input dtype. )"); + m.def( + "quantized_matmul_multi", + &mlx_kquant::quantized_matmul_multi, + "x"_a, + "weights"_a, + "kquant_type"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + Quantized matmuls of several weights that share x: + [x @ W_i^T for W_i in weights]. + + On a NAX GPU, at 2 to 8 rows of x where quantized_matmul would run + its verify_nax route on the widest weight, one dispatch per four + weights computes the products. Each weight is read from its own + buffer and each product is written to its own output, so the + weights need no joining copy. One split-K count covers the joint + grid. Every other case runs quantized_matmul per weight. + KQ_MATMUL_MULTI=0 runs quantized_matmul per weight on every call. + mx.grad gives the gradient with respect to x; the weights are + frozen, as in quantized_matmul. + + Args: + x (array): [..., K], float16/bfloat16. float32 runs in bfloat16. + weights (list[array]): 2-D uint8 wire bytes [N_i, K bytes] of + one codec. + kquant_type (str): the codec of every weight. + + Returns: + list[array]: [..., N_i] per weight, in the dtype + quantized_matmul gives for x. + )"); + m.def( "quantized_matmul_glu", &mlx_kquant::quantized_matmul_glu, diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index 1378f89..9994f5d 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -443,4 +443,41 @@ instantiate_kquant_verify_nax(bfloat16_t, 256, 3, q3_k) instantiate_kquant_verify_nax(float16_t, 256, 3, q3_k) instantiate_kquant_verify_nax(bfloat16_t, 256, 2, q2_k) instantiate_kquant_verify_nax(float16_t, 256, 2, q2_k) + +// Several weights that share x in one dispatch (kq_verify_nax.h). +#define instantiate_kquant_verify_nax_multi(type, gs, bits, codec) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_nax_multi_" #type "_gs_" #gs "_b_" \ + #bits, \ + kq_ ## codec ## _verify_nax_multi, \ + type, gs, bits) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_nax_multi_" #type "_gs_" #gs "_b_" \ + #bits "_pf32", \ + kq_ ## codec ## _verify_nax_multi, \ + type, gs, bits, float) +instantiate_kquant_verify_nax_multi(bfloat16_t, 128, 2, pq2_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 64, 2, q2_0) +instantiate_kquant_verify_nax_multi(float16_t, 128, 2, pq2_0) +instantiate_kquant_verify_nax_multi(float16_t, 64, 2, q2_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 4, q4_0) +instantiate_kquant_verify_nax_multi(float16_t, 32, 4, q4_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 4, q4_0_sb) +instantiate_kquant_verify_nax_multi(float16_t, 32, 4, q4_0_sb) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 8, q8_0) +instantiate_kquant_verify_nax_multi(float16_t, 32, 8, q8_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 4, q4_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 4, q4_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 4, iq4_xs) +instantiate_kquant_verify_nax_multi(float16_t, 256, 4, iq4_xs) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 4, iq4_nl) +instantiate_kquant_verify_nax_multi(float16_t, 32, 4, iq4_nl) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 5, q5_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 5, q5_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 6, q6_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 6, q6_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 3, q3_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 3, q3_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 2, q2_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 2, q2_k) // clang-format on diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index dc32757..b95fe0d 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -800,10 +800,10 @@ METAL_FUNC void kq_verify_nax_impl( const device uint8_t* w, const device T* x, device U* y, - const constant int& K, - const constant int& N, - const constant int& M, - const constant int& k_partition_size, + const int K, + const int N, + const int M, + const int k_partition_size, uint3 tid, uint simd_gid, uint simd_lid) { @@ -944,6 +944,54 @@ METAL_FUNC void kq_verify_nax_impl( simd_lid); \ } +// Several weights that share x in one dispatch: weight i owns the +// threadgroups [tg_end[i - 1], tg_end[i]) of grid x, each weight from its +// own buffer and into its own output, with one split count over grid z. +// part_stride[i] is output i's partial stride. Unused weight slots repeat +// slot 0 and own no threadgroup. +#define KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(CODEC, TRAITS) \ + template \ + [[kernel]] void kq_##CODEC##_verify_nax_multi( \ + const device uint8_t* w0, \ + const device uint8_t* w1, \ + const device uint8_t* w2, \ + const device uint8_t* w3, \ + const device T* x, \ + device U* y0, \ + device U* y1, \ + device U* y2, \ + device U* y3, \ + const constant int& K, \ + const constant int& M, \ + const constant int& k_partition_size, \ + const constant int4& n_rows, \ + const constant int4& tg_end, \ + const constant int4& part_stride, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + static_assert(group_size == TRAITS::group, #CODEC " block width"); \ + const int t = int(tid.x); \ + const short i = t < tg_end[0] ? 0 \ + : t < tg_end[1] ? 1 \ + : t < tg_end[2] ? 2 \ + : 3; \ + const int t0 = i == 0 ? 0 : tg_end[i - 1]; \ + const device uint8_t* w = i == 0 ? w0 : i == 1 ? w1 : i == 2 ? w2 : w3; \ + device U* y = i == 0 ? y0 : i == 1 ? y1 : i == 2 ? y2 : y3; \ + kq_verify_nax_impl( \ + w, \ + x, \ + y + tid.z * static_cast(part_stride[i]), \ + K, \ + n_rows[i], \ + M, \ + k_partition_size, \ + uint3(uint(t - t0), tid.y, tid.z), \ + simd_gid, \ + simd_lid); \ + } + KQ_DEFINE_VERIFY_NAX_KERNEL(pq2_0, KqPq2_0Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q2_0, KqQ2_0Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q4_0, KqQ4_0Nax) @@ -956,3 +1004,16 @@ KQ_DEFINE_VERIFY_NAX_KERNEL(q5_k, KqQ5_KNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q6_k, KqQ6_KNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q3_k, KqQ3_KNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q2_k, KqQ2_KNax) + +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(pq2_0, KqPq2_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q2_0, KqQ2_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q4_0, KqQ4_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q4_0_sb, KqQ4_0SbNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q8_0, KqQ8_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q4_k, KqQ4_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(iq4_xs, KqIq4XsNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(iq4_nl, KqIq4NlNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q5_k, KqQ5_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q6_k, KqQ6_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q3_k, KqQ3_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q2_k, KqQ2_KNax) diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index c757ce5..875665e 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -87,6 +87,7 @@ quantized_matmul, quantized_matmul_glu, quantized_matmul_lora, + quantized_matmul_multi, quantized_matmul_qmv_bias, quantized_matmul_residual, residency_commit, @@ -195,6 +196,7 @@ "quantize", "quantized_matmul", "quantized_matmul_glu", + "quantized_matmul_multi", "quantized_matmul_lora", "quantized_matmul_qmv_bias", "quantized_matmul_residual", diff --git a/src/kquant.h b/src/kquant.h index 27d53f0..7379f7b 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -1706,6 +1706,67 @@ class KQuantMatmulFused : public mx::Primitive { int act_; // 0 = silu, 1 = gelu (tanh form) }; +// Quantized matmuls of several weights that share x: y_i = x @ W_i^T. +// x [..., K] (float16/bfloat16; float32 runs in bfloat16), weights 2-D +// uint8 wire bytes of one codec, [N_i, K bytes] each, in separate buffers. +// Returns one [..., N_i] array per weight, in the dtype quantized_matmul +// gives for x. On a NAX GPU, at 2 to 8 rows of x where quantized_matmul +// would run its verify_nax route on the widest weight, one dispatch per +// four weights computes the products: each weight is read in place, each +// output is written in its own buffer, and one split-K count covers the +// joint grid. Every other case runs quantized_matmul per weight. +std::vector quantized_matmul_multi( + mx::array x, + std::vector weights, + const std::string& kquant_type, + mx::StreamOrDevice s = {}); + +// Primitive of quantized_matmul_multi: inputs {x, scales, w_0, ...}, +// one output per weight. vjp and jvp are those of the per-weight +// matmuls, wrt x only. +class KQuantMatmulMulti : public mx::Primitive { + public: + KQuantMatmulMulti( + mx::Stream stream, + std::string kquant_type, + int group_size, + int bits) + : mx::Primitive(stream), + kquant_type_(std::move(kquant_type)), + group_size_(group_size), + bits_(bits) {} + + 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 "KQuantMatmulMulti"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + std::vector jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) override; + + private: + std::string kquant_type_; + int group_size_; + int bits_; +}; + // Quantized matmul with a low-rank term: // y = x @ W^T + (strength * (x @ A)) @ B // x [..., K] (float16/bfloat16; float32 runs in bfloat16), w 2-D uint8 wire diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index f6ab404..4c14b1e 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -10,6 +10,7 @@ // decode-latency-critical path); the decode-only bias-fused fast path lives in // the KQuantQmvBias primitive below (qmv_bias), which reuses this file's qmv // dispatch helpers. +#include #include #include #include @@ -1432,6 +1433,29 @@ void verify_mma( ce.dispatch_threads(agrid, agroup); } +// The several-weight form of kq_verify_nax_kname. +static std::string kq_verify_nax_multi_kname( + const std::string& t, + const std::string& ts, + int gs, + int bits, + int kstep, + const std::string& suffix) { + std::string kname; + kname.reserve(72); + mx::concatenate( + kname, + (t == "q4_0" && kstep == 256 ? "kquant_q4_0_sb_" : kq_kname_prefix(t)) + + "verify_nax_multi_", + ts, + "_gs_", + gs, + "_b_", + bits, + suffix); + return kname; +} + // K per verify_nax iteration for this call. q4_0 runs its eight-block // kernel (q4_0_sb, 256 weights per step) where K is a multiple of 256, // and its two-block kernel (64) elsewhere. KQ_VERIFY_NAX_Q4_0_SB=0 keeps @@ -1539,6 +1563,106 @@ void verify_nax( ce.dispatch_threads(agrid, agroup); } +static int kq_verify_nax_splits(int N, int K, int kstep, const std::string& t); + +// verify_nax over up to four weights that share x, in one dispatch +// (kq_verify_nax_multi): weight i takes ceil(N_i / 128) threadgroups of +// grid x after those of the weights before it, and reads its own buffer. +// The split count is the one verify_nax takes for the joint grid. Split +// partials of weight i sit in one temporary at splits * M * (rows of the +// weights before it), and one fold per weight writes its output. +void verify_nax_multi( + const array& x, + const std::vector& ws, + std::vector& outs, + int group_size, + int bits, + int M, + int K, + int kstep, + Device& d, + const Stream& s, + const std::string& kquant_type) { + const int n = static_cast(ws.size()); + std::array n_rows{0, 0, 0, 0}; + std::array tg_end{0, 0, 0, 0}; + std::array part_stride{0, 0, 0, 0}; + std::array row_off{0, 0, 0, 0}; + int tgs = 0; + int64_t rows = 0; + for (int i = 0; i < 4; ++i) { + const int N = i < n ? ws[i].shape(0) : 0; + n_rows[i] = N; + row_off[i] = rows; + rows += N; + tgs += (N + 127) / 128; + tg_end[i] = tgs; + part_stride[i] = M * N; + } + const int splits = kq_verify_nax_splits(tgs * 128, K, kstep, kquant_type); + const int k_partition = (K / kstep / splits) * kstep; + + auto& ce = mx::metal::get_command_encoder(s); + const KqSplitkPartials pp = splits > 1 ? kq_splitk_partials(x.dtype()) + : KqSplitkPartials{x.dtype(), ""}; + array partials({splits, M, static_cast(rows)}, pp.dtype, nullptr, {}); + if (splits > 1) { + partials.set_data(mx::allocator::malloc(partials.nbytes())); + ce.add_temporary(partials); + } + + const std::string type_string = kq_type_string(x.dtype()); + const std::string kname = kq_verify_nax_multi_kname( + kquant_type, type_string, group_size, bits, kstep, pp.suffix); + auto kernel = kq_get_kernel(d, kname); + ce.set_compute_pipeline_state(kernel); + + int c = 0; + for (int i = 0; i < 4; ++i) { + ce.set_input_array(ws[i < n ? i : 0], c++); + } + ce.set_input_array(x, c++); + for (int i = 0; i < 4; ++i) { + const int j = i < n ? i : 0; + if (splits > 1) { + ce.set_output_array( + partials, c++, splits * M * row_off[j] * partials.itemsize()); + } else { + ce.set_output_array(outs[j], c++); + } + } + ce.set_bytes(K, c++); + ce.set_bytes(M, c++); + ce.set_bytes(k_partition, c++); + ce.set_bytes(n_rows.data(), 4, c++); + ce.set_bytes(tg_end.data(), 4, c++); + ce.set_bytes(part_stride.data(), 4, c++); + MTL::Size group_dims(128, 1, 1); + MTL::Size grid_dims(tgs, 1, splits); + kq_log_qmm(kname, M, static_cast(rows), K, grid_dims); + ce.dispatch_threadgroups(grid_dims, group_dims); + if (splits == 1) { + return; + } + + const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); + auto accum = kq_get_kernel(d, aname); + for (int i = 0; i < n; ++i) { + ce.set_compute_pipeline_state(accum); + const int n_elems = M * n_rows[i]; + c = 0; + ce.set_input_array( + partials, c++, splits * M * row_off[i] * partials.itemsize()); + ce.set_output_array(outs[i], c++); + ce.set_bytes(n_elems, c++); + ce.set_bytes(splits, c++); + ce.set_bytes(part_stride[i], c++); + MTL::Size agrid(static_cast(n_elems), 1, 1); + MTL::Size agroup(256, 1, 1); + ce.dispatch_threads(agrid, agroup); + } +} + // Split-K qmm_t on the NAX BM=32 or BM=64 tile (KQ_QMM_SPLITK_NAX). Same // partial/fold shape as qmm_splitk, but slices run the tensor-core tile: // the steel splitk probe measured per-TG pipeline bound (~140-160 GB/s flat @@ -2596,6 +2720,79 @@ std::vector KQuantMatmul::vjp( #ifdef _METAL_ +// One verify_nax dispatch per four weights where the widest weight would +// take the verify_nax route at this M, and quantized_matmul per weight +// otherwise. The route levers of quantized_matmul apply: KQ_NAX_QMV, +// KQ_VERIFY_NAX and KQ_DISABLE_NAX through kq_is_nax_available. +void KQuantMatmulMulti::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + const auto& x = inputs[0]; + const auto& scales = inputs[1]; + const int n = static_cast(outputs.size()); + std::vector ws; + ws.reserve(n); + for (int i = 0; i < n; ++i) { + ws.push_back(kq_matrix_dense_gpu(inputs[2 + i], s)); + kq_check_weight_base(ws.back(), kquant_type_, "quantized_matmul_multi"); + } + const int K = x.shape(-1); + const int M = K > 0 ? static_cast(x.size() / K) : 0; + int n_max = 0; + for (const auto& w : ws) { + n_max = std::max(n_max, w.shape(0)); + } + const int kstep = kq_verify_nax_call_kstep(K, kquant_type_); + bool multi = M >= 2 && M <= 8 && n_max > 0 && x.flags().row_contiguous && + x.dtype() != mx::float32 && kq_is_nax_available() && + codec_has_nax(kquant_type_) && kstep > 0; + if (multi) { + int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); + if (qmv_max_m == -1) { + qmv_max_m = kq_nax_qmv_max_m(kquant_type_, n_max); + } + int vnax_min_m = kq_env_int("KQ_VERIFY_NAX", -1); + if (vnax_min_m == -1) { + vnax_min_m = + kq_verify_nax_min_m(kquant_type_, n_max, K, x.dtype() == mx::float16); + } + multi = M > qmv_max_m && vnax_min_m > 0 && M >= vnax_min_m; + for (const auto& w : ws) { + multi = multi && kq_verify_nax_fits(w, K, kstep, kquant_type_); + } + } + if (!multi) { + KQuantMatmul one(s, kquant_type_, group_size_, bits_, true); + for (int i = 0; i < n; ++i) { + std::vector o{outputs[i]}; + one.eval_gpu({x, ws[i], scales}, o); + } + return; + } + std::vector gw; + std::vector go; + for (int i = 0; i < n; ++i) { + outputs[i].set_data(mx::allocator::malloc(outputs[i].nbytes())); + if (outputs[i].size() == 0) { + continue; + } + gw.push_back(ws[i]); + go.push_back(outputs[i]); + if (gw.size() == 4 || i == n - 1) { + verify_nax_multi( + x, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); + gw.clear(); + go.clear(); + } + } + if (!gw.empty()) { + verify_nax_multi( + x, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); + } +} + void KQuantMatmul::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -3126,6 +3323,13 @@ void KQuantMatmul::eval_gpu( "[mlx_kquant] quantized_matmul has no GPU implementation."); } +void KQuantMatmulMulti::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant] quantized_matmul_multi has no GPU implementation."); +} + #endif // _METAL_ std::vector KQuantQmvBias::output_shapes( @@ -3328,6 +3532,7 @@ std::vector split_kernel_names() { for (int k : ksteps) { each_suffix(dt, [&](const std::string& sfx) { out.push_back(kq_verify_nax_kname(t, ts, gs, bits, k, sfx)); + out.push_back(kq_verify_nax_multi_kname(t, ts, gs, bits, k, sfx)); }); } } diff --git a/src/kquant_matmul_fused.cpp b/src/kquant_matmul_fused.cpp index 565aac6..b58464f 100644 --- a/src/kquant_matmul_fused.cpp +++ b/src/kquant_matmul_fused.cpp @@ -461,4 +461,181 @@ mx::array quantized_matmul_glu( {std::move(x_c), std::move(w_gate), std::move(w_up)}); } +// ----------------------------- several weights, one x ----------------- + +namespace { + +const char* kMultiOp = "[mlx_kquant.quantized_matmul_multi]"; + +// KQ_MATMUL_MULTI=0 runs quantized_matmul per weight on every call. Read +// live so an A/B can flip arms in one process. +bool multi_on() { + return kq_env_flag("KQ_MATMUL_MULTI", 1) != 0; +} + +std::vector multi_composed( + const mx::array& x, + const std::vector& weights, + const std::string& kquant_type, + mx::Stream s) { + auto scales = mx::zeros({1}, mx::uint8, s); + std::vector out; + out.reserve(weights.size()); + for (const auto& w : weights) { + out.push_back(quantized_matmul( + x, + w, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s)); + } + return out; +} + +} // namespace + +void KQuantMatmulMulti::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + std::string(kMultiOp) + " CPU streams build the per-weight ops."); +} + +std::vector KQuantMatmulMulti::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector&) { + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul_multi vjp: no gradient wrt the " + "quantized weights (the kquant base is frozen)."); + } + } + std::vector ws(primals.begin() + 2, primals.end()); + auto fun = [&](const std::vector& in) { + return multi_composed(in[0], ws, kquant_type_, stream()); + }; + auto vjps = + mx::vjp(fun, std::vector{primals[0]}, cotangents).second; + return std::vector(argnums.size(), vjps[0]); +} + +std::vector KQuantMatmulMulti::jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) { + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul_multi jvp: no derivative wrt the " + "quantized weights (the kquant base is frozen)."); + } + } + std::vector ws(primals.begin() + 2, primals.end()); + auto fun = [&](const std::vector& in) { + return multi_composed(in[0], ws, kquant_type_, stream()); + }; + return mx::jvp( + fun, + std::vector{primals[0]}, + std::vector{tangents[0]}) + .second; +} + +bool KQuantMatmulMulti::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return kquant_type_ == o.kquant_type_; +} + +std::vector KQuantMatmulMulti::output_shapes( + const std::vector& inputs) { + std::vector out; + for (size_t i = 2; i < inputs.size(); ++i) { + auto shape = inputs[0].shape(); + shape.back() = inputs[i].shape(0); + out.push_back(std::move(shape)); + } + return out; +} + +std::vector quantized_matmul_multi( + mx::array x, + std::vector weights, + const std::string& kquant_type, + mx::StreamOrDevice s_) { + const std::string op = kMultiOp; + const KQuantCodec* codec = codec_by_name(kquant_type); + if (codec == nullptr) { + throw std::invalid_argument( + op + " Unknown kquant_type: '" + kquant_type + "'."); + } + if (weights.empty()) { + throw std::invalid_argument(op + " weights must hold at least one weight."); + } + for (const auto& w : weights) { + if (w.dtype() != mx::uint8 || w.ndim() != 2) { + throw std::invalid_argument( + op + " weights must be 2-D uint8 (raw GGUF wire bytes)."); + } + if (w.shape(-1) != weights[0].shape(-1)) { + throw std::invalid_argument( + op + " weights must have one row length in bytes."); + } + } + if (weights[0].shape(-1) % codec->bytes_per_block != 0) { + throw std::invalid_argument( + op + " weight last dim (" + std::to_string(weights[0].shape(-1)) + + " bytes) is not a whole number of " + + std::to_string(codec->bytes_per_block) + "-byte " + codec->name + + " blocks."); + } + const int K = (weights[0].shape(-1) / codec->bytes_per_block) * + codec->weights_per_block; + if (x.ndim() < 1 || x.shape(-1) != K) { + throw std::invalid_argument( + op + " x last dim does not match the weight inner dim (" + + std::to_string(K) + ")."); + } + if (!mx::issubdtype(x.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " Only real floating x dtypes are supported."); + } + auto s = mx::to_stream(s_); + const mx::Dtype out_type = + x.dtype() == mx::float32 ? mx::bfloat16 : x.dtype(); + const int64_t rows = K > 0 ? static_cast(x.size()) / K : 0; + const bool fused = s.device == mx::Device::gpu && nax_available() && + codec->has_matmul_kernel && multi_on() && weights.size() >= 2 && + rows >= 2 && rows <= 8; + if (!fused) { + return multi_composed(x, weights, kquant_type, s); + } + std::vector inputs; + inputs.reserve(weights.size() + 2); + inputs.push_back(mx::contiguous(mx::astype(x, out_type, s), false, s)); + inputs.push_back(mx::zeros({1}, mx::uint8, s)); + std::vector shapes; + std::vector dtypes; + for (auto& w : weights) { + kq_check_weight_base_at_build(w, kquant_type, s, "quantized_matmul_multi"); + auto shape = x.shape(); + shape.back() = w.shape(0); + shapes.push_back(std::move(shape)); + dtypes.push_back(out_type); + inputs.push_back(kq_matrix_dense(w, s)); + } + return mx::array::make_arrays( + std::move(shapes), + dtypes, + std::make_shared( + s, kquant_type, codec->weights_per_block, codec->bits), + std::move(inputs)); +} + } // namespace mlx_kquant diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py new file mode 100644 index 0000000..ca4ff29 --- /dev/null +++ b/tests/test_matmul_multi.py @@ -0,0 +1,258 @@ +"""quantized_matmul_multi against quantized_matmul per weight. + +On a NAX GPU the op runs one verify_nax dispatch per four weights at 2 to +8 rows, where quantized_matmul would run verify_nax on the widest weight. +With the split count forced to one value, the verify_nax entry at M 2 and +the per-row qmv route off, quantized_matmul per weight runs the same +kernel body and fold, so the outputs match bit for bit. At the default +levers the joint grid can take another split count than each weight +alone, so the outputs are checked +against a float64 reference built from kq.dequantize weights. Every other +case runs quantized_matmul per weight, which the CPU stream checks. +""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from qmm_probe import force_cpu + +import mlx_kquant as kq + +DTYPES = [mx.bfloat16, mx.float16] +R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} +# Gate of quantized_matmul on the CPU stream, as in test_cpu_neon. +CPU_REL = 5e-2 +# Codecs with a verify_nax kernel, and an inner dim each serves. +CODECS = { + "q4_0": 512, + "q8_0": 512, + "iq4_nl": 512, + "q4_k": 512, + "q5_k": 512, + "q6_k": 512, + "q3_k": 512, + "q2_k": 512, + "iq4_xs": 512, +} +# Row counts of the joint grid's weights: tiny, partial row tiles, a +# 32-row remainder and a weight wider than the rest. +NS = [48, 160, 1000, 300] + +_W = {} + + +def _weight(codec, n, k, seed): + key = (codec, n, k, seed) + if key not in _W: + wf = mx.random.normal((n, k), key=mx.random.key(97 * seed + n + k)) * 0.1 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + wd = np.array(kq.dequantize(w, s, codec, dtype=mx.float32)).astype(np.float64) + _W[key] = (w, wd) + return _W[key] + + +def _weights(codec, ns, k): + return [_weight(codec, n, k, i) for i, n in enumerate(ns)] + + +def _x(m, k, dtype, seed=0): + return mx.random.normal((m, k), key=mx.random.key(seed + 7 * m)).astype(dtype) + + +def _per_weight(x, ws, codec): + sc = mx.zeros((1,), mx.uint8) + return [kq.quantized_matmul(x, w, sc, codec, transpose=True) for w in ws] + + +@pytest.fixture +def exact_levers(monkeypatch): + """One split count, verify_nax from M 2 and no per-row qmv route, so + quantized_matmul runs verify_nax on every weight the op joins.""" + + def set_splits(n): + monkeypatch.setenv("KQ_VERIFY_NAX_SPLITS", str(n)) + + monkeypatch.setenv("KQ_NAX_QMV", "0") + monkeypatch.setenv("KQ_VERIFY_NAX", "2") + return set_splits + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", sorted(CODECS)) +@pytest.mark.parametrize("splits", [1, 2]) +def test_matches_per_weight_bits(codec, dtype, splits, exact_levers): + exact_levers(splits) + k = CODECS[codec] + ws = [w for w, _ in _weights(codec, NS, k)] + for m in range(1, 9): + x = _x(m, k, dtype) + got = kq.quantized_matmul_multi(x, ws, codec) + want = _per_weight(x, ws, codec) + mx.eval(got, want) + assert len(got) == len(ws) + for g, w in zip(got, want, strict=True): + assert g.dtype == w.dtype and g.shape == w.shape + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ), f"{codec} M={m} N={w.shape[-1]}" + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", ["q4_0", "q4_k", "q6_k"]) +def test_reference_default_levers(codec, dtype): + k = CODECS[codec] + pairs = _weights(codec, NS, k) + ws = [w for w, _ in pairs] + for m in (1, 2, 3, 4, 6, 8, 9, 16): + x = _x(m, k, dtype, seed=3) + got = kq.quantized_matmul_multi(x, ws, codec) + mx.eval(got) + xd = np.array(x.astype(mx.float32)).astype(np.float64) + for g, (_, wd) in zip(got, pairs, strict=True): + ref = xd @ wd.T + err = np.linalg.norm(np.array(g.astype(mx.float32)) - ref) + rel = CPU_REL if force_cpu() else R1[dtype] + assert err <= rel * np.linalg.norm(ref), f"{codec} M={m}" + + +def test_q4_0_two_block_inner_dim(exact_levers): + # K not a multiple of 256 runs the two-block q4_0 kernel. + exact_levers(1) + ws = [w for w, _ in _weights("q4_0", [64, 200], 320)] + for m in (3, 5, 8): + x = _x(m, 320, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +@pytest.mark.parametrize("count", [1, 4, 5, 9]) +def test_weight_counts(count, exact_levers): + # More than four weights run one dispatch per four. + exact_levers(1) + ns = [64 + 32 * i for i in range(count)] + ws = [w for w, _ in _weights("q4_0", ns, 512)] + x = _x(6, 512, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +def test_batch_dims_views_and_float32(exact_levers): + exact_levers(1) + big, _ = _weight("q4_0", 512, 512, 11) + ws = [big[100:228], big[300:348], _weight("q4_0", 96, 512, 12)[0]] + x = _x(6, 512, mx.float32).reshape(2, 3, 512) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert g.shape == w.shape and g.dtype == mx.bfloat16 + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +def test_empty_weight(exact_levers): + exact_levers(1) + ws = [_weight("q4_0", 128, 512, 0)[0], mx.zeros((0, 288), mx.uint8)] + x = _x(4, 512, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + mx.eval(got) + assert got[1].shape == (4, 0) + want = _per_weight(x, ws[:1], "q4_0")[0] + assert np.array_equal( + np.array(got[0].astype(mx.float32)), np.array(want.astype(mx.float32)) + ) + + +def test_lever_off(monkeypatch): + monkeypatch.setenv("KQ_MATMUL_MULTI", "0") + ws = [w for w, _ in _weights("q4_0", NS, 512)] + x = _x(6, 512, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +def test_grad_matches_per_weight(): + ws = [w for w, _ in _weights("q4_0", NS, 512)] + x = _x(6, 512, mx.float32) + cot = [mx.random.normal((6, n), key=mx.random.key(n)) for n in NS] + + def f_multi(x): + return sum( + (y.astype(mx.float32) * c).sum() + for y, c in zip(kq.quantized_matmul_multi(x, ws, "q4_0"), cot, strict=True) + ) + + def f_each(x): + return sum( + (y.astype(mx.float32) * c).sum() + for y, c in zip(_per_weight(x, ws, "q4_0"), cot, strict=True) + ) + + g1 = mx.grad(f_multi)(x) + g2 = mx.grad(f_each)(x) + mx.eval(g1, g2) + assert np.allclose(np.array(g1), np.array(g2), rtol=2e-2, atol=2e-2) + + +def test_argument_errors(): + w, _ = _weight("q4_0", 64, 512, 0) + x = _x(4, 512, mx.bfloat16) + with pytest.raises(ValueError): + kq.quantized_matmul_multi(x, [], "q4_0") + with pytest.raises(ValueError): + kq.quantized_matmul_multi(x, [w, _weight("q4_0", 64, 256, 0)[0]], "q4_0") + with pytest.raises(ValueError): + kq.quantized_matmul_multi(_x(4, 256, mx.bfloat16), [w], "q4_0") + with pytest.raises(ValueError): + kq.quantized_matmul_multi(x, [w], "no_such_codec") + + +_LOG_CHILD = r""" +import os, sys +import mlx.core as mx +import mlx_kquant as kq +ns = (4096, 640, 48, 48) +w = [kq.quantize(mx.random.normal((n, 512)) * 0.1, "q4_0")[0] for n in ns] +x = mx.random.normal((6, 512)).astype(mx.bfloat16) +mx.eval(kq.quantized_matmul_multi(x, w, "q4_0")) +""" + + +@pytest.mark.skipif(force_cpu() or not kq.nax_available(), reason="NAX GPU route") +def test_one_dispatch_on_nax(): + env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + env["KQ_QMM_LOG"] = "1" + env["KQ_VERIFY_NAX_SPLITS"] = "1" + r = subprocess.run( + [sys.executable, "-c", _LOG_CHILD], + env=env, + capture_output=True, + text=True, + timeout=300, + ) + assert r.returncode == 0, r.stderr + lines = [ln for ln in r.stderr.splitlines() if ln.startswith("[kq qmm]")] + assert len(lines) == 1 and "verify_nax_multi" in lines[0], r.stderr From 540a640bfc96ed29e6257e680baad4ddde443074 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 03:31:01 -0700 Subject: [PATCH 02/10] fix(matmul): quantized_matmul_multi takes the route quantized_matmul takes, under every lever --- CHANGELOG.md | 7 + bindings.cpp | 7 +- docs/kernels.md | 41 +++- .../mlx/backend/metal/kernels/kq_verify_nax.h | 81 ++++--- mlx_kquant/__init__.py | 2 +- src/kquant_matmul.cpp | 174 ++++++++------- src/kquant_matmul_fused.cpp | 7 +- tests/qmm_probe.py | 24 ++ tests/test_matmul_empty.py | 13 +- tests/test_matmul_multi.py | 206 +++++++++++++++--- tests/test_split_kernel_dispatch.py | 39 +++- tests/test_weight_base.py | 41 ++++ tests/test_weight_views.py | 21 ++ 13 files changed, 510 insertions(+), 153 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0259f21..fe6b9bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ All notable changes to this project are documented here. The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project aims to adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- `kq.quantized_matmul_multi` computes the products of several packed weights + with one input, such as the q, k and v projections, in one dispatch on NAX + GPUs (M5) at speculative-verify widths of 2 to 8 rows. + ## [0.5.0] ### Added diff --git a/bindings.cpp b/bindings.cpp index edf8af9..7e3d066 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -683,8 +683,11 @@ NB_MODULE(_ext, m) { weights computes the products. Each weight is read from its own buffer and each product is written to its own output, so the weights need no joining copy. One split-K count covers the joint - grid. Every other case runs quantized_matmul per weight. - KQ_MATMUL_MULTI=0 runs quantized_matmul per weight on every call. + grid. The narrower weights take the widest weight's route, so + their outputs can differ from quantized_matmul in the last bits. + Every other case runs quantized_matmul per weight. + KQ_MATMUL_MULTI=0 runs quantized_matmul per weight on every call + (read once per process unless KQ_TUNE_LIVE=1). mx.grad gives the gradient with respect to x; the weights are frozen, as in quantized_matmul. diff --git a/docs/kernels.md b/docs/kernels.md index af3452c..5aa34a8 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -223,7 +223,7 @@ Tuning levers (defaults are right for normal use): the weights; `benchmarks/bench_dit_shapes.py` runs its tile arms this way. It reads the environment on every dispatch, so it is for single-threaded A/B runs that set these variables between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, `KQ_GLU_MIN_ROWS`), the - `quantized_matmul_lora` levers (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the + `quantized_matmul_multi` lever (`KQ_MATMUL_MULTI`), the `quantized_matmul_lora` levers (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the `quantized_matmul_residual` lever (`KQ_QMM_RESID`), the dequantize route lever (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are read when the op is called and builds its graph, not when a kernel is dispatched, so a compiled @@ -642,6 +642,45 @@ entries the kernel ran up to 1.10x slower at K 4096 and 1.21x slower at N 12288, entry under M 5 sits one row higher. At M 8 the kernel runs 1.05-1.39x faster than NAX split-K on the shapes of N 2048 and wider. +## Quantized matmuls that share an input + +- `quantized_matmul_multi` - `[x @ W^T for W in weights]` for packed weights of one codec and one + inner dim that read the same x, such as the q, k and v projections or the gate and up projections + of a transformer block. It returns one output per weight, as `quantized_matmul` would. + + On a NAX GPU (M5), at 2 to 8 activation rows, where `quantized_matmul` would take its verify_nax + route on the widest weight, one dispatch computes the products of up to four weights. The row + tiles of the weights lie end to end in one grid, each threadgroup owns 128 rows of one weight, + and the split-K partials of all of them share one temporary and one fold dispatch. More than four + weights run one dispatch per four. The split count is chosen for the joint grid. The call saves + the dispatch, the wait and the tail of every weight but the first, which counts most for narrow + weights next to a wide one. + + The narrower weights take the route of the widest. A weight that `quantized_matmul` alone would + run on per-row qmv or on another kernel runs on verify_nax here, so its output can differ from + `quantized_matmul` in the last bits, inside the same bound against a float64 reference. This + holds for every codec with a verify_nax kernel, including those whose verify_nax entry starts at + a width that the narrow weight does not reach. Where every weight would take verify_nax at one + split count, the outputs equal `quantized_matmul` bit for bit. + + Every other case runs `quantized_matmul` per weight: one row or more than 8, a single weight, + codecs without a verify_nax kernel, an inner dim that kernel does not step over, a weight whose + start it cannot read, CPU streams and GPUs without NAX. x is float16 or bfloat16, float32 x runs + in bfloat16 as in `quantized_matmul`, and leading batch dims fold into rows. + + The route levers of `quantized_matmul` decide the route on the widest weight the same way: + `KQ_NAX_QMV`, `KQ_VERIFY_NAX`, `KQ_VERIFY_MMA`, `KQ_QMM_SPLITK_NAX`, `KQ_QMM_SPLITK`, + `KQ_VERIFY_EXT` and `KQ_DISABLE_NAX`. A set `KQ_QMM_ROUTE` or `KQ_FORCE_QMM_MIN_M` runs + `quantized_matmul` per weight, so a forced route times the route it names. + + Lever, read when the op builds its graph, once per process unless `KQ_TUNE_LIVE=1`: + + - `KQ_MATMUL_MULTI` - `0` builds `quantized_matmul` per weight on every call. Unset (on) by + default. + + `mx.grad` gives the gradient with respect to x through `quantized_matmul` per weight. The weights + are frozen, as in `quantized_matmul`. + ## GLU product of two quantized matmuls - `quantized_matmul_glu` - `y = act(x @ Wg^T) * (x @ Wu^T)` for two packed weights of one codec and diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index b95fe0d..57469b9 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -949,47 +949,46 @@ METAL_FUNC void kq_verify_nax_impl( // own buffer and into its own output, with one split count over grid z. // part_stride[i] is output i's partial stride. Unused weight slots repeat // slot 0 and own no threadgroup. -#define KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(CODEC, TRAITS) \ - template \ - [[kernel]] void kq_##CODEC##_verify_nax_multi( \ - const device uint8_t* w0, \ - const device uint8_t* w1, \ - const device uint8_t* w2, \ - const device uint8_t* w3, \ - const device T* x, \ - device U* y0, \ - device U* y1, \ - device U* y2, \ - device U* y3, \ - const constant int& K, \ - const constant int& M, \ - const constant int& k_partition_size, \ - const constant int4& n_rows, \ - const constant int4& tg_end, \ - const constant int4& part_stride, \ - uint3 tid [[threadgroup_position_in_grid]], \ - uint simd_gid [[simdgroup_index_in_threadgroup]], \ - uint simd_lid [[thread_index_in_simdgroup]]) { \ - static_assert(group_size == TRAITS::group, #CODEC " block width"); \ - const int t = int(tid.x); \ - const short i = t < tg_end[0] ? 0 \ - : t < tg_end[1] ? 1 \ - : t < tg_end[2] ? 2 \ - : 3; \ - const int t0 = i == 0 ? 0 : tg_end[i - 1]; \ - const device uint8_t* w = i == 0 ? w0 : i == 1 ? w1 : i == 2 ? w2 : w3; \ - device U* y = i == 0 ? y0 : i == 1 ? y1 : i == 2 ? y2 : y3; \ - kq_verify_nax_impl( \ - w, \ - x, \ - y + tid.z * static_cast(part_stride[i]), \ - K, \ - n_rows[i], \ - M, \ - k_partition_size, \ - uint3(uint(t - t0), tid.y, tid.z), \ - simd_gid, \ - simd_lid); \ +#define KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(CODEC, TRAITS) \ + template \ + [[kernel]] void kq_##CODEC##_verify_nax_multi( \ + const device uint8_t* w0, \ + const device uint8_t* w1, \ + const device uint8_t* w2, \ + const device uint8_t* w3, \ + const device T* x, \ + device U* y0, \ + device U* y1, \ + device U* y2, \ + device U* y3, \ + const constant int& K, \ + const constant int& M, \ + const constant int& k_partition_size, \ + const constant int4& n_rows, \ + const constant int4& tg_end, \ + const constant int4& part_stride, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + static_assert(group_size == TRAITS::group, #CODEC " block width"); \ + const int t = int(tid.x); \ + /* tg_end is cumulative, so the weight is the ends passed. */ \ + const short i = \ + short(t >= tg_end[0]) + short(t >= tg_end[1]) + short(t >= tg_end[2]); \ + const int t0 = i == 0 ? 0 : tg_end[i - 1]; \ + const device uint8_t* w = i == 0 ? w0 : i == 1 ? w1 : i == 2 ? w2 : w3; \ + device U* y = i == 0 ? y0 : i == 1 ? y1 : i == 2 ? y2 : y3; \ + kq_verify_nax_impl( \ + w, \ + x, \ + y + tid.z * static_cast(part_stride[i]), \ + K, \ + n_rows[i], \ + M, \ + k_partition_size, \ + uint3(uint(t - t0), tid.y, tid.z), \ + simd_gid, \ + simd_lid); \ } KQ_DEFINE_VERIFY_NAX_KERNEL(pq2_0, KqPq2_0Nax) diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index 875665e..2436ce5 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -196,8 +196,8 @@ "quantize", "quantized_matmul", "quantized_matmul_glu", - "quantized_matmul_multi", "quantized_matmul_lora", + "quantized_matmul_multi", "quantized_matmul_qmv_bias", "quantized_matmul_residual", "rmsnorm2_add", diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 4c14b1e..6cbee28 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -2720,10 +2720,92 @@ std::vector KQuantMatmul::vjp( #ifdef _METAL_ -// One verify_nax dispatch per four weights where the widest weight would -// take the verify_nax route at this M, and quantized_matmul per weight -// otherwise. The route levers of quantized_matmul apply: KQ_NAX_QMV, -// KQ_VERIFY_NAX and KQ_DISABLE_NAX through kq_is_nax_available. +// Probe lever: KQ_FORCE_QMM_MIN_M= routes transpose shapes with M >= m +// straight to qmm (-> NAX), bypassing the mv_ext/verify_qmv claims, for +// qmm-vs-mv_ext crossover measurement below M=13. Read once per process. +static int kq_force_qmm_min_m() { + static const int m = []() { + const char* e = std::getenv("KQ_FORCE_QMM_MIN_M"); + return e != nullptr ? std::atoi(e) : 0; + }(); + return m; +} + +// KQ_QMM_ROUTE= is set (see eval_gpu_base). +static bool kq_qmm_route_forced() { + const char* e = std::getenv("KQ_QMM_ROUTE"); + return e != nullptr && e[0] != '\0'; +} + +enum class KqSmallMRoute { other, qmv, verify_nax }; + +// The NAX small-M route quantized_matmul takes for a transposed, +// non-batched call on a float16 or bfloat16 x of M rows, once the forced +// KQ_FORCE_QMM_MIN_M and KQ_QMM_ROUTE routes are out of the way: per-row +// qmv, verify_nax (still subject to kq_verify_nax_fits), or a later route. +// quantized_matmul_multi joins weights into one verify_nax dispatch only +// where this returns verify_nax for the widest weight. +// +// Per-row qmv on NAX GPUs, M 2 through kq_nax_qmv_max_m for the call's +// N. KQ_NAX_QMV= sets the limit to m at every N, and 0 disables the +// route. The default yields to a forced KQ_VERIFY_NAX or KQ_VERIFY_MMA +// on a codec they serve, a forced KQ_QMM_SPLITK_NAX or KQ_QMM_SPLITK, +// and a set KQ_VERIFY_EXT, so their arms time the route they name. +// KQ_DISABLE_NAX turns it off, since the limits are measured against +// the NAX routes. KQ_NAX_QMV is read live so an A/B can flip arms on +// one generator. +// +// Register-fed NAX verify, from kq_verify_nax_min_m through M8; see +// verify_nax. KQ_VERIFY_NAX= forces the entry at M >= m, and 0 +// disables it. Read live so an A/B can flip arms on one generator. +// KQ_DISABLE_NAX turns it off with the other NAX routes. Decided ahead +// of verify_mma and NAX split-K, but the default entry yields to a +// forced KQ_VERIFY_MMA on a codec verify_mma serves, or a forced +// KQ_QMM_SPLITK_NAX, so their arms time the route they name. +static KqSmallMRoute +kq_nax_small_m_route(const std::string& t, int M, int N, int K, bool f16) { + if (M < 1 || M > 32 || !kq_is_nax_available() || !codec_has_nax(t)) { + return KqSmallMRoute::other; + } + const int vnax_env = kq_env_int("KQ_VERIFY_NAX", -1); + const int vmma_force = kq_env_int("KQ_VERIFY_MMA", -1); + const int sk_force = kq_env_int("KQ_QMM_SPLITK_NAX", -1); + if (M >= 2) { + int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); + if (qmv_max_m == -1) { + const int ext = kq_verify_ext_env(); + const bool forced = + (vnax_env >= 1 && M >= vnax_env && codec_verify_nax_kstep(t) > 0) || + (vmma_force >= 2 && M >= vmma_force && + codec_verify_mma_rows(t) > 0) || + sk_force >= 1 || kq_qmm_splitk_env() > 1 || ext == 0 || ext == 1; + qmv_max_m = forced ? 0 : kq_nax_qmv_max_m(t, N); + } + if (M <= qmv_max_m) { + return KqSmallMRoute::qmv; + } + } + if (M <= 8) { + const bool other_forced = + (vmma_force >= 2 && M >= vmma_force && codec_verify_mma_rows(t) > 0) || + sk_force >= 1; + int vnax_min_m = vnax_env; + if (vnax_env == -1) { + vnax_min_m = other_forced ? 0 : kq_verify_nax_min_m(t, N, K, f16); + } + if (vnax_min_m > 0 && M >= vnax_min_m) { + return KqSmallMRoute::verify_nax; + } + } + return KqSmallMRoute::other; +} + +// One verify_nax dispatch per four weights where quantized_matmul would +// take the verify_nax route on the widest weight at this M, and +// quantized_matmul per weight otherwise. The narrower weights ride along on +// that route. Every route lever of quantized_matmul applies, through +// kq_nax_small_m_route, and a forced KQ_QMM_ROUTE or KQ_FORCE_QMM_MIN_M +// runs quantized_matmul per weight. void KQuantMatmulMulti::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -2745,20 +2827,14 @@ void KQuantMatmulMulti::eval_gpu( n_max = std::max(n_max, w.shape(0)); } const int kstep = kq_verify_nax_call_kstep(K, kquant_type_); - bool multi = M >= 2 && M <= 8 && n_max > 0 && x.flags().row_contiguous && - x.dtype() != mx::float32 && kq_is_nax_available() && - codec_has_nax(kquant_type_) && kstep > 0; + const int force_min_m = kq_force_qmm_min_m(); + bool multi = n_max > 0 && kstep > 0 && x.flags().row_contiguous && + x.dtype() != mx::float32 && !kq_qmm_route_forced() && + !(force_min_m > 0 && M >= force_min_m) && + kq_nax_small_m_route( + kquant_type_, M, n_max, K, x.dtype() == mx::float16) == + KqSmallMRoute::verify_nax; if (multi) { - int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); - if (qmv_max_m == -1) { - qmv_max_m = kq_nax_qmv_max_m(kquant_type_, n_max); - } - int vnax_min_m = kq_env_int("KQ_VERIFY_NAX", -1); - if (vnax_min_m == -1) { - vnax_min_m = - kq_verify_nax_min_m(kquant_type_, n_max, K, x.dtype() == mx::float16); - } - multi = M > qmv_max_m && vnax_min_m > 0 && M >= vnax_min_m; for (const auto& w : ws) { multi = multi && kq_verify_nax_fits(w, K, kstep, kquant_type_); } @@ -2858,13 +2934,8 @@ void KQuantMatmul::eval_gpu_base( // through to dispatch_qmv below rather than throw; a qmv_quad kernel would be // a perf-only path for an essentially-dead case. - // Probe lever: KQ_FORCE_QMM_MIN_M= routes transpose shapes with M >= m - // straight to qmm (-> NAX), bypassing the mv_ext/verify_qmv claims, for - // qmm-vs-mv_ext crossover measurement below M=13. - static const int force_qmm_min_m = []() { - const char* e = std::getenv("KQ_FORCE_QMM_MIN_M"); - return e != nullptr ? std::atoi(e) : 0; - }(); + // Probe lever KQ_FORCE_QMM_MIN_M (kq_force_qmm_min_m). + const int force_qmm_min_m = kq_force_qmm_min_m(); if (force_qmm_min_m > 0 && transpose_ && M >= force_qmm_min_m) { qmm(x, w, @@ -2919,60 +2990,17 @@ void KQuantMatmul::eval_gpu_base( } } - // Per-row qmv on NAX GPUs, M 2 through kq_nax_qmv_max_m for the call's - // N. KQ_NAX_QMV= sets the limit to m at every N, and 0 disables the - // route. The default yields to a forced KQ_VERIFY_NAX or KQ_VERIFY_MMA - // on a codec they serve, a forced KQ_QMM_SPLITK_NAX or KQ_QMM_SPLITK, - // and a set KQ_VERIFY_EXT, so their arms time the route they name. - // KQ_DISABLE_NAX turns it off, since the limits are measured against - // the NAX routes. KQ_NAX_QMV is read live so an A/B can flip arms on - // one generator. - if (transpose_ && non_batched && M >= 2 && M <= 32 && - x.dtype() != mx::float32 && kq_is_nax_available() && - codec_has_nax(kquant_type_)) { - int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); - if (qmv_max_m == -1) { - const int vnax_force = kq_env_int("KQ_VERIFY_NAX", -1); - const int vmma_force = kq_env_int("KQ_VERIFY_MMA", -1); - const int sk_force = kq_env_int("KQ_QMM_SPLITK_NAX", -1); - const int ext = kq_verify_ext_env(); - const bool forced = (vnax_force >= 1 && M >= vnax_force && - codec_verify_nax_kstep(kquant_type_) > 0) || - (vmma_force >= 2 && M >= vmma_force && - codec_verify_mma_rows(kquant_type_) > 0) || - sk_force >= 1 || kq_qmm_splitk_env() > 1 || ext == 0 || ext == 1; - qmv_max_m = forced ? 0 : kq_nax_qmv_max_m(kquant_type_, N); - } - if (M <= qmv_max_m) { + // Per-row qmv and verify_nax on NAX GPUs: kq_nax_small_m_route. + if (transpose_ && non_batched && x.dtype() != mx::float32) { + const KqSmallMRoute route = + kq_nax_small_m_route(kquant_type_, M, N, K, x.dtype() == mx::float16); + if (route == KqSmallMRoute::qmv) { dispatch_qmv( x, w, scales, out, group_size_, bits_, M, N, K, d, s, kquant_type_); return; } - } - - // Register-fed NAX verify, from kq_verify_nax_min_m through M8; see - // verify_nax. KQ_VERIFY_NAX= forces the entry at M >= m, and 0 - // disables it. Read live so an A/B can flip arms on one generator. - // KQ_DISABLE_NAX turns it off with the other NAX routes. Decided ahead - // of verify_mma and NAX split-K, but the default entry yields to a - // forced KQ_VERIFY_MMA on a codec verify_mma serves, or a forced - // KQ_QMM_SPLITK_NAX, so their arms time the route they name. - if (transpose_ && non_batched && M <= 8 && x.dtype() != mx::float32 && - kq_is_nax_available() && codec_has_nax(kquant_type_)) { const int kstep = kq_verify_nax_call_kstep(K, kquant_type_); - const int vnax_env = kq_env_int("KQ_VERIFY_NAX", -1); - const int vmma_force = kq_env_int("KQ_VERIFY_MMA", -1); - const int sk_force = kq_env_int("KQ_QMM_SPLITK_NAX", -1); - const bool other_forced = (vmma_force >= 2 && M >= vmma_force && - codec_verify_mma_rows(kquant_type_) > 0) || - sk_force >= 1; - int vnax_min_m = vnax_env; - if (vnax_env == -1) { - vnax_min_m = other_forced - ? 0 - : kq_verify_nax_min_m(kquant_type_, N, K, x.dtype() == mx::float16); - } - if (vnax_min_m > 0 && M >= vnax_min_m && + if (route == KqSmallMRoute::verify_nax && kq_verify_nax_fits(w, K, kstep, kquant_type_)) { const int sp = kq_verify_nax_splits(N, K, kstep, kquant_type_); verify_nax( diff --git a/src/kquant_matmul_fused.cpp b/src/kquant_matmul_fused.cpp index b58464f..7b9275f 100644 --- a/src/kquant_matmul_fused.cpp +++ b/src/kquant_matmul_fused.cpp @@ -467,10 +467,11 @@ namespace { const char* kMultiOp = "[mlx_kquant.quantized_matmul_multi]"; -// KQ_MATMUL_MULTI=0 runs quantized_matmul per weight on every call. Read -// live so an A/B can flip arms in one process. +// KQ_MATMUL_MULTI=0 builds quantized_matmul per weight on every call. Read +// once per process, or on every call under KQ_TUNE_LIVE=1, as KQ_GLU. bool multi_on() { - return kq_env_flag("KQ_MATMUL_MULTI", 1) != 0; + static const bool cached = kq_env_flag("KQ_MATMUL_MULTI", 1) != 0; + return glu_tune_live() ? kq_env_flag("KQ_MATMUL_MULTI", 1) != 0 : cached; } std::vector multi_composed( diff --git a/tests/qmm_probe.py b/tests/qmm_probe.py index d4e108b..01c33ee 100644 --- a/tests/qmm_probe.py +++ b/tests/qmm_probe.py @@ -6,6 +6,8 @@ and returns per call the ``KQ_QMM_LOG`` lines it printed (kernel name and grid), its error if it raised, and its relative Frobenius error against a float64 reference built from ``kq.dequantize`` weights. A call with +``"multi": [n, ...]`` runs ``quantized_matmul_multi`` over its weight and +one more of each listed N, and gets the largest error over the outputs. A call with ``"r2": true`` also gets ``rel2`` against those weights rounded to the activation dtype, and ``floor1``/``floor2``, the error of rounding each reference to the output dtype. A call with ``"ulp_vs": i`` gets ``ulp``, @@ -244,6 +246,28 @@ def ref_of(lo): ) os.write(2, f"[probe] {i}\n".encode()) res = {"i": i, "error": None, "rel": None, "ulp": None} + if c.get("multi"): + # quantized_matmul_multi over weights of N n and each of "multi"; + # rel is the largest over the outputs. + ns = [c["n"], *c["multi"]] + wl = [weights(c["codec"], n, c["shape"][-1], design)[0] for n in ns] + try: + ys = kq.quantized_matmul_multi(x, wl, c["codec"]) + mx.eval(ys) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + outs.append(None) + results.append(res) + continue + xn = as64(x) + res["rel"] = max( + rel(as64(y), xn @ dequantized(c["codec"], n, c["shape"][-1], design).T) + for y, n in zip(ys, ns, strict=True) + ) + outs.append(None) + results.append(res) + mx.clear_cache() + continue try: y = kq.quantized_matmul(x, w, s, c["codec"], transpose=True) mx.eval(y) diff --git a/tests/test_matmul_empty.py b/tests/test_matmul_empty.py index 6ddc4ef..1e6b185 100644 --- a/tests/test_matmul_empty.py +++ b/tests/test_matmul_empty.py @@ -1,4 +1,5 @@ -"""quantized_matmul and quantized_matmul_glu on empty inputs. +"""quantized_matmul, quantized_matmul_glu and quantized_matmul_multi on +empty inputs. No rows or no output columns give an empty array of the product's shape, and an inner dim of 0 gives zeros. The cases run on the default device, so @@ -27,11 +28,19 @@ def _ops(codec): return { "matmul": lambda x, w: kq.quantized_matmul(x, w, SCALES, codec, True), "glu": lambda x, w: kq.quantized_matmul_glu(x, w, w, codec), + "multi": lambda x, w: _multi(x, w, codec), } +def _multi(x, w, codec): + # Two weights through one call; both outputs are the same product. + a, b = kq.quantized_matmul_multi(x, [w, w], codec) + assert a.shape == b.shape and a.dtype == b.dtype + return mx.concatenate([a[None], b[None]]).max(axis=0) + + @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) -@pytest.mark.parametrize("op", ["matmul", "glu"]) +@pytest.mark.parametrize("op", ["matmul", "glu", "multi"]) @pytest.mark.parametrize("codec", ["q4_k", "q8_0", "q3_k"]) def test_empty_inputs(codec, op, dtype): f = _ops(codec)[op] diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py index ca4ff29..fd84f1a 100644 --- a/tests/test_matmul_multi.py +++ b/tests/test_matmul_multi.py @@ -2,9 +2,9 @@ On a NAX GPU the op runs one verify_nax dispatch per four weights at 2 to 8 rows, where quantized_matmul would run verify_nax on the widest weight. -With the split count forced to one value, the verify_nax entry at M 2 and -the per-row qmv route off, quantized_matmul per weight runs the same -kernel body and fold, so the outputs match bit for bit. At the default +With the split count forced to one value and the verify_nax entry forced +at M 2, quantized_matmul per weight runs the same kernel body and fold, so +the outputs match bit for bit. At the default levers the joint grid can take another split count than each weight alone, so the outputs are checked against a float64 reference built from kq.dequantize weights. Every other @@ -13,6 +13,7 @@ from __future__ import annotations +import json import os import subprocess import sys @@ -20,9 +21,11 @@ import mlx.core as mx import numpy as np import pytest +from kqref import synth_wire from qmm_probe import force_cpu import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY DTYPES = [mx.bfloat16, mx.float16] R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} @@ -73,17 +76,31 @@ def _per_weight(x, ws, codec): @pytest.fixture def exact_levers(monkeypatch): - """One split count, verify_nax from M 2 and no per-row qmv route, so - quantized_matmul runs verify_nax on every weight the op joins.""" + """One split count and verify_nax forced from M 2, which also turns the + default per-row qmv route off, so quantized_matmul runs verify_nax on + every weight the op joins.""" def set_splits(n): monkeypatch.setenv("KQ_VERIFY_NAX_SPLITS", str(n)) - monkeypatch.setenv("KQ_NAX_QMV", "0") monkeypatch.setenv("KQ_VERIFY_NAX", "2") return set_splits +def _synth(codec, n, k, seed): + # Codecs kq.quantize does not encode: random valid wire. + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + rng = np.random.default_rng(seed) + wire = synth_wire(rng, codec, bpb, n * (k // wpb)) + return mx.array(wire.reshape(n, (k // wpb) * bpb)) + + +def _equal(a, b): + return np.array_equal( + np.array(a.astype(mx.float32)), np.array(b.astype(mx.float32)) + ) + + @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("codec", sorted(CODECS)) @pytest.mark.parametrize("splits", [1, 2]) @@ -104,6 +121,21 @@ def test_matches_per_weight_bits(codec, dtype, splits, exact_levers): ), f"{codec} M={m} N={w.shape[-1]}" +@pytest.mark.parametrize("codec", ["pq2_0", "q2_0"]) +@pytest.mark.parametrize("splits", [1, 2, 4]) +def test_synth_codecs_match_per_weight_bits(codec, splits, exact_levers): + exact_levers(splits) + ws = [_synth(codec, n, 1024, i) for i, n in enumerate(NS)] + for dtype in DTYPES: + for m in range(1, 9): + x = _x(m, 1024, dtype) + got = kq.quantized_matmul_multi(x, ws, codec) + want = _per_weight(x, ws, codec) + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert _equal(g, w), f"{codec} M={m} N={w.shape[-1]}" + + @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("codec", ["q4_0", "q4_k", "q6_k"]) def test_reference_default_levers(codec, dtype): @@ -181,17 +213,23 @@ def test_empty_weight(exact_levers): ) -def test_lever_off(monkeypatch): - monkeypatch.setenv("KQ_MATMUL_MULTI", "0") - ws = [w for w, _ in _weights("q4_0", NS, 512)] +def test_misaligned_base_falls_back(exact_levers): + # verify_nax reads q4_0 in 4-byte words: a weight 2 bytes past a 4-byte + # boundary sends the whole group per weight, off verify_nax. + exact_levers(1) + w = _weight("q4_0", 64, 512, 3)[0] + flat = mx.concatenate([mx.zeros((2,), mx.uint8), w.reshape(-1)]) + mx.eval(flat) + wv = flat[2:].reshape(w.shape) + mx.eval(wv) + ws = [_weight("q4_0", 128, 512, 4)[0], wv] x = _x(6, 512, mx.bfloat16) got = kq.quantized_matmul_multi(x, ws, "q4_0") want = _per_weight(x, ws, "q4_0") mx.eval(got, want) for g, w in zip(got, want, strict=True): - assert np.array_equal( - np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) - ) + assert np.isfinite(np.array(g.astype(mx.float32))).all() + assert _equal(g, w) def test_grad_matches_per_weight(): @@ -231,28 +269,146 @@ def test_argument_errors(): _LOG_CHILD = r""" -import os, sys +import json, os, sys import mlx.core as mx import mlx_kquant as kq -ns = (4096, 640, 48, 48) -w = [kq.quantize(mx.random.normal((n, 512)) * 0.1, "q4_0")[0] for n in ns] -x = mx.random.normal((6, 512)).astype(mx.bfloat16) -mx.eval(kq.quantized_matmul_multi(x, w, "q4_0")) +codec, m, k, ns, steps = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), \ + json.loads(sys.argv[4]), json.loads(sys.argv[5]) +w = [kq.quantize(mx.random.normal((n, k), key=mx.random.key(i)) * 0.1, codec)[0] + for i, n in enumerate(ns)] +x = mx.random.normal((m, k), key=mx.random.key(9)).astype(mx.bfloat16) +sc = mx.zeros((1,), mx.uint8) +mx.eval(w, x) +for i, (op, env) in enumerate(steps): + for name, value in env.items(): + os.environ[name] = value + sys.stderr.write(f"=== {i}\n") + sys.stderr.flush() + if op == "multi": + out = kq.quantized_matmul_multi(x, w, codec) + else: + out = [kq.quantized_matmul(x, wi, sc, codec, transpose=True) for wi in w] + mx.eval(out) """ -@pytest.mark.skipif(force_cpu() or not kq.nax_available(), reason="NAX GPU route") +def _census(steps, codec="q4_0", m=6, k=512, ns=(4096, 640, 48, 48), env=None): + """KQ_QMM_LOG kernel names per step of a child process. Each step is + (op, env): the env is set before the step's call, which is + quantized_matmul_multi ("multi") or quantized_matmul per weight.""" + child = {kk: v for kk, v in os.environ.items() if not kk.startswith("KQ_")} + child["KQ_QMM_LOG"] = "1" + child.update(env or {}) + r = subprocess.run( + [ + sys.executable, + "-c", + _LOG_CHILD, + codec, + str(m), + str(k), + json.dumps(list(ns)), + json.dumps(steps), + ], + env=child, + capture_output=True, + text=True, + timeout=300, + ) + assert r.returncode == 0, r.stderr + out = [] + for part in r.stderr.split("=== ")[1:]: + lines = part.splitlines()[1:] + out.append([ln.split()[2] for ln in lines if ln.startswith("[kq qmm]")]) + return out + + +NAX_ONLY = pytest.mark.skipif( + force_cpu() or not kq.nax_available(), reason="NAX GPU route" +) + + +@NAX_ONLY def test_one_dispatch_on_nax(): - env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} - env["KQ_QMM_LOG"] = "1" - env["KQ_VERIFY_NAX_SPLITS"] = "1" + (names,) = _census([["multi", {}]], env={"KQ_VERIFY_NAX_SPLITS": "1"}) + assert len(names) == 1 and "verify_nax_multi" in names[0], names + + +@NAX_ONLY +@pytest.mark.parametrize( + "lever", + [ + {"KQ_VERIFY_MMA": "2"}, + {"KQ_QMM_SPLITK_NAX": "1"}, + {"KQ_QMM_ROUTE": "qmv"}, + {"KQ_QMM_ROUTE": "verify_mma"}, + {"KQ_FORCE_QMM_MIN_M": "2"}, + {"KQ_VERIFY_NAX": "0"}, + {"KQ_NAX_QMV": "8"}, + ], + ids=lambda e: ",".join(f"{k}={v}" for k, v in e.items()), +) +def test_route_levers_are_followed(lever): + # A lever that sends quantized_matmul off verify_nax sends the op to + # quantized_matmul per weight, which then runs the same kernels. + ns = (4096, 4096) + per, multi = _census([["per", {}], ["multi", {}]], ns=ns, env=lever) + assert not any("verify_nax" in n for n in per), per + assert multi == per + + +@NAX_ONLY +def test_forced_verify_nax_fuses_narrow_weights(): + # KQ_VERIFY_NAX=2 alone takes quantized_matmul to verify_nax at widths + # whose default is per-row qmv; the op follows and joins the weights. + ns = (512, 256) + per, multi = _census( + [["per", {}], ["multi", {}]], m=4, ns=ns, env={"KQ_VERIFY_NAX": "2"} + ) + assert len(per) == 2 and all("_verify_nax_" in n for n in per), per + assert len(multi) == 1 and "verify_nax_multi" in multi[0], multi + + +@NAX_ONLY +def test_lever_read_once_per_process(): + # KQ_MATMUL_MULTI keeps its first read, as KQ_GLU, unless KQ_TUNE_LIVE=1. + steps = [["multi", {}], ["multi", {"KQ_MATMUL_MULTI": "0"}]] + first, second = _census(steps) + assert first == second and "verify_nax_multi" in first[0], (first, second) + first, second = _census(steps, env={"KQ_TUNE_LIVE": "1"}) + assert "verify_nax_multi" in first[0], first + assert not any("verify_nax_multi" in n for n in second), second + + +def test_lever_off(): + # KQ_MATMUL_MULTI=0 builds quantized_matmul per weight: the same bits. r = subprocess.run( - [sys.executable, "-c", _LOG_CHILD], - env=env, + [sys.executable, "-c", _LEVER_OFF_CHILD], + env={**os.environ, "KQ_MATMUL_MULTI": "0"}, capture_output=True, text=True, timeout=300, ) assert r.returncode == 0, r.stderr - lines = [ln for ln in r.stderr.splitlines() if ln.startswith("[kq qmm]")] - assert len(lines) == 1 and "verify_nax_multi" in lines[0], r.stderr + assert r.stdout.split()[-1] == "EQUAL", r.stdout + + +_LEVER_OFF_CHILD = r""" +import os +import mlx.core as mx +import numpy as np +if os.environ.get("KQUANT_FORCE_CPU"): + mx.set_default_device(mx.cpu) +import mlx_kquant as kq +ns = (48, 160, 1000, 300) +w = [kq.quantize(mx.random.normal((n, 512), key=mx.random.key(n)) * 0.1, "q4_0")[0] + for n in ns] +x = mx.random.normal((6, 512), key=mx.random.key(1)).astype(mx.bfloat16) +sc = mx.zeros((1,), mx.uint8) +got = kq.quantized_matmul_multi(x, w, "q4_0") +want = [kq.quantized_matmul(x, wi, sc, "q4_0", transpose=True) for wi in w] +mx.eval(got, want) +f32 = [np.array(a.astype(mx.float32)) for a in got + want] +same = all(np.array_equal(a, b) for a, b in zip(f32[: len(w)], f32[len(w) :])) +print("EQUAL" if same else "DIFFER") +""" diff --git a/tests/test_split_kernel_dispatch.py b/tests/test_split_kernel_dispatch.py index 5d6e375..fdaf692 100644 --- a/tests/test_split_kernel_dispatch.py +++ b/tests/test_split_kernel_dispatch.py @@ -4,11 +4,12 @@ route at the row counts that select each of its tiles, on both activation dtypes, on an N that is and is not a multiple of the column tile, with partials in the activation dtype and in float (KQ_SPLITK_F32), and with the -half verify_mma kernels (KQ_VERIFY_MMA_HALF). KQ_QMM_LOG names the kernel -each call ran. Together the calls must run every name the host enumerator -lists for the codec, and each result must meet the R1 bound against float64 -with exact kq.dequantize weights, so a kernel that compiles to nothing, or -that no call reaches, fails here. +half verify_mma kernels (KQ_VERIFY_MMA_HALF), and run quantized_matmul_multi +for its joined verify_nax kernels. KQ_QMM_LOG names the kernel each call +ran. Together the calls must run every name the host enumerator lists for +the codec, and each result must meet the R1 bound against float64 with exact +kq.dequantize weights, so a kernel that compiles to nothing, or that no +call reaches, fails here. """ from __future__ import annotations @@ -48,6 +49,9 @@ def _no_nax(): "KQ_QMM_ROUTE": None, "KQ_QMM_ROUTE_STRICT": None, "KQ_VERIFY_MMA_HALF": None, + "KQ_NAX_QMV": None, + "KQ_VERIFY_NAX": None, + "KQ_VERIFY_NAX_SPLITS": None, } # (family the arm may be declined for, rows, K, env). The split-K tile is # BM=8 through M 8, BM=16 through M 16 and the plain tile above; the NAX @@ -70,6 +74,11 @@ def _no_nax(): # K off the 256 step, for the codecs whose block divides it: q4_0 runs its # 64-step verify_nax kernel there. OFF_STEP = ("_verify_nax_", 8, K + 64, {"KQ_QMM_ROUTE": "verify_nax", **STRICT}) +# quantized_matmul_multi runs the verify_nax_multi kernels: verify_nax from +# M 2 and more than one split, so the partials dtype shows in the name. The +# split count divides the K steps: 16 of 256 at K, 65 of 64 at K + 64. +MULTI = {"KQ_NAX_QMV": "0", "KQ_VERIFY_NAX": "2"} +MULTI_ARMS = [(8, K, "2"), (8, K + 64, "5")] def _host_names(): @@ -119,6 +128,26 @@ def test_every_split_kernel_runs(codec): } ) families.append(family) + for m, k, splits in MULTI_ARMS: + if k % CODEC_GEOMETRY[codec][3]: + continue + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "multi": [n + 64], + "dtype": dtype, + "seed": 1, + "env": { + **CLEAR, + **MULTI, + "KQ_VERIFY_NAX_SPLITS": splits, + "KQ_SPLITK_F32": f32, + }, + } + ) + families.append(None) res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"}) ran = set() for c, family, r in zip(calls, families, res, strict=True): diff --git a/tests/test_weight_base.py b/tests/test_weight_base.py index af7a200..415853f 100644 --- a/tests/test_weight_base.py +++ b/tests/test_weight_base.py @@ -196,6 +196,47 @@ def call(o, j=j, x=x): assert _equal(kq.quantized_matmul_glu(x, _shift(wg, 16), wu, codec), ref) +@GPU +@pytest.mark.parametrize("codec", ["q4_k", "q6_k", "q8_0", "q4_0"]) +def test_gpu_weight_base_multi(codec): + """quantized_matmul_multi checks the start of every weight, at the rows + that join them in one verify kernel and at those that do not.""" + a = _align(codec) + ws = [_wire(codec, (N, K), seed=j)[0] for j in range(3)] + # verify_nax declines these starts, so the group runs another route. + vnax = VNAX_BASE.get(codec, 1) + for m in (1, 6): + x = (mx.random.normal((m, K), key=mx.random.key(m)) * 0.5).astype(mx.bfloat16) + # One row builds quantized_matmul per weight, which names itself. + op = ( + "quantized_matmul_multi" + if m > 1 and kq.nax_available() + else ("quantized_matmul") + ) + ref = kq.quantized_matmul_multi(x, ws, codec) + mx.eval(ref) + for j in range(3): + for off in sorted({a, 16}): + moved = [_shift(w, off) if i == j else w for i, w in enumerate(ws)] + got = kq.quantized_matmul_multi(x, moved, codec) + same = _close if m > 1 and off % vnax else _equal + for g, r in zip(got, ref, strict=True): + assert same(g, r), (m, j, off) + moved = [_shift(w, a // 2) if i == j else w for i, w in enumerate(ws)] + msg = _msg(op, codec, a // 2) + with pytest.raises(ValueError, match=msg): + kq.quantized_matmul_multi(x, moved, codec) + # A weight that is not yet evaluated raises from the eval instead. + lazy = [ws[0], _shift(ws[1], a // 2, evaluate=False), ws[2]] + ys = kq.quantized_matmul_multi(x, lazy, codec) + with pytest.raises(ValueError, match=_msg(op, codec, a // 2)): + mx.eval(ys) + # The stream keeps working. + got = kq.quantized_matmul_multi(x, ws, codec) + for g, r in zip(got, ref, strict=True): + assert _equal(g, r) + + @GPU def test_gpu_weight_base_qmv_bias(): codec = "q8_0" diff --git a/tests/test_weight_views.py b/tests/test_weight_views.py index 50cdd37..2226f2a 100644 --- a/tests/test_weight_views.py +++ b/tests/test_weight_views.py @@ -106,6 +106,27 @@ def test_quantized_matmul_weight_views(codec, state): assert mx.array_equal(y, ref).item(), (codec, m, name, state) +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) +@pytest.mark.parametrize("codec", ["q4_k", "q8_0", "q4_0"]) +def test_multi_weight_views(codec, dtype, state): + """1 row runs quantized_matmul per weight, 4 and 6 rows the joined + verify kernel where there is one.""" + views = _views(codec) + for m in (1, 4, 6): + x = _x(m, K, dtype) + for pair in (("even_rows", "odd_rows"), ("row_band", "left_cols")): + makes = [views[name] for name in pair] + [views["whole"]] + ref = kq.quantized_matmul_multi(x, [_dense(f()) for f in makes], codec) + ws = [f() for f in makes] + if state == "evaluated": + mx.eval(ws) + ys = kq.quantized_matmul_multi(x, ws, codec) + mx.eval(ys, ref) + for y, r in zip(ys, ref, strict=True): + assert mx.array_equal(y, r).item(), (codec, m, pair, state) + + @pytest.mark.parametrize("state", ["lazy", "evaluated"]) @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) @pytest.mark.parametrize("m", [64, 1100]) From 2ebec0109ea517ea44559e63b667fd8ce987f13d Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 03:41:22 -0700 Subject: [PATCH 03/10] feat(sdpa): sdpa_fa_verify runs its key pass on the nax tensor-op units at head dims 128 and 256 --- CHANGELOG.md | 4 + docs/kernels.md | 10 +- metal/kq_sdpa_nax.metal | 21 ++ metal/mlx/backend/metal/kernels/kq_sdpa_nax.h | 258 ++++++++++++++++++ src/kquant_sdpa.cpp | 64 ++++- tests/test_sdpa.py | 123 +++++++-- 6 files changed, 459 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fe6b9bd..96822bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). with one input, such as the q, k and v projections, in one dispatch on NAX GPUs (M5) at speculative-verify widths of 2 to 8 rows. +### Changed +- `kq.sdpa_fa_verify` runs its key pass on the NAX tensor-op units (M5) at head + dims 128 and 256, so verify attention over a deep cache takes less time. + ## [0.5.0] ### Added diff --git a/docs/kernels.md b/docs/kernels.md index 5aa34a8..354a877 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -1058,7 +1058,15 @@ into a packed layout first. blocks instead of the full key axis and no `[L, S]` mask is materialized. Causality inside the diagonal page is enforced in-kernel. - **`sdpa_fa_verify`** - speculative-verify attention on the matrix units for a GQA-folded query tile. - Head dims 64 through 512; `return_lse` as above. + Head dims 64 through 512; `return_lse` as above. On NAX GPUs, at head dims 128 and 256 with up to + 64 folded rows on plain K and V, the key pass runs on the tensor-op units: each 16-row strip of + the fold is one row of simdgroups, and the head dim splits over 2 simdgroups (4 at head dim 256 + for folds of up to 32 rows), which add their partial scores in threadgroup memory. It writes the + partials of the simdgroup-matrix kernel to the same merge pass. Its automatic split count aims at + about 64 threadgroups over the KV heads, between 4 and 64, and halves while a split would hold + fewer than 64 keys. The two kernels round differently, so their outputs are not bit-identical. + `KQ_FA_VERIFY_NAX=0` keeps the simdgroup-matrix kernel, and `KQ_FA_VERIFY_NAX_WN` sets the head + dim split (2 or 4 at 256, 1 or 2 at 128); both are read at each call. - **`sdpa_fa_indexed`** - index-gathered attention over one shared K/V latent at head dim 512, the absorbed-MLA sparse decode step: query j of `q [1, Hq, Q, 512]` attends the rows of `kv [1, 1, N, 512]` that `idx[j]` (int32, -1 pads) lists, so the selected rows are read once from the latent diff --git a/metal/kq_sdpa_nax.metal b/metal/kq_sdpa_nax.metal index 81c354e..04c2ef6 100644 --- a/metal/kq_sdpa_nax.metal +++ b/metal/kq_sdpa_nax.metal @@ -20,3 +20,24 @@ instantiate_kq_sdpa_fa_indexed_nax(float16_t) instantiate_kq_sdpa_gqa_nax(bfloat16_t) instantiate_kq_sdpa_gqa_nax(float16_t) + +#define instantiate_kq_sdpa_fa_verify_nax(type, D, WM, WN) \ + instantiate_kernel( \ + "kq_sdpa_fa_verify_2pass_1_nax_" #type "_" #D "_wm" #WM "_wn" #WN, \ + kq_sdpa_fa_verify_2pass_1_nax, \ + type, D, WM, WN) + +#define instantiate_kq_sdpa_fa_verify_nax_rows(type, D, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 1, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 2, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 3, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 4, WN) + +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 256, 2) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 256, 2) +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 256, 4) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 256, 4) +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 128, 1) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 128, 1) +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 128, 2) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 128, 2) diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h b/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h index b7a5dc6..82ca870 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h @@ -607,3 +607,261 @@ template } } } + +// Speculative-verify attention on the tensor-op units, pass 1 of the +// sdpa_fa_verify split-K form, after MLX's head-dim split NAX attention +// (attention_nax_dsplit). It writes the partials of +// kq_sdpa_fa_verify_2pass_1 (kq_sdpa.h), so kq_sdpa_gqa_2pass_2 merges +// either. The caller folds the GQA group into the query rows: q is +// [B, Hkv, n_rows, D], and folded row r attends keys up to +// N - q_len + (r % q_len). +// +// One threadgroup per (kv head, batch row, key split), WM x WN simdgroups. +// Simdgroup (row group, d group) owns query rows 16 * row group onward and +// the d group's D / WN columns: its Q slice stays in registers, it reads +// its columns of each 32-key K and V tile from device memory into +// fragments, and keeps O for its columns only. The WN partial scores of a +// row group are summed in threadgroup memory in a fixed order, and every +// d group runs the same row softmax on the sum. A tile past a row's limit +// leaves the row's running max at finite_min and its P row at zero; an +// empty split writes O = 0, sum = 0 and max = finite_min. +template +[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +kq_sdpa_fa_verify_2pass_1_nax( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device float* out [[buffer(3)]], + device float* sums [[buffer(4)]], + device float* maxs [[buffer(5)]], + const constant int& N [[buffer(6)]], + const constant size_t& k_head_stride [[buffer(7)]], + const constant size_t& k_seq_stride [[buffer(8)]], + const constant size_t& v_head_stride [[buffer(9)]], + const constant size_t& v_seq_stride [[buffer(10)]], + const constant float& scale [[buffer(11)]], + const constant int& q_len [[buffer(12)]], + const constant int& n_rows [[buffer(13)]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tpg [[threadgroups_per_grid]]) { + using namespace kq_sdpa_nax; + constexpr short kU = 16; + constexpr int BK = 32; // keys per tile + constexpr short TK = BK / kU; + constexpr short TDh = D / kU / WN; // column fragments per simdgroup + constexpr int BDh = D / WN; + static_assert(TDh % 2 == 0, "P @ V accumulates column fragments in pairs"); + + using otile_t = NAXTile; + using stile_t = NAXTile; + constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; + constexpr short kRowsPT = otile_t::kRowsPerThread; + + threadgroup float s_xchg[WM][WN][TK * kEPF * 32]; + + const short row_group = simd_group_id / WN; + const short d_group = simd_group_id % WN; + const int n_kv_heads = tpg.x; + const int split_idx = tid.z; + const size_t qbh = (size_t)tid.y * n_kv_heads + tid.x; + const int tm = kU * row_group; + + const int chunk = ((N + gqa_splits * BK - 1) / (gqa_splits * BK)) * BK; + const int k0 = split_idx * chunk; + const int k1 = min(k0 + chunk, N); + + const int k_ld = int(k_seq_stride); + const int v_ld = int(v_seq_stride); + const device T* K = + keys + qbh * k_head_stride + (size_t)k0 * k_seq_stride + d_group * BDh; + const device T* V = + values + qbh * v_head_stride + (size_t)k0 * v_seq_stride + d_group * BDh; + + const short2 sc = otile_t::NAXFrag_t::get_coord(); + const short sm = sc.y; + const short sn = sc.x; + + // Q slice resident in registers; rows past n_rows load as zero. + const short q_rows = short(metal::clamp(n_rows - tm, 0, int(kU))); + NAXTile Qtiles[TDh]; + { + const device T* qs = queries + (qbh * n_rows + tm) * D + d_group * BDh; + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + Qtiles[id].load_rows(qs + id * kU, D, q_rows); + } + } + + otile_t Otile; + Otile.clear(); + metal::vec max_score; + metal::vec sum_score{0}; + // Highest key of this lane's rows (tm + sm and tm + sm + 8). + int lim[kRowsPT]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = Limits::finite_min; + lim[i] = N - q_len + ((tm + sm + i * 8) % q_len); + } + const int lim_min = N - q_len; // every row reaches it + const float scale2 = scale * M_LOG2E_F; + + for (int kt = k0; kt < k1; kt += BK) { + const int krem = min(k1 - kt, BK); + const bool full = krem == BK; + + stile_t Stile; + Stile.clear(); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + NAXTile Ktile; + const device T* kp = K + ik * kU * k_ld + id * kU; + if (full) { + Ktile.load(kp, k_ld); + } else { + Ktile.load_rows(kp, k_ld, short(krem - ik * kU)); + } + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), + Stile.frag_at(0, ik + 1), + Qtiles[id].frag_at(0, 0), + metal::false_type{}, + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), + metal::true_type{}); + } + } + + if (WN > 1) { + // Sum the WN partial scores of the row group in a fixed order. + threadgroup float* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + slot[base + ik * kEPF + i] = Stile.frag_at(0, ik)[i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short p = 1; p < WN; p++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += s_xchg[row_group][p][base + ik * kEPF + i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { + Stile.elems()[ii] *= scale2; + } + // Only the split tail or a tile reaching past N - q_len masks. + if (!full || kt + BK - 1 > lim_min) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& fg = Stile.frag_at(0, ik); + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const int c = kt + ik * kU + sn + jj; + const short loc = ii * stile_t::kFragThrCols + jj; + if (c >= k1 || c > lim[ii]) { + fg[loc] = Limits::finite_min; + } + } + } + } + } + + // Online softmax (exp2 space). A row with no valid key yet keeps its + // max at finite_min and contributes a zero P row. + metal::vec new_max = max_score; + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + metal::vec factor; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + if (new_max[i] > Limits::finite_min) { + factor[i] = fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + } else { + factor[i] = 1.0f; + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + Stile.frag_at(0, ik)[i * stile_t::kFragThrCols + jj] = 0.0f; + } + } + } + } + sum_score = sum_score * factor; + Stile.template row_reduce(sum_score); + Otile.template row_bin_op(factor); + + simdgroup_barrier(mem_flags::mem_none); + + // O += P @ V over this simdgroup's columns. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id += 2) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + NAXTile Vtile; + const device T* vp = V + ik * kU * v_ld + id * kU; + if (full) { + Vtile.load(vp, v_ld); + } else { + Vtile.load_rows(vp, v_ld, short(krem - ik * kU)); + } + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), + Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); + } + } + + K += BK * k_ld; + V += BK * v_ld; + } + + // Unnormalized partials [B, Hkv, n_rows, splits, D]; d group 0 writes the + // row stats, with the max in natural log for the shared merge. + if (q_rows > 0) { + device float* dst = out + + ((qbh * n_rows + tm) * gqa_splits + split_idx) * D + d_group * BDh; + Otile.store_rows(dst, gqa_splits * D, q_rows); + if (d_group == 0 && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = tm + sm + i * 8; + if (row < n_rows) { + const size_t po = (qbh * n_rows + row) * gqa_splits + split_idx; + sums[po] = sum_score[i]; + maxs[po] = max_score[i] == Limits::finite_min + ? Limits::finite_min + : max_score[i] * M_LN2_F; + } + } + } + } +} diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 1e07d03..6deae10 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -191,6 +191,23 @@ int kq_fa_verify_splits( return std::min({splits, cap, 128}); } +// Split count for sdpa_fa_verify on the tensor-op pass 1 when the caller +// passes 0. That kernel runs a split in a fraction of the time of the +// simdgroup-matrix kernel, so the float32 partials and their merge weigh +// more: the count aims at about 64 threadgroups over the kv heads, at least +// 4 and at most 64, and drops while a split would hold under 64 keys. Tuned +// on M5 Max at 1 to 32 kv heads. +int kq_fa_verify_nax_splits(int n_heads, int kL) { + int splits = 64; + for (int h = 1; h < n_heads && splits > 4; h *= 2) { + splits /= 2; + } + while (splits > 4 && kL < 64 * splits) { + splits /= 2; + } + return splits; +} + // Split count for sdpa_fa_indexed when the caller passes 0: the decode // bucket over the M listed keys. The NAX kernel on the s, c and d classes // writes 32 heads by 512 float32 partials per split, so once about 512 @@ -689,9 +706,19 @@ void KQuantSDPAFAVerify::eval_gpu( int D = q.shape(3); int kL = kvarn ? (kvarn_n_attend_ ? kvarn_n_attend_ : kvarn_n_) : k.shape(2); int q_len = q_len_; + // Pass 1 on the tensor-op units (kq_sdpa_fa_verify_2pass_1_nax) at + // head_dim 128 and 256 on plain K/V: a 16-row strip of the fold per + // simdgroup row group, the head dim split over WN simdgroups. It writes + // the partials of the simdgroup-matrix kernel below. WN is 2 at 128, and + // at 256 it is 4 for folds of up to 32 rows and 2 above. + // KQ_FA_VERIFY_NAX=0 keeps that kernel, and KQ_FA_VERIFY_NAX_WN sets WN + // (2 or 4 at 256, 1 or 2 at 128). Both are read per call. + const bool nax_pass1 = !kvarn && (D == 256 || D == 128) && n_rows <= 64 && + kq_env_flag("KQ_FA_VERIFY_NAX", 1) != 0 && kq_is_nax_available(); int splits = splits_; if (splits == 0) { - splits = kq_fa_verify_splits(B, n_kv_heads, n_rows, D, kL, d); + splits = nax_pass1 ? kq_fa_verify_nax_splits(B * n_kv_heads, kL) + : kq_fa_verify_splits(B, n_kv_heads, n_rows, D, kL, d); } size_t k_head_stride = @@ -743,7 +770,40 @@ void KQuantSDPAFAVerify::eval_gpu( // per 8 tile rows ((BQ/8)*32 threads, BQ the smallest of 32, 48, 64 that // holds the fold); 512 runs the 256-thread d-split variant (BQ 32 only, // capped at the op). - { + // Pass 1 on the tensor-op units, set up above. + if (nax_pass1) { + const int wm = (n_rows + 15) / 16; + int wn = D == 128 ? 2 : wm <= 2 ? 4 : 2; + const int wn_env = kq_env_lever("KQ_FA_VERIFY_NAX_WN", 0); + if ((D == 256 && (wn_env == 2 || wn_env == 4)) || + (D == 128 && (wn_env == 1 || wn_env == 2))) { + wn = wn_env; + } + std::string kname = "kq_sdpa_fa_verify_2pass_1_nax_" + ts + "_" + + std::to_string(D) + "_wm" + std::to_string(wm) + "_wn" + + std::to_string(wn); + std::string hash = kname + "_s" + std::to_string(splits); + auto kernel = kq_get_kernel(d, kname, hash, fc); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(q, 0); + ce.set_input_array(k, 1); + ce.set_input_array(v, 2); + ce.set_output_array(partials, 3); + ce.set_output_array(sums, 4); + ce.set_output_array(maxs, 5); + ce.set_bytes(kL, 6); + ce.set_bytes(k_head_stride, 7); + ce.set_bytes(k_seq_stride, 8); + ce.set_bytes(v_head_stride, 9); + ce.set_bytes(v_seq_stride, 10); + ce.set_bytes(scale, 11); + ce.set_bytes(q_len, 12); + ce.set_bytes(n_rows, 13); + MTL::Size group_dims(32, wm * wn, 1); + MTL::Size grid_dims(n_kv_heads, B, splits); + ce.dispatch_threadgroups(grid_dims, group_dims); + } + if (!nax_pass1) { const int bq = n_rows <= 32 ? 32 : n_rows <= 48 ? 48 : 64; std::string kname = "kq_sdpa_fa_verify_2pass_1_" + ts + "_" + std::to_string(D) + "_bq" + std::to_string(bq); diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 47a79b7..9c45627 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -997,10 +997,22 @@ def _check_fa(D, qL, kL, dtype, Hkv=4, G=6, strided=False, splits=0): assert got.shape == q.shape +# At head_dim 128 and 256 with up to 64 folded rows, sdpa_fa_verify runs +# pass 1 on the tensor-op units on NAX GPUs, and KQ_FA_VERIFY_NAX=0 keeps +# the simdgroup-matrix kernel. The op reads the variable at each call, so +# these cases run both kernels against the reference. +@pytest.fixture(params=["0", "1"], ids=["simdgroup", "nax"]) +def fa_pass1(request, monkeypatch): + if request.param == "1" and not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", request.param) + return request.param + + @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -@pytest.mark.parametrize("qL", [2, 3, 4, 5, 6]) -def test_sdpa_fa_verify(D, qL, dtype): +@pytest.mark.parametrize("qL", [2, 3, 4, 5, 6, 7, 8]) +def test_sdpa_fa_verify(D, qL, dtype, fa_pass1): _check_fa(D, qL, kL=4096, dtype=dtype, G=4) @@ -1011,7 +1023,7 @@ def test_sdpa_fa_verify_qwen_geometry(): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_decode(D, dtype): +def test_sdpa_fa_verify_decode(D, dtype, fa_pass1): # q_len == 1: plain GQA decode on the matrix units (every folded row # attends the full KV). 122b shape: G16 x qL1 at hd256, 2 kv heads; # same fold on the hd512 d-split kernel. @@ -1024,39 +1036,39 @@ def test_sdpa_fa_verify_gemma_geometry(): @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_full_tile(D): +def test_sdpa_fa_verify_full_tile(D, fa_pass1): # n_rows == 32 fills the tile exactly (no padding rows), qL at the cap _check_fa(D, 8, kL=4096, dtype=mx.bfloat16, Hkv=2, G=4) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_partial_warp(D): +def test_sdpa_fa_verify_partial_warp(D, fa_pass1): # n_rows == 18: the third row strip covers rows 16..17 plus padding _check_fa(D, 6, kL=2048, dtype=mx.bfloat16, G=3) @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_strided_unaligned(D, dtype): +def test_sdpa_fa_verify_strided_unaligned(D, dtype, fa_pass1): # strided KV-cache prefix + a key length off every tile/split boundary _check_fa(D, 4, kL=3071, dtype=dtype, strided=True, splits=16) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_short_kv(D): +def test_sdpa_fa_verify_short_kv(D, fa_pass1): # kL small enough that most splits stage zero keys: their empty partials # (max = finite_min, sum = 0) must merge as weight zero _check_fa(D, 4, kL=17, dtype=mx.bfloat16, splits=16, G=6) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_min_kv(D): +def test_sdpa_fa_verify_min_kv(D, fa_pass1): # kL == qL floor: row 0 attends exactly one key, every row masked hard _check_fa(D, 4, kL=4, dtype=mx.bfloat16, G=6) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_causal_split_straddle(D): +def test_sdpa_fa_verify_causal_split_straddle(D, fa_pass1): # the last qL keys straddle a split boundary (splits=128, kL=4098 puts # keys 4096..4097 alone in the final split): that split is entirely past # the low rows' causal limits, exercising the dead-row guard in a @@ -1065,38 +1077,38 @@ def test_sdpa_fa_verify_causal_split_straddle(D): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) -def test_sdpa_fa_verify_bq64_full_tile(dtype): +def test_sdpa_fa_verify_bq64_full_tile(dtype, fa_pass1): # gqa16 x qL4 = 64 rows (qwen3.5-122b-a10b verify fold): fills the BQ=64 # tile exactly _check_fa(256, 4, kL=4096, dtype=dtype, Hkv=2, G=16) @pytest.mark.parametrize("G,qL", [(10, 6), (7, 7), (14, 4)]) -def test_sdpa_fa_verify_bq64_padded(G, qL): +def test_sdpa_fa_verify_bq64_padded(G, qL, fa_pass1): # 49..63 rows: BQ=64 with padding rows in the upper simdgroups _check_fa(256, qL, kL=2048, dtype=mx.bfloat16, Hkv=2, G=G) @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256]) -def test_sdpa_fa_verify_bq48_full_tile(D, dtype): +def test_sdpa_fa_verify_bq48_full_tile(D, dtype, fa_pass1): # gqa6 x qL8 = 48 rows (qwen3.x full attention at a DFlash 2 block of # 8): fills the BQ=48 tile exactly _check_fa(D, 8, kL=4096, dtype=dtype, Hkv=4, G=6) @pytest.mark.parametrize("G,qL", [(8, 5), (6, 6), (6, 7), (12, 4)]) -def test_sdpa_fa_verify_bq48_padded(G, qL): +def test_sdpa_fa_verify_bq48_padded(G, qL, fa_pass1): # 33..48 rows: BQ=48, padding rows in the top simdgroup below 48 _check_fa(256, qL, kL=2048, dtype=mx.bfloat16, Hkv=2, G=G) -def test_sdpa_fa_verify_bq48_strided_split_straddle(): +def test_sdpa_fa_verify_bq48_strided_split_straddle(fa_pass1): # the last qL keys alone in the final split, over a strided cache view _check_fa(256, 8, kL=4098, dtype=mx.bfloat16, Hkv=4, G=6, strided=True, splits=128) -def test_sdpa_fa_verify_bq64_strided_kv(): +def test_sdpa_fa_verify_bq64_strided_kv(fa_pass1): _check_fa(256, 4, kL=3071, dtype=mx.bfloat16, Hkv=2, G=16, strided=True, splits=16) @@ -1112,11 +1124,17 @@ def test_sdpa_fa_verify_bq64_strided_kv(): (64, 8, 4, 8, 8192, 32), # 1024 simdgroups at head_dim <= 128 ], ) -def test_sdpa_fa_verify_default_splits(D, Hkv, G, qL, kL, want): +def test_sdpa_fa_verify_default_splits(D, Hkv, G, qL, kL, want, monkeypatch): # On the large GPU classes splits=0 raises the count past the decode - # bucket to fill the GPU; elsewhere it keeps the bucket. + # bucket to fill the GPU; elsewhere it keeps the bucket. Counts of the + # simdgroup-matrix pass 1. + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") if mx.device_info()["architecture"][-1] not in "scd": want = 16 if kL <= 8192 else 32 + _check_default_splits(D, Hkv, G, qL, kL, want) + + +def _check_default_splits(D, Hkv, G, qL, kL, want): q, k, v = _make(1, Hkv, Hkv, G * qL, kL, D, mx.bfloat16, seed=kL + D, strided=False) scale = 1.0 / (D**0.5) got = kq.sdpa_fa_verify(q, k, v, scale, q_len=qL) @@ -1127,9 +1145,78 @@ def test_sdpa_fa_verify_default_splits(D, Hkv, G, qL, kL, want): assert _rel(got, ref) < REL_BOUND[mx.bfloat16] +@pytest.mark.parametrize( + "D,Hkv,G,qL,kL,want", + [ + (256, 1, 8, 4, 8192, 64), # 64 threadgroups on one kv head + (256, 2, 12, 2, 4096, 32), + (256, 4, 6, 8, 16385, 16), # no rise past 16k keys + (128, 8, 4, 4, 4096, 8), + (128, 16, 2, 4, 4096, 4), # never under 4 + (128, 32, 1, 8, 4096, 4), + (256, 4, 6, 4, 1023, 8), # at least 64 keys per split + (256, 2, 12, 2, 700, 8), + (256, 4, 6, 6, 17, 4), + (128, 3, 8, 2, 4096, 16), # 3 kv heads count as 4 + ], +) +def test_sdpa_fa_verify_nax_default_splits(D, Hkv, G, qL, kL, want): + # splits=0 on the tensor-op pass 1 aims at about 64 threadgroups. + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + _check_default_splits(D, Hkv, G, qL, kL, want) + + +# KQ_FA_VERIFY_NAX_WN sets the simdgroups across the head dim of the +# tensor-op pass 1: every value it takes, over row strips 1 to 4, partial +# strips, split counts that leave splits without keys, and a strided cache +# view with a key count off every tile. +@pytest.mark.parametrize( + "D,wn", [(128, "1"), (128, "2"), (256, "2"), (256, "4")], ids=str +) +@pytest.mark.parametrize( + "G,qL,kL,splits,strided", + [ + (6, 1, 3071, 0, True), + (6, 4, 3071, 16, True), + (6, 6, 4098, 128, False), + (6, 8, 300, 1, False), + (16, 4, 2048, 2, True), + (3, 7, 17, 16, False), + (1, 5, 5, 0, False), + (64, 1, 4300, 0, True), + ], +) +def test_sdpa_fa_verify_nax_wn(D, wn, G, qL, kL, splits, strided, monkeypatch): + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "1") + monkeypatch.setenv("KQ_FA_VERIFY_NAX_WN", wn) + for dtype in (mx.bfloat16, mx.float16): + _check_fa(D, qL, kL=kL, dtype=dtype, Hkv=2, G=G, strided=strided, splits=splits) + + +@pytest.mark.parametrize("D", [128, 256]) +@pytest.mark.parametrize("qL", [1, 2, 3, 4, 5, 6, 7, 8]) +def test_sdpa_fa_verify_nax_matches_simdgroup(D, qL, monkeypatch): + # Both pass 1 kernels write the same partials layout to the same merge: + # at one split count they agree to bfloat16 rounding of the output. + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + G = 6 + q, k, v = _make(1, 4, 4, G * qL, 4300, D, mx.bfloat16, seed=qL + D, strided=True) + scale = 1.0 / (D**0.5) + out = {} + for lever in ("0", "1"): + monkeypatch.setenv("KQ_FA_VERIFY_NAX", lever) + out[lever] = kq.sdpa_fa_verify(q, k, v, scale, q_len=qL, splits=16) + _eval_or_skip(out[lever]) + assert _rel(out["1"], out["0"]) < REL_BOUND[mx.bfloat16] + + @pytest.mark.parametrize("D", [64, 128]) @pytest.mark.parametrize("G", [32, 64]) -def test_sdpa_fa_verify_cascade_fold(D, G): +def test_sdpa_fa_verify_cascade_fold(D, G, fa_pass1): # shared-prefix cascade: B batch rows folded kv-head-major into the row # axis at q_len 1 (every row attends the full prefix); G=64 fills the # BQ=64 tile at the register-resident head dims From 846912e62c1db752bed4a10f81764ac540ec8c21 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 03:59:57 -0700 Subject: [PATCH 04/10] test(kvarn): the materialize reference keeps the simdgroup pass 1 of sdpa_fa_verify --- tests/test_kvarn_sdpa.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_kvarn_sdpa.py b/tests/test_kvarn_sdpa.py index ac68fad..9e576bf 100644 --- a/tests/test_kvarn_sdpa.py +++ b/tests/test_kvarn_sdpa.py @@ -6,8 +6,8 @@ therefore materializes the cache (sink + dequantized records + live) into a plain fp16 KV and runs sdpa_decode_gqa with identical splits/tile_c: outputs must match bit for bit, including lse. The reference keeps the simdgroup -pass 1 (KQ_GQA_NAX=0), since at head dim 512 the plain op may otherwise take -the tensor-op pass, whose rounding differs. +pass 1 (KQ_GQA_NAX=0, KQ_FA_VERIFY_NAX=0), since on NAX GPUs the plain ops +may otherwise take a tensor-op pass, whose rounding differs. """ from __future__ import annotations @@ -32,6 +32,7 @@ @pytest.fixture(autouse=True) def _simdgroup_reference(monkeypatch): monkeypatch.setenv("KQ_GQA_NAX", "0") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") def quantize_head(x, bits, kind): From 421d7e6a7d5b415239746f163659b3d4feed9092 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:49:33 -0700 Subject: [PATCH 05/10] fix(sdpa): test the kvarn gate of the nax verify pass, qualify kvarn bit-exactness, warn on an out-of-set KQ_FA_VERIFY_NAX_WN --- bindings.cpp | 10 +++++--- docs/kernels.md | 16 +++++++----- .../mlx/backend/metal/kernels/kq_verify_nax.h | 5 ++-- src/kquant_sdpa.cpp | 11 +++++--- tests/test_kvarn_sdpa.py | 25 ++++++++++++++++++- tests/test_sdpa.py | 22 ++++++++++++++++ 6 files changed, 73 insertions(+), 16 deletions(-) diff --git a/bindings.cpp b/bindings.cpp index 7e3d066..f15632d 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -1866,9 +1866,13 @@ NB_MODULE(_ext, m) { GQA fold as sdpa_fa_verify (q [1, Hkv, G*q_len, D], kv-major heads, original q_len; G*q_len <= 64, or 32 at head_dim 512) in the rotated domain, and the sealed records dequantize at tile stage - exactly as in sdpa_decode_gqa_kvarn, so the output matches - sdpa_fa_verify over the materialized cache bit for bit. Operands, - n, n_attend and full_visibility follow sdpa_decode_gqa_kvarn; the + exactly as in sdpa_decode_gqa_kvarn. The key pass runs on the + simdgroup-matrix kernel at every head dim, so the output matches + sdpa_fa_verify over the materialized cache bit for bit where that + call runs the same kernel (KQ_FA_VERIFY_NAX=0). Where + sdpa_fa_verify takes its NAX key pass (head_dim 128 and 256 on NAX + GPUs), the two agree to rounding. Operands, n, n_attend and + full_visibility follow sdpa_decode_gqa_kvarn; the result is rotated-domain like the query (un-rotate it). This is the verify-width route: the vector kernel's cost climbs steeply past two queries, while the matrix tile prices extra rows at nearly diff --git a/docs/kernels.md b/docs/kernels.md index 354a877..00d8d3a 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -223,9 +223,9 @@ Tuning levers (defaults are right for normal use): the weights; `benchmarks/bench_dit_shapes.py` runs its tile arms this way. It reads the environment on every dispatch, so it is for single-threaded A/B runs that set these variables between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, `KQ_GLU_MIN_ROWS`), the - `quantized_matmul_multi` lever (`KQ_MATMUL_MULTI`), the `quantized_matmul_lora` levers (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the - `quantized_matmul_residual` lever (`KQ_QMM_RESID`), the dequantize route lever - (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are + `quantized_matmul_multi` lever (`KQ_MATMUL_MULTI`), the `quantized_matmul_lora` levers + (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the `quantized_matmul_residual` lever (`KQ_QMM_RESID`), + the dequantize route lever (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are read when the op is called and builds its graph, not when a kernel is dispatched, so a compiled function keeps the kernel choice of its trace. Unset (off) by default, and then each of them keeps its first read for the life of the process. Read once per process. @@ -1066,7 +1066,8 @@ into a packed layout first. about 64 threadgroups over the KV heads, between 4 and 64, and halves while a split would hold fewer than 64 keys. The two kernels round differently, so their outputs are not bit-identical. `KQ_FA_VERIFY_NAX=0` keeps the simdgroup-matrix kernel, and `KQ_FA_VERIFY_NAX_WN` sets the head - dim split (2 or 4 at 256, 1 or 2 at 128); both are read at each call. + dim split (2 or 4 at 256, 1 or 2 at 128; another value warns once on stderr and keeps the + default); both are read at each call. - **`sdpa_fa_indexed`** - index-gathered attention over one shared K/V latent at head dim 512, the absorbed-MLA sparse decode step: query j of `q [1, Hq, Q, 512]` attends the rows of `kv [1, 1, N, 512]` that `idx[j]` (int32, -1 pads) lists, so the selected rows are read once from the latent @@ -1096,8 +1097,11 @@ their cache layout. row's region map follows its own end and `tail_rows` is the per-row form of `n_attend`, so a ragged batch merges per-row tails the same way. Head dims 128, 256 and 512, q_len 1 to 4. - **`sdpa_fa_verify_kvarn`** - the same matrix-unit verify pass over a KVarN cache: sealed records - dequantize at tile stage through the loaders `sdpa_decode_gqa_kvarn` uses, so the result matches - `sdpa_fa_verify` over the materialized cache bit for bit. The verify-width route (q_len 2 to 8): + dequantize at tile stage through the loaders `sdpa_decode_gqa_kvarn` uses, and its key pass runs + on the simdgroup-matrix kernel at every head dim. The result matches `sdpa_fa_verify` over the + materialized cache bit for bit where that call runs the same kernel (`KQ_FA_VERIFY_NAX=0`). Where + `sdpa_fa_verify` takes its NAX key pass (head dims 128 and 256 on NAX GPUs), the two agree to + rounding. The verify-width route (q_len 2 to 8): the vector kernel's cost climbs steeply past two queries, the matrix tile prices extra rows at nearly zero. Head dims 128, 256 and 512; `n_attend` / `full_visibility` / `return_lse` as the decode op. diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index 57469b9..8be02b3 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -947,8 +947,9 @@ METAL_FUNC void kq_verify_nax_impl( // Several weights that share x in one dispatch: weight i owns the // threadgroups [tg_end[i - 1], tg_end[i]) of grid x, each weight from its // own buffer and into its own output, with one split count over grid z. -// part_stride[i] is output i's partial stride. Unused weight slots repeat -// slot 0 and own no threadgroup. +// part_stride[i] is output i's partial stride. Unused weight slots bind +// slot 0's buffers, and their tg_end is the grid end, so they own no +// threadgroup. #define KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(CODEC, TRAITS) \ template \ [[kernel]] void kq_##CODEC##_verify_nax_multi( \ diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 6deae10..f7666c6 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -774,10 +774,13 @@ void KQuantSDPAFAVerify::eval_gpu( if (nax_pass1) { const int wm = (n_rows + 15) / 16; int wn = D == 128 ? 2 : wm <= 2 ? 4 : 2; - const int wn_env = kq_env_lever("KQ_FA_VERIFY_NAX_WN", 0); - if ((D == 256 && (wn_env == 2 || wn_env == 4)) || - (D == 128 && (wn_env == 1 || wn_env == 2))) { - wn = wn_env; + // Another value warns once and keeps the default. + static const std::vector wn_128 = {"1", "2"}; + static const std::vector wn_256 = {"2", "4"}; + const auto& wn_set = D == 128 ? wn_128 : wn_256; + const int wn_env = kq_env_choice("KQ_FA_VERIFY_NAX_WN", wn_set, -1); + if (wn_env >= 0) { + wn = std::stoi(wn_set[wn_env]); } std::string kname = "kq_sdpa_fa_verify_2pass_1_nax_" + ts + "_" + std::to_string(D) + "_wm" + std::to_string(wm) + "_wn" + diff --git a/tests/test_kvarn_sdpa.py b/tests/test_kvarn_sdpa.py index 9e576bf..51474fe 100644 --- a/tests/test_kvarn_sdpa.py +++ b/tests/test_kvarn_sdpa.py @@ -773,7 +773,8 @@ def test_dwide_rejects_wrong_tile_c_and_stale_shapes(d): # sdpa_fa_verify_kvarn stages the same values as the decode kernel through # the shared kvarn loaders, and everything downstream is sdpa_fa_verify # unchanged, so it must match sdpa_fa_verify over the materialized twin bit -# for bit, lse included. +# for bit, lse included, where that call runs the simdgroup pass 1 (the +# autouse fixture above). def fold(q, hkv, ql): @@ -857,6 +858,28 @@ def test_fa_verify_explicit_splits(): assert_bit_equal(out, ref) +@pytest.mark.parametrize("d", [128, 256]) +def test_fa_verify_default_lever_keeps_kvarn_kernel(d, monkeypatch): + # At the default KQ_FA_VERIFY_NAX, plain sdpa_fa_verify takes the NAX + # pass 1 at these head dims, which reads plain K/V. The kvarn op keeps + # the simdgroup kernel and its split rule, so it gives the bits it gives + # with the lever off, and those of the simdgroup reference. The lever + # is read when the kernel is dispatched, so each arm is evaluated + # before the lever changes. + st = build_state(1, 2, 1000, 6, 6, seed=123 + d, d=d) + q = make_q(1, 8, 4, d=d) + monkeypatch.delenv("KQ_FA_VERIFY_NAX") + out_default, ref_nax = run_fa_pair(st, q, 1000, 6, 6, 4, return_lse=True) + mx.eval(out_default, ref_nax) + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") + out_off, ref = run_fa_pair(st, q, 1000, 6, 6, 4, return_lse=True) + mx.eval(out_off, ref) + # The plain op did take the NAX pass at the default. + assert not np.array_equal(np.array(ref_nax[0]), np.array(ref[0])) + assert_bit_equal(out_default, out_off) + assert_bit_equal(out_default, ref) + + def test_fa_verify_n_attend_full_visibility(): # Body walk of a tail merge at verify width: the first 800 of 1000 keys # with the clamp lifted equals a q_len-1 fold (every row sees every diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 9c45627..0110361 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -1196,6 +1196,28 @@ def test_sdpa_fa_verify_nax_wn(D, wn, G, qL, kL, splits, strided, monkeypatch): _check_fa(D, qL, kL=kL, dtype=dtype, Hkv=2, G=G, strided=strided, splits=splits) +# A KQ_FA_VERIFY_NAX_WN value outside the head dim's set warns once on +# stderr and keeps the default split: the same bits as the lever unset. +@pytest.mark.parametrize("D,bad", [(128, "4"), (256, "1"), (256, "3"), (128, "x")]) +def test_sdpa_fa_verify_nax_wn_out_of_set(D, bad, monkeypatch, capfd): + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "1") + q, k, v = _make(1, 2, 2, 6 * 4, 1000, D, mx.bfloat16, seed=D, strided=False) + scale = 1.0 / (D**0.5) + out = {} + for value in (None, bad): + if value is None: + monkeypatch.delenv("KQ_FA_VERIFY_NAX_WN", raising=False) + else: + monkeypatch.setenv("KQ_FA_VERIFY_NAX_WN", value) + out[value] = kq.sdpa_fa_verify(q, k, v, scale, q_len=4, splits=8) + _eval_or_skip(out[value]) + assert mx.array_equal(out[None], out[bad]).item() + err = capfd.readouterr().err + assert f"KQ_FA_VERIFY_NAX_WN={bad} is not" in err, err + + @pytest.mark.parametrize("D", [128, 256]) @pytest.mark.parametrize("qL", [1, 2, 3, 4, 5, 6, 7, 8]) def test_sdpa_fa_verify_nax_matches_simdgroup(D, qL, monkeypatch): From 3c272399f3081b44390e699a069e6e5aee886682 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 04:43:23 -0700 Subject: [PATCH 06/10] perf(matmul): verify_nax reads two activation steps per 16-byte load where x is 16-byte aligned --- CHANGELOG.md | 3 + docs/kernels.md | 21 ++-- .../mlx/backend/metal/kernels/kq_verify_nax.h | 98 +++++++++++++++---- src/kquant_matmul.cpp | 22 ++++- tests/qmm_probe.py | 22 ++++- tests/test_matmul_multi.py | 36 ++++++- tests/test_nax_smallm.py | 57 +++++++++++ 7 files changed, 226 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 96822bd..ed2c7da 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Changed - `kq.sdpa_fa_verify` runs its key pass on the NAX tensor-op units (M5) at head dims 128 and 256, so verify attention over a deep cache takes less time. +- The `verify_nax` matmul kernels of `q4_0`, `q8_0`, the k-quants and `iq4_*` + read two steps of activations per load, so speculative verify at 2 to 8 rows + takes less time. ## [0.5.0] diff --git a/docs/kernels.md b/docs/kernels.md index 00d8d3a..b8373a1 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -217,12 +217,12 @@ the two kernels are within 3% of each other on every shape measured. Tuning levers (defaults are right for normal use): - `KQ_TUNE_LIVE` - `1` makes the NAX tile levers (`KQ_NAX_SMALL_BM`, `KQ_NAX_BM128`, - `KQ_NAX_BM128_ODD_M`, `KQ_NAX_DB64`, the `_BK128` levers, `KQ_NAX_SWIZZLE`), the tile overrides - (`KQ_NAX_TILE`, `KQ_QMM_TILE`, `KQ_CONV_TILE`), `KQ_SPLITK_F32` and `KQ_QMM_LOG` read their - variables on every dispatch, so one process can switch tiles between calls on one resident copy of - the weights; `benchmarks/bench_dit_shapes.py` runs its tile arms this way. It reads the - environment on every dispatch, so it is for single-threaded A/B runs that set these variables - between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, `KQ_GLU_MIN_ROWS`), the + `KQ_NAX_BM128_ODD_M`, `KQ_NAX_DB64`, the `_BK128` levers, `KQ_NAX_SWIZZLE`, `KQ_VNAX_X16`), the + tile overrides (`KQ_NAX_TILE`, `KQ_QMM_TILE`, `KQ_CONV_TILE`), `KQ_SPLITK_F32` and `KQ_QMM_LOG` + read their variables on every dispatch, so one process can switch tiles between calls on one + resident copy of the weights; `benchmarks/bench_dit_shapes.py` runs its tile arms this way. It + reads the environment on every dispatch, so it is for single-threaded A/B runs that set these + variables between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, `KQ_GLU_MIN_ROWS`), the `quantized_matmul_multi` lever (`KQ_MATMUL_MULTI`), the `quantized_matmul_lora` levers (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the `quantized_matmul_residual` lever (`KQ_QMM_RESID`), the dequantize route lever (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are @@ -397,6 +397,10 @@ Tuning levers (defaults are right for normal use): float kernels only. Unset (off) by default. Read live per call. - `KQ_VERIFY_NAX_Q4_0_SB` - probe lever: `0` keeps `q4_0` on the two-block `verify_nax` kernel at a K that is a multiple of 256, where the eight-block kernel runs by default. Read live per call. +- `KQ_VNAX_X16` - probe lever: `0` keeps the `verify_nax` kernels that load their activations per + NAX op on one 8-byte load per op where the activations start on a 16-byte boundary, in place of + one 16-byte load per two ops. The output is the same. Read once per process unless + `KQ_TUNE_LIVE=1`. - `KQ_MV_EXT_SB` / `KQ_MV_EXT_NX` / `KQ_MV_EXT_HD` - `mv_ext` activation-traffic experiments: shuffle-broadcast (`1`), wide nxpsg (`16`/`32`), half-precision chunk dots (`1`). q6_k M 4-12 only. `HD` measured +4-5% at M 8; the rest flat to negative on M5 Max. Kept as probes. Default @@ -554,7 +558,10 @@ eight-block kernel at a K that is a multiple of 256 and the two-block kernel els kernels fold each 16- or 32-weight sub-block's scale, and its min where the codec has one, into the half weights. Loading a whole K step of activations ahead measured faster for `pq2_0` and the two-block `q4_0` kernel, while the kernels with larger decoded words load them one NAX op at a time, -since loading ahead ran `q4_k` 1.6-2x slower. +since loading ahead ran `q4_k` 1.6-2x slower. Two consecutive NAX ops read 8 contiguous activations +of a row, so where the activations start on a 16-byte boundary these kernels read them in one +16-byte load per two NAX ops, and elsewhere in one 8-byte load per op. The fewer loads take less +time in a chain of dependent calls, and both forms give the same bits. Split counts for `pq2_0`, `q8_0` and the two-block `q4_0` kernel are the largest divisor of the K steps at or under their target, at most 16. The two-block kernel's K walk of 640 per simdgroup diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index 8be02b3..c39c2f0 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -42,6 +42,16 @@ // pair(words, f): fragment f's element pair, block scale applied // xstep(xb, q, s): left elements 0..3 of step s, xb the unit's // first activation of the lane's row +// xoff(q, s): where xstep reads, in elements from xb. A codec that loads +// one step at a time reads steps s and s + 1 (s even) from 8 +// contiguous elements at a 16-byte boundary of the row, so where x +// starts on one (kq_vnax_x16) it loads both in one 16-byte load. + +// Function constant: x starts on a 16-byte boundary. Unset keeps the +// 8-byte load per step. +constant bool kq_vnax_x16_fc [[function_constant(40)]]; +constant bool kq_vnax_x16 = + is_function_constant_defined(kq_vnax_x16_fc) && kq_vnax_x16_fc; MLX_MTL_CONST int KQ_VNAX_NSG = 4; @@ -53,6 +63,10 @@ struct KqPq2_0Nax { static constant constexpr int block_bytes = KQ_PQ2_0_BLOCK_BYTES; static constant constexpr int ub = 1; static constant constexpr bool x_ahead = true; + // Unused: x_ahead codecs load their steps through xstep. + static METAL_FUNC int xoff(short, short) { + return 0; + } struct Words { uint w[2]; half d; @@ -98,6 +112,10 @@ struct KqQ2_0Nax { static constant constexpr int block_bytes = KQ_Q2_0_BLOCK_BYTES; static constant constexpr int ub = 2; static constant constexpr bool x_ahead = true; + // Unused: x_ahead codecs load their steps through xstep. + static METAL_FUNC int xoff(short, short) { + return 0; + } struct Words { uint w; half d; @@ -163,9 +181,12 @@ struct KqQ4_0Nax { half2(1032.0h)) * o.d; } + static METAL_FUNC int xoff(short q, short s) { + return 16 * s + 4 * q; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 16 * s + 4 * q); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -216,9 +237,12 @@ struct KqQ4_0SbNax { half2(1032.0h)) * o.d[b]; } + static METAL_FUNC int xoff(short q, short s) { + return 64 * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 64 * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -266,10 +290,12 @@ struct KqQ8_0Nax { half2(1152.0h)) * o.d; } + static METAL_FUNC int xoff(short q, short s) { + return KQ_Q8_0_GROUP * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + KQ_Q8_0_GROUP * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -346,10 +372,12 @@ struct KqQ4_KNax { const uint wv = (hi ? (o.w[s & 7] >> 4) : o.w[s & 7]) & 0x0F0F0F0Fu; return kq_vnax_kpair(wv >> (8 * (f & 1)), o.ds, o.dm, hi); } + static METAL_FUNC int xoff(short q, short s) { + return 64 * q + 32 * (s >> 3) + 4 * (s & 7); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 64 * q + 32 * (s >> 3) + 4 * (s & 7)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -409,11 +437,12 @@ struct KqQ5_KNax { lo4 | (((o.hb >> (2 * i + (hi ? 1 : 0))) & 0x01010101u) << 4); return kq_vnax_kpair(q5 >> (8 * (f & 1)), o.ds, o.dm, hi); } + static METAL_FUNC int xoff(short q, short s) { + return 64 * (q >> 1) + 16 * (q & 1) + 32 * (s >> 2) + 4 * (s & 3); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 64 * (q >> 1) + 16 * (q & 1) + - 32 * (s >> 2) + 4 * (s & 3)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -485,11 +514,12 @@ struct KqQ6_KNax { half2(1056.0h); return v * (hi ? o.sc.y : o.sc.x); } + static METAL_FUNC int xoff(short q, short s) { + return 32 * (q & 1) + 16 * (q >> 1) + 64 * (s >> 2) + 4 * (s & 3); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 32 * (q & 1) + 16 * (q >> 1) + - 64 * (s >> 2) + 4 * (s & 3)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -561,10 +591,12 @@ struct KqQ3_KNax { const half2 sc = o.sc[j >> 1]; return v * ((j & 1) ? sc.y : sc.x); } + static METAL_FUNC int xoff(short q, short s) { + return 32 * (s >> 1) + 8 * q + 4 * (s & 1); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 32 * (s >> 1) + 8 * q + 4 * (s & 1)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -624,10 +656,12 @@ struct KqQ2_KNax { const half m = (j & 1) ? o.dm[k].y : o.dm[k].x; return fma(v, half2(d), -half2(m)); } + static METAL_FUNC int xoff(short q, short s) { + return 32 * (s >> 1) + 8 * q + 4 * (s & 1); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 32 * (s >> 1) + 8 * q + 4 * (s & 1)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -742,9 +776,12 @@ struct KqIq4XsNax { const uint src = (((s >> 2) & 1) ? (w >> 4) : w) >> (8 * (f & 1)); return kq_vnax_iq4pair(src, o.dl[b]); } + static METAL_FUNC int xoff(short q, short s) { + return 64 * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 64 * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -788,9 +825,12 @@ struct KqIq4NlNax { const uint src = (((s >> 2) & 1) ? (w >> 4) : w) >> (8 * (f & 1)); return kq_vnax_iq4pair(src, o.d[b]); } + static METAL_FUNC int xoff(short q, short s) { + return 64 * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 64 * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -811,6 +851,7 @@ METAL_FUNC void kq_verify_nax_impl( constexpr int BK = Codec::block_k; constexpr int UB = Codec::ub; constexpr int SPB = BK / 16; + static_assert(Codec::x_ahead || SPB % 2 == 0, "x16 loads step pairs"); const int n0 = (tid.x * KQ_VNAX_NSG + simd_gid) * 32; if (n0 >= N) { return; @@ -874,13 +915,28 @@ METAL_FUNC void kq_verify_nax_impl( } } } + uint4 xu = uint4(0); #pragma unroll for (short u = 0; u < UB; ++u) { #pragma unroll for (short s = 0; s < SPB; ++s) { - const vec xv = Codec::x_ahead ? xa[u][s] - : xrow ? Codec::template xstep(xr + kb + u * BK, q, s) - : vec(0); + vec xv; + if (Codec::x_ahead) { + xv = xa[u][s]; + } else if (kq_vnax_x16) { + // Steps s and s + 1 from one 16-byte load, in xstep's order. + if ((s & 1) == 0) { + xu = xrow + ? *(const device uint4*)(xr + kb + u * BK + Codec::xoff(q, s)) + : uint4(0); + } + const vec a = as_type>((s & 1) ? xu.z : xu.x); + const vec c = as_type>((s & 1) ? xu.w : xu.y); + xv = vec(a[0], c[0], a[1], c[1]); + } else { + xv = xrow ? Codec::template xstep(xr + kb + u * BK, q, s) + : vec(0); + } ca[0] = xv[0]; ca[1] = xv[1]; ca[2] = xv[2]; diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 6cbee28..bb66d95 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -1496,6 +1496,24 @@ kq_verify_nax_fits(const array& w, int K, int kstep, const std::string& t) { // device memory per block and nothing staged. Same partial/fold shape as // verify_mma. Non-batched transpose shapes, bf16 or f16 activations, NAX // GPUs. +// KQ_VNAX_X16=0 keeps verify_nax on one 8-byte activation load per step. +// Read once per process unless KQ_TUNE_LIVE=1. +static bool kq_vnax_x16_enabled() { + static const bool cached = kq_env_int("KQ_VNAX_X16", 1) != 0; + return kq_tune_live() ? kq_env_int("KQ_VNAX_X16", 1) != 0 : cached; +} + +// verify_nax pipeline: the kernels that load one activation step at a +// time load two in one 16-byte load where x starts on a 16-byte boundary +// (function constant 40, kq_vnax_x16). +static MTL::ComputePipelineState* +kq_verify_nax_pipeline(Device& d, const std::string& kname, const array& x) { + bool x16 = reinterpret_cast(x.data()) % 16 == 0 && + kq_vnax_x16_enabled(); + mx::metal::MTLFCList fc = {{&x16, MTL::DataType::DataTypeBool, 40}}; + return kq_get_kernel(d, kname, kname + (x16 ? "_x16" : "_x8"), fc); +} + void verify_nax( const array& x, const array& w, @@ -1527,7 +1545,7 @@ void verify_nax( const std::string kname = kq_verify_nax_kname( kquant_type, type_string, group_size, bits, kstep, pp.suffix); - auto kernel = kq_get_kernel(d, kname); + auto kernel = kq_verify_nax_pipeline(d, kname, x); ce.set_compute_pipeline_state(kernel); int c = 0; @@ -1614,7 +1632,7 @@ void verify_nax_multi( const std::string type_string = kq_type_string(x.dtype()); const std::string kname = kq_verify_nax_multi_kname( kquant_type, type_string, group_size, bits, kstep, pp.suffix); - auto kernel = kq_get_kernel(d, kname); + auto kernel = kq_verify_nax_pipeline(d, kname, x); ce.set_compute_pipeline_state(kernel); int c = 0; diff --git a/tests/qmm_probe.py b/tests/qmm_probe.py index 01c33ee..edaf995 100644 --- a/tests/qmm_probe.py +++ b/tests/qmm_probe.py @@ -18,8 +18,9 @@ comparison over every row of x's flattened leading dims, B rows at a time, so a call with many rows holds float64 copies of B rows; it reports ``checked``, the number of rows it read. Every call gets -``digest``, a hash of its output bits, to compare calls across probe -processes. +``digest``, a hash of its output bits (of every output on a ``"multi"`` +call), to compare calls across probe processes. ``"x_offset": E`` passes +x as a row-contiguous view E elements into its buffer. ``"design": "lead"`` draws nonnegative weights and activations whose first K column is 1024. ``"design": "big"`` draws small weights and dense @@ -244,6 +245,13 @@ def ref_of(lo): x = activations( list(c["shape"]), dtype, c.get("layout", "dense"), c.get("seed", i), design ) + if c.get("x_offset"): + # The same values as a row-contiguous view x_offset elements into + # its buffer. + off = c["x_offset"] + flat = mx.concatenate([mx.zeros((off,), dtype), x.reshape(-1)]) + mx.eval(flat) + x = flat[off:].reshape(x.shape) os.write(2, f"[probe] {i}\n".encode()) res = {"i": i, "error": None, "rel": None, "ulp": None} if c.get("multi"): @@ -259,6 +267,16 @@ def ref_of(lo): outs.append(None) results.append(res) continue + bits = [y.view(mx.uint16) if y.dtype in DT.values() else y for y in ys] + h = hashlib.sha1() + for b in bits: + h.update(memoryview(np.array(b))) + res["digest"] = h.hexdigest() + if not spec.get("check", True): + outs.append(None) + results.append(res) + mx.clear_cache() + continue xn = as64(x) res["rel"] = max( rel(as64(y), xn @ dequantized(c["codec"], n, c["shape"][-1], design).T) diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py index fd84f1a..4937d82 100644 --- a/tests/test_matmul_multi.py +++ b/tests/test_matmul_multi.py @@ -22,7 +22,7 @@ import numpy as np import pytest from kqref import synth_wire -from qmm_probe import force_cpu +from qmm_probe import force_cpu, run_probe import mlx_kquant as kq from mlx_kquant.codec_geometry import CODEC_GEOMETRY @@ -369,6 +369,40 @@ def test_forced_verify_nax_fuses_narrow_weights(): assert len(multi) == 1 and "verify_nax_multi" in multi[0], multi +@NAX_ONLY +@pytest.mark.parametrize("codec", ["q4_0", "q4_k", "q6_k", "iq4_xs"]) +def test_x16_load_bits(codec): + # x at offset 0 (paired 16-byte activation loads), at offset 0 with + # KQ_VNAX_X16=0 and at 8 bytes (one load per step), and at 16 bytes: + # the same bits from one verify_nax_multi dispatch. + arms = [(0, None), (0, "0"), (4, None), (8, None)] + cases = [(m, sp) for m in range(2, 9) for sp in ("1", "2")] + calls = [ + { + "codec": codec, + "shape": [m, 1024], + "n": 1000, + "multi": [48, 300], + "seed": m, + "x_offset": off, + "env": {"KQ_VNAX_X16": lever, "KQ_VERIFY_NAX_SPLITS": sp}, + } + for m, sp in cases + for off, lever in arms + ] + env = {"KQ_QMM_LOG": "1", "KQ_TUNE_LIVE": "1", "KQ_VERIFY_NAX": "2"} + res = run_probe(calls, env=env, check=False) + for i, (m, sp) in enumerate(cases): + group = res[len(arms) * i : len(arms) * (i + 1)] + for (off, lever), r in zip(arms, group, strict=True): + tag = (codec, m, sp, off, lever) + assert r["error"] is None, (tag, r["error"]) + assert [n for n, _ in r["kernels"]] and all( + "verify_nax_multi" in n for n, _ in r["kernels"] + ), (tag, r["kernels"]) + assert r["digest"] == group[0]["digest"], tag + + @NAX_ONLY def test_lever_read_once_per_process(): # KQ_MATMUL_MULTI keeps its first read, as KQ_GLU, unless KQ_TUNE_LIVE=1. diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index 943e653..740ac2e 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -1321,6 +1321,63 @@ def test_verify_nax_disable_nax(codec, monkeypatch): _sweep(codec, w, s, ref_w, 1000, ms=[3, 5, 8]) +# The verify_nax codecs that read one activation step at a time read two +# steps in one 16-byte load where x starts on a 16-byte boundary. Each arm +# multiplies the same values in the same order: x at offset 0 (the paired +# load), at offset 0 with KQ_VNAX_X16=0 and at 8 bytes (one load per +# step), and at 16 bytes (the paired load off the buffer start). +X16_ARMS = [(0, None), (0, "0"), (4, None), (8, None)] +X16_SPLITS = [None, "3"] +X16_DTYPES = ["bfloat16", "float16"] + + +@functools.cache +def _x16_digests(): + cases = [ + (codec, m, dtype, splits) + for codec in VNAX_CODECS + VNAX_IQ4 + for m in range(2, 9) + for dtype in X16_DTYPES + for splits in X16_SPLITS + ] + calls = [ + { + "codec": codec, + "shape": [m, K], + "n": 1000, + "dtype": dtype, + "seed": m, + "x_offset": off, + "env": {"KQ_VNAX_X16": lever, "KQ_VERIFY_NAX_SPLITS": splits}, + } + for codec, m, dtype, splits in cases + for off, lever in X16_ARMS + ] + env = { + "KQ_QMM_LOG": "1", + "KQ_TUNE_LIVE": "1", + "KQ_QMM_ROUTE": "verify_nax", + "KQ_QMM_ROUTE_STRICT": "1", + } + res = run_probe(calls, env=env, check=False) + n = len(X16_ARMS) + return {c: res[n * i : n * i + n] for i, c in enumerate(cases)} + + +@pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") +@pytest.mark.parametrize("codec", VNAX_CODECS + VNAX_IQ4) +def test_verify_nax_x16_bits(codec): + for (c, m, dtype, splits), arms in _x16_digests().items(): + if c != codec: + continue + tag = f"{codec} M{m} {dtype} splits {splits}" + for (off, lever), r in zip(X16_ARMS, arms, strict=True): + assert r["error"] is None, (tag, off, lever, r["error"]) + ((name, _),) = r["kernels"] + assert "_verify_nax_" in name, (tag, name) + assert r["digest"] == arms[0]["digest"], (tag, off, lever) + + # verify_nax folds the block scale into half weights and accumulates in # float32, so one activation channel near the top of the half range stays # finite and within the contract. From f4d788d9525fa91df7d17b160079382abc38ac7e Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:07:33 -0700 Subject: [PATCH 07/10] fix(matmul): test verify_nax x16 at every 2-byte x offset, name the x16/x8 variant in KQ_QMM_LOG, warn on an out-of-set KQ_VNAX_X16, check the load pair geometry at compile time --- CHANGELOG.md | 6 +-- docs/kernels.md | 29 +++++++------ .../mlx/backend/metal/kernels/kq_verify_nax.h | 43 +++++++++++++------ src/kquant_matmul.cpp | 30 ++++++++----- tests/qmm_probe.py | 4 +- tests/test_env_levers.py | 28 ++++++++++++ tests/test_matmul_multi.py | 23 ++++++---- tests/test_nax_smallm.py | 16 +++++-- 8 files changed, 127 insertions(+), 52 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ed2c7da..efb833b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,9 +14,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Changed - `kq.sdpa_fa_verify` runs its key pass on the NAX tensor-op units (M5) at head dims 128 and 256, so verify attention over a deep cache takes less time. -- The `verify_nax` matmul kernels of `q4_0`, `q8_0`, the k-quants and `iq4_*` - read two steps of activations per load, so speculative verify at 2 to 8 rows - takes less time. +- On NAX GPUs (M5), the `verify_nax` matmul kernels of `q4_0` at a K that is a + multiple of 256, `q8_0`, the k-quants and `iq4_*` read two steps of activations + per load, so speculative verify at 2 to 8 rows takes less time. ## [0.5.0] diff --git a/docs/kernels.md b/docs/kernels.md index b8373a1..065a24b 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -217,18 +217,19 @@ the two kernels are within 3% of each other on every shape measured. Tuning levers (defaults are right for normal use): - `KQ_TUNE_LIVE` - `1` makes the NAX tile levers (`KQ_NAX_SMALL_BM`, `KQ_NAX_BM128`, - `KQ_NAX_BM128_ODD_M`, `KQ_NAX_DB64`, the `_BK128` levers, `KQ_NAX_SWIZZLE`, `KQ_VNAX_X16`), the - tile overrides (`KQ_NAX_TILE`, `KQ_QMM_TILE`, `KQ_CONV_TILE`), `KQ_SPLITK_F32` and `KQ_QMM_LOG` - read their variables on every dispatch, so one process can switch tiles between calls on one - resident copy of the weights; `benchmarks/bench_dit_shapes.py` runs its tile arms this way. It - reads the environment on every dispatch, so it is for single-threaded A/B runs that set these - variables between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, `KQ_GLU_MIN_ROWS`), the - `quantized_matmul_multi` lever (`KQ_MATMUL_MULTI`), the `quantized_matmul_lora` levers - (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the `quantized_matmul_residual` lever (`KQ_QMM_RESID`), - the dequantize route lever (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are - read when the op is called and builds its graph, not when a kernel is dispatched, so a compiled - function keeps the kernel choice of its trace. Unset (off) by default, and then each of them keeps - its first read for the life of the process. Read once per process. + `KQ_NAX_BM128_ODD_M`, `KQ_NAX_DB64`, the `_BK128` levers, `KQ_NAX_SWIZZLE`), the `verify_nax` load + lever (`KQ_VNAX_X16`), the tile overrides (`KQ_NAX_TILE`, `KQ_QMM_TILE`, `KQ_CONV_TILE`), + `KQ_SPLITK_F32` and `KQ_QMM_LOG` read their variables on every dispatch, so one process can switch + tiles between calls on one resident copy of the weights; `benchmarks/bench_dit_shapes.py` runs its + tile arms this way. It reads the environment on every dispatch, so it is for single-threaded A/B + runs that set these variables between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, + `KQ_GLU_MIN_ROWS`), the `quantized_matmul_multi` lever (`KQ_MATMUL_MULTI`), the + `quantized_matmul_lora` levers (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the + `quantized_matmul_residual` lever (`KQ_QMM_RESID`), the dequantize route lever + (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are read when the op is called and + builds its graph, not when a kernel is dispatched, so a compiled function keeps the kernel choice + of its trace. Unset (off) by default, and then each of them keeps its first read for the life of + the process. Read once per process. - `KQ_NAX_TILE` - `bm,bn,wm,wn` pins the NAX `qmm_t` tile of every transpose call that reaches it, in place of the BM=32, BM=64 and BM=128 policy. BK stays 64, except where a `_BK128` lever sets it. The `_db` variant still follows `KQ_NAX_DB64` on the bm64 bn64 wm2 wn2 tile. A tile the @@ -399,8 +400,8 @@ Tuning levers (defaults are right for normal use): a K that is a multiple of 256, where the eight-block kernel runs by default. Read live per call. - `KQ_VNAX_X16` - probe lever: `0` keeps the `verify_nax` kernels that load their activations per NAX op on one 8-byte load per op where the activations start on a 16-byte boundary, in place of - one 16-byte load per two ops. The output is the same. Read once per process unless - `KQ_TUNE_LIVE=1`. + one 16-byte load per two ops. The output is the same. In `KQ_QMM_LOG` the kernel name + carries the variant, `_x16` or `_x8`. Read once per process unless `KQ_TUNE_LIVE=1`. - `KQ_MV_EXT_SB` / `KQ_MV_EXT_NX` / `KQ_MV_EXT_HD` - `mv_ext` activation-traffic experiments: shuffle-broadcast (`1`), wide nxpsg (`16`/`32`), half-precision chunk dots (`1`). q6_k M 4-12 only. `HD` measured +4-5% at M 8; the rest flat to negative on M5 Max. Kept as probes. Default diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index c39c2f0..4efae33 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -55,6 +55,21 @@ constant bool kq_vnax_x16 = MLX_MTL_CONST int KQ_VNAX_NSG = 4; +// The 16-byte load's geometry: for every lane-quad q and even step s, +// steps s and s + 1 sit in 8 contiguous elements at an 8-element boundary. +template +constexpr bool kq_vnax_x16_pairs() { + for (short q = 0; q < 8; ++q) { + for (short s = 0; s < SPB; s += 2) { + if (Codec::xoff(q, s) % 8 != 0 || + Codec::xoff(q, s + 1) != Codec::xoff(q, s) + 4) { + return false; + } + } + } + return true; +} + // PQ2_0: lane-quad q owns the two 16-code words at qs bytes 8q..8q+7. // Pair f is code f % 8 of word f / 8 in each half (KqPq2_0Mma). struct KqPq2_0Nax { @@ -64,7 +79,7 @@ struct KqPq2_0Nax { static constant constexpr int ub = 1; static constant constexpr bool x_ahead = true; // Unused: x_ahead codecs load their steps through xstep. - static METAL_FUNC int xoff(short, short) { + static constexpr METAL_FUNC int xoff(short, short) { return 0; } struct Words { @@ -113,7 +128,7 @@ struct KqQ2_0Nax { static constant constexpr int ub = 2; static constant constexpr bool x_ahead = true; // Unused: x_ahead codecs load their steps through xstep. - static METAL_FUNC int xoff(short, short) { + static constexpr METAL_FUNC int xoff(short, short) { return 0; } struct Words { @@ -181,7 +196,7 @@ struct KqQ4_0Nax { half2(1032.0h)) * o.d; } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 16 * s + 4 * q; } template @@ -237,7 +252,7 @@ struct KqQ4_0SbNax { half2(1032.0h)) * o.d[b]; } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 64 * q + 4 * s; } template @@ -290,7 +305,7 @@ struct KqQ8_0Nax { half2(1152.0h)) * o.d; } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return KQ_Q8_0_GROUP * q + 4 * s; } template @@ -372,7 +387,7 @@ struct KqQ4_KNax { const uint wv = (hi ? (o.w[s & 7] >> 4) : o.w[s & 7]) & 0x0F0F0F0Fu; return kq_vnax_kpair(wv >> (8 * (f & 1)), o.ds, o.dm, hi); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 64 * q + 32 * (s >> 3) + 4 * (s & 7); } template @@ -437,7 +452,7 @@ struct KqQ5_KNax { lo4 | (((o.hb >> (2 * i + (hi ? 1 : 0))) & 0x01010101u) << 4); return kq_vnax_kpair(q5 >> (8 * (f & 1)), o.ds, o.dm, hi); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 64 * (q >> 1) + 16 * (q & 1) + 32 * (s >> 2) + 4 * (s & 3); } template @@ -514,7 +529,7 @@ struct KqQ6_KNax { half2(1056.0h); return v * (hi ? o.sc.y : o.sc.x); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 32 * (q & 1) + 16 * (q >> 1) + 64 * (s >> 2) + 4 * (s & 3); } template @@ -591,7 +606,7 @@ struct KqQ3_KNax { const half2 sc = o.sc[j >> 1]; return v * ((j & 1) ? sc.y : sc.x); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 32 * (s >> 1) + 8 * q + 4 * (s & 1); } template @@ -656,7 +671,7 @@ struct KqQ2_KNax { const half m = (j & 1) ? o.dm[k].y : o.dm[k].x; return fma(v, half2(d), -half2(m)); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 32 * (s >> 1) + 8 * q + 4 * (s & 1); } template @@ -776,7 +791,7 @@ struct KqIq4XsNax { const uint src = (((s >> 2) & 1) ? (w >> 4) : w) >> (8 * (f & 1)); return kq_vnax_iq4pair(src, o.dl[b]); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 64 * q + 4 * s; } template @@ -825,7 +840,7 @@ struct KqIq4NlNax { const uint src = (((s >> 2) & 1) ? (w >> 4) : w) >> (8 * (f & 1)); return kq_vnax_iq4pair(src, o.d[b]); } - static METAL_FUNC int xoff(short q, short s) { + static constexpr METAL_FUNC int xoff(short q, short s) { return 64 * q + 4 * s; } template @@ -851,7 +866,9 @@ METAL_FUNC void kq_verify_nax_impl( constexpr int BK = Codec::block_k; constexpr int UB = Codec::ub; constexpr int SPB = BK / 16; - static_assert(Codec::x_ahead || SPB % 2 == 0, "x16 loads step pairs"); + static_assert( + Codec::x_ahead || (SPB % 2 == 0 && kq_vnax_x16_pairs()), + "x16 loads step pairs"); const int n0 = (tid.x * KQ_VNAX_NSG + simd_gid) * 32; if (n0 >= N) { return; diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index bb66d95..8c1a64d 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -1499,19 +1499,27 @@ kq_verify_nax_fits(const array& w, int K, int kstep, const std::string& t) { // KQ_VNAX_X16=0 keeps verify_nax on one 8-byte activation load per step. // Read once per process unless KQ_TUNE_LIVE=1. static bool kq_vnax_x16_enabled() { - static const bool cached = kq_env_int("KQ_VNAX_X16", 1) != 0; - return kq_tune_live() ? kq_env_int("KQ_VNAX_X16", 1) != 0 : cached; + static const bool cached = kq_env_flag("KQ_VNAX_X16", 1) != 0; + return kq_tune_live() ? kq_env_flag("KQ_VNAX_X16", 1) != 0 : cached; } // verify_nax pipeline: the kernels that load one activation step at a // time load two in one 16-byte load where x starts on a 16-byte boundary -// (function constant 40, kq_vnax_x16). -static MTL::ComputePipelineState* -kq_verify_nax_pipeline(Device& d, const std::string& kname, const array& x) { +// (function constant 40, kq_vnax_x16). vname, kname with the variant +// (_x16 or _x8) ahead of its partials suffix, names the pipeline and the +// KQ_QMM_LOG line. +static MTL::ComputePipelineState* kq_verify_nax_pipeline( + Device& d, + const std::string& kname, + const std::string& suffix, + const array& x, + std::string& vname) { bool x16 = reinterpret_cast(x.data()) % 16 == 0 && kq_vnax_x16_enabled(); mx::metal::MTLFCList fc = {{&x16, MTL::DataType::DataTypeBool, 40}}; - return kq_get_kernel(d, kname, kname + (x16 ? "_x16" : "_x8"), fc); + vname = kname.substr(0, kname.size() - suffix.size()) + + (x16 ? "_x16" : "_x8") + suffix; + return kq_get_kernel(d, kname, vname, fc); } void verify_nax( @@ -1545,7 +1553,8 @@ void verify_nax( const std::string kname = kq_verify_nax_kname( kquant_type, type_string, group_size, bits, kstep, pp.suffix); - auto kernel = kq_verify_nax_pipeline(d, kname, x); + std::string vname; + auto kernel = kq_verify_nax_pipeline(d, kname, pp.suffix, x, vname); ce.set_compute_pipeline_state(kernel); int c = 0; @@ -1560,7 +1569,7 @@ void verify_nax( ce.set_bytes(part_stride, c++); MTL::Size group_dims(128, 1, 1); MTL::Size grid_dims((N + 127) / 128, 1, splits); - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(vname, M, N, K, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; @@ -1632,7 +1641,8 @@ void verify_nax_multi( const std::string type_string = kq_type_string(x.dtype()); const std::string kname = kq_verify_nax_multi_kname( kquant_type, type_string, group_size, bits, kstep, pp.suffix); - auto kernel = kq_verify_nax_pipeline(d, kname, x); + std::string vname; + auto kernel = kq_verify_nax_pipeline(d, kname, pp.suffix, x, vname); ce.set_compute_pipeline_state(kernel); int c = 0; @@ -1657,7 +1667,7 @@ void verify_nax_multi( ce.set_bytes(part_stride.data(), 4, c++); MTL::Size group_dims(128, 1, 1); MTL::Size grid_dims(tgs, 1, splits); - kq_log_qmm(kname, M, static_cast(rows), K, grid_dims); + kq_log_qmm(vname, M, static_cast(rows), K, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; diff --git a/tests/qmm_probe.py b/tests/qmm_probe.py index edaf995..72d6dae 100644 --- a/tests/qmm_probe.py +++ b/tests/qmm_probe.py @@ -247,8 +247,10 @@ def ref_of(lo): ) if c.get("x_offset"): # The same values as a row-contiguous view x_offset elements into - # its buffer. + # its buffer. An empty cache keeps the buffer from carrying the + # slack that makes the op copy the view to a fresh one. off = c["x_offset"] + mx.clear_cache() flat = mx.concatenate([mx.zeros((off,), dtype), x.reshape(-1)]) mx.eval(flat) x = flat[off:].reshape(x.shape) diff --git a/tests/test_env_levers.py b/tests/test_env_levers.py index 8dee9aa..4d1e77c 100644 --- a/tests/test_env_levers.py +++ b/tests/test_env_levers.py @@ -9,6 +9,8 @@ from __future__ import annotations +import re + import pytest from qmm_probe import force_cpu, run_probe @@ -86,3 +88,29 @@ def test_verify_nax_q4_0_sb_values(): ((name, _),) = r["kernels"] assert name.startswith("kquant_q4_0_sb_verify_nax_") == sb, (v, name) assert _warnings(res[2], "KQ_VERIFY_NAX_Q4_0_SB", "typo"), res[2]["stderr"] + + +# KQ_VNAX_X16 is on unless 0, read live under KQ_TUNE_LIVE: x at a 16-byte +# boundary takes the _x16 pipeline unless it is off. Values outside 0 and 1 +# warn and keep it on. +@pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") +def test_vnax_x16_values(): + values = [ + (None, True, False), + ("", True, False), + ("on", True, False), + ("1", True, False), + ("off", False, False), + ("0", False, False), + ("2", True, True), + ("x", True, True), + ] + call = {"codec": "q4_k", "shape": [8, 3840], "n": 1000, "dtype": "bfloat16"} + calls = [dict(call, env={"KQ_VNAX_X16": v}) for v, _, _ in values] + res = run_probe(calls, env={**LOG, "KQ_QMM_ROUTE": "verify_nax"}) + for (v, x16, warns), r in zip(values, res, strict=True): + assert r["error"] is None, (v, r["error"]) + ((name, _),) = r["kernels"] + assert "_verify_nax_" in name, (v, name) + assert re.search("_x16(_pf32)?$" if x16 else "_x8(_pf32)?$", name), (v, name) + assert bool(_warnings(r, "KQ_VNAX_X16", v)) == warns, (v, r["stderr"]) diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py index 4937d82..3d568f9 100644 --- a/tests/test_matmul_multi.py +++ b/tests/test_matmul_multi.py @@ -15,6 +15,7 @@ import json import os +import re import subprocess import sys @@ -370,12 +371,14 @@ def test_forced_verify_nax_fuses_narrow_weights(): @NAX_ONLY -@pytest.mark.parametrize("codec", ["q4_0", "q4_k", "q6_k", "iq4_xs"]) -def test_x16_load_bits(codec): +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", [*CODECS, "pq2_0", "q2_0"]) +def test_x16_load_bits(codec, dtype): # x at offset 0 (paired 16-byte activation loads), at offset 0 with - # KQ_VNAX_X16=0 and at 8 bytes (one load per step), and at 16 bytes: - # the same bits from one verify_nax_multi dispatch. - arms = [(0, None), (0, "0"), (4, None), (8, None)] + # KQ_VNAX_X16=0, at every 2-byte offset from 2 to 14 bytes (one load + # per step) and at 16 bytes: the same bits from one verify_nax_multi + # dispatch. + arms = [(0, None), (0, "0"), *[(e, None) for e in range(1, 9)]] cases = [(m, sp) for m in range(2, 9) for sp in ("1", "2")] calls = [ { @@ -383,6 +386,7 @@ def test_x16_load_bits(codec): "shape": [m, 1024], "n": 1000, "multi": [48, 300], + "dtype": dtype, "seed": m, "x_offset": off, "env": {"KQ_VNAX_X16": lever, "KQ_VERIFY_NAX_SPLITS": sp}, @@ -397,9 +401,11 @@ def test_x16_load_bits(codec): for (off, lever), r in zip(arms, group, strict=True): tag = (codec, m, sp, off, lever) assert r["error"] is None, (tag, r["error"]) - assert [n for n, _ in r["kernels"]] and all( - "verify_nax_multi" in n for n, _ in r["kernels"] - ), (tag, r["kernels"]) + end = "_x16(_pf32)?$" if lever is None and off % 8 == 0 else "_x8(_pf32)?$" + names = [n for n, _ in r["kernels"]] + assert names and all( + "verify_nax_multi" in n and re.search(end, n) for n in names + ), (tag, names) assert r["digest"] == group[0]["digest"], tag @@ -429,6 +435,7 @@ def test_lever_off(): _LEVER_OFF_CHILD = r""" import os +import re import mlx.core as mx import numpy as np if os.environ.get("KQUANT_FORCE_CPU"): diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index 740ac2e..8e37003 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -24,6 +24,7 @@ import functools import json import os +import re import subprocess import sys @@ -1324,9 +1325,17 @@ def test_verify_nax_disable_nax(codec, monkeypatch): # The verify_nax codecs that read one activation step at a time read two # steps in one 16-byte load where x starts on a 16-byte boundary. Each arm # multiplies the same values in the same order: x at offset 0 (the paired -# load), at offset 0 with KQ_VNAX_X16=0 and at 8 bytes (one load per -# step), and at 16 bytes (the paired load off the buffer start). -X16_ARMS = [(0, None), (0, "0"), (4, None), (8, None)] +# load), at offset 0 with KQ_VNAX_X16=0, at every 2-byte offset from 2 to +# 14 bytes (one load per step) and at 16 bytes (the paired load off the +# buffer start). +X16_ARMS = [(0, None), (0, "0"), *[(e, None) for e in range(1, 9)]] + + +def x16_variant(off, lever): + # The KQ_QMM_LOG name end of the pipeline an arm takes. + return "_x16(_pf32)?$" if lever is None and off % 8 == 0 else "_x8(_pf32)?$" + + X16_SPLITS = [None, "3"] X16_DTYPES = ["bfloat16", "float16"] @@ -1375,6 +1384,7 @@ def test_verify_nax_x16_bits(codec): assert r["error"] is None, (tag, off, lever, r["error"]) ((name, _),) = r["kernels"] assert "_verify_nax_" in name, (tag, name) + assert re.search(x16_variant(off, lever), name), (tag, off, lever, name) assert r["digest"] == arms[0]["digest"], (tag, off, lever) From 86d2262faae4c8171aceb9b3192074b98bd59b41 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:25:52 -0700 Subject: [PATCH 08/10] test(matmul): split kernel dispatch runs every verify_nax kernel in its x16 and x8 pipelines --- tests/test_split_kernel_dispatch.py | 87 ++++++++++++++++++----------- 1 file changed, 53 insertions(+), 34 deletions(-) diff --git a/tests/test_split_kernel_dispatch.py b/tests/test_split_kernel_dispatch.py index fdaf692..bd2e3a5 100644 --- a/tests/test_split_kernel_dispatch.py +++ b/tests/test_split_kernel_dispatch.py @@ -5,11 +5,14 @@ dtypes, on an N that is and is not a multiple of the column tile, with partials in the activation dtype and in float (KQ_SPLITK_F32), and with the half verify_mma kernels (KQ_VERIFY_MMA_HALF), and run quantized_matmul_multi -for its joined verify_nax kernels. KQ_QMM_LOG names the kernel each call -ran. Together the calls must run every name the host enumerator lists for -the codec, and each result must meet the R1 bound against float64 with exact -kq.dequantize weights, so a kernel that compiles to nothing, or that no -call reaches, fails here. +for its joined verify_nax kernels. Every verify_nax call also runs with x +2 bytes into its buffer, so each verify_nax kernel runs both its 16-byte +(_x16) and 8-byte (_x8) x load pipelines. KQ_QMM_LOG names the pipeline +each call ran. Together the calls must run every name the host enumerator +lists for the codec, each verify_nax name in both pipelines, and each +result must meet the R1 bound against float64 with exact kq.dequantize +weights, so a kernel that compiles to nothing, or that no call reaches, +fails here. """ from __future__ import annotations @@ -52,6 +55,7 @@ def _no_nax(): "KQ_NAX_QMV": None, "KQ_VERIFY_NAX": None, "KQ_VERIFY_NAX_SPLITS": None, + "KQ_VNAX_X16": None, } # (family the arm may be declined for, rows, K, env). The split-K tile is # BM=8 through M 8, BM=16 through M 16 and the plain tile above; the NAX @@ -91,6 +95,16 @@ def _host_names(): return out +def _pipelines(name): + """The pipeline names KQ_QMM_LOG gives a host kernel name: a verify_nax + kernel has an _x16 and an _x8 pipeline, named ahead of its partials + suffix.""" + if "_verify_nax_" not in name: + return {name} + base, suffix = (name[:-5], "_pf32") if name.endswith("_pf32") else (name, "") + return {base + "_x16" + suffix, base + "_x8" + suffix} + + HOST = {} if _no_nax() else _host_names() @@ -117,37 +131,41 @@ def test_every_split_kernel_runs(codec): for f32 in ("0", "1"): for n in (256, 250): for family, m, k, env in arms: - calls.append( - { - "codec": codec, - "shape": [m, k], - "n": n, - "dtype": dtype, - "seed": 1, - "env": {**CLEAR, **env, "KQ_SPLITK_F32": f32}, - } - ) - families.append(family) + for off in (0, 1) if family == "_verify_nax_" else (0,): + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "dtype": dtype, + "seed": 1, + "x_offset": off, + "env": {**CLEAR, **env, "KQ_SPLITK_F32": f32}, + } + ) + families.append(family) for m, k, splits in MULTI_ARMS: if k % CODEC_GEOMETRY[codec][3]: continue - calls.append( - { - "codec": codec, - "shape": [m, k], - "n": n, - "multi": [n + 64], - "dtype": dtype, - "seed": 1, - "env": { - **CLEAR, - **MULTI, - "KQ_VERIFY_NAX_SPLITS": splits, - "KQ_SPLITK_F32": f32, - }, - } - ) - families.append(None) + for off in (0, 1): + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "multi": [n + 64], + "dtype": dtype, + "seed": 1, + "x_offset": off, + "env": { + **CLEAR, + **MULTI, + "KQ_VERIFY_NAX_SPLITS": splits, + "KQ_SPLITK_F32": f32, + }, + } + ) + families.append(None) res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"}) ran = set() for c, family, r in zip(calls, families, res, strict=True): @@ -159,4 +177,5 @@ def test_every_split_kernel_runs(codec): half = c["env"].get("KQ_VERIFY_MMA_HALF") == "1" assert r["rel"] < (HALF if half else R1[c["dtype"]]), (c, r["rel"]) ran |= {name for name, _ in r["kernels"]} - assert sorted(host - ran) == [] + want = set().union(*(_pipelines(n) for n in host)) + assert sorted(want - ran) == [] From cc498664e793b4788bc338241ec6c1b4f0d1a963 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:24:05 -0700 Subject: [PATCH 09/10] fix(kernels): read operands that start or step off their kernel's vector boundary from an aligned copy, verify_nax reads x in place on an 8-byte boundary, log the copies and matmul x alignment in KQ_QMM_LOG --- CHANGELOG.md | 4 + benchmarks/dit_common.py | 2 +- docs/integration.md | 3 +- docs/kernels.md | 26 +- src/kquant_conv_nax.cpp | 5 +- src/kquant_dense.cpp | 61 ++ src/kquant_dsa_indexer.cpp | 11 +- src/kquant_dsa_sparse_attn.cpp | 10 +- src/kquant_gather.cpp | 15 +- src/kquant_hc_glue.cpp | 47 +- src/kquant_internal.h | 32 + src/kquant_latent_fp4.cpp | 7 +- src/kquant_matmul.cpp | 154 +++- src/kquant_matmul_fused.cpp | 7 +- src/kquant_matmul_lora.cpp | 7 +- src/kquant_matmul_resid.cpp | 8 +- src/kquant_metal_internal.h | 8 + src/kquant_moe_glu.cpp | 33 +- src/kquant_moe_mix.cpp | 6 +- src/kquant_norm_fused.cpp | 13 +- src/kquant_sdpa.cpp | 88 ++- src/kquant_sdpa_sparse_decode.cpp | 33 +- src/kquant_skinny_mv.cpp | 5 +- tests/qmm_probe.py | 41 +- tests/test_matmul_multi.py | 48 +- tests/test_nax_smallm.py | 36 +- tests/test_prism_mv_range.py | 5 +- tests/test_split_kernel_dispatch.py | 6 +- tests/test_splitk_cap.py | 10 +- tests/test_x_offset_views.py | 1030 +++++++++++++++++++++++++++ 30 files changed, 1582 insertions(+), 179 deletions(-) create mode 100644 tests/test_x_offset_views.py diff --git a/CHANGELOG.md b/CHANGELOG.md index efb833b..17aa076 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). fewer rows by GPU core count, for A/B runs on wide ones such as a vocab head. - `KQ_DEQUANT_ROUTE` turns the dequantize route off or forces it on any GPU, and `KQ_QMM_LOG` logs each call that takes it. +- `KQ_QMM_LOG` also logs the mat-vec, gather and MoE dispatches with `xa`, the + activation's byte offset modulo 16, and each operand an op copies to align it. - `KQ_SPLITK_F32=1` makes bfloat16 split-K matmuls, MTP verify among them, round once instead of once per K slice. `0` stores float16 partials again, which can return inf. @@ -113,6 +115,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - After `patch_mlx_lm_load()`, `mlx_lm.utils.load_model` builds a kquant checkpoint with the caller's `get_model_classes` instead of the stock class. `mlx_kquant.loader.load` takes it too. +- float16 and bfloat16 inputs that start off a 16-byte boundary, such as a row + slice, no longer give wrong output from matmul, norm, attention and DSA ops. ## [0.4.18] diff --git a/benchmarks/dit_common.py b/benchmarks/dit_common.py index 0d0db26..cc77395 100644 --- a/benchmarks/dit_common.py +++ b/benchmarks/dit_common.py @@ -134,7 +134,7 @@ class KernelLog: while a block runs, by pointing fd 2 at a temporary file.""" LINE = re.compile( - r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+)$" + r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+) xa=\d+$" ) def __enter__(self): diff --git a/docs/integration.md b/docs/integration.md index 77476d4..e79d772 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -38,7 +38,8 @@ boundary. It raises from the call when the weight is evaluated and row-contiguou `mx.eval` otherwise. Arrays that MLX allocates and tensors from `kq.load_gguf` always meet the boundary, as do slices of them that start on a block boundary. A view that starts inside a block can miss it, and a copy in a new buffer, such as `mx.array(np.array(w))`, runs. The CPU kernels accept -any start. +any start. Activations, KV caches and the other float operands can start anywhere: a GPU op copies +one that starts off the boundary of its kernel's vector loads to a new buffer first. ## Codecs diff --git a/docs/kernels.md b/docs/kernels.md index 065a24b..881341c 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -249,14 +249,19 @@ Tuning levers (defaults are right for normal use): included, for the codecs, rows and range of its table, so that tests run its code on any machine. Any other value warns once and reads as `auto`. `auto` by default. Read when the op builds its graph, once per process unless `KQ_TUNE_LIVE=1`. -- `KQ_QMM_LOG` - set to any value, logs the kernel name, shape and grid of every NAX and ALU qmm - tile dispatch, every split-K dispatch, every `verify_mma` and `verify_nax` dispatch, split or not, - and every fused `quantized_matmul_glu`, `quantized_matmul_lora` and `quantized_matmul_residual` - dispatch, and every `conv2d_3x3` dispatch to stderr, one `[kq qmm]` line each. The grid depth is - the K slice count. A call that takes the dequantize route logs one - `[kq dequant] M= N= K=` line when the op builds its graph. The mat-vec - routes (`qmv`, `verify_qmv`, `mv_ext`) do not log. Unset (off) by default. Read once per process - unless `KQ_TUNE_LIVE=1`. +- `KQ_QMM_LOG` - set to any value, logs every quantized matmul dispatch to stderr, one + `[kq qmm] M= N= K= grid=,, xa=` line each: the NAX and ALU qmm + tiles, the mat-vec routes (`qmv`, `qvm`, `verify_qmv`, `mv_ext`), split-K, `verify_mma` and + `verify_nax`, split or not, the fused `quantized_matmul_glu`, `quantized_matmul_lora` and + `quantized_matmul_residual` kernels, `skinny_matmul`, the `gather_qmm` and `gather_qmm_seg` + routes, the fused MoE gathers, `gather_mix` and `conv2d_3x3`. The grid depth is the K slice + count. `xa` is the byte offset of the activation's start modulo 16, as the kernel reads it. A + gather or MoE line gives the rows and last dim of the output as M and N. Any GPU op that copies + an operand to an aligned buffer, because the operand starts or steps off the boundary of its + kernel's vector loads, logs one `[kq align] xin=` line, where `n` is the + operand's byte offset modulo that boundary. A call that takes the dequantize route logs one + `[kq dequant] M= N= K=` line when the op builds its graph. Unset (off) by + default. Read once per process unless `KQ_TUNE_LIVE=1`. - `KQ_SPLITK_F32` - partial dtype of the split-K routes (`splitk`, `nax_splitk`, the `KQ_QMM_MIDM` band, and `verify_mma` and `verify_nax` when they split K). Float partials let the fold round the result once, as the unsplit kernels do. Partials in the activation dtype round @@ -561,8 +566,9 @@ half weights. Loading a whole K step of activations ahead measured faster for `p two-block `q4_0` kernel, while the kernels with larger decoded words load them one NAX op at a time, since loading ahead ran `q4_k` 1.6-2x slower. Two consecutive NAX ops read 8 contiguous activations of a row, so where the activations start on a 16-byte boundary these kernels read them in one -16-byte load per two NAX ops, and elsewhere in one 8-byte load per op. The fewer loads take less -time in a chain of dependent calls, and both forms give the same bits. +16-byte load per two NAX ops, and where they start on an 8-byte boundary in one 8-byte load per op. +Activations that start elsewhere are read from a copy on a 16-byte boundary. The fewer loads take +less time in a chain of dependent calls, and both forms give the same bits. Split counts for `pq2_0`, `q8_0` and the two-block `q4_0` kernel are the largest divisor of the K steps at or under their target, at most 16. The two-block kernel's K walk of 640 per simdgroup diff --git a/src/kquant_conv_nax.cpp b/src/kquant_conv_nax.cpp index 3d42e73..7d4350a 100644 --- a/src/kquant_conv_nax.cpp +++ b/src/kquant_conv_nax.cpp @@ -138,14 +138,15 @@ void KQuantConv3x3::eval_gpu( if (conv_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, O, 9 * C, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname); diff --git a/src/kquant_dense.cpp b/src/kquant_dense.cpp index fdeaac6..cebe267 100644 --- a/src/kquant_dense.cpp +++ b/src/kquant_dense.cpp @@ -17,6 +17,7 @@ // reads and a row-major copy otherwise. An op uses them for an operand that // is unevaluated on most calls, such as an activation, which then costs no // node when it is dense. +#include #include #include #include @@ -152,6 +153,56 @@ mx::array kq_row_major_gpu(const mx::array& a, const mx::Stream& s) { return kq_dense_copy_gpu(a, s); } +namespace { + +// The copy of an operand that does not start on its kernels' alignment, +// logged under KQ_QMM_LOG with the byte offset it came at. +mx::array aligned_copy( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align) { + if (kq_qmm_log_on()) { + std::fprintf( + stderr, + "[kq align] %s xin=%d\n", + what, + static_cast(a.offset() % align)); + } + return kq_dense_copy_gpu(a, s); +} + +} // namespace + +mx::array kq_x_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align) { + if (a.size() == 0 || (a.flags().row_contiguous && a.offset() % align == 0)) { + return a; + } + return aligned_copy(a, s, what, align); +} + +mx::array kq_strides_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align) { + if (a.size() == 0) { + return a; + } + const int nd = static_cast(a.ndim()); + bool ok = a.offset() % align == 0 && + (nd == 0 || a.shape(nd - 1) == 1 || a.strides()[nd - 1] == 1); + for (int i = 0; ok && i + 1 < nd; i++) { + ok = a.shape(i) == 1 || + (a.strides()[i] * static_cast(a.itemsize())) % align == 0; + } + return ok ? a : aligned_copy(a, s, what, align); +} + mx::array kq_unit_stride_gpu(const mx::array& a, const mx::Stream& s) { if (a.size() == 0 || a.ndim() == 0 || a.strides().back() == 1) { return a; @@ -186,6 +237,16 @@ mx::array kq_row_major_gpu(const mx::array&, const mx::Stream&) { throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); } +mx::array +kq_x_aligned_gpu(const mx::array&, const mx::Stream&, const char*, int) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +mx::array +kq_strides_aligned_gpu(const mx::array&, const mx::Stream&, const char*, int) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + mx::array kq_unit_stride_gpu(const mx::array&, const mx::Stream&) { throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); } diff --git a/src/kquant_dsa_indexer.cpp b/src/kquant_dsa_indexer.cpp index 875fa91..940b08b 100644 --- a/src/kquant_dsa_indexer.cpp +++ b/src/kquant_dsa_indexer.cpp @@ -230,11 +230,16 @@ void KQDsaIndexerScoreDecode::eval_gpu( auto& d = mx::metal::device(s.device); auto& ce = mx::metal::get_command_encoder(s); - const auto& q = inputs[0]; + // q and k are read as 8-wide vectors. + const int align = 8 * static_cast(inputs[0].itemsize()); + const mx::array q = + kq_x_aligned_gpu(inputs[0], s, "dsa_indexer_score_decode q", align); // [B, P * D], the row-major key block. A batch is row-major from the // builder. One sequence comes through a reshape that can keep a view - // with an element stride other than 1, which is copied here. - const mx::array k = kq_unit_stride_gpu(inputs[1], s); + // with an element stride other than 1, which is copied here, as is a + // view whose start or batch stride is off the vector alignment. + const mx::array k = + kq_strides_aligned_gpu(inputs[1], s, "dsa_indexer_score_decode k", align); const auto& weights = inputs[2]; auto& out = outputs[0]; diff --git a/src/kquant_dsa_sparse_attn.cpp b/src/kquant_dsa_sparse_attn.cpp index 415ac71..383d6b5 100644 --- a/src/kquant_dsa_sparse_attn.cpp +++ b/src/kquant_dsa_sparse_attn.cpp @@ -36,9 +36,13 @@ void KQDsaSparseAttention::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& q = inputs[0]; - const auto& local_kv = inputs[1]; - const auto& pooled = inputs[2]; + // q and the key rows are read as 16-byte vectors through their strides. + const mx::array q = + kq_strides_aligned_gpu(inputs[0], s, "dsa_sparse_attention q"); + const mx::array local_kv = + kq_strides_aligned_gpu(inputs[1], s, "dsa_sparse_attention local_kv"); + const mx::array pooled = + kq_strides_aligned_gpu(inputs[2], s, "dsa_sparse_attention pooled"); const auto& topk = inputs[3]; const auto& sinks = inputs[4]; auto& o = outputs[0]; diff --git a/src/kquant_gather.cpp b/src/kquant_gather.cpp index 8bf2b40..219e1a5 100644 --- a/src/kquant_gather.cpp +++ b/src/kquant_gather.cpp @@ -121,6 +121,7 @@ void gather_qmm_nax( c = add_strides_and_shapes(ce, false, x, w, scales, c); add_gather_strides_and_shapes(ce, lhs_indices, rhs_indices, c); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -199,6 +200,7 @@ void gather_qmm( c = add_strides_and_shapes(ce, false, x, w, scales, c); add_gather_strides_and_shapes(ce, lhs_indices, rhs_indices, c); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -255,6 +257,7 @@ void gather_qmv( c = add_strides_and_shapes(ce, false, x, w, scales, c); add_gather_strides_and_shapes(ce, lhs_indices, rhs_indices, c); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -358,6 +361,7 @@ void gather_qmm_rhs_nax( ce.set_bytes(N, c++); ce.set_bytes(K, c++); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -438,6 +442,7 @@ void gather_qmm_rhs( ce.set_bytes(N, c++); ce.set_bytes(K, c++); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -825,7 +830,8 @@ void KQuantGatherQMM::eval_gpu( // inputs are row-contiguous (ensured by the op): // x (float), w (uint8), scales, lhs_indices (uint32), rhs_indices (uint32). - const auto& x = inputs[0]; + // x is bound from a 16-byte boundary (kq_x_aligned_gpu). + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "gather_qmm x"); const auto& w = inputs[1]; const auto& scales = inputs[2]; const auto& lhs_indices = inputs[3]; @@ -1050,8 +1056,9 @@ void KQuantGatherQMMSeg::eval_gpu( auto& out = outputs[0]; // inputs are row-contiguous (ensured by the op): x [R, K] (float), - // w [E, N, bpr] (uint8), scales, map [T, 3] (uint32), counts [1]. - const auto& x = inputs[0]; + // w [E, N, bpr] (uint8), scales, map [T, 3] (uint32), counts [1]. x is + // bound from a 16-byte boundary (kq_x_aligned_gpu). + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "gather_qmm_seg x"); const auto& w = inputs[1]; const auto& scales = inputs[2]; const auto& map = inputs[3]; @@ -1107,6 +1114,7 @@ void KQuantGatherQMMSeg::eval_gpu( ce.set_bytes(K, c++); MTL::Size group_dims(32, 2, 2); MTL::Size grid_dims((N + bn - 1) / bn, map.shape(0), 1); + kq_log_dispatch(nname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); return; } @@ -1135,6 +1143,7 @@ void KQuantGatherQMMSeg::eval_gpu( ce.set_bytes(N, c++); MTL::Size group_dims(32, 2, 2); MTL::Size grid_dims((N + bn - 1) / bn, map.shape(0), 1); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/src/kquant_hc_glue.cpp b/src/kquant_hc_glue.cpp index b4458ff..cdacc40 100644 --- a/src/kquant_hc_glue.cpp +++ b/src/kquant_hc_glue.cpp @@ -89,6 +89,12 @@ void check_fn(const mx::array& fn, int D, const char* op) { namespace { +// The kernels read the stream, residual and collapsed rows, fn, the mixes +// and the base biases as vectors of four elements. +mx::array hc_vec(const mx::array& a, const mx::Stream& s, const char* what) { + return kq_x_aligned_gpu(a, s, what, 4 * static_cast(a.itemsize())); +} + // Threads per threadgroup for one launch: the smaller of the count the // kernel prefers and what this pipeline supports, rounded down to a whole // simdgroup. Register pressure puts the pipeline limit below the preferred @@ -175,7 +181,10 @@ void KQuantHcFrontReduce::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& x = inputs[0]; + // The operands the kernels read as vectors are bound from boundaries of + // four of their elements (hc_vec), here and in the evals below. + const mx::array x = hc_vec(inputs[0], s, "hc_front_reduce x"); + const mx::array fn = hc_vec(inputs[1], s, "hc_front_reduce fn"); int D = x.shape(-1); int rows = int(x.size() / (HC * D)); @@ -184,7 +193,7 @@ void KQuantHcFrontReduce::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); ce.set_input_array(x, 0); - ce.set_input_array(inputs[1], 1); + ce.set_input_array(fn, 1); ce.set_output_array(outputs[0], 2); ce.set_output_array(outputs[1], 3); ce.set_bytes(D, 4); @@ -201,7 +210,9 @@ void KQuantHcFrontExpandReduce::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& resid = inputs[1]; + const mx::array x_sub = hc_vec(inputs[0], s, "hc_front_expand_reduce x_sub"); + const mx::array resid = hc_vec(inputs[1], s, "hc_front_expand_reduce resid"); + const mx::array fn = hc_vec(inputs[4], s, "hc_front_expand_reduce fn"); int D = resid.shape(-1); int rows = int(resid.size() / (HC * D)); @@ -210,11 +221,11 @@ void KQuantHcFrontExpandReduce::eval_gpu( 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(inputs[0], 0); + ce.set_input_array(x_sub, 0); ce.set_input_array(resid, 1); ce.set_input_array(inputs[2], 2); ce.set_input_array(inputs[3], 3); - ce.set_input_array(inputs[4], 4); + ce.set_input_array(fn, 4); ce.set_output_array(outputs[0], 5); ce.set_output_array(outputs[1], 6); ce.set_output_array(outputs[2], 7); @@ -232,7 +243,9 @@ void KQuantHcSinkhornCollapse::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& x = inputs[0]; + const mx::array x = hc_vec(inputs[0], s, "hc_sinkhorn_collapse x"); + const mx::array mixes = hc_vec(inputs[1], s, "hc_sinkhorn_collapse mixes"); + const mx::array base = hc_vec(inputs[4], s, "hc_sinkhorn_collapse base"); int D = x.shape(-1); int rows = int(x.size() / (HC * D)); @@ -241,10 +254,10 @@ void KQuantHcSinkhornCollapse::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); ce.set_input_array(x, 0); - ce.set_input_array(inputs[1], 1); + ce.set_input_array(mixes, 1); ce.set_input_array(inputs[2], 2); ce.set_input_array(inputs[3], 3); - ce.set_input_array(inputs[4], 4); + ce.set_input_array(base, 4); ce.set_input_array(inputs[5], 5); ce.set_output_array(outputs[0], 6); ce.set_output_array(outputs[1], 7); @@ -265,7 +278,12 @@ void KQuantHcFrontExpandCollapse::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& resid = inputs[1]; + const mx::array x_sub = + hc_vec(inputs[0], s, "hc_front_expand_collapse x_sub"); + const mx::array resid = + hc_vec(inputs[1], s, "hc_front_expand_collapse resid"); + const mx::array fn = hc_vec(inputs[4], s, "hc_front_expand_collapse fn"); + const mx::array base = hc_vec(inputs[6], s, "hc_front_expand_collapse base"); int D = resid.shape(-1); int rows = int(resid.size() / (HC * D)); @@ -293,13 +311,13 @@ void KQuantHcFrontExpandCollapse::eval_gpu( "kq_hc_front_expand_collapse_" + kq_type_string(resid.dtype()); auto kernel = kq_get_kernel(d, kname); ce.set_compute_pipeline_state(kernel); - ce.set_input_array(inputs[0], 0); + ce.set_input_array(x_sub, 0); ce.set_input_array(resid, 1); ce.set_input_array(inputs[2], 2); ce.set_input_array(inputs[3], 3); - ce.set_input_array(inputs[4], 4); + ce.set_input_array(fn, 4); ce.set_input_array(inputs[5], 5); - ce.set_input_array(inputs[6], 6); + ce.set_input_array(base, 6); ce.set_input_array(inputs[7], 7); ce.set_output_array(outputs[0], 8); ce.set_output_array(mixes_raw, 9); @@ -324,7 +342,8 @@ void KQuantHcExpand::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& resid = inputs[1]; + const mx::array x = hc_vec(inputs[0], s, "hc_expand x"); + const mx::array resid = hc_vec(inputs[1], s, "hc_expand resid"); int D = resid.shape(-1); int rows = int(resid.size() / (HC * D)); @@ -332,7 +351,7 @@ void KQuantHcExpand::eval_gpu( 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(inputs[0], 0); + ce.set_input_array(x, 0); ce.set_input_array(resid, 1); ce.set_input_array(inputs[2], 2); ce.set_input_array(inputs[3], 3); diff --git a/src/kquant_internal.h b/src/kquant_internal.h index 2c5e85e..e0ee98d 100644 --- a/src/kquant_internal.h +++ b/src/kquant_internal.h @@ -61,6 +61,12 @@ int kq_dequant_route_mode(); // only. bool kq_qmm_log_on(); +// Byte offset of x's start modulo 16, for KQ_QMM_LOG. A buffer starts on a +// 16-byte boundary, so this is the alignment the kernels see. +inline int kq_x_align(const mx::array& x) { + return static_cast(x.offset() % 16); +} + // KVarN key/value widths the record codec supports. inline bool kq_kvarn_bits_ok(int bits) { return bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || @@ -179,6 +185,32 @@ mx::array kq_dense_copy_gpu(const mx::array& a, const mx::Stream& s); // this in its eval and adds no node for the operand. mx::array kq_row_major_gpu(const mx::array& a, const mx::Stream& s); +// An operand as kernels that read it as vectors need it: a itself when it +// is row-major and starts on an align-byte boundary, else the row-major +// copy, whose buffer starts on a 256-byte boundary. Metal defines a vector +// load only at an address aligned to the vector's size, and a row-major +// view can start at any element of its buffer. The ops bind activations +// and float operands they read as vectors through this, with align the +// widest load (16 bytes unless noted at the call). what names the op and +// operand for KQ_QMM_LOG, which prints "[kq align] xin=" for +// each copy, n being the byte offset of a's start modulo align. +mx::array kq_x_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align = 16); + +// The same for an operand a kernel reads with a unit stride in its last +// dim and the operand's own strides in the dims before it, such as a KV +// cache: a itself when its last stride is 1 and its start and the stride +// of each dim longer than 1 fall on align-byte boundaries, else the +// row-major copy. +mx::array kq_strides_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align = 16); + // The same for an operand a kernel reads with a unit stride in its last // dim and the operand's own strides in the dims before it: a itself when // the last stride is 1, else the row-major copy. diff --git a/src/kquant_latent_fp4.cpp b/src/kquant_latent_fp4.cpp index d4c24ed..eb46114 100644 --- a/src/kquant_latent_fp4.cpp +++ b/src/kquant_latent_fp4.cpp @@ -29,7 +29,12 @@ void KQLatentFp4Pack::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& x = inputs[0]; + // The kernel reads each group of 16 as two 8-wide vectors. + const mx::array x = kq_x_aligned_gpu( + inputs[0], + s, + "latent_fp4_pack x", + 8 * static_cast(inputs[0].itemsize())); auto& codes = outputs[0]; auto& scales = outputs[1]; codes.set_data(mx::allocator::malloc(codes.nbytes())); diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 8c1a64d..4007153 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -749,8 +749,8 @@ static KqSplitkPartials kq_splitk_partials(mx::Dtype dt) { return kq_splitk_partials(dt, mode < 0 ? dt == mx::float16 : mode != 0); } -// KQ_QMM_LOG set (any value): the qmm tiles log the kernel name and shape -// of each dispatch to stderr. +// KQ_QMM_LOG set (any value): the quantized matmul and mat-vec kernels log +// the kernel name and shape of each dispatch to stderr. static bool kq_qmm_log() { static const bool cached = std::getenv("KQ_QMM_LOG") != nullptr; return kq_tune_live() ? std::getenv("KQ_QMM_LOG") != nullptr : cached; @@ -814,24 +814,27 @@ static void kq_require_tile_kernel( kq_require_kernel(d, std::string(lever) + "=" + kq_tile_string(t), kname); } -// One KQ_QMM_LOG line: kernel name, shape and grid of a qmm tile dispatch. +// One KQ_QMM_LOG line: kernel name, shape and grid of a dispatch, and xa, +// the byte offset of x's start modulo 16. static void kq_log_qmm( const std::string& kname, int M, int N, int K, - const MTL::Size& grid) { + const MTL::Size& grid, + const array& x) { if (kq_qmm_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, K, static_cast(grid.width), static_cast(grid.height), - static_cast(grid.depth)); + static_cast(grid.depth), + kq_x_align(x)); } } @@ -1013,7 +1016,7 @@ void qmm_nax( if (tile) { kq_require_tile_kernel(d, "KQ_NAX_TILE", *tile, kname); } - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -1111,7 +1114,7 @@ void qmm( if (tile) { kq_require_tile_kernel(d, "KQ_QMM_TILE", *tile, kname); } - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -1332,7 +1335,7 @@ void qmm_splitk( ce.set_bytes(part_stride, c++); MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, splits); - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); @@ -1412,7 +1415,7 @@ void verify_mma( ce.set_bytes(part_stride, c++); MTL::Size group_dims(256, 1, 1); MTL::Size grid_dims((N + rows - 1) / rows, 1, splits); - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; @@ -1569,7 +1572,7 @@ void verify_nax( ce.set_bytes(part_stride, c++); MTL::Size group_dims(128, 1, 1); MTL::Size grid_dims((N + 127) / 128, 1, splits); - kq_log_qmm(vname, M, N, K, grid_dims); + kq_log_qmm(vname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; @@ -1667,7 +1670,7 @@ void verify_nax_multi( ce.set_bytes(part_stride.data(), 4, c++); MTL::Size group_dims(128, 1, 1); MTL::Size grid_dims(tgs, 1, splits); - kq_log_qmm(vname, M, static_cast(rows), K, grid_dims); + kq_log_qmm(vname, M, static_cast(rows), K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; @@ -1745,7 +1748,7 @@ void qmm_nax_splitk( ce.set_bytes(part_stride, c++); MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, splits); - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); @@ -1797,6 +1800,7 @@ void qvm( bits, B > 1 ? "_batch_1" : "_batch_0"); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -1854,6 +1858,7 @@ static bool qmv_splitk( if (int(kernel->maxTotalThreadsPerThreadgroup()) < 32 * nsg) { return false; } + kq_log_qmm(kname, 1, N, K, MTL::Size(1, N / rps, 1), x); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); int c = 0; @@ -1930,6 +1935,7 @@ void qmv( bits, B > 1 ? "_batch_1" : "_batch_0"); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -1994,6 +2000,7 @@ void qmv_bias( bits, "_batch_0"); + kq_log_qmm(kname, 1, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -2073,6 +2080,7 @@ void verify_qmv( bits, "_batch_0"); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -2171,6 +2179,7 @@ void verify_mv_ext( : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") : (use_hd ? "_hd" : "")))); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -2491,7 +2500,7 @@ static int kq_verify_nax_splits(int N, int K, int kstep, const std::string& t) { // non-f32 shapes with M <= 32 only (the caller checks). bool kq_forced_route( const char* route, - const array& x, + const array& x_in, const array& w, const array& scales, array& out, @@ -2504,23 +2513,25 @@ bool kq_forced_route( const Stream& s, const std::string& t) { const std::string r(route); + // x is bound once the route serves the call, as in eval_gpu_base. + auto x16 = [&] { return kq_x_aligned_gpu(x_in, s, "quantized_matmul x"); }; const KQuantCodec* codec = codec_by_name(t); const int wpb = codec ? codec->weights_per_block : 256; if (r == "qmv") { - dispatch_qmv(x, w, scales, out, group_size, bits, M, N, K, d, s, t); + dispatch_qmv(x16(), w, scales, out, group_size, bits, M, N, K, d, s, t); return true; } if (r == "verify_qmv") { if (M >= 2 && M <= verify_qmv_max_rows() && N % kquant_qmv_bn(t) == 0 && K % qmv_fast_k_align() == 0 && codec_has_verify_qmv(t)) { - verify_qmv(x, w, scales, out, group_size, bits, M, N, K, d, s, t); + verify_qmv(x16(), w, scales, out, group_size, bits, M, N, K, d, s, t); return true; } return false; } if (r == "mv_ext") { if (kq_codec_has_mv_ext(t) && M >= 2 && M <= 12 && K % wpb == 0) { - verify_mv_ext(x, w, scales, out, group_size, bits, M, N, K, d, s, t); + verify_mv_ext(x16(), w, scales, out, group_size, bits, M, N, K, d, s, t); return true; } return false; @@ -2529,7 +2540,7 @@ bool kq_forced_route( const int rows = codec_verify_mma_rows(t); if (rows > 0 && M <= 8) { const int sp = kq_verify_mma_splits(N, K, group_size, rows); - verify_mma(x, w, scales, out, group_size, bits, M, N, K, sp, d, s, t); + verify_mma(x16(), w, scales, out, group_size, bits, M, N, K, sp, d, s, t); return true; } return false; @@ -2540,7 +2551,20 @@ bool kq_forced_route( kq_verify_nax_fits(w, K, kstep, t)) { const int sp = kq_verify_nax_splits(N, K, kstep, t); verify_nax( - x, w, scales, out, group_size, bits, M, N, K, kstep, sp, d, s, t); + kq_x_aligned_gpu(x_in, s, "quantized_matmul x", 8), + w, + scales, + out, + group_size, + bits, + M, + N, + K, + kstep, + sp, + d, + s, + t); return true; } return false; @@ -2550,7 +2574,7 @@ bool kq_forced_route( const int target = (e != nullptr && std::atoi(e) > 1) ? std::atoi(e) : 16; const int sp = kq_split_count(target, K / group_size); if (kq_splitk_codec(t) && sp > 1) { - qmm_splitk(x, w, scales, out, group_size, bits, M, N, K, sp, d, s, t); + qmm_splitk(x16(), w, scales, out, group_size, bits, M, N, K, sp, d, s, t); return true; } return false; @@ -2559,7 +2583,7 @@ bool kq_forced_route( kq_is_nax_available() && codec_has_nax(t) && (K % 64 == 0); if (r == "nax") { if (nax_ok) { - qmm(x, w, scales, out, true, group_size, bits, M, N, K, d, s, t); + qmm(x16(), w, scales, out, true, group_size, bits, M, N, K, d, s, t); return true; } return false; @@ -2574,7 +2598,20 @@ bool kq_forced_route( const int sp = kq_nax_split(target, K / q, per); if (nax_ok && sp > 1) { qmm_nax_splitk( - x, w, scales, out, group_size, bits, M, N, K, sp, d, s, t, per * q); + x16(), + w, + scales, + out, + group_size, + bits, + M, + N, + K, + sp, + d, + s, + t, + per * q); return true; } return false; @@ -2605,6 +2642,18 @@ bool kq_qmm_log_on() { return kq_qmm_log(); } +void kq_log_dispatch( + const std::string& kname, + const array& x, + const array& out, + const MTL::Size& grid) { + if (kq_qmm_log()) { + const int N = out.ndim() > 0 ? out.shape(-1) : 1; + const int M = N > 0 ? static_cast(out.size() / N) : 0; + kq_log_qmm(kname, M, N, x.ndim() > 0 ? x.shape(-1) : 1, grid, x); + } +} + #endif // _METAL_ std::vector KQuantMatmul::output_shapes( @@ -2868,13 +2917,18 @@ void KQuantMatmulMulti::eval_gpu( } } if (!multi) { + // x is bound once, so that each per-weight call reads it in place. + const array x16 = kq_x_aligned_gpu(x, s, "quantized_matmul_multi x"); KQuantMatmul one(s, kquant_type_, group_size_, bits_, true); for (int i = 0; i < n; ++i) { std::vector o{outputs[i]}; - one.eval_gpu({x, ws[i], scales}, o); + one.eval_gpu({x16, ws[i], scales}, o); } return; } + // verify_nax reads x where it lands on an 8-byte boundary and from the + // aligned copy otherwise, as in KQuantMatmul::eval_gpu_base. + const array xv = kq_x_aligned_gpu(x, s, "quantized_matmul_multi x", 8); std::vector gw; std::vector go; for (int i = 0; i < n; ++i) { @@ -2886,14 +2940,14 @@ void KQuantMatmulMulti::eval_gpu( go.push_back(outputs[i]); if (gw.size() == 4 || i == n - 1) { verify_nax_multi( - x, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); + xv, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); gw.clear(); go.clear(); } } if (!gw.empty()) { verify_nax_multi( - x, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); + xv, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); } } @@ -2922,7 +2976,14 @@ void KQuantMatmul::eval_gpu_base( auto& out = outputs[0]; // inputs are row-contiguous (ensured by the op): x, w (uint8), scales. - const auto& x = inputs[0]; + // x is bound once the route is chosen, from a 16-byte boundary + // (kq_x_aligned_gpu), except by verify_nax: it reads x where it lands on + // an 8-byte boundary, picking its x16 or x8 load pipeline by the start of + // x (kq_verify_nax_pipeline), and from the aligned copy otherwise. + const auto& x_in = inputs[0]; + auto bind_x = [&](int align) { + return kq_x_aligned_gpu(x_in, s, "quantized_matmul x", align); + }; const auto& w = inputs[1]; const auto& scales = inputs[2]; kq_check_weight_base(w, kquant_type_, "quantized_matmul"); @@ -2931,16 +2992,16 @@ void KQuantMatmul::eval_gpu_base( return; } - bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; - int K = x.shape(-1); - int M = non_batched ? static_cast(x.size()) / K : x.shape(-2); + bool non_batched = w.ndim() == 2 && x_in.flags().row_contiguous; + int K = x_in.shape(-1); + int M = non_batched ? static_cast(x_in.size()) / K : x_in.shape(-2); int N = out.shape(-1); int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; // KQuant special cases. if (!transpose_ && M < vector_limit) { - qmm(x, + qmm(bind_x(16), w, scales, out, @@ -2965,7 +3026,7 @@ void KQuantMatmul::eval_gpu_base( // Probe lever KQ_FORCE_QMM_MIN_M (kq_force_qmm_min_m). const int force_qmm_min_m = kq_force_qmm_min_m(); if (force_qmm_min_m > 0 && transpose_ && M >= force_qmm_min_m) { - qmm(x, + qmm(bind_x(16), w, scales, out, @@ -2991,10 +3052,10 @@ void KQuantMatmul::eval_gpu_base( if (transpose_) { const char* route_e = std::getenv("KQ_QMM_ROUTE"); if (route_e != nullptr && route_e[0] != '\0') { - if (non_batched && M <= 32 && x.dtype() != mx::float32 && + if (non_batched && M <= 32 && x_in.dtype() != mx::float32 && kq_forced_route( route_e, - x, + x_in, w, scales, out, @@ -3019,12 +3080,23 @@ void KQuantMatmul::eval_gpu_base( } // Per-row qmv and verify_nax on NAX GPUs: kq_nax_small_m_route. - if (transpose_ && non_batched && x.dtype() != mx::float32) { - const KqSmallMRoute route = - kq_nax_small_m_route(kquant_type_, M, N, K, x.dtype() == mx::float16); + if (transpose_ && non_batched && x_in.dtype() != mx::float32) { + const KqSmallMRoute route = kq_nax_small_m_route( + kquant_type_, M, N, K, x_in.dtype() == mx::float16); if (route == KqSmallMRoute::qmv) { dispatch_qmv( - x, w, scales, out, group_size_, bits_, M, N, K, d, s, kquant_type_); + bind_x(16), + w, + scales, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + kquant_type_); return; } const int kstep = kq_verify_nax_call_kstep(K, kquant_type_); @@ -3032,7 +3104,7 @@ void KQuantMatmul::eval_gpu_base( kq_verify_nax_fits(w, K, kstep, kquant_type_)) { const int sp = kq_verify_nax_splits(N, K, kstep, kquant_type_); verify_nax( - x, + bind_x(8), w, scales, out, @@ -3050,6 +3122,9 @@ void KQuantMatmul::eval_gpu_base( } } + // Every other route reads x from a 16-byte boundary. + const mx::array x = bind_x(16); + // Register-resident MMA verify, from the per-device entry // (kq_verify_mma_min_m, kq_verify_mma_min_m_nax) through M8. // KQ_VERIFY_MMA= forces the entry at M >= m, 0 or 1 disables; read @@ -3477,7 +3552,8 @@ void KQuantQmvBias::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; - const auto& x = inputs[0]; + const mx::array x = + kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_qmv_bias x"); const auto& w = inputs[1]; const auto& scales = inputs[2]; const auto& bias = inputs[3]; diff --git a/src/kquant_matmul_fused.cpp b/src/kquant_matmul_fused.cpp index 7b9275f..6b56583 100644 --- a/src/kquant_matmul_fused.cpp +++ b/src/kquant_matmul_fused.cpp @@ -216,7 +216,7 @@ void KQuantMatmulFused::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& x = inputs[0]; + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_glu x"); // A weight whose rows are dense binds as it is, at its own offset, so a // row band of a wider weight is read in place. Any other view is copied // to dense rows first. @@ -279,14 +279,15 @@ void KQuantMatmulFused::eval_gpu( if (glu_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, K, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname, hash_name, func_consts); diff --git a/src/kquant_matmul_lora.cpp b/src/kquant_matmul_lora.cpp index eb74db4..9982537 100644 --- a/src/kquant_matmul_lora.cpp +++ b/src/kquant_matmul_lora.cpp @@ -271,7 +271,7 @@ void KQuantMatmulLora::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& x = inputs[0]; + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_lora x"); // A weight whose rows are dense binds as it is, at its own offset. Any // other view is copied to dense rows first. z and B bind as they are // when they are row-major. @@ -312,7 +312,7 @@ void KQuantMatmulLora::eval_gpu( if (lora_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d R=%d bands=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d R=%d bands=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, @@ -321,7 +321,8 @@ void KQuantMatmulLora::eval_gpu( n_bands, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname); diff --git a/src/kquant_matmul_resid.cpp b/src/kquant_matmul_resid.cpp index 0d91fae..d4240f2 100644 --- a/src/kquant_matmul_resid.cpp +++ b/src/kquant_matmul_resid.cpp @@ -118,7 +118,8 @@ void KQuantMatmulResid::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const mx::array x = kq_row_major_gpu(inputs[0], s); + const mx::array x = + kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_residual x"); // A weight whose rows are dense binds as it is, at its own offset. Any // other view is copied to dense rows first. The residual and the gate // bind as they are when they are row-major. @@ -160,7 +161,7 @@ void KQuantMatmulResid::eval_gpu( if (resid_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d B=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d B=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, @@ -168,7 +169,8 @@ void KQuantMatmulResid::eval_gpu( B, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname); diff --git a/src/kquant_metal_internal.h b/src/kquant_metal_internal.h index b238b3f..26c031c 100644 --- a/src/kquant_metal_internal.h +++ b/src/kquant_metal_internal.h @@ -388,6 +388,14 @@ inline int add_gather_strides_and_shapes( return offset; } +// One KQ_QMM_LOG line for a gather or MoE dispatch: the kernel name, M and +// N of out's rows, K of x's last dim, the grid and xa. +void kq_log_dispatch( + const std::string& kname, + const array& x, + const array& out, + const MTL::Size& grid); + // Fetch a kq kernel from the bundled metallib. // // The library handle is resolved exactly once and cached in a function-local diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index e1e34e5..bc5e3dd 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -265,6 +265,7 @@ inline void kq_gather_qmv_pair_gpu( const int fsg = kq_moe_fit_sg(kernel, sg); MTL::Size group_dims(32, fsg, 1); MTL::Size grid_dims(N / (fsg * 32 / nx), 1, (T + 1) / 2); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } #endif // _METAL_ @@ -392,8 +393,9 @@ void KQuantMoEGLU::eval_gpu( const auto& indices = inputs[7]; moe_check_stack("moe_glu_gather", gw, gs); moe_check_stack("moe_glu_gather", uw, us); - // x is read as dense rows: a view with other strides is copied first. - const mx::array x = kq_row_major_gpu(inputs[6], s); + // x is read as dense rows from a 16-byte boundary: any other view is + // copied first (kq_x_aligned_gpu). + const mx::array x = kq_x_aligned_gpu(inputs[6], s, "moe_glu_gather x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -423,6 +425,7 @@ void KQuantMoEGLU::eval_gpu( ce.set_bytes(limit, 12); MTL::Size group_dims(32, 2, 1); MTL::Size grid_dims(N / 8, R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -438,7 +441,7 @@ void KQuantGatherQMVBias::eval_gpu( const auto& b = inputs[2]; const auto& indices = inputs[4]; moe_check_stack("gather_qmv_bias", w, sc); - const mx::array x = kq_row_major_gpu(inputs[3], s); + const mx::array x = kq_x_aligned_gpu(inputs[3], s, "gather_qmv_bias x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -463,6 +466,7 @@ void KQuantGatherQMVBias::eval_gpu( ce.set_bytes(N, 7); MTL::Size group_dims(32, 2, 1); MTL::Size grid_dims(N / (fine ? 2 : 8), R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -479,7 +483,7 @@ void KQuantGatherQMVMixBias::eval_gpu( const auto& indices = inputs[4]; const auto& scores = inputs[5]; moe_check_stack("gather_qmv_mix_bias", w, sc); - const mx::array x = kq_row_major_gpu(inputs[3], s); + const mx::array x = kq_x_aligned_gpu(inputs[3], s, "gather_qmv_mix_bias x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -506,6 +510,7 @@ void KQuantGatherQMVMixBias::eval_gpu( ce.set_bytes(S, 9); MTL::Size group_dims(32, 2, 1); MTL::Size grid_dims(N / (fine ? 2 : 8), 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -525,7 +530,8 @@ void KQuantMoEGLUKQ::eval_gpu( moe_check_stack("moe_glu_gather_kq", uw); kq_check_weight_base(gw, kquant_type_, "moe_glu_gather_kq"); kq_check_weight_base(uw, kquant_type_, "moe_glu_gather_kq"); - const mx::array x = kq_row_major_gpu(inputs[biased ? 4 : 2], s); + const mx::array x = + kq_x_aligned_gpu(inputs[biased ? 4 : 2], s, "moe_glu_gather_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -580,6 +586,7 @@ void KQuantMoEGLUKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sg * 32 / nx), R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -597,7 +604,8 @@ void KQuantGatherQMVKQ::eval_gpu( const auto& indices = inputs[n_base - 1]; moe_check_stack("gather_qmv_kq", w); kq_check_weight_base(w, kquant_type_, "gather_qmv_kq"); - const mx::array x = kq_row_major_gpu(inputs[biased ? 2 : 1], s); + const mx::array x = + kq_x_aligned_gpu(inputs[biased ? 2 : 1], s, "gather_qmv_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -645,6 +653,7 @@ void KQuantGatherQMVKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (fine ? 2 : (sg * 32 / nx)), R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); if (lora_flags_ & KQ_LORA_PRESENT) { kq_lora_epilogue_rows_gpu( @@ -679,7 +688,8 @@ void KQuantMoEGLUShexpKQ::eval_gpu( kq_check_weight_base(uw, kquant_type_, "moe_glu_gather_shexp_kq"); kq_check_weight_base(sgw, shexp_type_, "moe_glu_gather_shexp_kq"); kq_check_weight_base(suw, shexp_up_type_, "moe_glu_gather_shexp_kq"); - const mx::array x = kq_row_major_gpu(inputs[4], s); + const mx::array x = + kq_x_aligned_gpu(inputs[4], s, "moe_glu_gather_shexp_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -714,6 +724,7 @@ void KQuantMoEGLUShexpKQ::eval_gpu( const int sg = kq_moe_fit_sg(kernel, kq_moe_pick_sg(kquant_type_, N, nx)); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sg * 32 / nx), R + 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); return; } @@ -762,6 +773,7 @@ void KQuantMoEGLUShexpKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sg * 32 / nx), R + 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -780,7 +792,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( moe_check_shexp("gather_qmv_mix_kq", sw); kq_check_weight_base(w, kquant_type_, "gather_qmv_mix_kq"); kq_check_weight_base(sw, shexp_type_, "gather_qmv_mix_kq"); - const mx::array x = kq_row_major_gpu(inputs[2], s); + const mx::array x = kq_x_aligned_gpu(inputs[2], s, "gather_qmv_mix_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = x.shape(0); @@ -826,6 +838,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( ce.set_bytes(SC, 9); MTL::Size group_dims(32, 2 * S, 1); MTL::Size grid_dims(N / 8, 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); return; } @@ -898,6 +911,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (fine ? 2 : (sg * 32 / nx)), 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -913,7 +927,7 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( const auto& scores = inputs[3]; moe_check_stack("gather_qmv_mix_ns_kq", w); kq_check_weight_base(w, kquant_type_, "gather_qmv_mix_ns_kq"); - const mx::array x = kq_row_major_gpu(inputs[1], s); + const mx::array x = kq_x_aligned_gpu(inputs[1], s, "gather_qmv_mix_ns_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = x.shape(0); @@ -992,6 +1006,7 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( sp ? 2 * S : kq_moe_fit_sg(kernel, kq_moe_pick_sg(kquant_type_, N, nx)); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sp ? 8 : (sg * 32 / nx)), 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } if (lora_z) { diff --git a/src/kquant_moe_mix.cpp b/src/kquant_moe_mix.cpp index 5bdabd6..85d4761 100644 --- a/src/kquant_moe_mix.cpp +++ b/src/kquant_moe_mix.cpp @@ -92,7 +92,8 @@ void KQuantGatherMix::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& y = inputs[0]; + // y is the activation, bound from a 16-byte boundary. + const mx::array y = kq_x_aligned_gpu(inputs[0], s, "gather_mix y"); const int N = y.shape(1); const int Tn = inputs[2].shape(0); const int K = inputs[2].shape(1); @@ -100,7 +101,7 @@ void KQuantGatherMix::eval_gpu( 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(inputs[0], 0); + ce.set_input_array(y, 0); ce.set_input_array(inputs[1], 1); ce.set_input_array(inputs[2], 2); ce.set_output_array(out, 3); @@ -113,6 +114,7 @@ void KQuantGatherMix::eval_gpu( } MTL::Size group_dims(threads, 1, 1); MTL::Size grid_dims(Tn, 1, 1); + kq_log_dispatch(kname, y, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/src/kquant_norm_fused.cpp b/src/kquant_norm_fused.cpp index 7ff2b29..81fe4a8 100644 --- a/src/kquant_norm_fused.cpp +++ b/src/kquant_norm_fused.cpp @@ -756,7 +756,12 @@ void KQuantRMSNormGate::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& x = inputs[0]; + // x, w and gate are read as vectors of up to four elements. + const int align = 4 * static_cast(inputs[0].itemsize()); + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "rmsnorm_gate x", align); + const mx::array w = kq_x_aligned_gpu(inputs[1], s, "rmsnorm_gate w", align); + const mx::array gate = + kq_x_aligned_gpu(inputs[2], s, "rmsnorm_gate gate", align); const int D = x.shape(-1); const int nrows = int(x.size() / D); std::string kname = "kq_rmsnorm_gate_" + kq_type_string(x.dtype()) + "_" + @@ -764,9 +769,9 @@ void KQuantRMSNormGate::eval_gpu( 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(inputs[0], 0); - ce.set_input_array(inputs[1], 1); - ce.set_input_array(inputs[2], 2); + ce.set_input_array(x, 0); + ce.set_input_array(w, 1); + ce.set_input_array(gate, 2); ce.set_output_array(out, 3); ce.set_bytes(eps_, 4); ce.set_bytes(nrows, 5); diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index f7666c6..671bc4e 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -56,6 +56,18 @@ void kq_sdpa_check_layout( } } +// The kernels read q, k, v and the KVarN stages as vectors of four +// elements (four uint32 words of a q8 cache), so each is bound from a +// boundary of four of its elements, k, v and the stages with their strides +// on such boundaries too. +array sdpa_q(const array& a, const Stream& s, const char* what) { + return kq_x_aligned_gpu(a, s, what, 4 * static_cast(a.itemsize())); +} + +array sdpa_rows(const array& a, const Stream& s, const char* what) { + return kq_strides_aligned_gpu(a, s, what, 4 * static_cast(a.itemsize())); +} + // Mirror of the Metal-side KQKvarnMeta (kq_sdpa.h, buffer 28); the layouts // must match field for field. struct KQKvarnMeta { @@ -271,10 +283,12 @@ void KQuantSDPA::eval_gpu( out.set_data(mx::allocator::malloc(out.nbytes())); // q row-contiguous [B, Hq, qL, D]; k/v [B, Hkv, kL, D], D contiguous. - const auto& q = inputs[0]; - const auto& k = inputs[1]; - const auto& v = inputs[2]; - kq_sdpa_check_layout("sdpa_vector", q, k, v); + kq_sdpa_check_layout("sdpa_vector", inputs[0], inputs[1], inputs[2]); + // q, k and v are bound from vector boundaries (sdpa_q, sdpa_rows), here + // and in the evals below. + const mx::array q = sdpa_q(inputs[0], s, "sdpa_vector q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_vector k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_vector v"); int B = q.shape(0); int n_q_heads = q.shape(1); @@ -385,9 +399,10 @@ void KQuantSDPAGQA::eval_gpu( outputs[1].set_data(mx::allocator::malloc(outputs[1].nbytes())); } - const auto& q = inputs[0]; - const auto& k = inputs[1]; - const auto& v = inputs[2]; + kq_sdpa_check_layout("sdpa_decode_gqa", inputs[0], inputs[1], inputs[2]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_decode_gqa q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_decode_gqa k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_decode_gqa v"); const bool sinks = has_sinks_; const bool starts = has_starts_; const bool kv_q8 = has_kv_q8_; @@ -397,7 +412,12 @@ void KQuantSDPAGQA::eval_gpu( const size_t starts_idx = 3 + (sinks ? 1 : 0); const size_t ends_idx = starts_idx + (starts ? 1 : 0); const size_t qkv_idx = ends_idx + (ends ? 1 : 0); - kq_sdpa_check_layout("sdpa_decode_gqa", q, k, v); + // The KVarN stages are key and value rows too. A copy dispatches on the + // encoder, so each operand is bound before the first kernel is set up. + const mx::array stk = + kvarn ? sdpa_rows(inputs[qkv_idx + 2], s, "sdpa_decode_gqa stage_k") : q; + const mx::array stv = + kvarn ? sdpa_rows(inputs[qkv_idx + 3], s, "sdpa_decode_gqa stage_v") : q; if (kv_q8 || kvarn) { for (int i = 0; i < 4; i++) { kq_sdpa_check_batch_heads("sdpa_decode_gqa", inputs[qkv_idx + i]); @@ -612,8 +632,6 @@ void KQuantSDPAGQA::eval_gpu( // KVarN axes/stage operands and region map (dummies when compiled out). KQKvarnMeta kvm = {}; if (kvarn) { - const auto& stk = inputs[qkv_idx + 2]; - const auto& stv = inputs[qkv_idx + 3]; const int sink_cap = stk.shape(2) - KQ_KVARN_GROUP; kvm.sink_rows = std::min(kvarn_n_, sink_cap); kvm.n_rec_keys = @@ -691,15 +709,19 @@ void KQuantSDPAFAVerify::eval_gpu( // q is the GQA-folded query tile [B, Hkv, n_rows, D], row-contiguous; // k/v [B, Hkv, kL, D] are read in place via their head/seq strides. - const auto& q = inputs[0]; - const auto& k = inputs[1]; - const auto& v = inputs[2]; - kq_sdpa_check_layout("sdpa_fa_verify", q, k, v); + kq_sdpa_check_layout("sdpa_fa_verify", inputs[0], inputs[1], inputs[2]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_fa_verify q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_fa_verify k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_fa_verify v"); // With kvarn, k/v are the code slabs and the key count comes from the // op; a tail-merge body walks only the first n_attend keys while the // region map stays on the full layout (as sdpa_decode_gqa). const bool kvarn = has_kvarn_; + const mx::array stk = + kvarn ? sdpa_rows(inputs[5], s, "sdpa_fa_verify stage_k") : q; + const mx::array stv = + kvarn ? sdpa_rows(inputs[6], s, "sdpa_fa_verify stage_v") : q; int B = q.shape(0); int n_kv_heads = k.shape(1); int n_rows = q.shape(2); @@ -863,8 +885,6 @@ void KQuantSDPAFAVerify::eval_gpu( // KVarN axes/stage operands and region map (dummies when compiled out). KQKvarnMeta kvm = {}; if (kvarn) { - const auto& stk = inputs[5]; - const auto& stv = inputs[6]; const int sink_cap = stk.shape(2) - KQ_KVARN_GROUP; kvm.sink_rows = std::min(kvarn_n_, sink_cap); kvm.n_rec_keys = @@ -933,10 +953,10 @@ void KQuantSDPAFAIndexed::eval_gpu( // q [1, Hq, Q, D] row-contiguous; kv [1, 1, N, D] read in place through // its seq stride; idx [Q, M] int32 row-contiguous. - const auto& q = inputs[0]; - const auto& kv = inputs[1]; + kq_sdpa_check_layout("sdpa_fa_indexed", inputs[0], inputs[1], inputs[1]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_fa_indexed q"); + const mx::array kv = sdpa_rows(inputs[1], s, "sdpa_fa_indexed kv"); const auto& idx = inputs[2]; - kq_sdpa_check_layout("sdpa_fa_indexed", q, kv, kv); int n_heads = q.shape(1); int n_queries = q.shape(2); @@ -1068,19 +1088,23 @@ void KQuantSDPACascade::eval_gpu( outputs[1].set_data(mx::allocator::malloc(outputs[1].nbytes())); } - const auto& q = inputs[0]; - const auto& qf = inputs[1]; - const auto& k_sh = inputs[2]; - const auto& v_sh = inputs[3]; - const auto& k_pr = inputs[4]; - const auto& v_pr = inputs[5]; + kq_sdpa_check_layout( + "sdpa_decode_gqa_cascade", inputs[0], inputs[4], inputs[5]); + kq_sdpa_check_layout( + "sdpa_decode_gqa_cascade", inputs[1], inputs[2], inputs[3]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_decode_gqa_cascade q"); + const mx::array qf = sdpa_q(inputs[1], s, "sdpa_decode_gqa_cascade q_folded"); + const mx::array k_sh = + sdpa_rows(inputs[2], s, "sdpa_decode_gqa_cascade k_shared"); + const mx::array v_sh = + sdpa_rows(inputs[3], s, "sdpa_decode_gqa_cascade v_shared"); + const mx::array k_pr = sdpa_rows(inputs[4], s, "sdpa_decode_gqa_cascade k"); + const mx::array v_pr = sdpa_rows(inputs[5], s, "sdpa_decode_gqa_cascade v"); const bool starts = has_starts_; const bool kv_q8 = has_kv_q8_; // Optional trailing inputs: [starts], then the eight q8 scale/bias // arrays (shared k/v pairs first, then private). const int q8_base = 6 + int(starts); - kq_sdpa_check_layout("sdpa_decode_gqa_cascade", q, k_pr, v_pr); - kq_sdpa_check_layout("sdpa_decode_gqa_cascade", qf, k_sh, v_sh); if (kv_q8) { for (int i = 4; i < 8; i++) { kq_sdpa_check_batch_heads("sdpa_decode_gqa_cascade", inputs[q8_base + i]); @@ -1346,14 +1370,16 @@ void KQuantSDPABSPrefill::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& q = inputs[0]; // folded [1, Hkv, n_qt, 48, D] - const auto& k = inputs[1]; - const auto& v = inputs[2]; + kq_sdpa_check_layout( + "sdpa_prefill_block_sparse", inputs[0], inputs[1], inputs[2]); + // q folded [1, Hkv, n_qt, 48, D]. + const mx::array q = sdpa_q(inputs[0], s, "sdpa_prefill_block_sparse q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_prefill_block_sparse k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_prefill_block_sparse v"); const auto& pages = inputs[3]; const auto& pmask = inputs[4]; const auto& counts = inputs[5]; auto& out = outputs[0]; - kq_sdpa_check_layout("sdpa_prefill_block_sparse", q, k, v); out.set_data(mx::allocator::malloc(out.nbytes())); int n_kv_heads = q.shape(1); diff --git a/src/kquant_sdpa_sparse_decode.cpp b/src/kquant_sdpa_sparse_decode.cpp index 9ab3c2c..eaff51c 100644 --- a/src/kquant_sdpa_sparse_decode.cpp +++ b/src/kquant_sdpa_sparse_decode.cpp @@ -52,7 +52,8 @@ mx::array rows_flat(const mx::array& a, const char* op, mx::StreamOrDevice s) { // The window, pool and pool scales as the kernels read them. A reshape to // [B, rows * D] keeps a view whenever the dims of its input fold into the // two of its output, so a row can come out with an element stride other -// than 1, which is copied here. The kernels step the batches of a packed +// than 1, which is copied here, as is a window or pool that starts or +// steps off a 16-byte boundary. The kernels step the batches of a packed // pool by the stride of its codes and take an eighth of it for the scales, // so codes and scales that step differently are copied too. struct SparseRows { @@ -65,10 +66,19 @@ SparseRows sparse_rows( const mx::array& win, const mx::array& pool, const mx::array* scales, - const mx::Stream& s) { + const mx::Stream& s, + bool prefill) { + // The kernels read window and pool rows as 16-byte vectors (4-byte words + // for packed codes), so the rows start on 16-byte boundaries too. SparseRows r{ - kq_unit_stride_gpu(win, s), - kq_unit_stride_gpu(pool, s), + kq_strides_aligned_gpu( + win, + s, + prefill ? "sdpa_sparse_prefill win" : "sdpa_sparse_decode win"), + kq_strides_aligned_gpu( + pool, + s, + prefill ? "sdpa_sparse_prefill pool" : "sdpa_sparse_decode pool"), scales ? kq_unit_stride_gpu(*scales, s) : pool}; if (scales && r.pool.size() > 0 && r.pool.shape(0) != 1 && r.pool.strides(0) != 8 * r.scales.strides(0)) { @@ -89,11 +99,13 @@ void KQSdpaSparseDecode::eval_gpu( auto& d = mx::metal::device(s.device); auto& ce = mx::metal::get_command_encoder(s); - const auto& q = inputs[0]; + // q is read as 16-byte vectors through its strides. + const mx::array q = + kq_strides_aligned_gpu(inputs[0], s, "sdpa_sparse_decode q"); const auto& idx = inputs[3]; size_t next = 4; - const SparseRows rows = - sparse_rows(inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s); + const SparseRows rows = sparse_rows( + inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s, false); const auto& win = rows.win; // [B|1, W * D], rows D apart const auto& pool = rows.pool; // [B|1, P * D] (packed: [B|1, P * D / 2] codes) const mx::array* pscales = packed_ ? &rows.scales : nullptr; @@ -481,11 +493,12 @@ void KQSdpaSparsePrefill::eval_gpu( auto& d = mx::metal::device(s.device); auto& ce = mx::metal::get_command_encoder(s); - const auto& q = inputs[0]; + const mx::array q = + kq_strides_aligned_gpu(inputs[0], s, "sdpa_sparse_prefill q"); const auto& idx = inputs[3]; size_t next = 4; - const SparseRows rows = - sparse_rows(inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s); + const SparseRows rows = sparse_rows( + inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s, true); const auto& win = rows.win; // [B|1, S * D], rows D apart const auto& pool = rows.pool; // [B|1, P * D] (packed: [B|1, P * D / 2] codes) const mx::array* pscales = packed_ ? &rows.scales : nullptr; diff --git a/src/kquant_skinny_mv.cpp b/src/kquant_skinny_mv.cpp index ec6f523..aa89e12 100644 --- a/src/kquant_skinny_mv.cpp +++ b/src/kquant_skinny_mv.cpp @@ -35,8 +35,8 @@ void KQuantSkinnyMV::eval_gpu( auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& x = inputs[0]; - const auto& w = inputs[1]; + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "skinny_matmul x"); + const mx::array w = kq_x_aligned_gpu(inputs[1], s, "skinny_matmul w"); int K = x.shape(-1); int M = x.shape(-2); int N = w.shape(0); @@ -55,6 +55,7 @@ void KQuantSkinnyMV::eval_gpu( ce.set_bytes(K, 5); MTL::Size group_dims(32 * 8, 1, 1); MTL::Size grid_dims((N + 7) / 8, T, 1); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/tests/qmm_probe.py b/tests/qmm_probe.py index 72d6dae..17a5afe 100644 --- a/tests/qmm_probe.py +++ b/tests/qmm_probe.py @@ -20,7 +20,9 @@ ``checked``, the number of rows it read. Every call gets ``digest``, a hash of its output bits (of every output on a ``"multi"`` call), to compare calls across probe processes. ``"x_offset": E`` passes -x as a row-contiguous view E elements into its buffer. +x as a row-contiguous view E elements into its buffer, and ``copies`` +lists the ``[kq align]`` lines of the call, (operand, xin) for each +operand an op copied to an aligned buffer. ``"design": "lead"`` draws nonnegative weights and activations whose first K column is 1024. ``"design": "big"`` draws small weights and dense @@ -247,12 +249,17 @@ def ref_of(lo): ) if c.get("x_offset"): # The same values as a row-contiguous view x_offset elements into - # its buffer. An empty cache keeps the buffer from carrying the - # slack that makes the op copy the view to a fresh one. + # its buffer. The buffer is allocated with the cache off: a reused + # one can carry the slack that makes the op copy the view to a + # fresh one, and buffers freed late can refill an emptied cache. off = c["x_offset"] - mx.clear_cache() - flat = mx.concatenate([mx.zeros((off,), dtype), x.reshape(-1)]) - mx.eval(flat) + limit = mx.set_cache_limit(0) + try: + mx.clear_cache() + flat = mx.concatenate([mx.zeros((off,), dtype), x.reshape(-1)]) + mx.eval(flat) + finally: + mx.set_cache_limit(limit) x = flat[off:].reshape(x.shape) os.write(2, f"[probe] {i}\n".encode()) res = {"i": i, "error": None, "rel": None, "ulp": None} @@ -358,14 +365,19 @@ def ref_of(lo): print(json.dumps(results)) """ -_LOG = re.compile(r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+)$") +_LOG = re.compile( + r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+) xa=(\d+)$" +) +_ALIGN = re.compile(r"^\[kq align\] (.+) xin=(\d+)$") def run_probe(calls, env=None, check=True, timeout=900): """Run ``calls`` in a fresh interpreter whose environment keeps no ``KQ_*`` variable but those in ``env``. Returns one dict per call with - ``kernels`` (list of (name, grid) from KQ_QMM_LOG), ``stderr`` (its - other stderr lines), ``error``, ``rel`` and ``ulp``.""" + ``kernels`` (list of (name, grid) from KQ_QMM_LOG), ``xa`` (the logged + byte offset of x modulo 16 per kernel), ``copies`` (the operands copied + to an aligned buffer, with the byte offset they came at), ``stderr`` + (its other stderr lines), ``error``, ``rel`` and ``ulp``.""" child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} child_env.update(env or {}) spec = {"tests_dir": TESTS_DIR, "calls": calls, "check": check} @@ -379,12 +391,16 @@ def run_probe(calls, env=None, check=True, timeout=900): assert r.returncode == 0, r.stderr[-6000:] results = json.loads(r.stdout.strip().splitlines()[-1]) logs = {} + aligns = {} + copies = {} other = {} cur = None for line in r.stderr.splitlines(): if line.startswith("[probe] "): cur = int(line.split()[1]) logs[cur] = [] + aligns[cur] = [] + copies[cur] = [] other[cur] = [] continue if cur is None: @@ -393,10 +409,17 @@ def run_probe(calls, env=None, check=True, timeout=900): if m: grid = tuple(int(g) for g in m.group(5, 6, 7)) logs[cur].append((m.group(1), grid)) + aligns[cur].append(int(m.group(8))) + continue + m = _ALIGN.match(line) + if m: + copies[cur].append((m.group(1), int(m.group(2)))) else: other[cur].append(line) for res in results: res["kernels"] = logs.get(res["i"], []) + res["xa"] = aligns.get(res["i"], []) + res["copies"] = copies.get(res["i"], []) res["stderr"] = other.get(res["i"], []) return results diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py index 3d568f9..7c92d47 100644 --- a/tests/test_matmul_multi.py +++ b/tests/test_matmul_multi.py @@ -375,9 +375,9 @@ def test_forced_verify_nax_fuses_narrow_weights(): @pytest.mark.parametrize("codec", [*CODECS, "pq2_0", "q2_0"]) def test_x16_load_bits(codec, dtype): # x at offset 0 (paired 16-byte activation loads), at offset 0 with - # KQ_VNAX_X16=0, at every 2-byte offset from 2 to 14 bytes (one load - # per step) and at 16 bytes: the same bits from one verify_nax_multi - # dispatch. + # KQ_VNAX_X16=0, at every 2-byte offset from 2 to 14 bytes (8 bytes: one + # load per step; the others: an aligned copy, paired) and at 16 bytes: + # the same bits from one verify_nax_multi dispatch. arms = [(0, None), (0, "0"), *[(e, None) for e in range(1, 9)]] cases = [(m, sp) for m in range(2, 9) for sp in ("1", "2")] calls = [ @@ -401,14 +401,54 @@ def test_x16_load_bits(codec, dtype): for (off, lever), r in zip(arms, group, strict=True): tag = (codec, m, sp, off, lever) assert r["error"] is None, (tag, r["error"]) - end = "_x16(_pf32)?$" if lever is None and off % 8 == 0 else "_x8(_pf32)?$" + b = 2 * off + x8 = lever is not None or b % 16 == 8 + end = "_x8(_pf32)?$" if x8 else "_x16(_pf32)?$" names = [n for n, _ in r["kernels"]] assert names and all( "verify_nax_multi" in n and re.search(end, n) for n in names ), (tag, names) + copied = [("quantized_matmul_multi x", b % 8)] if b % 8 else [] + assert r["copies"] == copied, (tag, r["copies"]) + assert set(r["xa"]) == {0 if b % 8 else b % 16}, (tag, r["xa"]) assert r["digest"] == group[0]["digest"], tag +@pytest.mark.skipif(force_cpu(), reason="GPU copies") +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize( + "lever", + [{"KQ_VERIFY_NAX": "0"}, {"KQ_QMM_ROUTE": "qmv"}], + ids=lambda e: ",".join(f"{k}={v}" for k, v in e.items()), +) +def test_per_weight_copies_x_once(lever, dtype): + # Off verify_nax, the op runs quantized_matmul per weight on one aligned + # copy of x: one copy for the three weights, every kernel reads x on a + # 16-byte boundary, and the bits match x at offset 0. + calls = [ + { + "codec": "q4_k", + "shape": [4, 1024], + "n": 1024, + "multi": [256, 512], + "dtype": dtype, + "seed": 4, + "x_offset": off, + } + for off in range(9) + ] + res = run_probe(calls, env={"KQ_QMM_LOG": "1", **lever}, check=False) + for off, r in enumerate(res): + b = 2 * off + assert r["error"] is None, (off, r["error"]) + names = [n for n, _ in r["kernels"]] + assert names and not any("multi" in n for n in names), (off, names) + copied = [("quantized_matmul_multi x", b % 16)] if b % 16 else [] + assert r["copies"] == copied, (off, r["copies"]) + assert set(r["xa"]) == {0}, (off, r["xa"]) + assert r["digest"] == res[0]["digest"], off + + @NAX_ONLY def test_lever_read_once_per_process(): # KQ_MATMUL_MULTI keeps its first read, as KQ_GLU, unless KQ_TUNE_LIVE=1. diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index 8e37003..0687e5f 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -682,8 +682,9 @@ def _depth(entry): return _one_kernel(entry)[1][2] -# Kernel name fragment of each route that logs. The mat-vec routes (qmv, -# mv_ext, verify_qmv) log nothing. +# Kernel name fragment of each route but the mat-vec routes (qmv, mv_ext, +# verify_qmv), whose kernels hold one of MATVEC_KERNEL. +MATVEC_KERNEL = ("_qmv_", "_mv_ext_") ROUTE_KERNEL = { "verify_nax": "_verify_nax_", "verify_mma": "_verify_mma_", @@ -693,14 +694,14 @@ def _depth(entry): def _route_of(entry): - """The route an entry of _logged ran: a ROUTE_KERNEL key, "tile" for - another logged kernel, or "matvec" when nothing logged.""" + """The route an entry of _logged ran: "matvec" for a mat-vec kernel, a + ROUTE_KERNEL key, or "tile" for another kernel.""" error, kernels = entry assert error is None, error - if not kernels: - return "matvec" assert len(kernels) == 1, kernels name = kernels[0][0] + if any(frag in name for frag in MATVEC_KERNEL): + return "matvec" return next((r for r, frag in ROUTE_KERNEL.items() if frag in name), "tile") @@ -1323,17 +1324,26 @@ def test_verify_nax_disable_nax(codec, monkeypatch): # The verify_nax codecs that read one activation step at a time read two -# steps in one 16-byte load where x starts on a 16-byte boundary. Each arm -# multiplies the same values in the same order: x at offset 0 (the paired -# load), at offset 0 with KQ_VNAX_X16=0, at every 2-byte offset from 2 to -# 14 bytes (one load per step) and at 16 bytes (the paired load off the -# buffer start). +# steps in one 16-byte load where x starts on a 16-byte boundary, one step +# per 8-byte load where it starts 8 bytes past one, and x from an aligned +# copy elsewhere. Each arm multiplies the same values in the same order: x +# at offset 0 (the paired load), at offset 0 with KQ_VNAX_X16=0, at every +# 2-byte offset from 2 to 14 bytes (8 bytes: one load per step; the others: +# the copy, paired) and at 16 bytes (the paired load off the buffer start). X16_ARMS = [(0, None), (0, "0"), *[(e, None) for e in range(1, 9)]] def x16_variant(off, lever): # The KQ_QMM_LOG name end of the pipeline an arm takes. - return "_x16(_pf32)?$" if lever is None and off % 8 == 0 else "_x8(_pf32)?$" + x8 = lever is not None or (2 * off) % 16 == 8 + return "_x8(_pf32)?$" if x8 else "_x16(_pf32)?$" + + +def x16_reads(off, what): + # (copies, xa): the copy an arm makes of x and the x offset its kernel + # logs. x off an 8-byte boundary is copied; x8 reads x 8 bytes in. + b = 2 * off + return ([(what, b % 8)], [0]) if b % 8 else ([], [b % 16]) X16_SPLITS = [None, "3"] @@ -1385,6 +1395,8 @@ def test_verify_nax_x16_bits(codec): ((name, _),) = r["kernels"] assert "_verify_nax_" in name, (tag, name) assert re.search(x16_variant(off, lever), name), (tag, off, lever, name) + copies, xa = x16_reads(off, "quantized_matmul x") + assert (r["copies"], r["xa"]) == (copies, xa), (tag, off, lever) assert r["digest"] == arms[0]["digest"], (tag, off, lever) diff --git a/tests/test_prism_mv_range.py b/tests/test_prism_mv_range.py index 82d1c12..8144799 100644 --- a/tests/test_prism_mv_range.py +++ b/tests/test_prism_mv_range.py @@ -75,7 +75,8 @@ def test_q2_0_mv_range(design, dtype, m, route): def _check(r, dtype): assert r["error"] is None, r["error"] - # The mat-vec routes log no tile kernel. - assert r["kernels"] == [], r["kernels"] + # One mat-vec kernel, no tile. + ((name, _),) = r["kernels"] + assert "_qmv_" in name, name assert r["nonfinite"] == 0, r["nonfinite"] assert _excess(r["rel"], r["floor1"]) < EXCESS[dtype], (r["rel"], r["floor1"]) diff --git a/tests/test_split_kernel_dispatch.py b/tests/test_split_kernel_dispatch.py index bd2e3a5..20bbcd3 100644 --- a/tests/test_split_kernel_dispatch.py +++ b/tests/test_split_kernel_dispatch.py @@ -6,7 +6,7 @@ partials in the activation dtype and in float (KQ_SPLITK_F32), and with the half verify_mma kernels (KQ_VERIFY_MMA_HALF), and run quantized_matmul_multi for its joined verify_nax kernels. Every verify_nax call also runs with x -2 bytes into its buffer, so each verify_nax kernel runs both its 16-byte +8 bytes into its buffer, so each verify_nax kernel runs both its 16-byte (_x16) and 8-byte (_x8) x load pipelines. KQ_QMM_LOG names the pipeline each call ran. Together the calls must run every name the host enumerator lists for the codec, each verify_nax name in both pipelines, and each @@ -131,7 +131,7 @@ def test_every_split_kernel_runs(codec): for f32 in ("0", "1"): for n in (256, 250): for family, m, k, env in arms: - for off in (0, 1) if family == "_verify_nax_" else (0,): + for off in (0, 4) if family == "_verify_nax_" else (0,): calls.append( { "codec": codec, @@ -147,7 +147,7 @@ def test_every_split_kernel_runs(codec): for m, k, splits in MULTI_ARMS: if k % CODEC_GEOMETRY[codec][3]: continue - for off in (0, 1): + for off in (0, 4): calls.append( { "codec": codec, diff --git a/tests/test_splitk_cap.py b/tests/test_splitk_cap.py index b03cd59..0f04f83 100644 --- a/tests/test_splitk_cap.py +++ b/tests/test_splitk_cap.py @@ -9,8 +9,7 @@ is no fallthrough here. Each capped call must match the call with the split route off bit for bit, and float64 within the R1 bound. A partial cap cuts the slice count, a forced target is not capped, and 0 keeps the -fixed target. The mat-vec routes do not log, so the output digest shows -which of them ran. +fixed target. The logged kernel shows which route ran. Unset, the cap keeps the fixed target on the ALU route, and the NAX route takes its default: one slice at a grid of 16 threadgroups per core or more @@ -107,10 +106,10 @@ def test_nax_cap_falls_through(codec, m): _split(default, "_qmm_t_nax_splitk_") assert default["rel"] < R1, default["rel"] _same(capped, off) + ((name, grid),) = capped["kernels"] if (codec, m) in MV_EXT: - assert capped["kernels"] == [], capped["kernels"] + assert "_mv_ext_" in name, name else: - ((name, grid),) = capped["kernels"] assert tile_of(name)[0] == 32 and grid[2] == 1, (name, grid) @@ -139,7 +138,8 @@ def test_nax_cap_falls_through_to_qmv(m): forced, capped, off = _nax_qmv_fallthrough()[m] _split(forced, "_qmm_t_nax_splitk_") _same(capped, off) - assert capped["kernels"] == [] + ((name, _),) = capped["kernels"] + assert "_qmv_" in name, name # N 1024 at M 16 is 16 BM=32 tiles, and K 5120 is 20 units of 256 weights, diff --git a/tests/test_x_offset_views.py b/tests/test_x_offset_views.py new file mode 100644 index 0000000..6dbffd4 --- /dev/null +++ b/tests/test_x_offset_views.py @@ -0,0 +1,1030 @@ +"""Operands at an element offset into their buffer. + +A row-contiguous view can start at any element of its buffer, and an op's +contiguous node passes it on in place. Kernels load activations, caches and +some float operands as vectors, and Metal defines a vector load only at an +address aligned to its size. So each op binds such an operand that starts +off the boundary of its kernel's widest load from a copy, and a strided +operand whose rows step off it too; the matmul ops bind x from a 16-byte +boundary. Each case runs an op on views at offsets of 1 to 8 elements and +requires the output bits of the same values in their own buffer. +KQ_QMM_LOG prints a "[kq align] xin=" line for each copy, +and every matmul dispatch logs x on a 16-byte boundary (xa=0), but +verify_nax's: it reads x where it lands on an 8-byte boundary (xa=8) and +from the copy otherwise. The cases require all three: each view that +starts off the boundary reaches its op there and is copied or read on its +kernel's boundary, each view on the boundary is read in place, and no +matmul kernel reads x off it. + +- quantized_matmul for every codec with a matmul kernel, at 1 to 64 rows, + on the default routes, with the NAX mat-vec off (KQ_NAX_QMV=0), with NAX + off (KQ_DISABLE_NAX=1, the routes of a GPU without NAX), and on each + forced route (KQ_QMM_ROUTE); +- the transpose=False, GLU, LoRA (also with the fused LoRA kernel, + KQ_QMM_LORA=2), residual and bias matmuls, skinny_matmul, gather_qmm, + gather_qmm_seg, the fused MoE gathers and gather_mix, in a child per route + environment; +- rmsnorm_gate, latent_fp4_pack, dsa_indexer_score_decode, the hc glue ops, + the attention ops and the sparse attention ops, in a child with NAX on and + off, each operand read as vectors at the offsets and, for the keys and + values the attention ops read through their strides, as rows that step + one element past their length. + +Float32 activations of the matmul ops are cast to bfloat16 into a new +buffer, so their views reach the kernels aligned. Under KQUANT_FORCE_CPU +the quantized_matmul cases run on the CPU stream, which needs no alignment; +the other ops are Metal-only kernels with no CPU eval, and are skipped. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from kqref import is_synth, synth_wire +from qmm_probe import TESTS_DIR, force_cpu, run_probe + +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY + +OFFSETS = tuple(range(1, 9)) +DTYPES = ("float16", "bfloat16", "float32") +ROWS = (1, 2, 4, 8, 12, 16, 32, 64) +STRICT = {"KQ_QMM_ROUTE_STRICT": "1"} +# (name, env, rows). The forced routes run at the rows they serve; a route +# that declines a codec or shape raises, on the offset views too. +ROUTES = [ + ("default", {}, ROWS), + ("nax_qmv_off", {"KQ_NAX_QMV": "0"}, ROWS), + ("nax_off", {"KQ_DISABLE_NAX": "1"}, ROWS), + ("qmv", {"KQ_QMM_ROUTE": "qmv", **STRICT}, (1, 2, 8)), + ("verify_qmv", {"KQ_QMM_ROUTE": "verify_qmv", **STRICT}, (2, 8)), + ("mv_ext", {"KQ_QMM_ROUTE": "mv_ext", **STRICT}, (2, 8, 12)), + ("verify_mma", {"KQ_QMM_ROUTE": "verify_mma", **STRICT}, (2, 8)), + ("verify_nax", {"KQ_QMM_ROUTE": "verify_nax", **STRICT}, (2, 8)), + ("splitk", {"KQ_QMM_ROUTE": "splitk", **STRICT}, (8, 32)), + ("nax", {"KQ_QMM_ROUTE": "nax", **STRICT}, (8, 32)), + ("nax_splitk", {"KQ_QMM_ROUTE": "nax_splitk", **STRICT}, (8, 32)), +] +CPU_ROUTES = [("default", {}, ROWS)] +MATMUL_CODECS = [c for c in kq.codecs() if kq.codec_has_matmul(c)] +ITEMSIZE = {"float16": 2, "bfloat16": 2, "float32": 4} + + +def _unaligned_x(name, xa): + """Whether a kernel read x off its boundary: 16 bytes, or 8 for + verify_nax, which reads x where it lands on an 8-byte boundary.""" + return xa % (8 if "_verify_nax_" in name else 16) != 0 + + +metal_only = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="Metal-only kernel; no CPU path.", +) + + +def quantized_matmul_offsets(codec): + """(bad, unaligned, unreached, copied) over the routes: the offset calls + whose output or error differs from the call on x in its own buffer, the + kernels that read x off a 16-byte boundary with the byte offset they + logged, the calls on a view off the boundary that log neither a copy of + x nor a kernel reading x there, and the calls on a view on the boundary + that log a copy.""" + bad, unaligned, unreached, copied = [], set(), [], [] + for route, env, rows in CPU_ROUTES if force_cpu() else ROUTES: + calls, cells = [], [] + for dtype in DTYPES: + # float32 activations are cast to bfloat16 into a new buffer: + # the default route covers them. + if dtype == "float32" and route != "default": + continue + for m in rows: + for off in (0, *OFFSETS): + calls.append( + { + "codec": codec, + "shape": [m, 4096], + "n": 256, + "dtype": dtype, + "seed": m, + "x_offset": off, + } + ) + cells.append((dtype, m, off)) + res = run_probe(calls, env={"KQ_QMM_LOG": "1", **env}, check=False) + base = {} + for (dtype, m, off), r in zip(cells, res, strict=True): + got = r["error"] or r["digest"] + names = [k for k, _ in r["kernels"]] + if off == 0: + base[(dtype, m)] = got + elif got != base[(dtype, m)]: + bad.append((route, dtype, m, off, names, r["xa"])) + for name, xa in zip(names, r["xa"], strict=True): + if _unaligned_x(name, xa): + unaligned.add((name, xa)) + off_boundary = (off * ITEMSIZE[dtype]) % 16 != 0 + if not off_boundary and r["copies"]: + copied.append((route, dtype, m, off, r["copies"])) + if ( + off_boundary + and dtype != "float32" + and not force_cpu() + and not r["error"] + and not any(xin for _, xin in r["copies"]) + and not any(r["xa"]) + ): + unreached.append((route, dtype, m, off, names)) + return bad, sorted(unaligned), unreached, copied + + +@pytest.mark.parametrize("codec", MATMUL_CODECS) +def test_quantized_matmul_x_offsets(codec): + """Bit-exact output at every offset, every kernel reads x from a + 16-byte boundary, every view off it reaches the op there, and every + view on it is read in place.""" + bad, unaligned, unreached, copied = quantized_matmul_offsets(codec) + assert bad == [], bad + assert unaligned == [], unaligned + assert unreached == [], unreached + assert copied == [], copied + + +# -------------------------------- other ops -------------------------------- +# Each group runs in a child interpreter per environment, since the route +# levers are read once per process. The child logs its dispatches and +# copies, and marks each offset call on stderr with the copies it needs. + +ENVS = { + "default": {}, + "nax_qmv_off": {"KQ_NAX_QMV": "0"}, + "nax_off": {"KQ_DISABLE_NAX": "1"}, + "lora2": {"KQ_QMM_LORA": "2"}, +} +K, N = 768, 256 +OP_CODECS = ("q4_0", "q4_k", "q6_k", "q8_0", "iq4_nl", "iq2_xxs") +MOE_CODECS = ("q4_k", "q8_0", "iq4_nl", "iq2_xxs") +MX_DT = {"float16": mx.float16, "bfloat16": mx.bfloat16, "float32": mx.float32} +HALF = ("float16", "bfloat16") +_W = {} + + +def _own_buffer(make): + """make() evaluated into a buffer of its own size. A buffer reused from + the cache can run more than 16 KiB past the view, and the op's + contiguous node then copies the view before the eval sees it; buffers + freed late by earlier work could refill an emptied cache, so the cache + stays off while the buffer is allocated.""" + limit = mx.set_cache_limit(0) + mx.clear_cache() + try: + a = make() + mx.eval(a) + finally: + mx.set_cache_limit(limit) + return a + + +def _at_offset(x, off): + """The values of x as a row-contiguous view off elements into a buffer + of their own.""" + flat = _own_buffer( + lambda: mx.concatenate([mx.zeros((off,), x.dtype), x.reshape(-1)]) + ) + return flat[off:].reshape(x.shape) + + +def _rows_view(x): + """The values of x in rows that step one element past their length, so + that every row after the first starts off a 16-byte boundary.""" + pad = mx.zeros((*x.shape[:-1], 1), x.dtype) + return _own_buffer(lambda: mx.concatenate([x, pad], axis=-1))[..., :-1] + + +def _mark(label, need, names=()): + """need: "off" when the call must copy each operand in names (or, with + no names, log a copy or a kernel reading x) at a nonzero offset, "copy" + when it must copy each of them at any offset, "none" when it must copy + no operand, "any" for no requirement.""" + print(f"[arm] {label} | {need} | {','.join(names)}", file=sys.stderr, flush=True) + + +def _listify(out): + return list(out) if isinstance(out, (list, tuple)) else [out] + + +def _same(got, want): + return all( + g.dtype == w.dtype and mx.array_equal(g, w).item() + for g, w in zip(got, want, strict=True) + ) + + +def _x(shape, dtype, seed): + x = (mx.random.normal(shape, key=mx.random.key(seed)) * 0.5).astype(MX_DT[dtype]) + mx.eval(x) + return x + + +def _weight(codec, n, k, seed=0, experts=None): + """Weights [n, k], or [experts, n, k], and scales. The codecs that + quantize only with an importance matrix get random valid wire.""" + key = (codec, n, k, seed, experts) + if key not in _W: + e = experts or 1 + if is_synth(codec): + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + rng = np.random.default_rng(seed * 131 + n + k) + wire = synth_wire(rng, codec, bpb, e * n * (k // wpb)) + w = mx.array(wire.reshape(e, n, (k // wpb) * bpb)) + w = w if experts else w[0] + s = mx.zeros((1,), dtype=mx.uint8) + else: + shape = (n, k) if experts is None else (experts, n, k) + wf = mx.random.normal(shape, key=mx.random.key(seed)) * 0.1 + w, s = kq.quantize(wf, codec) + w = mx.contiguous(w) + mx.eval(w, s) + _W[key] = (w, s) + return _W[key] + + +def _offset_failures(call, shape, dtype, seed=0, label="", names=()): + """The offsets at which call(view) differs from call(x) in its own + buffer. A call the op refuses is skipped. A float32 view is cast into a + new buffer, so only the 16-bit views must reach the op off the + boundary; every view on it must reach the op uncopied.""" + x = _x(shape, dtype, seed) + try: + want = _listify(call(x)) + mx.eval(*want) + except (ValueError, RuntimeError): + return [] + bad = [] + for off in OFFSETS: + view = _at_offset(x, off) + need = "none" + if (off * ITEMSIZE[dtype]) % 16: + need = "off" if dtype in HALF else "any" + _mark(f"{label} {dtype} {list(shape)} off{off}", need, names) + got = _listify(call(view)) + mx.eval(*got) + if not _same(got, want): + bad.append(off) + return bad + + +def _operand_failures(label, call, ops, cases): + """call(**ops) with one operand at a time replaced by its views. cases: + (operand, the copies it needs, vec, strided), vec the elements of the + kernel's widest vector load of the operand. Each view at an offset off + a boundary of vec elements must be copied, and each view on one must + not; a strided operand also runs as rows that step one element past + that boundary, which must be copied too.""" + want = _listify(call(**ops)) + mx.eval(*want) + bad = [] + for name, names, vec, strided in cases: + a = ops[name] + arms = [(f"off{off}", off, "off" if off % vec else "none") for off in OFFSETS] + if strided: + arms.append(("rows", None, "copy")) + for arm, off, need in arms: + view = _rows_view(a) if off is None else _at_offset(a, off) + _mark(f"{label} {a.dtype} {name} {arm}", need, names) + got = _listify(call(**{**ops, name: view})) + mx.eval(*got) + if not _same(got, want): + bad.append([label, str(a.dtype), name, arm]) + return bad + + +def _matmul_ops(): + """transpose=False, GLU, LoRA and residual matmuls.""" + bad = [] + a = (mx.random.normal((K, 8), key=mx.random.key(5)) / np.sqrt(K)).astype( + mx.bfloat16 + ) + b = (mx.random.normal((8, N), key=mx.random.key(6)) * 0.3).astype(mx.bfloat16) + mx.eval(a, b) + for codec in OP_CODECS: + w, _ = _weight(codec, N, K) + wg, _ = _weight(codec, N, K, 1) + wt, st = _weight(codec, K, N, 2) + + def resid(x, w=w, codec=codec): + out = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + r = (mx.random.normal((x.shape[0], N), key=mx.random.key(7)) * 0.5).astype( + out + ) + return kq.quantized_matmul_residual(x, w, codec, r) + + def lora(x, w=w, codec=codec): + dt = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + return kq.quantized_matmul_lora(x, w, codec, a.astype(dt), b.astype(dt)) + + def glu(x, w=w, wg=wg, codec=codec): + return kq.quantized_matmul_glu(x, wg, w, codec) + + def matmul_n(x, wt=wt, st=st, codec=codec): + return kq.quantized_matmul(x, wt, st, codec, False) + + ops = [ + ("matmul_n", (1, 8, 64), matmul_n), + ("glu", (1, 8, 64, 1024), glu), + ("lora", (1, 8, 64, 1024), lora), + ("residual", (1, 8, 64, 1024), resid), + ] + for name, rows, call in ops: + for dtype in DTYPES: + for m in rows: + offs = _offset_failures(call, (m, K), dtype, m, f"{name} {codec}") + if offs: + bad.append([name, codec, dtype, m, offs]) + return bad + + +def _bias_and_skinny(): + """The q8_0 bias mat-vec and skinny_matmul, whose dense weight is read + as vectors too.""" + bad = [] + w, s = _weight("q8_0", N, K) + bias = mx.random.normal((N,), key=mx.random.key(3)).astype(mx.float32) + for dtype in DTYPES: + offs = _offset_failures( + lambda x: kq.quantized_matmul_qmv_bias(x, w, s, bias, "q8_0"), + (1, K), + dtype, + label="qmv_bias", + ) + if offs: + bad.append(["qmv_bias", dtype, 1, offs]) + wd = (mx.random.normal((64, K), key=mx.random.key(4)) * 0.5).astype( + MX_DT[dtype] + ) + mx.eval(wd) + for m in (1, 4, 16): + offs = _offset_failures( + lambda x, wd=wd: kq.skinny_matmul(x, wd), (m, K), dtype, label="skinny" + ) + if offs: + bad.append(["skinny", dtype, m, offs]) + if dtype in HALF: + x = _x((m, K), dtype, m) + bad += _operand_failures( + f"skinny w m{m}", + lambda x, w: kq.skinny_matmul(x, w), + {"x": x, "w": wd}, + [("w", ("skinny_matmul w",), 8, False)], + ) + return bad + + +def _gathers(): + """gather_qmm: the mat-vec, the tiled gather, the sorted per-expert + GEMM and transpose=False; and gather_qmm_seg.""" + experts, bad = 4, [] + rng = np.random.default_rng(4) + for codec in MOE_CODECS: + wt, st = _weight(codec, N, K, 3, experts) + wn, sn = _weight(codec, K, N, 4, experts) + for rows, m, transpose, sort in [ + (8, 1, True, False), + (4, 32, True, False), + (256, 1, True, True), + (8, 1, False, False), + ]: + idx = rng.integers(0, experts, rows).astype(np.uint32) + rhs = mx.array(np.sort(idx) if sort else idx) + w, s = (wt, st) if transpose else (wn, sn) + + def call(x, w=w, s=s, rhs=rhs, transpose=transpose, sort=sort, c=codec): + return kq.gather_qmm( + x, + w, + s, + c, + rhs_indices=rhs, + transpose=transpose, + sorted_indices=sort, + ) + + for dtype in DTYPES: + offs = _offset_failures( + call, (rows, m, K), dtype, label=f"gather_qmm {codec}" + ) + if offs: + bad.append(["gather_qmm", codec, dtype, rows, m, transpose, offs]) + idx = np.sort(rng.integers(0, experts, 200).astype(np.uint32)) + maps = kq.expert_tile_map(mx.array(idx), experts) + for dtype in DTYPES: + offs = _offset_failures( + lambda x, wt=wt, st=st, c=codec, mp=maps: kq.gather_qmm_seg( + x, wt, st, c, *mp + ), + (len(idx), K), + dtype, + label=f"gather_qmm_seg {codec}", + ) + if offs: + bad.append(["gather_qmm_seg", codec, dtype, len(idx), offs]) + return bad + + +T, R = 3, 2 +S = R + 1 +ME, MN, MK = 4, 64, 512 + + +def _routing(seed, slots): + rng = np.random.default_rng(seed) + inds = mx.array(rng.integers(0, ME, size=(T, slots)).astype(np.uint32)) + sc = mx.array(rng.uniform(0.05, 0.9, size=(T, slots)).astype(np.float32)) + mx.eval(inds, sc) + return inds, sc + + +def _kq_moe_ops(codec): + """name -> (shape of x, call(x)) for the K-quant fused MoE ops.""" + gw, uw = _weight(codec, MN, MK, 1, ME)[0], _weight(codec, MN, MK, 2, ME)[0] + shexp = _weight(codec, MN, MK, 3, ME)[0] + sgw, suw = mx.contiguous(shexp[0]), mx.contiguous(shexp[1]) + inds, _ = _routing(3, R) + inds_s, sc_s = _routing(4, S) + la = (mx.random.normal((ME, MK, 4), key=mx.random.key(8)) * 0.05).astype(mx.float16) + lb = (mx.random.normal((ME, 4, MN), key=mx.random.key(9)) * 0.05).astype(mx.float16) + mx.eval(sgw, suw, la, lb) + return { + "moe_glu_gather_kq": ( + (T, MK), + lambda x: kq.moe_glu_gather_kq(x, gw, uw, codec, inds), + ), + "gather_qmv_kq": ((T, R, MK), lambda x: kq.gather_qmv_kq(x, gw, codec, inds)), + "gather_qmv_kq_lora": ( + (T, R, MK), + lambda x: kq.gather_qmv_kq(x, gw, codec, inds, lora_a=la, lora_b=lb), + ), + "moe_glu_gather_shexp_kq": ( + (T, MK), + lambda x: kq.moe_glu_gather_shexp_kq(x, gw, uw, sgw, suw, codec, inds), + ), + "gather_qmv_mix_kq": ( + (T, S, MK), + lambda x: kq.gather_qmv_mix_kq(x, gw, sgw, codec, inds, sc_s), + ), + "gather_qmv_mix_ns_kq": ( + (T, S, MK), + lambda x: kq.gather_qmv_mix_ns_kq(x, gw, codec, inds_s, sc_s), + ), + "gather_qmv_mix_ns_kq_lora": ( + (T, S, MK), + lambda x: kq.gather_qmv_mix_ns_kq( + x, gw, codec, inds_s, sc_s, lora_a=la, lora_b=lb + ), + ), + } + + +def _mxfp4_moe_ops(): + """name -> (shape of x, call(x)) for the mxfp4 fused MoE ops.""" + rng = np.random.default_rng(1) + + def stack(): + w = mx.array(rng.integers(0, 2**32, size=(ME, MN, MK // 8), dtype=np.uint32)) + s = mx.array(rng.integers(118, 124, size=(ME, MN, MK // 32), dtype=np.uint8)) + b = mx.array(rng.uniform(-0.2, 0.2, (ME, MN)).astype(np.float32)) + mx.eval(w, s, b) + return w, s, b + + gw, gs, gb = stack() + uw, us, ub = stack() + inds, sc = _routing(3, R) + return { + "moe_glu_gather": ( + (T, MK), + lambda x: kq.moe_glu_gather(x, gw, gs, gb, uw, us, ub, inds), + ), + "gather_qmv_bias": ( + (T, R, MK), + lambda x: kq.gather_qmv_bias(x, gw, gs, gb, inds), + ), + "gather_qmv_mix_bias": ( + (T, R, MK), + lambda x: kq.gather_qmv_mix_bias(x, gw, gs, gb, inds, sc), + ), + } + + +def _moe(): + """The fused MoE gathers, K-quant and mxfp4.""" + bad = [] + for codec in (*MOE_CODECS, "mxfp4"): + ops = _mxfp4_moe_ops() if codec == "mxfp4" else _kq_moe_ops(codec) + for name, (shape, call) in ops.items(): + for dtype in DTYPES: + offs = _offset_failures(call, shape, dtype, label=f"{name} {codec}") + if offs: + bad.append([name, codec, dtype, offs]) + return bad + + +def _gather_mix(): + """gather_mix reads the expert outputs y, its activation, as vectors.""" + tokens, k, n = 5, 4, 256 + order = np.random.default_rng(6).permutation(tokens * k).astype(np.uint32) + inv = mx.array(order) + sc = mx.array( + np.random.default_rng(7).uniform(0.1, 0.9, (tokens, k)).astype(np.float32) + ) + mx.eval(inv, sc) + bad = [] + for dtype in HALF: + offs = _offset_failures( + lambda y: kq.gather_mix(y, inv, sc), + (tokens * k, n), + dtype, + label="gather_mix", + ) + if offs: + bad.append(["gather_mix", dtype, offs]) + return bad + + +def _rand(shape, dtype, seed, scale=0.5): + a = (mx.random.normal(shape, key=mx.random.key(seed)) * scale).astype(dtype) + mx.eval(a) + return a + + +def _norm_latent_indexer(): + """rmsnorm_gate, latent_fp4_pack and dsa_indexer_score_decode.""" + bad = [] + for dt in (mx.float16, mx.bfloat16): + for d in (64, 256): + ops = { + "x": _rand((6, d), dt, 1), + "w": _rand((d,), dt, 2) + 1.0, + "gate": _rand((6, d), dt, 3), + } + bad += _operand_failures( + f"rmsnorm_gate d{d}", + lambda x, w, gate: kq.rmsnorm_gate(x, w, gate, 1e-6), + ops, + [(n, (f"rmsnorm_gate {n}",), 4, False) for n in ("x", "w", "gate")], + ) + rng = np.random.default_rng(11) + for heads, ql, rows, qo, ratio in ((4, 1, 64, 63, 1), (64, 2, 1027, 4106, 4)): + ops = { + "q": mx.array(rng.standard_normal((1, heads, ql, 128)) * 0.4).astype( + dt + ), + "k": mx.array(rng.standard_normal((1, rows, 128)) * 0.4).astype(dt), + "w": mx.array(rng.standard_normal((1, ql, heads)) * 0.2).astype(dt), + } + bad += _operand_failures( + f"dsa_indexer_score_decode h{heads}", + lambda q, k, w, qo=qo, ratio=ratio: kq.dsa_indexer_score_decode( + q, k, w, qo, ratio + ), + ops, + [(n, (f"dsa_indexer_score_decode {n}",), 8, False) for n in ("q", "k")], + ) + for dt in (mx.float16, mx.bfloat16, mx.float32): + bad += _operand_failures( + "latent_fp4_pack", + lambda x: kq.latent_fp4_pack(x), + {"x": _rand((8, 64), dt, 4)}, + [("x", ("latent_fp4_pack x",), 8, False)], + ) + return bad + + +def _hc(): + """The hc glue ops: the stream, residual and collapsed rows, the + projection fn, the mixes and the base biases.""" + return [b for dt in (mx.float16, mx.bfloat16) for b in _hc_dtype(dt)] + + +def _hc_dtype(dt): + from test_hc_glue import HC, HC_EPS, ITERS, MIX, NORM_EPS + + d = 1024 + sk = {"iters": ITERS, "hc_eps": HC_EPS, "norm_eps": NORM_EPS} + bad = [] + x = _rand((1, 1, HC, d), dt, 1, 0.05) + fn = _rand((MIX, HC * d), mx.float32, 2, 0.02) + scale = mx.array([1.1, 0.9, 1.3], dtype=mx.float32) + base = _rand((MIX,), mx.float32, 3, 0.1) + w = _rand((d,), dt, 4, 0.1) + 1.0 + mr, ssq = kq.hc_front_reduce(x, fn) + col, post, comb = kq.hc_sinkhorn_collapse(x, mr, ssq, scale, base, w, **sk) + mx.eval(scale, w, mr, ssq, col, post, comb) + + def case(op, *names): + return [(n, (f"{op} {n}",), 4, False) for n in names] + + bad += _operand_failures( + "hc_front_reduce", + lambda x, fn: kq.hc_front_reduce(x, fn), + {"x": x, "fn": fn}, + case("hc_front_reduce", "x", "fn"), + ) + bad += _operand_failures( + "hc_sinkhorn_collapse", + lambda x, mixes, base: kq.hc_sinkhorn_collapse( + x, mixes, ssq, scale, base, w, **sk + ), + {"x": x, "mixes": mr, "base": base}, + case("hc_sinkhorn_collapse", "x", "mixes", "base"), + ) + bad += _operand_failures( + "hc_expand", + lambda x, resid: kq.hc_expand(x, resid, post, comb), + {"x": col, "resid": x}, + case("hc_expand", "x", "resid"), + ) + bad += _operand_failures( + "hc_front_expand_reduce", + lambda x_sub, resid, fn: kq.hc_front_expand_reduce( + x_sub, resid, post, comb, fn + ), + {"x_sub": col, "resid": x, "fn": fn}, + case("hc_front_expand_reduce", "x_sub", "resid", "fn"), + ) + bad += _operand_failures( + "hc_front_expand_collapse", + lambda x_sub, resid, fn, base: kq.hc_front_expand_collapse( + x_sub, resid, post, comb, fn, scale, base, w, **sk + ), + {"x_sub": col, "resid": x, "fn": fn, "base": base}, + case("hc_front_expand_collapse", "x_sub", "resid", "fn", "base"), + ) + return bad + + +def _qkv(op, strided=("k", "v"), q="q"): + """Cases for an attention op's query and its keys and values.""" + return [(q, (f"{op} {q}",), 4, False)] + [ + (n, (f"{op} {n}",), 4, True) for n in strided + ] + + +def _sdpa(): + """The attention ops: q, k and v, the q8 cache words and the KVarN + stages.""" + bad = [b for dt in (mx.float16, mx.bfloat16) for b in _sdpa_dtype(dt)] + return bad + _sdpa_kvarn() + + +def _sdpa_dtype(dt): + from test_sdpa import _bs_prefill_lists, _make_indexed + + bad = [] + q = _rand((1, 8, 1, 128), dt, 1) + k = _rand((1, 2, 512, 128), dt, 2) + v = _rand((1, 2, 512, 128), dt, 3) + qkv = {"q": q, "k": k, "v": v} + bad += _operand_failures( + "sdpa_vector", + lambda q, k, v: kq.sdpa_vector(q, k, v, 0.0625), + { + "q": _rand((1, 8, 2, 256), dt, 15), + "k": _rand((1, 2, 300, 256), dt, 16), + "v": _rand((1, 2, 300, 256), dt, 17), + }, + _qkv("sdpa_vector"), + ) + bad += _operand_failures( + "sdpa_decode_gqa", + lambda q, k, v: kq.sdpa_decode_gqa(q, k, v, 0.0884), + qkv, + _qkv("sdpa_decode_gqa"), + ) + pages = mx.array(np.broadcast_to(np.arange(16, dtype=np.int32), (1, 2, 16))) + bad += _operand_failures( + "sdpa_decode_gqa_paged", + lambda q, k, v: kq.sdpa_decode_gqa_paged(q, k, v, 0.0884, pages), + qkv, + _qkv("sdpa_decode_gqa"), + ) + # Head dim 512 at the depth that runs pass 1 on NAX. + wide = { + "q": _rand((1, 8, 1, 512), dt, 4), + "k": _rand((1, 1, 4096, 512), dt, 5), + "v": _rand((1, 1, 4096, 512), dt, 6), + } + bad += _operand_failures( + "sdpa_decode_gqa d512", + lambda q, k, v: kq.sdpa_decode_gqa(q, k, v, 0.0442), + wide, + _qkv("sdpa_decode_gqa"), + ) + # A q8 cache: k and v are the packed words. + kw, ks, kb = mx.quantize(k, group_size=64, bits=8) + vw, vs, vb = mx.quantize(v, group_size=64, bits=8) + mx.eval(kw, ks, kb, vw, vs, vb) + bad += _operand_failures( + "sdpa_decode_gqa q8", + lambda k, v: kq.sdpa_decode_gqa( + q, k, v, 0.0884, k_scales=ks, k_biases=kb, v_scales=vs, v_biases=vb + ), + {"k": kw, "v": vw}, + _qkv("sdpa_decode_gqa")[1:], + ) + qf = _rand((1, 2, 12, 128), dt, 7) + bad += _operand_failures( + "sdpa_fa_verify", + lambda q, k, v: kq.sdpa_fa_verify(q, k, v, 0.0884, q_len=3), + {"q": qf, "k": k, "v": v}, + _qkv("sdpa_fa_verify"), + ) + qi, kv, idx = _make_indexed(8, 1, 100, 64, 0, dt, seed=9) + bad += _operand_failures( + "sdpa_fa_indexed", + lambda q, kv: kq.sdpa_fa_indexed(q, kv, idx, 0.0442), + {"q": qi, "kv": kv}, + _qkv("sdpa_fa_indexed", strided=("kv",)), + ) + # One sequence: the shared pass's folded query is q itself. + bad += _operand_failures( + "sdpa_decode_gqa_cascade", + lambda q, k_shared, v_shared, k, v: kq.sdpa_decode_gqa_cascade( + q, k_shared, v_shared, k, v, 0.0884 + ), + { + "q": q, + "k_shared": k, + "v_shared": v, + "k": _rand((1, 2, 64, 128), dt, 10), + "v": _rand((1, 2, 64, 128), dt, 11), + }, + [ + ( + "q", + ("sdpa_decode_gqa_cascade q", "sdpa_decode_gqa_cascade q_folded"), + 4, + False, + ) + ] + + _qkv("sdpa_decode_gqa_cascade", ("k_shared", "v_shared", "k", "v"))[1:], + ) + # A q8 cache: the shared pass stages the words in the fa_verify kernel. + kp, kps, kpb = mx.quantize(_rand((1, 2, 64, 128), dt, 12), group_size=64, bits=8) + vp, vps, vpb = mx.quantize(_rand((1, 2, 64, 128), dt, 13), group_size=64, bits=8) + mx.eval(kp, kps, kpb, vp, vps, vpb) + bad += _operand_failures( + "sdpa_decode_gqa_cascade q8", + lambda k_shared, v_shared, k, v: kq.sdpa_decode_gqa_cascade( + q, + k_shared, + v_shared, + k, + v, + 0.0884, + k_shared_scales=ks, + k_shared_biases=kb, + v_shared_scales=vs, + v_shared_biases=vb, + k_priv_scales=kps, + k_priv_biases=kpb, + v_priv_scales=vps, + v_priv_biases=vpb, + ), + {"k_shared": kw, "v_shared": vw, "k": kp, "v": vp}, + _qkv("sdpa_decode_gqa_cascade", ("k_shared", "v_shared", "k", "v"))[1:], + ) + # Four queries: the folded query is a view of q. + length, s, topk = 4, 1024, 16 + offset = s - length + sel = np.zeros((length, topk), dtype=np.int64) + rng = np.random.default_rng(3) + for t in range(length): + sel[t] = rng.choice((offset + t + 1) // 4, size=topk, replace=False) + lists = [mx.array(a) for a in _bs_prefill_lists(sel, offset, s)] + bad += _operand_failures( + "sdpa_prefill_block_sparse", + lambda q, k, v: kq.sdpa_prefill_block_sparse(q, k, v, 0.0625, *lists, offset), + { + "q": _rand((1, 24, length, 256), dt, 12), + "k": _rand((1, 2, s, 256), dt, 13), + "v": _rand((1, 2, s, 256), dt, 14), + }, + _qkv("sdpa_prefill_block_sparse"), + ) + return bad + + +def _sdpa_kvarn(): + """The KVarN stages, float16.""" + from test_kvarn_sdpa import build_state, make_q + + bad = [] + st = build_state(1, 2, 300, 6, 6, seed=21) + codes = {n: st[n] for n in ("codes_k", "axes_k", "codes_v", "axes_v")} + qk = mx.array(make_q(1, 8, 1)).astype(mx.float16) + qv = mx.array(make_q(1, 8, 4)).astype(mx.float16).reshape(1, 2, 16, 128) + stages = {"stage_k": st["stage_k"], "stage_v": st["stage_v"]} + bad += _operand_failures( + "sdpa_decode_gqa_kvarn", + lambda stage_k, stage_v: kq.sdpa_decode_gqa_kvarn( + qk, + codes["codes_k"], + codes["axes_k"], + codes["codes_v"], + codes["axes_v"], + stage_k, + stage_v, + 300, + 0.0884, + 6, + 6, + ), + stages, + [(n, (f"sdpa_decode_gqa {n}",), 4, False) for n in stages], + ) + bad += _operand_failures( + "sdpa_fa_verify_kvarn", + lambda stage_k, stage_v: kq.sdpa_fa_verify_kvarn( + qv, + codes["codes_k"], + codes["axes_k"], + codes["codes_v"], + codes["axes_v"], + stage_k, + stage_v, + 300, + 0.0884, + 6, + 6, + 4, + ), + stages, + [(n, (f"sdpa_fa_verify {n}",), 4, False) for n in stages], + ) + return bad + + +def _sparse(): + """The sparse attention ops: q and the window, pool and key rows.""" + return [b for dt in (mx.float16, mx.bfloat16) for b in _sparse_dtype(dt)] + + +def _sparse_dtype(dt): + from test_dsa_sparse_attn import _build as dsa_build + from test_sdpa_sparse_decode import _case as decode_case + from test_sdpa_sparse_decode import _on_grid_pool + from test_sdpa_sparse_prefill import _case as prefill_case + + bad = [] + (q, win, pool, idx, scale, kw), _ = decode_case( + 1, 16, 2, 128, 24, 300, 40, dt, np.int32, True, True, True + ) + mx.eval(q, win, pool, idx) + names = [(n, (f"sdpa_sparse_decode {n}",), 8, False) for n in ("q", "win", "pool")] + bad += _operand_failures( + "sdpa_sparse_decode", + lambda q, win, pool: kq.sdpa_sparse_decode(q, win, pool, idx, scale, **kw), + {"q": q, "win": win, "pool": pool}, + names, + ) + if dt == mx.float16: + codes, scales = kq.latent_fp4_pack(_on_grid_pool(pool)) + mx.eval(codes, scales) + bad += _operand_failures( + "sdpa_sparse_decode packed", + lambda pool: kq.sdpa_sparse_decode( + q, win, pool, idx, scale, pool_scales=scales, **kw + ), + {"pool": codes}, + [("pool", ("sdpa_sparse_decode pool",), 16, False)], + ) + (q, win, pool, idx, scale, band, kw), _ = prefill_case( + 1, 12, 9, 128, 30, 96, 24, 8, dt, np.int32, True, True, True + ) + mx.eval(q, win, pool, idx) + bad += _operand_failures( + "sdpa_sparse_prefill", + lambda q, win, pool: kq.sdpa_sparse_prefill( + q, win, pool, idx, scale, band, **kw + ), + {"q": q, "win": win, "pool": pool}, + [(n, (f"sdpa_sparse_prefill {n}",), 8, False) for n in ("q", "win", "pool")], + ) + q, local_kv, pooled, topk, sinks, scale = dsa_build( + 1, 128, 640, 512, 2560, 4, 128, dt, 5, "random" + ) + mx.eval(q, local_kv, pooled, topk, sinks) + bad += _operand_failures( + "dsa_sparse_attention", + lambda q, local_kv, pooled: kq.dsa_sparse_attention( + q, local_kv, pooled, topk, sinks, scale, 2560, 4, 128 + ), + {"q": q, "local_kv": local_kv, "pooled": pooled}, + [ + (n, (f"dsa_sparse_attention {n}",), 8, False) + for n in ("q", "local_kv", "pooled") + ], + ) + return bad + + +GROUPS = { + "matmul_ops": _matmul_ops, + "bias_skinny": _bias_and_skinny, + "gathers": _gathers, + "moe": _moe, + "gather_mix": _gather_mix, + "norm_latent_indexer": _norm_latent_indexer, + "hc": _hc, + "sdpa": _sdpa, + "sparse": _sparse, +} +# The matmul groups run under each route environment, the fused LoRA kernel +# (KQ_QMM_LORA=2) with the LoRA matmul; the other ops with NAX on and off. +MATMUL_ENVS = ("default", "nax_qmv_off", "nax_off") +GROUP_ENVS = [ + *((g, e) for g in list(GROUPS)[:5] for e in MATMUL_ENVS), + ("matmul_ops", "lora2"), + *((g, e) for g in list(GROUPS)[5:] for e in ("default", "nax_off")), +] + +_CHILD = r""" +import json, sys +sys.path.insert(0, sys.argv[1]) +import test_x_offset_views as t +print("RESULT " + json.dumps(t.GROUPS[sys.argv[2]]())) +""" +_XA = re.compile(r"^\[kq qmm\] (\S+) .* xa=(\d+)$") +_ALIGN = re.compile(r"^\[kq align\] (.+) xin=(\d+)$") +_ARM = re.compile(r"^\[arm\] (.*) \| (\S+) \| (.*)$") + + +def _unreached(arm, need, names, copies, xa): + """True when an arm that needs copies has not logged them.""" + if need in ("none", "any"): + return False + want = [n for n in names.split(",") if n] + if need == "copy": + return any(n not in copies for n in want) if want else not copies + off = {n for n, xin in copies.items() if xin} + if want: + return any(n not in off for n in want) + return not off and not any(xa) + + +def ops_offsets(group, env): + """(bad, unaligned, unreached, copied) of an op group in a child under + one environment, as quantized_matmul_offsets gives them.""" + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update({"KQ_QMM_LOG": "1", **ENVS[env]}) + r = subprocess.run( + [sys.executable, "-c", _CHILD, TESTS_DIR, group], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert r.returncode == 0, r.stderr[-6000:] + line = [ln for ln in r.stdout.splitlines() if ln.startswith("RESULT ")][-1] + bad = json.loads(line[len("RESULT ") :]) + unaligned, unreached, copied = set(), [], [] + arm = None + for ln in [*r.stderr.splitlines(), "[arm] end | any | "]: + m = _ARM.match(ln) + if m: + if arm and _unreached(*arm): + unreached.append(arm[0]) + if arm and arm[1] == "none" and arm[3]: + copied.append((arm[0], sorted(arm[3]))) + arm = (m.group(1), m.group(2), m.group(3), {}, []) + continue + m = _XA.match(ln) + if m: + if _unaligned_x(m.group(1), int(m.group(2))): + unaligned.add((m.group(1), int(m.group(2)))) + if arm: + arm[4].append(int(m.group(2))) + continue + m = _ALIGN.match(ln) + if m and arm: + arm[3][m.group(1)] = max(arm[3].get(m.group(1), 0), int(m.group(2))) + return bad, sorted(unaligned), unreached, copied + + +@metal_only +@pytest.mark.parametrize("group,env", GROUP_ENVS) +def test_ops_x_offsets(group, env): + """Bit-exact output at every offset, every matmul kernel reads x from a + 16-byte boundary, every view off the boundary reaches its op there and + is copied, and every view on it is read in place.""" + bad, unaligned, unreached, copied = ops_offsets(group, env) + assert bad == [], bad + assert unaligned == [], unaligned + assert unreached == [], unreached[:40] + assert copied == [], copied[:40] From 411ba218da887cc3fb375afc755b047089c6af2d Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:41:46 -0700 Subject: [PATCH 10/10] test(verify): the fa_verify default lever and multi per-weight copy tests expect the nax routes on nax gpus only --- tests/test_kvarn_sdpa.py | 15 ++++++++------- tests/test_matmul_multi.py | 8 ++++---- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/test_kvarn_sdpa.py b/tests/test_kvarn_sdpa.py index 51474fe..649d011 100644 --- a/tests/test_kvarn_sdpa.py +++ b/tests/test_kvarn_sdpa.py @@ -861,11 +861,11 @@ def test_fa_verify_explicit_splits(): @pytest.mark.parametrize("d", [128, 256]) def test_fa_verify_default_lever_keeps_kvarn_kernel(d, monkeypatch): # At the default KQ_FA_VERIFY_NAX, plain sdpa_fa_verify takes the NAX - # pass 1 at these head dims, which reads plain K/V. The kvarn op keeps - # the simdgroup kernel and its split rule, so it gives the bits it gives - # with the lever off, and those of the simdgroup reference. The lever - # is read when the kernel is dispatched, so each arm is evaluated - # before the lever changes. + # pass 1 at these head dims on NAX GPUs, which reads plain K/V. The + # kvarn op keeps the simdgroup kernel and its split rule, so it gives + # the bits it gives with the lever off, and those of the simdgroup + # reference. The lever is read when the kernel is dispatched, so each + # arm is evaluated before the lever changes. st = build_state(1, 2, 1000, 6, 6, seed=123 + d, d=d) q = make_q(1, 8, 4, d=d) monkeypatch.delenv("KQ_FA_VERIFY_NAX") @@ -874,8 +874,9 @@ def test_fa_verify_default_lever_keeps_kvarn_kernel(d, monkeypatch): monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") out_off, ref = run_fa_pair(st, q, 1000, 6, 6, 4, return_lse=True) mx.eval(out_off, ref) - # The plain op did take the NAX pass at the default. - assert not np.array_equal(np.array(ref_nax[0]), np.array(ref[0])) + # The plain op takes the NAX pass at the default on NAX GPUs only. + took_nax = not np.array_equal(np.array(ref_nax[0]), np.array(ref[0])) + assert took_nax == kq.nax_available() assert_bit_equal(out_default, out_off) assert_bit_equal(out_default, ref) diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py index 7c92d47..bba316e 100644 --- a/tests/test_matmul_multi.py +++ b/tests/test_matmul_multi.py @@ -414,7 +414,7 @@ def test_x16_load_bits(codec, dtype): assert r["digest"] == group[0]["digest"], tag -@pytest.mark.skipif(force_cpu(), reason="GPU copies") +@NAX_ONLY @pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) @pytest.mark.parametrize( "lever", @@ -422,9 +422,9 @@ def test_x16_load_bits(codec, dtype): ids=lambda e: ",".join(f"{k}={v}" for k, v in e.items()), ) def test_per_weight_copies_x_once(lever, dtype): - # Off verify_nax, the op runs quantized_matmul per weight on one aligned - # copy of x: one copy for the three weights, every kernel reads x on a - # 16-byte boundary, and the bits match x at offset 0. + # Off verify_nax on a NAX GPU, the op runs quantized_matmul per weight + # on one aligned copy of x: one copy for the three weights, every kernel + # reads x on a 16-byte boundary, and the bits match x at offset 0. calls = [ { "codec": "q4_k",