GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 93.3% 139 / 0 / 149
Functions: 100.0% 6 / 0 / 6
Branches: 57.8% 89 / 0 / 154
src/cpu/self_attention_bias.cpp
Line Branch Exec Source
1 // ─── CPU self-attention with additive pre-softmax bias ─────────────────────
2 //
3 // FP32 scalar host implementation. Ports src/cuda/self_attention_bias.cu —
4 // FP32 path only (the CPU backend is FP32-only per CLAUDE.md).
5 //
6 // Multi-head self-attention with an optional per-head (L, L) additive bias:
7 // S[h,q,k] = scale * (Q_h[q] . K_h[k]) + attn_bias[h*L+q, k]
8 // O = concat_h( softmax_k(S[h]) @ V_h ) @ Wo
9 //
10 // Wq/Wk/Wv/Wo are (D, D); per-head split takes contiguous weight rows
11 // hh*dh .. hh*dh+dh. d_mask is a length-L key-validity buffer that also gates
12 // padded query rows (their output row is zeroed). attn_bias is FP32 or null.
13 // O is fully overwritten.
14
15 #include <brotensor/tensor.h>
16 #include <brotensor/detail/cpu/thread_pool.h>
17
18 #include <algorithm>
19 #include <cmath>
20 #include <stdexcept>
21 #include <string>
22 #include <vector>
23
24 namespace brotensor::detail::cpu {
25
26 namespace {
27
28 143 inline void check_fp32(const ::brotensor::Tensor& t,
29 const char* name) {
30
1/2
✓ Branch 0 taken 143 times.
✗ Branch 1 not taken.
 143 if (t.dtype != Dtype::FP32) {
31 ✗ throw std::runtime_error(std::string("self_attention_bias_forward: ") +
32 ✗ name + " must be FP32 (CPU backend is FP32-only)");
33 }
34 143 }
35
36 } // namespace
37
38 21 void self_attention_bias_forward(const ::brotensor::Tensor& X,
39 const ::brotensor::Tensor& Wq,
40 const ::brotensor::Tensor& Wk,
41 const ::brotensor::Tensor& Wv,
42 const ::brotensor::Tensor& Wo,
43 const ::brotensor::Tensor* bq,
44 const ::brotensor::Tensor* bk,
45 const ::brotensor::Tensor* bv,
46 const ::brotensor::Tensor* bo,
47 const float* d_mask,
48 const ::brotensor::Tensor* attn_bias,
49 int num_heads, float scale,
50 ::brotensor::Tensor& O) {
51 21 check_fp32(X, "X");
52 21 check_fp32(Wq, "Wq"); check_fp32(Wk, "Wk");
53 21 check_fp32(Wv, "Wv"); check_fp32(Wo, "Wo");
54 21 const int L = X.rows;
55 21 const int D = X.cols;
56
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (num_heads <= 0 || D % num_heads != 0) {
57 ✗ throw std::runtime_error("self_attention_bias_forward: num_heads must divide D");
58 }
59
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 42 if (Wq.rows != D || Wq.cols != D || Wk.rows != D || Wk.cols != D ||
60 21 Wv.rows != D || Wv.cols != D || Wo.rows != D || Wo.cols != D) {
61 ✗ throw std::runtime_error("self_attention_bias_forward: Wq/Wk/Wv/Wo must be (D, D)");
62 }
63 21 const int H = num_heads;
64 21 const int dh = D / H;
65 21 const float* bias = nullptr;
66
3/4
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 16 times.
 21 if (attn_bias && attn_bias->data) {
67 16 check_fp32(*attn_bias, "attn_bias");
68
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (attn_bias->size() != H * L * L) {
69 ✗ throw std::runtime_error("self_attention_bias_forward: attn_bias must be (num_heads*L, L)");
70 }
71 16 bias = attn_bias->host_f32();
72 16 }
73 // Optional length-D projection biases (added post-projection).
74 105 auto bias_ptr = [&](const ::brotensor::Tensor* b, const char* name) -> const float* {
75
3/4
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 68 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 16 times.
 84 if (!b || !b->data) return nullptr;
76 16 check_fp32(*b, name);
77
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (b->size() != D)
78 ✗ throw std::runtime_error(std::string("self_attention_bias_forward: ") +
79 ✗ name + " must have D entries");
80 16 return b->host_f32();
81 84 };
82 21 const float* bqp = bias_ptr(bq, "bq");
83 21 const float* bkp = bias_ptr(bk, "bk");
84 21 const float* bvp = bias_ptr(bv, "bv");
85 21 const float* bop = bias_ptr(bo, "bo");
86
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 21 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 21 if (O.rows != L || O.cols != D || O.dtype != Dtype::FP32) {
87 21 O.resize(L, D, Dtype::FP32);
88 21 }
89
2/4
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 21 times.
 21 if (L == 0 || D == 0) return;
90
91 21 const float* Xp = X.host_f32();
92 21 const float* Wqp = Wq.host_f32();
93 21 const float* Wkp = Wk.host_f32();
94 21 const float* Wvp = Wv.host_f32();
95 21 const float* Wop = Wo.host_f32();
96 21 float* Op = O.host_f32_mut();
97
98 // Per-head projections: Qh / Kh / Vh laid out (H*L, dh).
99 21 std::vector<float> Qh(static_cast<size_t>(H) * L * dh);
100
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> Kh(static_cast<size_t>(H) * L * dh);
101
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> Vh(static_cast<size_t>(H) * L * dh);
102 84 auto project = [&](const float* W, const float* b, std::vector<float>& Out) {
103
2/2
✓ Branch 0 taken 315 times.
✓ Branch 1 taken 63 times.
 378 for (int hh = 0; hh < H; ++hh) {
104
2/2
✓ Branch 0 taken 5427 times.
✓ Branch 1 taken 315 times.
 5742 for (int i = 0; i < L; ++i) {
105 5427 const float* xr = Xp + static_cast<size_t>(i) * D;
106
2/2
✓ Branch 0 taken 77016 times.
✓ Branch 1 taken 5427 times.
 82443 for (int j = 0; j < dh; ++j) {
107 77016 const int o = hh * dh + j;
108 77016 const float* wr = W + static_cast<size_t>(o) * D;
109
2/2
✓ Branch 0 taken 6336 times.
✓ Branch 1 taken 70680 times.
 77016 float acc = b ? b[o] : 0.0f;
110
2/2
✓ Branch 0 taken 10889280 times.
✓ Branch 1 taken 77016 times.
 10966296 for (int k = 0; k < D; ++k) acc += xr[k] * wr[k];
111 77016 Out[(static_cast<size_t>(hh) * L + i) * dh + j] = acc;
112 77016 }
113 5427 }
114 315 }
115 63 };
116
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 project(Wqp, bqp, Qh);
117
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 project(Wkp, bkp, Kh);
118
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 project(Wvp, bvp, Vh);
119
120 // Yconcat (L, D): per-head attention output, concatenated.
121
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> Yc(static_cast<size_t>(L) * D, 0.0f);
122
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> srow(L);
123
124
2/2
✓ Branch 0 taken 105 times.
✓ Branch 1 taken 21 times.
 126 for (int hh = 0; hh < H; ++hh) {
125
2/2
✓ Branch 0 taken 1809 times.
✓ Branch 1 taken 105 times.
 1914 for (int i = 0; i < L; ++i) {
126
4/4
✓ Branch 0 taken 160 times.
✓ Branch 1 taken 1649 times.
✓ Branch 2 taken 68 times.
✓ Branch 3 taken 92 times.
 1809 if (d_mask && d_mask[i] < 0.5f) continue; // padded query → 0 row
127 // scores = scale * Q.K + bias
128 1717 const float* qr = &Qh[(static_cast<size_t>(hh) * L + i) * dh];
129 1717 float row_max = -1e30f;
130
2/2
✓ Branch 0 taken 54715 times.
✓ Branch 1 taken 1717 times.
 56432 for (int j = 0; j < L; ++j) {
131
4/4
✓ Branch 0 taken 992 times.
✓ Branch 1 taken 53723 times.
✓ Branch 2 taken 428 times.
✓ Branch 3 taken 564 times.
 54715 if (d_mask && d_mask[j] < 0.5f) { srow[j] = 0.0f; continue; }
132 54151 const float* kr = &Kh[(static_cast<size_t>(hh) * L + j) * dh];
133 54151 float s = 0.0f;
134
2/2
✓ Branch 0 taken 836408 times.
✓ Branch 1 taken 54151 times.
 890559 for (int k = 0; k < dh; ++k) s += qr[k] * kr[k];
135 54151 s *= scale;
136
2/2
✓ Branch 0 taken 2496 times.
✓ Branch 1 taken 51655 times.
 54151 if (bias) s += bias[(static_cast<size_t>(hh) * L + i) * L + j];
137 54151 srow[j] = s;
138
2/2
✓ Branch 0 taken 47507 times.
✓ Branch 1 taken 6644 times.
 54151 if (s > row_max) row_max = s;
139 54151 }
140 // softmax over valid keys
141 1717 float sum = 0.0f;
142
2/2
✓ Branch 0 taken 54715 times.
✓ Branch 1 taken 1717 times.
 56432 for (int j = 0; j < L; ++j) {
143
4/4
✓ Branch 0 taken 992 times.
✓ Branch 1 taken 53723 times.
✓ Branch 2 taken 428 times.
✓ Branch 3 taken 564 times.
 54715 if (d_mask && d_mask[j] < 0.5f) { srow[j] = 0.0f; continue; }
144 54151 const float e = std::exp(srow[j] - row_max);
145 54151 srow[j] = e;
146 54151 sum += e;
147 54151 }
148
1/2
✓ Branch 0 taken 1717 times.
✗ Branch 1 not taken.
 1717 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
149 // weighted sum of V
150
2/2
✓ Branch 0 taken 24936 times.
✓ Branch 1 taken 1717 times.
 26653 for (int k = 0; k < dh; ++k) {
151 24936 float acc = 0.0f;
152
2/2
✓ Branch 0 taken 840920 times.
✓ Branch 1 taken 24936 times.
 865856 for (int j = 0; j < L; ++j) {
153
4/4
✓ Branch 0 taken 7936 times.
✓ Branch 1 taken 832984 times.
✓ Branch 2 taken 3424 times.
✓ Branch 3 taken 4512 times.
 840920 if (d_mask && d_mask[j] < 0.5f) continue;
154 1672816 acc += srow[j] * inv *
155 836408 Vh[(static_cast<size_t>(hh) * L + j) * dh + k];
156 836408 }
157 24936 Yc[static_cast<size_t>(i) * D + (hh * dh + k)] = acc;
158 24936 }
159 1717 }
160 105 }
161
162 // Output projection O = Yconcat @ Wo^T, query-mask gated.
163
2/2
✓ Branch 0 taken 335 times.
✓ Branch 1 taken 21 times.
 356 for (int i = 0; i < L; ++i) {
164
4/4
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 311 times.
✓ Branch 2 taken 10 times.
✓ Branch 3 taken 14 times.
 335 if (d_mask && d_mask[i] < 0.5f) {
165
2/2
✓ Branch 0 taken 736 times.
✓ Branch 1 taken 14 times.
 750 for (int c = 0; c < D; ++c) Op[static_cast<size_t>(i) * D + c] = 0.0f;
166 14 continue;
167 }
168 321 const float* yr = &Yc[static_cast<size_t>(i) * D];
169
2/2
✓ Branch 0 taken 24936 times.
✓ Branch 1 taken 321 times.
 25257 for (int c = 0; c < D; ++c) {
170 24936 const float* wr = Wop + static_cast<size_t>(c) * D;
171
2/2
✓ Branch 0 taken 2112 times.
✓ Branch 1 taken 22824 times.
 24936 float acc = bop ? bop[c] : 0.0f;
172
2/2
✓ Branch 0 taken 3587776 times.
✓ Branch 1 taken 24936 times.
 3612712 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
173 24936 Op[static_cast<size_t>(i) * D + c] = acc;
174 24936 }
175 321 }
176 21 }
177
178 // ─── Transformer-XL relative-position bias ─────────────────────────────────
179 //
180 // Bias[h*T + q, k] = sum_d Qv[q, h*dk + d] * Pk[(T-1-q) + k, h*dk + d].
181 //
182 // The rel_shift is expressed as the column offset (T-1-q) rather than by
183 // building matrix_bd and shifting it: the shifted read is the same cost and
184 // there is no (T, 2T-1) intermediate. Parallelised over (head, query) rows,
185 // which is num_heads*T of them — hundreds, so the pool has something to spread.
186 3 void rel_pos_bias_xl_forward(const ::brotensor::Tensor& Qv,
187 const ::brotensor::Tensor& Pk,
188 int num_heads, int head_dim,
189 ::brotensor::Tensor& Bias) {
190 3 check_fp32(Qv, "Qv");
191 3 check_fp32(Pk, "Pk");
192 3 const int T = Qv.rows;
193 3 const int D = Qv.cols;
194
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (num_heads <= 0 || head_dim <= 0 || num_heads * head_dim != D)
195 ✗ throw std::runtime_error("rel_pos_bias_xl_forward: num_heads*head_dim "
196 "must equal Qv.cols");
197
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (Pk.cols != D || Pk.rows != 2 * T - 1)
198 ✗ throw std::runtime_error("rel_pos_bias_xl_forward: Pk must be "
199 "(2*Qv.rows - 1, Qv.cols)");
200
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 3 if (Bias.rows != num_heads * T || Bias.cols != T ||
201 ✗ Bias.dtype != Dtype::FP32) {
202 3 Bias = ::brotensor::Tensor::empty_on(Qv.device, num_heads * T, T,
203 Dtype::FP32);
204 3 }
205
206 3 const float* qv = static_cast<const float*>(Qv.data);
207 3 const float* pk = static_cast<const float*>(Pk.data);
208 3 float* out = static_cast<float*>(Bias.data);
209
210
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 parallel_for(static_cast<std::size_t>(num_heads) * static_cast<std::size_t>(T),
211 1844 [&](std::size_t row) {
212 1841 const int h = static_cast<int>(row / static_cast<std::size_t>(T));
213 1841 const int q = static_cast<int>(row % static_cast<std::size_t>(T));
214 1841 const int co = h * head_dim;
215 1841 const float* qrow = qv + static_cast<std::size_t>(q) * D + co;
216 1841 const int base = T - 1 - q;
217 1841 float* orow = out + row * static_cast<std::size_t>(T);
218
2/2
✓ Branch 0 taken 397922 times.
✓ Branch 1 taken 1841 times.
 399763 for (int k = 0; k < T; ++k) {
219 397922 const float* prow =
220 397922 pk + static_cast<std::size_t>(base + k) * D + co;
221 397922 float s = 0.0f;
222
2/2
✓ Branch 0 taken 17655135 times.
✓ Branch 1 taken 397922 times.
 18053057 for (int d = 0; d < head_dim; ++d) s += qrow[d] * prow[d];
223 397922 orow[k] = s;
224 397922 }
225 1841 });
226 3 }
227
228 } // namespace brotensor::detail::cpu
229