Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,196 @@
|
||||
//
|
||||
// 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 <metal_stdlib>
|
||||
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];
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user