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.
 252 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
66 198 const int hh = static_cast<int>(hhi);
67 198 const int row_off = hh * dh;
68
2/2
✓ Branch 0 taken 1883 times.
✓ Branch 1 taken 198 times.
 2081 for (int i = 0; i < L; ++i) {
69 1883 const float* xr = In + static_cast<std::size_t>(i) * Din;
70 1883 const std::size_t out_row =
71 1883 (static_cast<std::size_t>(hh) * L + i) * dh;
72
2/2
✓ Branch 0 taken 12101 times.
✓ Branch 1 taken 1883 times.
 13984 for (int j = 0; j < dh; ++j) {
73 12101 const float* wr = W + static_cast<std::size_t>(row_off + j) * Din;
74 12101 float acc = 0.0f;
75
2/2
✓ Branch 0 taken 351872 times.
✓ Branch 1 taken 12101 times.
 363973 for (int k = 0; k < Din; ++k) acc += xr[k] * wr[k];
76 12101 Out[out_row + j] = acc;
77 12101 }
78 1883 }
79 198 });
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.
 56 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
144 43 const int hh = static_cast<int>(hhi);
145
2/2
✓ Branch 0 taken 345 times.
✓ Branch 1 taken 43 times.
 388 for (int i = 0; i < Lq; ++i) {
146 345 float* arow = Ap + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
147
6/6
✓ Branch 0 taken 119 times.
✓ Branch 1 taken 226 times.
✓ Branch 2 taken 30 times.
✓ Branch 3 taken 89 times.
✓ Branch 4 taken 22 times.
✓ Branch 5 taken 8 times.
 345 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 337 const float* qr = Qp + (static_cast<std::size_t>(hh) * Lq + i) * dh;
152 337 float m = -1e30f;
153
2/2
✓ Branch 0 taken 3267 times.
✓ Branch 1 taken 337 times.
 3604 for (int j = 0; j < Lk; ++j) {
154
4/4
✓ Branch 0 taken 457 times.
✓ Branch 1 taken 2810 times.
✓ Branch 2 taken 282 times.
✓ Branch 3 taken 175 times.
 3267 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
155 3092 const float* kr =
156 3092 Kp + (static_cast<std::size_t>(hh) * Lk + j) * dh;
157 3092 float s = 0.0f;
158
2/2
✓ Branch 0 taken 14927 times.
✓ Branch 1 taken 3092 times.
 18019 for (int k = 0; k < dh; ++k) s += qr[k] * kr[k];
159 3092 s *= inv_sqrtdh;
160 3092 arow[j] = s;
161
2/2
✓ Branch 0 taken 2110 times.
✓ Branch 1 taken 982 times.
 3092 if (s > m) m = s;
162 3092 }
163 337 float sum = 0.0f;
164
2/2
✓ Branch 0 taken 3385 times.
✓ Branch 1 taken 337 times.
 3722 for (int j = 0; j < Lk; ++j) {
165
4/4
✓ Branch 0 taken 530 times.
✓ Branch 1 taken 2855 times.
✓ Branch 2 taken 365 times.
✓ Branch 3 taken 165 times.
 3385 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
166 3220 const float e = std::exp(arow[j] - m);
167 3220 arow[j] = e;
168 3220 sum += e;
169 3220 }
170
1/2
✓ Branch 0 taken 337 times.
✗ Branch 1 not taken.
 337 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
171
2/2
✓ Branch 0 taken 3848 times.
✓ Branch 1 taken 337 times.
 4185 for (int j = 0; j < Lk; ++j) arow[j] *= inv;
172 337 }
173 43 });
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.
 56 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
179 43 const int hh = static_cast<int>(hhi);
180
2/2
✓ Branch 0 taken 361 times.
✓ Branch 1 taken 43 times.
 404 for (int i = 0; i < Lq; ++i) {
181 361 const float* arow =
182 361 Ap + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
183
2/2
✓ Branch 0 taken 2184 times.
✓ Branch 1 taken 361 times.
 2545 for (int k = 0; k < dh; ++k) {
184 2184 float acc = 0.0f;
185
2/2
✓ Branch 0 taken 21930 times.
✓ Branch 1 taken 2184 times.
 24114 for (int j = 0; j < Lk; ++j) {
186 21930 const float vv =
187 21930 Vp[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
188 21930 acc += arow[j] * vv;
189 21930 }
190 2184 Yp[static_cast<std::size_t>(i) * D + (hh * dh + k)] = acc;
191 2184 }
192 361 }
193 43 });
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.
 93 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
