Skip to content

Extend SDPA vector 2-pass kernel to query sequence lengths up to 16 #3838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
katlun-lgtm wants to merge 1 commit into
base: main
Choose a base branch
from
+319 −106
Open

Extend SDPA vector 2-pass kernel to query sequence lengths up to 16 #3838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
289 changes: 217 additions & 72 deletions mlx/backend/metal/kernels/sdpa_vector.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename T, int D, int V = D>
[[kernel]] void sdpa_vector(
Expand Down Expand Up @@ -184,6 +185,7 @@ template <typename T, int D, int V = D>
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)]],
Expand All @@ -207,113 +209,256 @@ template <typename T, int D, int V = D>
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<U>(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<U>::finite_min;
U sum_exp_score = 0;
if (has_sinks && block_idx == 0) {
max_score = static_cast<U>(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<T>::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<U>(scale) * queries[i];
}

U max_score = Limits<U>::finite_min;
U sum_exp_score = 0;
if (has_sinks && block_idx == 0) {
max_score = static_cast<U>(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<T>::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<T>(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<U>::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<U>(scale) * qr[i];
}
if (has_sinks && block_idx == 0) {
max_score[r] = static_cast<U>(sinks[q_head_idx]);
sum_exp_score[r] = 1;
}
}
}

for (int i = 0; i < v_per_thread; i++) {
out[i] = static_cast<T>(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<T>::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<U>(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<T>(o[r][j]);
}
}
}
}

Expand Down
Loading