src/cpu/flash_attention.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // CPU backend — flash-attention family (CHUNK 6). | ||
| 2 | // | ||
| 3 | // Ground truth: src/cuda/flash_attention.cu and | ||
| 4 | // src/cuda/flash_attention_backward.cu. | ||
| 5 | // | ||
| 6 | // Six ops: | ||
| 7 | // * flash_attention_forward — SDPA over pre-projected QKV. | ||
| 8 | // * flash_attention_qkvo_forward — fused QKVO projections + SDPA. | ||
| 9 | // * flash_attention_qkvo_backward — full backward of the above. | ||
| 10 | // * flash_attention_backward — backward over pre-projected QKV. | ||
| 11 | // * flash_attention_project_kv — project ctx → K_out, V_out. | ||
| 12 | // * flash_attention_q_with_kv_cached_forward — project X→Q, attend cached K/V, | ||
| 13 | // project through Wo. | ||
| 14 | // | ||
| 15 | // DTYPE DECISIONS | ||
| 16 | // Every CUDA op in this family runs FP16 internally (Q/K/V/weights all FP16, | ||
| 17 | // attention core in FP16). The CPU backend is FP32-only (per CLAUDE.md), so | ||
| 18 | // all six CPU impls run the straightforward FP32 scalar math that produces | ||
| 19 | // the SAME mathematical result as the GPU's tiled flash kernels — flash | ||
| 20 | // tiling is a memory optimisation, not a different result, so a plain | ||
| 21 | // materialised attention is correct. The parity tests quantise inputs | ||
| 22 | // through FP16 (so both backends start identical), feed FP16 to the GPU and | ||
| 23 | // FP32 to the CPU, and compare with a loose FP16-scale tolerance. | ||
| 24 | // | ||
| 25 | // CONVENTIONS (verified against flash_attention.cu / *_backward.cu) | ||
| 26 | // * Q/K/V are (Lq|Lk, D) with D = num_heads * head_dim, the head dimension | ||
| 27 | // contiguous within each row: element (l, h*head_dim + d). | ||
| 28 | // * Softmax scale: 1/sqrt(head_dim). | ||
| 29 | // * Mask: length-Lk key-validity buffer (1=valid, 0=invalid). mask[k] <= 0.5 | ||
| 30 | // drops key k from the softmax (score → -inf). A fully-masked row yields a | ||
| 31 | // zero output row (rsum == 0 ⇒ inv == 0). Mask never gates query rows. | ||
| 32 | // * Causal: key k contributes to query q only when k <= q (requires Lq==Lk). | ||
| 33 | // * Projection weight layout (linear_forward_batched_fp16): Wq/Wo are (D, D); | ||
| 34 | // Wk/Wv are (D, D_ctx). out(i, n) = sum_k In(i, k) * W(n, k) + b(n) — i.e. | ||
| 35 | // In @ W^T, optional bias broadcast over rows. | ||
| 36 | // * flash_attention_qkvo_forward: Ctx==null ⇒ self-attention (kv_src = X); | ||
| 37 | // non-null ⇒ cross-attention. | ||
| 38 | // | ||
| 39 | // ACCUMULATION (verified against the CUDA backward kernels) | ||
| 40 | // * flash_attention_backward: dQ/dK/dV OVERWRITTEN (CUDA zeros then | ||
| 41 | // fills each per-head slot). | ||
| 42 | // * flash_attention_qkvo_backward: dX OVERWRITTEN (CUDA zeros it, then the | ||
| 43 | // single fa_fp16_add of dX_from_Q makes | ||
| 44 | // it equal that path; self-attn adds the | ||
| 45 | // K/V paths on top). dCtx OVERWRITTEN | ||
| 46 | // (zeroed then K+V paths added). | ||
| 47 | // dWq/dWk/dWv/dWo and dbq/dbk/dbv/dbo | ||
| 48 | // ACCUMULATE (+=) — linear_backward_batched | ||
| 49 | // folds into the caller's grad buffers. | ||
| 50 | |||
| 51 | #include <brotensor/tensor.h> | ||
| 52 | #include <brotensor/detail/cpu/thread_pool.h> | ||
| 53 | |||
| 54 | #include <cmath> | ||
| 55 | #include <stdexcept> | ||
| 56 | #include <vector> | ||
| 57 | |||
| 58 | namespace brotensor::detail::cpu { | ||
| 59 | |||
| 60 | namespace { | ||
| 61 | |||
| 62 | using ::brotensor::Tensor; | ||
| 63 | using ::brotensor::Dtype; | ||
| 64 | |||
| 65 | // out(i, n) = sum_k In(i, k) * W(n, k) + (bias ? bias[n] : 0). | ||
| 66 | // In : (M, Din) W : (Dout, Din) bias : (Dout,) optional out : (M, Dout) | ||
| 67 | // Each token i owns row i of Out exclusively (In/W/bias read-only), so the | ||
| 68 | // token axis parallelizes with no cross-thread writes. | ||
| 69 | 107 | void linear_proj(const float* In, const float* W, const float* bias, | |
| 70 | float* Out, int M, int Din, int Dout) { | ||
| 71 |
1/2✓ Branch 0 taken 107 times.
✗ Branch 1 not taken.
|
1027 | parallel_for(static_cast<std::size_t>(M), [&](std::size_t ii) { |
| 72 | 920 | const int i = static_cast<int>(ii); | |
| 73 | 920 | const float* xr = In + static_cast<std::size_t>(i) * Din; | |
| 74 | 920 | float* orow = Out + static_cast<std::size_t>(i) * Dout; | |
| 75 |
2/2✓ Branch 0 taken 30433 times.
✓ Branch 1 taken 920 times.
|
31353 | for (int n = 0; n < Dout; ++n) { |
| 76 | 30433 | const float* wr = W + static_cast<std::size_t>(n) * Din; | |
| 77 |
2/2✓ Branch 0 taken 13786 times.
✓ Branch 1 taken 16647 times.
|
30433 | float acc = bias ? bias[n] : 0.0f; |
| 78 |
2/2✓ Branch 0 taken 877549 times.
✓ Branch 1 taken 30433 times.
|
907982 | for (int k = 0; k < Din; ++k) acc += xr[k] * wr[k]; |
| 79 | 30433 | orow[n] = acc; | |
| 80 | 30433 | } | |
| 81 | 920 | }); | |
| 82 | 107 | } | |
| 83 | |||
| 84 | // Backward of linear_proj. forward: Out = In @ W^T + b. | ||
| 85 | // dIn(i, k) = sum_n dOut(i, n) * W(n, k) — accumulate into dIn. | ||
| 86 | // dW(n, k) += sum_i dOut(i, n) * In(i, k) — accumulate. | ||
| 87 | // db(n) += sum_i dOut(i, n) — accumulate (if db != null). | ||
| 88 | 36 | void linear_proj_backward(const float* In, const float* W, const float* dOut, | |
| 89 | float* dIn, float* dW, float* db, | ||
| 90 | int M, int Din, int Dout) { | ||
| 91 | // Each i owns row i of dIn exclusively: the "+=" here accumulates onto | ||
| 92 | // that row's pre-existing value (from before this call), never onto | ||
| 93 | // another i's row, so this parallelizes over i (the token axis). | ||
| 94 |
1/2✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
|
337 | parallel_for(static_cast<std::size_t>(M), [&](std::size_t ii) { |
| 95 | 301 | const int i = static_cast<int>(ii); | |
| 96 | 301 | const float* dor = dOut + static_cast<std::size_t>(i) * Dout; | |
| 97 | 301 | float* dir = dIn + static_cast<std::size_t>(i) * Din; | |
| 98 |
2/2✓ Branch 0 taken 8656 times.
✓ Branch 1 taken 301 times.
|
8957 | for (int k = 0; k < Din; ++k) { |
| 99 | 8656 | float acc = 0.0f; | |
| 100 |
2/2✓ Branch 0 taken 217566 times.
✓ Branch 1 taken 8656 times.
|
226222 | for (int n = 0; n < Dout; ++n) |
| 101 | 217566 | acc += dor[n] * W[static_cast<std::size_t>(n) * Din + k]; | |
| 102 | 8656 | dir[k] += acc; | |
| 103 | 8656 | } | |
| 104 | 301 | }); | |
| 105 | // dW/db: n-outer with a private serial sum over i (the reduction axis) — | ||
| 106 | // each n owns row n of dW and scalar db[n] exclusively, written exactly | ||
| 107 | // once, so this parallelizes over n (Dout). | ||
| 108 |
1/2✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
|
1360 | parallel_for(static_cast<std::size_t>(Dout), [&](std::size_t ni) { |
| 109 | 1324 | const int n = static_cast<int>(ni); | |
| 110 | 1324 | float dbacc = 0.0f; | |
| 111 |
2/2✓ Branch 0 taken 34859 times.
✓ Branch 1 taken 1324 times.
|
36183 | for (int k = 0; k < Din; ++k) { |
| 112 | 34859 | float dw = 0.0f; | |
| 113 |
2/2✓ Branch 0 taken 240473 times.
✓ Branch 1 taken 34859 times.
|
275332 | for (int i = 0; i < M; ++i) |
| 114 | 480946 | dw += dOut[static_cast<std::size_t>(i) * Dout + n] * | |
| 115 | 240473 | In[static_cast<std::size_t>(i) * Din + k]; | |
| 116 | 34859 | dW[static_cast<std::size_t>(n) * Din + k] += dw; | |
| 117 | 34859 | } | |
| 118 |
2/2✓ Branch 0 taken 10966 times.
✓ Branch 1 taken 1324 times.
|
12290 | for (int i = 0; i < M; ++i) |
| 119 | 10966 | dbacc += dOut[static_cast<std::size_t>(i) * Dout + n]; | |
| 120 |
2/2✓ Branch 0 taken 433 times.
✓ Branch 1 taken 891 times.
|
1324 | if (db) db[n] += dbacc; |
| 121 | 1324 | }); | |
| 122 | 36 | } | |
| 123 | |||
| 124 | // Materialised multi-head attention over pre-projected Q/K/V. | ||
| 125 | // Q : (Lq, D) K, V : (Lk, D) D = H * hd. | ||
| 126 | // O : (Lq, D) — written. | ||
| 127 | // P (optional, may be null): (H*Lq, Lk) per-head softmax probabilities. | ||
| 128 | // scale = 1/sqrt(hd); key mask + causal applied during the row softmax. | ||
| 129 | // window > 0 adds a sliding causal band: query q attends only keys | ||
| 130 | // k in [q-window+1, q] (window <= 0 leaves the band unbounded). | ||
| 131 | // Q/O are H heads wide (Dq = H*hd); K/V may be grouped (Dkv = (H/group)*hd, | ||
| 132 | // GQA) — query head h reads K/V head h/group. group == 1 / Dkv < 0 is plain MHA. | ||
| 133 | // scratch_scores, if non-null, must point at a caller-owned buffer of at | ||
| 134 | // least Lk floats and is reused verbatim (no allocation here) — callers | ||
| 135 | // looping over many short sequences (e.g. varlen attention) hoist one buffer | ||
| 136 | // sized to the batch's max Lk instead of paying a heap allocation per call. | ||
| 137 | // A null scratch_scores falls back to a local per-call allocation, matching | ||
| 138 | // the original single-call behavior. | ||
| 139 | // | ||
| 140 | // THREADING: when scratch_scores is null (every caller except the varlen | ||
| 141 | // family), the outer q loop is safe to parallelize because each q's row of | ||
| 142 | // O/P is written exclusively by that q (see per-iteration comments below); | ||
| 143 | // the only shared mutable state is the `scores` scratch buffer, so the | ||
| 144 | // parallel path gives every q its own private Lk-wide slice of a | ||
| 145 | // freshly-allocated Lq*Lk buffer instead of one buffer reused across q. | ||
| 146 | // When scratch_scores is non-null, the caller (flash_attention_varlen_*) | ||
| 147 | // has explicitly hoisted ONE Lk-sized buffer across many sequential calls | ||
| 148 | // to this function (one per sequence in a batch) specifically to avoid | ||
| 149 | // per-sequence allocation; parallelizing the q loop in that case would have | ||
| 150 | // every thread's q race on that single shared buffer, so this path is left | ||
| 151 | // exactly as it was (serial, reusing the caller's buffer verbatim) — see | ||
| 152 | // flash_attention_varlen_forward's own comment for why it stays | ||
| 153 | // single-threaded. | ||
| 154 | 88 | void attention_core(const float* Q, const float* K, const float* V, | |
| 155 | const float* mask, int Lq, int Lk, int D, int H, | ||
| 156 | bool causal, float* O, float* P, int window = 0, | ||
| 157 | int q_offset = 0, int Dkv = -1, int group = 1, | ||
| 158 | float* scratch_scores = nullptr) { | ||
| 159 | 88 | const int hd = D / H; | |
| 160 |
2/2✓ Branch 0 taken 9 times.
✓ Branch 1 taken 79 times.
|
88 | if (Dkv < 0) Dkv = D; |
| 161 | 88 | const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(hd)); | |
| 162 | |||
| 163 | // Per-q body. `scores` must be a private Lk-wide buffer when running | ||
| 164 | // concurrently with other q's (see threading note above) — each q | ||
| 165 | // writes only row q of O (and, if requested, the q-th row of every | ||
| 166 | // head's P segment), so distinct q's never touch the same output bytes. | ||
| 167 | 1543 | auto process_q = [&](int q, float* scores) { | |
| 168 | 1455 | const int aq = q + q_offset; // absolute causal position of this query | |
| 169 |
2/2✓ Branch 0 taken 4386 times.
✓ Branch 1 taken 1455 times.
|
5841 | for (int h = 0; h < H; ++h) { |
| 170 | 4386 | const int off = h * hd; // Q/O head offset (Dq-wide) | |
| 171 | 4386 | const int off_kv = (h / group) * hd; // K/V head offset (Dkv-wide) | |
| 172 | 4386 | float m = -1e30f; | |
| 173 |
2/2✓ Branch 0 taken 159552 times.
✓ Branch 1 taken 4386 times.
|
163938 | for (int k = 0; k < Lk; ++k) { |
| 174 |
4/4✓ Branch 0 taken 71005 times.
✓ Branch 1 taken 88547 times.
✓ Branch 2 taken 50093 times.
✓ Branch 3 taken 20912 times.
|
159552 | if (mask && mask[k] <= 0.5f) { scores[k] = -1e30f; continue; } |
| 175 |
4/4✓ Branch 0 taken 3419 times.
✓ Branch 1 taken 135221 times.
✓ Branch 2 taken 1997 times.
✓ Branch 3 taken 1422 times.
|
138640 | if (causal && k > aq) { scores[k] = -1e30f; continue; } |
| 176 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 137218 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
137218 | if (window > 0 && k <= aq - window) { scores[k] = -1e30f; continue; } |
| 177 | 137218 | const float* qr = Q + static_cast<std::size_t>(q) * D + off; | |
| 178 | 137218 | const float* kr = K + static_cast<std::size_t>(k) * Dkv + off_kv; | |
| 179 | 137218 | float dot = 0.0f; | |
| 180 |
2/2✓ Branch 0 taken 4162284 times.
✓ Branch 1 taken 137218 times.
|
4299502 | for (int d = 0; d < hd; ++d) dot += qr[d] * kr[d]; |
| 181 | 137218 | const float s = dot * inv_sqrt; | |
| 182 | 137218 | scores[k] = s; | |
| 183 |
2/2✓ Branch 0 taken 121536 times.
✓ Branch 1 taken 15682 times.
|
137218 | if (s > m) m = s; |
| 184 | 137218 | } | |
| 185 | 4386 | const bool empty = (m <= -1e29f); | |
| 186 | 4386 | float sum = 0.0f; | |
| 187 |
2/2✓ Branch 0 taken 177303 times.
✓ Branch 1 taken 4386 times.
|
181689 | for (int k = 0; k < Lk; ++k) { |
| 188 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 177303 times.
|
177303 | const float e = empty ? 0.0f : std::exp(scores[k] - m); |
| 189 | 177303 | scores[k] = e; | |
| 190 | 177303 | sum += e; | |
| 191 | 177303 | } | |
| 192 |
1/2✓ Branch 0 taken 4386 times.
✗ Branch 1 not taken.
|
4386 | const float inv = sum > 0.0f ? 1.0f / sum : 0.0f; |
| 193 | 4386 | float* orow = O + static_cast<std::size_t>(q) * D + off; | |
| 194 |
2/2✓ Branch 0 taken 92890 times.
✓ Branch 1 taken 4386 times.
|
97276 | for (int d = 0; d < hd; ++d) { |
| 195 | 92890 | float acc = 0.0f; | |
| 196 |
2/2✓ Branch 0 taken 4226756 times.
✓ Branch 1 taken 92890 times.
|
4319646 | for (int k = 0; k < Lk; ++k) { |
| 197 | 8453512 | acc += scores[k] * inv * | |
| 198 | 4226756 | V[static_cast<std::size_t>(k) * Dkv + off_kv + d]; | |
| 199 | 4226756 | } | |
| 200 | 92890 | orow[d] = acc; | |
| 201 | 92890 | } | |
| 202 |
1/2✓ Branch 0 taken 4386 times.
✗ Branch 1 not taken.
|
4386 | if (P) { |
| 203 | ✗ | float* prow = | |
| 204 | ✗ | P + (static_cast<std::size_t>(h) * Lq + q) * Lk; | |
| 205 | ✗ | for (int k = 0; k < Lk; ++k) prow[k] = scores[k] * inv; | |
| 206 | ✗ | } | |
| 207 | 4386 | } | |
| 208 | 1455 | }; | |
| 209 | |||
| 210 |
2/2✓ Branch 0 taken 33 times.
✓ Branch 1 taken 55 times.
|
88 | if (scratch_scores) { |
| 211 | // Varlen path: caller owns a single Lk-sized buffer reused across | ||
| 212 | // many sequential calls to this function. Stay serial (see the | ||
| 213 | // threading note above). | ||
| 214 |
2/2✓ Branch 0 taken 225 times.
✓ Branch 1 taken 33 times.
|
258 | for (int q = 0; q < Lq; ++q) process_q(q, scratch_scores); |
| 215 | 33 | } else { | |
| 216 | // Every other caller: parallelize over q, each q privately slicing | ||
| 217 | // a freshly allocated Lq*Lk buffer so concurrent q's never share a | ||
| 218 | // `scores` row. | ||
| 219 | 110 | std::vector<float> scores_all(static_cast<std::size_t>(Lq) * | |
| 220 | 55 | static_cast<std::size_t>(Lk)); | |
| 221 |
2/4✓ Branch 0 taken 55 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 55 times.
✗ Branch 3 not taken.
|
1285 | parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t qi) { |
| 222 | 1230 | const int q = static_cast<int>(qi); | |
| 223 | 1230 | process_q(q, scores_all.data() + static_cast<std::size_t>(q) * Lk); | |
| 224 | 1230 | }); | |
| 225 | 55 | } | |
| 226 | 88 | } | |
| 227 | |||
| 228 | // Per-head flash-attention backward over pre-projected Q/K/V. Recompute P, | ||
| 229 | // then dV = P^T·dO_attn, dP = dO_attn·V^T, dS = P*(dP - D_q)*scale, | ||
| 230 | // dQ = dS·K, dK = dS^T·Q. dQ/dK/dV are OVERWRITTEN. | ||
| 231 | // | ||
| 232 | // scratch_P/scratch_dP/scratch_dS, if all non-null, must each point at a | ||
| 233 | // caller-owned buffer of at least Lq*Lk floats and are reused verbatim (no | ||
| 234 | // allocation here) — callers looping over many short sequences (e.g. varlen | ||
| 235 | // attention backward) hoist one triple of buffers sized to the batch's max | ||
| 236 | // Lq*Lk instead of paying three heap allocations per call. Leaving any of | ||
| 237 | // them null falls back to a local per-call allocation, matching the original | ||
| 238 | // single-call behavior. | ||
| 239 | // | ||
| 240 | // THREADING: when the scratch triple is null (every caller except the | ||
| 241 | // varlen family), the outer h loop is safe to parallelize — each head h | ||
| 242 | // writes only its own column range [h*hd, h*hd+hd) of dQ/dK/dV (never | ||
| 243 | // touched by another h), so the only shared mutable state is P/dP/dS, which | ||
| 244 | // the parallel path gives each h its own private Lq*Lk slice of a | ||
| 245 | // freshly-allocated H*Lq*Lk buffer instead of one triple reused across h. | ||
| 246 | // When the scratch triple is supplied, the caller | ||
| 247 | // (flash_attention_varlen_backward) has explicitly hoisted ONE Lq*Lk-sized | ||
| 248 | // triple across many sequential calls to this function (one per sequence in | ||
| 249 | // a batch); parallelizing the h loop in that case would have every thread's | ||
| 250 | // h race on that single shared triple, so this path is left exactly as it | ||
| 251 | // was (serial, reusing the caller's buffers verbatim). | ||
| 252 | 50 | void attention_core_backward(const float* Q, const float* K, const float* V, | |
| 253 | const float* dO, const float* mask, | ||
| 254 | int Lq, int Lk, int D, int H, bool causal, | ||
| 255 | float* dQ, float* dK, float* dV, | ||
| 256 | float* scratch_P = nullptr, | ||
| 257 | float* scratch_dP = nullptr, | ||
| 258 | float* scratch_dS = nullptr) { | ||
| 259 | 50 | const int hd = D / H; | |
| 260 | 50 | const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(hd)); | |
| 261 |
2/2✓ Branch 0 taken 10440 times.
✓ Branch 1 taken 50 times.
|
10490 | for (std::size_t i = 0; i < static_cast<std::size_t>(Lq) * D; ++i) |
| 262 | 10440 | dQ[i] = 0.0f; | |
| 263 |
2/2✓ Branch 0 taken 11432 times.
✓ Branch 1 taken 50 times.
|
11482 | for (std::size_t i = 0; i < static_cast<std::size_t>(Lk) * D; ++i) { |
| 264 | 11432 | dK[i] = 0.0f; | |
| 265 | 11432 | dV[i] = 0.0f; | |
| 266 | 11432 | } | |
| 267 | |||
| 268 | // Per-head body. P/dP/dS must be a private Lq*Lk buffer when running | ||
| 269 | // concurrently with other heads (see threading note above) — each h | ||
| 270 | // writes only its own column range of dQ/dK/dV, so distinct h's never | ||
| 271 | // touch the same output bytes. | ||
| 272 | 196 | auto process_head = [&](int h, float* P, float* dP, float* dS) { | |
| 273 | 146 | const int off = h * hd; | |
| 274 | // Recompute P for this head. | ||
| 275 |
2/2✓ Branch 0 taken 1111 times.
✓ Branch 1 taken 146 times.
|
1257 | for (int q = 0; q < Lq; ++q) { |
| 276 | 1111 | float m = -1e30f; | |
| 277 | 1111 | float* prow = P + static_cast<std::size_t>(q) * Lk; | |
| 278 |
2/2✓ Branch 0 taken 10911 times.
✓ Branch 1 taken 1111 times.
|
12022 | for (int k = 0; k < Lk; ++k) { |
| 279 |
4/4✓ Branch 0 taken 1165 times.
✓ Branch 1 taken 9746 times.
✓ Branch 2 taken 823 times.
✓ Branch 3 taken 342 times.
|
10911 | if (mask && mask[k] <= 0.5f) { prow[k] = -1e30f; continue; } |
| 280 |
4/4✓ Branch 0 taken 1717 times.
✓ Branch 1 taken 8852 times.
✓ Branch 2 taken 974 times.
✓ Branch 3 taken 743 times.
|
10569 | if (causal && k > q) { prow[k] = -1e30f; continue; } |
| 281 | 9826 | const float* qr = Q + static_cast<std::size_t>(q) * D + off; | |
| 282 | 9826 | const float* kr = K + static_cast<std::size_t>(k) * D + off; | |
| 283 | 9826 | float dot = 0.0f; | |
| 284 |
2/2✓ Branch 0 taken 107282 times.
✓ Branch 1 taken 9826 times.
|
117108 | for (int d = 0; d < hd; ++d) dot += qr[d] * kr[d]; |
| 285 | 9826 | prow[k] = dot * inv_sqrt; | |
| 286 |
2/2✓ Branch 0 taken 6868 times.
✓ Branch 1 taken 2958 times.
|
9826 | if (prow[k] > m) m = prow[k]; |
| 287 | 9826 | } | |
| 288 | 1111 | const bool empty = (m <= -1e29f); | |
| 289 | 1111 | float sum = 0.0f; | |
| 290 |
2/2✓ Branch 0 taken 10947 times.
✓ Branch 1 taken 1111 times.
|
12058 | for (int k = 0; k < Lk; ++k) { |
| 291 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10947 times.
|
10947 | const float e = empty ? 0.0f : std::exp(prow[k] - m); |
| 292 | 10947 | prow[k] = e; | |
| 293 | 10947 | sum += e; | |
| 294 | 10947 | } | |
| 295 |
1/2✓ Branch 0 taken 1111 times.
✗ Branch 1 not taken.
|
1111 | const float inv = sum > 0.0f ? 1.0f / sum : 0.0f; |
| 296 |
2/2✓ Branch 0 taken 11076 times.
✓ Branch 1 taken 1111 times.
|
12187 | for (int k = 0; k < Lk; ++k) prow[k] *= inv; |
| 297 | 1111 | } | |
| 298 | // dV[k, off+d] = sum_q P[q,k] * dO[q, off+d]. | ||
| 299 |
2/2✓ Branch 0 taken 1240 times.
✓ Branch 1 taken 146 times.
|
1386 | for (int k = 0; k < Lk; ++k) { |
| 300 | 1240 | float* dvr = dV + static_cast<std::size_t>(k) * D + off; | |
| 301 |
2/2✓ Branch 0 taken 11093 times.
✓ Branch 1 taken 1240 times.
|
12333 | for (int d = 0; d < hd; ++d) { |
| 302 | 11093 | float acc = 0.0f; | |
| 303 |
2/2✓ Branch 0 taken 116490 times.
✓ Branch 1 taken 11093 times.
|
127583 | for (int q = 0; q < Lq; ++q) |
| 304 | 232980 | acc += P[static_cast<std::size_t>(q) * Lk + k] * | |
| 305 | 116490 | dO[static_cast<std::size_t>(q) * D + off + d]; | |
| 306 | 11093 | dvr[d] = acc; | |
| 307 | 11093 | } | |
| 308 | 1240 | } | |
| 309 | // dP[q,k] = sum_d dO[q, off+d] * V[k, off+d]. | ||
| 310 |
2/2✓ Branch 0 taken 1113 times.
✓ Branch 1 taken 146 times.
|
1259 | for (int q = 0; q < Lq; ++q) { |
| 311 |
2/2✓ Branch 0 taken 10940 times.
✓ Branch 1 taken 1113 times.
|
12053 | for (int k = 0; k < Lk; ++k) { |
| 312 | 10940 | float acc = 0.0f; | |
| 313 |
2/2✓ Branch 0 taken 117507 times.
✓ Branch 1 taken 10940 times.
|
128447 | for (int d = 0; d < hd; ++d) |
| 314 | 235014 | acc += dO[static_cast<std::size_t>(q) * D + off + d] * | |
| 315 | 117507 | V[static_cast<std::size_t>(k) * D + off + d]; | |
| 316 | 10940 | dP[static_cast<std::size_t>(q) * Lk + k] = acc; | |
| 317 | 10940 | } | |
| 318 | 1113 | } | |
| 319 | // dS[q,k] = P[q,k] * (dP[q,k] - D_q) * inv_sqrt. | ||
| 320 |
2/2✓ Branch 0 taken 1117 times.
✓ Branch 1 taken 146 times.
|
1263 | for (int q = 0; q < Lq; ++q) { |
| 321 | 1117 | const float* prow = P + static_cast<std::size_t>(q) * Lk; | |
| 322 | 1117 | const float* dpr = dP + static_cast<std::size_t>(q) * Lk; | |
| 323 | 1117 | float* dsr = dS + static_cast<std::size_t>(q) * Lk; | |
| 324 | 1117 | float Dq = 0.0f; | |
| 325 |
2/2✓ Branch 0 taken 11061 times.
✓ Branch 1 taken 1117 times.
|
12178 | for (int k = 0; k < Lk; ++k) Dq += prow[k] * dpr[k]; |
| 326 |
2/2✓ Branch 0 taken 11074 times.
✓ Branch 1 taken 1117 times.
|
12191 | for (int k = 0; k < Lk; ++k) |
| 327 | 11074 | dsr[k] = prow[k] * (dpr[k] - Dq) * inv_sqrt; | |
| 328 | 1117 | } | |
| 329 | // dQ[q, off+d] = sum_k dS[q,k] * K[k, off+d]. | ||
| 330 |
2/2✓ Branch 0 taken 1112 times.
✓ Branch 1 taken 146 times.
|
1258 | for (int q = 0; q < Lq; ++q) { |
| 331 | 1112 | float* dqr = dQ + static_cast<std::size_t>(q) * D + off; | |
| 332 | 1112 | const float* dsr = dS + static_cast<std::size_t>(q) * Lk; | |
| 333 |
2/2✓ Branch 0 taken 10233 times.
✓ Branch 1 taken 1112 times.
|
11345 | for (int d = 0; d < hd; ++d) { |
| 334 | 10233 | float acc = 0.0f; | |
| 335 |
2/2✓ Branch 0 taken 117595 times.
✓ Branch 1 taken 10233 times.
|
127828 | for (int k = 0; k < Lk; ++k) |
| 336 | 235190 | acc += dsr[k] * | |
| 337 | 117595 | K[static_cast<std::size_t>(k) * D + off + d]; | |
| 338 | 10233 | dqr[d] = acc; | |
| 339 | 10233 | } | |
| 340 | 1112 | } | |
| 341 | // dK[k, off+d] = sum_q dS[q,k] * Q[q, off+d]. | ||
| 342 |
2/2✓ Branch 0 taken 1246 times.
✓ Branch 1 taken 146 times.
|
1392 | for (int k = 0; k < Lk; ++k) { |
| 343 | 1246 | float* dkr = dK + static_cast<std::size_t>(k) * D + off; | |
| 344 |
2/2✓ Branch 0 taken 11297 times.
✓ Branch 1 taken 1246 times.
|
12543 | for (int d = 0; d < hd; ++d) { |
| 345 | 11297 | float acc = 0.0f; | |
| 346 |
2/2✓ Branch 0 taken 117729 times.
✓ Branch 1 taken 11297 times.
|
129026 | for (int q = 0; q < Lq; ++q) |
| 347 | 235458 | acc += dS[static_cast<std::size_t>(q) * Lk + k] * | |
| 348 | 117729 | Q[static_cast<std::size_t>(q) * D + off + d]; | |
| 349 | 11297 | dkr[d] = acc; | |
| 350 | 11297 | } | |
| 351 | 1246 | } | |
| 352 | 146 | }; | |
| 353 | |||
| 354 |
4/6✓ Branch 0 taken 33 times.
✓ Branch 1 taken 17 times.
✓ Branch 2 taken 33 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 33 times.
|
50 | if (scratch_P && scratch_dP && scratch_dS) { |
| 355 | // Varlen path: caller owns a single Lq*Lk-sized triple reused across | ||
| 356 | // many sequential calls to this function. Stay serial (see the | ||
| 357 | // threading note above). | ||
| 358 |
2/2✓ Branch 0 taken 76 times.
✓ Branch 1 taken 33 times.
|
109 | for (int h = 0; h < H; ++h) |
| 359 | 76 | process_head(h, scratch_P, scratch_dP, scratch_dS); | |
| 360 | 33 | } else { | |
| 361 | // Every other caller: parallelize over h, each h privately slicing | ||
| 362 | // a freshly allocated H*Lq*Lk buffer so concurrent heads never | ||
| 363 | // share a P/dP/dS segment. | ||
| 364 | 17 | const std::size_t n = static_cast<std::size_t>(Lq) * Lk; | |
| 365 | 17 | std::vector<float> P_all(static_cast<std::size_t>(H) * n); | |
| 366 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | std::vector<float> dP_all(static_cast<std::size_t>(H) * n); |
| 367 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | std::vector<float> dS_all(static_cast<std::size_t>(H) * n); |
| 368 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
87 | parallel_for(static_cast<std::size_t>(H), [&](std::size_t hi) { |
| 369 | 70 | const int h = static_cast<int>(hi); | |
| 370 | 140 | process_head(h, P_all.data() + static_cast<std::size_t>(h) * n, | |
| 371 | 70 | dP_all.data() + static_cast<std::size_t>(h) * n, | |
| 372 | 70 | dS_all.data() + static_cast<std::size_t>(h) * n); | |
| 373 | 70 | }); | |
| 374 | 17 | } | |
| 375 | 50 | } | |
| 376 | |||
| 377 | 158 | inline void ensure_f32(Tensor& t, int r, int c) { | |
| 378 |
4/6✓ Branch 0 taken 12 times.
✓ Branch 1 taken 146 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
|
158 | if (t.rows != r || t.cols != c || t.dtype != Dtype::FP32) |
| 379 | 146 | t.resize(r, c, Dtype::FP32); | |
| 380 | 158 | } | |
| 381 | |||
| 382 | } // namespace | ||
| 383 | |||
| 384 | // ─── flash_attention_forward ─────────────────────────────────────────────── | ||
| 385 | |||
| 386 | 16 | void flash_attention_forward(const ::brotensor::Tensor& Q, | |
| 387 | const ::brotensor::Tensor& K, | ||
| 388 | const ::brotensor::Tensor& V, | ||
| 389 | const float* d_mask, | ||
| 390 | int num_heads, | ||
| 391 | bool causal, | ||
| 392 | ::brotensor::Tensor& O) { | ||
| 393 | 16 | const int Lq = Q.rows; | |
| 394 | 16 | const int Lk = K.rows; | |
| 395 | 16 | const int D = Q.cols; | |
| 396 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | if (K.cols != D || V.cols != D || V.rows != Lk) |
| 397 | ✗ | throw std::runtime_error("flash_attention_forward: shape mismatch"); | |
| 398 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | if (num_heads <= 0 || D % num_heads != 0) |
| 399 | ✗ | throw std::runtime_error("flash_attention_forward: num_heads must divide D"); | |
| 400 |
3/4✓ Branch 0 taken 3 times.
✓ Branch 1 taken 13 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
16 | if (causal && Lq != Lk) |
| 401 | ✗ | throw std::runtime_error("flash_attention_forward: causal requires Lq == Lk"); | |
| 402 | 16 | ensure_f32(O, Lq, D); | |
| 403 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 16 times.
|
16 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 404 | |||
| 405 | 32 | attention_core(Q.host_f32(), K.host_f32(), V.host_f32(), d_mask, | |
| 406 | 16 | Lq, Lk, D, num_heads, causal, O.host_f32_mut(), nullptr); | |
| 407 | 16 | } | |
| 408 | |||
| 409 | // ─── flash_attention_windowed_forward ────────────────────────────────────── | ||
| 410 | // | ||
| 411 | // Sliding-window causal self-attention (FP32, inference-only). Always causal. | ||
| 412 | // The Lq queries occupy the last Lq positions of a length-Lk causal sequence | ||
| 413 | // (q_offset = Lk - Lq): query row r sits at absolute position r + q_offset and | ||
| 414 | // attends keys [max(0, pos-window+1), pos]. window <= 0 is unbounded causal. | ||
| 415 | // - Lq == Lk (q_offset 0): plain causal / sliding window (e.g. the codec). | ||
| 416 | // - Lq < Lk: incremental decode — a short query block (often a single token) | ||
| 417 | // attending the whole K/V cache, the autoregressive step. Lq == 1 with | ||
| 418 | // window <= 0 attends every cached key (full attention over the cache). | ||
| 419 | // Lk >= Lq required. | ||
| 420 | ✗ | void flash_attention_windowed_forward(const ::brotensor::Tensor& Q, | |
| 421 | const ::brotensor::Tensor& K, | ||
| 422 | const ::brotensor::Tensor& V, | ||
| 423 | const float* d_mask, | ||
| 424 | int num_heads, | ||
| 425 | int window, | ||
| 426 | ::brotensor::Tensor& O) { | ||
| 427 | ✗ | const int Lq = Q.rows; | |
| 428 | ✗ | const int Lk = K.rows; | |
| 429 | ✗ | const int D = Q.cols; // Dq = num_heads * head_dim | |
| 430 | ✗ | const int Dkv = K.cols; // n_kv * head_dim (GQA when < D) | |
| 431 | ✗ | if (V.cols != Dkv || V.rows != Lk) | |
| 432 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: shape mismatch"); | |
| 433 | ✗ | if (num_heads <= 0 || D % num_heads != 0) | |
| 434 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: num_heads must divide D"); | |
| 435 | ✗ | const int head_dim = D / num_heads; | |
| 436 | ✗ | if (Dkv == 0 || Dkv % head_dim != 0) | |
| 437 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: K/V width must be a head_dim multiple"); | |
| 438 | ✗ | const int n_kv = Dkv / head_dim; | |
| 439 | ✗ | if (num_heads % n_kv != 0) | |
| 440 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: num_heads must be a multiple of n_kv"); | |
| 441 | ✗ | if (Lk < Lq) | |
| 442 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: requires Lk >= Lq"); | |
| 443 | ✗ | ensure_f32(O, Lq, D); | |
| 444 | ✗ | if (Lq == 0 || Lk == 0 || D == 0) return; | |
| 445 | |||
| 446 | ✗ | attention_core(Q.host_f32(), K.host_f32(), V.host_f32(), d_mask, | |
| 447 | ✗ | Lq, Lk, D, num_heads, /*causal=*/true, O.host_f32_mut(), | |
| 448 | ✗ | nullptr, window, /*q_offset=*/Lk - Lq, Dkv, | |
| 449 | ✗ | /*group=*/num_heads / n_kv); | |
| 450 | ✗ | } | |
| 451 | |||
| 452 | // ─── flash_attention_gqa_forward ─────────────────────────────────────────── | ||
| 453 | // | ||
| 454 | // GQA generalisation of flash_attention_forward, causal OR bidirectional. Q is | ||
| 455 | // num_q_heads-wide, K/V are num_kv_heads-wide (GQA group = num_q_heads/num_kv_heads; | ||
| 456 | // equal == MHA). causal == false attends every key (the bidirectional encoder | ||
| 457 | // prefill). attention_core already carries the group mapping, the causal flag, | ||
| 458 | // and the key mask, so this is a validate-and-dispatch shim — q_offset is 0 | ||
| 459 | // because the queries occupy positions [0, Lq) (Lq == Lk when causal). | ||
| 460 | 9 | void flash_attention_gqa_forward(const ::brotensor::Tensor& Q, | |
| 461 | const ::brotensor::Tensor& K, | ||
| 462 | const ::brotensor::Tensor& V, | ||
| 463 | const float* d_mask, | ||
| 464 | int num_q_heads, | ||
| 465 | int num_kv_heads, | ||
| 466 | bool causal, | ||
| 467 | ::brotensor::Tensor& O) { | ||
| 468 | 9 | const int Lq = Q.rows; | |
| 469 | 9 | const int Lk = K.rows; | |
| 470 | 9 | const int Dq = Q.cols; // num_q_heads * head_dim | |
| 471 | 9 | const int Dkv = K.cols; // num_kv_heads * head_dim | |
| 472 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (num_q_heads <= 0 || num_kv_heads <= 0) |
| 473 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: head counts must be positive"); | |
| 474 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (Dq % num_q_heads != 0) |
| 475 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: num_q_heads must divide Q.cols"); | |
| 476 | 9 | const int head_dim = Dq / num_q_heads; | |
| 477 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (V.cols != Dkv || V.rows != Lk) |
| 478 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: shape mismatch"); | |
| 479 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (Dkv != num_kv_heads * head_dim) |
| 480 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: K/V width must be num_kv_heads*head_dim"); | |
| 481 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (num_q_heads % num_kv_heads != 0) |
| 482 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: num_kv_heads must divide num_q_heads"); | |
| 483 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (Lk < Lq) |
| 484 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: requires Lk >= Lq"); | |
| 485 |
3/4✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
9 | if (causal && Lq != Lk) |
| 486 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: causal requires Lq == Lk"); | |
| 487 | 9 | ensure_f32(O, Lq, Dq); | |
| 488 |
3/6✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 9 times.
|
9 | if (Lq == 0 || Lk == 0 || Dq == 0) return; |
| 489 | |||
| 490 | 18 | attention_core(Q.host_f32(), K.host_f32(), V.host_f32(), d_mask, | |
| 491 | 9 | Lq, Lk, Dq, num_q_heads, causal, O.host_f32_mut(), | |
| 492 | 9 | /*P=*/nullptr, /*window=*/0, /*q_offset=*/0, Dkv, | |
| 493 | 9 | /*group=*/num_q_heads / num_kv_heads); | |
| 494 | 9 | } | |
| 495 | |||
| 496 | // ─── flash_attention_varlen_forward ──────────────────────────────────────── | ||
| 497 | // | ||
| 498 | // Packed variable-length attention (Qwen3-VL window attention). Q/K/V are one | ||
| 499 | // big (total_tokens, num_heads*head_dim) tensor each; cu_seqlens_q/k are | ||
| 500 | // length B+1 INT32 prefix sums delimiting per-sequence row ranges. Sequence b | ||
| 501 | // runs attention_core over its own Q[Lq_b, D] and K/V[Lk_b, D] slice — no | ||
| 502 | // cross-sequence attention. | ||
| 503 | // | ||
| 504 | // On CPU cu_seqlens_q/k are raw host pointers (matches the existing d_mask | ||
| 505 | // convention: device pointer on GPU, host pointer on CPU). max_seqlen_q/k are | ||
| 506 | // only used by the GPU kernel for block sizing and are ignored here. | ||
| 507 | 17 | void flash_attention_varlen_forward(const ::brotensor::Tensor& Q, | |
| 508 | const ::brotensor::Tensor& K, | ||
| 509 | const ::brotensor::Tensor& V, | ||
| 510 | const int32_t* cu_seqlens_q, | ||
| 511 | const int32_t* cu_seqlens_k, | ||
| 512 | int batch_size, | ||
| 513 | int /*max_seqlen_q*/, | ||
| 514 | int /*max_seqlen_k*/, | ||
| 515 | int num_heads, | ||
| 516 | int head_dim, | ||
| 517 | bool causal, | ||
| 518 | ::brotensor::Tensor& O) { | ||
| 519 | 17 | const int total_q = Q.rows; | |
| 520 | 17 | const int total_k = K.rows; | |
| 521 | 17 | const int D = num_heads * head_dim; | |
| 522 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (Q.cols != D || K.cols != D || V.cols != D || V.rows != total_k) |
| 523 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: shape mismatch"); | |
| 524 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (num_heads <= 0 || head_dim <= 0) |
| 525 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: num_heads/head_dim must be positive"); | |
| 526 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (batch_size < 0) |
| 527 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: batch_size must be non-negative"); | |
| 528 |
2/4✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
17 | if (batch_size > 0 && (!cu_seqlens_q || !cu_seqlens_k)) |
| 529 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_q/k required when batch_size > 0"); | |
| 530 | 17 | ensure_f32(O, total_q, D); | |
| 531 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 17 times.
|
17 | if (total_q == 0 || D == 0) return; |
| 532 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
17 | if (cu_seqlens_q && cu_seqlens_q[0] != 0) |
| 533 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_q[0] must be 0"); | |
| 534 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
17 | if (cu_seqlens_k && cu_seqlens_k[0] != 0) |
| 535 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_k[0] must be 0"); | |
| 536 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
|
17 | if (batch_size > 0) { |
| 537 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (cu_seqlens_q[batch_size] != total_q) |
| 538 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_q[B] != total_tokens_q"); | |
| 539 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (cu_seqlens_k[batch_size] != total_k) |
| 540 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_k[B] != total_tokens_k"); | |
| 541 | 17 | } | |
| 542 | |||
| 543 | 17 | const float* Qp = Q.host_f32(); | |
| 544 | 17 | const float* Kp = K.host_f32(); | |
| 545 | 17 | const float* Vp = V.host_f32(); | |
| 546 | 17 | float* Op = O.host_f32_mut(); | |
| 547 | |||
| 548 | // Hoist attention_core's per-call `scores` scratch (size Lk) to one | ||
| 549 | // buffer sized for the widest sequence in the batch, reused for every | ||
| 550 | // sequence below — batches of many short sequences would otherwise pay | ||
| 551 | // one small heap allocation per sequence. Negative per-sequence lengths | ||
| 552 | // (malformed cu_seqlens) are skipped here; the loop below throws on them | ||
| 553 | // before attention_core is ever called for that b. | ||
| 554 | 17 | int max_lk = 0; | |
| 555 |
2/2✓ Branch 0 taken 33 times.
✓ Branch 1 taken 17 times.
|
50 | for (int b = 0; b < batch_size; ++b) { |
| 556 | 33 | const int lk = cu_seqlens_k[b + 1] - cu_seqlens_k[b]; | |
| 557 |
2/2✓ Branch 0 taken 11 times.
✓ Branch 1 taken 22 times.
|
33 | if (lk > max_lk) max_lk = lk; |
| 558 | 33 | } | |
| 559 | 17 | std::vector<float> scratch_scores(static_cast<std::size_t>(max_lk)); | |
| 560 | |||
| 561 |
2/2✓ Branch 0 taken 17 times.
✓ Branch 1 taken 33 times.
|
50 | for (int b = 0; b < batch_size; ++b) { |
| 562 | 33 | const int q_beg = cu_seqlens_q[b]; | |
| 563 | 33 | const int q_end = cu_seqlens_q[b + 1]; | |
| 564 | 33 | const int k_beg = cu_seqlens_k[b]; | |
| 565 | 33 | const int k_end = cu_seqlens_k[b + 1]; | |
| 566 | 33 | const int Lq = q_end - q_beg; | |
| 567 | 33 | const int Lk = k_end - k_beg; | |
| 568 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (Lq < 0 || Lk < 0) |
| 569 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens must be non-decreasing"); | |
| 570 |
3/4✓ Branch 0 taken 15 times.
✓ Branch 1 taken 18 times.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
|
33 | if (causal && Lq != Lk) |
| 571 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: causal requires per-sequence Lq == Lk"); | |
| 572 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (Lq == 0) continue; |
| 573 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 33 times.
|
33 | if (Lk == 0) { |
| 574 | // Fully empty K range — output rows are zero (no keys to attend to). | ||
| 575 | ✗ | for (std::size_t i = 0; i < static_cast<std::size_t>(Lq) * D; ++i) | |
| 576 | ✗ | Op[static_cast<std::size_t>(q_beg) * D + i] = 0.0f; | |
| 577 | ✗ | continue; | |
| 578 | } | ||
| 579 |
2/4✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 33 times.
✗ Branch 3 not taken.
|
66 | attention_core(Qp + static_cast<std::size_t>(q_beg) * D, |
| 580 | 33 | Kp + static_cast<std::size_t>(k_beg) * D, | |
| 581 | 33 | Vp + static_cast<std::size_t>(k_beg) * D, | |
| 582 | /*mask=*/nullptr, | ||
| 583 | 33 | Lq, Lk, D, num_heads, causal, | |
| 584 | 33 | Op + static_cast<std::size_t>(q_beg) * D, | |
| 585 | /*P=*/nullptr, /*window=*/0, /*q_offset=*/0, /*Dkv=*/-1, | ||
| 586 | 33 | /*group=*/1, scratch_scores.data()); | |
| 587 | 33 | } | |
| 588 | 17 | } | |
| 589 | |||
| 590 | // ─── flash_attention_project_kv ──────────────────────────────────────────── | ||
| 591 | |||
| 592 | 4 | void flash_attention_project_kv(const ::brotensor::Tensor& ctx, | |
| 593 | const ::brotensor::Tensor& Wk, | ||
| 594 | const ::brotensor::Tensor* bk, | ||
| 595 | const ::brotensor::Tensor& Wv, | ||
| 596 | const ::brotensor::Tensor* bv, | ||
| 597 | ::brotensor::Tensor& K_out, | ||
| 598 | ::brotensor::Tensor& V_out) { | ||
| 599 | 4 | const int Lk = ctx.rows; | |
| 600 | 4 | const int D_ctx = ctx.cols; | |
| 601 | 4 | const int D = Wk.rows; | |
| 602 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | if (Wk.cols != D_ctx || Wv.rows != D || Wv.cols != D_ctx) |
| 603 | ✗ | throw std::runtime_error("flash_attention_project_kv: Wk/Wv shape mismatch"); | |
| 604 | 4 | ensure_f32(K_out, Lk, D); | |
| 605 | 4 | ensure_f32(V_out, Lk, D); | |
| 606 |
2/4✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
|
4 | if (Lk == 0 || D == 0) return; |
| 607 | |||
| 608 | 8 | linear_proj(ctx.host_f32(), Wk.host_f32(), | |
| 609 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
|
4 | bk ? bk->host_f32() : nullptr, |
| 610 | 4 | K_out.host_f32_mut(), Lk, D_ctx, D); | |
| 611 | 8 | linear_proj(ctx.host_f32(), Wv.host_f32(), | |
| 612 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
|
4 | bv ? bv->host_f32() : nullptr, |
| 613 | 4 | V_out.host_f32_mut(), Lk, D_ctx, D); | |
| 614 | 4 | } | |
| 615 | |||
| 616 | // ─── flash_attention_q_with_kv_cached_forward ────────────────────────────── | ||
| 617 | |||
| 618 | 6 | void flash_attention_q_with_kv_cached_forward(const ::brotensor::Tensor& X, | |
| 619 | const ::brotensor::Tensor& K, | ||
| 620 | const ::brotensor::Tensor& V, | ||
| 621 | const ::brotensor::Tensor& Wq, | ||
| 622 | const ::brotensor::Tensor* bq, | ||
| 623 | const ::brotensor::Tensor& Wo, | ||
| 624 | const ::brotensor::Tensor* bo, | ||
| 625 | const float* d_mask, | ||
| 626 | int num_heads, | ||
| 627 | bool causal, | ||
| 628 | ::brotensor::Tensor& O) { | ||
| 629 | 6 | const int Lq = X.rows; | |
| 630 | 6 | const int D = X.cols; | |
| 631 | 6 | const int Lk = K.rows; | |
| 632 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (K.cols != D || V.rows != Lk || V.cols != D) |
| 633 | ✗ | throw std::runtime_error("flash_attention_q_with_kv_cached_forward: K/V shape mismatch"); | |
| 634 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (Wq.rows != D || Wq.cols != D || Wo.rows != D || Wo.cols != D) |
| 635 | ✗ | throw std::runtime_error("flash_attention_q_with_kv_cached_forward: Wq/Wo shape mismatch"); | |
| 636 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (num_heads <= 0 || D % num_heads != 0) |
| 637 | ✗ | throw std::runtime_error("flash_attention_q_with_kv_cached_forward: num_heads must divide D"); | |
| 638 | 6 | ensure_f32(O, Lq, D); | |
| 639 |
3/6✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
|
6 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 640 | |||
| 641 | 6 | std::vector<float> Qp(static_cast<std::size_t>(Lq) * D); | |
| 642 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | std::vector<float> Op(static_cast<std::size_t>(Lq) * D); |
| 643 |
7/12✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 4 times.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 6 times.
✗ Branch 11 not taken.
|
6 | linear_proj(X.host_f32(), Wq.host_f32(), bq ? bq->host_f32() : nullptr, |
| 644 | 6 | Qp.data(), Lq, D, D); | |
| 645 |
4/8✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
|
6 | attention_core(Qp.data(), K.host_f32(), V.host_f32(), d_mask, |
| 646 | 6 | Lq, Lk, D, num_heads, causal, Op.data(), nullptr); | |
| 647 |
5/8✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
|
6 | linear_proj(Op.data(), Wo.host_f32(), bo ? bo->host_f32() : nullptr, |
| 648 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | O.host_f32_mut(), Lq, D, D); |
| 649 | 6 | } | |
| 650 | |||
| 651 | // ─── flash_attention_qkvo_forward ────────────────────────────────────────── | ||
| 652 | |||
| 653 | 15 | void flash_attention_qkvo_forward(const ::brotensor::Tensor& X, | |
| 654 | const ::brotensor::Tensor* Ctx, | ||
| 655 | const ::brotensor::Tensor& Wq, | ||
| 656 | const ::brotensor::Tensor* bq, | ||
| 657 | const ::brotensor::Tensor& Wk, | ||
| 658 | const ::brotensor::Tensor* bk, | ||
| 659 | const ::brotensor::Tensor& Wv, | ||
| 660 | const ::brotensor::Tensor* bv, | ||
| 661 | const ::brotensor::Tensor& Wo, | ||
| 662 | const ::brotensor::Tensor* bo, | ||
| 663 | const float* d_mask, | ||
| 664 | int num_heads, | ||
| 665 | bool causal, | ||
| 666 | ::brotensor::Tensor& O) { | ||
| 667 | 15 | const int Lq = X.rows; | |
| 668 | 15 | const int D = X.cols; | |
| 669 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 12 times.
|
15 | const Tensor& kv_src = Ctx ? *Ctx : X; |
| 670 | 15 | const int Lk = kv_src.rows; | |
| 671 | 15 | const int D_ctx = kv_src.cols; | |
| 672 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
30 | if (Wq.rows != D || Wq.cols != D || |
| 673 | 15 | Wk.rows != D || Wk.cols != D_ctx || | |
| 674 | 15 | Wv.rows != D || Wv.cols != D_ctx || | |
| 675 | 15 | Wo.rows != D || Wo.cols != D) | |
| 676 | ✗ | throw std::runtime_error("flash_attention_qkvo_forward: shape mismatch"); | |
| 677 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | if (num_heads <= 0 || D % num_heads != 0) |
| 678 | ✗ | throw std::runtime_error("flash_attention_qkvo_forward: num_heads must divide D"); | |
| 679 | 15 | ensure_f32(O, Lq, D); | |
| 680 |
3/6✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 15 times.
|
15 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 681 | |||
| 682 | 15 | std::vector<float> Qp(static_cast<std::size_t>(Lq) * D); | |
| 683 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | std::vector<float> Kp(static_cast<std::size_t>(Lk) * D); |
| 684 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | std::vector<float> Vp(static_cast<std::size_t>(Lk) * D); |
| 685 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | std::vector<float> Op(static_cast<std::size_t>(Lq) * D); |
| 686 | |||
| 687 |
7/12✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✓ Branch 5 taken 9 times.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 15 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 15 times.
✗ Branch 11 not taken.
|
15 | linear_proj(X.host_f32(), Wq.host_f32(), bq ? bq->host_f32() : nullptr, |
| 688 | 15 | Qp.data(), Lq, D, D); | |
| 689 |
3/6✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 15 times.
✗ Branch 5 not taken.
|
30 | linear_proj(kv_src.host_f32(), Wk.host_f32(), |
| 690 |
3/4✓ Branch 0 taken 6 times.
✓ Branch 1 taken 9 times.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
15 | bk ? bk->host_f32() : nullptr, Kp.data(), Lk, D_ctx, D); |
| 691 |
3/6✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 15 times.
✗ Branch 5 not taken.
|
30 | linear_proj(kv_src.host_f32(), Wv.host_f32(), |
| 692 |
3/4✓ Branch 0 taken 6 times.
✓ Branch 1 taken 9 times.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
15 | bv ? bv->host_f32() : nullptr, Vp.data(), Lk, D_ctx, D); |
| 693 |
2/4✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
|
30 | attention_core(Qp.data(), Kp.data(), Vp.data(), d_mask, |
| 694 | 15 | Lq, Lk, D, num_heads, causal, Op.data(), nullptr); | |
| 695 |
5/8✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 9 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 15 times.
✗ Branch 7 not taken.
|
15 | linear_proj(Op.data(), Wo.host_f32(), bo ? bo->host_f32() : nullptr, |
| 696 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | O.host_f32_mut(), Lq, D, D); |
| 697 | 15 | } | |
| 698 | |||
| 699 | // ─── flash_attention_backward ────────────────────────────────────────────── | ||
| 700 | |||
| 701 | 8 | void flash_attention_backward(const ::brotensor::Tensor& Q, | |
| 702 | const ::brotensor::Tensor& K, | ||
| 703 | const ::brotensor::Tensor& V, | ||
| 704 | const ::brotensor::Tensor& O, | ||
| 705 | const ::brotensor::Tensor& dO, | ||
| 706 | const float* d_mask, | ||
| 707 | int num_heads, | ||
| 708 | bool causal, | ||
| 709 | ::brotensor::Tensor& dQ, | ||
| 710 | ::brotensor::Tensor& dK, | ||
| 711 | ::brotensor::Tensor& dV) { | ||
| 712 | 8 | (void)O; // recompute-based; O retained in API for symmetry. | |
| 713 | 8 | const int Lq = Q.rows; | |
| 714 | 8 | const int Lk = K.rows; | |
| 715 | 8 | const int D = Q.cols; | |
| 716 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (K.cols != D || V.cols != D || V.rows != Lk) |
| 717 | ✗ | throw std::runtime_error("flash_attention_backward: Q/K/V shape mismatch"); | |
| 718 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (dO.rows != Lq || dO.cols != D) |
| 719 | ✗ | throw std::runtime_error("flash_attention_backward: dO shape mismatch"); | |
| 720 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (num_heads <= 0 || D % num_heads != 0) |
| 721 | ✗ | throw std::runtime_error("flash_attention_backward: num_heads must divide D"); | |
| 722 |
3/4✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
|
8 | if (causal && Lq != Lk) |
| 723 | ✗ | throw std::runtime_error("flash_attention_backward: causal requires Lq == Lk"); | |
| 724 | 8 | ensure_f32(dQ, Lq, D); | |
| 725 | 8 | ensure_f32(dK, Lk, D); | |
| 726 | 8 | ensure_f32(dV, Lk, D); | |
| 727 |
3/6✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
|
8 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 728 | |||
| 729 | 16 | attention_core_backward(Q.host_f32(), K.host_f32(), V.host_f32(), | |
| 730 | 8 | dO.host_f32(), d_mask, Lq, Lk, D, num_heads, | |
| 731 | 8 | causal, dQ.host_f32_mut(), dK.host_f32_mut(), | |
| 732 | 8 | dV.host_f32_mut()); | |
| 733 | 8 | } | |
| 734 | |||
| 735 | // ─── flash_attention_varlen_backward ─────────────────────────────────────── | ||
| 736 | // | ||
| 737 | // Per-sequence backward over the packed (total_tokens, D) layout. Same | ||
| 738 | // recompute math as flash_attention_backward (attention_core_backward) run | ||
| 739 | // once per sequence on its [cu_seqlens_q[b], cu_seqlens_q[b+1]) Q slice and | ||
| 740 | // [cu_seqlens_k[b], cu_seqlens_k[b+1]) K/V slice. dQ/dK/dV are OVERWRITTEN — | ||
| 741 | // attention_core_backward zeros its per-call output region, and rows outside | ||
| 742 | // any sequence's range (which can't happen for a well-formed cu_seqlens) plus | ||
| 743 | // rows whose K range is empty are explicitly zeroed here so the contract | ||
| 744 | // holds globally. | ||
| 745 | 17 | void flash_attention_varlen_backward(const ::brotensor::Tensor& Q, | |
| 746 | const ::brotensor::Tensor& K, | ||
| 747 | const ::brotensor::Tensor& V, | ||
| 748 | const ::brotensor::Tensor& O, | ||
| 749 | const ::brotensor::Tensor& dO, | ||
| 750 | const int32_t* cu_seqlens_q, | ||
| 751 | const int32_t* cu_seqlens_k, | ||
| 752 | int batch_size, | ||
| 753 | int /*max_seqlen_q*/, | ||
| 754 | int /*max_seqlen_k*/, | ||
| 755 | int num_heads, | ||
| 756 | int head_dim, | ||
| 757 | bool causal, | ||
| 758 | ::brotensor::Tensor& dQ, | ||
| 759 | ::brotensor::Tensor& dK, | ||
| 760 | ::brotensor::Tensor& dV) { | ||
| 761 | 17 | (void)O; // recompute-based; O retained in API for symmetry. | |
| 762 | 17 | const int total_q = Q.rows; | |
| 763 | 17 | const int total_k = K.rows; | |
| 764 | 17 | const int D = num_heads * head_dim; | |
| 765 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (Q.cols != D || K.cols != D || V.cols != D || V.rows != total_k) |
| 766 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: Q/K/V shape mismatch"); | |
| 767 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (dO.rows != total_q || dO.cols != D) |
| 768 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: dO shape mismatch"); | |
| 769 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (num_heads <= 0 || head_dim <= 0) |
| 770 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: num_heads/head_dim must be positive"); | |
| 771 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (batch_size < 0) |
| 772 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: batch_size must be non-negative"); | |
| 773 |
2/4✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
17 | if (batch_size > 0 && (!cu_seqlens_q || !cu_seqlens_k)) |
| 774 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_q/k required when batch_size > 0"); | |
| 775 | 17 | ensure_f32(dQ, total_q, D); | |
| 776 | 17 | ensure_f32(dK, total_k, D); | |
| 777 | 17 | ensure_f32(dV, total_k, D); | |
| 778 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
17 | if (total_q == 0 && total_k == 0) return; |
| 779 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (D == 0) return; |
| 780 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
17 | if (cu_seqlens_q && cu_seqlens_q[0] != 0) |
| 781 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_q[0] must be 0"); | |
| 782 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
17 | if (cu_seqlens_k && cu_seqlens_k[0] != 0) |
| 783 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_k[0] must be 0"); | |
| 784 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
|
17 | if (batch_size > 0) { |
| 785 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (cu_seqlens_q[batch_size] != total_q) |
| 786 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_q[B] != total_tokens_q"); | |
| 787 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (cu_seqlens_k[batch_size] != total_k) |
| 788 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_k[B] != total_tokens_k"); | |
| 789 | 17 | } | |
| 790 | |||
| 791 | // Zero all grads upfront so any rows not covered by a sequence (or covered | ||
| 792 | // by an empty-K sequence) end up at exactly 0.0f without depending on the | ||
| 793 | // per-sequence path to touch them. | ||
| 794 | 17 | float* dQp = dQ.host_f32_mut(); | |
| 795 | 17 | float* dKp = dK.host_f32_mut(); | |
| 796 | 17 | float* dVp = dV.host_f32_mut(); | |
| 797 |
2/2✓ Branch 0 taken 5672 times.
✓ Branch 1 taken 17 times.
|
5689 | for (std::size_t i = 0; i < static_cast<std::size_t>(total_q) * D; ++i) dQp[i] = 0.0f; |
| 798 |
2/2✓ Branch 0 taken 5704 times.
✓ Branch 1 taken 17 times.
|
5721 | for (std::size_t i = 0; i < static_cast<std::size_t>(total_k) * D; ++i) { |
| 799 | 5704 | dKp[i] = 0.0f; | |
| 800 | 5704 | dVp[i] = 0.0f; | |
| 801 | 5704 | } | |
| 802 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (batch_size == 0) return; |
| 803 | |||
| 804 | 17 | const float* Qp = Q.host_f32(); | |
| 805 | 17 | const float* Kp = K.host_f32(); | |
| 806 | 17 | const float* Vp = V.host_f32(); | |
| 807 | 17 | const float* dOp = dO.host_f32(); | |
| 808 | |||
| 809 | // Hoist attention_core_backward's per-call P/dP/dS scratch (each size | ||
| 810 | // Lq*Lk) to one triple of buffers sized for the widest sequence in the | ||
| 811 | // batch, reused for every sequence below — batches of many short | ||
| 812 | // sequences would otherwise pay three small heap allocations per | ||
| 813 | // sequence. Malformed (negative) per-sequence lengths are skipped in this | ||
| 814 | // sizing pass; the loop below throws on them before attention_core_backward | ||
| 815 | // is ever called for that b, and casting a negative length to size_t here | ||
| 816 | // could otherwise wrap to a huge value and provoke a bogus allocation. | ||
| 817 | 17 | std::size_t max_qk = 0; | |
| 818 |
2/2✓ Branch 0 taken 33 times.
✓ Branch 1 taken 17 times.
|
50 | for (int b = 0; b < batch_size; ++b) { |
| 819 | 33 | const int lq = cu_seqlens_q[b + 1] - cu_seqlens_q[b]; | |
| 820 | 33 | const int lk = cu_seqlens_k[b + 1] - cu_seqlens_k[b]; | |
| 821 |
2/4✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 33 times.
|
33 | if (lq < 0 || lk < 0) continue; |
| 822 | 33 | const std::size_t qk = | |
| 823 | 33 | static_cast<std::size_t>(lq) * static_cast<std::size_t>(lk); | |
| 824 |
2/2✓ Branch 0 taken 11 times.
✓ Branch 1 taken 22 times.
|
33 | if (qk > max_qk) max_qk = qk; |
| 825 | 33 | } | |
| 826 | 17 | std::vector<float> scratch_P(max_qk); | |
| 827 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
|
17 | std::vector<float> scratch_dP(max_qk); |
| 828 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | std::vector<float> scratch_dS(max_qk); |
| 829 | |||
| 830 |
2/2✓ Branch 0 taken 17 times.
✓ Branch 1 taken 33 times.
|
50 | for (int b = 0; b < batch_size; ++b) { |
| 831 | 33 | const int q_beg = cu_seqlens_q[b]; | |
| 832 | 33 | const int q_end = cu_seqlens_q[b + 1]; | |
| 833 | 33 | const int k_beg = cu_seqlens_k[b]; | |
| 834 | 33 | const int k_end = cu_seqlens_k[b + 1]; | |
| 835 | 33 | const int Lq = q_end - q_beg; | |
| 836 | 33 | const int Lk = k_end - k_beg; | |
| 837 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (Lq < 0 || Lk < 0) |
| 838 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens must be non-decreasing"); | |
| 839 |
3/4✓ Branch 0 taken 15 times.
✓ Branch 1 taken 18 times.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
|
33 | if (causal && Lq != Lk) |
| 840 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: causal requires per-sequence Lq == Lk"); | |
| 841 |
2/4✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 33 times.
|
33 | if (Lq == 0 || Lk == 0) continue; // grad rows already zero. |
| 842 | |||
| 843 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | attention_core_backward( |
| 844 | 33 | Qp + static_cast<std::size_t>(q_beg) * D, | |
| 845 | 33 | Kp + static_cast<std::size_t>(k_beg) * D, | |
| 846 | 33 | Vp + static_cast<std::size_t>(k_beg) * D, | |
| 847 | 33 | dOp + static_cast<std::size_t>(q_beg) * D, | |
| 848 | /*mask=*/nullptr, | ||
| 849 | 33 | Lq, Lk, D, num_heads, causal, | |
| 850 | 33 | dQp + static_cast<std::size_t>(q_beg) * D, | |
| 851 | 33 | dKp + static_cast<std::size_t>(k_beg) * D, | |
| 852 | 33 | dVp + static_cast<std::size_t>(k_beg) * D, | |
| 853 | 33 | scratch_P.data(), scratch_dP.data(), scratch_dS.data()); | |
| 854 | 33 | } | |
| 855 | 17 | } | |
| 856 | |||
| 857 | // ─── flash_attention_qkvo_backward ───────────────────────────────────────── | ||
| 858 | |||
| 859 | 9 | void flash_attention_qkvo_backward(const ::brotensor::Tensor& X, | |
| 860 | const ::brotensor::Tensor* Ctx, | ||
| 861 | const ::brotensor::Tensor& Wq, | ||
| 862 | const ::brotensor::Tensor* bq, | ||
| 863 | const ::brotensor::Tensor& Wk, | ||
| 864 | const ::brotensor::Tensor* bk, | ||
| 865 | const ::brotensor::Tensor& Wv, | ||
| 866 | const ::brotensor::Tensor* bv, | ||
| 867 | const ::brotensor::Tensor& Wo, | ||
| 868 | const ::brotensor::Tensor* bo, | ||
| 869 | const float* d_mask, | ||
| 870 | int num_heads, | ||
| 871 | bool causal, | ||
| 872 | const ::brotensor::Tensor& dO, | ||
| 873 | ::brotensor::Tensor& dX, | ||
| 874 | ::brotensor::Tensor* dCtx, | ||
| 875 | ::brotensor::Tensor& dWq, | ||
| 876 | ::brotensor::Tensor* dbq, | ||
| 877 | ::brotensor::Tensor& dWk, | ||
| 878 | ::brotensor::Tensor* dbk, | ||
| 879 | ::brotensor::Tensor& dWv, | ||
| 880 | ::brotensor::Tensor* dbv, | ||
| 881 | ::brotensor::Tensor& dWo, | ||
| 882 | ::brotensor::Tensor* dbo) { | ||
| 883 | 9 | const bool self_attn = (Ctx == nullptr); | |
| 884 |
3/4✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
9 | if (self_attn && dCtx != nullptr) |
| 885 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: dCtx must be null when Ctx is null"); | |
| 886 |
3/4✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
9 | if (!self_attn && dCtx == nullptr) |
| 887 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: dCtx must be non-null when Ctx is non-null"); | |
| 888 | 45 | auto bias_ok = [](const Tensor* b, const Tensor* db) { | |
| 889 | 36 | return static_cast<bool>(b) == static_cast<bool>(db); | |
| 890 | }; | ||
| 891 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
18 | if (!bias_ok(bq, dbq) || !bias_ok(bk, dbk) || |
| 892 | 9 | !bias_ok(bv, dbv) || !bias_ok(bo, dbo)) | |
| 893 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: bias/grad-bias presence mismatch"); | |
| 894 | |||
| 895 | 9 | const int Lq = X.rows; | |
| 896 | 9 | const int D = X.cols; | |
| 897 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
|
9 | const Tensor& kv_src = Ctx ? *Ctx : X; |
| 898 | 9 | const int Lk = kv_src.rows; | |
| 899 | 9 | const int D_ctx = kv_src.cols; | |
| 900 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
18 | if (Wq.rows != D || Wq.cols != D || |
| 901 | 9 | Wk.rows != D || Wk.cols != D_ctx || | |
| 902 | 9 | Wv.rows != D || Wv.cols != D_ctx || | |
| 903 | 9 | Wo.rows != D || Wo.cols != D) | |
| 904 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: shape mismatch"); | |
| 905 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (dO.rows != Lq || dO.cols != D) |
| 906 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: dO shape mismatch"); | |
| 907 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (num_heads <= 0 || D % num_heads != 0) |
| 908 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: num_heads must divide D"); | |
| 909 |
3/4✓ Branch 0 taken 2 times.
✓ Branch 1 taken 7 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
|
9 | if (causal && Lq != Lk) |
| 910 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: causal requires Lq == Lk"); | |
| 911 | |||
| 912 | 9 | ensure_f32(dX, Lq, D); | |
| 913 | 9 | dX.zero(); | |
| 914 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
|
9 | if (!self_attn) { |
| 915 | 3 | ensure_f32(*dCtx, Lk, D_ctx); | |
| 916 | 3 | dCtx->zero(); | |
| 917 | 3 | } | |
| 918 |
3/6✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 9 times.
|
9 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 919 | |||
| 920 | 9 | const float* Xp = X.host_f32(); | |
| 921 | 9 | const float* kvp = kv_src.host_f32(); | |
| 922 | |||
| 923 | // ── 1. Recompute forward projections. ── | ||
| 924 | 9 | std::vector<float> Q(static_cast<std::size_t>(Lq) * D); | |
| 925 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> K(static_cast<std::size_t>(Lk) * D); |
| 926 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> V(static_cast<std::size_t>(Lk) * D); |
| 927 |
6/10✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 3 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 9 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 9 times.
✗ Branch 9 not taken.
|
9 | linear_proj(Xp, Wq.host_f32(), bq ? bq->host_f32() : nullptr, |
| 928 | 9 | Q.data(), Lq, D, D); | |
| 929 |
6/10✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 3 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 9 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 9 times.
✗ Branch 9 not taken.
|
9 | linear_proj(kvp, Wk.host_f32(), bk ? bk->host_f32() : nullptr, |
| 930 | 9 | K.data(), Lk, D_ctx, D); | |
| 931 |
6/10✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 3 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 9 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 9 times.
✗ Branch 9 not taken.
|
9 | linear_proj(kvp, Wv.host_f32(), bv ? bv->host_f32() : nullptr, |
| 932 | 9 | V.data(), Lk, D_ctx, D); | |
| 933 | |||
| 934 | // ── 2. Recompute O_attn (post-attention, pre-Wo). ── | ||
| 935 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> O_attn(static_cast<std::size_t>(Lq) * D); |
| 936 |
2/4✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
|
18 | attention_core(Q.data(), K.data(), V.data(), d_mask, |
| 937 | 9 | Lq, Lk, D, num_heads, causal, O_attn.data(), nullptr); | |
| 938 | |||
| 939 | // ── 3. Wo backward: dO_attn = dO·Wo, dWo += dO^T·O_attn, dbo += colsum. ── | ||
| 940 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dO_attn(static_cast<std::size_t>(Lq) * D, 0.0f); |
| 941 |
3/6✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 9 times.
✗ Branch 5 not taken.
|
18 | linear_proj_backward(O_attn.data(), Wo.host_f32(), dO.host_f32(), |
| 942 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | dO_attn.data(), dWo.host_f32_mut(), |
| 943 |
3/4✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
9 | dbo ? dbo->host_f32_mut() : nullptr, |
| 944 | 9 | Lq, D, D); | |
| 945 | |||
| 946 | // ── 4. Attention core backward → dQ, dK, dV. ── | ||
| 947 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dQ(static_cast<std::size_t>(Lq) * D); |
| 948 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dK(static_cast<std::size_t>(Lk) * D); |
| 949 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dV(static_cast<std::size_t>(Lk) * D); |
| 950 |
2/4✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
|
18 | attention_core_backward(Q.data(), K.data(), V.data(), dO_attn.data(), |
| 951 | 9 | d_mask, Lq, Lk, D, num_heads, causal, | |
| 952 | 9 | dQ.data(), dK.data(), dV.data()); | |
| 953 | |||
| 954 | // ── 5. Q/K/V projection backward. ── | ||
| 955 | // dX accumulates the Q path (and, for self-attn, the K and V paths). | ||
| 956 | // dWq/dWk/dWv and dbq/dbk/dbv accumulate. | ||
| 957 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | float* dXp = dX.host_f32_mut(); |
| 958 |
2/4✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
|
18 | linear_proj_backward(Xp, Wq.host_f32(), dQ.data(), dXp, |
| 959 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | dWq.host_f32_mut(), |
| 960 |
3/4✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
9 | dbq ? dbq->host_f32_mut() : nullptr, |
| 961 | 9 | Lq, D, D); | |
| 962 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
|
9 | if (self_attn) { |
| 963 |
2/4✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
12 | linear_proj_backward(Xp, Wk.host_f32(), dK.data(), dXp, |
| 964 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | dWk.host_f32_mut(), |
| 965 |
3/4✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
|
6 | dbk ? dbk->host_f32_mut() : nullptr, |
| 966 | 6 | Lk, D_ctx, D); | |
| 967 |
2/4✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
12 | linear_proj_backward(Xp, Wv.host_f32(), dV.data(), dXp, |
| 968 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | dWv.host_f32_mut(), |
| 969 |
3/4✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
|
6 | dbv ? dbv->host_f32_mut() : nullptr, |
| 970 | 6 | Lk, D_ctx, D); | |
| 971 | 6 | } else { | |
| 972 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | float* dCp = dCtx->host_f32_mut(); |
| 973 |
2/4✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
6 | linear_proj_backward(kvp, Wk.host_f32(), dK.data(), dCp, |
| 974 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | dWk.host_f32_mut(), |
| 975 |
3/4✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
|
3 | dbk ? dbk->host_f32_mut() : nullptr, |
| 976 | 3 | Lk, D_ctx, D); | |
| 977 |
2/4✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
6 | linear_proj_backward(kvp, Wv.host_f32(), dV.data(), dCp, |
| 978 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | dWv.host_f32_mut(), |
| 979 |
3/4✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
|
3 | dbv ? dbv->host_f32_mut() : nullptr, |
| 980 | 3 | Lk, D_ctx, D); | |
| 981 | } | ||
| 982 | 9 | } | |
| 983 | |||
| 984 | } // namespace brotensor::detail::cpu | ||
| 985 |