198 80 const int i = static_cast<int>(ii);
199 80 float* orow = Op + static_cast<std::size_t>(i) * D;
200
6/6
✓ Branch 0 taken 29 times.
✓ Branch 1 taken 51 times.
✓ Branch 2 taken 7 times.
✓ Branch 3 taken 22 times.
✓ Branch 4 taken 5 times.
✓ Branch 5 taken 2 times.
 80 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 78 const float* yr = Yp + static_cast<std::size_t>(i) * D;
205
2/2
✓ Branch 0 taken 78 times.
✓ Branch 1 taken 2518 times.
 2596 for (int c = 0; c < D; ++c) {
206 2518 const float* wr = Wop + static_cast<std::size_t>(c) * D;
207 2518 float acc = 0.0f;
208
2/2
✓ Branch 0 taken 63604 times.
✓ Branch 1 taken 2518 times.
 66122 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
209 2518 orow[c] = acc;
210 2518 }
211 80 });
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.
 23 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
307 18 const int hh = static_cast<int>(hhi);
308
2/2
✓ Branch 0 taken 134 times.
✓ Branch 1 taken 18 times.
 152 for (int i = 0; i < Lq; ++i) {
309 134 float* arow = Attn.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
310 134 const float* qr = Q.data() + (static_cast<std::size_t>(hh) * Lq + i) * dh;
311 134 float m = -1e30f;
312
2/2
✓ Branch 0 taken 1016 times.
✓ Branch 1 taken 134 times.
 1150 for (int j = 0; j < Lk; ++j) {
313 1016 const float* kr =
314 1016 Kbuf.data() + (static_cast<std::size_t>(hh) * Lk + j) * dh;
315 1016 float s = 0.0f;
316
2/2
✓ Branch 0 taken 6538 times.
✓ Branch 1 taken 1016 times.
 7554 for (int k = 0; k < dh; ++k) s += qr[k] * kr[k];
317 1016 s *= inv_sqrtdh;
318
2/2
✓ Branch 0 taken 588 times.
✓ Branch 1 taken 428 times.
 1016 if (bias) s += bias[static_cast<std::size_t>(i) * Lk + j];
319 1016 arow[j] = s;
320
4/4
✓ Branch 0 taken 283 times.
✓ Branch 1 taken 733 times.
✓ Branch 2 taken 181 times.
✓ Branch 3 taken 102 times.
 1016 if (d_mask && d_mask[j] < 0.5f) continue;
321
2/2
✓ Branch 0 taken 572 times.
✓ Branch 1 taken 342 times.
 914 if (s > m) m = s;
322 914 }
323 134 float sum = 0.0f;
324
2/2
✓ Branch 0 taken 1107 times.
✓ Branch 1 taken 134 times.
 1241 for (int j = 0; j < Lk; ++j) {
325
4/4
✓ Branch 0 taken 331 times.
✓ Branch 1 taken 776 times.
✓ Branch 2 taken 218 times.
✓ Branch 3 taken 113 times.
 1107 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
326 994 const float e = std::exp(arow[j] - m);
327 994 arow[j] = e;
328 994 sum += e;
329 994 }
330
1/2
✓ Branch 0 taken 134 times.
✗ Branch 1 not taken.
 134 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
331
2/2
✓ Branch 0 taken 1106 times.
✓ Branch 1 taken 134 times.
 1240 for (int j = 0; j < Lk; ++j) arow[j] *= inv;
332 134 }
333 18 });
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.
 29 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
339 24 const int i = static_cast<int>(ii);
340
2/2
✓ Branch 0 taken 183 times.
✓ Branch 1 taken 24 times.
 207 for (int j = 0; j < Lk; ++j) {
341 183 float acc = 0.0f;
342
2/2
✓ Branch 0 taken 638 times.
✓ Branch 1 taken 183 times.
 821 for (int hh = 0; hh < H; ++hh)
343 638 acc += Attn[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
344 183 AAp[static_cast<std::size_t>(i) * Lk + j] =
345 183 acc / static_cast<float>(H);
346 183 }
347 24 });
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.
 23 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
