GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 94.0% 109 / 0 / 116
Functions: 100.0% 4 / 0 / 4
Branches: 60.6% 80 / 0 / 132
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
17 #include <algorithm>
18 #include <cmath>
19 #include <stdexcept>
20 #include <string>
21 #include <vector>
22
23 namespace brotensor::detail::cpu {
24
25 namespace {
26
27 137 inline void check_fp32(const ::brotensor::Tensor& t,
28 const char* name) {
29
1/2
✓ Branch 0 taken 137 times.
✗ Branch 1 not taken.
 137 if (t.dtype != Dtype::FP32) {
30 ✗ throw std::runtime_error(std::string("self_attention_bias_forward: ") +
31 ✗ name + " must be FP32 (CPU backend is FP32-only)");
32 }
33 137 }
34
35 } // namespace
36
37 21 void self_attention_bias_forward(const ::brotensor::Tensor& X,
38 const ::brotensor::Tensor& Wq,
39 const ::brotensor::Tensor& Wk,
40 const ::brotensor::Tensor& Wv,
41 const ::brotensor::Tensor& Wo,
42 const ::brotensor::Tensor* bq,
43 const ::brotensor::Tensor* bk,
44 const ::brotensor::Tensor* bv,
45 const ::brotensor::Tensor* bo,
46 const float* d_mask,
47 const ::brotensor::Tensor* attn_bias,
48 int num_heads, float scale,
49 ::brotensor::Tensor& O) {
50 21 check_fp32(X, "X");
51 21 check_fp32(Wq, "Wq"); check_fp32(Wk, "Wk");
52 21 check_fp32(Wv, "Wv"); check_fp32(Wo, "Wo");
53 21 const int L = X.rows;
54 21 const int D = X.cols;
55
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (num_heads <= 0 || D % num_heads != 0) {
56 ✗ throw std::runtime_error("self_attention_bias_forward: num_heads must divide D");
57 }
58
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 ||
59 21 Wv.rows != D || Wv.cols != D || Wo.rows != D || Wo.cols != D) {
60 ✗ throw std::runtime_error("self_attention_bias_forward: Wq/Wk/Wv/Wo must be (D, D)");
61 }
62 21 const int H = num_heads;
63 21 const int dh = D / H;
64 21 const float* bias = nullptr;
65
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) {
66 16 check_fp32(*attn_bias, "attn_bias");
67
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (attn_bias->size() != H * L * L) {
68 ✗ throw std::runtime_error("self_attention_bias_forward: attn_bias must be (num_heads*L, L)");
69 }
70 16 bias = attn_bias->host_f32();
71 16 }
72 // Optional length-D projection biases (added post-projection).
73 105 auto bias_ptr = [&](const ::brotensor::Tensor* b, const char* name) -> const float* {
74
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;
75 16 check_fp32(*b, name);
76
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (b->size() != D)
77 ✗ throw std::runtime_error(std::string("self_attention_bias_forward: ") +
78 ✗ name + " must have D entries");
79 16 return b->host_f32();
80 84 };
81 21 const float* bqp = bias_ptr(bq, "bq");
82 21 const float* bkp = bias_ptr(bk, "bk");
83 21 const float* bvp = bias_ptr(bv, "bv");
84 21 const float* bop = bias_ptr(bo, "bo");
85
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) {
86 21 O.resize(L, D, Dtype::FP32);
87 21 }
88
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;
89
90 21 const float* Xp = X.host_f32();
91 21 const float* Wqp = Wq.host_f32();
92 21 const float* Wkp = Wk.host_f32();
93 21 const float* Wvp = Wv.host_f32();
94 21 const float* Wop = Wo.host_f32();
95 21 float* Op = O.host_f32_mut();
96
97 // Per-head projections: Qh / Kh / Vh laid out (H*L, dh).
98 21 std::vector<float> Qh(static_cast<size_t>(H) * L * dh);
99
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> Kh(static_cast<size_t>(H) * L * dh);
100
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> Vh(static_cast<size_t>(H) * L * dh);
101 84 auto project = [&](const float* W, const float* b, std::vector<float>& Out) {
102
2/2
✓ Branch 0 taken 315 times.
✓ Branch 1 taken 63 times.
 378 for (int hh = 0; hh < H; ++hh) {
103
2/2
✓ Branch 0 taken 5427 times.
✓ Branch 1 taken 315 times.
 5742 for (int i = 0; i < L; ++i) {
104 5427 const float* xr = Xp + static_cast<size_t>(i) * D;
105
2/2
✓ Branch 0 taken 77016 times.
✓ Branch 1 taken 5427 times.
 82443 for (int j = 0; j < dh; ++j) {
106 77016 const int o = hh * dh + j;
107 77016 const float* wr = W + static_cast<size_t>(o) * D;
108
2/2
✓ Branch 0 taken 6336 times.
✓ Branch 1 taken 70680 times.
 77016 float acc = b ? b[o] : 0.0f;
109
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];
110 77016 Out[(static_cast<size_t>(hh) * L + i) * dh + j] = acc;
111 77016 }
112 5427 }
113 315 }
114 63 };
115
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 project(Wqp, bqp, Qh);
116
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 project(Wkp, bkp, Kh);
117
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 project(Wvp, bvp, Vh);
118
119 // Yconcat (L, D): per-head attention output, concatenated.
120
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> Yc(static_cast<size_t>(L) * D, 0.0f);
121
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 std::vector<float> srow(L);
122
123
2/2
✓ Branch 0 taken 105 times.
✓ Branch 1 taken 21 times.
 126 for (int hh = 0; hh < H; ++hh) {
124
2/2
✓ Branch 0 taken 1809 times.
✓ Branch 1 taken 105 times.
 1914 for (int i = 0; i < L; ++i) {
125
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
126 // scores = scale * Q.K + bias
127 1717 const float* qr = &Qh[(static_cast<size_t>(hh) * L + i) * dh];
128 1717 float row_max = -1e30f;
129
2/2
✓ Branch 0 taken 54715 times.
✓ Branch 1 taken 1717 times.
 56432 for (int j = 0; j < L; ++j) {
130
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; }
131 54151 const float* kr = &Kh[(static_cast<size_t>(hh) * L + j) * dh];
132 54151 float s = 0.0f;
133
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];
134 54151 s *= scale;
135
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];
136 54151 srow[j] = s;
137
2/2
✓ Branch 0 taken 47507 times.
✓ Branch 1 taken 6644 times.
 54151 if (s > row_max) row_max = s;
