spb/forge Public MIT
Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.
C++ 61.2%
C 23%
Python 7.6%
TeX 7.2%
CMake 1.1%
1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Tiled flash attention forward using simdgroup_matrix 8x8 fragments.4// Same algorithm and same outputs as flash_attn_fwd in flash_attention.metal5// (online softmax, stores only L = m + log(l)); the difference is that QK^T6// and PV are matrix multiplies over threadgroup-staged tiles instead of7// per-thread scalar dot products, which removes the serial FMA dependency8// chain that capped the scalar kernel at ~1.6 TFLOPs.9//10// Layout (BQ=32, BK=16, 4 simdgroups = 128 threads), split-Q per FA2:11// simdgroup s owns query rows [8s, 8s+8) — exactly one fragment row — and12// holds Q and the O accumulator in registers for the whole kernel13// (2 * HD/8 fragments = 32 floats/lane at HD=64). K/V tiles pass through14// threadgroup memory; the S/P tile round-trips through it so the softmax15// reductions can run on plain threads.16//17// THE PER-ROW RESCALE. Online softmax needs O <- diag(corr) * O every KV18// block, but MSL leaves the element->lane mapping of a simdgroup_matrix19// unspecified, so a lane cannot know which row its registers belong to.20// MLX reverse-engineers the mapping; instead this builds an 8x8 diagonal21// matrix holding corr in threadgroup memory and applies it with an MMA.22// That is spec-clean, costs HD/8 extra MMAs per block (~25% more MMA work,23// measured cheaper than the alternatives), and keeps O in registers — the24// point of the exercise, since staging O in threadgroup memory would add25// 8 KB and halve residency.26//27// Threadgroup memory ~12.7 KB at HD=64: K and V tiles (padded +4 floats per28// row against bank conflicts), the S/P tile, per-row softmax state, and one29// 8x8 diagonal scratch per simdgroup. The K/V tiles are reused as the output30// staging buffer at the end, once they are dead.31#include <metal_stdlib>32#include <metal_simdgroup_matrix>33using namespace metal;3435struct FlashParams {36 uint B, T, H, HKV;37 float scale;38 uint causal;39 uint window; // unused here: the host routes window > 0 to the scalar kernels40};4142// Enumerators, never `constant constexpr`: see RESEARCH.md 7a — the latter is43// a constant-address-space variable, which blocks unrolling and spills every44// fragment to the stack.45enum : uint {46 BQ = 32, // query rows per threadgroup47 BK = 16, // KV rows per iteration48 NSG = BQ / 8, // one 8-row fragment per simdgroup49 MMA_THREADS = NSG * 32,50 KF = BK / 8, // S fragments along the KV axis51 PADF = 4, // 16 bytes52 LDS = BK + PADF, // S/P tile row stride53};5455template <uint HD>56kernel void flash_attn_fwd_mma(device const float* Q [[buffer(0)]],57 device const float* K [[buffer(1)]],58 device const float* V [[buffer(2)]],59 device float* O [[buffer(3)]],60 device float* L [[buffer(4)]],61 constant FlashParams& p [[buffer(5)]],62 uint tgid [[threadgroup_position_in_grid]],63 uint tid [[thread_index_in_threadgroup]],64 uint lane [[thread_index_in_simdgroup]],65 uint sgid [[simdgroup_index_in_threadgroup]]) {66 constexpr uint DF = HD / 8; // fragments along head_dim67 constexpr uint LDKV = HD + PADF;6869 // K and V tiles, reused as output staging after the KV loop.70 threadgroup float kv[2 * BK * LDKV];71 threadgroup float* Ks = kv;72 threadgroup float* Vs = kv + BK * LDKV;73 threadgroup float Sbuf[BQ * LDS];74 threadgroup float row_m[BQ], row_l[BQ], row_corr[BQ];75 threadgroup float diag[NSG * 64];7677 const uint q_blocks = (p.T + BQ - 1) / BQ;78 const uint qb = tgid % q_blocks;79 const uint h = (tgid / q_blocks) % p.H;80 const uint b = tgid / (q_blocks * p.H);8182 const uint hkv = h / (p.H / p.HKV);83 const uint Cq = p.H * HD;84 const uint Ckv = p.HKV * HD;85 const uint q0 = qb * BQ; // first query row of this threadgroup86 const uint sg_row = sgid * 8; // first query row of this simdgroup8788 if (tid < BQ) {89 row_m[tid] = -FLT_MAX;90 row_l[tid] = 0.0f;91 }9293 // Q stays in registers for the whole kernel. It is staged through the94 // (not yet used) K/V tile first so a query block straddling the end of95 // the sequence gets zero-filled rows instead of an out-of-bounds read —96 // loading fragments straight from device would need the block to be a97 // full 8 rows.98 simdgroup_float8x8 Qf[DF], Of[DF];99 {100 threadgroup_barrier(mem_flags::mem_threadgroup);101 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {102 const uint r = e / HD;103 const uint d = e % HD;104 const uint i = q0 + r;105 kv[r * HD + d] =106 (i < p.T) ? Q[(ulong(b) * p.T + i) * Cq + h * HD + d] : 0.0f;107 }108 threadgroup_barrier(mem_flags::mem_threadgroup);109#pragma clang loop unroll(full)110 for (uint d = 0; d < DF; ++d) {111 simdgroup_load(Qf[d], kv + sg_row * HD + d * 8, HD);112 Of[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);113 }114 }115116 // Causal: only KV blocks up to this threadgroup's last query row matter.117 const uint row_max = min(q0 + BQ - 1, p.T - 1);118 const uint j_end = p.causal ? (row_max + 1) : p.T;119120 for (uint jb = 0; jb < j_end; jb += BK) {121 const uint block_len = min(BK, j_end - jb);122123 threadgroup_barrier(mem_flags::mem_threadgroup);124 // Stage K/V; rows past the block length are zeroed so masked columns125 // can never contribute a NaN through 0 * garbage.126 for (uint e = tid; e < BK * HD; e += MMA_THREADS) {127 const uint jr = e / HD;128 const uint d = e % HD;129 const bool ok = jr < block_len;130 const ulong src = (ulong(b) * p.T + jb + jr) * Ckv + hkv * HD + d;131 Ks[jr * LDKV + d] = ok ? K[src] : 0.0f;132 Vs[jr * LDKV + d] = ok ? V[src] : 0.0f;133 }134 threadgroup_barrier(mem_flags::mem_threadgroup);135136 // S = Q @ K^T, accumulated over head_dim fragments.137 simdgroup_float8x8 Sf[KF];138#pragma clang loop unroll(full)139 for (uint j = 0; j < KF; ++j) Sf[j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);140#pragma clang loop unroll(full)141 for (uint d = 0; d < DF; ++d) {142#pragma clang loop unroll(full)143 for (uint j = 0; j < KF; ++j) {144 simdgroup_float8x8 KTf;145 // transpose=true turns the K tile into K^T[d-block][j-block]146 simdgroup_load(KTf, Ks + (j * 8) * LDKV + d * 8, LDKV, 0, true);147 simdgroup_multiply_accumulate(Sf[j], Qf[d], KTf, Sf[j]);148 }149 }150#pragma clang loop unroll(full)151 for (uint j = 0; j < KF; ++j)152 simdgroup_store(Sf[j], Sbuf + sg_row * LDS + j * 8, LDS);153 threadgroup_barrier(mem_flags::mem_threadgroup);154155 // Softmax on plain threads, one row each. This is O(BQ*BK) against156 // O(BQ*BK*HD) of MMA work, so leaving 96 of 128 threads idle costs157 // ~1% and keeps the reduction free of any lane-mapping assumption.158 if (tid < BQ) {159 const uint i = q0 + tid;160 threadgroup float* srow = Sbuf + tid * LDS;161 const uint valid =162 (i >= p.T) ? 0u163 : (p.causal ? min(block_len, (i >= jb) ? (i - jb + 1) : 0u)164 : block_len);165 float bmax = -FLT_MAX;166 for (uint j = 0; j < valid; ++j) bmax = max(bmax, srow[j] * p.scale);167168 const float m_old = row_m[tid];169 const float m_new = max(m_old, bmax);170 const float corr = (m_old == -FLT_MAX) ? 0.0f : exp(m_old - m_new);171 float bsum = 0.0f;172 for (uint j = 0; j < BK; ++j) {173 if (j < valid) {174 const float e = exp(srow[j] * p.scale - m_new);175 srow[j] = e;176 bsum += e;177 } else {178 srow[j] = 0.0f; // masked / out of range179 }180 }181 row_m[tid] = (valid > 0) ? m_new : m_old;182 row_l[tid] = row_l[tid] * ((valid > 0) ? corr : 1.0f) + bsum;183 row_corr[tid] = (valid > 0) ? corr : 1.0f;184 }185 threadgroup_barrier(mem_flags::mem_threadgroup);186187 // O <- diag(corr) @ O, then O += P @ V.188 threadgroup float* dg = diag + sgid * 64;189 for (uint e = lane; e < 64; e += 32) dg[e] = 0.0f;190 simdgroup_barrier(mem_flags::mem_threadgroup);191 if (lane < 8) dg[lane * 8 + lane] = row_corr[sg_row + lane];192 simdgroup_barrier(mem_flags::mem_threadgroup);193194 simdgroup_float8x8 Dg;195 simdgroup_load(Dg, dg, 8);196#pragma clang loop unroll(full)197 for (uint d = 0; d < DF; ++d) {198 simdgroup_float8x8 scaled;199 simdgroup_multiply(scaled, Dg, Of[d]);200 Of[d] = scaled;201 }202#pragma clang loop unroll(full)203 for (uint j = 0; j < KF; ++j) {204 simdgroup_float8x8 Pf;205 simdgroup_load(Pf, Sbuf + sg_row * LDS + j * 8, LDS);206#pragma clang loop unroll(full)207 for (uint d = 0; d < DF; ++d) {208 simdgroup_float8x8 Vf;209 simdgroup_load(Vf, Vs + (j * 8) * LDKV + d * 8, LDKV);210 simdgroup_multiply_accumulate(Of[d], Pf, Vf, Of[d]);211 }212 }213 }214215 // Divide by l and write out. K/V are dead now, so their tile doubles as216 // the [BQ x HD] staging buffer (2 * BK * LDKV >= BQ * HD).217 threadgroup_barrier(mem_flags::mem_threadgroup);218 threadgroup float* Ostage = kv;219#pragma clang loop unroll(full)220 for (uint d = 0; d < DF; ++d)221 simdgroup_store(Of[d], Ostage + sg_row * HD + d * 8, HD);222 threadgroup_barrier(mem_flags::mem_threadgroup);223224 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {225 const uint r = e / HD;226 const uint d = e % HD;227 const uint i = q0 + r;228 if (i >= p.T) continue;229 const float l = row_l[r];230 O[(ulong(b) * p.T + i) * Cq + h * HD + d] =231 (l > 0.0f) ? (Ostage[r * HD + d] / l) : 0.0f;232 }233 if (tid < BQ && (q0 + tid) < p.T) {234 const float l = row_l[tid];235 L[(ulong(b) * p.H + h) * p.T + q0 + tid] =236 (l > 0.0f) ? (row_m[tid] + log(l)) : 0.0f;237 }238}239240#define INSTANTIATE_FLASH_MMA(HD) \241 template [[host_name("flash_attn_fwd_mma_f32_hd" #HD)]] kernel void \242 flash_attn_fwd_mma<HD>(device const float*, device const float*, \243 device const float*, device float*, device float*, \244 constant FlashParams&, uint, uint, uint, uint);245246INSTANTIATE_FLASH_MMA(16)247INSTANTIATE_FLASH_MMA(32)248INSTANTIATE_FLASH_MMA(48)249INSTANTIATE_FLASH_MMA(64)250INSTANTIATE_FLASH_MMA(80)251INSTANTIATE_FLASH_MMA(96)252INSTANTIATE_FLASH_MMA(128)253254// ============================================================ backward (MMA)255//256// Backward needs no online rescale — L from the forward already fixes the257// softmax normalisation — so there is no diagonal-matrix trick here, just258// three matmuls per block with an elementwise step between them.259//260// dQ kernel (parallel over queries, one threadgroup per BQ query rows):261// S = Q @ K^T ; P = exp(S*scale - L_i)262// dP = dO @ V^T ; dS = P * (dP - D_i) * scale263// dQ += dS @ K264// dKV kernel (parallel over KV rows, one threadgroup per BKV rows):265// everything transposed, obtained by swapping the operand roles:266// S^T = K @ Q^T and dP^T = V @ dO^T, so267// dV += P^T @ dO and dK += dS^T @ Q.268//269// Splitting this way keeps dK/dV accumulation private to a threadgroup — no270// atomics, and the result is deterministic. Fragments cut register pressure271// by 4x versus the scalar kernels: an 8x8 fragment is 2 floats per lane, so272// holding K, V, dK, dV as fragments is 64 floats/lane at hd=64 where the273// scalar version needed 256 and spilled.274275enum : uint {276 BQB = 16, // query block consumed per iteration by dKV277 LDQB = BQB + PADF,278};279280template <uint HD>281kernel void flash_attn_bwd_dq_mma(device const float* Q [[buffer(0)]],282 device const float* K [[buffer(1)]],283 device const float* V [[buffer(2)]],284 device const float* dO [[buffer(3)]],285 device const float* L [[buffer(4)]],286 device const float* D [[buffer(5)]],287 device float* dQ [[buffer(6)]],288 constant FlashParams& p [[buffer(7)]],289 uint tgid [[threadgroup_position_in_grid]],290 uint tid [[thread_index_in_threadgroup]],291 uint sgid [[simdgroup_index_in_threadgroup]]) {292 constexpr uint DF = HD / 8;293 constexpr uint LDKV = HD + PADF;294295 threadgroup float kv[2 * BK * LDKV];296 threadgroup float* Ks = kv;297 threadgroup float* Vs = kv + BK * LDKV;298 threadgroup float Sbuf[BQ * LDS];299 threadgroup float Pbuf[BQ * LDS];300 threadgroup float row_L[BQ], row_D[BQ];301302 const uint q_blocks = (p.T + BQ - 1) / BQ;303 const uint qb = tgid % q_blocks;304 const uint h = (tgid / q_blocks) % p.H;305 const uint b = tgid / (q_blocks * p.H);306307 const uint hkv = h / (p.H / p.HKV);308 const uint Cq = p.H * HD;309 const uint Ckv = p.HKV * HD;310 const uint q0 = qb * BQ;311 const uint sg_row = sgid * 8;312313 if (tid < BQ) {314 const uint i = q0 + tid;315 row_L[tid] = (i < p.T) ? L[(ulong(b) * p.H + h) * p.T + i] : 0.0f;316 row_D[tid] = (i < p.T) ? D[(ulong(b) * p.H + h) * p.T + i] : 0.0f;317 }318319 // Q and dO into registers (staged so ragged tails zero-fill).320 simdgroup_float8x8 Qf[DF], dOf[DF], dQf[DF];321 for (uint pass = 0; pass < 2; ++pass) {322 threadgroup_barrier(mem_flags::mem_threadgroup);323 device const float* src = (pass == 0) ? Q : dO;324 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {325 const uint r = e / HD, d = e % HD;326 const uint i = q0 + r;327 kv[r * HD + d] =328 (i < p.T) ? src[(ulong(b) * p.T + i) * Cq + h * HD + d] : 0.0f;329 }330 threadgroup_barrier(mem_flags::mem_threadgroup);331#pragma clang loop unroll(full)332 for (uint d = 0; d < DF; ++d) {333 if (pass == 0) simdgroup_load(Qf[d], kv + sg_row * HD + d * 8, HD);334 else simdgroup_load(dOf[d], kv + sg_row * HD + d * 8, HD);335 }336 }337#pragma clang loop unroll(full)338 for (uint d = 0; d < DF; ++d) dQf[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);339340 const uint row_max = min(q0 + BQ - 1, p.T - 1);341 const uint j_end = p.causal ? (row_max + 1) : p.T;342343 for (uint jb = 0; jb < j_end; jb += BK) {344 const uint block_len = min(BK, j_end - jb);345346 threadgroup_barrier(mem_flags::mem_threadgroup);347 for (uint e = tid; e < BK * HD; e += MMA_THREADS) {348 const uint jr = e / HD, d = e % HD;349 const bool ok = jr < block_len;350 const ulong src = (ulong(b) * p.T + jb + jr) * Ckv + hkv * HD + d;351 Ks[jr * LDKV + d] = ok ? K[src] : 0.0f;352 Vs[jr * LDKV + d] = ok ? V[src] : 0.0f;353 }354 threadgroup_barrier(mem_flags::mem_threadgroup);355356 // S = Q @ K^T and dP = dO @ V^T357 simdgroup_float8x8 Sf[KF], Pf[KF];358#pragma clang loop unroll(full)359 for (uint j = 0; j < KF; ++j) {360 Sf[j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);361 Pf[j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);362 }363#pragma clang loop unroll(full)364 for (uint d = 0; d < DF; ++d) {365#pragma clang loop unroll(full)366 for (uint j = 0; j < KF; ++j) {367 simdgroup_float8x8 KTf, VTf;368 simdgroup_load(KTf, Ks + (j * 8) * LDKV + d * 8, LDKV, 0, true);369 simdgroup_load(VTf, Vs + (j * 8) * LDKV + d * 8, LDKV, 0, true);370 simdgroup_multiply_accumulate(Sf[j], Qf[d], KTf, Sf[j]);371 simdgroup_multiply_accumulate(Pf[j], dOf[d], VTf, Pf[j]);372 }373 }374#pragma clang loop unroll(full)375 for (uint j = 0; j < KF; ++j) {376 simdgroup_store(Sf[j], Sbuf + sg_row * LDS + j * 8, LDS);377 simdgroup_store(Pf[j], Pbuf + sg_row * LDS + j * 8, LDS);378 }379 threadgroup_barrier(mem_flags::mem_threadgroup);380381 // dS = P * (dP - D_i) * scale, written back over Sbuf382 if (tid < BQ) {383 const uint i = q0 + tid;384 threadgroup float* srow = Sbuf + tid * LDS;385 threadgroup const float* prow = Pbuf + tid * LDS;386 const uint valid =387 (i >= p.T) ? 0u388 : (p.causal ? min(block_len, (i >= jb) ? (i - jb + 1) : 0u)389 : block_len);390 const float li = row_L[tid], di = row_D[tid];391 for (uint j = 0; j < BK; ++j) {392 if (j < valid) {393 const float prob = exp(srow[j] * p.scale - li);394 srow[j] = prob * (prow[j] - di) * p.scale;395 } else {396 srow[j] = 0.0f;397 }398 }399 }400 threadgroup_barrier(mem_flags::mem_threadgroup);401402 // dQ += dS @ K403#pragma clang loop unroll(full)404 for (uint j = 0; j < KF; ++j) {405 simdgroup_float8x8 dSf;406 simdgroup_load(dSf, Sbuf + sg_row * LDS + j * 8, LDS);407#pragma clang loop unroll(full)408 for (uint d = 0; d < DF; ++d) {409 simdgroup_float8x8 Kf;410 simdgroup_load(Kf, Ks + (j * 8) * LDKV + d * 8, LDKV);411 simdgroup_multiply_accumulate(dQf[d], dSf, Kf, dQf[d]);412 }413 }414 }415416 threadgroup_barrier(mem_flags::mem_threadgroup);417#pragma clang loop unroll(full)418 for (uint d = 0; d < DF; ++d)419 simdgroup_store(dQf[d], kv + sg_row * HD + d * 8, HD);420 threadgroup_barrier(mem_flags::mem_threadgroup);421 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {422 const uint r = e / HD, d = e % HD;423 const uint i = q0 + r;424 if (i >= p.T) continue;425 dQ[(ulong(b) * p.T + i) * Cq + h * HD + d] += kv[r * HD + d];426 }427}428429// dK and dV are separate kernels. Combined, one thread holds K, V, dK and dV430// as fragments and the compiler spills: gpudebug reported 4352 spilled bytes431// and 111 temp registers for the fused version, which made it the single most432// expensive backward kernel. Split, dV carries K+dV and dK carries K+V+dK, and433// both fit. The cost is recomputing S^T = K @ Q^T in each.434435template <uint HD>436kernel void flash_attn_bwd_dv_mma(device const float* Q [[buffer(0)]],437 device const float* K [[buffer(1)]],438 device const float* dO [[buffer(2)]],439 device const float* L [[buffer(3)]],440 device float* dV [[buffer(4)]],441 constant FlashParams& p [[buffer(5)]],442 uint tgid [[threadgroup_position_in_grid]],443 uint tid [[thread_index_in_threadgroup]],444 uint sgid [[simdgroup_index_in_threadgroup]]) {445 constexpr uint DF = HD / 8;446 constexpr uint LDQ = HD + PADF;447 constexpr uint IF = BQB / 8;448449 threadgroup float qo[2 * BQB * LDQ];450 threadgroup float* Qs = qo;451 threadgroup float* dOs = qo + BQB * LDQ;452 threadgroup float PTbuf[BQ * LDQB];453 threadgroup float col_L[BQB];454455 const uint kv_blocks = (p.T + BQ - 1) / BQ;456 const uint jbk = tgid % kv_blocks;457 const uint hkv = (tgid / kv_blocks) % p.HKV;458 const uint b = tgid / (kv_blocks * p.HKV);459460 const uint rep = p.H / p.HKV;461 const uint Cq = p.H * HD;462 const uint Ckv = p.HKV * HD;463 const uint j0 = jbk * BQ;464 const uint sg_row = sgid * 8;465466 simdgroup_float8x8 Kf[DF], dVf[DF];467 {468 threadgroup_barrier(mem_flags::mem_threadgroup);469 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {470 const uint r = e / HD, d = e % HD;471 const uint j = j0 + r;472 qo[r * HD + d] =473 (j < p.T) ? K[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] : 0.0f;474 }475 threadgroup_barrier(mem_flags::mem_threadgroup);476#pragma clang loop unroll(full)477 for (uint d = 0; d < DF; ++d) {478 simdgroup_load(Kf[d], qo + sg_row * HD + d * 8, HD);479 dVf[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);480 }481 }482483 const uint i_start = p.causal ? (j0 / BQB) * BQB : 0;484 for (uint r = 0; r < rep; ++r) {485 const uint h = hkv * rep + r;486 device const float* Lh = L + (ulong(b) * p.H + h) * p.T;487 for (uint ib = i_start; ib < p.T; ib += BQB) {488 const uint blk = min(BQB, p.T - ib);489 threadgroup_barrier(mem_flags::mem_threadgroup);490 for (uint e = tid; e < BQB * HD; e += MMA_THREADS) {491 const uint ir = e / HD, d = e % HD;492 const bool ok = ir < blk;493 const ulong src = (ulong(b) * p.T + ib + ir) * Cq + h * HD + d;494 Qs[ir * LDQ + d] = ok ? Q[src] : 0.0f;495 dOs[ir * LDQ + d] = ok ? dO[src] : 0.0f;496 }497 if (tid < BQB) col_L[tid] = (tid < blk) ? Lh[ib + tid] : 0.0f;498 threadgroup_barrier(mem_flags::mem_threadgroup);499500 simdgroup_float8x8 STf[IF];501#pragma clang loop unroll(full)502 for (uint i = 0; i < IF; ++i)503 STf[i] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);504#pragma clang loop unroll(full)505 for (uint d = 0; d < DF; ++d) {506#pragma clang loop unroll(full)507 for (uint i = 0; i < IF; ++i) {508 simdgroup_float8x8 QTf;509 simdgroup_load(QTf, Qs + (i * 8) * LDQ + d * 8, LDQ, 0, true);510 simdgroup_multiply_accumulate(STf[i], Kf[d], QTf, STf[i]);511 }512 }513#pragma clang loop unroll(full)514 for (uint i = 0; i < IF; ++i)515 simdgroup_store(STf[i], PTbuf + sg_row * LDQB + i * 8, LDQB);516 threadgroup_barrier(mem_flags::mem_threadgroup);517518 if (tid < BQ) {519 const uint j = j0 + tid;520 threadgroup float* prow = PTbuf + tid * LDQB;521 for (uint i = 0; i < BQB; ++i) {522 const uint iq = ib + i;523 const bool ok = (i < blk) && (j < p.T) && (!p.causal || iq >= j);524 prow[i] = ok ? exp(prow[i] * p.scale - col_L[i]) : 0.0f;525 }526 }527 threadgroup_barrier(mem_flags::mem_threadgroup);528529#pragma clang loop unroll(full)530 for (uint i = 0; i < IF; ++i) {531 simdgroup_float8x8 PTx;532 simdgroup_load(PTx, PTbuf + sg_row * LDQB + i * 8, LDQB);533#pragma clang loop unroll(full)534 for (uint d = 0; d < DF; ++d) {535 simdgroup_float8x8 dOf;536 simdgroup_load(dOf, dOs + (i * 8) * LDQ + d * 8, LDQ);537 simdgroup_multiply_accumulate(dVf[d], PTx, dOf, dVf[d]);538 }539 }540 }541 }542543 threadgroup_barrier(mem_flags::mem_threadgroup);544#pragma clang loop unroll(full)545 for (uint d = 0; d < DF; ++d)546 simdgroup_store(dVf[d], qo + sg_row * HD + d * 8, HD);547 threadgroup_barrier(mem_flags::mem_threadgroup);548 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {549 const uint r = e / HD, d = e % HD;550 const uint j = j0 + r;551 if (j >= p.T) continue;552 dV[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] += qo[r * HD + d];553 }554}555556template <uint HD>557kernel void flash_attn_bwd_dk_mma(device const float* Q [[buffer(0)]],558 device const float* K [[buffer(1)]],559 device const float* V [[buffer(2)]],560 device const float* dO [[buffer(3)]],561 device const float* L [[buffer(4)]],562 device const float* D [[buffer(5)]],563 device float* dK [[buffer(6)]],564 constant FlashParams& p [[buffer(7)]],565 uint tgid [[threadgroup_position_in_grid]],566 uint tid [[thread_index_in_threadgroup]],567 uint sgid [[simdgroup_index_in_threadgroup]]) {568 constexpr uint DF = HD / 8;569 constexpr uint LDQ = HD + PADF;570 constexpr uint IF = BQB / 8;571572 threadgroup float qo[2 * BQB * LDQ];573 threadgroup float* Qs = qo;574 threadgroup float* dOs = qo + BQB * LDQ;575 threadgroup float STbuf[BQ * LDQB];576 threadgroup float PTbuf[BQ * LDQB];577 threadgroup float col_L[BQB], col_D[BQB];578579 const uint kv_blocks = (p.T + BQ - 1) / BQ;580 const uint jbk = tgid % kv_blocks;581 const uint hkv = (tgid / kv_blocks) % p.HKV;582 const uint b = tgid / (kv_blocks * p.HKV);583584 const uint rep = p.H / p.HKV;585 const uint Cq = p.H * HD;586 const uint Ckv = p.HKV * HD;587 const uint j0 = jbk * BQ;588 const uint sg_row = sgid * 8;589590 simdgroup_float8x8 Kf[DF], Vf[DF], dKf[DF];591 for (uint pass = 0; pass < 2; ++pass) {592 threadgroup_barrier(mem_flags::mem_threadgroup);593 device const float* src = (pass == 0) ? K : V;594 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {595 const uint r = e / HD, d = e % HD;596 const uint j = j0 + r;597 qo[r * HD + d] =598 (j < p.T) ? src[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] : 0.0f;599 }600 threadgroup_barrier(mem_flags::mem_threadgroup);601#pragma clang loop unroll(full)602 for (uint d = 0; d < DF; ++d) {603 if (pass == 0) simdgroup_load(Kf[d], qo + sg_row * HD + d * 8, HD);604 else simdgroup_load(Vf[d], qo + sg_row * HD + d * 8, HD);605 }606 }607#pragma clang loop unroll(full)608 for (uint d = 0; d < DF; ++d) dKf[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);609610 const uint i_start = p.causal ? (j0 / BQB) * BQB : 0;611 for (uint r = 0; r < rep; ++r) {612 const uint h = hkv * rep + r;613 device const float* Lh = L + (ulong(b) * p.H + h) * p.T;614 device const float* Dh = D + (ulong(b) * p.H + h) * p.T;615 for (uint ib = i_start; ib < p.T; ib += BQB) {616 const uint blk = min(BQB, p.T - ib);617 threadgroup_barrier(mem_flags::mem_threadgroup);618 for (uint e = tid; e < BQB * HD; e += MMA_THREADS) {619 const uint ir = e / HD, d = e % HD;620 const bool ok = ir < blk;621 const ulong src = (ulong(b) * p.T + ib + ir) * Cq + h * HD + d;622 Qs[ir * LDQ + d] = ok ? Q[src] : 0.0f;623 dOs[ir * LDQ + d] = ok ? dO[src] : 0.0f;624 }625 if (tid < BQB) {626 col_L[tid] = (tid < blk) ? Lh[ib + tid] : 0.0f;627 col_D[tid] = (tid < blk) ? Dh[ib + tid] : 0.0f;628 }629 threadgroup_barrier(mem_flags::mem_threadgroup);630631 simdgroup_float8x8 STf[IF], PTf[IF];632#pragma clang loop unroll(full)633 for (uint i = 0; i < IF; ++i) {634 STf[i] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);635 PTf[i] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);636 }637#pragma clang loop unroll(full)638 for (uint d = 0; d < DF; ++d) {639#pragma clang loop unroll(full)640 for (uint i = 0; i < IF; ++i) {641 simdgroup_float8x8 QTf, dOTf;642 simdgroup_load(QTf, Qs + (i * 8) * LDQ + d * 8, LDQ, 0, true);643 simdgroup_load(dOTf, dOs + (i * 8) * LDQ + d * 8, LDQ, 0, true);644 simdgroup_multiply_accumulate(STf[i], Kf[d], QTf, STf[i]);645 simdgroup_multiply_accumulate(PTf[i], Vf[d], dOTf, PTf[i]);646 }647 }648#pragma clang loop unroll(full)649 for (uint i = 0; i < IF; ++i) {650 simdgroup_store(STf[i], STbuf + sg_row * LDQB + i * 8, LDQB);651 simdgroup_store(PTf[i], PTbuf + sg_row * LDQB + i * 8, LDQB);652 }653 threadgroup_barrier(mem_flags::mem_threadgroup);654655 if (tid < BQ) {656 const uint j = j0 + tid;657 threadgroup float* srow = STbuf + tid * LDQB;658 threadgroup const float* prow = PTbuf + tid * LDQB;659 for (uint i = 0; i < BQB; ++i) {660 const uint iq = ib + i;661 const bool ok = (i < blk) && (j < p.T) && (!p.causal || iq >= j);662 srow[i] = ok ? exp(srow[i] * p.scale - col_L[i]) *663 (prow[i] - col_D[i]) * p.scale664 : 0.0f;665 }666 }667 threadgroup_barrier(mem_flags::mem_threadgroup);668669#pragma clang loop unroll(full)670 for (uint i = 0; i < IF; ++i) {671 simdgroup_float8x8 STx;672 simdgroup_load(STx, STbuf + sg_row * LDQB + i * 8, LDQB);673#pragma clang loop unroll(full)674 for (uint d = 0; d < DF; ++d) {675 simdgroup_float8x8 Qf2;676 simdgroup_load(Qf2, Qs + (i * 8) * LDQ + d * 8, LDQ);677 simdgroup_multiply_accumulate(dKf[d], STx, Qf2, dKf[d]);678 }679 }680 }681 }682683 threadgroup_barrier(mem_flags::mem_threadgroup);684#pragma clang loop unroll(full)685 for (uint d = 0; d < DF; ++d)686 simdgroup_store(dKf[d], qo + sg_row * HD + d * 8, HD);687 threadgroup_barrier(mem_flags::mem_threadgroup);688 for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {689 const uint r = e / HD, d = e % HD;690 const uint j = j0 + r;691 if (j >= p.T) continue;692 dK[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] += qo[r * HD + d];693 }694}695696#define INSTANTIATE_FLASH_MMA_BWD(HD) \697 template [[host_name("flash_attn_bwd_dq_mma_f32_hd" #HD)]] kernel void \698 flash_attn_bwd_dq_mma<HD>(device const float*, device const float*, \699 device const float*, device const float*, \700 device const float*, device const float*, \701 device float*, constant FlashParams&, uint, uint, \702 uint); \703 template [[host_name("flash_attn_bwd_dv_mma_f32_hd" #HD)]] kernel void \704 flash_attn_bwd_dv_mma<HD>(device const float*, device const float*, \705 device const float*, device const float*, \706 device float*, constant FlashParams&, uint, uint, \707 uint); \708 template [[host_name("flash_attn_bwd_dk_mma_f32_hd" #HD)]] kernel void \709 flash_attn_bwd_dk_mma<HD>(device const float*, device const float*, \710 device const float*, device const float*, \711 device const float*, device const float*, \712 device float*, constant FlashParams&, uint, uint, \713 uint);714715INSTANTIATE_FLASH_MMA_BWD(16)716INSTANTIATE_FLASH_MMA_BWD(32)717INSTANTIATE_FLASH_MMA_BWD(48)718INSTANTIATE_FLASH_MMA_BWD(64)719INSTANTIATE_FLASH_MMA_BWD(80)720INSTANTIATE_FLASH_MMA_BWD(96)721INSTANTIATE_FLASH_MMA_BWD(128)722