352 18 const int hh = static_cast<int>(hhi);
353
2/2
✓ Branch 0 taken 134 times.
✓ Branch 1 taken 18 times.
 152 for (int i = 0; i < Lq; ++i) {
354 134 const float* arow =
355 134 Attn.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
356
2/2
✓ Branch 0 taken 772 times.
✓ Branch 1 taken 134 times.
 906 for (int k = 0; k < dh; ++k) {
357 772 float acc = 0.0f;
358
2/2
✓ Branch 0 taken 4140 times.
✓ Branch 1 taken 772 times.
 4912 for (int j = 0; j < Lk; ++j) {
359 4140 const float vv =
360 4140 Vbuf[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
361 4140 acc += arow[j] * vv;
362 4140 }
363 772 Yc[static_cast<std::size_t>(i) * D + (hh * dh + k)] = acc;
364 772 }
365 134 }
366 18 });
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.
 31 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
371 26 const int i = static_cast<int>(ii);
372 26 const float* yr = Yc.data() + static_cast<std::size_t>(i) * D;
373 26 float* orow = Op + static_cast<std::size_t>(i) * D;
374
2/2
✓ Branch 0 taken 732 times.
✓ Branch 1 taken 26 times.
 758 for (int c = 0; c < D; ++c) {
375 732 const float* wr = Wop + static_cast<std::size_t>(c) * D;
376 732 float acc = 0.0f;
377
2/2
✓ Branch 0 taken 16210 times.
✓ Branch 1 taken 732 times.
 16942 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
378 732 orow[c] = acc;
379 732 }
380 26 });
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.
 45 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
451 40 const int i = static_cast<int>(ii);
452
4/4
✓ Branch 0 taken 15 times.
✓ Branch 1 taken 25 times.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 7 times.
 40 const bool valid = !(gate_query && d_mask && d_mask[i] < 0.5f);
453
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 1616 times.
 1656 for (int k = 0; k < D; ++k) {
454 1616 float acc = 0.0f;
455
2/2
✓ Branch 0 taken 1576 times.
✓ Branch 1 taken 40 times.
 1616 if (valid) {
456
2/2
✓ Branch 0 taken 67588 times.
✓ Branch 1 taken 1576 times.
 69164 for (int c = 0; c < D; ++c)
457 135176 acc += Wop[static_cast<std::size_t>(c) * D + k] *
458 67588 dOp[static_cast<std::size_t>(i) * D + c];
459 1576 }
460 1616 dYc[static_cast<std::size_t>(i) * D + k] = acc;
461 1616 }
462 40 });
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.
 222 parallel_for(static_cast<std::size_t>(D), [&](std::size_t ci) {
469 217 const int c = static_cast<int>(ci);
470
2/2
✓ Branch 0 taken 8027 times.
✓ Branch 1 taken 217 times.
 8244 for (int k = 0; k < D; ++k) {
471 8027 float acc = 0.0f;
472
2/2
✓ Branch 0 taken 53591 times.
✓ Branch 1 taken 8027 times.
 61618 for (int i = 0; i < Lq; ++i) {
473
6/6
✓ Branch 0 taken 10700 times.
✓ Branch 1 taken 42891 times.
✓ Branch 2 taken 5922 times.
✓ Branch 3 taken 4778 times.
✓ Branch 4 taken 4257 times.
✓ Branch 5 taken 1665 times.
 53591 if (gate_query && d_mask && d_mask[i] < 0.5f) continue;
474 103852 acc += dOp[static_cast<std::size_t>(i) * D + c] *
475 51926 Yp[static_cast<std::size_t>(i) * D + k];
476 51926 }
477 8027 dWop[static_cast<std::size_t>(c) * D + k] += acc;
478 8027 }
479 217 });
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.
 31 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
485 26 const int hh = static_cast<int>(hhi);
486
2/2
✓ Branch 0 taken 220 times.
✓ Branch 1 taken 26 times.
 246 for (int i = 0; i < Lq; ++i) {
487
2/2
✓ Branch 0 taken 2401 times.
✓ Branch 1 taken 220 times.
 2621 for (int j = 0; j < Lk; ++j) {
488 2401 float acc = 0.0f;
489
2/2
✓ Branch 0 taken 5671 times.
✓ Branch 1 taken 2401 times.
 8072 for (int k = 0; k < dh; ++k) {
490 5671 const float dy =
491 5671 dYc[static_cast<std::size_t>(i) * D + (hh * dh + k)];
492 5671 const float vv =
493 5671 Vp[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
494 5671 acc += dy * vv;
495 5671 }
496 2401 dAttn[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j] = acc;
497 2401 }
498 220 }
499
2/2
✓ Branch 0 taken 337 times.
✓ Branch 1 taken 26 times.
 363 for (int j = 0; j < Lk; ++j) {
500
2/2
✓ Branch 0 taken 2487 times.
✓ Branch 1 taken 337 times.
 2824 for (int k = 0; k < dh; ++k) {
501 2487 float acc = 0.0f;
502
2/2
✓ Branch 0 taken 22632 times.
✓ Branch 1 taken 2487 times.
 25119 for (int i = 0; i < Lq; ++i) {
503 22632 const float a =
504 22632 Ap[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
505 22632 const float dy =
506 22632 dYc[static_cast<std::size_t>(i) * D + (hh * dh + k)];
507 22632 acc += a * dy;
508 22632 }
509 2487 dVh[(static_cast<std::size_t>(hh) * Lk + j) * dh + k] = acc;
510 2487 }
511 337 }
512 26 });
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.
 31 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
