// // Flash Attention Forward Pass - Metal Shader // // Implements the Flash Attention algorithm with online softmax // for memory-efficient attention computation. // #include 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)); } }