GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 99.5% 423 / 0 / 425
Functions: 100.0% 23 / 0 / 23
Branches: 69.8% 247 / 0 / 354
src/cpu/cross_attention.cpp
Line Branch Exec Source
1 // CPU backend — cross-attention family (CHUNK 5).
2 //
3 // Ground truth: src/cuda/cross_attention.cu and
4 // src/cuda/cross_attention_with_attn.cu.
5 //
6 // Four ops:
7 // * cross_attention_forward — public forward, O only.
8 // * cross_attention_forward_with_attn — forward + head-averaged AttnAvg,
9 // optional pre-softmax logit bias.
10 // * cross_attention_forward_train — forward, emits per-head caches.
11 // * cross_attention_backward — backward, accumulates dW*.
12 //
13 // DTYPE DECISIONS
14 // The op_table groups these as "FP16 inference + FP32 train". The CPU
15 // backend is FP32-only (per CLAUDE.md). All four CPU impls run the FP32
16 // scalar math identical to cross_attention.cu's FP32 train-core kernels.
17 // - cross_attention_forward: the CUDA op delegates to the FP16 flash
18 // path for FP16 inputs and to the FP32 train-core for FP32 inputs. The
19 // CPU op always runs the FP32 train-core math. The parity test feeds
20 // FP32 to CPU and FP16 to GPU, comparing with a loose FP16-scale
21 // tolerance.
22 // - cross_attention_forward_with_attn: CUDA op is FP16-only. CPU runs
23 // FP32; parity test quantises through FP16 and uses loose tolerance.
24 // - cross_attention_forward_train / cross_attention_backward: FP32 on
25 // both backends — straightforward FP32<->FP32 parity.
26 //
27 // CONVENTIONS (verified against cross_attention.cu)
28 // * Weight layout: Wq/Wo are (D, D); Wk/Wv are (D, D_ctx). Projection is
29 // out(hh,i,j) = sum_k In(i,k) * W(hh*dh + j, k) — i.e. In @ W^T with the
30 // head slice taken as contiguous rows hh*dh..hh*dh+dh of W.
31 // * Per-head split: Qh/Kh/Vh are (H*L, dh), row-major by (head, token).
32 // Yconcat is (Lq, D) with head hh occupying columns hh*dh..hh*dh+dh.
33 // * Softmax scale: 1/sqrt(dh).
34 // * Mask: length-Lk key-validity buffer (1=valid, 0=invalid). Masked keys
35 // are excluded from the softmax denominator and forced to 0 probability.
36 // Query-side gating is enabled only when Lq == Lk (gate_query): an
37 // invalid query row produces a zero Attn row and a zero O row.
38 // cross_attention_forward_with_attn does NOT do query-side gating (it
39 // follows cxa_row_softmax_kernel which only masks keys).
40 // * Intermediates for the *_train pair:
41 // Qh (H*Lq, dh) per-head Q projection of X
42 // Kh (H*Lk, dh) per-head K projection of Ctx
43 // Vh (H*Lk, dh) per-head V projection of Ctx
44 // Attnh (H*Lq, Lk) per-head softmax probabilities
45 // Yconcat (Lq, D) per-head Attn@V written into head-strided columns
46 // O (Lq, D) Yconcat @ Wo^T (query-gated)
47 // * Backward accumulation: dWq/dWk/dWv/dWo ACCUMULATE (+=). dX/dCtx are
48 // OVERWRITTEN. Matches cx_dW_proj_kernel (+=) and cx_dX/cx_dCtx kernels.
49
50 #include <brotensor/tensor.h>
51 #include <brotensor/detail/cpu/thread_pool.h>
52
53 #include <cmath>
54 #include <vector>
55
56 namespace brotensor::detail::cpu {
57
58 namespace {
59
60 // out(hh,i,j) = sum_k In(i,k) * W(hh*dh + j, k). In: (L, Din). W: (D, Din).
61 // Out: (H*L, dh). Each hh owns rows [hh*L, hh*L+L) of Out exclusively (In/W
62 // read-only), so the head axis parallelizes with no cross-thread writes.
63 54 void cx_proj(const float* In, const float* W, float* Out,
64 int L, int Din, int H, int dh) {
65
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 265 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
66 211 const int hh = static_cast<int>(hhi);
67 211 const int row_off = hh * dh;
68
2/2
✓ Branch 0 taken 1840 times.
✓ Branch 1 taken 211 times.
 2051 for (int i = 0; i < L; ++i) {
69 1840 const float* xr = In + static_cast<std::size_t>(i) * Din;
70 1840 const std::size_t out_row =
71 1840 (static_cast<std::size_t>(hh) * L + i) * dh;
72
2/2
✓ Branch 0 taken 11785 times.
✓ Branch 1 taken 1840 times.
 13625 for (int j = 0; j < dh; ++j) {
73 11785 const float* wr = W + static_cast<std::size_t>(row_off + j) * Din;
74 11785 float acc = 0.0f;
75
2/2
✓ Branch 0 taken 279829 times.
✓ Branch 1 taken 11785 times.
 291614 for (int k = 0; k < Din; ++k) acc += xr[k] * wr[k];
76 11785 Out[out_row + j] = acc;
77 11785 }
78 1840 }
79 211 });
80 54 }
81
82 // Core FP32 forward shared by cross_attention_forward and
83 // cross_attention_forward_train.
84 13 void cross_attention_forward_core(const ::brotensor::Tensor& X,
85 const ::brotensor::Tensor& Ctx,
86 const ::brotensor::Tensor& Wq,
87 const ::brotensor::Tensor& Wk,
88 const ::brotensor::Tensor& Wv,
89 const ::brotensor::Tensor& Wo,
90 const float* d_mask,
91 int num_heads,
92 ::brotensor::Tensor& Qh,
93 ::brotensor::Tensor& Kh,
94 ::brotensor::Tensor& Vh,
95 ::brotensor::Tensor& Attnh,
96 ::brotensor::Tensor& Yconcat,
97 ::brotensor::Tensor& O) {
98 using ::brotensor::Dtype;
99 13 const int Lq = X.rows;
100 13 const int D = X.cols;
101 13 const int Lk = Ctx.rows;
102 13 const int Dctx = Ctx.cols;
103 13 const int H = num_heads;
104
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 const int dh = (H > 0) ? D / H : 0;
105
106
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 13 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 13 if (Qh.rows != H * Lq || Qh.cols != dh || Qh.dtype != Dtype::FP32)
107 13 Qh.resize(H * Lq, dh, Dtype::FP32);
108
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 13 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 13 if (Kh.rows != H * Lk || Kh.cols != dh || Kh.dtype != Dtype::FP32)
109 13 Kh.resize(H * Lk, dh, Dtype::FP32);
110
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 13 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 13 if (Vh.rows != H * Lk || Vh.cols != dh || Vh.dtype != Dtype::FP32)
111 13 Vh.resize(H * Lk, dh, Dtype::FP32);
112
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 13 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 13 if (Attnh.rows != H * Lq || Attnh.cols != Lk || Attnh.dtype != Dtype::FP32)
113 13 Attnh.resize(H * Lq, Lk, Dtype::FP32);
114
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 13 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 13 if (Yconcat.rows != Lq || Yconcat.cols != D || Yconcat.dtype != Dtype::FP32)
115 13 Yconcat.resize(Lq, D, Dtype::FP32);
116
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 13 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 13 if (O.rows != Lq || O.cols != D || O.dtype != Dtype::FP32)
117 13 O.resize(Lq, D, Dtype::FP32);
118
4/8
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 13 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 13 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 13 times.
 13 if (Lq == 0 || Lk == 0 || D == 0 || H == 0) return;
119
120 13 const int gate_query = (Lq == Lk) ? 1 : 0;
121 13 const float inv_sqrtdh = 1.0f / std::sqrt(static_cast<float>(dh));
122
123 13 const float* Xp = X.host_f32();
124 13 const float* Ctxp = Ctx.host_f32();
125 13 const float* Wqp = Wq.host_f32();
126 13 const float* Wkp = Wk.host_f32();
127 13 const float* Wvp = Wv.host_f32();
128 13 const float* Wop = Wo.host_f32();
129 13 float* Qp = Qh.host_f32_mut();
130 13 float* Kp = Kh.host_f32_mut();
131 13 float* Vp = Vh.host_f32_mut();
132 13 float* Ap = Attnh.host_f32_mut();
133 13 float* Yp = Yconcat.host_f32_mut();
134 13 float* Op = O.host_f32_mut();
135
136 // Q from X (D, D); K, V from Ctx (D, Dctx).
137 13 cx_proj(Xp, Wqp, Qp, Lq, D, H, dh);
138 13 cx_proj(Ctxp, Wkp, Kp, Lk, Dctx, H, dh);
139 13 cx_proj(Ctxp, Wvp, Vp, Lk, Dctx, H, dh);
140
141 // Scores -> masked row softmax -> Attnh. Each hh owns rows
142 // [hh*Lq, hh*Lq+Lq) of Attnh exclusively, so parallelizes over hh.
143
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 61 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
144 48 const int hh = static_cast<int>(hhi);
145
2/2
✓ Branch 0 taken 335 times.
✓ Branch 1 taken 48 times.
 383 for (int i = 0; i < Lq; ++i) {
146 335 float* arow = Ap + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
147
6/6
✓ Branch 0 taken 120 times.
✓ Branch 1 taken 215 times.
✓ Branch 2 taken 28 times.
✓ Branch 3 taken 92 times.
✓ Branch 4 taken 20 times.
✓ Branch 5 taken 8 times.
 335 if (gate_query && d_mask && d_mask[i] < 0.5f) {
148
2/2
✓ Branch 0 taken 64 times.
✓ Branch 1 taken 8 times.
 72 for (int j = 0; j < Lk; ++j) arow[j] = 0.0f;
149 8 continue;
150 }
151 327 const float* qr = Qp + (static_cast<std::size_t>(hh) * Lq + i) * dh;
152 327 float m = -1e30f;
153
2/2
✓ Branch 0 taken 2514 times.
✓ Branch 1 taken 327 times.
 2841 for (int j = 0; j < Lk; ++j) {
154
4/4
✓ Branch 0 taken 394 times.
✓ Branch 1 taken 2120 times.
✓ Branch 2 taken 215 times.
✓ Branch 3 taken 179 times.
 2514 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
155 2335 const float* kr =
156 2335 Kp + (static_cast<std::size_t>(hh) * Lk + j) * dh;
157 2335 float s = 0.0f;
158
2/2
✓ Branch 0 taken 17400 times.
✓ Branch 1 taken 2335 times.
 19735 for (int k = 0; k < dh; ++k) s += qr[k] * kr[k];
159 2335 s *= inv_sqrtdh;
160 2335 arow[j] = s;
161
2/2
✓ Branch 0 taken 1480 times.
✓ Branch 1 taken 855 times.
 2335 if (s > m) m = s;
162 2335 }
163 327 float sum = 0.0f;
164
2/2
✓ Branch 0 taken 3572 times.
✓ Branch 1 taken 327 times.
 3899 for (int j = 0; j < Lk; ++j) {
165
4/4
✓ Branch 0 taken 590 times.
✓ Branch 1 taken 2982 times.
✓ Branch 2 taken 420 times.
✓ Branch 3 taken 170 times.
 3572 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
166 3402 const float e = std::exp(arow[j] - m);
167 3402 arow[j] = e;
168 3402 sum += e;
169 3402 }
170
1/2
✓ Branch 0 taken 327 times.
✗ Branch 1 not taken.
 327 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
171
2/2
✓ Branch 0 taken 3423 times.
✓ Branch 1 taken 327 times.
 3750 for (int j = 0; j < Lk; ++j) arow[j] *= inv;
172 327 }
173 48 });
174
175 // Attn @ V -> Yconcat(i, hh*dh+k). Each hh writes only its own column
176 // range [hh*dh, hh*dh+dh) of every row of Yconcat, so parallelizes
177 // over hh.
178
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 57 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
179 44 const int hh = static_cast<int>(hhi);
180
2/2
✓ Branch 0 taken 333 times.
✓ Branch 1 taken 44 times.
 377 for (int i = 0; i < Lq; ++i) {
181 333 const float* arow =
182 333 Ap + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
183
2/2
✓ Branch 0 taken 1706 times.
✓ Branch 1 taken 333 times.
 2039 for (int k = 0; k < dh; ++k) {
184 1706 float acc = 0.0f;
185
2/2
✓ Branch 0 taken 15397 times.
✓ Branch 1 taken 1706 times.
 17103 for (int j = 0; j < Lk; ++j) {
186 15397 const float vv =
187 15397 Vp[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
188 15397 acc += arow[j] * vv;
189 15397 }
190 1706 Yp[static_cast<std::size_t>(i) * D + (hh * dh + k)] = acc;
191 1706 }
192 333 }
193 44 });
194
195 // Output projection O = Yconcat @ Wo^T, query-gated. Each i owns row i
196 // of O exclusively, so parallelizes over i.
197
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 89 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
198 76 const int i = static_cast<int>(ii);
199 76 float* orow = Op + static_cast<std::size_t>(i) * D;
200
6/6
✓ Branch 0 taken 31 times.
✓ Branch 1 taken 45 times.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 23 times.
✓ Branch 4 taken 6 times.
✓ Branch 5 taken 2 times.
 76 if (gate_query && d_mask && d_mask[i] < 0.5f) {
201
2/2
✓ Branch 0 taken 64 times.
✓ Branch 1 taken 2 times.
 66 for (int c = 0; c < D; ++c) orow[c] = 0.0f;
202 2 return;
203 }
204 74 const float* yr = Yp + static_cast<std::size_t>(i) * D;
205
2/2
✓ Branch 0 taken 74 times.
✓ Branch 1 taken 1686 times.
 1760 for (int c = 0; c < D; ++c) {
206 1686 const float* wr = Wop + static_cast<std::size_t>(c) * D;
207 1686 float acc = 0.0f;
208
2/2
✓ Branch 0 taken 50404 times.
✓ Branch 1 taken 1686 times.
 52090 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
209 1686 orow[c] = acc;
210 1686 }
211 76 });
212 13 }
213
214 } // namespace
215
216 // ─── cross_attention_forward (FP32 scalar; O only) ─────────────────────────
217
218 8 void cross_attention_forward(const ::brotensor::Tensor& X,
219 const ::brotensor::Tensor& Ctx,
220 const ::brotensor::Tensor& Wq,
221 const ::brotensor::Tensor& Wk,
222 const ::brotensor::Tensor& Wv,
223 const ::brotensor::Tensor& Wo,
224 const float* d_mask,
225 int num_heads,
226 ::brotensor::Tensor& O) {
227 using ::brotensor::Tensor;
228 using ::brotensor::Dtype;
229 8 Tensor Qh, Kh, Vh, Attnh, Yconcat;
230
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 16 cross_attention_forward_core(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads,
231 8 Qh, Kh, Vh, Attnh, Yconcat, O);
232 8 }
233
234 // ─── cross_attention_forward_train (FP32; emits per-head caches) ────────────
235
236 5 void cross_attention_forward_train(const ::brotensor::Tensor& X,
237 const ::brotensor::Tensor& Ctx,
238 const ::brotensor::Tensor& Wq,
239 const ::brotensor::Tensor& Wk,
240 const ::brotensor::Tensor& Wv,
241 const ::brotensor::Tensor& Wo,
242 const float* d_mask,
243 int num_heads,
244 ::brotensor::Tensor& Qh,
245 ::brotensor::Tensor& Kh,
246 ::brotensor::Tensor& Vh,
247 ::brotensor::Tensor& Attnh,
248 ::brotensor::Tensor& Yconcat,
249 ::brotensor::Tensor& O) {
250 10 cross_attention_forward_core(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads,
251 5 Qh, Kh, Vh, Attnh, Yconcat, O);
252 5 }
253
254 // ─── cross_attention_forward_with_attn (FP32; + head-averaged AttnAvg) ─────
255
256 5 void cross_attention_forward_with_attn(const ::brotensor::Tensor& X,
257 const ::brotensor::Tensor& Ctx,
258 const ::brotensor::Tensor& Wq,
259 const ::brotensor::Tensor& Wk,
260 const ::brotensor::Tensor& Wv,
261 const ::brotensor::Tensor& Wo,
262 const float* d_mask,
263 const ::brotensor::Tensor* attn_logit_bias,
264 int num_heads,
265 ::brotensor::Tensor& O,
266 ::brotensor::Tensor& AttnAvg) {
267 using ::brotensor::Dtype;
268 5 const int Lq = X.rows;
269 5 const int D = X.cols;
270 5 const int Lk = Ctx.rows;
271 5 const int Dctx = Ctx.cols;
272 5 const int H = num_heads;
273
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 const int dh = (H > 0) ? D / H : 0;
274
275
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 5 if (O.rows != Lq || O.cols != D || O.dtype != Dtype::FP32)
276 5 O.resize(Lq, D, Dtype::FP32);
277
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 5 if (AttnAvg.rows != Lq || AttnAvg.cols != Lk || AttnAvg.dtype != Dtype::FP32)
278 5 AttnAvg.resize(Lq, Lk, Dtype::FP32);
279
4/8
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 5 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 5 times.
 5 if (Lq == 0 || Lk == 0 || D == 0 || H == 0) return;
280
281 5 const float inv_sqrtdh = 1.0f / std::sqrt(static_cast<float>(dh));
282
283 5 const float* Xp = X.host_f32();
284 5 const float* Ctxp = Ctx.host_f32();
285 5 const float* Wqp = Wq.host_f32();
286 5 const float* Wkp = Wk.host_f32();
287 5 const float* Wvp = Wv.host_f32();
288 5 const float* Wop = Wo.host_f32();
289
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 3 times.
 5 const float* bias = attn_logit_bias ? attn_logit_bias->host_f32() : nullptr;
290 5 float* Op = O.host_f32_mut();
291 5 float* AAp = AttnAvg.host_f32_mut();
292
293 5 std::vector<float> Q(static_cast<std::size_t>(H) * Lq * dh, 0.0f);
294
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> Kbuf(static_cast<std::size_t>(H) * Lk * dh, 0.0f);
295
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> Vbuf(static_cast<std::size_t>(H) * Lk * dh, 0.0f);
296
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> Attn(static_cast<std::size_t>(H) * Lq * Lk, 0.0f);
297
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> Yc(static_cast<std::size_t>(Lq) * D, 0.0f);
298
299
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 cx_proj(Xp, Wqp, Q.data(), Lq, D, H, dh);
300
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 cx_proj(Ctxp, Wkp, Kbuf.data(), Lk, Dctx, H, dh);
301
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 cx_proj(Ctxp, Wvp, Vbuf.data(), Lk, Dctx, H, dh);
302
303 // Scores (+ optional bias) -> key-masked row softmax. No query gating —
304 // matches cxa_row_softmax_kernel which only masks keys. Each hh owns
305 // rows [hh*Lq, hh*Lq+Lq) of Attn exclusively, so parallelizes over hh.
306
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 25 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
307 20 const int hh = static_cast<int>(hhi);
308
2/2
✓ Branch 0 taken 127 times.
✓ Branch 1 taken 20 times.
 147 for (int i = 0; i < Lq; ++i) {
309 127 float* arow = Attn.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
310 127 const float* qr = Q.data() + (static_cast<std::size_t>(hh) * Lq + i) * dh;
311 127 float m = -1e30f;
312
2/2
✓ Branch 0 taken 840 times.
✓ Branch 1 taken 127 times.
 967 for (int j = 0; j < Lk; ++j) {
313 840 const float* kr =
314 840 Kbuf.data() + (static_cast<std::size_t>(hh) * Lk + j) * dh;
315 840 float s = 0.0f;
316
2/2
✓ Branch 0 taken 5368 times.
✓ Branch 1 taken 840 times.
 6208 for (int k = 0; k < dh; ++k) s += qr[k] * kr[k];
317 840 s *= inv_sqrtdh;
318
2/2
✓ Branch 0 taken 302 times.
✓ Branch 1 taken 538 times.
 840 if (bias) s += bias[static_cast<std::size_t>(i) * Lk + j];
319 840 arow[j] = s;
320
4/4
✓ Branch 0 taken 217 times.
✓ Branch 1 taken 623 times.
✓ Branch 2 taken 111 times.
✓ Branch 3 taken 106 times.
 840 if (d_mask && d_mask[j] < 0.5f) continue;
321
2/2
✓ Branch 0 taken 427 times.
✓ Branch 1 taken 307 times.
 734 if (s > m) m = s;
322 734 }
323 127 float sum = 0.0f;
324
2/2
✓ Branch 0 taken 1068 times.
✓ Branch 1 taken 127 times.
 1195 for (int j = 0; j < Lk; ++j) {
325
4/4
✓ Branch 0 taken 373 times.
✓ Branch 1 taken 695 times.
✓ Branch 2 taken 258 times.
✓ Branch 3 taken 115 times.
 1068 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
326 953 const float e = std::exp(arow[j] - m);
327 953 arow[j] = e;
328 953 sum += e;
329 953 }
330
1/2
✓ Branch 0 taken 127 times.
✗ Branch 1 not taken.
 127 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
331
2/2
✓ Branch 0 taken 1081 times.
✓ Branch 1 taken 127 times.
 1208 for (int j = 0; j < Lk; ++j) arow[j] *= inv;
332 127 }
333 20 });
334
335 // AttnAvg(i, j) = (1/H) sum_h Attn(h, i, j). Each (i, j) writes exactly
336 // once from a private serial sum over hh, so parallelizes over i (each i
337 // owns row i of AttnAvg exclusively).
338
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 32 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
339 27 const int i = static_cast<int>(ii);
340
2/2
✓ Branch 0 taken 167 times.
✓ Branch 1 taken 27 times.
 194 for (int j = 0; j < Lk; ++j) {
341 167 float acc = 0.0f;
342
2/2
✓ Branch 0 taken 700 times.
✓ Branch 1 taken 167 times.
 867 for (int hh = 0; hh < H; ++hh)
343 700 acc += Attn[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
344 167 AAp[static_cast<std::size_t>(i) * Lk + j] =
345 167 acc / static_cast<float>(H);
346 167 }
347 27 });
348
349 // Attn @ V -> Yconcat(i, hh*dh+k). Each hh writes only its own column
350 // range of every row of Yc, so parallelizes over hh.
351
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 27 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
352 22 const int hh = static_cast<int>(hhi);
353
2/2
✓ Branch 0 taken 123 times.
✓ Branch 1 taken 22 times.
 145 for (int i = 0; i < Lq; ++i) {
354 123 const float* arow =
355 123 Attn.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
356
2/2
✓ Branch 0 taken 644 times.
✓ Branch 1 taken 123 times.
 767 for (int k = 0; k < dh; ++k) {
357 644 float acc = 0.0f;
358
2/2
✓ Branch 0 taken 4906 times.
✓ Branch 1 taken 644 times.
 5550 for (int j = 0; j < Lk; ++j) {
359 4906 const float vv =
360 4906 Vbuf[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
361 4906 acc += arow[j] * vv;
362 4906 }
363 644 Yc[static_cast<std::size_t>(i) * D + (hh * dh + k)] = acc;
364 644 }
365 123 }
366 22 });
367
368 // Output projection O = Yconcat @ Wo^T (no query gating). Each i owns
369 // row i of O exclusively, so parallelizes over i.
370
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 33 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
371 28 const int i = static_cast<int>(ii);
372 28 const float* yr = Yc.data() + static_cast<std::size_t>(i) * D;
373 28 float* orow = Op + static_cast<std::size_t>(i) * D;
374
2/2
✓ Branch 0 taken 697 times.
✓ Branch 1 taken 28 times.
 725 for (int c = 0; c < D; ++c) {
375 697 const float* wr = Wop + static_cast<std::size_t>(c) * D;
376 697 float acc = 0.0f;
377
2/2
✓ Branch 0 taken 19150 times.
✓ Branch 1 taken 697 times.
 19847 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
378 697 orow[c] = acc;
379 697 }
380 28 });
381 5 }
382
383 // ─── cross_attention_backward (FP32; accumulates dW*, overwrites dX/dCtx) ──
384
385 5 void cross_attention_backward(const ::brotensor::Tensor& dO,
386 const ::brotensor::Tensor& X,
387 const ::brotensor::Tensor& Ctx,
388 const ::brotensor::Tensor& Qh,
389 const ::brotensor::Tensor& Kh,
390 const ::brotensor::Tensor& Vh,
391 const ::brotensor::Tensor& Attnh,
392 const ::brotensor::Tensor& Yconcat,
393 const ::brotensor::Tensor& Wq,
394 const ::brotensor::Tensor& Wk,
395 const ::brotensor::Tensor& Wv,
396 const ::brotensor::Tensor& Wo,
397 const float* d_mask,
398 int num_heads,
399 ::brotensor::Tensor& dX,
400 ::brotensor::Tensor& dCtx,
401 ::brotensor::Tensor& dWq,
402 ::brotensor::Tensor& dWk,
403 ::brotensor::Tensor& dWv,
404 ::brotensor::Tensor& dWo) {
405 using ::brotensor::Dtype;
406 5 const int Lq = X.rows;
407 5 const int D = X.cols;
408 5 const int Lk = Ctx.rows;
409 5 const int Dctx = Ctx.cols;
410 5 const int H = num_heads;
411
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 const int dh = (H > 0) ? D / H : 0;
412
413
3/6
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 5 times.
 5 if (dX.rows != Lq || dX.cols != D || dX.dtype != Dtype::FP32)
414 ✗ dX.resize(Lq, D, Dtype::FP32);
415
3/6
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 5 times.
 5 if (dCtx.rows != Lk || dCtx.cols != Dctx || dCtx.dtype != Dtype::FP32)
416 ✗ dCtx.resize(Lk, Dctx, Dtype::FP32);
417
4/8
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 5 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 5 times.
 5 if (Lq == 0 || Lk == 0 || D == 0 || H == 0) return;
418
419 5 const int gate_query = (Lq == Lk) ? 1 : 0;
420 5 const float inv_sqrtdh = 1.0f / std::sqrt(static_cast<float>(dh));
421
422 5 const float* dOp = dO.host_f32();
423 5 const float* Xp = X.host_f32();
424 5 const float* Cp = Ctx.host_f32();
425 5 const float* Qp = Qh.host_f32();
426 5 const float* Kp = Kh.host_f32();
427 5 const float* Vp = Vh.host_f32();
428 5 const float* Ap = Attnh.host_f32();
429 5 const float* Yp = Yconcat.host_f32();
430 5 const float* Wqp = Wq.host_f32();
431 5 const float* Wkp = Wk.host_f32();
432 5 const float* Wvp = Wv.host_f32();
433 5 const float* Wop = Wo.host_f32();
434 5 float* dXp = dX.host_f32_mut();
435 5 float* dCp = dCtx.host_f32_mut();
436 5 float* dWqp = dWq.host_f32_mut();
437 5 float* dWkp = dWk.host_f32_mut();
438 5 float* dWvp = dWv.host_f32_mut();
439 5 float* dWop = dWo.host_f32_mut();
440
441 5 std::vector<float> dYc(static_cast<std::size_t>(Lq) * D, 0.0f);
442
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> dAttn(static_cast<std::size_t>(H) * Lq * Lk, 0.0f);
443
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> dVh(static_cast<std::size_t>(H) * Lk * dh, 0.0f);
444
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> dScores(static_cast<std::size_t>(H) * Lq * Lk, 0.0f);
445
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> dQh(static_cast<std::size_t>(H) * Lq * dh, 0.0f);
446
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 std::vector<float> dKh(static_cast<std::size_t>(H) * Lk * dh, 0.0f);
447
448 // dYconcat = dO @ Wo (query-gated; overwrite). Each i owns row i of dYc
449 // exclusively, so parallelizes over i.
450
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 44 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
451 39 const int i = static_cast<int>(ii);
452
4/4
✓ Branch 0 taken 15 times.
✓ Branch 1 taken 24 times.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 7 times.
 39 const bool valid = !(gate_query && d_mask && d_mask[i] < 0.5f);
453
2/2
✓ Branch 0 taken 39 times.
✓ Branch 1 taken 1488 times.
 1527 for (int k = 0; k < D; ++k) {
454 1488 float acc = 0.0f;
455
2/2
✓ Branch 0 taken 1458 times.
✓ Branch 1 taken 30 times.
 1488 if (valid) {
456
2/2
✓ Branch 0 taken 42886 times.
✓ Branch 1 taken 1458 times.
 44344 for (int c = 0; c < D; ++c)
457 85772 acc += Wop[static_cast<std::size_t>(c) * D + k] *
458 42886 dOp[static_cast<std::size_t>(i) * D + c];
459 1458 }
460 1488 dYc[static_cast<std::size_t>(i) * D + k] = acc;
461 1488 }
462 39 });
463 // dWo accumulates. c-outer with a private serial sum over i — each
464 // (c, k) cell is written exactly once, so this parallelizes over c
465 // (unlike ops_impl.cpp's mha_backward/attention_backward dWo, which use
466 // a token-outermost accumulate-in-place order and are left
467 // single-threaded there).
468
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 219 parallel_for(static_cast<std::size_t>(D), [&](std::size_t ci) {
469 214 const int c = static_cast<int>(ci);
470
2/2
✓ Branch 0 taken 6036 times.
✓ Branch 1 taken 214 times.
 6250 for (int k = 0; k < D; ++k) {
471 6036 float acc = 0.0f;
472
2/2
✓ Branch 0 taken 41627 times.
✓ Branch 1 taken 6036 times.
 47663 for (int i = 0; i < Lq; ++i) {
473
6/6
✓ Branch 0 taken 5744 times.
✓ Branch 1 taken 35883 times.
✓ Branch 2 taken 2899 times.
✓ Branch 3 taken 2845 times.
✓ Branch 4 taken 1820 times.
✓ Branch 5 taken 1079 times.
 41627 if (gate_query && d_mask && d_mask[i] < 0.5f) continue;
474 81096 acc += dOp[static_cast<std::size_t>(i) * D + c] *
475 40548 Yp[static_cast<std::size_t>(i) * D + k];
476 40548 }
477 6036 dWop[static_cast<std::size_t>(c) * D + k] += acc;
478 6036 }
479 214 });
480
481 // Per-head dAttn and dVh. Each hh owns rows [hh*Lq, hh*Lq+Lq) of dAttn
482 // and [hh*Lk, hh*Lk+Lk) of dVh exclusively (dVh's inner reduction over i
483 // sums only into this head's own private slice), so parallelizes over hh.
484
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 32 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
485 27 const int hh = static_cast<int>(hhi);
486
2/2
✓ Branch 0 taken 223 times.
✓ Branch 1 taken 27 times.
 250 for (int i = 0; i < Lq; ++i) {
487
2/2
✓ Branch 0 taken 1803 times.
✓ Branch 1 taken 223 times.
 2026 for (int j = 0; j < Lk; ++j) {
488 1803 float acc = 0.0f;
489
2/2
✓ Branch 0 taken 12122 times.
✓ Branch 1 taken 1803 times.
 13925 for (int k = 0; k < dh; ++k) {
490 12122 const float dy =
491 12122 dYc[static_cast<std::size_t>(i) * D + (hh * dh + k)];
492 12122 const float vv =
493 12122 Vp[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
494 12122 acc += dy * vv;
495 12122 }
496 1803 dAttn[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j] = acc;
497 1803 }
498 223 }
499
2/2
✓ Branch 0 taken 340 times.
✓ Branch 1 taken 27 times.
 367 for (int j = 0; j < Lk; ++j) {
500
2/2
✓ Branch 0 taken 1753 times.
✓ Branch 1 taken 340 times.
 2093 for (int k = 0; k < dh; ++k) {
501 1753 float acc = 0.0f;
502
2/2
✓ Branch 0 taken 12158 times.
✓ Branch 1 taken 1753 times.
 13911 for (int i = 0; i < Lq; ++i) {
503 12158 const float a =
504 12158 Ap[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
505 12158 const float dy =
506 12158 dYc[static_cast<std::size_t>(i) * D + (hh * dh + k)];
507 12158 acc += a * dy;
508 12158 }
509 1753 dVh[(static_cast<std::size_t>(hh) * Lk + j) * dh + k] = acc;
510 1753 }
511 340 }
512 27 });
513
514 // Per-head row-softmax backward -> dScores. Each hh owns rows
515 // [hh*Lq, hh*Lq+Lq) of dScores exclusively, so parallelizes over hh.
516
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 32 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
517 27 const int hh = static_cast<int>(hhi);
518
2/2
✓ Branch 0 taken 171 times.
✓ Branch 1 taken 27 times.
 198 for (int i = 0; i < Lq; ++i) {
519 171 const float* prow =
520 171 Ap + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
521 171 const float* dprow =
522 171 dAttn.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
523 171 float* drow =
524 171 dScores.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
525
6/6
✓ Branch 0 taken 39 times.
✓ Branch 1 taken 132 times.
✓ Branch 2 taken 18 times.
✓ Branch 3 taken 21 times.
✓ Branch 4 taken 10 times.
✓ Branch 5 taken 8 times.
 171 if (gate_query && d_mask && d_mask[i] < 0.5f) {
526
2/2
✓ Branch 0 taken 64 times.
✓ Branch 1 taken 8 times.
 72 for (int j = 0; j < Lk; ++j) drow[j] = 0.0f;
527 8 continue;
528 }
529 163 float dot = 0.0f;
530
2/2
✓ Branch 0 taken 1718 times.
✓ Branch 1 taken 163 times.
 1881 for (int j = 0; j < Lk; ++j) dot += dprow[j] * prow[j];
531
2/2
✓ Branch 0 taken 1749 times.
✓ Branch 1 taken 163 times.
 1912 for (int j = 0; j < Lk; ++j) {
532
4/4
✓ Branch 0 taken 256 times.
✓ Branch 1 taken 1493 times.
✓ Branch 2 taken 178 times.
✓ Branch 3 taken 78 times.
 1749 if (d_mask && d_mask[j] < 0.5f) drow[j] = 0.0f;
533 1671 else drow[j] = prow[j] * (dprow[j] - dot) * inv_sqrtdh;
534 1749 }
535 163 }
536 27 });
537
538 // Per-head dQh, dKh. Each hh owns rows [hh*Lq, hh*Lq+Lq) of dQh and
539 // [hh*Lk, hh*Lk+Lk) of dKh exclusively, so parallelizes over hh.
540
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 32 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
541 27 const int hh = static_cast<int>(hhi);
542
2/2
✓ Branch 0 taken 223 times.
✓ Branch 1 taken 27 times.
 250 for (int i = 0; i < Lq; ++i) {
543
2/2
✓ Branch 0 taken 1240 times.
✓ Branch 1 taken 223 times.
 1463 for (int k = 0; k < dh; ++k) {
544 1240 float acc = 0.0f;
545
2/2
✓ Branch 0 taken 12049 times.
✓ Branch 1 taken 1240 times.
 13289 for (int j = 0; j < Lk; ++j) {
546 12049 const float ds =
547 12049 dScores[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
548 12049 const float kk =
549 12049 Kp[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
550 12049 acc += ds * kk;
551 12049 }
552 1240 dQh[(static_cast<std::size_t>(hh) * Lq + i) * dh + k] = acc;
553 1240 }
554 223 }
555
2/2
✓ Branch 0 taken 329 times.
✓ Branch 1 taken 27 times.
 356 for (int j = 0; j < Lk; ++j) {
556
2/2
✓ Branch 0 taken 1738 times.
✓ Branch 1 taken 329 times.
 2067 for (int k = 0; k < dh; ++k) {
557 1738 float acc = 0.0f;
558
2/2
✓ Branch 0 taken 12098 times.
✓ Branch 1 taken 1738 times.
 13836 for (int i = 0; i < Lq; ++i) {
559 12098 const float ds =
560 12098 dScores[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
561 12098 const float qq =
562 12098 Qp[(static_cast<std::size_t>(hh) * Lq + i) * dh + k];
563 12098 acc += ds * qq;
564 12098 }
565 1738 dKh[(static_cast<std::size_t>(hh) * Lk + j) * dh + k] = acc;
566 1738 }
567 329 }
568 27 });
569
570 // dWq (D, D) accumulates against X.
571 // dWq(hh*dh+j, k) += sum_i dQh(hh,i,j) * X(i,k).
572 // wrow-outer with a private serial sum over i — each (wrow, k) cell is
573 // written exactly once, so this parallelizes over wrow.
574
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 222 parallel_for(static_cast<std::size_t>(D), [&](std::size_t wrowi) {
575 217 const int wrow = static_cast<int>(wrowi);
576 217 const int hh = wrow / dh;
577 217 const int j = wrow % dh;
578
2/2
✓ Branch 0 taken 5980 times.
✓ Branch 1 taken 217 times.
 6197 for (int k = 0; k < D; ++k) {
579 5980 float acc = 0.0f;
580
2/2
✓ Branch 0 taken 41489 times.
✓ Branch 1 taken 5980 times.
 47469 for (int i = 0; i < Lq; ++i) {
581 41489 const float xv = Xp[static_cast<std::size_t>(i) * D + k];
582 41489 acc += dQh[(static_cast<std::size_t>(hh) * Lq + i) * dh + j] * xv;
583 41489 }
584 5980 dWqp[static_cast<std::size_t>(wrow) * D + k] += acc;
585 5980 }
586 217 });
587 // dWk, dWv (D, Dctx) accumulate against Ctx. Same wrow-outer
588 // private-sum-inner pattern, so parallelizes over wrow.
589
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 223 parallel_for(static_cast<std::size_t>(D), [&](std::size_t wrowi) {
590 218 const int wrow = static_cast<int>(wrowi);
591 218 const int hh = wrow / dh;
592 218 const int j = wrow % dh;
593
2/2
✓ Branch 0 taken 5455 times.
✓ Branch 1 taken 218 times.
 5673 for (int k = 0; k < Dctx; ++k) {
594 5455 float ak = 0.0f, av = 0.0f;
595
2/2
✓ Branch 0 taken 42348 times.
✓ Branch 1 taken 5455 times.
 47803 for (int i = 0; i < Lk; ++i) {
596 42348 const float cv = Cp[static_cast<std::size_t>(i) * Dctx + k];
597 42348 ak += dKh[(static_cast<std::size_t>(hh) * Lk + i) * dh + j] * cv;
598 42348 av += dVh[(static_cast<std::size_t>(hh) * Lk + i) * dh + j] * cv;
599 42348 }
600 5455 const std::size_t idx = static_cast<std::size_t>(wrow) * Dctx + k;
601 5455 dWkp[idx] += ak;
602 5455 dWvp[idx] += av;
603 5455 }
604 218 });
605
606 // dX(i,k) = sum over heads, j: dQh(hh,i,j) * Wq(hh*dh+j, k). Overwrite.
607 // Each i owns row i of dX exclusively, so parallelizes over i.
608
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 42 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
609 37 const int i = static_cast<int>(ii);
610
2/2
✓ Branch 0 taken 1640 times.
✓ Branch 1 taken 37 times.
 1677 for (int k = 0; k < D; ++k) {
611 1640 float acc = 0.0f;
612
2/2
✓ Branch 0 taken 6124 times.
✓ Branch 1 taken 1640 times.
 7764 for (int hh = 0; hh < H; ++hh) {
613
2/2
✓ Branch 0 taken 41422 times.
✓ Branch 1 taken 6124 times.
 47546 for (int j = 0; j < dh; ++j) {
614 41422 const int wrow = hh * dh + j;
615 41422 const float gq =
616 41422 dQh[(static_cast<std::size_t>(hh) * Lq + i) * dh + j];
617 41422 acc += gq * Wqp[static_cast<std::size_t>(wrow) * D + k];
618 41422 }
619 6124 }
620 1640 dXp[static_cast<std::size_t>(i) * D + k] = acc;
621 1640 }
622 37 });
623
624 // dCtx(j,k) = sum over heads, m: dKh*Wk + dVh*Wv at (hh*dh+m, k).
625 // Overwrite. Each j owns row j of dCtx exclusively, so parallelizes
626 // over j.
627
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 60 parallel_for(static_cast<std::size_t>(Lk), [&](std::size_t ji) {
628 55 const int j = static_cast<int>(ji);
629
2/2
✓ Branch 0 taken 1828 times.
✓ Branch 1 taken 55 times.
 1883 for (int k = 0; k < Dctx; ++k) {
630 1828 float acc = 0.0f;
631
2/2
✓ Branch 0 taken 7845 times.
✓ Branch 1 taken 1828 times.
 9673 for (int hh = 0; hh < H; ++hh) {
632
2/2
✓ Branch 0 taken 44103 times.
✓ Branch 1 taken 7845 times.
 51948 for (int m = 0; m < dh; ++m) {
633 44103 const int wrow = hh * dh + m;
634 44103 const std::size_t widx =
635 44103 static_cast<std::size_t>(wrow) * Dctx + k;
636 44103 const float gk =
637 44103 dKh[(static_cast<std::size_t>(hh) * Lk + j) * dh + m];
638 44103 const float gv =
639 44103 dVh[(static_cast<std::size_t>(hh) * Lk + j) * dh + m];
640 44103 acc += gk * Wkp[widx] + gv * Wvp[widx];
641 44103 }
642 7845 }
643 1828 dCp[static_cast<std::size_t>(j) * Dctx + k] = acc;
644 1828 }
645 55 });
646 5 }
647
648 } // namespace brotensor::detail::cpu
649