From 4241aefa48173f48565cb7757e0d6b71e6c26b3f Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:59:10 -0700 Subject: [PATCH 01/10] feat(sdpa): per-row key start offsets for batched left-padded decode --- CHANGELOG.md | 3 + bindings.cpp | 8 +- metal/mlx/backend/metal/kernels/kq_sdpa.h | 28 ++++-- src/kquant.h | 34 +++++-- src/kquant_sdpa.cpp | 33 ++++++- tests/test_sdpa.py | 114 ++++++++++++++++++++++ 6 files changed, 200 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a70634c..9c0f9be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - `route_shed(indices, scores, slot_table)`: GPU-side routed-expert slot remap plus residency shed for streamed MoE decode; non-resident experts are shed and reported (miss ids and scores) without a host sync. +- `sdpa_decode_gqa` optional `starts` (int32 [B]): per-batch-row key start + offsets for left-padded batched KV caches; padded-out key chunks are + skipped, not staged. ## [0.3.7] diff --git a/bindings.cpp b/bindings.cpp index 633a4e3..c2f8132 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -207,6 +207,7 @@ NB_MODULE(_ext, m) { "sinks"_a = nb::none(), "splits"_a = 0, "tile_c"_a = 0, + "starts"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), R"( @@ -215,7 +216,9 @@ NB_MODULE(_ext, m) { chunk is streamed through threadgroup-staged K/V tiles shared by the whole GQA group, so device memory reads the KV once per kv-head. At qL 2..4 (speculative-verify width) every query also shares the staged - tiles, causally clamped to its own trailing position. + tiles, causally clamped to its own trailing position. With `starts`, + batch row b attends keys [starts[b], kL) -- a left-padded batched KV + cache -- and fully padded-out key chunks are skipped, not staged. Args: q (array): queries [B, n_q_heads, qL, D], float16/bfloat16; @@ -229,6 +232,9 @@ NB_MODULE(_ext, m) { splits (int): key-axis split count; 0 picks the default. tile_c (int): staged tile height, 8/16/32; 0 (default) picks by head_dim (32 up to D=128, 16 at D=256, 8 at D=512). + starts (array, optional): per-batch-row key start offsets, + int32 [B], each in [0, kL - qL]; row b attends [starts[b], + kL). Out-of-range values read as an empty row (zero output). Returns: array: attention output [B, n_q_heads, qL, D]. diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h index 6f215ce..3fe5ace 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h @@ -15,6 +15,7 @@ constant bool do_causal [[function_constant(0)]]; constant int blocks [[function_constant(1)]]; constant int gqa_splits [[function_constant(2)]]; constant bool gqa_has_sinks [[function_constant(3)]]; +constant bool gqa_has_starts [[function_constant(4)]]; template [[kernel]] void kq_sdpa_vector_2pass_1( @@ -211,6 +212,7 @@ template const constant size_t& v_seq_stride [[buffer(10)]], const constant float& scale [[buffer(11)]], const constant int& q_len [[buffer(12)]], + const device int* starts [[buffer(13)]], uint3 tptg [[threads_per_threadgroup]], uint3 tidtg [[thread_position_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -242,6 +244,18 @@ template const int k0 = split_idx * chunk; const int k1 = min(k0 + chunk, N); + // Per-row key start (left-padded batched KV cache): row batch_idx attends + // keys [row_start, N). Whole tiles below the start are skipped outright -- + // the pad region's bytes are never staged -- and a partially padded tile + // masks per key below. A chunk entirely below the start writes an empty + // partial (max finite_min, sum 0), which pass 2 merges at zero weight. + int row_start = 0; + int kt0 = k0; + if (gqa_has_starts) { + row_start = max(0, starts[batch_idx]); + kt0 = max(k0, (row_start / C) * C); + } + const device T* kbase = keys + (size_t)(batch_idx * num_kv_heads + kv_head_idx) * k_head_stride; const device T* vbase = @@ -277,7 +291,7 @@ template const int flat = (tidtg.z * gqa_factor + tidtg.y) * 32 + lane; const int n_threads = 32 * gqa_factor * tptg.z; - for (int kt = k0; kt < k1; kt += C) { + for (int kt = kt0; kt < k1; kt += C) { threadgroup_barrier(mem_flags::mem_threadgroup); // Cooperative tile load; zero-fill the tail so stale threadgroup data // can never reach the accumulators. @@ -319,17 +333,19 @@ template s[p] += simd_shuffle_down(s[p], off); } s[p] = simd_shuffle(s[p], NL * ty); - const bool valid = kg < k1 && kg <= lim[p]; + const bool valid = + kg < k1 && kg <= lim[p] && (!gqa_has_starts || kg >= row_start); mqk[p][cc] = valid ? s[p] : Limits::finite_min; m_tile[p] = max(m_tile[p], mqk[p][cc]); } } // Online softmax per query; each lane sums its ty-group's keys, so the - // simd_sum counts every key NL times. A tile entirely beyond a query's - // causal limit is skipped outright: with the running max still - // finite_min, exp(finite_min - finite_min) == 1 would poison the sum - // (can only happen at verify width; a decode query attends every key). + // simd_sum counts every key NL times. A tile with no valid key is + // skipped outright: with the running max still finite_min, + // exp(finite_min - finite_min) == 1 would poison the sum (happens past + // a query's causal limit at verify width, or below row_start on a + // left-padded row). float vs[QPS][C / NE]; for (short p = 0; p < QPS; p++) { m_tile[p] = simd_max(m_tile[p]); diff --git a/src/kquant.h b/src/kquant.h index cea057e..33f8793 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -153,14 +153,18 @@ mx::array sdpa_vector( bool causal = true, mx::StreamOrDevice s = {}); -// Decode-time (qL == 1) GQA attention tuned for long KV: fixed coarse +// Decode/verify (qL 1..4) GQA attention tuned for long KV: fixed coarse // contiguous key splits plus threadgroup-staged K/V tiles shared across the // GQA group, so device memory reads the KV once per kv-head. Optional // per-q-head attention sinks (an extra softmax logit with no value row). -// q [B, n_q_heads, 1, D], k/v [B, n_kv_heads, kL, D] with contiguous D; -// head/seq strides are read in place. head_dim 64 only; gqa_factor 2..8. -// `splits` 0 picks the default (32); `tile_c` is the staged tile height -// (32 or 16). Metal-only. +// Optional per-batch-row key start offsets `starts` (int32 [B]) restrict row +// b's attention to keys [starts[b], kL) -- a left-padded batched KV cache -- +// with fully padded-out key chunks skipped, not just masked. Values must be +// in [0, kL - qL]; out-of-range starts read as empty rows (zero output). +// q [B, n_q_heads, qL, D], k/v [B, n_kv_heads, kL, D] with contiguous D; +// head/seq strides are read in place. head_dim 64/128/256/512; gqa_factor +// <= 16. `splits` 0 picks the default; `tile_c` is the staged tile height +// (0 picks by head_dim). Metal-only. mx::array sdpa_decode_gqa( mx::array q, mx::array k, @@ -169,6 +173,7 @@ mx::array sdpa_decode_gqa( const std::optional& sinks = std::nullopt, int splits = 0, int tile_c = 32, + const std::optional& starts = std::nullopt, mx::StreamOrDevice s = {}); // Speculative-verify attention on the GPU matrix units for a GQA-folded query @@ -684,15 +689,24 @@ class KQuantSDPA : public mx::Primitive { bool causal_; }; -// Decode-time GQA attention (see sdpa_decode_gqa). Sinks presence is encoded -// in the input count (q, k, v[, sinks]). Inference-only. +// Decode-time GQA attention (see sdpa_decode_gqa). Optional inputs follow +// q, k, v in a fixed order (sinks, then starts); presence flags are carried +// here, not inferred from the input count. Inference-only. class KQuantSDPAGQA : public mx::Primitive { public: - explicit KQuantSDPAGQA(mx::Stream stream, float scale, int splits, int tile_c) + explicit KQuantSDPAGQA( + mx::Stream stream, + float scale, + int splits, + int tile_c, + bool has_sinks, + bool has_starts) : mx::Primitive(stream), scale_(scale), splits_(splits), - tile_c_(tile_c) {} + tile_c_(tile_c), + has_sinks_(has_sinks), + has_starts_(has_starts) {} void eval_cpu( const std::vector& inputs, @@ -713,6 +727,8 @@ class KQuantSDPAGQA : public mx::Primitive { float scale_; int splits_; int tile_c_; + bool has_sinks_; + bool has_starts_; }; // Simdgroup-matrix FA verify attention (see sdpa_fa_verify). Inference-only. diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 047bbb7..81a39d7 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -206,7 +206,9 @@ void KQuantSDPAGQA::eval_gpu( const auto& q = inputs[0]; const auto& k = inputs[1]; const auto& v = inputs[2]; - const bool sinks = inputs.size() == 4; + const bool sinks = has_sinks_; + const bool starts = has_starts_; + const size_t starts_idx = 3 + (sinks ? 1 : 0); kq_sdpa_check_layout("sdpa_decode_gqa", q, k, v); int B = q.shape(0); @@ -249,9 +251,11 @@ void KQuantSDPAGQA::eval_gpu( std::string ts = kq_type_string(q.dtype()); bool has_sinks = sinks; + bool has_starts = starts; mx::metal::MTLFCList fc = { {&splits, MTL::DataType::DataTypeInt, 2}, {&has_sinks, MTL::DataType::DataTypeBool, 3}, + {&has_starts, MTL::DataType::DataTypeBool, 4}, }; // Pass 1: one threadgroup per (kv-head, batch, split); the whole GQA group @@ -261,7 +265,8 @@ void KQuantSDPAGQA::eval_gpu( { std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) + "_c" + std::to_string(tile_c_) + (qL > 1 ? "_p2" : ""); - std::string hash = kname + "_s" + std::to_string(splits); + std::string hash = + kname + "_s" + std::to_string(splits) + (has_starts ? "_st1" : "_st0"); auto kernel = kq_get_kernel(d, kname, hash, fc); // Register-heavy pipeline: some GPUs cap it below the dispatch width, and // Metal turns an oversized dispatch into silent garbage, not an error. @@ -287,6 +292,9 @@ void KQuantSDPAGQA::eval_gpu( ce.set_bytes(v_seq_stride, 10); ce.set_bytes(scale, 11); ce.set_bytes(qL, 12); + // Metal wants every buffer bound; without starts, rebind sums as a dummy + // (the read is compiled out via the function constant). + ce.set_input_array(starts ? inputs[starts_idx] : sums, 13); MTL::Size group_dims(32, gqa_factor, qL > 1 ? (qL + 1) / 2 : 1); MTL::Size grid_dims(n_kv_heads, B, splits); ce.dispatch_threadgroups(grid_dims, group_dims); @@ -555,7 +563,8 @@ std::vector KQuantSDPAGQA::output_shapes( bool KQuantSDPAGQA::is_equivalent(const mx::Primitive& other) const { const auto& o = static_cast(other); - return scale_ == o.scale_ && splits_ == o.splits_ && tile_c_ == o.tile_c_; + return scale_ == o.scale_ && splits_ == o.splits_ && tile_c_ == o.tile_c_ && + has_sinks_ == o.has_sinks_ && has_starts_ == o.has_starts_; } mx::array sdpa_decode_gqa( @@ -566,6 +575,7 @@ mx::array sdpa_decode_gqa( const std::optional& sinks, int splits, int tile_c, + const std::optional& starts, mx::StreamOrDevice s_) { auto s = mx::to_stream(s_); @@ -653,12 +663,27 @@ mx::array sdpa_decode_gqa( sk = mx::astype(mx::reshape(sk, {n_q_heads}, s), mx::float32, s); inputs.push_back(mx::contiguous(sk, false, s)); } + if (starts.has_value()) { + auto st = *starts; + if (st.size() != static_cast(q.shape(0))) { + throw std::invalid_argument( + "[mlx_kquant.sdpa_decode_gqa] starts must have one element per " + "batch row."); + } + if (st.dtype() != mx::int32) { + throw std::invalid_argument( + "[mlx_kquant.sdpa_decode_gqa] starts must be int32."); + } + st = mx::reshape(st, {q.shape(0)}, s); + inputs.push_back(mx::contiguous(st, false, s)); + } auto out_shape = q.shape(); return mx::array( std::move(out_shape), dt, - std::make_shared(s, scale, splits, tile_c), + std::make_shared( + s, scale, splits, tile_c, sinks.has_value(), starts.has_value()), std::move(inputs)); } diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 79f4b98..8c54a1a 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -242,6 +242,120 @@ def test_sdpa_gqa_verify_short_kv(): _check_gqa(64, kL=17, dtype=mx.bfloat16, splits=16, qL=4) +def _ref_sdpa_starts(q, k, v, scale, pads, qL): + # per-row f32 reference on the visible tail [pads[b], kL) + outs = [] + for b in range(q.shape[0]): + p = int(pads[b]) + outs.append( + _ref_sdpa( + q[b : b + 1], + k[b : b + 1, :, p:, :], + v[b : b + 1, :, p:, :], + scale, + causal=qL > 1, + ) + ) + return mx.concatenate(outs, axis=0) + + +def _check_gqa_starts( + D, + kL, + dtype, + B=4, + Hq=24, + Hkv=4, + qL=1, + pads=None, + strided=False, + splits=0, +): + scale = 1.0 / (D**0.5) + q, k, v = _make(B, Hq, Hkv, qL, kL, D, dtype, seed=kL + D + B, strided=strided) + if pads is None: + pads = [(b * (kL - qL)) // B for b in range(B)] + starts = mx.array(pads, dtype=mx.int32) + mx.eval(starts) + got = kq.sdpa_decode_gqa(q, k, v, scale, splits=splits, starts=starts) + ref = _ref_sdpa_starts(q, k, v, scale, pads, qL) + _eval_or_skip(got, ref) + rel = _rel(got, ref) + bound = REL_BOUND[dtype] + print( + f" [gqa-starts] D={D} qL={qL} kL={kL} B={B} Hq/Hkv={Hq}/{Hkv} " + f"pads={pads} {str(dtype)[9:]:>9}: rel={rel:.3e}" + ) + assert rel < bound, f"D={D} kL={kL} pads={pads} rel {rel:.3e} >= {bound:.0e}" + assert got.shape == q.shape + + +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) +def test_sdpa_decode_gqa_starts(D, dtype): + # left-padded batched rows: row b attends keys [pads[b], kL) + _check_gqa_starts(D, kL=4096, dtype=dtype) + + +@pytest.mark.parametrize("D", [64, 512]) +def test_sdpa_decode_gqa_batched_nostarts(D): + # B > 1 without starts: the plain batched grid the ragged route + # degenerates to at pad-0 (and a pin for batched use on its own) + _check_gqa(D, kL=4096, dtype=mx.bfloat16) + scale = 1.0 / (D**0.5) + q, k, v = _make(8, 24, 4, 1, 4096, D, mx.bfloat16, seed=D, strided=False) + got = kq.sdpa_decode_gqa(q, k, v, scale) + ref = _ref_sdpa_sinks(q, k, v, scale, None) + _eval_or_skip(got, ref) + assert _rel(got, ref) < REL_BOUND[mx.bfloat16] + + +def test_sdpa_decode_gqa_starts_zero_matches_plain(): + # all-zero starts must match the no-starts call on the same inputs + scale = 1.0 / (512**0.5) + q, k, v = _make(4, 24, 4, 1, 2048, 512, mx.bfloat16, seed=3, strided=False) + starts = mx.zeros((4,), dtype=mx.int32) + a = kq.sdpa_decode_gqa(q, k, v, scale, starts=starts) + b = kq.sdpa_decode_gqa(q, k, v, scale) + _eval_or_skip(a, b) + assert _rel(a, b) < 1e-6 + + +@pytest.mark.parametrize("qL", [2, 4]) +def test_sdpa_decode_gqa_starts_verify(qL): + # verify width on left-padded rows: the block occupies the last qL + # positions of every row regardless of its pad (end-aligned causal) + _check_gqa_starts(512, kL=4096, dtype=mx.bfloat16, qL=qL) + + +def test_sdpa_decode_gqa_starts_edges(): + # pad 0, a pad on a tile boundary, a pad mid-tile, and the maximum + # in-contract pad (one visible key at qL=1: output equals that value row) + _check_gqa_starts( + 512, + kL=3071, + dtype=mx.bfloat16, + B=4, + pads=[0, 1024, 1543, 3070], + strided=True, + splits=16, + ) + + +@pytest.mark.parametrize("Hq,Hkv", [(32, 4), (16, 1)]) +def test_sdpa_decode_gqa_starts_gemma_geometry(Hq, Hkv): + # gemma-4 31b (gqa 8) and 12b (gqa 16, single kv head) global layers + _check_gqa_starts(512, kL=8192, dtype=mx.bfloat16, B=8, Hq=Hq, Hkv=Hkv) + + +def test_sdpa_decode_gqa_starts_validation(): + q, k, v = _make(4, 24, 4, 1, 512, 64, mx.bfloat16, seed=5, strided=False) + with pytest.raises(ValueError, match="one element per batch row"): + kq.sdpa_decode_gqa(q, k, v, 0.125, starts=mx.zeros((3,), mx.int32)) + with pytest.raises(ValueError, match="int32"): + kq.sdpa_decode_gqa(q, k, v, 0.125, starts=mx.zeros((4,), mx.int64)) + + def _ref_sdpa_fold(q, k, v, scale, q_len): """f32 reference for the GQA-folded verify layout: q [B, Hkv, G*qL, D] attends its own kv head directly; folded row r is causally clamped to From e1ac45cf08f85ea9afd4efc981dd91ff8d3e22d6 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:25:06 -0700 Subject: [PATCH 02/10] feat(matmul): env-gated split-k qmm_t dispatch + partial-fold accum kernel --- metal/kq_quantized.metal | 35 ++++++++++- src/kquant_matmul.cpp | 133 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 159 insertions(+), 9 deletions(-) diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index 919543f..641c796 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -863,4 +863,37 @@ instantiate_kquant_gather_qmm_rhs_codec(256, 1, iq1_m) map[3u * slot + 1u] = r; map[3u * slot + 2u] = min(64u, seg_len - rank); } -// clang-format on + +// Split-K partial fold for qmm_t_splitk: y[i] = (T)(f32 sum over the +// splits axis of partials[z * stride + i]). Partials carry one T rounding +// per slice; the fold accumulates in f32 and rounds once more. One thread +// per output element; cost is noise next to the GEMM pass. +template +[[kernel]] void kq_qmm_splitk_accum( + const device T* partials [[buffer(0)]], + device T* y [[buffer(1)]], + const constant int& n_elems [[buffer(2)]], + const constant int& splits [[buffer(3)]], + const constant int& stride [[buffer(4)]], + uint gid [[thread_position_in_grid]]) { + if (gid >= static_cast(n_elems)) { + return; + } + float acc = 0.0f; + for (int z = 0; z < splits; ++z) { + acc += static_cast( + partials[static_cast(z) * stride + gid]); + } + y[gid] = static_cast(acc); +} + +#define instantiate_kq_qmm_splitk_accum(type) \ + instantiate_kernel( \ + "kquant_qmm_splitk_accum_" #type, \ + kq_qmm_splitk_accum, \ + type) + +instantiate_kq_qmm_splitk_accum(float) +instantiate_kq_qmm_splitk_accum(bfloat16_t) +instantiate_kq_qmm_splitk_accum(float16_t) + // clang-format on diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index df152d1..efb69d8 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -2,12 +2,13 @@ // kernels (qmm / qmm_nax / qvm / qmv) from the bundled metallib via // d.get_kernel(name, lib); the op guarantees row-contiguity before dispatch and // kernel-name type tokens come from kq_type_string. NAX (tensor-core) -// availability is probed via kq_is_nax_available. The split-k paths -// (qmm_splitk / qvm_split_k) are omitted - plain qmm/qvm produce identical -// results with less parallelism. KQuantMatmul itself never carries a bias (a -// separate elementwise add is fine off the 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. +// availability is probed via kq_is_nax_available. qmm_splitk (env-gated, +// KQ_QMM_SPLITK) partitions K for the small-M band; qvm_split_k stays +// omitted - plain qvm is identical with less parallelism. KQuantMatmul itself +// never carries a bias (a separate elementwise add is fine off the +// 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 @@ -329,8 +330,8 @@ void qmm_nax( ce.dispatch_threadgroups(grid_dims, group_dims); } -// Tiled quantized GEMM (no biases). The split-k variant is omitted; plain qmm -// is correct with less parallelism. +// Tiled quantized GEMM (no biases). The split-k variant (qmm_splitk below) +// covers the small-M occupancy hole; plain qmm is the general path. void qmm( const array& x, const array& w, @@ -405,6 +406,86 @@ void qmm( ce.dispatch_threadgroups(grid_dims, group_dims); } +// Split-K qmm_t for the small-M decode band. The plain tile grid is only +// ceil(N/64) x 1 threadgroups at decode shapes with M <= 32 (84 at +// [5376 x 21504]) and the in-tile K walk serializes, capping qmm/NAX at +// 160-257 GB/s while mv_ext decays past M~4 on L2 activation re-reads. +// Partitioning K across grid.z multiplies threadgroup count by `splits`; +// each slice writes a T partial tile and a second tiny pass folds them in +// f32. Slice starts must land on wire-block boundaries, so the partition +// is a multiple of group_size (the caller guarantees splits divides +// K / group_size). Non-batched transpose shapes only. +void qmm_splitk( + const array& x, + const array& w, + const array& scales, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + int splits, + Device& d, + const Stream& s, + const std::string& kquant_type) { + constexpr int bm = 32, bn = 32; + constexpr int wm = 2, wn = 2; + const int k_partition = (K / group_size / splits) * group_size; + const int part_stride = M * N; + + array partials({splits, M, N}, x.dtype(), nullptr, {}); + partials.set_data(mx::allocator::malloc(partials.nbytes())); + + auto& ce = mx::metal::get_command_encoder(s); + ce.add_temporary(partials); + + std::string type_string = kq_type_string(x.dtype()); + bool aligned = N % bn == 0; + std::string kname; + kname.reserve(64); + mx::concatenate( + kname, + kq_kname_prefix(kquant_type) + "qmm_t_splitk_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + aligned ? "_alN_true" : "_alN_false"); + + auto kernel = kq_get_kernel(d, kname); + ce.set_compute_pipeline_state(kernel); + + int c = 0; + ce.set_input_array(w, c++); + ce.set_input_array(scales, c++); + ce.set_input_array(x, c++); + ce.set_output_array(partials, c++); + ce.set_bytes(K, c++); + ce.set_bytes(N, c++); + ce.set_bytes(M, c++); + ce.set_bytes(k_partition, c++); + 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); + ce.dispatch_threadgroups(grid_dims, group_dims); + + std::string aname = "kquant_qmm_splitk_accum_" + type_string; + auto accum = kq_get_kernel(d, aname); + ce.set_compute_pipeline_state(accum); + const int n_elems = M * N; + c = 0; + ce.set_input_array(partials, c++); + ce.set_output_array(out, c++); + ce.set_bytes(n_elems, c++); + ce.set_bytes(splits, c++); + ce.set_bytes(part_stride, c++); + MTL::Size agrid(static_cast(n_elems), 1, 1); + MTL::Size agroup(256, 1, 1); + ce.dispatch_threads(agrid, agroup); +} + // Vector-times-matrix quantized kernel dispatch (no biases). void qvm( const array& x, @@ -934,6 +1015,42 @@ void KQuantMatmul::eval_gpu( return; } + // Split-K qmm experiment (KQ_QMM_SPLITK=, 0 = off, read + // once): the small-M band's occupancy lever; see qmm_splitk. Routes when + // a >1 divisor of the wire-block count exists at or under the target. + // K-quants + q8_0 only for now (instantiation coverage). + static const int qmm_splitk_env = []() { + const char* e = std::getenv("KQ_QMM_SPLITK"); + return e != nullptr ? std::atoi(e) : 0; + }(); + if (qmm_splitk_env > 1 && transpose_ && non_batched && M <= 32 && + (kquant_type_ == "q6_k" || kquant_type_ == "q5_k" || + kquant_type_ == "q4_k" || kquant_type_ == "q3_k" || + kquant_type_ == "q2_k" || kquant_type_ == "q8_0")) { + const int nblk = K / group_size_; + int sp = std::min(qmm_splitk_env, nblk); + while (sp > 1 && nblk % sp != 0) { + --sp; + } + if (sp > 1) { + qmm_splitk( + x, + w, + scales, + out, + group_size_, + bits_, + M, + N, + K, + sp, + d, + s, + kquant_type_); + return; + } + } + // Small-M qmm route: the double-buffered BM=32 NAX tile beats the mv // paths' wide-M decay above a per-codec crossover (kq_smallbm_policy; // measured cold-stream, e.g. q6_k M>=9 274-305 vs 221-254, q4_k M>=7 From 52de6033023d41bd10d05a58e1e8c706fd5dd324 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:35:09 -0700 Subject: [PATCH 03/10] feat(matmul): mv_ext sb + x16/x32 experiment kernels (env-gated, measured losers on m5) --- metal/kq_quantized.metal | 53 +++++++++++ .../mlx/backend/metal/kernels/kq_quantized.h | 93 +++++++++++++++++++ .../metal/kernels/kq_quantized_kquants.h | 16 ++++ src/kquant_matmul.cpp | 25 ++++- 4 files changed, 185 insertions(+), 2 deletions(-) diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index 641c796..6c7fc5a 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -397,6 +397,59 @@ instantiate_mv_ext_all(iq1_m, 256, 1) instantiate_mv_ext_nr2_for_type(codec, gs, bits, float16_t) instantiate_mv_ext_nr2_all(q6_k, 256, 6) +// Shuffle-broadcast experiment (KQ_MV_EXT_SB=1): the four ty-lanes sharing +// a tx column exchange float4 quarters of the activation window over +// simd_shuffle instead of each loading all 16 elements -- activation cache +// traffic / 4, no barriers, no extra accumulators. q6_k only, M 4-12 where +// the nr0=1 activation decay bites. Suffix _sb. +#define instantiate_mv_ext_sb(codec, type, gs, bits, m) \ + instantiate_kernel( \ + "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m \ + "_sb", \ + kq_ ## codec ## _mv_ext_sb, type, m, 2, 8) +#define instantiate_mv_ext_sb_for_type(codec, gs, bits, type) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 4) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 5) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 6) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 7) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 8) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 9) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 10) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 11) \ + instantiate_mv_ext_sb(codec, type, gs, bits, 12) +#define instantiate_mv_ext_sb_all(codec, gs, bits) \ + instantiate_mv_ext_sb_for_type(codec, gs, bits, float) \ + instantiate_mv_ext_sb_for_type(codec, gs, bits, bfloat16_t) \ + instantiate_mv_ext_sb_for_type(codec, gs, bits, float16_t) +instantiate_mv_ext_sb_all(q6_k, 256, 6) + +// Wide-nxpsg experiment (KQ_MV_EXT_NX=16|32): more K lanes per simdgroup +// means fewer output rows per simdgroup, so each activation element has +// nypsg*nsg = 4 (x16) or 2 (x32) redundant readers instead of 8 -- the +// L1-capacity pressure behind the M4->8 decay -- and the grid gains +// threadgroups. Same impl, different template params. q6_k M 4-12. +#define instantiate_mv_ext_nx(codec, type, gs, bits, m, nx) \ + instantiate_kernel( \ + "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m \ + "_x" #nx, \ + kq_ ## codec ## _mv_ext, type, m, 2, nx) +#define instantiate_mv_ext_nx_for_type(codec, gs, bits, type, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 4, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 5, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 6, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 7, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 8, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 9, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 10, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 11, nx) \ + instantiate_mv_ext_nx(codec, type, gs, bits, 12, nx) +#define instantiate_mv_ext_nx_all(codec, gs, bits, nx) \ + instantiate_mv_ext_nx_for_type(codec, gs, bits, float, nx) \ + instantiate_mv_ext_nx_for_type(codec, gs, bits, bfloat16_t, nx) \ + instantiate_mv_ext_nx_for_type(codec, gs, bits, float16_t, nx) +instantiate_mv_ext_nx_all(q6_k, 256, 6, 16) +instantiate_mv_ext_nx_all(q6_k, 256, 6, 32) + #define instantiate_kquant_q3_k_for_type(type) \ instantiate_kquant_batched(verify_qmv, type, 256, 3, 0, q3_k) \ instantiate_kquant_batched(qmv_fast, type, 256, 3, 0, q3_k) \ diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index 7b69309..9a468d9 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -791,6 +791,99 @@ METAL_FUNC void kq_mv_ext_nr_impl( } } +// Shuffle-broadcast variant of kq_mv_ext_impl. The chunk index depends only +// on tx, so the nypsg ty-lanes sharing a tx column process identical +// activation windows and the nr0=1 kernel reads every activation element +// nypsg times per simdgroup (the M * N * K traffic that dominates past +// M~5). Here each ty-lane loads ONE float4 quarter of the 16-element window +// and the four lanes exchange quarters over simd_shuffle: activation cache +// traffic / 4 with no barriers, no threadgroup memory, and no extra +// accumulators (the register cost that sank the nr0=2 variant). Divergence +// safety: the K loop exits on a tx-uniform condition, so the four shuffle +// partners (same tx, ty 0-3) are always uniformly active, including OOB +// output rows, which run the loop on a clamped w row and mask the store. +// Requires nxpsg == 8 so ty spans exactly the four quarters. +template +METAL_FUNC void kq_mv_ext_sb_impl( + const device uint8_t* w, + const device T* x, + device T* y, + const constant int& in_vec_size, // K + const constant int& out_vec_size, // N + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + static_assert(nxpsg == 8, "shuffle-broadcast needs ty to span 4 quarters"); + constexpr short nypsg = 32 / nxpsg; // output rows per simdgroup + constexpr short chpb = Codec::superblock / 16; // 16-weight chunks per block + const short tx = tiisg % nxpsg; // K position within the row group + const short ty = tiisg / nxpsg; // which of nypsg rows this thread owns + + const int i01 = tgpig.x * (nypsg * nsg) + nypsg * sgitg + ty; // output row + const int i11 = tgpig.y * r1ptg; // first activation column + + const int nb = in_vec_size / Codec::superblock; + const int row_bytes = nb * Codec::block_bytes; + // Clamp OOB rows to row 0 for a valid read; the store is masked below. + const device uint8_t* w_row = + (i01 < out_vec_size) ? w + static_cast(i01) * row_bytes : w; + + const device T* y_col[r1ptg]; +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y_col[ir1] = x + static_cast(i11 + ir1) * in_vec_size + tx * 16; + } + + float sumf[r1ptg]; +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + sumf[ir1] = 0.0f; + } + + const ushort lane0 = tx; + for (int ich = tx; 16 * ich < in_vec_size; ich += nxpsg) { + const int ib = ich / chpb; // super-block index + const short cch = ich % chpb; // chunk within super-block + const device uint8_t* block = + w_row + static_cast(ib) * Codec::block_bytes; + float4x4 lx; + Codec::deq_chunk16(block, cch, lx); +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + const device T* yp = y_col[ir1]; + const float4 mine = float4(*(const device vec*)(yp + 4 * ty)); + const float4 a0 = simd_shuffle(mine, lane0); + const float4 a1 = simd_shuffle(mine, static_cast(lane0 + 8)); + const float4 a2 = simd_shuffle(mine, static_cast(lane0 + 16)); + const float4 a3 = simd_shuffle(mine, static_cast(lane0 + 24)); + sumf[ir1] += + dot(lx[0], a0) + dot(lx[1], a1) + dot(lx[2], a2) + dot(lx[3], a3); + y_col[ir1] += nxpsg * 16; + } + } + +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + if (nxpsg >= 8) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 4); + } + if (nxpsg >= 4) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 2); + } + if (nxpsg >= 2) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 1); + } + } + + if (tx == 0 && i01 < out_vec_size) { +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y[static_cast(i11 + ir1) * out_vec_size + i01] = + static_cast(sumf[ir1]); + } + } +} + // Q8_0: 34 bytes/32 weights. [fp16 d][int8 q[32]]. w[i] = d * q[i]. MLX_MTL_CONST int KQ_Q8_0_GROUP = 32; diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h index 95950c3..01eac0f 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h @@ -2882,6 +2882,22 @@ template w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); } +template +[[kernel]] void kq_q6_k_mv_ext_sb( + const device uint8_t* w, + const device uint8_t* /* scales */, + const device T* x, + device T* y, + const constant int& in_vec_size, // K + const constant int& out_vec_size, // N + const constant int& /* vm */, // == r1ptg + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + kq_mv_ext_sb_impl( + w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); +} + template [[kernel]] void kq_q6_k_qmv( const device uint8_t* w, diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index efb69d8..88daa15 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -777,7 +777,26 @@ void verify_mv_ext( return e != nullptr ? std::atoi(e) : 1; }(); const bool use_nr2 = mv_ext_nr == 2 && M >= 5 && kquant_type == "q6_k"; - const int rows_per_tg = (32 / nxpsg) * nsg * (use_nr2 ? 2 : 1); + // Shuffle-broadcast experiment (KQ_MV_EXT_SB=1): ty-lanes exchange + // activation quarters over simd_shuffle instead of each loading the full + // window -- activation cache traffic / 4, same grid. q6_k M 4-12 only. + static const bool mv_ext_sb = []() { + const char* e = std::getenv("KQ_MV_EXT_SB"); + return e != nullptr && std::atoi(e) == 1; + }(); + const bool use_sb = !use_nr2 && mv_ext_sb && M >= 4 && kquant_type == "q6_k"; + // Wide-nxpsg experiment (KQ_MV_EXT_NX=16|32): fewer redundant activation + // readers per element (nypsg*nsg drops 8 -> 4 -> 2) + more threadgroups. + // q6_k M 4-12 only; wins here would generalize per codec. + static const int mv_ext_nx = []() { + const char* e = std::getenv("KQ_MV_EXT_NX"); + const int v = e != nullptr ? std::atoi(e) : 0; + return (v == 16 || v == 32) ? v : 0; + }(); + const bool use_nx = + !use_nr2 && !use_sb && mv_ext_nx != 0 && M >= 4 && kquant_type == "q6_k"; + const int nxpsg_eff = use_nx ? mv_ext_nx : nxpsg; + const int rows_per_tg = (32 / nxpsg_eff) * nsg * (use_nr2 ? 2 : 1); MTL::Size group_dims(32, nsg, 1); MTL::Size grid_dims((N + rows_per_tg - 1) / rows_per_tg, 1, 1); @@ -794,7 +813,9 @@ void verify_mv_ext( bits, "_m", M, - use_nr2 ? "_nr2" : ""); + use_nr2 ? "_nr2" + : (use_sb ? "_sb" + : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") : ""))); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); From d8f19c797fb998e00dadd660c3c82ff705a871b3 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:42:49 -0700 Subject: [PATCH 04/10] feat(matmul): env-gated nax split-k qmm_t dispatch (bm32 tile, k-slice partials) --- metal/kq_quantized_nax.metal | 20 +++ .../backend/metal/kernels/kq_quantized_nax.h | 94 +++++++++++++- src/kquant_matmul.cpp | 120 +++++++++++++++++- 3 files changed, 229 insertions(+), 5 deletions(-) diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index 6a9f912..3e3432e 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -116,6 +116,26 @@ instantiate_kquant_nax_qmm_t_smallbm(codec, float16_t, gs, bits) \ instantiate_kquant_nax_qmm_t_smallbm(codec, bfloat16_t, gs, bits) +// Split-K qmm_t on the NAX BM=32 tile (KQ_QMM_SPLITK_NAX experiment): +// grid.z K-slices into T partials + shared accum fold. The plain small-M +// grid is TG-count starved (ceil(N/64) x 1 threadgroups at decode shapes); +// splitting K multiplies occupancy without touching the fragment shape. +// q6_k + q8_0 only; no batched or float x variants (route gates match +// qmm_nax and non_batched). +#define instantiate_kquant_nax_qmm_t_splitk(type, gs, bits, aligned_N, codec) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_splitk_" #type "_gs_" #gs "_b_" #bits \ + "_bm32_bn64_bk64_wm2_wn2_alN_" #aligned_N, \ + kq_ ## codec ## _qmm_t_nax_splitk, \ + type, gs, bits, aligned_N, 32, 64, 2, 2) +#define instantiate_kquant_nax_splitk(codec, gs, bits) \ + instantiate_kquant_nax_qmm_t_splitk(float16_t, gs, bits, true, codec) \ + instantiate_kquant_nax_qmm_t_splitk(float16_t, gs, bits, false, codec) \ + instantiate_kquant_nax_qmm_t_splitk(bfloat16_t, gs, bits, true, codec) \ + instantiate_kquant_nax_qmm_t_splitk(bfloat16_t, gs, bits, false, codec) +instantiate_kquant_nax_splitk(q6_k, 256, 6) +instantiate_kquant_nax_splitk(q8_0, 32, 8) + // Double-buffered BM=64 qmm_t, name-suffixed _db: dispatched by the host // solely for the M33-64 decode band (kq_smallbm_policy db64 + KQ_NAX_DB64). // As a blanket BM=64 default the doubled Ws cut occupancy (M96+ -3-15%, diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index d47c565..446bc90 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h @@ -33,7 +33,8 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]]) { + uint simd_lid [[thread_index_in_simdgroup]], + const int k_len = -1) { static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -41,6 +42,10 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( constexpr int BK_padded = (BK + 16 / sizeof(T)); + // k_len < 0 walks the full reduction; the split-k wrapper passes a slice + // length while K keeps supplying the (full) row strides. + const int k_bound = k_len < 0 ? K : k_len; + const int K_w = (K / LoaderW::weights_per_block) * LoaderW::bytes_per_block; const int y_row = tid.y * BM; const int y_col = tid.x * BN; @@ -109,7 +114,7 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( constexpr int WS_STRIDE = BN * BK_padded; LoaderW loader_w1(wl, K, Ws + WS_STRIDE, simd_gid, simd_lid); loader_w1.next(); - const int n_steps = K / BK; + const int n_steps = k_bound / BK; if constexpr (kAlignedN.value) { loader_w.load_unsafe(); } else { @@ -172,7 +177,7 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( threadgroup_barrier(mem_flags::mem_threadgroup); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < k_bound; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); if constexpr (kAlignedN.value) { loader_w.load_unsafe(); @@ -3320,6 +3325,89 @@ KQ_NAX_DEFINE_KERNELS(q6_k, 256, 6, KqNaxQ6_KBlockLoader) KQ_NAX_DEFINE_KERNELS(q3_k, 256, 3, KqNaxQ3_KBlockLoader) KQ_NAX_DEFINE_KERNELS(q2_k, 256, 2, KqNaxQ2_KBlockLoader) +// Split-K qmm_t on the NAX tile (KQ_QMM_SPLITK_NAX, small-M experiment): +// grid.z indexes K-slices; each slice walks k_partition_size weights from a +// superblock-aligned start and stores a T partial tile at +// tid.z * split_k_partition_stride. The shared kquant_qmm_splitk_accum pass +// folds slices in f32. The host guarantees k_partition_size is a multiple +// of both the codec superblock and BK, so every slice starts the loader at +// kt_base 0. Non-batched transpose shapes only; no swizzle (grid.y is a +// single row tile in the target band). +#define KQ_NAX_DEFINE_SPLITK_KERNEL(codec, GROUP_CONST, bits_val, LOADER) \ + template < \ + typename T, \ + int group_size, \ + int bits, \ + bool aligned_N, \ + int BM, \ + int BN, \ + int WM, \ + int WN> \ + [[kernel]] void kq_##codec##_qmm_t_nax_splitk( \ + const device uint8_t* w, \ + const device uint8_t* /* scales */, \ + const device T* x, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + const constant int& k_partition_size, \ + const constant int& split_k_partition_stride, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint lid [[thread_index_in_threadgroup]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + static_assert( \ + group_size == GROUP_CONST, \ + #codec " NAX kernel requires group_size=" #GROUP_CONST); \ + static_assert( \ + bits == bits_val, #codec " NAX kernel requires bits=" #bits_val); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + constexpr int BK = 64; \ + constexpr int BK_padded = (BK + 16 / sizeof(T)); \ + using LoaderW = LOADER< \ + T, \ + BN, \ + BK, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + constexpr int kWsBufs = (LoaderW::db_safe && BM == 32) ? 2 : 1; \ + threadgroup T Ws[kWsBufs * BN * BK_padded]; \ + const int k_start = int(tid.z) * k_partition_size; \ + x += k_start; \ + auto wl = w; \ + wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; \ + y += int(tid.z) * static_cast(split_k_partition_stride); \ + kq_qmm_t_nax_tgp_impl< \ + T, \ + LoaderW, \ + aligned_N, \ + BM, \ + BK, \ + BN, \ + WM, \ + WN, \ + kWsBufs == 2>( \ + wl, \ + x, \ + y, \ + Ws, \ + K, \ + N, \ + M, \ + tid, \ + lid, \ + simd_gid, \ + simd_lid, \ + k_partition_size); \ + } + +KQ_NAX_DEFINE_SPLITK_KERNEL(q6_k, 256, 6, KqNaxQ6_KBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q8_0, 32, 8, KqNaxQ8_0BlockLoader) + template < typename T, typename LoaderW, diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 88daa15..d738bb1 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -2,8 +2,9 @@ // kernels (qmm / qmm_nax / qvm / qmv) from the bundled metallib via // d.get_kernel(name, lib); the op guarantees row-contiguity before dispatch and // kernel-name type tokens come from kq_type_string. NAX (tensor-core) -// availability is probed via kq_is_nax_available. qmm_splitk (env-gated, -// KQ_QMM_SPLITK) partitions K for the small-M band; qvm_split_k stays +// availability is probed via kq_is_nax_available. qmm_splitk / qmm_nax_splitk +// (env-gated, KQ_QMM_SPLITK / KQ_QMM_SPLITK_NAX) partition K for the +// small-M band on the steel and NAX tiles respectively; qvm_split_k stays // omitted - plain qvm is identical with less parallelism. KQuantMatmul itself // never carries a bias (a separate elementwise add is fine off the // decode-latency-critical path); the decode-only bias-fused fast path lives in @@ -486,6 +487,85 @@ void qmm_splitk( ce.dispatch_threads(agrid, agroup); } +// Split-K qmm_t on the NAX BM=32 tile (KQ_QMM_SPLITK_NAX experiment). 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 +// in splits), while the NAX small-M cap is TG-count starvation -- the lever +// splitk actually multiplies. Slice starts must be superblock-aligned so +// every loader instance begins at kt_base 0; the caller guarantees splits +// divides K / max(superblock, BK). Non-batched transpose shapes only. +void qmm_nax_splitk( + const array& x, + const array& w, + const array& scales, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + int splits, + Device& d, + const Stream& s, + const std::string& kquant_type) { + constexpr int bm = 32, bn = 64; + constexpr int wm = 2, wn = 2; + const int k_partition = K / splits; + const int part_stride = M * N; + + array partials({splits, M, N}, x.dtype(), nullptr, {}); + partials.set_data(mx::allocator::malloc(partials.nbytes())); + + auto& ce = mx::metal::get_command_encoder(s); + ce.add_temporary(partials); + + std::string type_string = kq_type_string(x.dtype()); + bool aligned = N % bn == 0; + std::string kname; + kname.reserve(80); + mx::concatenate( + kname, + kq_kname_prefix(kquant_type) + "qmm_t_nax_splitk_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_bm32_bn64_bk64_wm2_wn2", + aligned ? "_alN_true" : "_alN_false"); + + auto kernel = kq_get_kernel(d, kname); + ce.set_compute_pipeline_state(kernel); + + int c = 0; + ce.set_input_array(w, c++); + ce.set_input_array(scales, c++); + ce.set_input_array(x, c++); + ce.set_output_array(partials, c++); + ce.set_bytes(K, c++); + ce.set_bytes(N, c++); + ce.set_bytes(M, c++); + ce.set_bytes(k_partition, c++); + 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); + ce.dispatch_threadgroups(grid_dims, group_dims); + + std::string aname = "kquant_qmm_splitk_accum_" + type_string; + auto accum = kq_get_kernel(d, aname); + ce.set_compute_pipeline_state(accum); + const int n_elems = M * N; + c = 0; + ce.set_input_array(partials, c++); + ce.set_output_array(out, c++); + ce.set_bytes(n_elems, c++); + ce.set_bytes(splits, c++); + ce.set_bytes(part_stride, c++); + MTL::Size agrid(static_cast(n_elems), 1, 1); + MTL::Size agroup(256, 1, 1); + ce.dispatch_threads(agrid, agroup); +} + // Vector-times-matrix quantized kernel dispatch (no biases). void qvm( const array& x, @@ -1036,6 +1116,42 @@ void KQuantMatmul::eval_gpu( return; } + // NAX split-K qmm experiment (KQ_QMM_SPLITK_NAX=, 0 = off, + // read once): K-slices on the tensor-core BM=32 tile; see qmm_nax_splitk. + // Slice quantum is max(superblock, BK) so every slice starts a loader at + // kt_base 0. q6_k + q8_0 instantiations only. + static const int qmm_splitk_nax_env = []() { + const char* e = std::getenv("KQ_QMM_SPLITK_NAX"); + return e != nullptr ? std::atoi(e) : 0; + }(); + if (qmm_splitk_nax_env > 1 && transpose_ && non_batched && M <= 32 && + (kquant_type_ == "q6_k" || kquant_type_ == "q8_0") && + kq_is_nax_available() && (K % 64 == 0) && x.dtype() != mx::float32) { + const int sliceq = std::max(group_size_, 64); + const int nblk = K / sliceq; + int sp = std::min(qmm_splitk_nax_env, nblk); + while (sp > 1 && nblk % sp != 0) { + --sp; + } + if (sp > 1) { + qmm_nax_splitk( + x, + w, + scales, + out, + group_size_, + bits_, + M, + N, + K, + sp, + d, + s, + kquant_type_); + return; + } + } + // Split-K qmm experiment (KQ_QMM_SPLITK=, 0 = off, read // once): the small-M band's occupancy lever; see qmm_splitk. Routes when // a >1 divisor of the wire-block count exists at or under the target. From cc09993b60ddff347b7221d3e52829a820492442 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:48:07 -0700 Subject: [PATCH 05/10] feat(matmul): mv_ext hd t-precision-dot experiment kernel (env-gated, +4-5% m8 only) --- metal/kq_quantized.metal | 25 +++++ .../mlx/backend/metal/kernels/kq_quantized.h | 101 ++++++++++++++++++ .../metal/kernels/kq_quantized_kquants.h | 16 +++ src/kquant_matmul.cpp | 12 ++- 4 files changed, 153 insertions(+), 1 deletion(-) diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index 6c7fc5a..5267871 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -450,6 +450,31 @@ instantiate_mv_ext_sb_all(q6_k, 256, 6) instantiate_mv_ext_nx_all(q6_k, 256, 6, 16) instantiate_mv_ext_nx_all(q6_k, 256, 6, 32) +// T-precision-dot experiment (KQ_MV_EXT_HD=1): the FMA-issue-bound band's +// ALU lever. Dequanted chunk converts float->T once (amortized over M rows), +// activations load at native T width with no per-row convert, and the +// 16-term chunk dot runs at half/bfloat issue rate before an f32 fold. +// q6_k only, M 4-12; no float x variant (no rate advantage). Suffix _hd. +#define instantiate_mv_ext_hd(codec, type, gs, bits, m) \ + instantiate_kernel( \ + "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m \ + "_hd", \ + kq_ ## codec ## _mv_ext_hd, type, m, 2, 8) +#define instantiate_mv_ext_hd_for_type(codec, gs, bits, type) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 4) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 5) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 6) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 7) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 8) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 9) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 10) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 11) \ + instantiate_mv_ext_hd(codec, type, gs, bits, 12) +#define instantiate_mv_ext_hd_all(codec, gs, bits) \ + instantiate_mv_ext_hd_for_type(codec, gs, bits, bfloat16_t) \ + instantiate_mv_ext_hd_for_type(codec, gs, bits, float16_t) +instantiate_mv_ext_hd_all(q6_k, 256, 6) + #define instantiate_kquant_q3_k_for_type(type) \ instantiate_kquant_batched(verify_qmv, type, 256, 3, 0, q3_k) \ instantiate_kquant_batched(qmv_fast, type, 256, 3, 0, q3_k) \ diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index 9a468d9..108b806 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -669,6 +669,107 @@ METAL_FUNC void kq_mv_ext_impl( } } +// T-precision-dot variant of kq_mv_ext_impl (suffix _hd). Past M~4 the base +// kernel is FMA-issue-bound, and its inner loop runs entirely at float rate: +// 16 activation T->float converts plus 4 float4 dots per chunk per row. Here +// the dequanted chunk converts float->T ONCE (amortized over r1ptg rows), +// activations load at native T width with no convert, and the 16-term chunk +// dot runs in vec arithmetic (2x issue rate for half/bfloat on M5) +// before folding into the f32 accumulator. Numerics: products and the +// 3-deep in-chunk adds round at T precision; cross-chunk accumulation stays +// f32, so error does not grow with K beyond the base kernel's. +template +METAL_FUNC void kq_mv_ext_hd_impl( + const device uint8_t* w, + const device T* x, + device T* y, + const constant int& in_vec_size, // K + const constant int& out_vec_size, // N + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + constexpr short nypsg = 32 / nxpsg; // output rows per simdgroup + constexpr short chpb = Codec::superblock / 16; // 16-weight chunks per block + const short tx = tiisg % nxpsg; // K position within the row group + const short ty = tiisg / nxpsg; // which of nypsg rows this thread owns + + const int i01 = tgpig.x * (nypsg * nsg) + nypsg * sgitg + ty; // output row + const int i11 = tgpig.y * r1ptg; // first activation column (grid.y==1 -> 0) + + const int nb = in_vec_size / Codec::superblock; + const int row_bytes = nb * Codec::block_bytes; + // Clamp OOB rows to row 0 for a valid read; the store is masked below. + const device uint8_t* w_row = + (i01 < out_vec_size) ? w + static_cast(i01) * row_bytes : w; + + const device T* y_col[r1ptg]; +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y_col[ir1] = x + static_cast(i11 + ir1) * in_vec_size + tx * 16; + } + + float sumf[r1ptg]; +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + sumf[ir1] = 0.0f; + } + + for (int ich = tx; 16 * ich < in_vec_size; ich += nxpsg) { + const int ib = ich / chpb; // super-block index + const short cch = ich % chpb; // chunk within super-block + const device uint8_t* block = + w_row + static_cast(ib) * Codec::block_bytes; + float4x4 lx; + Codec::deq_chunk16(block, cch, lx); + const vec lt0 = vec(lx[0]); + const vec lt1 = vec(lx[1]); + const vec lt2 = vec(lx[2]); + const vec lt3 = vec(lx[3]); +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + const device T* yp = y_col[ir1]; + const vec a0 = *(const device vec*)(yp + 0); + const vec a1 = *(const device vec*)(yp + 4); + const vec a2 = *(const device vec*)(yp + 8); + const vec a3 = *(const device vec*)(yp + 12); + vec pa = lt0 * a0; + pa += lt1 * a1; + pa += lt2 * a2; + pa += lt3 * a3; + const vec p2 = pa.xy + pa.zw; + sumf[ir1] += float(p2.x) + float(p2.y); + y_col[ir1] += nxpsg * 16; + } + } + +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + if (nxpsg >= 32) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 16); + } + if (nxpsg >= 16) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 8); + } + if (nxpsg >= 8) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 4); + } + if (nxpsg >= 4) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 2); + } + if (nxpsg >= 2) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 1); + } + } + + if (tx == 0 && i01 < out_vec_size) { +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y[static_cast(i11 + ir1) * out_vec_size + i01] = + static_cast(sumf[ir1]); + } + } +} + // Wide-M variant of kq_mv_ext_impl: each thread owns nr0 CONSECUTIVE output // rows instead of one. The nr0=1 kernel re-loads all r1ptg activation columns // per 16-weight chunk per row, so activation cache traffic scales as diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h index 01eac0f..3c48fd5 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h @@ -2898,6 +2898,22 @@ template w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); } +template +[[kernel]] void kq_q6_k_mv_ext_hd( + const device uint8_t* w, + const device uint8_t* /* scales */, + const device T* x, + device T* y, + const constant int& in_vec_size, // K + const constant int& out_vec_size, // N + const constant int& /* vm */, // == r1ptg + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + kq_mv_ext_hd_impl( + w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); +} + template [[kernel]] void kq_q6_k_qmv( const device uint8_t* w, diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index d738bb1..cf88d4a 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -875,6 +875,15 @@ void verify_mv_ext( }(); const bool use_nx = !use_nr2 && !use_sb && mv_ext_nx != 0 && M >= 4 && kquant_type == "q6_k"; + // T-precision-dot experiment (KQ_MV_EXT_HD=1): chunk dots in half/bfloat + // at 2x issue rate + no per-row activation converts, f32 fold per chunk. + // q6_k M 4-12, half/bfloat x only. + static const bool mv_ext_hd = []() { + const char* e = std::getenv("KQ_MV_EXT_HD"); + return e != nullptr && std::atoi(e) == 1; + }(); + const bool use_hd = !use_nr2 && !use_sb && !use_nx && mv_ext_hd && M >= 4 && + kquant_type == "q6_k" && x.dtype() != mx::float32; const int nxpsg_eff = use_nx ? mv_ext_nx : nxpsg; const int rows_per_tg = (32 / nxpsg_eff) * nsg * (use_nr2 ? 2 : 1); MTL::Size group_dims(32, nsg, 1); @@ -895,7 +904,8 @@ void verify_mv_ext( M, use_nr2 ? "_nr2" : (use_sb ? "_sb" - : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") : ""))); + : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") + : (use_hd ? "_hd" : "")))); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); From 8aa60a5e6ae374fbf118d2021471aae928ca4971 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:32:43 -0700 Subject: [PATCH 06/10] feat(sdpa): fused q8-kv operands for sdpa_decode_gqa (affine wire dequant on tile stage) --- bindings.cpp | 7 ++ metal/kq_quantized.metal | 25 +++++ .../mlx/backend/metal/kernels/kq_quantized.h | 105 ++++++++++++++++++ .../metal/kernels/kq_quantized_kquants.h | 17 +++ metal/mlx/backend/metal/kernels/kq_sdpa.h | 55 ++++++++- src/kquant.h | 11 +- src/kquant_matmul.cpp | 28 +++-- src/kquant_sdpa.cpp | 83 ++++++++++++-- 8 files changed, 311 insertions(+), 20 deletions(-) diff --git a/bindings.cpp b/bindings.cpp index c2f8132..0fc2b52 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -208,6 +208,10 @@ NB_MODULE(_ext, m) { "splits"_a = 0, "tile_c"_a = 0, "starts"_a = nb::none(), + "k_scales"_a = nb::none(), + "k_biases"_a = nb::none(), + "v_scales"_a = nb::none(), + "v_biases"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), R"( @@ -219,6 +223,9 @@ NB_MODULE(_ext, m) { tiles, causally clamped to its own trailing position. With `starts`, batch row b attends keys [starts[b], kL) -- a left-padded batched KV cache -- and fully padded-out key chunks are skipped, not staged. + With k_scales/k_biases/v_scales/v_biases (all four), k and v are + mlx affine-quantized wire (uint32, bits 8, group 64) and dequant is + fused into the tile stage. Args: q (array): queries [B, n_q_heads, qL, D], float16/bfloat16; diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index 5267871..da9cff9 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -475,6 +475,31 @@ instantiate_mv_ext_nx_all(q6_k, 256, 6, 32) instantiate_mv_ext_hd_for_type(codec, gs, bits, float16_t) instantiate_mv_ext_hd_all(q6_k, 256, 6) +// Staged-activation experiment (KQ_MV_EXT_TS=1): the M x 128 activation +// window stages into threadgroup memory once per K-step and 8 simdgroups +// (32 output rows) share it -- the activation-path lever the sb/nr2/nx +// falsifications never isolated. Dot math identical to base. q6_k M 4-12. +#define instantiate_mv_ext_ts(codec, type, gs, bits, m) \ + instantiate_kernel( \ + "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m \ + "_ts", \ + kq_ ## codec ## _mv_ext_ts, type, m, 8, 8) +#define instantiate_mv_ext_ts_for_type(codec, gs, bits, type) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 4) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 5) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 6) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 7) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 8) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 9) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 10) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 11) \ + instantiate_mv_ext_ts(codec, type, gs, bits, 12) +#define instantiate_mv_ext_ts_all(codec, gs, bits) \ + instantiate_mv_ext_ts_for_type(codec, gs, bits, float) \ + instantiate_mv_ext_ts_for_type(codec, gs, bits, bfloat16_t) \ + instantiate_mv_ext_ts_for_type(codec, gs, bits, float16_t) +instantiate_mv_ext_ts_all(q6_k, 256, 6) + #define instantiate_kquant_q3_k_for_type(type) \ instantiate_kquant_batched(verify_qmv, type, 256, 3, 0, q3_k) \ instantiate_kquant_batched(qmv_fast, type, 256, 3, 0, q3_k) \ diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index 108b806..1eb90c0 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -770,6 +770,111 @@ METAL_FUNC void kq_mv_ext_hd_impl( } } +// Staged-activation variant of kq_mv_ext_impl (suffix _ts). Weights stay in +// registers exactly as in the base kernel; the change is the activation +// path, the one lever the falsified variants never isolated: sb swapped +// loads for shuffles (worse throughput), nr2 amortized via registers +// (spilled), x16/x32 only changed which thread issues the loads. Here the +// M x (nxpsg*16) activation window stages into threadgroup memory once per +// K-step via a cooperative load, every row-thread dots from on-core SRAM, +// and the threadgroup carries nsg_ts simdgroups (32 rows at nsg_ts=8) so +// cross-TG device activation traffic drops rows_per_tg/8-fold vs the base +// kernel. Dot arithmetic is bit-identical to base (staging is a copy). +// K must be a multiple of nxpsg*16 (q6_k superblock 256 guarantees it). +template +METAL_FUNC void kq_mv_ext_ts_impl( + const device uint8_t* w, + const device T* x, + device T* y, + threadgroup T* staged, // r1ptg * nxpsg * 16 elements + const constant int& in_vec_size, // K + const constant int& out_vec_size, // N + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + constexpr short nypsg = 32 / nxpsg; // output rows per simdgroup + constexpr short chpb = Codec::superblock / 16; // 16-weight chunks per block + constexpr short stage_w = nxpsg * 16; // staged K-window elements per row + const short tx = tiisg % nxpsg; // K position within the row group + const short ty = tiisg / nxpsg; // which of nypsg rows this thread owns + + const int i01 = tgpig.x * (nypsg * nsg) + nypsg * sgitg + ty; // output row + const int i11 = tgpig.y * r1ptg; // first activation column (grid.y==1 -> 0) + + const int nb = in_vec_size / Codec::superblock; + const int row_bytes = nb * Codec::block_bytes; + // Clamp OOB rows to row 0 for a valid read; the store is masked below. + const device uint8_t* w_row = + (i01 < out_vec_size) ? w + static_cast(i01) * row_bytes : w; + + const short lin = sgitg * 32 + tiisg; // linear thread id in the TG + constexpr short tg_threads = nsg * 32; + + float sumf[r1ptg]; +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + sumf[ir1] = 0.0f; + } + + for (int base = 0; 16 * base < in_vec_size; base += nxpsg) { + // Cooperative stage of the M x stage_w activation window. + threadgroup_barrier(mem_flags::mem_threadgroup); + const int kw = 16 * base; + for (short f = lin; f < r1ptg * stage_w; f += tg_threads) { + const short ir1 = f / stage_w; + const short j = f % stage_w; + staged[f] = x[static_cast(i11 + ir1) * in_vec_size + kw + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int ich = base + tx; + const int ib = ich / chpb; // super-block index + const short cch = ich % chpb; // chunk within super-block + const device uint8_t* block = + w_row + static_cast(ib) * Codec::block_bytes; + float4x4 lx; + Codec::deq_chunk16(block, cch, lx); + const threadgroup T* sp = staged + tx * 16; +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + const threadgroup T* yp = sp + ir1 * stage_w; + const float4 a0 = float4(*(const threadgroup vec*)(yp + 0)); + const float4 a1 = float4(*(const threadgroup vec*)(yp + 4)); + const float4 a2 = float4(*(const threadgroup vec*)(yp + 8)); + const float4 a3 = float4(*(const threadgroup vec*)(yp + 12)); + sumf[ir1] += + dot(lx[0], a0) + dot(lx[1], a1) + dot(lx[2], a2) + dot(lx[3], a3); + } + } + +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + if (nxpsg >= 32) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 16); + } + if (nxpsg >= 16) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 8); + } + if (nxpsg >= 8) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 4); + } + if (nxpsg >= 4) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 2); + } + if (nxpsg >= 2) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 1); + } + } + + if (tx == 0 && i01 < out_vec_size) { +#pragma unroll + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y[static_cast(i11 + ir1) * out_vec_size + i01] = + static_cast(sumf[ir1]); + } + } +} + // Wide-M variant of kq_mv_ext_impl: each thread owns nr0 CONSECUTIVE output // rows instead of one. The nr0=1 kernel re-loads all r1ptg activation columns // per 16-weight chunk per row, so activation cache traffic scales as diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h index 3c48fd5..214bf83 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h @@ -2898,6 +2898,23 @@ template w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); } +template +[[kernel]] void kq_q6_k_mv_ext_ts( + const device uint8_t* w, + const device uint8_t* /* scales */, + const device T* x, + device T* y, + const constant int& in_vec_size, // K + const constant int& out_vec_size, // N + const constant int& /* vm */, // == r1ptg + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + threadgroup T staged[r1ptg * nxpsg * 16]; + kq_mv_ext_ts_impl( + w, x, y, staged, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); +} + template [[kernel]] void kq_q6_k_mv_ext_hd( const device uint8_t* w, diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h index 3fe5ace..e64a146 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h @@ -16,6 +16,11 @@ constant int blocks [[function_constant(1)]]; constant int gqa_splits [[function_constant(2)]]; constant bool gqa_has_sinks [[function_constant(3)]]; constant bool gqa_has_starts [[function_constant(4)]]; +// Quantized KV operands (mlx affine wire, bits 8 / group 64): keys/values +// bind as packed uint32 words plus per-group scale/bias arrays, and the +// cooperative tile stage dequants into sK/sV -- everything downstream of +// the stage is unchanged. Compiled out when false. +constant bool gqa_kv_q8 [[function_constant(5)]]; template [[kernel]] void kq_sdpa_vector_2pass_1( @@ -213,6 +218,14 @@ template const constant float& scale [[buffer(11)]], const constant int& q_len [[buffer(12)]], const device int* starts [[buffer(13)]], + const device T* k_scales [[buffer(14)]], + const device T* k_biases [[buffer(15)]], + const device T* v_scales [[buffer(16)]], + const device T* v_biases [[buffer(17)]], + const constant size_t& ks_head_stride [[buffer(18)]], + const constant size_t& ks_seq_stride [[buffer(19)]], + const constant size_t& vs_head_stride [[buffer(20)]], + const constant size_t& vs_seq_stride [[buffer(21)]], uint3 tptg [[threads_per_threadgroup]], uint3 tidtg [[thread_position_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -220,6 +233,7 @@ template constexpr int D4 = D / 4; constexpr int NL = 32 / NE; // lanes per in-flight key constexpr int DP4 = D4 / NL; // float4s per lane per key row + constexpr int GP4 = 64 / 4; // packed words per quant group (group_size 64) using T4 = metal::vec; threadgroup T4 sK[C * D4]; @@ -256,10 +270,21 @@ template kt0 = max(k0, (row_start / C) * C); } + // With gqa_kv_q8 the k/v buffers hold packed uint32 wire and the strides + // arrive in WORDS (D/4 per row); otherwise they hold T elements. const device T* kbase = keys + (size_t)(batch_idx * num_kv_heads + kv_head_idx) * k_head_stride; const device T* vbase = values + (size_t)(batch_idx * num_kv_heads + kv_head_idx) * v_head_stride; + const size_t kv_hb = (size_t)(batch_idx * num_kv_heads + kv_head_idx); + const device uint32_t* kwbase = + (const device uint32_t*)keys + kv_hb * k_head_stride; + const device uint32_t* vwbase = + (const device uint32_t*)values + kv_hb * v_head_stride; + const device T* ksb = k_scales + kv_hb * ks_head_stride; + const device T* kbb = k_biases + kv_hb * ks_head_stride; + const device T* vsb = v_scales + kv_hb * vs_head_stride; + const device T* vbb = v_biases + kv_hb * vs_head_stride; // Pre-scaled query slices for this lane's key-row columns // ([B, Hq, q_len, D], row-contiguous). A simdgroup past the runtime query @@ -294,14 +319,38 @@ template for (int kt = kt0; kt < k1; kt += C) { threadgroup_barrier(mem_flags::mem_threadgroup); // Cooperative tile load; zero-fill the tail so stale threadgroup data - // can never reach the accumulators. + // can never reach the accumulators. On the q8 path each uint32 word + // holds 4 codes; dequant (scale * q + bias, per 64-element group) lands + // in sK/sV at the same T4 precision as the fp path. for (int i = flat; i < C * D4; i += n_threads) { const int row = i / D4; const int col = i % D4; const int kg = kt + row; if (kg < k1) { - sK[i] = ((const device T4*)(kbase + (size_t)kg * k_seq_stride))[col]; - sV[i] = ((const device T4*)(vbase + (size_t)kg * v_seq_stride))[col]; + if (gqa_kv_q8) { + const uint32_t kw = kwbase[(size_t)kg * k_seq_stride + col]; + const uint32_t vw = vwbase[(size_t)kg * v_seq_stride + col]; + const size_t gs = col / GP4; + const float ksc = float(ksb[(size_t)kg * ks_seq_stride + gs]); + const float kbi = float(kbb[(size_t)kg * ks_seq_stride + gs]); + const float vsc = float(vsb[(size_t)kg * vs_seq_stride + gs]); + const float vbi = float(vbb[(size_t)kg * vs_seq_stride + gs]); + const float4 kq4 = float4( + float(kw & 0xff), + float((kw >> 8) & 0xff), + float((kw >> 16) & 0xff), + float(kw >> 24)); + const float4 vq4 = float4( + float(vw & 0xff), + float((vw >> 8) & 0xff), + float((vw >> 16) & 0xff), + float(vw >> 24)); + sK[i] = T4(kq4 * ksc + kbi); + sV[i] = T4(vq4 * vsc + vbi); + } else { + sK[i] = ((const device T4*)(kbase + (size_t)kg * k_seq_stride))[col]; + sV[i] = ((const device T4*)(vbase + (size_t)kg * v_seq_stride))[col]; + } } else { sK[i] = T4(T(0)); sV[i] = T4(T(0)); diff --git a/src/kquant.h b/src/kquant.h index 33f8793..6c1d98c 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -174,6 +174,10 @@ mx::array sdpa_decode_gqa( int splits = 0, int tile_c = 32, const std::optional& starts = std::nullopt, + const std::optional& k_scales = std::nullopt, + const std::optional& k_biases = std::nullopt, + const std::optional& v_scales = std::nullopt, + const std::optional& v_biases = std::nullopt, mx::StreamOrDevice s = {}); // Speculative-verify attention on the GPU matrix units for a GQA-folded query @@ -700,13 +704,15 @@ class KQuantSDPAGQA : public mx::Primitive { int splits, int tile_c, bool has_sinks, - bool has_starts) + bool has_starts, + bool has_kv_q8 = false) : mx::Primitive(stream), scale_(scale), splits_(splits), tile_c_(tile_c), has_sinks_(has_sinks), - has_starts_(has_starts) {} + has_starts_(has_starts), + has_kv_q8_(has_kv_q8) {} void eval_cpu( const std::vector& inputs, @@ -729,6 +735,7 @@ class KQuantSDPAGQA : public mx::Primitive { int tile_c_; bool has_sinks_; bool has_starts_; + bool has_kv_q8_; }; // Simdgroup-matrix FA verify attention (see sdpa_fa_verify). Inference-only. diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index cf88d4a..6df497f 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -884,9 +884,19 @@ void verify_mv_ext( }(); const bool use_hd = !use_nr2 && !use_sb && !use_nx && mv_ext_hd && M >= 4 && kquant_type == "q6_k" && x.dtype() != mx::float32; + // Staged-activation experiment (KQ_MV_EXT_TS=1): cooperative TG-memory + // stage of the activation window, 8 simdgroups / 32 rows per TG. q6_k + // M 4-12; K must cover a full 128-element window (q6_k geometry does). + static const bool mv_ext_ts = []() { + const char* e = std::getenv("KQ_MV_EXT_TS"); + return e != nullptr && std::atoi(e) == 1; + }(); + const bool use_ts = !use_nr2 && !use_sb && !use_nx && !use_hd && mv_ext_ts && + M >= 4 && kquant_type == "q6_k" && K % 128 == 0; + const int nsg_eff = use_ts ? 8 : nsg; const int nxpsg_eff = use_nx ? mv_ext_nx : nxpsg; - const int rows_per_tg = (32 / nxpsg_eff) * nsg * (use_nr2 ? 2 : 1); - MTL::Size group_dims(32, nsg, 1); + const int rows_per_tg = (32 / nxpsg_eff) * nsg_eff * (use_nr2 ? 2 : 1); + MTL::Size group_dims(32, nsg_eff, 1); MTL::Size grid_dims((N + rows_per_tg - 1) / rows_per_tg, 1, 1); std::string type_string = kq_type_string(x.dtype()); @@ -905,7 +915,7 @@ void verify_mv_ext( use_nr2 ? "_nr2" : (use_sb ? "_sb" : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") - : (use_hd ? "_hd" : "")))); + : (use_hd ? "_hd" : (use_ts ? "_ts" : ""))))); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -1127,13 +1137,15 @@ void KQuantMatmul::eval_gpu( } // NAX split-K qmm experiment (KQ_QMM_SPLITK_NAX=, 0 = off, - // read once): K-slices on the tensor-core BM=32 tile; see qmm_nax_splitk. + // 1 = auto target 32; read LIVE per call so --ab-env can flip it on one + // generator): K-slices on the tensor-core BM=32 tile; see qmm_nax_splitk. // Slice quantum is max(superblock, BK) so every slice starts a loader at // kt_base 0. q6_k + q8_0 instantiations only. - static const int qmm_splitk_nax_env = []() { - const char* e = std::getenv("KQ_QMM_SPLITK_NAX"); - return e != nullptr ? std::atoi(e) : 0; - }(); + const char* sk_nax_e = std::getenv("KQ_QMM_SPLITK_NAX"); + int qmm_splitk_nax_env = sk_nax_e != nullptr ? std::atoi(sk_nax_e) : 0; + if (qmm_splitk_nax_env == 1) { + qmm_splitk_nax_env = 32; + } if (qmm_splitk_nax_env > 1 && transpose_ && non_batched && M <= 32 && (kquant_type_ == "q6_k" || kquant_type_ == "q8_0") && kq_is_nax_available() && (K % 64 == 0) && x.dtype() != mx::float32) { diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 81a39d7..1514cf5 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -208,7 +208,9 @@ void KQuantSDPAGQA::eval_gpu( const auto& v = inputs[2]; const bool sinks = has_sinks_; const bool starts = has_starts_; + const bool kv_q8 = has_kv_q8_; const size_t starts_idx = 3 + (sinks ? 1 : 0); + const size_t qkv_idx = starts_idx + (starts ? 1 : 0); kq_sdpa_check_layout("sdpa_decode_gqa", q, k, v); int B = q.shape(0); @@ -252,10 +254,12 @@ void KQuantSDPAGQA::eval_gpu( std::string ts = kq_type_string(q.dtype()); bool has_sinks = sinks; bool has_starts = starts; + bool has_kv_q8 = kv_q8; mx::metal::MTLFCList fc = { {&splits, MTL::DataType::DataTypeInt, 2}, {&has_sinks, MTL::DataType::DataTypeBool, 3}, {&has_starts, MTL::DataType::DataTypeBool, 4}, + {&has_kv_q8, MTL::DataType::DataTypeBool, 5}, }; // Pass 1: one threadgroup per (kv-head, batch, split); the whole GQA group @@ -265,8 +269,8 @@ void KQuantSDPAGQA::eval_gpu( { std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) + "_c" + std::to_string(tile_c_) + (qL > 1 ? "_p2" : ""); - std::string hash = - kname + "_s" + std::to_string(splits) + (has_starts ? "_st1" : "_st0"); + std::string hash = kname + "_s" + std::to_string(splits) + + (has_starts ? "_st1" : "_st0") + (has_kv_q8 ? "_q8" : ""); auto kernel = kq_get_kernel(d, kname, hash, fc); // Register-heavy pipeline: some GPUs cap it below the dispatch width, and // Metal turns an oversized dispatch into silent garbage, not an error. @@ -295,6 +299,27 @@ void KQuantSDPAGQA::eval_gpu( // Metal wants every buffer bound; without starts, rebind sums as a dummy // (the read is compiled out via the function constant). ce.set_input_array(starts ? inputs[starts_idx] : sums, 13); + // Quantized-KV scale/bias operands (dummies when compiled out). Scale + // strides are shared per K/V side (scales and biases are congruent). + size_t ks_head_stride = 0, ks_seq_stride = 0; + size_t vs_head_stride = 0, vs_seq_stride = 0; + if (kv_q8) { + const auto& ksc = inputs[qkv_idx]; + const auto& vsc = inputs[qkv_idx + 2]; + ks_head_stride = static_cast( + ksc.shape(1) == 1 ? ksc.strides(0) : ksc.strides(1)); + ks_seq_stride = static_cast(ksc.strides(2)); + vs_head_stride = static_cast( + vsc.shape(1) == 1 ? vsc.strides(0) : vsc.strides(1)); + vs_seq_stride = static_cast(vsc.strides(2)); + } + for (int i = 0; i < 4; i++) { + ce.set_input_array(kv_q8 ? inputs[qkv_idx + i] : sums, 14 + i); + } + ce.set_bytes(ks_head_stride, 18); + ce.set_bytes(ks_seq_stride, 19); + ce.set_bytes(vs_head_stride, 20); + ce.set_bytes(vs_seq_stride, 21); MTL::Size group_dims(32, gqa_factor, qL > 1 ? (qL + 1) / 2 : 1); MTL::Size grid_dims(n_kv_heads, B, splits); ce.dispatch_threadgroups(grid_dims, group_dims); @@ -564,7 +589,8 @@ std::vector KQuantSDPAGQA::output_shapes( bool KQuantSDPAGQA::is_equivalent(const mx::Primitive& other) const { const auto& o = static_cast(other); return scale_ == o.scale_ && splits_ == o.splits_ && tile_c_ == o.tile_c_ && - has_sinks_ == o.has_sinks_ && has_starts_ == o.has_starts_; + has_sinks_ == o.has_sinks_ && has_starts_ == o.has_starts_ && + has_kv_q8_ == o.has_kv_q8_; } mx::array sdpa_decode_gqa( @@ -576,20 +602,51 @@ mx::array sdpa_decode_gqa( int splits, int tile_c, const std::optional& starts, + const std::optional& k_scales, + const std::optional& k_biases, + const std::optional& v_scales, + const std::optional& v_biases, mx::StreamOrDevice s_) { auto s = mx::to_stream(s_); + const int n_qkv = int(k_scales.has_value()) + int(k_biases.has_value()) + + int(v_scales.has_value()) + int(v_biases.has_value()); + const bool kv_q8 = n_qkv == 4; + if (n_qkv != 0 && n_qkv != 4) { + throw std::invalid_argument( + "[mlx_kquant.sdpa_decode_gqa] quantized KV needs all four of " + "k_scales/k_biases/v_scales/v_biases."); + } + if (q.ndim() != 4 || k.ndim() != 4 || v.ndim() != 4) { throw std::invalid_argument( "[mlx_kquant.sdpa_decode_gqa] q, k, v must be 4-D [B, heads, L, D]."); } int D = q.shape(-1); - if ((D != 64 && D != 128 && D != 256 && D != 512) || v.shape(-1) != D || - k.shape(-1) != D) { + if ((D != 64 && D != 128 && D != 256 && D != 512) || + (!kv_q8 && (v.shape(-1) != D || k.shape(-1) != D))) { throw std::invalid_argument( "[mlx_kquant.sdpa_decode_gqa] only head_dim 64/128/256/512 is " "supported."); } + if (kv_q8) { + // mlx affine wire, bits 8 / group 64: packed uint32 words, one + // scale/bias per 64-element group ([B, Hkv, S, D/64], q's dtype). + if (k.dtype() != mx::uint32 || v.dtype() != mx::uint32 || + k.shape(-1) != D / 4 || v.shape(-1) != D / 4) { + throw std::invalid_argument( + "[mlx_kquant.sdpa_decode_gqa] quantized k/v must be uint32 wire " + "with last dim head_dim / 4 (bits 8)."); + } + for (const auto& a : {*k_scales, *k_biases, *v_scales, *v_biases}) { + if (a.dtype() != q.dtype() || a.ndim() != 4 || a.shape(-1) != D / 64 || + a.shape(2) != k.shape(2)) { + throw std::invalid_argument( + "[mlx_kquant.sdpa_decode_gqa] quantized KV scales/biases must " + "be [B, n_kv_heads, S, head_dim / 64] in q's dtype (group 64)."); + } + } + } int qL = q.shape(2); if (qL < 1 || qL > 4) { throw std::invalid_argument( @@ -601,7 +658,7 @@ mx::array sdpa_decode_gqa( throw std::invalid_argument( "[mlx_kquant.sdpa_decode_gqa] q must be float16 or bfloat16."); } - if (k.dtype() != dt || v.dtype() != dt) { + if (!kv_q8 && (k.dtype() != dt || v.dtype() != dt)) { throw std::invalid_argument( "[mlx_kquant.sdpa_decode_gqa] q, k, v must share a dtype."); } @@ -677,13 +734,25 @@ mx::array sdpa_decode_gqa( st = mx::reshape(st, {q.shape(0)}, s); inputs.push_back(mx::contiguous(st, false, s)); } + if (kv_q8) { + for (const auto& a : {*k_scales, *k_biases, *v_scales, *v_biases}) { + inputs.push_back( + a.strides().back() == 1 ? a : mx::contiguous(a, false, s)); + } + } auto out_shape = q.shape(); return mx::array( std::move(out_shape), dt, std::make_shared( - s, scale, splits, tile_c, sinks.has_value(), starts.has_value()), + s, + scale, + splits, + tile_c, + sinks.has_value(), + starts.has_value(), + kv_q8), std::move(inputs)); } From a9638850790973807cf5a6b0e68fadcf5cc3ce4c Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:56:32 -0700 Subject: [PATCH 07/10] perf(sdpa): vectorize q8 kv stage to uint4 units with hoisted scales --- metal/mlx/backend/metal/kernels/kq_sdpa.h | 84 +++++++++++++++-------- 1 file changed, 57 insertions(+), 27 deletions(-) diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h index e64a146..343d708 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h @@ -319,41 +319,71 @@ template for (int kt = kt0; kt < k1; kt += C) { threadgroup_barrier(mem_flags::mem_threadgroup); // Cooperative tile load; zero-fill the tail so stale threadgroup data - // can never reach the accumulators. On the q8 path each uint32 word - // holds 4 codes; dequant (scale * q + bias, per 64-element group) lands - // in sK/sV at the same T4 precision as the fp path. - for (int i = flat; i < C * D4; i += n_threads) { - const int row = i / D4; - const int col = i % D4; - const int kg = kt + row; - if (kg < k1) { - if (gqa_kv_q8) { - const uint32_t kw = kwbase[(size_t)kg * k_seq_stride + col]; - const uint32_t vw = vwbase[(size_t)kg * v_seq_stride + col]; - const size_t gs = col / GP4; + // can never reach the accumulators. The q8 path works in uint4 units + // (16 codes): one vectorized wire load and ONE scale/bias pair per + // unit -- 16 consecutive elements never straddle a 64-element group -- + // instead of per-word scalar traffic, which measured ~40% below the + // fp16 kernel's bandwidth. + if (gqa_kv_q8) { + constexpr int DU = D / 16; // uint4 units per row + for (int u = flat; u < C * DU; u += n_threads) { + const int row = u / DU; + const int c16 = u % DU; + const int kg = kt + row; + const int i = row * D4 + c16 * 4; + if (kg < k1) { + const uint4 kw = + ((const device uint4*)(kwbase + (size_t)kg * k_seq_stride))[c16]; + const uint4 vw = + ((const device uint4*)(vwbase + (size_t)kg * v_seq_stride))[c16]; + const size_t gs = (c16 * 16) / 64; const float ksc = float(ksb[(size_t)kg * ks_seq_stride + gs]); const float kbi = float(kbb[(size_t)kg * ks_seq_stride + gs]); const float vsc = float(vsb[(size_t)kg * vs_seq_stride + gs]); const float vbi = float(vbb[(size_t)kg * vs_seq_stride + gs]); - const float4 kq4 = float4( - float(kw & 0xff), - float((kw >> 8) & 0xff), - float((kw >> 16) & 0xff), - float(kw >> 24)); - const float4 vq4 = float4( - float(vw & 0xff), - float((vw >> 8) & 0xff), - float((vw >> 16) & 0xff), - float(vw >> 24)); - sK[i] = T4(kq4 * ksc + kbi); - sV[i] = T4(vq4 * vsc + vbi); +#pragma unroll + for (short w = 0; w < 4; w++) { + const uint32_t kx = w == 0 ? kw.x + : w == 1 ? kw.y + : w == 2 ? kw.z + : kw.w; + const uint32_t vx = w == 0 ? vw.x + : w == 1 ? vw.y + : w == 2 ? vw.z + : vw.w; + const float4 kq4 = float4( + float(kx & 0xff), + float((kx >> 8) & 0xff), + float((kx >> 16) & 0xff), + float(kx >> 24)); + const float4 vq4 = float4( + float(vx & 0xff), + float((vx >> 8) & 0xff), + float((vx >> 16) & 0xff), + float(vx >> 24)); + sK[i + w] = T4(kq4 * ksc + kbi); + sV[i + w] = T4(vq4 * vsc + vbi); + } } else { +#pragma unroll + for (short w = 0; w < 4; w++) { + sK[i + w] = T4(T(0)); + sV[i + w] = T4(T(0)); + } + } + } + } else { + for (int i = flat; i < C * D4; i += n_threads) { + const int row = i / D4; + const int col = i % D4; + const int kg = kt + row; + if (kg < k1) { sK[i] = ((const device T4*)(kbase + (size_t)kg * k_seq_stride))[col]; sV[i] = ((const device T4*)(vbase + (size_t)kg * v_seq_stride))[col]; + } else { + sK[i] = T4(T(0)); + sV[i] = T4(T(0)); } - } else { - sK[i] = T4(T(0)); - sV[i] = T4(T(0)); } } threadgroup_barrier(mem_flags::mem_threadgroup); From 73d503cb2baab35f340264a15884bce75d3709a4 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:23:49 -0700 Subject: [PATCH 08/10] changelog: smallm-splitk entries under unreleased --- CHANGELOG.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c0f9be..34cec9a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - `sdpa_decode_gqa` optional `starts` (int32 [B]): per-batch-row key start offsets for left-padded batched KV caches; padded-out key chunks are skipped, not staged. +- `sdpa_decode_gqa` q8 KV operands (affine wire, bits 8, group 64): batched + decode attends over quantized KV directly, dequantizing on the staged + tile; up to 1.9x/call at depth vs dequantize-then-attend. +- Env-gated small-M qmm experiment kernels (`KQ_QMM_SPLITK`, + `KQ_QMM_SPLITK_NAX`, `KQ_MV_EXT_SB`, `KQ_MV_EXT_NX`, `KQ_MV_EXT_HD`): + the NAX split-K path lifts the collapsed M9-16 band 65-76%; the rest + measured flat to negative on M5 and stay off by default. ## [0.3.7] From 2ba9a92c61fe6e2581653a5b88e2828e7515151b Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:07:38 -0700 Subject: [PATCH 09/10] tests: skip bf16 bk256 dsa cases on paravirtual metal (transient backend-compile flake) --- tests/test_dsa_sparse_attn.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_dsa_sparse_attn.py b/tests/test_dsa_sparse_attn.py index 548411f..c0afcc3 100644 --- a/tests/test_dsa_sparse_attn.py +++ b/tests/test_dsa_sparse_attn.py @@ -36,6 +36,18 @@ reason="kq.dsa_sparse_attention is a Metal-only kernel; no CPU path.", ) +# Hosted-CI macOS runners expose a paravirtualized Metal device whose backend +# compiler transiently fails the heaviest pipeline build in this suite (the +# non-split bf16 bk256 instantiation) at first touch: "Unable to load kernel +# ... Compilation failed", while the same pipeline compiles fine moments +# later. Real hardware never reports a paravirtual device name, so this +# gates only virtualized runners; every case keeps real-GPU coverage. +try: + _DEVICE_NAME = str(mx.device_info().get("device_name", "")) +except Exception: + _DEVICE_NAME = "" +_VIRTUAL_GPU = "paravirtual" in _DEVICE_NAME.lower() + REL_BOUND = {mx.bfloat16: 5e-3, mx.float16: 2e-3} H, D = 64, 512 # kernel geometry (fixed by the instantiations) @@ -147,6 +159,14 @@ def _case(qL, localL, P, topk_n, q_offset, ratio, window, dtype, seed, mode="ran @pytest.mark.parametrize("case", CASES, ids=[c[0] for c in CASES]) def test_dsa_sparse_attention(case, dtype): name, qL, localL, P, topk_n, q_offset, ratio, window, mode = case + # qL > 4 forces the non-split route; topk_n > 128 selects the bk256 + # tile: together the exact set of cases that first-touch the pipeline + # the paravirtual backend flakes on. + if _VIRTUAL_GPU and dtype == mx.bfloat16 and topk_n > 128 and qL > 4: + pytest.skip( + "paravirtual Metal: transient backend-compile failure " + "on the bf16 bk256 pipeline; covered on real GPUs" + ) rel = _case( qL, localL, From d981c305119bf5397d4d2918307b3e3545524a3b Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:42:21 -0700 Subject: [PATCH 10/10] tests: prime the bk256 dsa pipeline with retries on paravirtual metal (covers the forced non-split route) --- tests/test_dsa_sparse_attn.py | 54 ++++++++++++++++++++++++++++------- 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/tests/test_dsa_sparse_attn.py b/tests/test_dsa_sparse_attn.py index c0afcc3..73fceb0 100644 --- a/tests/test_dsa_sparse_attn.py +++ b/tests/test_dsa_sparse_attn.py @@ -24,6 +24,7 @@ import os import sys +import time import mlx.core as mx import numpy as np @@ -38,16 +39,47 @@ # Hosted-CI macOS runners expose a paravirtualized Metal device whose backend # compiler transiently fails the heaviest pipeline build in this suite (the -# non-split bf16 bk256 instantiation) at first touch: "Unable to load kernel -# ... Compilation failed", while the same pipeline compiles fine moments -# later. Real hardware never reports a paravirtual device name, so this -# gates only virtualized runners; every case keeps real-GPU coverage. +# non-split bk256 instantiation) at first touch: "Unable to load kernel ... +# Compilation failed", while the same pipeline compiles fine moments later. +# On such runners, settle that pipeline once with retries before any test +# relies on it; tests skip only if it never compiles. Real hardware never +# reports a paravirtual device name, so this is a no-op there. try: _DEVICE_NAME = str(mx.device_info().get("device_name", "")) except Exception: _DEVICE_NAME = "" _VIRTUAL_GPU = "paravirtual" in _DEVICE_NAME.lower() +_PRIMED: dict[str, bool] = {} + + +def _heavy_pipeline_ready(dtype): + """Prime the non-split bk256 pipeline for dtype (qL 8 forces the + non-split route regardless of KQ_DSA_SPLIT). True on real hardware + without compiling anything.""" + if not _VIRTUAL_GPU: + return True + key = str(dtype) + if key not in _PRIMED: + q, local_kv, pooled, topk, sinks, scale = _build( + 8, 136, 700, 512, 2792, 4, 128, dtype, seed=7, mode="random" + ) + ok = False + for _ in range(4): + try: + mx.eval( + kq.dsa_sparse_attention( + q, local_kv, pooled, topk, sinks, scale, 2792, 4, 128 + ) + ) + ok = True + break + except RuntimeError: + time.sleep(2) + _PRIMED[key] = ok + return _PRIMED[key] + + REL_BOUND = {mx.bfloat16: 5e-3, mx.float16: 2e-3} H, D = 64, 512 # kernel geometry (fixed by the instantiations) @@ -160,13 +192,9 @@ def _case(qL, localL, P, topk_n, q_offset, ratio, window, dtype, seed, mode="ran def test_dsa_sparse_attention(case, dtype): name, qL, localL, P, topk_n, q_offset, ratio, window, mode = case # qL > 4 forces the non-split route; topk_n > 128 selects the bk256 - # tile: together the exact set of cases that first-touch the pipeline - # the paravirtual backend flakes on. - if _VIRTUAL_GPU and dtype == mx.bfloat16 and topk_n > 128 and qL > 4: - pytest.skip( - "paravirtual Metal: transient backend-compile failure " - "on the bf16 bk256 pipeline; covered on real GPUs" - ) + # tile: the pipeline the paravirtual backend flakes on. + if topk_n > 128 and qL > 4 and not _heavy_pipeline_ready(dtype): + pytest.skip("paravirtual Metal never compiled the bk256 pipeline") rel = _case( qL, localL, @@ -204,6 +232,10 @@ def test_dsa_sparse_attention_split_matches_base(case, dtype, monkeypatch): under the reference tolerance. """ name, qL, localL, P, topk_n, q_offset, ratio, window, mode = case + # The forced non-split arm (KQ_DSA_SPLIT=0) dispatches the bk256 + # pipeline whenever topk_n > 128, at any qL. + if topk_n > 128 and not _heavy_pipeline_ready(dtype): + pytest.skip("paravirtual Metal never compiled the bk256 pipeline") q, local_kv, pooled, topk, sinks, scale = _build( qL, localL, P, topk_n, q_offset, ratio, window, dtype, seed=11, mode=mode )