Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,183 @@
|
||||
//
|
||||
// Flash Attention Forward Pass - Metal Shader
|
||||
//
|
||||
// Implements the Flash Attention algorithm with online softmax
|
||||
// for memory-efficient attention computation.
|
||||
//
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// Preprocessor macros set during compilation:
|
||||
// BLOCK_Q - Number of Q rows per threadgroup
|
||||
// BLOCK_KV - Number of K/V rows per tile
|
||||
// HEAD_DIM - Maximum head dimension
|
||||
|
||||
// Apple Silicon GPUs have 32KB threadgroup memory limit
|
||||
// Block sizes tuned for performance while staying under limit
|
||||
|
||||
#ifndef BLOCK_Q
|
||||
#define BLOCK_Q 16
|
||||
#endif
|
||||
|
||||
#ifndef BLOCK_KV
|
||||
#define BLOCK_KV 16
|
||||
#endif
|
||||
|
||||
#ifndef HEAD_DIM
|
||||
#define HEAD_DIM 64
|
||||
#endif
|
||||
|
||||
/// Parameters for Flash Attention forward pass
|
||||
struct FlashAttentionParams {
|
||||
uint batch_size; // Number of sequences in batch
|
||||
uint num_heads; // Number of attention heads
|
||||
uint seq_len_q; // Query sequence length
|
||||
uint seq_len_kv; // Key/Value sequence length
|
||||
uint head_dim; // Dimension per head
|
||||
float softmax_scale; // Scaling factor (1/sqrt(head_dim))
|
||||
uint causal; // Whether to apply causal mask
|
||||
};
|
||||
|
||||
/// Flash Attention forward kernel
|
||||
///
|
||||
/// Computes: O = softmax(Q @ K^T * scale) @ V
|
||||
/// Using online softmax to avoid materializing the full attention matrix
|
||||
///
|
||||
/// Grid: (num_q_blocks, num_heads, batch_size)
|
||||
/// Threadgroup: (BLOCK_Q, 1, 1)
|
||||
kernel void flash_attention_forward(
|
||||
device const float* Q [[buffer(0)]], // [batch, heads, seq_q, head_dim]
|
||||
device const float* K [[buffer(1)]], // [batch, heads, seq_kv, head_dim]
|
||||
device const float* V [[buffer(2)]], // [batch, heads, seq_kv, head_dim]
|
||||
device float* O [[buffer(3)]], // [batch, heads, seq_q, head_dim]
|
||||
device float* LSE [[buffer(4)]], // [batch, heads, seq_q] - log-sum-exp for backward
|
||||
constant FlashAttentionParams& params [[buffer(5)]],
|
||||
uint3 tgid [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint simd_lane [[thread_index_in_simdgroup]]
|
||||
) {
|
||||
// Threadgroup shared memory for tiles
|
||||
threadgroup float Q_shared[BLOCK_Q * HEAD_DIM];
|
||||
threadgroup float K_shared[BLOCK_KV * HEAD_DIM];
|
||||
threadgroup float V_shared[BLOCK_KV * HEAD_DIM];
|
||||
|
||||
// Identify which batch/head/block this threadgroup handles
|
||||
uint batch_idx = tgid.z;
|
||||
uint head_idx = tgid.y;
|
||||
uint q_block = tgid.x;
|
||||
uint q_start = q_block * BLOCK_Q;
|
||||
uint q_idx = q_start + tid;
|
||||
|
||||
// Calculate memory strides
|
||||
uint stride_batch = params.num_heads * params.seq_len_q * params.head_dim;
|
||||
uint stride_head = params.seq_len_q * params.head_dim;
|
||||
uint base_qo = batch_idx * stride_batch + head_idx * stride_head;
|
||||
|
||||
// For K/V, seq_len might differ
|
||||
uint stride_kv_batch = params.num_heads * params.seq_len_kv * params.head_dim;
|
||||
uint stride_kv_head = params.seq_len_kv * params.head_dim;
|
||||
uint base_kv = batch_idx * stride_kv_batch + head_idx * stride_kv_head;
|
||||
|
||||
// Online softmax accumulators (per thread handles one Q row)
|
||||
float m_i = -INFINITY; // Running max
|
||||
float l_i = 0.0f; // Running sum of exp(scores - max)
|
||||
|
||||
// Output accumulator
|
||||
float o_acc[HEAD_DIM];
|
||||
for (uint d = 0; d < params.head_dim; d++) {
|
||||
o_acc[d] = 0.0f;
|
||||
}
|
||||
|
||||
// Load Q tile into shared memory
|
||||
if (q_idx < params.seq_len_q) {
|
||||
for (uint d = 0; d < params.head_dim; d++) {
|
||||
Q_shared[tid * HEAD_DIM + d] = Q[base_qo + q_idx * params.head_dim + d];
|
||||
}
|
||||
} else {
|
||||
// Pad with zeros for out-of-bounds threads
|
||||
for (uint d = 0; d < params.head_dim; d++) {
|
||||
Q_shared[tid * HEAD_DIM + d] = 0.0f;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Iterate over K/V blocks
|
||||
uint num_kv_blocks = (params.seq_len_kv + BLOCK_KV - 1) / BLOCK_KV;
|
||||
|
||||
for (uint kv_block = 0; kv_block < num_kv_blocks; kv_block++) {
|
||||
uint kv_start = kv_block * BLOCK_KV;
|
||||
|
||||
// Causal optimization: skip future blocks entirely
|
||||
if (params.causal != 0 && kv_start > q_start + BLOCK_Q - 1) {
|
||||
break;
|
||||
}
|
||||
|
||||
// Cooperatively load K and V tiles
|
||||
for (uint i = tid; i < BLOCK_KV * params.head_dim; i += BLOCK_Q) {
|
||||
uint kv_row = i / params.head_dim;
|
||||
uint d = i % params.head_dim;
|
||||
uint kv_idx = kv_start + kv_row;
|
||||
|
||||
if (kv_idx < params.seq_len_kv) {
|
||||
K_shared[kv_row * HEAD_DIM + d] = K[base_kv + kv_idx * params.head_dim + d];
|
||||
V_shared[kv_row * HEAD_DIM + d] = V[base_kv + kv_idx * params.head_dim + d];
|
||||
} else {
|
||||
K_shared[kv_row * HEAD_DIM + d] = 0.0f;
|
||||
V_shared[kv_row * HEAD_DIM + d] = 0.0f;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Compute attention scores for this K/V block
|
||||
if (q_idx < params.seq_len_q) {
|
||||
for (uint j = 0; j < BLOCK_KV; j++) {
|
||||
uint kv_idx = kv_start + j;
|
||||
|
||||
// Skip if out of bounds
|
||||
if (kv_idx >= params.seq_len_kv) continue;
|
||||
|
||||
// Apply causal mask
|
||||
if (params.causal != 0 && kv_idx > q_idx) continue;
|
||||
|
||||
// Compute dot product: Q[q_idx] @ K[kv_idx]^T
|
||||
float score = 0.0f;
|
||||
for (uint d = 0; d < params.head_dim; d++) {
|
||||
score += Q_shared[tid * HEAD_DIM + d] * K_shared[j * HEAD_DIM + d];
|
||||
}
|
||||
score *= params.softmax_scale;
|
||||
|
||||
// Online softmax update
|
||||
float m_new = max(m_i, score);
|
||||
float exp_diff = exp(m_i - m_new);
|
||||
float exp_score = exp(score - m_new);
|
||||
|
||||
// Update running sum and rescale accumulator
|
||||
l_i = l_i * exp_diff + exp_score;
|
||||
|
||||
// Update output accumulator: O += exp(score - max) * V[kv_idx]
|
||||
for (uint d = 0; d < params.head_dim; d++) {
|
||||
o_acc[d] = o_acc[d] * exp_diff + exp_score * V_shared[j * HEAD_DIM + 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) {
|
||||
float inv_l = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f;
|
||||
|
||||
for (uint d = 0; d < params.head_dim; d++) {
|
||||
O[base_qo + q_idx * params.head_dim + d] = o_acc[d] * inv_l;
|
||||
}
|
||||
|
||||
// Store log-sum-exp for backward pass
|
||||
uint lse_base = batch_idx * params.num_heads * params.seq_len_q +
|
||||
head_idx * params.seq_len_q;
|
||||
LSE[lse_base + q_idx] = m_i + log(max(l_i, 1e-10f));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user