diff --git a/CHANGELOG.md b/CHANGELOG.md index 34cec9a..c44db33 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). `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. +- `sdpa_decode_gqa_cascade`: fused shared-prefix batched decode; one KV + walk serves the shared prefix for every batch row, private suffixes read + per row. 1.6-4.2x vs per-row calls at P 14k-32k, hd128/hd256. +- `sdpa_fa_verify` head_dim 64/128 tiles; `return_lse` on + `sdpa_decode_gqa` and `sdpa_fa_verify`. ## [0.3.7] diff --git a/bindings.cpp b/bindings.cpp index 0fc2b52..91b9292 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -199,7 +199,52 @@ NB_MODULE(_ext, m) { m.def( "sdpa_decode_gqa", - &mlx_kquant::sdpa_decode_gqa, + [](mx::array q, + mx::array k, + mx::array v, + float scale, + const std::optional& sinks, + 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, + bool return_lse, + mx::StreamOrDevice s) -> nb::object { + if (return_lse) { + auto outs = mlx_kquant::sdpa_decode_gqa_lse( + std::move(q), + std::move(k), + std::move(v), + scale, + sinks, + splits, + tile_c, + starts, + k_scales, + k_biases, + v_scales, + v_biases, + s); + return nb::make_tuple(outs[0], outs[1]); + } + return nb::cast(mlx_kquant::sdpa_decode_gqa( + std::move(q), + std::move(k), + std::move(v), + scale, + sinks, + splits, + tile_c, + starts, + k_scales, + k_biases, + v_scales, + v_biases, + s)); + }, "q"_a, "k"_a, "v"_a, @@ -213,6 +258,7 @@ NB_MODULE(_ext, m) { "v_scales"_a = nb::none(), "v_biases"_a = nb::none(), nb::kw_only(), + "return_lse"_a = false, "stream"_a = nb::none(), R"( Decode/verify GQA attention tuned for long KV caches: the key axis @@ -244,12 +290,37 @@ NB_MODULE(_ext, m) { kL). Out-of-range values read as an empty row (zero output). Returns: - array: attention output [B, n_q_heads, qL, D]. + array: attention output [B, n_q_heads, qL, D]. With + ``return_lse=True``, a tuple ``(out, lse)`` where lse + [B, n_q_heads, qL] float32 is the natural-log softmax + normalizer per query row (the merge weight for combining + attention over disjoint key regions). )"); m.def( "sdpa_fa_verify", - &mlx_kquant::sdpa_fa_verify, + [](mx::array q, + mx::array k, + mx::array v, + float scale, + int q_len, + int splits, + bool return_lse, + mx::StreamOrDevice s) -> nb::object { + if (return_lse) { + auto outs = mlx_kquant::sdpa_fa_verify_lse( + std::move(q), + std::move(k), + std::move(v), + scale, + q_len, + splits, + s); + return nb::make_tuple(outs[0], outs[1]); + } + return nb::cast(mlx_kquant::sdpa_fa_verify( + std::move(q), std::move(k), std::move(v), scale, q_len, splits, s)); + }, "q"_a, "k"_a, "v"_a, @@ -257,6 +328,7 @@ NB_MODULE(_ext, m) { "q_len"_a, "splits"_a = 0, nb::kw_only(), + "return_lse"_a = false, "stream"_a = nb::none(), R"( Speculative-verify attention on the GPU matrix units for a GQA-folded @@ -272,8 +344,8 @@ NB_MODULE(_ext, m) { Args: q (array): folded queries [1, n_kv_heads, G*q_len, D], - float16/bfloat16; D = 256 or 512; G*q_len <= 64 at D=256, - <= 32 at D=512. + float16/bfloat16; D = 64, 128, 256 or 512; G*q_len <= 64 + except <= 32 at D=512. k (array): keys [1, n_kv_heads, kL, D]; head/seq strided is fine (read in place), the head_dim must be contiguous. v (array): values [1, n_kv_heads, kL, D]. @@ -284,7 +356,86 @@ NB_MODULE(_ext, m) { splits (int): key-axis split count; 0 picks the default. Returns: - array: attention output [1, n_kv_heads, G*q_len, D]. + array: attention output [1, n_kv_heads, G*q_len, D]. With + ``return_lse=True``, a tuple ``(out, lse)`` where lse + [1, n_kv_heads, G*q_len] float32 is the natural-log softmax + normalizer per folded row (cascade merge weight). + )"); + + m.def( + "sdpa_decode_gqa_cascade", + [](mx::array q, + mx::array k_shared, + mx::array v_shared, + mx::array k_priv, + mx::array v_priv, + float scale, + const std::optional& starts, + int splits_shared, + int splits_priv, + int tile_c, + bool return_lse, + mx::StreamOrDevice s) -> nb::object { + auto outs = mlx_kquant::sdpa_decode_gqa_cascade( + std::move(q), + std::move(k_shared), + std::move(v_shared), + std::move(k_priv), + std::move(v_priv), + scale, + starts, + splits_shared, + splits_priv, + tile_c, + return_lse, + s); + if (return_lse) { + return nb::make_tuple(outs[0], outs[1]); + } + return nb::cast(outs[0]); + }, + "q"_a, + "k_shared"_a, + "v_shared"_a, + "k_priv"_a, + "v_priv"_a, + "scale"_a, + "starts"_a = nb::none(), + "splits_shared"_a = 0, + "splits_priv"_a = 0, + "tile_c"_a = 0, + nb::kw_only(), + "return_lse"_a = false, + "stream"_a = nb::none(), + R"( + Fused shared-prefix (cascade) decode attention: every batch row + attends one COMMON prefix, stored once, plus its own private + suffix. The shared region is walked ONCE for all B*gqa query rows + on the matrix-unit row tile; the private region runs per row (with + optional left-pad ``starts``); both partial sets fold through a + single merge pass. Equivalent to ``sdpa_decode_gqa`` over the + concatenated KV, reading the prefix once instead of B times. + + Args: + q (array): queries [B, n_q_heads, 1, D], float16/bfloat16; + D in {64, 128, 256, 512}; gqa <= 16; B*gqa <= 64 (<= 32 at + D=512). + k_shared (array): shared prefix keys [1, n_kv_heads, P, D]. + v_shared (array): shared prefix values [1, n_kv_heads, P, D]. + k_priv (array): private suffix keys [B, n_kv_heads, Sp, D], + Sp >= 1. + v_priv (array): private suffix values [B, n_kv_heads, Sp, D]. + scale (float): query scale (typically 1/sqrt(D)). + starts (array, optional): int32 [B] per-row private-region key + start offsets (left-padded private suffixes). + splits_shared (int): shared-region split count; 0 = default. + splits_priv (int): private-region split count; 0 = default. + tile_c (int): private-pass staged tile height; 0 picks by + head_dim. + + Returns: + array: attention output [B, n_q_heads, 1, D]. With + ``return_lse=True``, a tuple ``(out, lse)``. )"); m.def( diff --git a/metal/kq_sdpa.metal b/metal/kq_sdpa.metal index e0572c6..72bea44 100644 --- a/metal/kq_sdpa.metal +++ b/metal/kq_sdpa.metal @@ -108,6 +108,14 @@ instantiate_kq_sdpa_gqa_p2(float16_t, 64, 32, 4) type, \ D) +instantiate_kq_sdpa_fa_verify(bfloat16_t, 64, 32) +instantiate_kq_sdpa_fa_verify(float16_t, 64, 32) +instantiate_kq_sdpa_fa_verify(bfloat16_t, 64, 64) +instantiate_kq_sdpa_fa_verify(float16_t, 64, 64) +instantiate_kq_sdpa_fa_verify(bfloat16_t, 128, 32) +instantiate_kq_sdpa_fa_verify(float16_t, 128, 32) +instantiate_kq_sdpa_fa_verify(bfloat16_t, 128, 64) +instantiate_kq_sdpa_fa_verify(float16_t, 128, 64) instantiate_kq_sdpa_fa_verify(bfloat16_t, 256, 32) instantiate_kq_sdpa_fa_verify(float16_t, 256, 32) instantiate_kq_sdpa_fa_verify(bfloat16_t, 256, 64) diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h index 343d708..a135afb 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h @@ -21,6 +21,12 @@ constant bool gqa_has_starts [[function_constant(4)]]; // 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)]]; +constant bool gqa_write_lse [[function_constant(6)]]; +// Cascade merge: a second partials set (the shared-prefix region, written by +// the fa row-tile pass in its folded [1, Hkv, B*gqa, splits2, D] layout) +// folds into the same online-softmax reduction as the primary set. Compiled +// out when false. +constant bool gqa_cascade [[function_constant(7)]]; template [[kernel]] void kq_sdpa_vector_2pass_1( @@ -1029,6 +1035,12 @@ template const device float* sinks [[buffer(3)]], device T* out [[buffer(4)]], const constant int& n_q_heads [[buffer(5)]], + device float* out_lse [[buffer(6)]], + const device float* partials2 [[buffer(7)]], + const device float* sums2 [[buffer(8)]], + const device float* maxs2 [[buffer(9)]], + const constant int& cascade_splits [[buffer(10)]], + const constant int& cascade_gqa [[buffer(11)]], uint3 tid [[threadgroup_position_in_grid]], uint3 tpg [[threadgroups_per_grid]], uint simd_lid [[thread_index_in_simdgroup]]) { @@ -1042,12 +1054,28 @@ template sums += base * gqa_splits; maxs += base * gqa_splits; + // Second set: fa folded layout, row = kv*(B*gqa) + b*gqa + g. + int splits2 = 0; + if (gqa_cascade) { + const int kv = head_idx / cascade_gqa; + const int g = head_idx % cascade_gqa; + const size_t row = ((size_t)kv * tpg.y + batch_idx) * cascade_gqa + g; + splits2 = cascade_splits; + partials2 += row * splits2 * D; + sums2 += row * splits2; + maxs2 += row * splits2; + } + threadgroup float ws[128]; + threadgroup float ws2[128]; float m = Limits::finite_min; for (int s = simd_lid; s < gqa_splits; s += 32) { m = max(m, maxs[s]); } + for (int s = simd_lid; s < splits2; s += 32) { + m = max(m, maxs2[s]); + } m = simd_max(m); if (gqa_has_sinks) { m = max(m, sinks[head_idx]); @@ -1059,6 +1087,11 @@ template ws[s] = w; denom += w * sums[s]; } + for (int s = simd_lid; s < splits2; s += 32) { + const float w = fast::exp(maxs2[s] - m); + ws2[s] = w; + denom += w * sums2[s]; + } denom = simd_sum(denom); if (gqa_has_sinks) { denom += fast::exp(sinks[head_idx] - m); @@ -1072,8 +1105,19 @@ template acc[e] += w * partials[s * D + e * 32 + simd_lid]; } } + for (int s = 0; s < splits2; s++) { + const float w = ws2[s]; + for (short e = 0; e < EPT; e++) { + acc[e] += w * partials2[s * D + e * 32 + simd_lid]; + } + } out += base * D; for (short e = 0; e < EPT; e++) { out[e * 32 + simd_lid] = static_cast(denom == 0 ? 0.0f : acc[e] / denom); } + // Natural-log softmax normalizer (sinks included when present): the + // cascade merge weight for combining disjoint key regions. + if (gqa_write_lse && simd_lid == 0) { + out_lse[base] = denom == 0 ? -INFINITY : (fast::log(denom) + m); + } } diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index 2fb0d67..ccfe9e6 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -64,6 +64,7 @@ rmsnorm_multi3, route_shed, sdpa_decode_gqa, + sdpa_decode_gqa_cascade, sdpa_fa_verify, sdpa_vector, shared_event_create, @@ -120,6 +121,7 @@ "rmsnorm_multi3", "route_shed", "sdpa_decode_gqa", + "sdpa_decode_gqa_cascade", "sdpa_fa_verify", "sdpa_vector", "shared_event_create", diff --git a/src/kquant.h b/src/kquant.h index 6c1d98c..41e1394 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -180,6 +180,25 @@ mx::array sdpa_decode_gqa( const std::optional& v_biases = std::nullopt, mx::StreamOrDevice s = {}); +// sdpa_decode_gqa returning {out, lse}: lse [B, n_q_heads, qL] float32 is +// the natural-log softmax normalizer per query row (sinks included when +// present) -- the merge weight for combining attention over disjoint key +// regions (shared-prefix cascade). +std::vector sdpa_decode_gqa_lse( + mx::array q, + mx::array k, + mx::array v, + float scale, + const std::optional& sinks = std::nullopt, + 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 // tile. The caller folds q [B, Hq, qL, D] -> [B, Hkv, G*qL, D] (kv-major // heads) and passes the original qL, so folded row r is causally clamped to @@ -199,6 +218,41 @@ mx::array sdpa_fa_verify( int splits = 0, mx::StreamOrDevice s = {}); +// Fused shared-prefix cascade decode attention. Every batch row attends one +// COMMON prefix (stored once, [1, Hkv, P, D]) plus its own private suffix +// ([B, Hkv, Sp, D], optional per-row `starts` for left padding). Internally: +// the fa row-tile pass walks the shared prefix ONCE for all B*gqa folded +// query rows, the decode pass covers the private region per row, and both +// partial sets fold through one merge pass (the splits machinery is the LSE +// merge). Equivalent to sdpa_decode_gqa over the concatenated KV, reading +// the prefix once instead of B times. q [B, Hq, 1, D]; B*gqa <= 64 (32 at +// head_dim 512); gqa <= 16; Sp >= 1. Returns {out} or {out, lse} with +// return_lse. Metal-only. +std::vector sdpa_decode_gqa_cascade( + mx::array q, + mx::array k_shared, + mx::array v_shared, + mx::array k_priv, + mx::array v_priv, + float scale, + const std::optional& starts = std::nullopt, + int splits_shared = 0, + int splits_priv = 0, + int tile_c = 0, + bool return_lse = false, + mx::StreamOrDevice s = {}); + +// sdpa_fa_verify returning {out, lse}: lse [B, Hkv, n_rows] float32 is the +// natural-log softmax normalizer per folded row (cascade merge weight). +std::vector sdpa_fa_verify_lse( + mx::array q, + mx::array k, + mx::array v, + float scale, + int q_len, + int splits = 0, + mx::StreamOrDevice s = {}); + // Fused MoE GLU gather on the MLX packed mxfp4 layout: gate and up expert // matvecs (sharing each activation load), expert biases, and the clamped // SwiGLU epilogue out = (min(g, limit) * sigmoid(alpha * g)) * (clip(u, @@ -705,14 +759,16 @@ class KQuantSDPAGQA : public mx::Primitive { int tile_c, bool has_sinks, bool has_starts, - bool has_kv_q8 = false) + bool has_kv_q8 = false, + bool return_lse = false) : mx::Primitive(stream), scale_(scale), splits_(splits), tile_c_(tile_c), has_sinks_(has_sinks), has_starts_(has_starts), - has_kv_q8_(has_kv_q8) {} + has_kv_q8_(has_kv_q8), + return_lse_(return_lse) {} void eval_cpu( const std::vector& inputs, @@ -736,6 +792,7 @@ class KQuantSDPAGQA : public mx::Primitive { bool has_sinks_; bool has_starts_; bool has_kv_q8_; + bool return_lse_; }; // Simdgroup-matrix FA verify attention (see sdpa_fa_verify). Inference-only. @@ -745,8 +802,13 @@ class KQuantSDPAFAVerify : public mx::Primitive { mx::Stream stream, float scale, int q_len, - int splits) - : mx::Primitive(stream), scale_(scale), q_len_(q_len), splits_(splits) {} + int splits, + bool return_lse = false) + : mx::Primitive(stream), + scale_(scale), + q_len_(q_len), + splits_(splits), + return_lse_(return_lse) {} void eval_cpu( const std::vector& inputs, @@ -767,6 +829,51 @@ class KQuantSDPAFAVerify : public mx::Primitive { float scale_; int q_len_; int splits_; + bool return_lse_; +}; + +// Fused shared-prefix cascade attention (see sdpa_decode_gqa_cascade). +// Inference-only. +class KQuantSDPACascade : public mx::Primitive { + public: + explicit KQuantSDPACascade( + mx::Stream stream, + float scale, + int splits_shared, + int splits_priv, + int tile_c, + bool has_starts, + bool return_lse) + : mx::Primitive(stream), + scale_(scale), + splits_shared_(splits_shared), + splits_priv_(splits_priv), + tile_c_(tile_c), + has_starts_(has_starts), + return_lse_(return_lse) {} + + 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 "KQuantSDPACascade"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + private: + float scale_; + int splits_shared_; + int splits_priv_; + int tile_c_; + bool has_starts_; + bool return_lse_; }; // Fused MoE GLU gather (see moe_glu_gather). Inference-only. diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 1514cf5..f8867c3 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -202,6 +202,10 @@ void KQuantSDPAGQA::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); + const bool write_lse = return_lse_; + if (write_lse) { + outputs[1].set_data(mx::allocator::malloc(outputs[1].nbytes())); + } const auto& q = inputs[0]; const auto& k = inputs[1]; @@ -255,11 +259,15 @@ void KQuantSDPAGQA::eval_gpu( bool has_sinks = sinks; bool has_starts = starts; bool has_kv_q8 = kv_q8; + bool has_lse = write_lse; + bool has_cascade = false; 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}, + {&has_lse, MTL::DataType::DataTypeBool, 6}, + {&has_cascade, MTL::DataType::DataTypeBool, 7}, }; // Pass 1: one threadgroup per (kv-head, batch, split); the whole GQA group @@ -329,8 +337,8 @@ void KQuantSDPAGQA::eval_gpu( // Grid z is the query axis. { std::string kname = "kq_sdpa_gqa_2pass_2_" + ts + "_" + std::to_string(D); - std::string hash = - kname + "_s" + std::to_string(splits) + (has_sinks ? "_k1" : "_k0"); + std::string hash = kname + "_s" + std::to_string(splits) + + (has_sinks ? "_k1" : "_k0") + (write_lse ? "_lse" : ""); auto kernel = kq_get_kernel(d, kname, hash, fc); ce.set_compute_pipeline_state(kernel); ce.set_input_array(partials, 0); @@ -341,6 +349,18 @@ void KQuantSDPAGQA::eval_gpu( ce.set_input_array(sinks ? inputs[3] : sums, 3); ce.set_output_array(out, 4); ce.set_bytes(n_q_heads, 5); + if (write_lse) { + ce.set_output_array(outputs[1], 6); + } else { + ce.set_input_array(sums, 6); + } + // Cascade compiled out: dummy second-set bindings. + const int czero = 0; + for (int i = 7; i <= 9; i++) { + ce.set_input_array(sums, i); + } + ce.set_bytes(czero, 10); + ce.set_bytes(czero, 11); MTL::Size group_dims(32, 1, 1); MTL::Size grid_dims(n_q_heads, B, qL); ce.dispatch_threadgroups(grid_dims, group_dims); @@ -354,6 +374,10 @@ void KQuantSDPAFAVerify::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); + const bool write_lse = return_lse_; + if (write_lse) { + outputs[1].set_data(mx::allocator::malloc(outputs[1].nbytes())); + } // 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. @@ -402,9 +426,13 @@ void KQuantSDPAFAVerify::eval_gpu( std::string ts = kq_type_string(q.dtype()); bool has_sinks = false; + bool has_lse = write_lse; + bool has_cascade = false; mx::metal::MTLFCList fc = { {&splits, MTL::DataType::DataTypeInt, 2}, {&has_sinks, MTL::DataType::DataTypeBool, 3}, + {&has_lse, MTL::DataType::DataTypeBool, 6}, + {&has_cascade, MTL::DataType::DataTypeBool, 7}, }; // Pass 1: one threadgroup per (kv-head, batch, split) streams its key @@ -449,7 +477,8 @@ void KQuantSDPAFAVerify::eval_gpu( // Pass 2: the shared kq_sdpa_gqa merge; grid z is the folded row axis. { std::string kname = "kq_sdpa_gqa_2pass_2_" + ts + "_" + std::to_string(D); - std::string hash = kname + "_s" + std::to_string(splits) + "_k0"; + std::string hash = kname + "_s" + std::to_string(splits) + "_k0" + + (write_lse ? "_lse" : ""); auto kernel = kq_get_kernel(d, kname, hash, fc); ce.set_compute_pipeline_state(kernel); ce.set_input_array(partials, 0); @@ -460,12 +489,225 @@ void KQuantSDPAFAVerify::eval_gpu( ce.set_input_array(sums, 3); ce.set_output_array(out, 4); ce.set_bytes(n_kv_heads, 5); + if (write_lse) { + ce.set_output_array(outputs[1], 6); + } else { + ce.set_input_array(sums, 6); + } + // Cascade compiled out: dummy second-set bindings. + const int czero = 0; + for (int i = 7; i <= 9; i++) { + ce.set_input_array(sums, i); + } + ce.set_bytes(czero, 10); + ce.set_bytes(czero, 11); MTL::Size group_dims(32, 1, 1); MTL::Size grid_dims(n_kv_heads, B, n_rows); ce.dispatch_threadgroups(grid_dims, group_dims); } } +void KQuantSDPACascade::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + auto& out = outputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + const bool write_lse = return_lse_; + if (write_lse) { + 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]; + const bool starts = has_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); + + int B = q.shape(0); + int n_q_heads = q.shape(1); + int D = q.shape(3); + int n_kv_heads = k_sh.shape(1); + int gqa_factor = n_q_heads / n_kv_heads; + int n_rows = B * gqa_factor; + int P = k_sh.shape(2); + int Sp = k_pr.shape(2); + int qL = 1; + float scale = scale_; + + int s_sh = splits_shared_; + if (s_sh == 0) { + s_sh = P <= 8192 ? 16 : P <= 24576 ? 32 : P <= 49152 ? 64 : 128; + } + int s_pr = splits_priv_; + if (s_pr == 0) { + s_pr = Sp <= 8192 ? 16 : Sp <= 24576 ? 32 : Sp <= 49152 ? 64 : 128; + } + + std::string ts = kq_type_string(q.dtype()); + auto& ce = mx::metal::get_command_encoder(s); + + // Private-region partials (decode layout) + shared-region partials (fa + // folded layout); one merge pass folds both. + mx::Shape p1_shape = {B, n_q_heads, qL, s_pr, D}; + mx::Shape r1_shape = {B, n_q_heads, qL, s_pr}; + array partials1(p1_shape, mx::float32, nullptr, {}); + array sums1(r1_shape, mx::float32, nullptr, {}); + array maxs1(r1_shape, mx::float32, nullptr, {}); + mx::Shape p2_shape = {1, n_kv_heads, n_rows, s_sh, D}; + mx::Shape r2_shape = {1, n_kv_heads, n_rows, s_sh}; + array partials2(p2_shape, mx::float32, nullptr, {}); + array sums2(r2_shape, mx::float32, nullptr, {}); + array maxs2(r2_shape, mx::float32, nullptr, {}); + for (array* a : {&partials1, &sums1, &maxs1, &partials2, &sums2, &maxs2}) { + a->set_data(mx::allocator::malloc(a->nbytes())); + ce.add_temporary(*a); + } + + // Pass 1a: private suffixes through the decode-gqa kernel (per-row grid, + // starts honored). + { + bool f = false; + bool has_starts = starts; + mx::metal::MTLFCList fc = { + {&s_pr, MTL::DataType::DataTypeInt, 2}, + {&f, MTL::DataType::DataTypeBool, 3}, + {&has_starts, MTL::DataType::DataTypeBool, 4}, + {&f, MTL::DataType::DataTypeBool, 5}, + }; + std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) + + "_c" + std::to_string(tile_c_); + std::string hash = kname + "_s" + std::to_string(s_pr) + + (has_starts ? "_st1" : "_st0") + "_casc"; + auto kernel = kq_get_kernel(d, kname, hash, fc); + const size_t tg = size_t(32) * gqa_factor; + if (tg > kernel->maxTotalThreadsPerThreadgroup()) { + throw std::runtime_error( + "[mlx_kquant.sdpa_decode_gqa_cascade] threadgroup of " + + std::to_string(tg) + " threads exceeds this GPU's pipeline limit."); + } + size_t k_head_stride = static_cast( + k_pr.shape(1) == 1 ? k_pr.strides(0) : k_pr.strides(1)); + size_t k_seq_stride = static_cast(k_pr.strides(2)); + size_t v_head_stride = static_cast( + v_pr.shape(1) == 1 ? v_pr.strides(0) : v_pr.strides(1)); + size_t v_seq_stride = static_cast(v_pr.strides(2)); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(q, 0); + ce.set_input_array(k_pr, 1); + ce.set_input_array(v_pr, 2); + ce.set_output_array(partials1, 3); + ce.set_output_array(sums1, 4); + ce.set_output_array(maxs1, 5); + ce.set_bytes(Sp, 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(qL, 12); + ce.set_input_array(starts ? inputs[6] : sums1, 13); + size_t zero = 0; + for (int i = 0; i < 4; i++) { + ce.set_input_array(sums1, 14 + i); + } + ce.set_bytes(zero, 18); + ce.set_bytes(zero, 19); + ce.set_bytes(zero, 20); + ce.set_bytes(zero, 21); + MTL::Size group_dims(32, gqa_factor, 1); + MTL::Size grid_dims(n_kv_heads, B, s_pr); + ce.dispatch_threadgroups(grid_dims, group_dims); + } + + // Pass 1b: the shared prefix through the fa row-tile kernel -- one KV walk + // serves all B*gqa folded rows. + { + bool f = false; + mx::metal::MTLFCList fc = { + {&s_sh, MTL::DataType::DataTypeInt, 2}, + {&f, MTL::DataType::DataTypeBool, 3}, + }; + const int bq = n_rows <= 32 ? 32 : 64; + std::string kname = "kq_sdpa_fa_verify_2pass_1_" + ts + "_" + + std::to_string(D) + "_bq" + std::to_string(bq); + std::string hash = kname + "_s" + std::to_string(s_sh) + "_casc"; + auto kernel = kq_get_kernel(d, kname, hash, fc); + const size_t tg = D == 512 ? 256 : (bq / 8) * 32; + if (tg > kernel->maxTotalThreadsPerThreadgroup()) { + throw std::runtime_error( + "[mlx_kquant.sdpa_decode_gqa_cascade] threadgroup of " + + std::to_string(tg) + " threads exceeds this GPU's pipeline limit."); + } + size_t k_head_stride = static_cast( + k_sh.shape(1) == 1 ? k_sh.strides(0) : k_sh.strides(1)); + size_t k_seq_stride = static_cast(k_sh.strides(2)); + size_t v_head_stride = static_cast( + v_sh.shape(1) == 1 ? v_sh.strides(0) : v_sh.strides(1)); + size_t v_seq_stride = static_cast(v_sh.strides(2)); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(qf, 0); + ce.set_input_array(k_sh, 1); + ce.set_input_array(v_sh, 2); + ce.set_output_array(partials2, 3); + ce.set_output_array(sums2, 4); + ce.set_output_array(maxs2, 5); + ce.set_bytes(P, 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(qL, 12); + ce.set_bytes(n_rows, 13); + MTL::Size group_dims(32, tg / 32, 1); + MTL::Size grid_dims(n_kv_heads, 1, s_sh); + ce.dispatch_threadgroups(grid_dims, group_dims); + } + + // Pass 2: one merge over both partial sets (the LSE merge). + { + bool f = false; + bool t = true; + bool has_lse = write_lse; + mx::metal::MTLFCList fc = { + {&s_pr, MTL::DataType::DataTypeInt, 2}, + {&f, MTL::DataType::DataTypeBool, 3}, + {&has_lse, MTL::DataType::DataTypeBool, 6}, + {&t, MTL::DataType::DataTypeBool, 7}, + }; + std::string kname = "kq_sdpa_gqa_2pass_2_" + ts + "_" + std::to_string(D); + std::string hash = kname + "_s" + std::to_string(s_pr) + "_k0_casc" + + std::to_string(s_sh) + (write_lse ? "_lse" : ""); + auto kernel = kq_get_kernel(d, kname, hash, fc); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(partials1, 0); + ce.set_input_array(sums1, 1); + ce.set_input_array(maxs1, 2); + ce.set_input_array(sums1, 3); + ce.set_output_array(out, 4); + ce.set_bytes(n_q_heads, 5); + if (write_lse) { + ce.set_output_array(outputs[1], 6); + } else { + ce.set_input_array(sums1, 6); + } + ce.set_input_array(partials2, 7); + ce.set_input_array(sums2, 8); + ce.set_input_array(maxs2, 9); + ce.set_bytes(s_sh, 10); + ce.set_bytes(gqa_factor, 11); + MTL::Size group_dims(32, 1, 1); + MTL::Size grid_dims(n_q_heads, B, qL); + ce.dispatch_threadgroups(grid_dims, group_dims); + } +} + #else // !_METAL_ void KQuantSDPA::eval_gpu( @@ -488,6 +730,13 @@ void KQuantSDPAFAVerify::eval_gpu( "[mlx_kquant.sdpa_fa_verify] requires a Metal build."); } +void KQuantSDPACascade::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.sdpa_decode_gqa_cascade] requires a Metal build."); +} + #endif void KQuantSDPA::eval_cpu( @@ -583,17 +832,22 @@ void KQuantSDPAGQA::eval_cpu( std::vector KQuantSDPAGQA::output_shapes( const std::vector& inputs) { - return {inputs[0].shape()}; + const auto& qs = inputs[0].shape(); + if (return_lse_) { + return {qs, {qs[0], qs[1], qs[2]}}; + } + return {qs}; } 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_kv_q8_ == o.has_kv_q8_; + has_kv_q8_ == o.has_kv_q8_ && return_lse_ == o.return_lse_; } -mx::array sdpa_decode_gqa( +static std::vector sdpa_decode_gqa_impl( + bool return_lse, mx::array q, mx::array k, mx::array v, @@ -741,19 +995,88 @@ mx::array sdpa_decode_gqa( } } + auto prim = std::make_shared( + s, + scale, + splits, + tile_c, + sinks.has_value(), + starts.has_value(), + kv_q8, + return_lse); 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(), - kv_q8), - std::move(inputs)); + if (return_lse) { + mx::Shape lse_shape = {q.shape(0), q.shape(1), q.shape(2)}; + return mx::array::make_arrays( + {std::move(out_shape), std::move(lse_shape)}, + {dt, mx::float32}, + std::move(prim), + std::move(inputs)); + } + return { + mx::array(std::move(out_shape), dt, std::move(prim), std::move(inputs))}; +} + +mx::array sdpa_decode_gqa( + mx::array q, + mx::array k, + mx::array v, + float scale, + const std::optional& sinks, + 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_) { + return sdpa_decode_gqa_impl( + false, + std::move(q), + std::move(k), + std::move(v), + scale, + sinks, + splits, + tile_c, + starts, + k_scales, + k_biases, + v_scales, + v_biases, + s_)[0]; +} + +std::vector sdpa_decode_gqa_lse( + mx::array q, + mx::array k, + mx::array v, + float scale, + const std::optional& sinks, + 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_) { + return sdpa_decode_gqa_impl( + true, + std::move(q), + std::move(k), + std::move(v), + scale, + sinks, + splits, + tile_c, + starts, + k_scales, + k_biases, + v_scales, + v_biases, + s_); } void KQuantSDPAFAVerify::eval_cpu( @@ -763,17 +1086,194 @@ void KQuantSDPAFAVerify::eval_cpu( "[mlx_kquant.sdpa_fa_verify] has no CPU implementation."); } +void KQuantSDPACascade::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.sdpa_decode_gqa_cascade] has no CPU implementation."); +} + +std::vector KQuantSDPACascade::output_shapes( + const std::vector& inputs) { + const auto& qs = inputs[0].shape(); + if (return_lse_) { + return {qs, {qs[0], qs[1], qs[2]}}; + } + return {qs}; +} + +bool KQuantSDPACascade::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return scale_ == o.scale_ && splits_shared_ == o.splits_shared_ && + splits_priv_ == o.splits_priv_ && tile_c_ == o.tile_c_ && + has_starts_ == o.has_starts_ && return_lse_ == o.return_lse_; +} + +std::vector sdpa_decode_gqa_cascade( + mx::array q, + mx::array k_shared, + mx::array v_shared, + mx::array k_priv, + mx::array v_priv, + float scale, + const std::optional& starts, + int splits_shared, + int splits_priv, + int tile_c, + bool return_lse, + mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.sdpa_decode_gqa_cascade] "; + + if (q.ndim() != 4 || k_shared.ndim() != 4 || v_shared.ndim() != 4 || + k_priv.ndim() != 4 || v_priv.ndim() != 4) { + throw std::invalid_argument( + std::string(op) + "q, k, v must be 4-D [B, heads, L, D]."); + } + int B = q.shape(0); + int n_q_heads = q.shape(1); + int D = q.shape(3); + int n_kv_heads = k_shared.shape(1); + auto dt = q.dtype(); + if (dt != mx::float16 && dt != mx::bfloat16) { + throw std::invalid_argument( + std::string(op) + "q must be float16 or bfloat16."); + } + for (const auto& a : {k_shared, v_shared, k_priv, v_priv}) { + if (a.dtype() != dt || a.shape(3) != D) { + throw std::invalid_argument( + std::string(op) + "k/v must share q's dtype and head_dim."); + } + } + if (D != 64 && D != 128 && D != 256 && D != 512) { + throw std::invalid_argument( + std::string(op) + "head_dim must be 64, 128, 256 or 512."); + } + if (q.shape(2) != 1) { + throw std::invalid_argument(std::string(op) + "q_len must be 1."); + } + if (k_shared.shape(0) != 1 || v_shared.shape(0) != 1 || + v_shared.shape(1) != n_kv_heads || + v_shared.shape(2) != k_shared.shape(2)) { + throw std::invalid_argument( + std::string(op) + + "shared k/v must be [1, n_kv_heads, P, D] with matching P."); + } + if (k_priv.shape(0) != B || v_priv.shape(0) != B || + k_priv.shape(1) != n_kv_heads || v_priv.shape(1) != n_kv_heads || + v_priv.shape(2) != k_priv.shape(2)) { + throw std::invalid_argument( + std::string(op) + + "private k/v must be [B, n_kv_heads, Sp, D] with matching Sp."); + } + if (k_priv.shape(2) < 1 || k_shared.shape(2) < 1) { + throw std::invalid_argument( + std::string(op) + "both key regions must be non-empty."); + } + if (n_q_heads % n_kv_heads != 0) { + throw std::invalid_argument( + std::string(op) + "q heads must be a multiple of kv heads."); + } + int gqa_factor = n_q_heads / n_kv_heads; + if (gqa_factor > 16) { + throw std::invalid_argument(std::string(op) + "gqa factor must be <= 16."); + } + int n_rows = B * gqa_factor; + int max_rows = D == 512 ? 32 : 64; + if (n_rows > max_rows) { + throw std::invalid_argument( + std::string(op) + "B * gqa must be <= " + std::to_string(max_rows) + + " at head_dim " + std::to_string(D) + "."); + } + if (tile_c == 0) { + tile_c = D <= 128 ? 32 : D == 256 ? 16 : 8; + } + const bool tile_ok = (D <= 128 && (tile_c == 32 || tile_c == 16)) || + (D == 256 && (tile_c == 16 || tile_c == 8)) || (D == 512 && tile_c == 8); + if (!tile_ok) { + throw std::invalid_argument( + std::string(op) + "tile_c not instantiated for this head_dim."); + } + if (splits_shared < 0 || splits_shared > 128 || splits_priv < 0 || + splits_priv > 128) { + throw std::invalid_argument( + std::string(op) + "splits must be in [0, 128]."); + } + + auto q_c = mx::contiguous(q, false, s); + // kv-head-major fold for the shared row-tile pass. + auto q_folded = mx::contiguous( + mx::reshape( + mx::transpose( + mx::reshape(q_c, {B, n_kv_heads, gqa_factor, D}, s), + {1, 0, 2, 3}, + s), + {1, n_kv_heads, n_rows, D}, + s), + false, + s); + auto contig_kv = [&](const mx::array& a) { + return a.strides().back() == 1 ? a : mx::contiguous(a, false, s); + }; + + std::vector inputs = { + std::move(q_c), + std::move(q_folded), + contig_kv(k_shared), + contig_kv(v_shared), + contig_kv(k_priv), + contig_kv(v_priv)}; + if (starts.has_value()) { + auto st = *starts; + if (st.size() != static_cast(B)) { + throw std::invalid_argument( + std::string(op) + "starts must have one element per batch row."); + } + if (st.dtype() != mx::int32) { + throw std::invalid_argument(std::string(op) + "starts must be int32."); + } + st = mx::reshape(st, {B}, s); + inputs.push_back(mx::contiguous(st, false, s)); + } + + auto prim = std::make_shared( + s, + scale, + splits_shared, + splits_priv, + tile_c, + starts.has_value(), + return_lse); + auto out_shape = q.shape(); + if (return_lse) { + mx::Shape lse_shape = {B, n_q_heads, 1}; + return mx::array::make_arrays( + {std::move(out_shape), std::move(lse_shape)}, + {dt, mx::float32}, + std::move(prim), + std::move(inputs)); + } + return { + mx::array(std::move(out_shape), dt, std::move(prim), std::move(inputs))}; +} + std::vector KQuantSDPAFAVerify::output_shapes( const std::vector& inputs) { - return {inputs[0].shape()}; + const auto& qs = inputs[0].shape(); + if (return_lse_) { + return {qs, {qs[0], qs[1], qs[2]}}; + } + return {qs}; } bool KQuantSDPAFAVerify::is_equivalent(const mx::Primitive& other) const { const auto& o = static_cast(other); - return scale_ == o.scale_ && q_len_ == o.q_len_ && splits_ == o.splits_; + return scale_ == o.scale_ && q_len_ == o.q_len_ && splits_ == o.splits_ && + return_lse_ == o.return_lse_; } -mx::array sdpa_fa_verify( +static std::vector sdpa_fa_verify_impl( + bool return_lse, mx::array q, mx::array k, mx::array v, @@ -788,9 +1288,11 @@ mx::array sdpa_fa_verify( "[mlx_kquant.sdpa_fa_verify] q, k, v must be 4-D [B, heads, L, D]."); } int D = q.shape(-1); - if ((D != 256 && D != 512) || k.shape(-1) != D || v.shape(-1) != D) { + if ((D != 64 && D != 128 && D != 256 && D != 512) || k.shape(-1) != D || + v.shape(-1) != D) { throw std::invalid_argument( - "[mlx_kquant.sdpa_fa_verify] only head_dim 256 or 512 is supported."); + "[mlx_kquant.sdpa_fa_verify] only head_dim 64, 128, 256 or 512 is " + "supported."); } auto dt = q.dtype(); if (dt != mx::float16 && dt != mx::bfloat16) { @@ -819,8 +1321,8 @@ mx::array sdpa_fa_verify( "[mlx_kquant.sdpa_fa_verify] q_len must be in [1, 8]."); } int n_rows = q.shape(2); - // The 64-row tile exists only for head_dim 256; the 512 d-split kernel is - // fixed at the 32-row tile. + // The 64-row tile exists for the register-resident head dims (64-256); + // the 512 d-split kernel is fixed at the 32-row tile. int max_rows = D == 512 ? 32 : 64; if (n_rows < q_len || n_rows > max_rows || n_rows % q_len != 0) { throw std::invalid_argument( @@ -844,12 +1346,52 @@ mx::array sdpa_fa_verify( auto k_c = k.strides().back() == 1 ? k : mx::contiguous(k, false, s); auto v_c = v.strides().back() == 1 ? v : mx::contiguous(v, false, s); + auto prim = + std::make_shared(s, scale, q_len, splits, return_lse); auto out_shape = q_c.shape(); - return mx::array( - std::move(out_shape), - dt, - std::make_shared(s, scale, q_len, splits), - {std::move(q_c), std::move(k_c), std::move(v_c)}); + std::vector inputs = { + std::move(q_c), std::move(k_c), std::move(v_c)}; + if (return_lse) { + mx::Shape lse_shape = {out_shape[0], out_shape[1], out_shape[2]}; + return mx::array::make_arrays( + {std::move(out_shape), std::move(lse_shape)}, + {dt, mx::float32}, + std::move(prim), + std::move(inputs)); + } + return { + mx::array(std::move(out_shape), dt, std::move(prim), std::move(inputs))}; +} + +mx::array sdpa_fa_verify( + mx::array q, + mx::array k, + mx::array v, + float scale, + int q_len, + int splits, + mx::StreamOrDevice s_) { + return sdpa_fa_verify_impl( + false, + std::move(q), + std::move(k), + std::move(v), + scale, + q_len, + splits, + s_)[0]; +} + +std::vector sdpa_fa_verify_lse( + mx::array q, + mx::array k, + mx::array v, + float scale, + int q_len, + int splits, + mx::StreamOrDevice s_) { + return sdpa_fa_verify_impl( + true, std::move(q), std::move(k), std::move(v), scale, q_len, splits, s_); } } // namespace mlx_kquant diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 8c54a1a..b388734 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -390,7 +390,7 @@ def _check_fa(D, qL, kL, dtype, Hkv=4, G=6, strided=False, splits=0): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) -@pytest.mark.parametrize("D", [256, 512]) +@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): _check_fa(D, qL, kL=4096, dtype=dtype, G=4) @@ -402,7 +402,7 @@ def test_sdpa_fa_verify_qwen_geometry(): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) -@pytest.mark.parametrize("D", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_decode(D, dtype): # 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; @@ -415,39 +415,39 @@ def test_sdpa_fa_verify_gemma_geometry(): _check_fa(512, 4, kL=8192, dtype=mx.bfloat16, Hkv=4, G=8) -@pytest.mark.parametrize("D", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_full_tile(D): # 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", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_partial_warp(D): # 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", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_strided_unaligned(D, dtype): # 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", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_short_kv(D): # 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", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_min_kv(D): # 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", [256, 512]) +@pytest.mark.parametrize("D", [64, 128, 256, 512]) def test_sdpa_fa_verify_causal_split_straddle(D): # 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 @@ -473,6 +473,15 @@ def test_sdpa_fa_verify_bq64_strided_kv(): _check_fa(256, 4, kL=3071, dtype=mx.bfloat16, Hkv=2, G=16, strided=True, splits=16) +@pytest.mark.parametrize("D", [64, 128]) +@pytest.mark.parametrize("G", [32, 64]) +def test_sdpa_fa_verify_cascade_fold(D, G): + # 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 + _check_fa(D, 1, kL=6144, dtype=mx.bfloat16, Hkv=2, G=G) + + def test_sdpa_fa_verify_lazy_strided_q(): # Regression: an UNEVALUATED strided q view (here a chunk of a folded # tile) must not be trusted as row-contiguous at op-build time; the op @@ -509,3 +518,148 @@ def main(): if __name__ == "__main__": sys.exit(main()) + + +def test_sdpa_return_lse_matches_reference(): + B, Hq, Hkv, D, S = 2, 32, 8, 128, 3001 + q, k, v = _make(B, Hq, Hkv, 1, S, D, mx.bfloat16, seed=11, strided=False) + scale = 1.0 / (D**0.5) + o, lse = kq.sdpa_decode_gqa(q, k, v, scale, return_lse=True) + o0 = kq.sdpa_decode_gqa(q, k, v, scale) + kr = mx.repeat(k, Hq // Hkv, axis=1).astype(mx.float32) + s_ref = (q.astype(mx.float32) * scale) @ kr.swapaxes(-1, -2) + lse_ref = mx.logsumexp(s_ref[:, :, 0, :], axis=-1)[..., None] + _eval_or_skip(o, lse, o0, lse_ref) + assert lse.shape == (B, Hq, 1) and lse.dtype == mx.float32 + assert float(mx.abs(o - o0).max()) == 0.0 + assert float(mx.abs(lse - lse_ref).max()) < 2e-2 + + +def test_sdpa_fa_verify_return_lse(): + q, k, v = _make(1, 8, 8, 32, 2048, 128, mx.bfloat16, seed=12, strided=False) + scale = 1.0 / (128**0.5) + o, lse = kq.sdpa_fa_verify(q, k, v, scale, 1, return_lse=True) + o0 = kq.sdpa_fa_verify(q, k, v, scale, 1) + s_ref = (q.astype(mx.float32) * scale) @ k.astype(mx.float32).swapaxes(-1, -2) + lse_ref = mx.logsumexp(s_ref, axis=-1) + _eval_or_skip(o, lse, o0, lse_ref) + assert lse.shape == (1, 8, 32) and lse.dtype == mx.float32 + assert float(mx.abs(o - o0).max()) == 0.0 + assert float(mx.abs(lse - lse_ref).max()) < 2e-2 + + +def test_sdpa_cascade_lse_merge(): + # shared-prefix cascade: fa fold over the shared block plus a per-row + # private call, LSE-merged, must match one call over the concatenated KV + B, Hq, Hkv, D = 4, 32, 8, 128 + gqa = Hq // Hkv + P, Sp = 4096, 384 + scale = 1.0 / (D**0.5) + q, k_sh, v_sh = _make(1, Hq, Hkv, 1, P, D, mx.bfloat16, seed=13, strided=False) + q = mx.random.normal((B, Hq, 1, D)).astype(mx.bfloat16) * 0.5 + _, k_pr, v_pr = _make(B, Hq, Hkv, 1, Sp, D, mx.bfloat16, seed=14, strided=False) + k_full = mx.concatenate([mx.broadcast_to(k_sh, (B, Hkv, P, D)), k_pr], axis=2) + v_full = mx.concatenate([mx.broadcast_to(v_sh, (B, Hkv, P, D)), v_pr], axis=2) + ref = kq.sdpa_decode_gqa(q, mx.contiguous(k_full), mx.contiguous(v_full), scale) + qf = mx.contiguous( + q[:, :, 0, :] + .reshape(B, Hkv, gqa, D) + .transpose(1, 0, 2, 3) + .reshape(1, Hkv, B * gqa, D) + ) + o_f, l_f = kq.sdpa_fa_verify(qf, k_sh, v_sh, scale, 1, return_lse=True) + o_sh = o_f.reshape(1, Hkv, B, gqa, D)[0].transpose(1, 0, 2, 3).reshape(B, Hq, 1, D) + l_sh = l_f.reshape(1, Hkv, B, gqa)[0].transpose(1, 0, 2).reshape(B, Hq, 1) + o_pr, l_pr = kq.sdpa_decode_gqa(q, k_pr, v_pr, scale, return_lse=True) + m = mx.maximum(l_sh, l_pr) + w_sh = mx.exp(l_sh - m)[..., None] + w_pr = mx.exp(l_pr - m)[..., None] + merged = ( + (o_sh.astype(mx.float32) * w_sh + o_pr.astype(mx.float32) * w_pr) + / (w_sh + w_pr) + ).astype(q.dtype) + _eval_or_skip(merged, ref) + rel = _rel(merged, ref) + assert rel < REL_BOUND[mx.bfloat16], f"cascade merge rel {rel:.3e}" + + +@pytest.mark.parametrize("D", [64, 128, 256]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) +def test_sdpa_cascade_fused_matches_concat(D, dtype): + # fused cascade op == one sdpa_decode_gqa call over the concatenated KV + B, Hq, Hkv = 4, 32, 8 + P, Sp = 3071, 257 + scale = 1.0 / (D**0.5) + _, k_sh, v_sh = _make(1, Hq, Hkv, 1, P, D, dtype, seed=21, strided=False) + q, k_pr, v_pr = _make(B, Hq, Hkv, 1, Sp, D, dtype, seed=22, strided=False) + k_full = mx.contiguous( + mx.concatenate([mx.broadcast_to(k_sh, (B, Hkv, P, D)), k_pr], axis=2) + ) + v_full = mx.contiguous( + mx.concatenate([mx.broadcast_to(v_sh, (B, Hkv, P, D)), v_pr], axis=2) + ) + ref = kq.sdpa_decode_gqa(q, k_full, v_full, scale) + got = kq.sdpa_decode_gqa_cascade(q, k_sh, v_sh, k_pr, v_pr, scale) + _eval_or_skip(got, ref) + rel = _rel(got, ref) + print(f" [cascade] fused D={D} {dtype}: rel={rel:.3e}") + assert rel < REL_BOUND[dtype], f"fused cascade rel {rel:.3e}" + + +def test_sdpa_cascade_fused_starts(): + # per-row private left-pad: keys below starts[b] in the PRIVATE region + # are excluded; the shared prefix is always fully attended + B, Hq, Hkv, D = 4, 32, 8, 128 + P, Sp = 2048, 384 + scale = 1.0 / (D**0.5) + _, k_sh, v_sh = _make(1, Hq, Hkv, 1, P, D, mx.bfloat16, seed=23, strided=False) + q, k_pr, v_pr = _make(B, Hq, Hkv, 1, Sp, D, mx.bfloat16, seed=24, strided=False) + starts = mx.array([0, 7, 133, Sp - 1], dtype=mx.int32) + got = kq.sdpa_decode_gqa_cascade(q, k_sh, v_sh, k_pr, v_pr, scale, starts=starts) + k_full = mx.concatenate([mx.broadcast_to(k_sh, (B, Hkv, P, D)), k_pr], axis=2) + v_full = mx.concatenate([mx.broadcast_to(v_sh, (B, Hkv, P, D)), v_pr], axis=2) + pos = mx.arange(P + Sp)[None, :] + keep = pos >= (starts[:, None] + P) + keep = mx.logical_or(pos < P, keep) + bias = mx.where(keep, mx.zeros(keep.shape), mx.full(keep.shape, -mx.inf)) + kr = mx.repeat(k_full, Hq // Hkv, axis=1).astype(mx.float32) + vr = mx.repeat(v_full, Hq // Hkv, axis=1).astype(mx.float32) + s_ref = (q.astype(mx.float32) * scale) @ kr.swapaxes(-1, -2) + s_ref = s_ref + bias[:, None, None, :] + ref = mx.softmax(s_ref, axis=-1) @ vr + _eval_or_skip(got, ref) + rel = _rel(got.astype(mx.float32), ref) + assert rel < REL_BOUND[mx.bfloat16], f"cascade starts rel {rel:.3e}" + + +def test_sdpa_cascade_fused_return_lse(): + B, Hq, Hkv, D = 2, 16, 8, 128 + P, Sp = 1023, 65 + scale = 1.0 / (D**0.5) + _, k_sh, v_sh = _make(1, Hq, Hkv, 1, P, D, mx.bfloat16, seed=25, strided=False) + q, k_pr, v_pr = _make(B, Hq, Hkv, 1, Sp, D, mx.bfloat16, seed=26, strided=False) + o, lse = kq.sdpa_decode_gqa_cascade( + q, k_sh, v_sh, k_pr, v_pr, scale, return_lse=True + ) + o0 = kq.sdpa_decode_gqa_cascade(q, k_sh, v_sh, k_pr, v_pr, scale) + k_full = mx.concatenate([mx.broadcast_to(k_sh, (B, Hkv, P, D)), k_pr], axis=2) + kr = mx.repeat(k_full, Hq // Hkv, axis=1).astype(mx.float32) + s_ref = (q.astype(mx.float32) * scale) @ kr.swapaxes(-1, -2) + lse_ref = mx.logsumexp(s_ref[:, :, 0, :], axis=-1)[..., None] + _eval_or_skip(o, lse, o0, lse_ref) + assert lse.shape == (B, Hq, 1) and lse.dtype == mx.float32 + assert float(mx.abs(o - o0).max()) == 0.0 + assert float(mx.abs(lse - lse_ref).max()) < 2e-2 + + +def test_sdpa_cascade_fused_validation(): + B, Hq, Hkv, D = 2, 16, 8, 128 + scale = 1.0 / (D**0.5) + _, k_sh, v_sh = _make(1, Hq, Hkv, 1, 512, D, mx.bfloat16, seed=27, strided=False) + q, k_pr, v_pr = _make(B, Hq, Hkv, 1, 64, D, mx.bfloat16, seed=28, strided=False) + with pytest.raises(ValueError): + kq.sdpa_decode_gqa_cascade(q, k_pr, v_pr, k_pr, v_pr, scale) # shared B != 1 + with pytest.raises(ValueError): + kq.sdpa_decode_gqa_cascade( + q, k_sh, v_sh, k_pr[:, :, :0, :], v_pr[:, :, :0, :], scale + ) # empty private region