517 26 const int hh = static_cast<int>(hhi);
518
2/2
✓ Branch 0 taken 222 times.
✓ Branch 1 taken 26 times.
 248 for (int i = 0; i < Lq; ++i) {
519 222 const float* prow =
520 222 Ap + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
521 222 const float* dprow =
522 222 dAttn.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
523 222 float* drow =
524 222 dScores.data() + (static_cast<std::size_t>(hh) * Lq + i) * Lk;
525
6/6
✓ Branch 0 taken 57 times.
✓ Branch 1 taken 165 times.
✓ Branch 2 taken 28 times.
✓ Branch 3 taken 29 times.
✓ Branch 4 taken 20 times.
✓ Branch 5 taken 8 times.
 222 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 214 float dot = 0.0f;
530
2/2
✓ Branch 0 taken 2796 times.
✓ Branch 1 taken 214 times.
 3010 for (int j = 0; j < Lk; ++j) dot += dprow[j] * prow[j];
531
2/2
✓ Branch 0 taken 2894 times.
✓ Branch 1 taken 214 times.
 3108 for (int j = 0; j < Lk; ++j) {
532
4/4
✓ Branch 0 taken 438 times.
✓ Branch 1 taken 2456 times.
✓ Branch 2 taken 298 times.
✓ Branch 3 taken 140 times.
 2894 if (d_mask && d_mask[j] < 0.5f) drow[j] = 0.0f;
533 2754 else drow[j] = prow[j] * (dprow[j] - dot) * inv_sqrtdh;
534 2894 }
535 214 }
536 26 });
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.
 31 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
