495 lines
18 KiB
Metal
495 lines
18 KiB
Metal
//! Flash Attention Metal Shaders
|
|
//!
|
|
//! Metal Shading Language implementation of Flash Attention for Apple Silicon
|
|
//! Based on the Flash Attention algorithm by Dao et al. (2022)
|
|
//!
|
|
//! Implements tiled attention with online softmax for O(N) memory complexity.
|
|
|
|
#include <metal_stdlib>
|
|
using namespace metal;
|
|
|
|
// Block sizes for tiling - tuned for Apple Silicon's 32KB threadgroup memory limit
|
|
// Forward kernel: 32x32 blocks → Q(8KB) + K(8KB) + V(8KB) + S(4KB) = 28KB ✓
|
|
// Backward kernels: 16x16 blocks → Q+K+V+dO+O (5 tiles * 4KB) = 20KB ✓
|
|
constant uint BLOCK_Q_FWD = 32;
|
|
constant uint BLOCK_KV_FWD = 32;
|
|
constant uint BLOCK_Q_BWD = 16;
|
|
constant uint BLOCK_KV_BWD = 16;
|
|
constant uint MAX_HEAD_DIM = 128;
|
|
|
|
/// Flash Attention kernel parameters
|
|
struct FlashAttentionParams {
|
|
uint batch_size;
|
|
uint num_heads;
|
|
uint seq_len_q;
|
|
uint seq_len_kv;
|
|
uint head_dim;
|
|
float softmax_scale;
|
|
uint causal; // 0 = false, 1 = true
|
|
uint block_size_q;
|
|
uint block_size_kv;
|
|
};
|
|
|
|
/// Helper: Simdgroup reduction for max
|
|
inline float simdgroup_reduce_max(float val) {
|
|
for (uint offset = 16; offset > 0; offset >>= 1) {
|
|
val = max(val, simd_shuffle_down(val, offset));
|
|
}
|
|
return simd_broadcast_first(val);
|
|
}
|
|
|
|
/// Helper: Simdgroup reduction for sum
|
|
inline float simdgroup_reduce_sum(float val) {
|
|
for (uint offset = 16; offset > 0; offset >>= 1) {
|
|
val += simd_shuffle_down(val, offset);
|
|
}
|
|
return simd_broadcast_first(val);
|
|
}
|
|
|
|
/// Flash Attention forward kernel
|
|
///
|
|
/// Implements the Flash Attention algorithm with:
|
|
/// - Tiled loading of Q, K, V into threadgroup memory
|
|
/// - Online softmax with running max/sum
|
|
/// - Causal masking support
|
|
/// - Optimized memory access patterns for Apple Silicon
|
|
kernel void flash_attention_forward(
|
|
device const half* Q [[buffer(0)]],
|
|
device const half* K [[buffer(1)]],
|
|
device const half* V [[buffer(2)]],
|
|
device half* O [[buffer(3)]],
|
|
device float* LSE [[buffer(4)]],
|
|
constant FlashAttentionParams& params [[buffer(5)]],
|
|
uint3 tgid [[threadgroup_position_in_grid]],
|
|
uint3 tid [[thread_position_in_threadgroup]],
|
|
uint3 tg_size [[threads_per_threadgroup]],
|
|
uint thread_idx [[thread_index_in_threadgroup]],
|
|
uint simd_lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]
|
|
) {
|
|
// Threadgroup memory for Q, K, V tiles (32KB limit on Apple Silicon)
|
|
// Forward: 8KB + 8KB + 8KB + 4KB = 28KB ✓
|
|
threadgroup half Q_tile[BLOCK_Q_FWD][MAX_HEAD_DIM];
|
|
threadgroup half K_tile[BLOCK_KV_FWD][MAX_HEAD_DIM];
|
|
threadgroup half V_tile[BLOCK_KV_FWD][MAX_HEAD_DIM];
|
|
threadgroup float S_tile[BLOCK_Q_FWD][BLOCK_KV_FWD]; // Attention scores
|
|
|
|
// Identify which batch, head, and Q block we're processing
|
|
uint batch_idx = tgid.z;
|
|
uint head_idx = tgid.y;
|
|
uint q_block_idx = tgid.x;
|
|
|
|
// Base offsets for this batch and head
|
|
uint qkv_stride = params.seq_len_q * params.head_dim;
|
|
uint head_offset = (batch_idx * params.num_heads + head_idx) * qkv_stride;
|
|
|
|
// Q block start position
|
|
uint q_start = q_block_idx * BLOCK_Q_FWD;
|
|
|
|
// Thread's row within the Q block
|
|
uint q_row = thread_idx;
|
|
uint q_idx = q_start + q_row;
|
|
|
|
// Online softmax state (per-thread for its Q row)
|
|
float m_i = -INFINITY; // Running max
|
|
float l_i = 0.0f; // Running sum of exp(scores - max)
|
|
|
|
// Output accumulator (per-thread for its Q row)
|
|
float o_acc[MAX_HEAD_DIM];
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
o_acc[d] = 0.0f;
|
|
}
|
|
|
|
// Load Q tile into threadgroup memory (each thread loads one row)
|
|
if (q_idx < params.seq_len_q && q_row < BLOCK_Q_FWD) {
|
|
uint q_offset = head_offset + q_idx * params.head_dim;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
Q_tile[q_row][d] = Q[q_offset + d];
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Number of K/V blocks to iterate over
|
|
uint num_kv_blocks = (params.seq_len_kv + BLOCK_KV_FWD - 1) / BLOCK_KV_FWD;
|
|
|
|
// For causal attention, limit K/V blocks based on Q position
|
|
uint max_kv_block = num_kv_blocks;
|
|
if (params.causal) {
|
|
// Only process K/V positions up to q_start + BLOCK_Q_FWD - 1
|
|
max_kv_block = min(num_kv_blocks, (q_start + BLOCK_Q_FWD + BLOCK_KV_FWD - 1) / BLOCK_KV_FWD);
|
|
}
|
|
|
|
// Iterate over K/V blocks
|
|
for (uint kv_block = 0; kv_block < max_kv_block; kv_block++) {
|
|
uint kv_start = kv_block * BLOCK_KV_FWD;
|
|
|
|
// Load K and V tiles (collaborative loading)
|
|
if (q_row < BLOCK_KV_FWD) {
|
|
uint kv_idx = kv_start + q_row;
|
|
if (kv_idx < params.seq_len_kv) {
|
|
uint kv_offset = head_offset + kv_idx * params.head_dim;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
K_tile[q_row][d] = K[kv_offset + d];
|
|
V_tile[q_row][d] = V[kv_offset + d];
|
|
}
|
|
} else {
|
|
// Padding for out-of-bounds
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
K_tile[q_row][d] = half(0.0);
|
|
V_tile[q_row][d] = half(0.0);
|
|
}
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Skip if this thread's Q row is out of bounds
|
|
if (q_idx >= params.seq_len_q || q_row >= BLOCK_Q_FWD) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
continue;
|
|
}
|
|
|
|
// Compute attention scores for this Q row against all K rows in the block
|
|
// S[i,j] = Q[i] @ K[j]^T * scale
|
|
float row_max = -INFINITY;
|
|
|
|
for (uint j = 0; j < BLOCK_KV_FWD; j++) {
|
|
uint kv_idx = kv_start + j;
|
|
|
|
// Skip if K/V position is out of bounds
|
|
if (kv_idx >= params.seq_len_kv) {
|
|
S_tile[q_row][j] = -INFINITY;
|
|
continue;
|
|
}
|
|
|
|
// Causal mask: future positions get -infinity
|
|
if (params.causal && kv_idx > q_idx) {
|
|
S_tile[q_row][j] = -INFINITY;
|
|
continue;
|
|
}
|
|
|
|
// Dot product Q[i] @ K[j]
|
|
float score = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
score += float(Q_tile[q_row][d]) * float(K_tile[j][d]);
|
|
}
|
|
score *= params.softmax_scale;
|
|
S_tile[q_row][j] = score;
|
|
row_max = max(row_max, score);
|
|
}
|
|
|
|
// Online softmax update
|
|
// m_new = max(m_i, row_max)
|
|
// l_new = l_i * exp(m_i - m_new) + sum(exp(S - m_new))
|
|
// O_new = O_i * exp(m_i - m_new) + sum(exp(S - m_new) * V)
|
|
|
|
float m_new = max(m_i, row_max);
|
|
float exp_diff = exp(m_i - m_new);
|
|
|
|
// Rescale existing accumulator
|
|
l_i *= exp_diff;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
o_acc[d] *= exp_diff;
|
|
}
|
|
|
|
// Accumulate new contributions
|
|
for (uint j = 0; j < BLOCK_KV_FWD; j++) {
|
|
float score = S_tile[q_row][j];
|
|
if (score > -INFINITY * 0.5f) { // Skip masked positions
|
|
float exp_score = exp(score - m_new);
|
|
l_i += exp_score;
|
|
|
|
// Accumulate weighted V
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
o_acc[d] += exp_score * float(V_tile[j][d]);
|
|
}
|
|
}
|
|
}
|
|
|
|
m_i = m_new;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write final output: O = acc / l_i
|
|
if (q_idx < params.seq_len_q && q_row < BLOCK_Q_FWD) {
|
|
float inv_l = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f;
|
|
uint o_offset = head_offset + q_idx * params.head_dim;
|
|
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
O[o_offset + d] = half(o_acc[d] * inv_l);
|
|
}
|
|
|
|
// Store log-sum-exp for backward pass: LSE = m + log(l)
|
|
uint lse_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q + q_idx;
|
|
LSE[lse_offset] = m_i + log(max(l_i, 1e-10f));
|
|
}
|
|
}
|
|
|
|
/// Flash Attention backward kernel for dQ
|
|
///
|
|
/// Computes gradient with respect to Q:
|
|
/// dQ[i] = sum_j( softmax_grad[i,j] * K[j] )
|
|
/// where softmax_grad[i,j] = P[i,j] * (dO[i] @ V[j]^T - D[i])
|
|
/// and D[i] = sum_j(P[i,j] * dO[i] @ V[j]^T) = O[i] @ dO[i]^T (element-wise then sum)
|
|
kernel void flash_attention_backward_dq(
|
|
device const half* dO [[buffer(0)]],
|
|
device const half* Q [[buffer(1)]],
|
|
device const half* K [[buffer(2)]],
|
|
device const half* V [[buffer(3)]],
|
|
device const half* O [[buffer(4)]],
|
|
device const float* LSE [[buffer(5)]],
|
|
device half* dQ [[buffer(6)]],
|
|
constant FlashAttentionParams& params [[buffer(7)]],
|
|
uint3 tgid [[threadgroup_position_in_grid]],
|
|
uint3 tid [[thread_position_in_threadgroup]],
|
|
uint thread_idx [[thread_index_in_threadgroup]]
|
|
) {
|
|
// Threadgroup memory for tiles (32KB limit on Apple Silicon)
|
|
// Backward: 5 tiles * 16 * 128 * 2 = 20KB ✓
|
|
threadgroup half Q_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
|
|
threadgroup half K_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
|
|
threadgroup half V_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
|
|
threadgroup half dO_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
|
|
threadgroup half O_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
|
|
threadgroup float D_tile[BLOCK_Q_BWD]; // D[i] = dO[i] @ O[i]
|
|
|
|
uint batch_idx = tgid.z;
|
|
uint head_idx = tgid.y;
|
|
uint q_block_idx = tgid.x;
|
|
|
|
uint head_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q * params.head_dim;
|
|
uint lse_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q;
|
|
|
|
uint q_start = q_block_idx * BLOCK_Q_BWD;
|
|
uint q_row = thread_idx;
|
|
uint q_idx = q_start + q_row;
|
|
|
|
// Output gradient accumulator
|
|
float dq_acc[MAX_HEAD_DIM];
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dq_acc[d] = 0.0f;
|
|
}
|
|
|
|
// Load Q, dO, O tiles and compute D = dO @ O (element-wise sum)
|
|
if (q_idx < params.seq_len_q && q_row < BLOCK_Q_BWD) {
|
|
uint offset = head_offset + q_idx * params.head_dim;
|
|
float d_val = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
Q_tile[q_row][d] = Q[offset + d];
|
|
dO_tile[q_row][d] = dO[offset + d];
|
|
O_tile[q_row][d] = O[offset + d];
|
|
d_val += float(dO_tile[q_row][d]) * float(O_tile[q_row][d]);
|
|
}
|
|
D_tile[q_row] = d_val;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Iterate over K/V blocks
|
|
uint num_kv_blocks = (params.seq_len_kv + BLOCK_KV_BWD - 1) / BLOCK_KV_BWD;
|
|
uint max_kv_block = num_kv_blocks;
|
|
if (params.causal) {
|
|
max_kv_block = min(num_kv_blocks, (q_start + BLOCK_Q_BWD + BLOCK_KV_BWD - 1) / BLOCK_KV_BWD);
|
|
}
|
|
|
|
for (uint kv_block = 0; kv_block < max_kv_block; kv_block++) {
|
|
uint kv_start = kv_block * BLOCK_KV_BWD;
|
|
|
|
// Load K, V tiles
|
|
if (q_row < BLOCK_KV_BWD) {
|
|
uint kv_idx = kv_start + q_row;
|
|
if (kv_idx < params.seq_len_kv) {
|
|
uint offset = head_offset + kv_idx * params.head_dim;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
K_tile[q_row][d] = K[offset + d];
|
|
V_tile[q_row][d] = V[offset + d];
|
|
}
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (q_idx >= params.seq_len_q || q_row >= BLOCK_Q_BWD) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
continue;
|
|
}
|
|
|
|
float lse_val = LSE[lse_offset + q_idx];
|
|
|
|
for (uint j = 0; j < BLOCK_KV_BWD; j++) {
|
|
uint kv_idx = kv_start + j;
|
|
if (kv_idx >= params.seq_len_kv) continue;
|
|
if (params.causal && kv_idx > q_idx) continue;
|
|
|
|
// Recompute attention score
|
|
float score = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
score += float(Q_tile[q_row][d]) * float(K_tile[j][d]);
|
|
}
|
|
score *= params.softmax_scale;
|
|
|
|
// Recompute P[i,j] = exp(score - LSE[i])
|
|
float p_ij = exp(score - lse_val);
|
|
|
|
// Compute dO[i] @ V[j]
|
|
float dov = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dov += float(dO_tile[q_row][d]) * float(V_tile[j][d]);
|
|
}
|
|
|
|
// softmax_grad = P * (dO @ V - D)
|
|
float softmax_grad = p_ij * (dov - D_tile[q_row]);
|
|
|
|
// dQ[i] += softmax_grad * K[j] * scale
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dq_acc[d] += softmax_grad * float(K_tile[j][d]) * params.softmax_scale;
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write dQ output
|
|
if (q_idx < params.seq_len_q && q_row < BLOCK_Q_BWD) {
|
|
uint offset = head_offset + q_idx * params.head_dim;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dQ[offset + d] = half(dq_acc[d]);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Flash Attention backward kernel for dK and dV
|
|
///
|
|
/// Computes gradients with respect to K and V:
|
|
/// dK[j] = sum_i( softmax_grad[i,j] * Q[i] )
|
|
/// dV[j] = sum_i( P[i,j] * dO[i] )
|
|
kernel void flash_attention_backward_dkv(
|
|
device const half* dO [[buffer(0)]],
|
|
device const half* Q [[buffer(1)]],
|
|
device const half* K [[buffer(2)]],
|
|
device const half* V [[buffer(3)]],
|
|
device const half* O [[buffer(4)]],
|
|
device const float* LSE [[buffer(5)]],
|
|
device half* dK [[buffer(6)]],
|
|
device half* dV [[buffer(7)]],
|
|
constant FlashAttentionParams& params [[buffer(8)]],
|
|
uint3 tgid [[threadgroup_position_in_grid]],
|
|
uint3 tid [[thread_position_in_threadgroup]],
|
|
uint thread_idx [[thread_index_in_threadgroup]]
|
|
) {
|
|
// Threadgroup memory for tiles (32KB limit on Apple Silicon)
|
|
// Backward: 5 tiles * 16 * 128 * 2 = 20KB ✓
|
|
threadgroup half Q_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
|
|
threadgroup half K_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
|
|
threadgroup half V_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
|
|
threadgroup half dO_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
|
|
threadgroup half O_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
|
|
threadgroup float D_tile[BLOCK_Q_BWD];
|
|
|
|
uint batch_idx = tgid.z;
|
|
uint head_idx = tgid.y;
|
|
uint kv_block_idx = tgid.x;
|
|
|
|
uint head_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q * params.head_dim;
|
|
uint lse_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q;
|
|
|
|
uint kv_start = kv_block_idx * BLOCK_KV_BWD;
|
|
uint kv_row = thread_idx;
|
|
uint kv_idx = kv_start + kv_row;
|
|
|
|
// Output gradient accumulators
|
|
float dk_acc[MAX_HEAD_DIM];
|
|
float dv_acc[MAX_HEAD_DIM];
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dk_acc[d] = 0.0f;
|
|
dv_acc[d] = 0.0f;
|
|
}
|
|
|
|
// Load K, V tiles
|
|
if (kv_idx < params.seq_len_kv && kv_row < BLOCK_KV_BWD) {
|
|
uint offset = head_offset + kv_idx * params.head_dim;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
K_tile[kv_row][d] = K[offset + d];
|
|
V_tile[kv_row][d] = V[offset + d];
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Iterate over Q blocks
|
|
uint num_q_blocks = (params.seq_len_q + BLOCK_Q_BWD - 1) / BLOCK_Q_BWD;
|
|
|
|
for (uint q_block = 0; q_block < num_q_blocks; q_block++) {
|
|
uint q_start = q_block * BLOCK_Q_BWD;
|
|
|
|
// For causal: skip Q blocks that come before this K/V position
|
|
if (params.causal && q_start + BLOCK_Q_BWD - 1 < kv_start) {
|
|
continue;
|
|
}
|
|
|
|
// Load Q, dO, O tiles and compute D
|
|
if (kv_row < BLOCK_Q_BWD) {
|
|
uint q_idx = q_start + kv_row;
|
|
if (q_idx < params.seq_len_q) {
|
|
uint offset = head_offset + q_idx * params.head_dim;
|
|
float d_val = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
Q_tile[kv_row][d] = Q[offset + d];
|
|
dO_tile[kv_row][d] = dO[offset + d];
|
|
O_tile[kv_row][d] = O[offset + d];
|
|
d_val += float(dO_tile[kv_row][d]) * float(O_tile[kv_row][d]);
|
|
}
|
|
D_tile[kv_row] = d_val;
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (kv_idx >= params.seq_len_kv || kv_row >= BLOCK_KV_BWD) {
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
continue;
|
|
}
|
|
|
|
for (uint i = 0; i < BLOCK_Q_BWD; i++) {
|
|
uint q_idx = q_start + i;
|
|
if (q_idx >= params.seq_len_q) continue;
|
|
if (params.causal && kv_idx > q_idx) continue;
|
|
|
|
float lse_val = LSE[lse_offset + q_idx];
|
|
|
|
// Recompute attention score
|
|
float score = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
score += float(Q_tile[i][d]) * float(K_tile[kv_row][d]);
|
|
}
|
|
score *= params.softmax_scale;
|
|
|
|
// Recompute P[i,j] = exp(score - LSE[i])
|
|
float p_ij = exp(score - lse_val);
|
|
|
|
// dV[j] += P[i,j] * dO[i]
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dv_acc[d] += p_ij * float(dO_tile[i][d]);
|
|
}
|
|
|
|
// Compute dO[i] @ V[j]
|
|
float dov = 0.0f;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dov += float(dO_tile[i][d]) * float(V_tile[kv_row][d]);
|
|
}
|
|
|
|
// softmax_grad = P * (dO @ V - D)
|
|
float softmax_grad = p_ij * (dov - D_tile[i]);
|
|
|
|
// dK[j] += softmax_grad * Q[i] * scale
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dk_acc[d] += softmax_grad * float(Q_tile[i][d]) * params.softmax_scale;
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write dK, dV outputs
|
|
if (kv_idx < params.seq_len_kv && kv_row < BLOCK_KV_BWD) {
|
|
uint offset = head_offset + kv_idx * params.head_dim;
|
|
for (uint d = 0; d < params.head_dim; d++) {
|
|
dK[offset + d] = half(dk_acc[d]);
|
|
dV[offset + d] = half(dv_acc[d]);
|
|
}
|
|
}
|
|
}
|