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.
|
1026 | parallel_for(static_cast<std::size_t>(M), [&](std::size_t ii) { |
| 72 | 919 | const int i = static_cast<int>(ii); | |
| 73 | 919 | const float* xr = In + static_cast<std::size_t>(i) * Din; | |
| 74 | 919 | float* orow = Out + static_cast<std::size_t>(i) * Dout; | |
| 75 |
2/2✓ Branch 0 taken 29639 times.
✓ Branch 1 taken 919 times.
|
30558 | for (int n = 0; n < Dout; ++n) { |
| 76 | 29639 | const float* wr = W + static_cast<std::size_t>(n) * Din; | |
| 77 |
2/2✓ Branch 0 taken 13192 times.
✓ Branch 1 taken 16447 times.
|
29639 | float acc = bias ? bias[n] : 0.0f; |
| 78 |
2/2✓ Branch 0 taken 860445 times.
✓ Branch 1 taken 29639 times.
|
890084 | for (int k = 0; k < Din; ++k) acc += xr[k] * wr[k]; |
| 79 | 29639 | orow[n] = acc; | |
| 80 | 29639 | } | |
| 81 | 919 | }); | |
| 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.
|
326 | parallel_for(static_cast<std::size_t>(M), [&](std::size_t ii) { |
| 95 | 290 | const int i = static_cast<int>(ii); | |
| 96 | 290 | const float* dor = dOut + static_cast<std::size_t>(i) * Dout; | |
| 97 | 290 | float* dir = dIn + static_cast<std::size_t>(i) * Din; | |
| 98 |
2/2✓ Branch 0 taken 7416 times.
✓ Branch 1 taken 290 times.
|
7706 | for (int k = 0; k < Din; ++k) { |
| 99 | 7416 | float acc = 0.0f; | |
| 100 |
2/2✓ Branch 0 taken 166949 times.
✓ Branch 1 taken 7416 times.
|
174365 | for (int n = 0; n < Dout; ++n) |
| 101 | 166949 | acc += dor[n] * W[static_cast<std::size_t>(n) * Din + k]; | |
| 102 | 7416 | dir[k] += acc; | |
| 103 | 7416 | } | |
| 104 | 290 | }); | |
| 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.
|
1346 | parallel_for(static_cast<std::size_t>(Dout), [&](std::size_t ni) { |
| 109 | 1310 | const int n = static_cast<int>(ni); | |
| 110 | 1310 | float dbacc = 0.0f; | |
| 111 |
2/2✓ Branch 0 taken 27199 times.
✓ Branch 1 taken 1310 times.
|
28509 | for (int k = 0; k < Din; ++k) { |
| 112 | 27199 | float dw = 0.0f; | |
| 113 |
2/2✓ Branch 0 taken 173300 times.
✓ Branch 1 taken 27199 times.
|
200499 | for (int i = 0; i < M; ++i) |
| 114 | 346600 | dw += dOut[static_cast<std::size_t>(i) * Dout + n] * | |
| 115 | 173300 | In[static_cast<std::size_t>(i) * Din + k]; | |
| 116 | 27199 | dW[static_cast<std::size_t>(n) * Din + k] += dw; | |
| 117 | 27199 | } | |
| 118 |
2/2✓ Branch 0 taken 10904 times.
✓ Branch 1 taken 1310 times.
|
12214 | for (int i = 0; i < M; ++i) |
| 119 | 10904 | dbacc += dOut[static_cast<std::size_t>(i) * Dout + n]; | |
| 120 |
2/2✓ Branch 0 taken 436 times.
✓ Branch 1 taken 874 times.
|
1310 | if (db) db[n] += dbacc; |
| 121 | 1310 | }); | |
| 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 | 180 | 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 | 180 | const int hd = D / H; | |
| 160 |
2/2✓ Branch 0 taken 46 times.
✓ Branch 1 taken 134 times.
|
180 | if (Dkv < 0) Dkv = D; |
| 161 | 180 | 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 | 9328 | auto process_q = [&](int q, float* scores) { | |
| 168 | 9148 | const int aq = q + q_offset; // absolute causal position of this query | |
| 169 |
2/2✓ Branch 0 taken 27508 times.
✓ Branch 1 taken 9148 times.
|
36656 | for (int h = 0; h < H; ++h) { |
| 170 | 27508 | const int off = h * hd; // Q/O head offset (Dq-wide) | |
| 171 | 27508 | const int off_kv = (h / group) * hd; // K/V head offset (Dkv-wide) | |
| 172 | 27508 | float m = -1e30f; | |
| 173 |
2/2✓ Branch 0 taken 4262439 times.
✓ Branch 1 taken 27508 times.
|
4289947 | for (int k = 0; k < Lk; ++k) { |
| 174 |
4/4✓ Branch 0 taken 310173 times.
✓ Branch 1 taken 3952266 times.
✓ Branch 2 taken 229307 times.
✓ Branch 3 taken 80866 times.
|
4262439 | if (mask && mask[k] <= 0.5f) { scores[k] = -1e30f; continue; } |
| 175 |
4/4✓ Branch 0 taken 1366901 times.
✓ Branch 1 taken 2814672 times.
✓ Branch 2 taken 682860 times.
✓ Branch 3 taken 684041 times.
|
4181573 | if (causal && k > aq) { scores[k] = -1e30f; continue; } |
| 176 |
6/6✓ Branch 0 taken 708993 times.
✓ Branch 1 taken 2788539 times.
✓ Branch 2 taken 62995 times.
✓ Branch 3 taken 645998 times.
✓ Branch 4 taken 29726 times.
✓ Branch 5 taken 33269 times.
|
3497532 | if (causal && window > 0 && k <= aq - window) { scores[k] = -1e30f; continue; } |
| 177 |
6/6✓ Branch 0 taken 2790948 times.
✓ Branch 1 taken 673315 times.
✓ Branch 2 taken 102494 times.
✓ Branch 3 taken 2688454 times.
✓ Branch 4 taken 32222 times.
✓ Branch 5 taken 70272 times.
|
3464263 | if (!causal && window > 0 && std::abs(k - aq) > window / 2) { scores[k] = -1e30f; continue; } |
| 178 | 3393991 | const float* qr = Q + static_cast<std::size_t>(q) * D + off; | |
| 179 | 3393991 | const float* kr = K + static_cast<std::size_t>(k) * Dkv + off_kv; | |
| 180 | 3393991 | float dot = 0.0f; | |
| 181 |
2/2✓ Branch 0 taken 206269430 times.
✓ Branch 1 taken 3393991 times.
|
209663421 | for (int d = 0; d < hd; ++d) dot += qr[d] * kr[d]; |
| 182 | 3393991 | const float s = dot * inv_sqrt; | |
| 183 | 3393991 | scores[k] = s; | |
| 184 |
2/2✓ Branch 0 taken 3233995 times.
✓ Branch 1 taken 159996 times.
|
3393991 | if (s > m) m = s; |
| 185 | 3393991 | } | |
| 186 | 27508 | const bool empty = (m <= -1e29f); | |
| 187 | 27508 | float sum = 0.0f; | |
| 188 |
2/2✓ Branch 0 taken 4346259 times.
✓ Branch 1 taken 27508 times.
|
4373767 | for (int k = 0; k < Lk; ++k) { |
| 189 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 4346259 times.
|
4346259 | const float e = empty ? 0.0f : std::exp(scores[k] - m); |
| 190 | 4346259 | scores[k] = e; | |
| 191 | 4346259 | sum += e; | |
| 192 | 4346259 | } | |
| 193 |
1/2✓ Branch 0 taken 27508 times.
✗ Branch 1 not taken.
|
27508 | const float inv = sum > 0.0f ? 1.0f / sum : 0.0f; |
| 194 | 27508 | float* orow = O + static_cast<std::size_t>(q) * D + off; | |
| 195 |
2/2✓ Branch 0 taken 1522340 times.
✓ Branch 1 taken 27508 times.
|
1549848 | for (int d = 0; d < hd; ++d) { |
| 196 | 1522340 | float acc = 0.0f; | |
| 197 |
2/2✓ Branch 0 taken 221728220 times.
✓ Branch 1 taken 1522340 times.
|
223250560 | for (int k = 0; k < Lk; ++k) { |
| 198 | 443456440 | acc += scores[k] * inv * | |
| 199 | 221728220 | V[static_cast<std::size_t>(k) * Dkv + off_kv + d]; | |
| 200 | 221728220 | } | |
| 201 | 1522340 | orow[d] = acc; | |
| 202 | 1522340 | } | |
| 203 |
1/2✓ Branch 0 taken 27508 times.
✗ Branch 1 not taken.
|
27508 | if (P) { |
| 204 | ✗ | float* prow = | |
| 205 | ✗ | P + (static_cast<std::size_t>(h) * Lq + q) * Lk; | |
| 206 | ✗ | for (int k = 0; k < Lk; ++k) prow[k] = scores[k] * inv; | |
| 207 | ✗ | } | |
| 208 | 27508 | } | |
| 209 | 9148 | }; | |
| 210 | |||
| 211 |
2/2✓ Branch 0 taken 69 times.
✓ Branch 1 taken 111 times.
|
180 | if (scratch_scores) { |
| 212 | // Varlen path: caller owns a single Lk-sized buffer reused across | ||
| 213 | // many sequential calls to this function. Stay serial (see the | ||
| 214 | // threading note above). | ||
| 215 |
2/2✓ Branch 0 taken 2714 times.
✓ Branch 1 taken 69 times.
|
2783 | for (int q = 0; q < Lq; ++q) process_q(q, scratch_scores); |
| 216 | 69 | } else { | |
| 217 | // Every other caller: parallelize over q, each q privately slicing | ||
| 218 | // a freshly allocated Lq*Lk buffer so concurrent q's never share a | ||
| 219 | // `scores` row. | ||
| 220 | 222 | std::vector<float> scores_all(static_cast<std::size_t>(Lq) * | |
| 221 | 111 | static_cast<std::size_t>(Lk)); | |
| 222 |
2/4✓ Branch 0 taken 111 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 111 times.
✗ Branch 3 not taken.
|
6544 | parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t qi) { |
| 223 | 6433 | const int q = static_cast<int>(qi); | |
| 224 | 6433 | process_q(q, scores_all.data() + static_cast<std::size_t>(q) * Lk); | |
| 225 | 6433 | }); | |
| 226 | 111 | } | |
| 227 | 180 | } | |
| 228 | |||
| 229 | // Per-head flash-attention backward over pre-projected Q/K/V. Recompute P, | ||
| 230 | // then dV = P^T·dO_attn, dP = dO_attn·V^T, dS = P*(dP - D_q)*scale, | ||
| 231 | // dQ = dS·K, dK = dS^T·Q. dQ/dK/dV are OVERWRITTEN. | ||
| 232 | // | ||
| 233 | // scratch_P/scratch_dP/scratch_dS, if all non-null, must each point at a | ||
| 234 | // caller-owned buffer of at least Lq*Lk floats and are reused verbatim (no | ||
| 235 | // allocation here) — callers looping over many short sequences (e.g. varlen | ||
| 236 | // attention backward) hoist one triple of buffers sized to the batch's max | ||
| 237 | // Lq*Lk instead of paying three heap allocations per call. Leaving any of | ||
| 238 | // them null falls back to a local per-call allocation, matching the original | ||
| 239 | // single-call behavior. | ||
| 240 | // | ||
| 241 | // THREADING: when the scratch triple is null (every caller except the | ||
| 242 | // varlen family), the outer h loop is safe to parallelize — each head h | ||
| 243 | // writes only its own column range [h*hd, h*hd+hd) of dQ/dK/dV (never | ||
| 244 | // touched by another h), so the only shared mutable state is P/dP/dS, which | ||
| 245 | // the parallel path gives each h its own private Lq*Lk slice of a | ||
| 246 | // freshly-allocated H*Lq*Lk buffer instead of one triple reused across h. | ||
| 247 | // When the scratch triple is supplied, the caller | ||
| 248 | // (flash_attention_varlen_backward) has explicitly hoisted ONE Lq*Lk-sized | ||
| 249 | // triple across many sequential calls to this function (one per sequence in | ||
| 250 | // a batch); parallelizing the h loop in that case would have every thread's | ||
| 251 | // h race on that single shared triple, so this path is left exactly as it | ||
| 252 | // was (serial, reusing the caller's buffers verbatim). | ||
| 253 | 50 | void attention_core_backward(const float* Q, const float* K, const float* V, | |
| 254 | const float* dO, const float* mask, | ||
| 255 | int Lq, int Lk, int D, int H, bool causal, | ||
| 256 | float* dQ, float* dK, float* dV, | ||
| 257 | float* scratch_P = nullptr, | ||
| 258 | float* scratch_dP = nullptr, | ||
| 259 | float* scratch_dS = nullptr) { | ||
| 260 | 50 | const int hd = D / H; | |
| 261 | 50 | const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(hd)); | |
| 262 |
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) |
| 263 | 10440 | dQ[i] = 0.0f; | |
| 264 |
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) { |
| 265 | 11432 | dK[i] = 0.0f; | |
| 266 | 11432 | dV[i] = 0.0f; | |
| 267 | 11432 | } | |
| 268 | |||
| 269 | // Per-head body. P/dP/dS must be a private Lq*Lk buffer when running | ||
| 270 | // concurrently with other heads (see threading note above) — each h | ||
| 271 | // writes only its own column range of dQ/dK/dV, so distinct h's never | ||
| 272 | // touch the same output bytes. | ||
| 273 | 196 | auto process_head = [&](int h, float* P, float* dP, float* dS) { | |
| 274 | 146 | const int off = h * hd; | |
| 275 | // Recompute P for this head. | ||
| 276 |
2/2✓ Branch 0 taken 1094 times.
✓ Branch 1 taken 146 times.
|
1240 | for (int q = 0; q < Lq; ++q) { |
| 277 | 1094 | float m = -1e30f; | |
| 278 | 1094 | float* prow = P + static_cast<std::size_t>(q) * Lk; | |
| 279 |
2/2✓ Branch 0 taken 10578 times.
✓ Branch 1 taken 1094 times.
|
11672 | for (int k = 0; k < Lk; ++k) { |
| 280 |
4/4✓ Branch 0 taken 968 times.
✓ Branch 1 taken 9610 times.
✓ Branch 2 taken 650 times.
✓ Branch 3 taken 318 times.
|
10578 | if (mask && mask[k] <= 0.5f) { prow[k] = -1e30f; continue; } |
| 281 |
4/4✓ Branch 0 taken 1662 times.
✓ Branch 1 taken 8598 times.
✓ Branch 2 taken 923 times.
✓ Branch 3 taken 739 times.
|
10260 | if (causal && k > q) { prow[k] = -1e30f; continue; } |
| 282 | 9521 | const float* qr = Q + static_cast<std::size_t>(q) * D + off; | |
| 283 | 9521 | const float* kr = K + static_cast<std::size_t>(k) * D + off; | |
| 284 | 9521 | float dot = 0.0f; | |
| 285 |
2/2✓ Branch 0 taken 104419 times.
✓ Branch 1 taken 9521 times.
|
113940 | for (int d = 0; d < hd; ++d) dot += qr[d] * kr[d]; |
| 286 | 9521 | prow[k] = dot * inv_sqrt; | |
| 287 |
2/2✓ Branch 0 taken 6630 times.
✓ Branch 1 taken 2891 times.
|
9521 | if (prow[k] > m) m = prow[k]; |
| 288 | 9521 | } | |
| 289 | 1094 | const bool empty = (m <= -1e29f); | |
| 290 | 1094 | float sum = 0.0f; | |
| 291 |
2/2✓ Branch 0 taken 10670 times.
✓ Branch 1 taken 1094 times.
|
11764 | for (int k = 0; k < Lk; ++k) { |
| 292 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10670 times.
|
10670 | const float e = empty ? 0.0f : std::exp(prow[k] - m); |
| 293 | 10670 | prow[k] = e; | |
| 294 | 10670 | sum += e; | |
| 295 | 10670 | } | |
| 296 |
1/2✓ Branch 0 taken 1094 times.
✗ Branch 1 not taken.
|
1094 | const float inv = sum > 0.0f ? 1.0f / sum : 0.0f; |
| 297 |
2/2✓ Branch 0 taken 10852 times.
✓ Branch 1 taken 1094 times.
|
11946 | for (int k = 0; k < Lk; ++k) prow[k] *= inv; |
| 298 | 1094 | } | |
| 299 | // dV[k, off+d] = sum_q P[q,k] * dO[q, off+d]. | ||
| 300 |
2/2✓ Branch 0 taken 1230 times.
✓ Branch 1 taken 146 times.
|
1376 | for (int k = 0; k < Lk; ++k) { |
| 301 | 1230 | float* dvr = dV + static_cast<std::size_t>(k) * D + off; | |
| 302 |
2/2✓ Branch 0 taken 10509 times.
✓ Branch 1 taken 1230 times.
|
11739 | for (int d = 0; d < hd; ++d) { |
| 303 | 10509 | float acc = 0.0f; | |
| 304 |
2/2✓ Branch 0 taken 111939 times.
✓ Branch 1 taken 10509 times.
|
122448 | for (int q = 0; q < Lq; ++q) |
| 305 | 223878 | acc += P[static_cast<std::size_t>(q) * Lk + k] * | |
| 306 | 111939 | dO[static_cast<std::size_t>(q) * D + off + d]; | |
| 307 | 10509 | dvr[d] = acc; | |
| 308 | 10509 | } | |
| 309 | 1230 | } | |
| 310 | // dP[q,k] = sum_d dO[q, off+d] * V[k, off+d]. | ||
| 311 |
2/2✓ Branch 0 taken 1113 times.
✓ Branch 1 taken 146 times.
|
1259 | for (int q = 0; q < Lq; ++q) { |
| 312 |
2/2✓ Branch 0 taken 10687 times.
✓ Branch 1 taken 1113 times.
|
11800 | for (int k = 0; k < Lk; ++k) { |
| 313 | 10687 | float acc = 0.0f; | |
| 314 |
2/2✓ Branch 0 taken 113636 times.
✓ Branch 1 taken 10687 times.
|
124323 | for (int d = 0; d < hd; ++d) |
| 315 | 227272 | acc += dO[static_cast<std::size_t>(q) * D + off + d] * | |
| 316 | 113636 | V[static_cast<std::size_t>(k) * D + off + d]; | |
| 317 | 10687 | dP[static_cast<std::size_t>(q) * Lk + k] = acc; | |
| 318 | 10687 | } | |
| 319 | 1113 | } | |
| 320 | // dS[q,k] = P[q,k] * (dP[q,k] - D_q) * inv_sqrt. | ||
| 321 |
2/2✓ Branch 0 taken 1090 times.
✓ Branch 1 taken 146 times.
|
1236 | for (int q = 0; q < Lq; ++q) { |
| 322 | 1090 | const float* prow = P + static_cast<std::size_t>(q) * Lk; | |
| 323 | 1090 | const float* dpr = dP + static_cast<std::size_t>(q) * Lk; | |
| 324 | 1090 | float* dsr = dS + static_cast<std::size_t>(q) * Lk; | |
| 325 | 1090 | float Dq = 0.0f; | |
| 326 |
2/2✓ Branch 0 taken 10707 times.
✓ Branch 1 taken 1090 times.
|
11797 | for (int k = 0; k < Lk; ++k) Dq += prow[k] * dpr[k]; |
| 327 |
2/2✓ Branch 0 taken 10692 times.
✓ Branch 1 taken 1090 times.
|
11782 | for (int k = 0; k < Lk; ++k) |
| 328 | 10692 | dsr[k] = prow[k] * (dpr[k] - Dq) * inv_sqrt; | |
| 329 | 1090 | } | |
| 330 | // dQ[q, off+d] = sum_k dS[q,k] * K[k, off+d]. | ||
| 331 |
2/2✓ Branch 0 taken 1104 times.
✓ Branch 1 taken 146 times.
|
1250 | for (int q = 0; q < Lq; ++q) { |
| 332 | 1104 | float* dqr = dQ + static_cast<std::size_t>(q) * D + off; | |
| 333 | 1104 | const float* dsr = dS + static_cast<std::size_t>(q) * Lk; | |
| 334 |
2/2✓ Branch 0 taken 9847 times.
✓ Branch 1 taken 1104 times.
|
10951 | for (int d = 0; d < hd; ++d) { |
| 335 | 9847 | float acc = 0.0f; | |
| 336 |
2/2✓ Branch 0 taken 113988 times.
✓ Branch 1 taken 9847 times.
|
123835 | for (int k = 0; k < Lk; ++k) |
| 337 | 227976 | acc += dsr[k] * | |
| 338 | 113988 | K[static_cast<std::size_t>(k) * D + off + d]; | |
| 339 | 9847 | dqr[d] = acc; | |
| 340 | 9847 | } | |
| 341 | 1104 | } | |
| 342 | // dK[k, off+d] = sum_q dS[q,k] * Q[q, off+d]. | ||
| 343 |
2/2✓ Branch 0 taken 1235 times.
✓ Branch 1 taken 146 times.
|
1381 | for (int k = 0; k < Lk; ++k) { |
| 344 | 1235 | float* dkr = dK + static_cast<std::size_t>(k) * D + off; | |
| 345 |
2/2✓ Branch 0 taken 10737 times.
✓ Branch 1 taken 1235 times.
|
11972 | for (int d = 0; d < hd; ++d) { |
| 346 | 10737 | float acc = 0.0f; | |
| 347 |
2/2✓ Branch 0 taken 114461 times.
✓ Branch 1 taken 10737 times.
|
125198 | for (int q = 0; q < Lq; ++q) |
| 348 | 228922 | acc += dS[static_cast<std::size_t>(q) * Lk + k] * | |
| 349 | 114461 | Q[static_cast<std::size_t>(q) * D + off + d]; | |
| 350 | 10737 | dkr[d] = acc; | |
| 351 | 10737 | } | |
| 352 | 1235 | } | |
| 353 | 146 | }; | |
| 354 | |||
| 355 |
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) { |
| 356 | // Varlen path: caller owns a single Lq*Lk-sized triple reused across | ||
| 357 | // many sequential calls to this function. Stay serial (see the | ||
| 358 | // threading note above). | ||
| 359 |
2/2✓ Branch 0 taken 76 times.
✓ Branch 1 taken 33 times.
|
109 | for (int h = 0; h < H; ++h) |
| 360 | 76 | process_head(h, scratch_P, scratch_dP, scratch_dS); | |
| 361 | 33 | } else { | |
| 362 | // Every other caller: parallelize over h, each h privately slicing | ||
| 363 | // a freshly allocated H*Lq*Lk buffer so concurrent heads never | ||
| 364 | // share a P/dP/dS segment. | ||
| 365 | 17 | const std::size_t n = static_cast<std::size_t>(Lq) * Lk; | |
| 366 | 17 | std::vector<float> P_all(static_cast<std::size_t>(H) * n); | |
| 367 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | std::vector<float> dP_all(static_cast<std::size_t>(H) * n); |
| 368 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | std::vector<float> dS_all(static_cast<std::size_t>(H) * n); |
| 369 |
2/4✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 17 times.
✗ Branch 3 not taken.
|
84 | parallel_for(static_cast<std::size_t>(H), [&](std::size_t hi) { |
| 370 | 67 | const int h = static_cast<int>(hi); | |
| 371 | 134 | process_head(h, P_all.data() + static_cast<std::size_t>(h) * n, | |
| 372 | 67 | dP_all.data() + static_cast<std::size_t>(h) * n, | |
| 373 | 67 | dS_all.data() + static_cast<std::size_t>(h) * n); | |
| 374 | 67 | }); | |
| 375 | 17 | } | |
| 376 | 50 | } | |
| 377 | |||
| 378 | 225 | inline void ensure_f32(Tensor& t, int r, int c) { | |
| 379 |
4/6✓ Branch 0 taken 13 times.
✓ Branch 1 taken 212 times.
✓ Branch 2 taken 13 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 13 times.
|
225 | if (t.rows != r || t.cols != c || t.dtype != Dtype::FP32) |
| 380 | 212 | t.resize(r, c, Dtype::FP32); | |
| 381 | 225 | } | |
| 382 | |||
| 383 | } // namespace | ||
| 384 | |||
| 385 | // ─── flash_attention_forward ─────────────────────────────────────────────── | ||
| 386 | |||
| 387 | 35 | void flash_attention_forward(const ::brotensor::Tensor& Q, | |
| 388 | const ::brotensor::Tensor& K, | ||
| 389 | const ::brotensor::Tensor& V, | ||
| 390 | const float* d_mask, | ||
| 391 | int num_heads, | ||
| 392 | bool causal, | ||
| 393 | ::brotensor::Tensor& O) { | ||
| 394 | 35 | const int Lq = Q.rows; | |
| 395 | 35 | const int Lk = K.rows; | |
| 396 | 35 | const int D = Q.cols; | |
| 397 |
1/2✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
|
35 | if (K.cols != D || V.cols != D || V.rows != Lk) |
| 398 | ✗ | throw std::runtime_error("flash_attention_forward: shape mismatch"); | |
| 399 |
1/2✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
|
35 | if (num_heads <= 0 || D % num_heads != 0) |
| 400 | ✗ | throw std::runtime_error("flash_attention_forward: num_heads must divide D"); | |
| 401 |
3/4✓ Branch 0 taken 18 times.
✓ Branch 1 taken 17 times.
✓ Branch 2 taken 18 times.
✗ Branch 3 not taken.
|
35 | if (causal && Lq != Lk) |
| 402 | ✗ | throw std::runtime_error("flash_attention_forward: causal requires Lq == Lk"); | |
| 403 | 35 | ensure_f32(O, Lq, D); | |
| 404 |
3/6✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 35 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 35 times.
|
35 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 405 | |||
| 406 | 70 | attention_core(Q.host_f32(), K.host_f32(), V.host_f32(), d_mask, | |
| 407 | 35 | Lq, Lk, D, num_heads, causal, O.host_f32_mut(), nullptr); | |
| 408 | 35 | } | |
| 409 | |||
| 410 | // ─── flash_attention_windowed_forward ────────────────────────────────────── | ||
| 411 | // | ||
| 412 | // Sliding-window causal self-attention (FP32, inference-only). Always causal. | ||
| 413 | // The Lq queries occupy the last Lq positions of a length-Lk causal sequence | ||
| 414 | // (q_offset = Lk - Lq): query row r sits at absolute position r + q_offset and | ||
| 415 | // attends keys [max(0, pos-window+1), pos]. window <= 0 is unbounded causal. | ||
| 416 | // - Lq == Lk (q_offset 0): plain causal / sliding window (e.g. the codec). | ||
| 417 | // - Lq < Lk: incremental decode — a short query block (often a single token) | ||
| 418 | // attending the whole K/V cache, the autoregressive step. Lq == 1 with | ||
| 419 | // window <= 0 attends every cached key (full attention over the cache). | ||
| 420 | // Lk >= Lq required. | ||
| 421 | 33 | void flash_attention_windowed_forward(const ::brotensor::Tensor& Q, | |
| 422 | const ::brotensor::Tensor& K, | ||
| 423 | const ::brotensor::Tensor& V, | ||
| 424 | const float* d_mask, | ||
| 425 | int num_heads, | ||
| 426 | int window, | ||
| 427 | ::brotensor::Tensor& O, | ||
| 428 | bool causal) { | ||
| 429 | 33 | const int Lq = Q.rows; | |
| 430 | 33 | const int Lk = K.rows; | |
| 431 | 33 | const int D = Q.cols; // Dq = num_heads * head_dim | |
| 432 | 33 | const int Dkv = K.cols; // n_kv * head_dim (GQA when < D) | |
| 433 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (V.cols != Dkv || V.rows != Lk) |
| 434 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: shape mismatch"); | |
| 435 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (num_heads <= 0 || D % num_heads != 0) |
| 436 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: num_heads must divide D"); | |
| 437 | 33 | const int head_dim = D / num_heads; | |
| 438 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (Dkv == 0 || Dkv % head_dim != 0) |
| 439 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: K/V width must be a head_dim multiple"); | |
| 440 | 33 | const int n_kv = Dkv / head_dim; | |
| 441 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (num_heads % n_kv != 0) |
| 442 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: num_heads must be a multiple of n_kv"); | |
| 443 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (Lk < Lq) |
| 444 | ✗ | throw std::runtime_error("flash_attention_windowed_forward: requires Lk >= Lq"); | |
| 445 | 33 | ensure_f32(O, Lq, D); | |
| 446 |
3/6✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 33 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 33 times.
|
33 | if (Lq == 0 || Lk == 0 || D == 0) return; |
| 447 | |||
| 448 | 66 | attention_core(Q.host_f32(), K.host_f32(), V.host_f32(), d_mask, | |
| 449 | 33 | Lq, Lk, D, num_heads, causal, O.host_f32_mut(), | |
| 450 | 33 | nullptr, window, /*q_offset=*/Lk - Lq, Dkv, | |
| 451 | 33 | /*group=*/num_heads / n_kv); | |
| 452 | 33 | } | |
| 453 | |||
| 454 | // ─── flash_attention_gqa_forward ─────────────────────────────────────────── | ||
| 455 | // | ||
| 456 | // GQA generalisation of flash_attention_forward, causal OR bidirectional. Q is | ||
| 457 | // num_q_heads-wide, K/V are num_kv_heads-wide (GQA group = num_q_heads/num_kv_heads; | ||
| 458 | // equal == MHA). causal == false attends every key (the bidirectional encoder | ||
| 459 | // prefill). attention_core already carries the group mapping, the causal flag, | ||
| 460 | // and the key mask, so this is a validate-and-dispatch shim — q_offset is 0 | ||
| 461 | // because the queries occupy positions [0, Lq) (Lq == Lk when causal). | ||
| 462 | 13 | void flash_attention_gqa_forward(const ::brotensor::Tensor& Q, | |
| 463 | const ::brotensor::Tensor& K, | ||
| 464 | const ::brotensor::Tensor& V, | ||
| 465 | const float* d_mask, | ||
| 466 | int num_q_heads, | ||
| 467 | int num_kv_heads, | ||
| 468 | bool causal, | ||
| 469 | ::brotensor::Tensor& O) { | ||
| 470 | 13 | const int Lq = Q.rows; | |
| 471 | 13 | const int Lk = K.rows; | |
| 472 | 13 | const int Dq = Q.cols; // num_q_heads * head_dim | |
| 473 | 13 | const int Dkv = K.cols; // num_kv_heads * head_dim | |
| 474 |
1/2✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
|
13 | if (num_q_heads <= 0 || num_kv_heads <= 0) |
| 475 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: head counts must be positive"); | |
| 476 |
1/2✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
|
13 | if (Dq % num_q_heads != 0) |
| 477 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: num_q_heads must divide Q.cols"); | |
| 478 | 13 | const int head_dim = Dq / num_q_heads; | |
| 479 |
1/2✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
|
13 | if (V.cols != Dkv || V.rows != Lk) |
| 480 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: shape mismatch"); | |
| 481 |
1/2✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
|
13 | if (Dkv != num_kv_heads * head_dim) |
| 482 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: K/V width must be num_kv_heads*head_dim"); | |
| 483 |
1/2✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
|
13 | if (num_q_heads % num_kv_heads != 0) |
| 484 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: num_kv_heads must divide num_q_heads"); | |
| 485 |
1/2✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
|
13 | if (Lk < Lq) |
| 486 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: requires Lk >= Lq"); | |
| 487 |
3/4✓ Branch 0 taken 3 times.
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
13 | if (causal && Lq != Lk) |
| 488 | ✗ | throw std::runtime_error("flash_attention_gqa_forward: causal requires Lq == Lk"); | |
| 489 | 13 | ensure_f32(O, Lq, Dq); | |
| 490 |
3/6✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 13 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 13 times.
|
13 | if (Lq == 0 || Lk == 0 || Dq == 0) return; |
| 491 | |||
| 492 | 26 | attention_core(Q.host_f32(), K.host_f32(), V.host_f32(), d_mask, | |
| 493 | 13 | Lq, Lk, Dq, num_q_heads, causal, O.host_f32_mut(), | |
| 494 | 13 | /*P=*/nullptr, /*window=*/0, /*q_offset=*/0, Dkv, | |
| 495 | 13 | /*group=*/num_q_heads / num_kv_heads); | |
| 496 | 13 | } | |
| 497 | |||
| 498 | // ─── flash_attention_varlen_forward ──────────────────────────────────────── | ||
| 499 | // | ||
| 500 | // Packed variable-length attention (Qwen3-VL window attention). Q/K/V are one | ||
| 501 | // big (total_tokens, num_heads*head_dim) tensor each; cu_seqlens_q/k are | ||
| 502 | // length B+1 INT32 prefix sums delimiting per-sequence row ranges. Sequence b | ||
| 503 | // runs attention_core over its own Q[Lq_b, D] and K/V[Lk_b, D] slice — no | ||
| 504 | // cross-sequence attention. | ||
| 505 | // | ||
| 506 | // On CPU cu_seqlens_q/k are raw host pointers (matches the existing d_mask | ||
| 507 | // convention: device pointer on GPU, host pointer on CPU). max_seqlen_q/k are | ||
| 508 | // only used by the GPU kernel for block sizing and are ignored here. | ||
| 509 | 28 | void flash_attention_varlen_forward(const ::brotensor::Tensor& Q, | |
| 510 | const ::brotensor::Tensor& K, | ||
| 511 | const ::brotensor::Tensor& V, | ||
| 512 | const int32_t* cu_seqlens_q, | ||
| 513 | const int32_t* cu_seqlens_k, | ||
| 514 | int batch_size, | ||
| 515 | int /*max_seqlen_q*/, | ||
| 516 | int /*max_seqlen_k*/, | ||
| 517 | int num_heads, | ||
| 518 | int head_dim, | ||
| 519 | bool causal, | ||
| 520 | ::brotensor::Tensor& O) { | ||
| 521 | 28 | const int total_q = Q.rows; | |
| 522 | 28 | const int total_k = K.rows; | |
| 523 | 28 | const int D = num_heads * head_dim; | |
| 524 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (Q.cols != D || K.cols != D || V.cols != D || V.rows != total_k) |
| 525 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: shape mismatch"); | |
| 526 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (num_heads <= 0 || head_dim <= 0) |
| 527 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: num_heads/head_dim must be positive"); | |
| 528 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (batch_size < 0) |
| 529 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: batch_size must be non-negative"); | |
| 530 |
2/4✗ Branch 0 not taken.
✓ Branch 1 taken 28 times.
✓ Branch 2 taken 28 times.
✗ Branch 3 not taken.
|
28 | if (batch_size > 0 && (!cu_seqlens_q || !cu_seqlens_k)) |
| 531 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_q/k required when batch_size > 0"); | |
| 532 | 28 | ensure_f32(O, total_q, D); | |
| 533 |
2/4✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 28 times.
|
28 | if (total_q == 0 || D == 0) return; |
| 534 |
2/4✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 28 times.
✗ Branch 3 not taken.
|
28 | if (cu_seqlens_q && cu_seqlens_q[0] != 0) |
| 535 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_q[0] must be 0"); | |
| 536 |
2/4✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 28 times.
✗ Branch 3 not taken.
|
28 | if (cu_seqlens_k && cu_seqlens_k[0] != 0) |
| 537 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_k[0] must be 0"); | |
| 538 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 28 times.
|
28 | if (batch_size > 0) { |
| 539 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (cu_seqlens_q[batch_size] != total_q) |
| 540 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_q[B] != total_tokens_q"); | |
| 541 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (cu_seqlens_k[batch_size] != total_k) |
| 542 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens_k[B] != total_tokens_k"); | |
| 543 | 28 | } | |
| 544 | |||
| 545 | 28 | const float* Qp = Q.host_f32(); | |
| 546 | 28 | const float* Kp = K.host_f32(); | |
| 547 | 28 | const float* Vp = V.host_f32(); | |
| 548 | 28 | float* Op = O.host_f32_mut(); | |
| 549 | |||
| 550 | // Hoist attention_core's per-call `scores` scratch (size Lk) to one | ||
| 551 | // buffer sized for the widest sequence in the batch, reused for every | ||
| 552 | // sequence below — batches of many short sequences would otherwise pay | ||
| 553 | // one small heap allocation per sequence. Negative per-sequence lengths | ||
| 554 | // (malformed cu_seqlens) are skipped here; the loop below throws on them | ||
| 555 | // before attention_core is ever called for that b. | ||
| 556 | 28 | int max_lk = 0; | |
| 557 |
2/2✓ Branch 0 taken 69 times.
✓ Branch 1 taken 28 times.
|
97 | for (int b = 0; b < batch_size; ++b) { |
| 558 | 69 | const int lk = cu_seqlens_k[b + 1] - cu_seqlens_k[b]; | |
| 559 |
2/2✓ Branch 0 taken 34 times.
✓ Branch 1 taken 35 times.
|
69 | if (lk > max_lk) max_lk = lk; |
| 560 | 69 | } | |
| 561 | 28 | std::vector<float> scratch_scores(static_cast<std::size_t>(max_lk)); | |
| 562 | |||
| 563 |
2/2✓ Branch 0 taken 28 times.
✓ Branch 1 taken 69 times.
|
97 | for (int b = 0; b < batch_size; ++b) { |
| 564 | 69 | const int q_beg = cu_seqlens_q[b]; | |
| 565 | 69 | const int q_end = cu_seqlens_q[b + 1]; | |
| 566 | 69 | const int k_beg = cu_seqlens_k[b]; | |
| 567 | 69 | const int k_end = cu_seqlens_k[b + 1]; | |
| 568 | 69 | const int Lq = q_end - q_beg; | |
| 569 | 69 | const int Lk = k_end - k_beg; | |
| 570 |
1/2✓ Branch 0 taken 69 times.
✗ Branch 1 not taken.
|
69 | if (Lq < 0 || Lk < 0) |
| 571 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: cu_seqlens must be non-decreasing"); | |
| 572 |
3/4✓ Branch 0 taken 21 times.
✓ Branch 1 taken 48 times.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
|
69 | if (causal && Lq != Lk) |
| 573 | ✗ | throw std::runtime_error("flash_attention_varlen_forward: causal requires per-sequence Lq == Lk"); | |
| 574 |
1/2✓ Branch 0 taken 69 times.
✗ Branch 1 not taken.
|
69 | if (Lq == 0) continue; |
| 575 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 69 times.
|
69 | if (Lk == 0) { |
| 576 | // Fully empty K range — output rows are zero (no keys to attend to). | ||
| 577 | ✗ | for (std::size_t i = 0; i < static_cast<std::size_t>(Lq) * D; ++i) | |
| 578 | ✗ | Op[static_cast<std::size_t>(q_beg) * D + i] = 0.0f; | |
| 579 | ✗ | continue; | |
| 580 | } | ||
| 581 |
2/4✓ Branch 0 taken 69 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 69 times.
✗ Branch 3 not taken.
|
138 | attention_core(Qp + static_cast<std::size_t>(q_beg) * D, |
| 582 | 69 | Kp + static_cast<std::size_t>(k_beg) * D, | |
| 583 | 69 | Vp + static_cast<std::size_t>(k_beg) * D, | |
| 584 | /*mask=*/nullptr, | ||
| 585 | 69 | Lq, Lk, D, num_heads, causal, | |
| 586 | 69 | Op + static_cast<std::size_t>(q_beg) * D, | |
| 587 | /*P=*/nullptr, /*window=*/0, /*q_offset=*/0, /*Dkv=*/-1, | ||
| 588 | 69 | /*group=*/1, scratch_scores.data()); | |
| 589 | 69 | } | |
| 590 | 28 | } | |
| 591 | |||
| 592 | // ─── flash_attention_project_kv ──────────────────────────────────────────── | ||
| 593 | |||
| 594 | 4 | void flash_attention_project_kv(const ::brotensor::Tensor& ctx, | |
| 595 | const ::brotensor::Tensor& Wk, | ||
| 596 | const ::brotensor::Tensor* bk, | ||
| 597 | const ::brotensor::Tensor& Wv, | ||
| 598 | const ::brotensor::Tensor* bv, | ||
| 599 | ::brotensor::Tensor& K_out, | ||
| 600 | ::brotensor::Tensor& V_out) { | ||
| 601 | 4 | const int Lk = ctx.rows; | |
| 602 | 4 | const int D_ctx = ctx.cols; | |
| 603 | 4 | const int D = Wk.rows; | |
| 604 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | if (Wk.cols != D_ctx || Wv.rows != D || Wv.cols != D_ctx) |
| 605 | ✗ | throw std::runtime_error("flash_attention_project_kv: Wk/Wv shape mismatch"); | |
| 606 | 4 | ensure_f32(K_out, Lk, D); | |
| 607 | 4 | ensure_f32(V_out, Lk, D); | |
| 608 |
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; |
| 609 | |||
| 610 | 8 | linear_proj(ctx.host_f32(), Wk.host_f32(), | |
| 611 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
|
4 | bk ? bk->host_f32() : nullptr, |
| 612 | 4 | K_out.host_f32_mut(), Lk, D_ctx, D); | |
| 613 | 8 | linear_proj(ctx.host_f32(), Wv.host_f32(), | |
| 614 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
|
4 | bv ? bv->host_f32() : nullptr, |
| 615 | 4 | V_out.host_f32_mut(), Lk, D_ctx, D); | |
| 616 | 4 | } | |
| 617 | |||
| 618 | // ─── flash_attention_q_with_kv_cached_forward ────────────────────────────── | ||
| 619 | |||
| 620 | 6 | void flash_attention_q_with_kv_cached_forward(const ::brotensor::Tensor& X, | |
| 621 | const ::brotensor::Tensor& K, | ||
| 622 | const ::brotensor::Tensor& V, | ||
| 623 | const ::brotensor::Tensor& Wq, | ||
| 624 | const ::brotensor::Tensor* bq, | ||
| 625 | const ::brotensor::Tensor& Wo, | ||
| 626 | const ::brotensor::Tensor* bo, | ||
| 627 | const float* d_mask, | ||
| 628 | int num_heads, | ||
| 629 | bool causal, | ||
| 630 | ::brotensor::Tensor& O) { | ||
| 631 | 6 | const int Lq = X.rows; | |
| 632 | 6 | const int D = X.cols; | |
| 633 | 6 | const int Lk = K.rows; | |
| 634 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (K.cols != D || V.rows != Lk || V.cols != D) |
| 635 | ✗ | throw std::runtime_error("flash_attention_q_with_kv_cached_forward: K/V shape mismatch"); | |
| 636 |
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) |
| 637 | ✗ | throw std::runtime_error("flash_attention_q_with_kv_cached_forward: Wq/Wo shape mismatch"); | |
| 638 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (num_heads <= 0 || D % num_heads != 0) |
| 639 | ✗ | throw std::runtime_error("flash_attention_q_with_kv_cached_forward: num_heads must divide D"); | |
| 640 | 6 | ensure_f32(O, Lq, D); | |
| 641 |
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; |
| 642 | |||
| 643 | 6 | std::vector<float> Qp(static_cast<std::size_t>(Lq) * D); | |
| 644 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | std::vector<float> Op(static_cast<std::size_t>(Lq) * D); |
| 645 |
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, |
| 646 | 6 | Qp.data(), Lq, D, D); | |
| 647 |
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, |
| 648 | 6 | Lq, Lk, D, num_heads, causal, Op.data(), nullptr); | |
| 649 |
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, |
| 650 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | O.host_f32_mut(), Lq, D, D); |
| 651 | 6 | } | |
| 652 | |||
| 653 | // ─── flash_attention_qkvo_forward ────────────────────────────────────────── | ||
| 654 | |||
| 655 | 15 | void flash_attention_qkvo_forward(const ::brotensor::Tensor& X, | |
| 656 | const ::brotensor::Tensor* Ctx, | ||
| 657 | const ::brotensor::Tensor& Wq, | ||
| 658 | const ::brotensor::Tensor* bq, | ||
| 659 | const ::brotensor::Tensor& Wk, | ||
| 660 | const ::brotensor::Tensor* bk, | ||
| 661 | const ::brotensor::Tensor& Wv, | ||
| 662 | const ::brotensor::Tensor* bv, | ||
| 663 | const ::brotensor::Tensor& Wo, | ||
| 664 | const ::brotensor::Tensor* bo, | ||
| 665 | const float* d_mask, | ||
| 666 | int num_heads, | ||
| 667 | bool causal, | ||
| 668 | ::brotensor::Tensor& O) { | ||
| 669 | 15 | const int Lq = X.rows; | |
| 670 | 15 | const int D = X.cols; | |
| 671 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 12 times.
|
15 | const Tensor& kv_src = Ctx ? *Ctx : X; |
| 672 | 15 | const int Lk = kv_src.rows; | |
| 673 | 15 | const int D_ctx = kv_src.cols; | |
| 674 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
30 | if (Wq.rows != D || Wq.cols != D || |
| 675 | 15 | Wk.rows != D || Wk.cols != D_ctx || | |
| 676 | 15 | Wv.rows != D || Wv.cols != D_ctx || | |
| 677 | 15 | Wo.rows != D || Wo.cols != D) | |
| 678 | ✗ | throw std::runtime_error("flash_attention_qkvo_forward: shape mismatch"); | |
| 679 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | if (num_heads <= 0 || D % num_heads != 0) |
| 680 | ✗ | throw std::runtime_error("flash_attention_qkvo_forward: num_heads must divide D"); | |
| 681 | 15 | ensure_f32(O, Lq, D); | |
| 682 |
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; |
| 683 | |||
| 684 | 15 | std::vector<float> Qp(static_cast<std::size_t>(Lq) * D); | |
| 685 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | std::vector<float> Kp(static_cast<std::size_t>(Lk) * D); |
| 686 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | std::vector<float> Vp(static_cast<std::size_t>(Lk) * D); |
| 687 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | std::vector<float> Op(static_cast<std::size_t>(Lq) * D); |
| 688 | |||
| 689 |
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, |
| 690 | 15 | Qp.data(), Lq, D, 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(), Wk.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 | bk ? bk->host_f32() : nullptr, Kp.data(), Lk, D_ctx, D); |
| 693 |
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(), |
| 694 |
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); |
| 695 |
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, |
| 696 | 15 | Lq, Lk, D, num_heads, causal, Op.data(), nullptr); | |
| 697 |
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, |
| 698 |
1/2✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
|
15 | O.host_f32_mut(), Lq, D, D); |
| 699 | 15 | } | |
| 700 | |||
| 701 | // ─── flash_attention_backward ────────────────────────────────────────────── | ||
| 702 | |||
| 703 | 8 | void flash_attention_backward(const ::brotensor::Tensor& Q, | |
| 704 | const ::brotensor::Tensor& K, | ||
| 705 | const ::brotensor::Tensor& V, | ||
| 706 | const ::brotensor::Tensor& O, | ||
| 707 | const ::brotensor::Tensor& dO, | ||
| 708 | const float* d_mask, | ||
| 709 | int num_heads, | ||
| 710 | bool causal, | ||
| 711 | ::brotensor::Tensor& dQ, | ||
| 712 | ::brotensor::Tensor& dK, | ||
| 713 | ::brotensor::Tensor& dV) { | ||
| 714 | 8 | (void)O; // recompute-based; O retained in API for symmetry. | |
| 715 | 8 | const int Lq = Q.rows; | |
| 716 | 8 | const int Lk = K.rows; | |
| 717 | 8 | const int D = Q.cols; | |
| 718 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (K.cols != D || V.cols != D || V.rows != Lk) |
| 719 | ✗ | throw std::runtime_error("flash_attention_backward: Q/K/V shape mismatch"); | |
| 720 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (dO.rows != Lq || dO.cols != D) |
| 721 | ✗ | throw std::runtime_error("flash_attention_backward: dO shape mismatch"); | |
| 722 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (num_heads <= 0 || D % num_heads != 0) |
| 723 | ✗ | throw std::runtime_error("flash_attention_backward: num_heads must divide D"); | |
| 724 |
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) |
| 725 | ✗ | throw std::runtime_error("flash_attention_backward: causal requires Lq == Lk"); | |
| 726 | 8 | ensure_f32(dQ, Lq, D); | |
| 727 | 8 | ensure_f32(dK, Lk, D); | |
| 728 | 8 | ensure_f32(dV, Lk, D); | |
| 729 |
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; |
| 730 | |||
| 731 | 16 | attention_core_backward(Q.host_f32(), K.host_f32(), V.host_f32(), | |
| 732 | 8 | dO.host_f32(), d_mask, Lq, Lk, D, num_heads, | |
| 733 | 8 | causal, dQ.host_f32_mut(), dK.host_f32_mut(), | |
| 734 | 8 | dV.host_f32_mut()); | |
| 735 | 8 | } | |
| 736 | |||
| 737 | // ─── flash_attention_varlen_backward ─────────────────────────────────────── | ||
| 738 | // | ||
| 739 | // Per-sequence backward over the packed (total_tokens, D) layout. Same | ||
| 740 | // recompute math as flash_attention_backward (attention_core_backward) run | ||
| 741 | // once per sequence on its [cu_seqlens_q[b], cu_seqlens_q[b+1]) Q slice and | ||
| 742 | // [cu_seqlens_k[b], cu_seqlens_k[b+1]) K/V slice. dQ/dK/dV are OVERWRITTEN — | ||
| 743 | // attention_core_backward zeros its per-call output region, and rows outside | ||
| 744 | // any sequence's range (which can't happen for a well-formed cu_seqlens) plus | ||
| 745 | // rows whose K range is empty are explicitly zeroed here so the contract | ||
| 746 | // holds globally. | ||
| 747 | 17 | void flash_attention_varlen_backward(const ::brotensor::Tensor& Q, | |
| 748 | const ::brotensor::Tensor& K, | ||
| 749 | const ::brotensor::Tensor& V, | ||
| 750 | const ::brotensor::Tensor& O, | ||
| 751 | const ::brotensor::Tensor& dO, | ||
| 752 | const int32_t* cu_seqlens_q, | ||
| 753 | const int32_t* cu_seqlens_k, | ||
| 754 | int batch_size, | ||
| 755 | int /*max_seqlen_q*/, | ||
| 756 | int /*max_seqlen_k*/, | ||
| 757 | int num_heads, | ||
| 758 | int head_dim, | ||
| 759 | bool causal, | ||
| 760 | ::brotensor::Tensor& dQ, | ||
| 761 | ::brotensor::Tensor& dK, | ||
| 762 | ::brotensor::Tensor& dV) { | ||
| 763 | 17 | (void)O; // recompute-based; O retained in API for symmetry. | |
| 764 | 17 | const int total_q = Q.rows; | |
| 765 | 17 | const int total_k = K.rows; | |
| 766 | 17 | const int D = num_heads * head_dim; | |
| 767 |
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) |
| 768 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: Q/K/V shape mismatch"); | |
| 769 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (dO.rows != total_q || dO.cols != D) |
| 770 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: dO shape mismatch"); | |
| 771 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (num_heads <= 0 || head_dim <= 0) |
| 772 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: num_heads/head_dim must be positive"); | |
| 773 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (batch_size < 0) |
| 774 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: batch_size must be non-negative"); | |
| 775 |
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)) |
| 776 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_q/k required when batch_size > 0"); | |
| 777 | 17 | ensure_f32(dQ, total_q, D); | |
| 778 | 17 | ensure_f32(dK, total_k, D); | |
| 779 | 17 | ensure_f32(dV, total_k, D); | |
| 780 |
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; |
| 781 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (D == 0) return; |
| 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_q && cu_seqlens_q[0] != 0) |
| 783 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_q[0] must be 0"); | |
| 784 |
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) |
| 785 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_k[0] must be 0"); | |
| 786 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
|
17 | if (batch_size > 0) { |
| 787 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (cu_seqlens_q[batch_size] != total_q) |
| 788 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_q[B] != total_tokens_q"); | |
| 789 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (cu_seqlens_k[batch_size] != total_k) |
| 790 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens_k[B] != total_tokens_k"); | |
| 791 | 17 | } | |
| 792 | |||
| 793 | // Zero all grads upfront so any rows not covered by a sequence (or covered | ||
| 794 | // by an empty-K sequence) end up at exactly 0.0f without depending on the | ||
| 795 | // per-sequence path to touch them. | ||
| 796 | 17 | float* dQp = dQ.host_f32_mut(); | |
| 797 | 17 | float* dKp = dK.host_f32_mut(); | |
| 798 | 17 | float* dVp = dV.host_f32_mut(); | |
| 799 |
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; |
| 800 |
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) { |
| 801 | 5704 | dKp[i] = 0.0f; | |
| 802 | 5704 | dVp[i] = 0.0f; | |
| 803 | 5704 | } | |
| 804 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (batch_size == 0) return; |
| 805 | |||
| 806 | 17 | const float* Qp = Q.host_f32(); | |
| 807 | 17 | const float* Kp = K.host_f32(); | |
| 808 | 17 | const float* Vp = V.host_f32(); | |
| 809 | 17 | const float* dOp = dO.host_f32(); | |
| 810 | |||
| 811 | // Hoist attention_core_backward's per-call P/dP/dS scratch (each size | ||
| 812 | // Lq*Lk) to one triple of buffers sized for the widest sequence in the | ||
| 813 | // batch, reused for every sequence below — batches of many short | ||
| 814 | // sequences would otherwise pay three small heap allocations per | ||
| 815 | // sequence. Malformed (negative) per-sequence lengths are skipped in this | ||
| 816 | // sizing pass; the loop below throws on them before attention_core_backward | ||
| 817 | // is ever called for that b, and casting a negative length to size_t here | ||
| 818 | // could otherwise wrap to a huge value and provoke a bogus allocation. | ||
| 819 | 17 | std::size_t max_qk = 0; | |
| 820 |
2/2✓ Branch 0 taken 33 times.
✓ Branch 1 taken 17 times.
|
50 | for (int b = 0; b < batch_size; ++b) { |
| 821 | 33 | const int lq = cu_seqlens_q[b + 1] - cu_seqlens_q[b]; | |
| 822 | 33 | const int lk = cu_seqlens_k[b + 1] - cu_seqlens_k[b]; | |
| 823 |
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; |
| 824 | 33 | const std::size_t qk = | |
| 825 | 33 | static_cast<std::size_t>(lq) * static_cast<std::size_t>(lk); | |
| 826 |
2/2✓ Branch 0 taken 11 times.
✓ Branch 1 taken 22 times.
|
33 | if (qk > max_qk) max_qk = qk; |
| 827 | 33 | } | |
| 828 | 17 | std::vector<float> scratch_P(max_qk); | |
| 829 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 17 times.
|
17 | std::vector<float> scratch_dP(max_qk); |
| 830 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | std::vector<float> scratch_dS(max_qk); |
| 831 | |||
| 832 |
2/2✓ Branch 0 taken 17 times.
✓ Branch 1 taken 33 times.
|
50 | for (int b = 0; b < batch_size; ++b) { |
| 833 | 33 | const int q_beg = cu_seqlens_q[b]; | |
| 834 | 33 | const int q_end = cu_seqlens_q[b + 1]; | |
| 835 | 33 | const int k_beg = cu_seqlens_k[b]; | |
| 836 | 33 | const int k_end = cu_seqlens_k[b + 1]; | |
| 837 | 33 | const int Lq = q_end - q_beg; | |
| 838 | 33 | const int Lk = k_end - k_beg; | |
| 839 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | if (Lq < 0 || Lk < 0) |
| 840 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: cu_seqlens must be non-decreasing"); | |
| 841 |
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) |
| 842 | ✗ | throw std::runtime_error("flash_attention_varlen_backward: causal requires per-sequence Lq == Lk"); | |
| 843 |
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. |
| 844 | |||
| 845 |
1/2✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
|
33 | attention_core_backward( |
| 846 | 33 | Qp + static_cast<std::size_t>(q_beg) * D, | |
| 847 | 33 | Kp + static_cast<std::size_t>(k_beg) * D, | |
| 848 | 33 | Vp + static_cast<std::size_t>(k_beg) * D, | |
| 849 | 33 | dOp + static_cast<std::size_t>(q_beg) * D, | |
| 850 | /*mask=*/nullptr, | ||
| 851 | 33 | Lq, Lk, D, num_heads, causal, | |
| 852 | 33 | dQp + static_cast<std::size_t>(q_beg) * D, | |
| 853 | 33 | dKp + static_cast<std::size_t>(k_beg) * D, | |
| 854 | 33 | dVp + static_cast<std::size_t>(k_beg) * D, | |
| 855 | 33 | scratch_P.data(), scratch_dP.data(), scratch_dS.data()); | |
| 856 | 33 | } | |
| 857 | 17 | } | |
| 858 | |||
| 859 | // ─── flash_attention_qkvo_backward ───────────────────────────────────────── | ||
| 860 | |||
| 861 | 9 | void flash_attention_qkvo_backward(const ::brotensor::Tensor& X, | |
| 862 | const ::brotensor::Tensor* Ctx, | ||
| 863 | const ::brotensor::Tensor& Wq, | ||
| 864 | const ::brotensor::Tensor* bq, | ||
| 865 | const ::brotensor::Tensor& Wk, | ||
| 866 | const ::brotensor::Tensor* bk, | ||
| 867 | const ::brotensor::Tensor& Wv, | ||
| 868 | const ::brotensor::Tensor* bv, | ||
| 869 | const ::brotensor::Tensor& Wo, | ||
| 870 | const ::brotensor::Tensor* bo, | ||
| 871 | const float* d_mask, | ||
| 872 | int num_heads, | ||
| 873 | bool causal, | ||
| 874 | const ::brotensor::Tensor& dO, | ||
| 875 | ::brotensor::Tensor& dX, | ||
| 876 | ::brotensor::Tensor* dCtx, | ||
| 877 | ::brotensor::Tensor& dWq, | ||
| 878 | ::brotensor::Tensor* dbq, | ||
| 879 | ::brotensor::Tensor& dWk, | ||
| 880 | ::brotensor::Tensor* dbk, | ||
| 881 | ::brotensor::Tensor& dWv, | ||
| 882 | ::brotensor::Tensor* dbv, | ||
| 883 | ::brotensor::Tensor& dWo, | ||
| 884 | ::brotensor::Tensor* dbo) { | ||
| 885 | 9 | const bool self_attn = (Ctx == nullptr); | |
| 886 |
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) |
| 887 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: dCtx must be null when Ctx is null"); | |
| 888 |
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) |
| 889 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: dCtx must be non-null when Ctx is non-null"); | |
| 890 | 45 | auto bias_ok = [](const Tensor* b, const Tensor* db) { | |
| 891 | 36 | return static_cast<bool>(b) == static_cast<bool>(db); | |
| 892 | }; | ||
| 893 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
18 | if (!bias_ok(bq, dbq) || !bias_ok(bk, dbk) || |
| 894 | 9 | !bias_ok(bv, dbv) || !bias_ok(bo, dbo)) | |
| 895 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: bias/grad-bias presence mismatch"); | |
| 896 | |||
| 897 | 9 | const int Lq = X.rows; | |
| 898 | 9 | const int D = X.cols; | |
| 899 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
|
9 | const Tensor& kv_src = Ctx ? *Ctx : X; |
| 900 | 9 | const int Lk = kv_src.rows; | |
| 901 | 9 | const int D_ctx = kv_src.cols; | |
| 902 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
18 | if (Wq.rows != D || Wq.cols != D || |
| 903 | 9 | Wk.rows != D || Wk.cols != D_ctx || | |
| 904 | 9 | Wv.rows != D || Wv.cols != D_ctx || | |
| 905 | 9 | Wo.rows != D || Wo.cols != D) | |
| 906 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: shape mismatch"); | |
| 907 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (dO.rows != Lq || dO.cols != D) |
| 908 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: dO shape mismatch"); | |
| 909 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | if (num_heads <= 0 || D % num_heads != 0) |
| 910 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: num_heads must divide D"); | |
| 911 |
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) |
| 912 | ✗ | throw std::runtime_error("flash_attention_qkvo_backward: causal requires Lq == Lk"); | |
| 913 | |||
| 914 | 9 | ensure_f32(dX, Lq, D); | |
| 915 | 9 | dX.zero(); | |
| 916 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
|
9 | if (!self_attn) { |
| 917 | 3 | ensure_f32(*dCtx, Lk, D_ctx); | |
| 918 | 3 | dCtx->zero(); | |
| 919 | 3 | } | |
| 920 |
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; |
| 921 | |||
| 922 | 9 | const float* Xp = X.host_f32(); | |
| 923 | 9 | const float* kvp = kv_src.host_f32(); | |
| 924 | |||
| 925 | // ── 1. Recompute forward projections. ── | ||
| 926 | 9 | std::vector<float> Q(static_cast<std::size_t>(Lq) * D); | |
| 927 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> K(static_cast<std::size_t>(Lk) * D); |
| 928 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> V(static_cast<std::size_t>(Lk) * 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(Xp, Wq.host_f32(), bq ? bq->host_f32() : nullptr, |
| 930 | 9 | Q.data(), Lq, D, 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, Wk.host_f32(), bk ? bk->host_f32() : nullptr, |
| 932 | 9 | K.data(), Lk, D_ctx, D); | |
| 933 |
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, |
| 934 | 9 | V.data(), Lk, D_ctx, D); | |
| 935 | |||
| 936 | // ── 2. Recompute O_attn (post-attention, pre-Wo). ── | ||
| 937 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> O_attn(static_cast<std::size_t>(Lq) * D); |
| 938 |
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, |
| 939 | 9 | Lq, Lk, D, num_heads, causal, O_attn.data(), nullptr); | |
| 940 | |||
| 941 | // ── 3. Wo backward: dO_attn = dO·Wo, dWo += dO^T·O_attn, dbo += colsum. ── | ||
| 942 |
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); |
| 943 |
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(), |
| 944 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | dO_attn.data(), dWo.host_f32_mut(), |
| 945 |
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, |
| 946 | 9 | Lq, D, D); | |
| 947 | |||
| 948 | // ── 4. Attention core backward → dQ, dK, dV. ── | ||
| 949 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dQ(static_cast<std::size_t>(Lq) * D); |
| 950 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dK(static_cast<std::size_t>(Lk) * D); |
| 951 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | std::vector<float> dV(static_cast<std::size_t>(Lk) * D); |
| 952 |
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(), |
| 953 | 9 | d_mask, Lq, Lk, D, num_heads, causal, | |
| 954 | 9 | dQ.data(), dK.data(), dV.data()); | |
| 955 | |||
| 956 | // ── 5. Q/K/V projection backward. ── | ||
| 957 | // dX accumulates the Q path (and, for self-attn, the K and V paths). | ||
| 958 | // dWq/dWk/dWv and dbq/dbk/dbv accumulate. | ||
| 959 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | float* dXp = dX.host_f32_mut(); |
| 960 |
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, |
| 961 |
1/2✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
|
9 | dWq.host_f32_mut(), |
| 962 |
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, |
| 963 | 9 | Lq, D, D); | |
| 964 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 3 times.
|
9 | if (self_attn) { |
| 965 |
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, |
| 966 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | dWk.host_f32_mut(), |
| 967 |
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, |
| 968 | 6 | Lk, D_ctx, D); | |
| 969 |
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, |
| 970 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | dWv.host_f32_mut(), |
| 971 |
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, |
| 972 | 6 | Lk, D_ctx, D); | |
| 973 | 6 | } else { | |
| 974 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | float* dCp = dCtx->host_f32_mut(); |
| 975 |
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, |
| 976 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | dWk.host_f32_mut(), |
| 977 |
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, |
| 978 | 3 | Lk, D_ctx, D); | |
| 979 |
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, |
| 980 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | dWv.host_f32_mut(), |
| 981 |
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, |
| 982 | 3 | Lk, D_ctx, D); | |
| 983 | } | ||
| 984 | 9 | } | |
| 985 | |||
| 986 | } // namespace brotensor::detail::cpu | ||
| 987 |