Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,167 @@
+//
+// Flash Attention Backward Pass - dQ Kernel
+//
+// Computes gradient with respect to Q:
+// dQ = scale * dP @ K
+// where dP = P * (dO @ V^T - D), D = rowsum(dO * O)
+//
+
+#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 dQ gradient
+///
+/// Grid: (num_q_blocks, num_heads, batch_size)
+/// Threadgroup: (BLOCK_Q, 1, 1)
+kernel void flash_attention_backward_dq(
+    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* dQ [[buffer(6)]],            // [batch, heads, seq_q, head_dim]
+    constant FlashBackwardParams& params [[buffer(7)]],
+    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];
+
+    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;
+
+    // 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;
+
+    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 Q, O, dO tiles
+    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];
+            O_shared[tid * HEAD_DIM + d] = O[base_qo + q_idx * params.head_dim + d];
+            dO_shared[tid * HEAD_DIM + d] = dO[base_qo + q_idx * params.head_dim + d];
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Compute D = rowsum(dO * O) for this Q row
+    float D_i = 0.0f;
+    if (q_idx < params.seq_len_q) {
+        for (uint d = 0; d < params.head_dim; d++) {
+            D_i += dO_shared[tid * HEAD_DIM + d] * O_shared[tid * HEAD_DIM + d];
+        }
+    }
+
+    // Get LSE for this row
+    float lse_i = (q_idx < params.seq_len_q) ? LSE[lse_base + q_idx] : 0.0f;
+
+    // Accumulator for dQ
+    float dq_acc[HEAD_DIM];
+    for (uint d = 0; d < params.head_dim; d++) {
+        dq_acc[d] = 0.0f;
+    }
+
+    // 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: skip future blocks
+        if (params.causal != 0 && kv_start > q_start + BLOCK_Q - 1) {
+            break;
+        }
+
+        // Load K, 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];
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Compute dQ contribution from 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;
+                if (kv_idx >= params.seq_len_kv) continue;
+                if (params.causal != 0 && kv_idx > q_idx) continue;
+
+                // Recompute attention score
+                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;
+
+                // Recompute P_ij = exp(score - lse)
+                float p_ij = exp(score - lse_i);
+
+                // 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[tid * HEAD_DIM + d] * V_shared[j * HEAD_DIM + d];
+                }
+                float ds_ij = p_ij * (dov - D_i);
+
+                // Accumulate dQ += dS @ K
+                for (uint d = 0; d < params.head_dim; d++) {
+                    dq_acc[d] += ds_ij * K_shared[j * HEAD_DIM + d];
+                }
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    // Write dQ with scale
+    if (q_idx < params.seq_len_q) {
+        for (uint d = 0; d < params.head_dim; d++) {
+            dQ[base_qo + q_idx * params.head_dim + d] = dq_acc[d] * params.softmax_scale;
+        }
+    }
+}