// // Flash Attention Backward Pass - dK/dV Kernel // // Computes gradients with respect to K and V: // dV = P^T @ dO // dK = scale * dS^T @ Q // where dS = P * (dO @ V^T - D) // #include using namespace metal; // Apple Silicon GPUs have 32KB threadgroup memory 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 backward pass struct FlashBackwardParams { uint batch_size; uint num_heads; uint seq_len_q; uint seq_len_kv; uint head_dim; float softmax_scale; uint causal; }; /// Compute dK and dV gradients /// /// Grid: (num_kv_blocks, num_heads, batch_size) /// Threadgroup: (BLOCK_KV, 1, 1) kernel void flash_attention_backward_dkv( 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 const float* O [[buffer(3)]], // [batch, heads, seq_q, head_dim] device const float* dO [[buffer(4)]], // [batch, heads, seq_q, head_dim] device const float* LSE [[buffer(5)]], // [batch, heads, seq_q] device float* dK [[buffer(6)]], // [batch, heads, seq_kv, head_dim] device float* dV [[buffer(7)]], // [batch, heads, seq_kv, head_dim] constant FlashBackwardParams& params [[buffer(8)]], uint3 tgid [[threadgroup_position_in_grid]], uint tid [[thread_index_in_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]; threadgroup float dO_shared[BLOCK_Q * HEAD_DIM]; threadgroup float O_shared[BLOCK_Q * HEAD_DIM]; threadgroup float LSE_shared[BLOCK_Q]; threadgroup float D_shared[BLOCK_Q]; uint batch_idx = tgid.z; uint head_idx = tgid.y; uint kv_block = tgid.x; uint kv_start = kv_block * BLOCK_KV; uint kv_idx = kv_start + tid; // Memory strides uint stride_q_batch = params.num_heads * params.seq_len_q * params.head_dim; uint stride_q_head = params.seq_len_q * params.head_dim; uint base_q = batch_idx * stride_q_batch + head_idx * stride_q_head; 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; uint lse_base = batch_idx * params.num_heads * params.seq_len_q + head_idx * params.seq_len_q; // Load K, V for this block if (kv_idx < params.seq_len_kv) { for (uint d = 0; d < params.head_dim; d++) { K_shared[tid * HEAD_DIM + d] = K[base_kv + kv_idx * params.head_dim + d]; V_shared[tid * HEAD_DIM + d] = V[base_kv + kv_idx * params.head_dim + d]; } } threadgroup_barrier(mem_flags::mem_threadgroup); // Accumulators for dK and dV float dk_acc[HEAD_DIM]; float dv_acc[HEAD_DIM]; for (uint d = 0; d < params.head_dim; d++) { dk_acc[d] = 0.0f; dv_acc[d] = 0.0f; } // Iterate over Q blocks uint num_q_blocks = (params.seq_len_q + BLOCK_Q - 1) / BLOCK_Q; for (uint q_block = 0; q_block < num_q_blocks; q_block++) { uint q_start = q_block * BLOCK_Q; // Causal: skip if all Q in this block come after our K/V if (params.causal != 0 && q_start > kv_start + BLOCK_KV - 1) { // For causal, Q rows before kv_idx attend to us // If q_start > kv_end, none of the Q rows in this block attend to us // But we need to continue because later Q blocks might // Actually, for dK/dV we accumulate from all Q that attend to K[kv_idx] // So if q_start > kv_idx, this Q block doesn't attend to our K row } // Causal early exit for dK/dV: if all Q in block are after this K, skip // For causal attention, Q[i] only attends to K[j] where j <= i // So K[kv_idx] receives attention from Q[i] where i >= kv_idx if (params.causal != 0 && q_start + BLOCK_Q - 1 < kv_start) { continue; // All Q in this block are before our K range } // Load Q, O, dO, LSE tiles cooperatively for (uint i = tid; i < BLOCK_Q * params.head_dim; i += BLOCK_KV) { uint q_row = i / params.head_dim; uint d = i % params.head_dim; uint q_i = q_start + q_row; if (q_i < params.seq_len_q) { Q_shared[q_row * HEAD_DIM + d] = Q[base_q + q_i * params.head_dim + d]; O_shared[q_row * HEAD_DIM + d] = O[base_q + q_i * params.head_dim + d]; dO_shared[q_row * HEAD_DIM + d] = dO[base_q + q_i * params.head_dim + d]; } } // Load LSE for Q block if (tid < BLOCK_Q) { uint q_i = q_start + tid; if (q_i < params.seq_len_q) { LSE_shared[tid] = LSE[lse_base + q_i]; // Compute D = rowsum(dO * O) for this Q row float d_val = 0.0f; for (uint d = 0; d < params.head_dim; d++) { d_val += dO_shared[tid * HEAD_DIM + d] * O_shared[tid * HEAD_DIM + d]; } D_shared[tid] = d_val; } } threadgroup_barrier(mem_flags::mem_threadgroup); // Compute dK/dV contributions from this Q block if (kv_idx < params.seq_len_kv) { for (uint i = 0; i < BLOCK_Q; i++) { uint q_i = q_start + i; if (q_i >= params.seq_len_q) continue; // Causal: only process if Q[q_i] attends to K[kv_idx] if (params.causal != 0 && kv_idx > q_i) continue; // Recompute attention score float score = 0.0f; for (uint d = 0; d < params.head_dim; d++) { score += Q_shared[i * HEAD_DIM + d] * K_shared[tid * HEAD_DIM + d]; } score *= params.softmax_scale; // Recompute P = exp(score - lse) float p_ij = exp(score - LSE_shared[i]); // dV += P^T @ dO = P_ij * dO[q_i] for (uint d = 0; d < params.head_dim; d++) { dv_acc[d] += p_ij * dO_shared[i * HEAD_DIM + d]; } // Compute dS_ij = P_ij * (dO @ V^T - D) float dov = 0.0f; for (uint d = 0; d < params.head_dim; d++) { dov += dO_shared[i * HEAD_DIM + d] * V_shared[tid * HEAD_DIM + d]; } float ds_ij = p_ij * (dov - D_shared[i]); // dK += dS^T @ Q = dS_ij * Q[q_i] for (uint d = 0; d < params.head_dim; d++) { dk_acc[d] += ds_ij * Q_shared[i * HEAD_DIM + d]; } } } threadgroup_barrier(mem_flags::mem_threadgroup); } // Write dK and dV with scale if (kv_idx < params.seq_len_kv) { for (uint d = 0; d < params.head_dim; d++) { dK[base_kv + kv_idx * params.head_dim + d] = dk_acc[d] * params.softmax_scale; dV[base_kv + kv_idx * params.head_dim + d] = dv_acc[d]; } } }