diff --git a/mlx/backend/metal/kernels/sdpa_vector.h b/mlx/backend/metal/kernels/sdpa_vector.h index 1eec72be31..ad343e8df2 100644 --- a/mlx/backend/metal/kernels/sdpa_vector.h +++ b/mlx/backend/metal/kernels/sdpa_vector.h @@ -11,6 +11,7 @@ constant bool bool_mask [[function_constant(23)]]; constant bool float_mask [[function_constant(24)]]; constant bool has_sinks [[function_constant(25)]]; constant int blocks [[function_constant(26)]]; +constant int n_per_simd [[function_constant(27)]]; template [[kernel]] void sdpa_vector( @@ -184,6 +185,7 @@ template device T* out [[buffer(3)]], device float* sums [[buffer(4)]], device float* maxs [[buffer(5)]], + const constant int& q_seq_len [[buffer(6)]], const constant int& N [[buffer(7)]], const constant size_t& k_head_stride [[buffer(8)]], const constant size_t& k_seq_stride [[buffer(9)]], @@ -207,113 +209,256 @@ template constexpr int BD = 32; constexpr int qk_per_thread = D / BD; constexpr int v_per_thread = V / BD; + // Upper bound on query rows per simdgroup, sized so the per-lane + // accumulator state stays register-resident. + constexpr int MAX_NQ = D <= 128 ? 4 : 2; typedef float U; - thread U q[qk_per_thread]; - thread U o[v_per_thread] = {0}; - // Adjust positions const int kv_head_idx = tid.x; const int batch_idx = tid.y; const int block_idx = tid.z; const int gqa_factor = tptg.y; - const int q_seq_len = tptg.z; - const int q_seq_idx = tidtg.z; const int q_head_idx = gqa_factor * kv_head_idx + tidtg.y; const int num_kv_heads = tpg.x; const int num_q_heads = num_kv_heads * gqa_factor; const int q_batch_head_idx = (batch_idx * num_q_heads + q_head_idx); - const int o_offset = q_batch_head_idx * q_seq_len + q_seq_idx; - const int q_offset = - query_transposed ? num_q_heads * q_seq_idx + q_batch_head_idx : o_offset; - - queries += q_offset * D + simd_lid * qk_per_thread; - const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx; + keys += kv_batch_head_idx * k_head_stride + block_idx * k_seq_stride + simd_lid * qk_per_thread; values += kv_batch_head_idx * v_head_stride + block_idx * v_seq_stride + simd_lid * v_per_thread; - out += o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread; - if (bool_mask) { - bmask += q_batch_head_idx * mask_head_stride + - block_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; - } - if (float_mask) { - fmask += q_batch_head_idx * mask_head_stride + - block_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; - } - sums += o_offset * blocks + block_idx; - maxs += o_offset * blocks + block_idx; - // Read the query - for (int i = 0; i < qk_per_thread; i++) { - q[i] = static_cast(scale) * queries[i]; - } + if (n_per_simd == 1) { + // One query row per simdgroup (decode shapes). Same as the original + // single-row kernel. + const int q_seq_idx = tidtg.z; + const int o_offset = q_batch_head_idx * q_seq_len + q_seq_idx; + const int q_offset = query_transposed + ? num_q_heads * q_seq_idx + q_batch_head_idx + : o_offset; - U max_score = Limits::finite_min; - U sum_exp_score = 0; - if (has_sinks && block_idx == 0) { - max_score = static_cast(sinks[q_head_idx]); - sum_exp_score = 1; - } + thread U q[qk_per_thread]; + thread U o[v_per_thread] = {0}; - // For each key - for (int i = block_idx; i < N; i += blocks) { - bool use_key = true; - if (do_causal) { - use_key = i <= (N - q_seq_len + int(q_seq_idx)); - } else if (bool_mask) { - use_key = bmask[0]; - } else if (float_mask) { - use_key = (fmask[0] >= Limits::finite_min); + queries += q_offset * D + simd_lid * qk_per_thread; + out += o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread; + if (bool_mask) { + bmask += q_batch_head_idx * mask_head_stride + + block_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } - if (use_key) { - // Compute the i-th score - U score = 0; - for (int i = 0; i < qk_per_thread; i++) { - score += q[i] * keys[i]; + if (float_mask) { + fmask += q_batch_head_idx * mask_head_stride + + block_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + } + sums += o_offset * blocks + block_idx; + maxs += o_offset * blocks + block_idx; + + // Read the query + for (int i = 0; i < qk_per_thread; i++) { + q[i] = static_cast(scale) * queries[i]; + } + + U max_score = Limits::finite_min; + U sum_exp_score = 0; + if (has_sinks && block_idx == 0) { + max_score = static_cast(sinks[q_head_idx]); + sum_exp_score = 1; + } + + // For each key + for (int i = block_idx; i < N; i += blocks) { + bool use_key = true; + if (do_causal) { + use_key = i <= (N - q_seq_len + int(q_seq_idx)); + } else if (bool_mask) { + use_key = bmask[0]; + } else if (float_mask) { + use_key = (fmask[0] >= Limits::finite_min); + } + if (use_key) { + // Compute the i-th score + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * keys[j]; + } + score = simd_sum(score); + + if (float_mask) { + score += fmask[0]; + } + + // Update the accumulators + U new_max = max(max_score, score); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(score - new_max); + + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + // Update the output accumulator + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * values[j]; + } } - score = simd_sum(score); + // Move the pointers to the next kv + keys += blocks * int(k_seq_stride); + values += blocks * int(v_seq_stride); + if (bool_mask) { + bmask += blocks * mask_kv_seq_stride; + } if (float_mask) { - score += fmask[0]; + fmask += blocks * mask_kv_seq_stride; } + } - // Update the accumulators - U new_max = max(max_score, score); - U factor = fast::exp(max_score - new_max); - U exp_score = fast::exp(score - new_max); - - max_score = new_max; - sum_exp_score = sum_exp_score * factor + exp_score; + // Write the sum and max and outputs + if (simd_lid == 0) { + sums[0] = sum_exp_score; + maxs[0] = max_score; + } - // Update the output accumulator - for (int i = 0; i < v_per_thread; i++) { - o[i] = o[i] * factor + exp_score * values[i]; - } + for (int i = 0; i < v_per_thread; i++) { + out[i] = static_cast(o[i]); } + } else { + // Several query rows per simdgroup (8 < qL <= 16, e.g. speculative + // decoding verify shapes). The K/V rows are read once per threadgroup + // and reused across the rows, keeping the block-parallel reduction + // over the key sequence that a single-row-per-simdgroup layout would + // lose to the threadgroup size limit. + const int nq = min(n_per_simd, MAX_NQ); + const int q_seq0 = tidtg.z * nq; + + thread U q[MAX_NQ][qk_per_thread]; + thread U o[MAX_NQ][v_per_thread]; + thread U kk[qk_per_thread]; + thread U vv[v_per_thread]; + thread U max_score[MAX_NQ]; + thread U sum_exp_score[MAX_NQ]; - // Move the pointers to the next kv - keys += blocks * int(k_seq_stride); - values += blocks * int(v_seq_stride); if (bool_mask) { - bmask += blocks * mask_kv_seq_stride; + bmask += q_batch_head_idx * mask_head_stride + + block_idx * mask_kv_seq_stride + q_seq0 * mask_q_seq_stride; } if (float_mask) { - fmask += blocks * mask_kv_seq_stride; + fmask += q_batch_head_idx * mask_head_stride + + block_idx * mask_kv_seq_stride + q_seq0 * mask_q_seq_stride; } - } - // Write the sum and max and outputs - if (simd_lid == 0) { - sums[0] = sum_exp_score; - maxs[0] = max_score; - } + // Read the queries and initialize the accumulators + for (int r = 0; r < nq; r++) { + const int q_seq_idx = q_seq0 + r; + max_score[r] = Limits::finite_min; + sum_exp_score[r] = 0; + for (int i = 0; i < v_per_thread; i++) { + o[r][i] = 0; + } + if (q_seq_idx < q_seq_len) { + const int o_offset = q_batch_head_idx * q_seq_len + q_seq_idx; + const int q_offset = query_transposed + ? num_q_heads * q_seq_idx + q_batch_head_idx + : o_offset; + const device T* qr = queries + q_offset * D + simd_lid * qk_per_thread; + for (int i = 0; i < qk_per_thread; i++) { + q[r][i] = static_cast(scale) * qr[i]; + } + if (has_sinks && block_idx == 0) { + max_score[r] = static_cast(sinks[q_head_idx]); + sum_exp_score[r] = 1; + } + } + } - for (int i = 0; i < v_per_thread; i++) { - out[i] = static_cast(o[i]); + // For each key + for (int i = block_idx; i < N; i += blocks) { + bool use_key[MAX_NQ]; + bool any_use_key = false; + for (int r = 0; r < nq; r++) { + const int q_seq_idx = q_seq0 + r; + bool use = q_seq_idx < q_seq_len; + if (use) { + if (do_causal) { + use = i <= (N - q_seq_len + q_seq_idx); + } else if (bool_mask) { + use = bmask[r * mask_q_seq_stride]; + } else if (float_mask) { + use = (fmask[r * mask_q_seq_stride] >= Limits::finite_min); + } + } + use_key[r] = use; + any_use_key |= use; + } + + if (any_use_key) { + // Read the key and value once and reuse across the query rows + for (int j = 0; j < qk_per_thread; j++) { + kk[j] = keys[j]; + } + for (int j = 0; j < v_per_thread; j++) { + vv[j] = values[j]; + } + + for (int r = 0; r < nq; r++) { + if (!use_key[r]) { + continue; + } + // Compute the i-th score for the r-th row + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q[r][j] * kk[j]; + } + score = simd_sum(score); + if (float_mask) { + score += static_cast(fmask[r * mask_q_seq_stride]); + } + + // Update the accumulators + U new_max = max(max_score[r], score); + U factor = fast::exp(max_score[r] - new_max); + U exp_score = fast::exp(score - new_max); + + max_score[r] = new_max; + sum_exp_score[r] = sum_exp_score[r] * factor + exp_score; + + // Update the output accumulator + for (int j = 0; j < v_per_thread; j++) { + o[r][j] = o[r][j] * factor + exp_score * vv[j]; + } + } + } + + // Move the pointers to the next kv + keys += blocks * int(k_seq_stride); + values += blocks * int(v_seq_stride); + if (bool_mask) { + bmask += blocks * mask_kv_seq_stride; + } + if (float_mask) { + fmask += blocks * mask_kv_seq_stride; + } + } + + // Write the sums, maxes and partial outputs + for (int r = 0; r < nq; r++) { + const int q_seq_idx = q_seq0 + r; + if (q_seq_idx >= q_seq_len) { + continue; + } + const int o_offset = q_batch_head_idx * q_seq_len + q_seq_idx; + if (simd_lid == 0) { + sums[o_offset * blocks + block_idx] = sum_exp_score[r]; + maxs[o_offset * blocks + block_idx] = max_score[r]; + } + device T* outr = + out + o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread; + for (int j = 0; j < v_per_thread; j++) { + outr[j] = static_cast(o[r][j]); + } + } } } diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index d6f74bc2a9..3ddf2f1bc1 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -415,6 +415,31 @@ void sdpa_vector( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +// Query rows processed serially by each simdgroup in the 2 pass vector +// kernel. Bounded by register pressure: larger head dims hold more per-lane +// accumulator state (see MAX_NQ in sdpa_vector.h). +inline int sdpa_vector_max_nq(int query_head_dim) { + return query_head_dim <= 128 ? 4 : 2; +} + +inline bool sdpa_vector_supported_shape( + int query_head_dim, + int value_head_dim, + int query_sequence_length, + int key_sequence_length, + int gqa_factor) { + const bool supported_head_dim = + (query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || + query_head_dim == 256)) || + (query_head_dim == 192 && value_head_dim == 128); + const int max_nq = sdpa_vector_max_nq(query_head_dim); + const int min_rows_per_tg = (query_sequence_length + max_nq - 1) / max_nq; + return supported_head_dim && query_sequence_length <= 16 && + query_sequence_length <= key_sequence_length && + gqa_factor * min_rows_per_tg <= 32; +} + void sdpa_vector_2pass( const Stream& s, metal::Device& d, @@ -438,7 +463,16 @@ void sdpa_vector_2pass( // Compute the necessary sizes int gqa_factor = q.shape(1) / k.shape(1); - int n_simds = gqa_factor * q.shape(2); + int q_seq_len = q.shape(2); + // Rows per simdgroup: grow nq until the threadgroup fits its simdgroups + // (32 lanes x gqa_factor x rows_per_tg <= 1024 threads). + int max_nq = sdpa_vector_max_nq(q.shape(-1)); + int nq = 1; + while (nq < max_nq && gqa_factor * ((q_seq_len + nq - 1) / nq) > 32) { + nq *= 2; + } + int rows_per_tg = (q_seq_len + nq - 1) / nq; + int n_simds = gqa_factor * rows_per_tg; char devc = d.get_architecture().back(); int N = k.shape(2); @@ -481,7 +515,7 @@ void sdpa_vector_2pass( size_t k_seq_stride = k.strides()[2]; size_t v_head_stride = v.shape(1) == 1 ? v.strides(0) : v.strides(1); size_t v_seq_stride = v.strides()[2]; - MTL::Size group_dims(32, gqa_factor, q.shape(2)); + MTL::Size group_dims(32, gqa_factor, rows_per_tg); MTL::Size grid_dims(k.shape(1), q.shape(0), blocks); // Allocate the intermediates @@ -516,6 +550,7 @@ void sdpa_vector_2pass( {&float_mask, MTL::DataType::DataTypeBool, 24}, {&has_sinks, MTL::DataType::DataTypeBool, 25}, {&blocks, MTL::DataType::DataTypeInt, 26}, + {&nq, MTL::DataType::DataTypeInt, 27}, }; std::string hash_name = kname; hash_name += has_mask ? (bool_mask ? "_boolmask" : "_floatmask") : "_nomask"; @@ -523,9 +558,26 @@ void sdpa_vector_2pass( hash_name += do_causal ? "_c" : "_nc"; hash_name += has_sinks ? "_sinks_" : "_nosinks_"; hash_name += std::to_string(blocks); + auto base_hash_name = hash_name; + hash_name += "_nq"; + hash_name += std::to_string(nq); // Get the kernel auto kernel = d.get_kernel(kname, hash_name, func_consts); + // High register pressure can lower the pipeline's threadgroup limit; + // trade threadgroup rows for more rows per simdgroup when that happens. + while (nq < max_nq && + kernel->maxTotalThreadsPerThreadgroup() < + static_cast(32 * n_simds)) { + nq *= 2; + rows_per_tg = (q_seq_len + nq - 1) / nq; + n_simds = gqa_factor * rows_per_tg; + group_dims = MTL::Size(32, gqa_factor, rows_per_tg); + hash_name = base_hash_name; + hash_name += "_nq"; + hash_name += std::to_string(nq); + kernel = d.get_kernel(kname, hash_name, func_consts); + } check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.set_compute_pipeline_state(kernel); @@ -537,6 +589,7 @@ void sdpa_vector_2pass( 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(q_seq_len, 6); compute_encoder.set_bytes(N, 7); compute_encoder.set_bytes(k_head_stride, 8); compute_encoder.set_bytes(k_seq_stride, 9); @@ -618,11 +671,6 @@ bool ScaledDotProductAttention::use_fallback( const int num_kv_heads = k.shape(1); const int gqa_factor = num_query_heads / num_kv_heads; - const bool sdpa_vector_supported_head_dim = - (query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || - query_head_dim == 256)) || - (query_head_dim == 192 && value_head_dim == 128); const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim && (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128); @@ -632,10 +680,12 @@ bool ScaledDotProductAttention::use_fallback( const bool supports_sdpa_full = query_sequence_length > 8 && sdpa_full_supported_mask && sdpa_full_supported_head_dim; - const bool supports_sdpa_vector = (query_sequence_length <= 8) && - (query_sequence_length <= key_sequence_length) && - sdpa_vector_supported_head_dim && - (query_sequence_length * gqa_factor) <= 32; + const bool supports_sdpa_vector = sdpa_vector_supported_shape( + query_head_dim, + value_head_dim, + query_sequence_length, + key_sequence_length, + gqa_factor); return !(supports_sdpa_full || supports_sdpa_vector); } @@ -682,8 +732,25 @@ void ScaledDotProductAttention::eval_gpu( } bool has_arr_mask = inputs.size() > (3 + has_sinks_); - // We are in vector mode ie single query - if (q_pre.shape(2) <= 8) { + // We are in vector mode ie a small number of query rows (decode and + // speculative-decoding verify shapes) + const int query_seq_len = q_pre.shape(2); + const bool vector_supported = sdpa_vector_supported_shape( + q_pre.shape(-1), + v_pre.shape(-1), + query_seq_len, + k_pre.shape(2), + q_pre.shape(1) / k_pre.shape(1)); + const bool full_supported_head_dim = q_pre.shape(-1) == v_pre.shape(-1) && + (q_pre.shape(-1) == 64 || q_pre.shape(-1) == 80 || + q_pre.shape(-1) == 128); + // For 8 < qL <= 16 route to the vector path when the keys are long enough + // that its block-parallel reduction over the sequence beats the full + // attention kernel's single query tile, or when the full kernel does not + // support the head dim. + if (vector_supported && + (query_seq_len <= 8 || k_pre.shape(2) >= 1024 || + !full_supported_head_dim)) { auto q_copy_unless = [](const array& arr) { if (arr.flags().row_contiguous) { return true; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 7bd867084e..158debb6db 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -280,28 +280,29 @@ def test_sdpa_few_query(self): L = 4096 scale = 1.0 mx.random.seed(0) - q = 5e-1 * mx.random.normal(shape=(1, Nq, Lq, D)) - k = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D)) - v = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D)) + for Lq, Nkv in product((8, 12, 16), (1, 2)): + q = 5e-1 * mx.random.normal(shape=(1, Nq, Lq, D)) + k = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D)) + v = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D)) - masks = [ - None, - mx.array(True), - mx.array([True] * (L - 10) + [False] * 10), - mx.random.uniform(shape=(Nq, 1, L)) > 0.2, - mx.random.uniform(shape=(L, 1, Nq)).T > 0.2, - "causal", - ] - for m in masks: - ref = mlx_primitives_sdpa(q, k, v, scale, mask=m) - out = mx.fast.scaled_dot_product_attention( - q, - k, - v, - scale=scale, - mask=m, - ) - self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + masks = [ + None, + mx.array(True), + mx.array([True] * (L - 10) + [False] * 10), + mx.random.uniform(shape=(Nq, 1, L)) > 0.2, + mx.random.uniform(shape=(L, 1, Nq)).T > 0.2, + "causal", + ] + for m in masks: + ref = mlx_ref_attn(q, k, v, scale, mask=m) + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + mask=m, + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) @unittest.skip("Different head and value dims is not enabled") def test_sdpa_vector_value_dims(self):