138 54151 }
139 // softmax over valid keys
140 1717 float sum = 0.0f;
141
2/2
✓ Branch 0 taken 54715 times.
✓ Branch 1 taken 1717 times.
 56432 for (int j = 0; j < L; ++j) {
142
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; }
143 54151 const float e = std::exp(srow[j] - row_max);
144 54151 srow[j] = e;
145 54151 sum += e;
146 54151 }
147
1/2
✓ Branch 0 taken 1717 times.
✗ Branch 1 not taken.
 1717 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
148 // weighted sum of V
149
2/2
✓ Branch 0 taken 24936 times.
✓ Branch 1 taken 1717 times.
 26653 for (int k = 0; k < dh; ++k) {
150 24936 float acc = 0.0f;
151
2/2
✓ Branch 0 taken 840920 times.
✓ Branch 1 taken 24936 times.
 865856 for (int j = 0; j < L; ++j) {
152
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;
153 1672816 acc += srow[j] * inv *
154 836408 Vh[(static_cast<size_t>(hh) * L + j) * dh + k];
155 836408 }
156 24936 Yc[static_cast<size_t>(i) * D + (hh * dh + k)] = acc;
157 24936 }
158 1717 }
159 105 }
160
161 // Output projection O = Yconcat @ Wo^T, query-mask gated.
162
2/2
✓ Branch 0 taken 335 times.
✓ Branch 1 taken 21 times.
 356 for (int i = 0; i < L; ++i) {
163
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) {
164
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;
165 14 continue;
166 }
167 321 const float* yr = &Yc[static_cast<size_t>(i) * D];
168
2/2
✓ Branch 0 taken 24936 times.
✓ Branch 1 taken 321 times.
 25257 for (int c = 0; c < D; ++c) {
169 24936 const float* wr = Wop + static_cast<size_t>(c) * D;
170
2/2
✓ Branch 0 taken 2112 times.
✓ Branch 1 taken 22824 times.
 24936 float acc = bop ? bop[c] : 0.0f;
171
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];
172 24936 Op[static_cast<size_t>(i) * D + c] = acc;
173 24936 }
174 321 }
175 21 }
176
177 } // namespace brotensor::detail::cpu
178