From 812fee5240a3c9cdf1a2ace3c0da241f41791713 Mon Sep 17 00:00:00 2001 From: Duhyeon Date: Sat, 8 Aug 2026 21:52:05 +0900 Subject: [PATCH 1/2] Read each K/V byte once in gqa-8 decode attention --- .../scaled_dot_product_attention.metal | 12 ++ mlx/backend/metal/kernels/sdpa_vector.h | 143 ++++++++++++++++++ .../metal/scaled_dot_product_attention.cpp | 47 ++++++ python/tests/test_fast_sdpa.py | 14 ++ 4 files changed, 216 insertions(+) diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 44ccb834e9..878831b986 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -28,12 +28,24 @@ using namespace metal; qk_dim, \ value_dim) +#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \ + sdpa_vector_2pass_1_gqa, \ + type, \ + qk_dim, \ + value_dim, \ + 8, \ + hpt) + #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ instantiate_sdpa_vector(type, 192, 128) \ instantiate_sdpa_vector(type, 256, 256) \ + instantiate_sdpa_vector_gqa(type, 64, 64, 8) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 4) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ diff --git a/mlx/backend/metal/kernels/sdpa_vector.h b/mlx/backend/metal/kernels/sdpa_vector.h index 1eec72be31..3f40dbd7c0 100644 --- a/mlx/backend/metal/kernels/sdpa_vector.h +++ b/mlx/backend/metal/kernels/sdpa_vector.h @@ -317,6 +317,149 @@ template } } +// Duplication-free variant for high gqa_factor decode: each simdgroup owns a +// contiguous token sub-chunk and computes HPT of its group's query heads, so +// each K/V byte is read G / HPT times instead of G times. Single-token +// queries without mask or sinks only; the partials layout matches +// sdpa_vector_2pass_2. +template +[[kernel]] void sdpa_vector_2pass_1_gqa( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device T* out [[buffer(3)]], + device float* sums [[buffer(4)]], + device float* maxs [[buffer(5)]], + const constant int& N [[buffer(7)]], + const constant size_t& k_head_stride [[buffer(8)]], + const constant size_t& k_seq_stride [[buffer(9)]], + const constant size_t& v_head_stride [[buffer(10)]], + const constant size_t& v_seq_stride [[buffer(11)]], + const constant float& scale [[buffer(12)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tpg [[threadgroups_per_grid]], + uint3 tidtg [[thread_position_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int BD = 32; + constexpr int qk_per_thread = D / BD; + constexpr int v_per_thread = V / BD; + constexpr int NT = G / HPT; + + typedef float U; + + const int kv_head_idx = tid.x; + const int batch_idx = tid.y; + const int block_idx = tid.z; + const int blocks = tpg.z; + const int g = tidtg.y; + const int cchunk = g / NT; + const int h0 = (g % NT) * HPT; + const int num_kv_heads = tpg.x; + const int num_q_heads = num_kv_heads * G; + const int base_head = batch_idx * num_q_heads + kv_head_idx * G; + + const int chunk = (N + blocks - 1) / blocks; + const int kstart = block_idx * chunk; + const int kend = min(N, kstart + chunk); + const int sub = (chunk + HPT - 1) / HPT; + const int s0 = kstart + cchunk * sub; + const int s1 = min(kend, s0 + sub); + + const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride + + simd_lid * qk_per_thread; + const device T* vp = values + kv_head_idx * v_head_stride + + s0 * v_seq_stride + simd_lid * v_per_thread; + + U q[HPT][qk_per_thread]; + for (int j = 0; j < HPT; j++) { + const device T* qp = + queries + (base_head + h0 + j) * D + simd_lid * qk_per_thread; + for (int i = 0; i < qk_per_thread; i++) { + q[j][i] = static_cast(scale) * qp[i]; + } + } + + U max_score[HPT]; + U sum_exp_score[HPT]; + U o[HPT][v_per_thread]; + for (int j = 0; j < HPT; j++) { + max_score[j] = Limits::finite_min; + sum_exp_score[j] = 0; + for (int i = 0; i < v_per_thread; i++) { + o[j][i] = 0; + } + } + + for (int t = s0; t < s1; t++) { + U kr[qk_per_thread]; + U vr[v_per_thread]; + for (int i = 0; i < qk_per_thread; i++) { + kr[i] = kp[i]; + } + for (int i = 0; i < v_per_thread; i++) { + vr[i] = vp[i]; + } + kp += k_seq_stride; + vp += v_seq_stride; + for (int j = 0; j < HPT; j++) { + U score = 0; + for (int i = 0; i < qk_per_thread; i++) { + score += q[j][i] * kr[i]; + } + score = simd_sum(score); + U new_max = max(max_score[j], score); + U factor = fast::exp(max_score[j] - new_max); + U exp_score = fast::exp(score - new_max); + max_score[j] = new_max; + sum_exp_score[j] = sum_exp_score[j] * factor + exp_score; + for (int i = 0; i < v_per_thread; i++) { + o[j][i] = o[j][i] * factor + exp_score * vr[i]; + } + } + } + + threadgroup U o_sh[G * HPT * V]; + threadgroup U se_sh[G * HPT]; + threadgroup U mx_sh[G * HPT]; + for (int j = 0; j < HPT; j++) { + int slot = (h0 + j) * HPT + cchunk; + U inv = sum_exp_score[j] > 0 ? 1 / sum_exp_score[j] : 0; + for (int i = 0; i < v_per_thread; i++) { + o_sh[slot * V + simd_lid * v_per_thread + i] = o[j][i] * inv; + } + if (simd_lid == 0) { + se_sh[slot] = sum_exp_score[j]; + mx_sh[slot] = max_score[j]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + U gmax = Limits::finite_min; + for (int s = 0; s < HPT; s++) { + gmax = max(gmax, mx_sh[g * HPT + s]); + } + U denom = 0; + U acc[v_per_thread] = {0}; + for (int s = 0; s < HPT; s++) { + U w = se_sh[g * HPT + s] * fast::exp(mx_sh[g * HPT + s] - gmax); + denom += w; + for (int i = 0; i < v_per_thread; i++) { + acc[i] += w * o_sh[(g * HPT + s) * V + simd_lid * v_per_thread + i]; + } + } + + const int o_offset = base_head + g; + device T* op = + out + o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread; + for (int i = 0; i < v_per_thread; i++) { + op[i] = static_cast(acc[i]); + } + if (simd_lid == 0) { + sums[o_offset * blocks + block_idx] = denom; + maxs[o_offset * blocks + block_idx] = gmax; + } +} + template [[kernel]] void sdpa_vector_2pass_2( const device T* partials [[buffer(0)]], diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index cb4af8523e..143b730929 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -506,6 +506,53 @@ void sdpa_vector_2pass( compute_encoder.add_temporary(sums); compute_encoder.add_temporary(maxs); + if (!mask && !sinks && q.shape(2) == 1 && gqa_factor == 8 && + q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) && + N >= 8192) { + std::string gname; + gname.reserve(64); + gname += "sdpa_vector_2pass_1_gqa_"; + gname += get_type_string(q.dtype()); + gname += "_"; + gname += std::to_string(q.shape(-1)); + gname += "_"; + gname += std::to_string(v.shape(-1)); + auto kernel = d.get_kernel(gname); + check_kernel_threadgroup_size(kernel, group_dims, gname); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_output_array(intermediate, 3); + compute_encoder.set_output_array(sums, 4); + compute_encoder.set_output_array(maxs, 5); + compute_encoder.set_bytes(N, 7); + compute_encoder.set_bytes(k_head_stride, 8); + compute_encoder.set_bytes(k_seq_stride, 9); + compute_encoder.set_bytes(v_head_stride, 10); + compute_encoder.set_bytes(v_seq_stride, 11); + compute_encoder.set_bytes(scale, 12); + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + + kname.clear(); + kname = "sdpa_vector_2pass_2_"; + kname += get_type_string(q.dtype()); + kname += "_"; + kname += std::to_string(v.shape(-1)); + kernel = d.get_kernel(kname); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(intermediate, 0); + compute_encoder.set_input_array(sums, 1); + compute_encoder.set_input_array(maxs, 2); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(blocks, 4); + group_dims = MTL::Size(1024, 1, 1); + grid_dims = MTL::Size(q.shape(0) * q.shape(1), q.shape(2), 1); + check_kernel_threadgroup_size(kernel, group_dims, kname); + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + return; + } + bool has_mask = mask.has_value(); bool bool_mask = has_mask && (*mask).dtype() == bool_; bool float_mask = has_mask && !bool_mask; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 2418997bc8..15f75a3d81 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -248,6 +248,20 @@ def test_sdpa_vector(self): ) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_vector_gqa_long(self): + scale = 1.0 + mx.random.seed(0) + for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64)]: + for L in [8192, 8201]: + q = 5e-1 * mx.random.normal(shape=(1, Nq, 1, D)) + k = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + v = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + kr = mx.repeat(k, Nq // Nkv, axis=1) + vr = mx.repeat(v, Nq // Nkv, axis=1) + ref = mlx_primitives_sdpa(q, kr, vr, scale) + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_fully_masked(self): Lkv = 8 mask = mx.array(False) From d1ab9ded8ee0b4172f0073cb0418bd3d0ac33c42 Mon Sep 17 00:00:00 2001 From: Duhyeon Date: Tue, 18 Aug 2026 16:20:38 +0900 Subject: [PATCH 2/2] Select the gqa kernel by name in the shared dispatch --- .../metal/scaled_dot_product_attention.cpp | 55 +++---------------- 1 file changed, 7 insertions(+), 48 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 143b730929..e49530a230 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -430,7 +430,13 @@ void sdpa_vector_2pass( // Set the kernel name std::string kname; kname.reserve(64); - kname += "sdpa_vector_2pass_1_"; + kname += "sdpa_vector_2pass_1"; + if (!mask && !sinks && q.shape(2) == 1 && q.shape(1) == 8 * k.shape(1) && + q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) && + k.shape(2) >= 8192) { + kname += "_gqa"; + } + kname += "_"; kname += get_type_string(q.dtype()); kname += "_"; kname += std::to_string(q.shape(-1)); @@ -506,53 +512,6 @@ void sdpa_vector_2pass( compute_encoder.add_temporary(sums); compute_encoder.add_temporary(maxs); - if (!mask && !sinks && q.shape(2) == 1 && gqa_factor == 8 && - q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) && - N >= 8192) { - std::string gname; - gname.reserve(64); - gname += "sdpa_vector_2pass_1_gqa_"; - gname += get_type_string(q.dtype()); - gname += "_"; - gname += std::to_string(q.shape(-1)); - gname += "_"; - gname += std::to_string(v.shape(-1)); - auto kernel = d.get_kernel(gname); - check_kernel_threadgroup_size(kernel, group_dims, gname); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(v, 2); - compute_encoder.set_output_array(intermediate, 3); - compute_encoder.set_output_array(sums, 4); - compute_encoder.set_output_array(maxs, 5); - compute_encoder.set_bytes(N, 7); - compute_encoder.set_bytes(k_head_stride, 8); - compute_encoder.set_bytes(k_seq_stride, 9); - compute_encoder.set_bytes(v_head_stride, 10); - compute_encoder.set_bytes(v_seq_stride, 11); - compute_encoder.set_bytes(scale, 12); - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - - kname.clear(); - kname = "sdpa_vector_2pass_2_"; - kname += get_type_string(q.dtype()); - kname += "_"; - kname += std::to_string(v.shape(-1)); - kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(intermediate, 0); - compute_encoder.set_input_array(sums, 1); - compute_encoder.set_input_array(maxs, 2); - compute_encoder.set_output_array(out, 3); - compute_encoder.set_bytes(blocks, 4); - group_dims = MTL::Size(1024, 1, 1); - grid_dims = MTL::Size(q.shape(0) * q.shape(1), q.shape(2), 1); - check_kernel_threadgroup_size(kernel, group_dims, kname); - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - return; - } - bool has_mask = mask.has_value(); bool bool_mask = has_mask && (*mask).dtype() == bool_; bool float_mask = has_mask && !bool_mask;