541 26 const int hh = static_cast<int>(hhi);
542
2/2
✓ Branch 0 taken 211 times.
✓ Branch 1 taken 26 times.
 237 for (int i = 0; i < Lq; ++i) {
543
2/2
✓ Branch 0 taken 1434 times.
✓ Branch 1 taken 211 times.
 1645 for (int k = 0; k < dh; ++k) {
544 1434 float acc = 0.0f;
545
2/2
✓ Branch 0 taken 6528 times.
✓ Branch 1 taken 1434 times.
 7962 for (int j = 0; j < Lk; ++j) {
546 6528 const float ds =
547 6528 dScores[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
548 6528 const float kk =
549 6528 Kp[(static_cast<std::size_t>(hh) * Lk + j) * dh + k];
550 6528 acc += ds * kk;
551 6528 }
552 1434 dQh[(static_cast<std::size_t>(hh) * Lq + i) * dh + k] = acc;
553 1434 }
554 211 }
555
2/2
✓ Branch 0 taken 329 times.
✓ Branch 1 taken 26 times.
 355 for (int j = 0; j < Lk; ++j) {
556
2/2
✓ Branch 0 taken 2452 times.
✓ Branch 1 taken 329 times.
 2781 for (int k = 0; k < dh; ++k) {
557 2452 float acc = 0.0f;
558
2/2
✓ Branch 0 taken 22610 times.
✓ Branch 1 taken 2452 times.
 25062 for (int i = 0; i < Lq; ++i) {
559 22610 const float ds =
560 22610 dScores[(static_cast<std::size_t>(hh) * Lq + i) * Lk + j];
561 22610 const float qq =
562 22610 Qp[(static_cast<std::size_t>(hh) * Lq + i) * dh + k];
563 22610 acc += ds * qq;
564 22610 }
565 2452 dKh[(static_cast<std::size_t>(hh) * Lk + j) * dh + k] = acc;
566 2452 }
567 329 }
568 26 });
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.
 226 parallel_for(static_cast<std::size_t>(D), [&](std::size_t wrowi) {
575 221 const int wrow = static_cast<int>(wrowi);
576 221 const int hh = wrow / dh;
577 221 const int j = wrow % dh;
578
2/2
✓ Branch 0 taken 7360 times.
✓ Branch 1 taken 221 times.
 7581 for (int k = 0; k < D; ++k) {
579 7360 float acc = 0.0f;
580
2/2
✓ Branch 0 taken 17125 times.
✓ Branch 1 taken 7360 times.
 24485 for (int i = 0; i < Lq; ++i) {
581 17125 const float xv = Xp[static_cast<std::size_t>(i) * D + k];
582 17125 acc += dQh[(static_cast<std::size_t>(hh) * Lq + i) * dh + j] * xv;
583 17125 }
584 7360 dWqp[static_cast<std::size_t>(wrow) * D + k] += acc;
585 7360 }
586 221 });
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.
 221 parallel_for(static_cast<std::size_t>(D), [&](std::size_t wrowi) {
590 216 const int wrow = static_cast<int>(wrowi);
591 216 const int hh = wrow / dh;
592 216 const int j = wrow % dh;
593
2/2
✓ Branch 0 taken 6350 times.
✓ Branch 1 taken 216 times.
 6566 for (int k = 0; k < Dctx; ++k) {
594 6350 float ak = 0.0f, av = 0.0f;
595
2/2
✓ Branch 0 taken 23567 times.
✓ Branch 1 taken 6350 times.
 29917 for (int i = 0; i < Lk; ++i) {
596 23567 const float cv = Cp[static_cast<std::size_t>(i) * Dctx + k];
597 23567 ak += dKh[(static_cast<std::size_t>(hh) * Lk + i) * dh + j] * cv;
598 23567 av += dVh[(static_cast<std::size_t>(hh) * Lk + i) * dh + j] * cv;
599 23567 }
600 6350 const std::size_t idx = static_cast<std::size_t>(wrow) * Dctx + k;
601 6350 dWkp[idx] += ak;
602 6350 dWvp[idx] += av;
603 6350 }
604 216 });
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.
 40 parallel_for(static_cast<std::size_t>(Lq), [&](std::size_t ii) {
609 35 const int i = static_cast<int>(ii);
610
2/2
✓ Branch 0 taken 1665 times.
✓ Branch 1 taken 35 times.
 1700 for (int k = 0; k < D; ++k) {
611 1665 float acc = 0.0f;
612
2/2
✓ Branch 0 taken 8157 times.
✓ Branch 1 taken 1665 times.
 9822 for (int hh = 0; hh < H; ++hh) {
613
2/2
✓ Branch 0 taken 20518 times.
✓ Branch 1 taken 8157 times.
 28675 for (int j = 0; j < dh; ++j) {
614 20518 const int wrow = hh * dh + j;
615 20518 const float gq =
616 20518 dQh[(static_cast<std::size_t>(hh) * Lq + i) * dh + j];
617 20518 acc += gq * Wqp[static_cast<std::size_t>(wrow) * D + k];
618 20518 }
619 8157 }
620 1665 dXp[static_cast<std::size_t>(i) * D + k] = acc;
621 1665 }
622 35 });
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.
 59 parallel_for(static_cast<std::size_t>(Lk), [&](std::size_t ji) {
628 54 const int j = static_cast<int>(ji);
629
2/2
✓ Branch 0 taken 1872 times.
✓ Branch 1 taken 54 times.
 1926 for (int k = 0; k < Dctx; ++k) {
630 1872 float acc = 0.0f;
631
2/2
✓ Branch 0 taken 9518 times.
✓ Branch 1 taken 1872 times.
 11390 for (int hh = 0; hh < H; ++hh) {
632
2/2
✓ Branch 0 taken 51060 times.
✓ Branch 1 taken 9518 times.
 60578 for (int m = 0; m < dh; ++m) {
633 51060 const int wrow = hh * dh + m;
634 51060 const std::size_t widx =
635 51060 static_cast<std::size_t>(wrow) * Dctx + k;
636 51060 const float gk =
637 51060 dKh[(static_cast<std::size_t>(hh) * Lk + j) * dh + m];
638 51060 const float gv =
639 51060 dVh[(static_cast<std::size_t>(hh) * Lk + j) * dh + m];
640 51060 acc += gk * Wkp[widx] + gv * Wvp[widx];
641 51060 }
642 9518 }
643 1872 dCp[static_cast<std::size_t>(j) * Dctx + k] = acc;
644 1872 }
645 54 });
646 5 }
647
648 } // namespace brotensor::detail::cpu
649