GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 86.6% 149 / 0 / 172
Functions: 100.0% 7 / 0 / 7
Branches: 40.9% 76 / 0 / 186
src/cpu/kv_cache.cpp
Line Branch Exec Source
1 // ─── CPU KV-cache append + flash-attention decode (CHUNK 4) ────────────────
2 //
3 // FP32 scalar host implementations. Ports src/cuda/kv_cache.cu.
4 //
5 // IMPORTANT — dtype: the GPU kv-cache ops run FP16 internally (their tensors
6 // must be FP16). The CPU backend is FP32-only, so the CPU impls require FP32
7 // tensors. CPU↔GPU parity feeds FP16 to the GPU and FP32 to the CPU and
8 // compares with a loose FP16-driven tolerance (see test_kv_cache_parity.cpp).
9 //
10 // ── kv_cache_append ──
11 // Copies the L_new rows of K_new / V_new into K_cache / V_cache starting at
12 // row `cur_len`. Layout is row-major (rows, D); the destination slice is
13 // rows [cur_len, cur_len+L_new). Caches are pre-allocated by the caller to
14 // their max length L_max; rows outside the written slice are untouched.
15 // The op OVERWRITES exactly the [cur_len, cur_len+L_new) slice.
16 //
17 // ── flash_attention_decode ──
18 // Causal multi-head attention of Lq query rows against the first valid_len
19 // rows of the K/V cache. Q/K/V/O are (rows, D) with D = num_heads*head_dim;
20 // head h occupies columns [h*head_dim, (h+1)*head_dim). Scale = 1/sqrt(hd).
21 // Query row q maps to absolute sequence position p_q = seq_offset + q where
22 // seq_offset = valid_len - Lq; key kg is attended iff kg <= p_q (causal).
23 // Softmax over the valid causal keys; O OVERWRITTEN. Numerically this is a
24 // plain (max-subtracted) softmax — equivalent to the GPU's online/streaming
25 // softmax, just without the tiling.
26
27 #include <brotensor/tensor.h>
28
29 #include <algorithm>
30 #include <cmath>
31 #include <cstring>
32 #include <stdexcept>
33 #include <string>
34 #include <vector>
35
36 namespace brotensor::detail::cpu {
37
38 namespace {
39
40 221 inline void check_fp32(const ::brotensor::Tensor& t,
41 const char* op, const char* name) {
42
1/2
✓ Branch 0 taken 221 times.
✗ Branch 1 not taken.
 221 if (t.dtype != Dtype::FP32) {
43 ✗ throw std::runtime_error(std::string(op) + ": " + name +
44 " must be FP32 (CPU backend is FP32-only)");
45 }
46 221 }
47
48 // Gemma-2 tanh logit soft-capping, applied to a raw (already 1/sqrt(hd)-scaled)
49 // score. softcap <= 0 disables it — returns s unchanged (bit-identical path).
50 456868 inline float apply_softcap(float s, float softcap) {
51
2/2
✓ Branch 0 taken 167 times.
✓ Branch 1 taken 456701 times.
 456868 if (softcap > 0.0f) return softcap * std::tanh(s / softcap);
52 456701 return s;
53 456868 }
54
55 // One query head against keys [0, nkeys) of one KV head: scores, stable
56 // softmax, weighted V sum into orow. Key kg takes part iff kg >= lo and, when
57 // a mask is given, mask[kg] > 0.5. Both decode ops run exactly this code, so
58 // masked decode with a valid-prefix mask is bit-identical to the unmasked one
59 // (the compiler's FMA contraction and vectorisation choices are made once,
60 // here, rather than separately per caller).
61 829 void attend_head(const float* qrow, const float* Kp, const float* Vp, int Dkv,
62 int kv_head_off, int head_dim, int nkeys, const float* mask, int lo,
63 float inv_sqrt, float softcap, std::vector<float>& scores, float* orow) {
64 829 scores.assign(static_cast<std::size_t>(nkeys), 0.0f);
65 829 float run_max = -1e30f;
66
2/2
✓ Branch 0 taken 524966 times.
✓ Branch 1 taken 829 times.
 525795 for (int kg = 0; kg < nkeys; ++kg) {
67
6/6
✓ Branch 0 taken 524634 times.
✓ Branch 1 taken 332 times.
✓ Branch 2 taken 102076 times.
✓ Branch 3 taken 422558 times.
✓ Branch 4 taken 67766 times.
✓ Branch 5 taken 34310 times.
 524966 if (kg < lo || (mask && mask[kg] <= 0.5f)) {
68 68098 scores[static_cast<std::size_t>(kg)] = -1e30f; // out of window / masked
69 68098 continue;
70 }
71 456868 const float* krow = Kp + static_cast<std::size_t>(kg) * Dkv + kv_head_off;
72 456868 float dot = 0.0f;
73
2/2
✓ Branch 0 taken 56744512 times.
✓ Branch 1 taken 456868 times.
 57201380 for (int d = 0; d < head_dim; ++d) dot += qrow[d] * krow[d];
74 456868 float s = dot * inv_sqrt;
75 456868 s = apply_softcap(s, softcap); // Gemma-2 tanh soft-cap
76 456868 scores[static_cast<std::size_t>(kg)] = s;
77
2/2
✓ Branch 0 taken 452906 times.
✓ Branch 1 taken 3962 times.
 456868 if (s > run_max) run_max = s;
78 456868 }
79 // Stable softmax.
80 829 float sum = 0.0f;
81
2/2
✓ Branch 0 taken 524966 times.
✓ Branch 1 taken 829 times.
 525795 for (int kg = 0; kg < nkeys; ++kg) {
82
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 524966 times.
 524966 const float e = (run_max <= -1e29f)
83 ? 0.0f
84 524966 : std::exp(scores[static_cast<std::size_t>(kg)] - run_max);
85 524966 scores[static_cast<std::size_t>(kg)] = e;
86 524966 sum += e;
87 524966 }
88
1/2
✓ Branch 0 taken 829 times.
✗ Branch 1 not taken.
 829 const float inv = (sum > 0.0f) ? (1.0f / sum) : 0.0f;
89 // Weighted sum of V (the KV head's V).
90
2/2
✓ Branch 0 taken 68120 times.
✓ Branch 1 taken 829 times.
 68949 for (int d = 0; d < head_dim; ++d) {
91 68120 float acc = 0.0f;
92
2/2
✓ Branch 0 taken 65192016 times.
✓ Branch 1 taken 68120 times.
 65260136 for (int kg = 0; kg < nkeys; ++kg) {
93 65192016 const float s = scores[static_cast<std::size_t>(kg)];
94 // Skip zero-weight keys: identical math for any finite V, and
95 // masked rows may hold garbage the multiply would propagate.
96
2/2
✓ Branch 0 taken 8447504 times.
✓ Branch 1 taken 56744512 times.
 65192016 if (s == 0.0f) continue;
97 56744512 acc += s * Vp[static_cast<std::size_t>(kg) * Dkv + kv_head_off + d];
98 56744512 }
99 68120 orow[d] = acc * inv;
100 68120 }
101 829 }
102
103 } // namespace
104
105 14 void kv_cache_append(const ::brotensor::Tensor& K_new,
106 const ::brotensor::Tensor& V_new,
107 int cur_len,
108 ::brotensor::Tensor& K_cache,
109 ::brotensor::Tensor& V_cache) {
110 14 check_fp32(K_new, "kv_cache_append", "K_new");
111 14 check_fp32(V_new, "kv_cache_append", "V_new");
112 14 check_fp32(K_cache, "kv_cache_append", "K_cache");
113 14 check_fp32(V_cache, "kv_cache_append", "V_cache");
114
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (K_new.cols != V_new.cols || K_new.cols != K_cache.cols ||
115 14 K_cache.cols != V_cache.cols) {
116 ✗ throw std::runtime_error("kv_cache_append: column mismatch");
117 }
118
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (K_new.rows != V_new.rows) {
119 ✗ throw std::runtime_error("kv_cache_append: K_new/V_new row mismatch");
120 }
121
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (K_cache.rows != V_cache.rows) {
122 ✗ throw std::runtime_error("kv_cache_append: K_cache/V_cache row mismatch");
123 }
124 14 const int L_new = K_new.rows;
125 14 const int L_max = K_cache.rows;
126 14 const int D = K_new.cols;
127
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (cur_len < 0 || cur_len + L_new > L_max) {
128 ✗ throw std::runtime_error("kv_cache_append: cur_len + L_new exceeds cache capacity");
129 }
130
2/4
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 14 times.
 14 if (L_new == 0 || D == 0) return;
131
132 14 const std::size_t n = static_cast<std::size_t>(L_new) * D * sizeof(float);
133 14 const std::size_t dst_off = static_cast<std::size_t>(cur_len) * D;
134 14 std::memcpy(K_cache.host_f32_mut() + dst_off, K_new.host_f32(), n);
135 14 std::memcpy(V_cache.host_f32_mut() + dst_off, V_new.host_f32(), n);
136 14 }
137
138 46 void flash_attention_decode(const ::brotensor::Tensor& Q,
139 const ::brotensor::Tensor& K_cache,
140 const ::brotensor::Tensor& V_cache,
141 int valid_len,
142 int num_q_heads, int num_kv_heads,
143 ::brotensor::Tensor& O,
144 float attn_softcap, int window) {
145 46 check_fp32(Q, "flash_attention_decode", "Q");
146 46 check_fp32(K_cache, "flash_attention_decode", "K_cache");
147 46 check_fp32(V_cache, "flash_attention_decode", "V_cache");
148 46 const int Lq = Q.rows;
149 46 const int Dq = Q.cols;
150 46 const int Dkv = K_cache.cols;
151
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (V_cache.cols != Dkv) {
152 ✗ throw std::runtime_error("flash_attention_decode: K_cache.cols != V_cache.cols");
153 }
154
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (valid_len < 0 || valid_len > K_cache.rows || valid_len > V_cache.rows) {
155 ✗ throw std::runtime_error("flash_attention_decode: invalid valid_len");
156 }
157
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (valid_len < Lq) {
158 ✗ throw std::runtime_error("flash_attention_decode: valid_len must be >= Lq");
159 }
160
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (num_q_heads <= 0 || num_kv_heads <= 0) {
161 ✗ throw std::runtime_error("flash_attention_decode: num_q_heads / num_kv_heads must be positive");
162 }
163
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (num_q_heads % num_kv_heads != 0) {
164 ✗ throw std::runtime_error("flash_attention_decode: num_kv_heads must divide num_q_heads");
165 }
166
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (Dq % num_q_heads != 0 || Dkv % num_kv_heads != 0) {
167 ✗ throw std::runtime_error("flash_attention_decode: head_dim does not divide cols cleanly");
168 }
169 46 const int head_dim = Dq / num_q_heads;
170
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (Dkv / num_kv_heads != head_dim) {
171 ✗ throw std::runtime_error("flash_attention_decode: head_dim mismatch between Q and K/V");
172 }
173 46 const int q_per_kv = num_q_heads / num_kv_heads;
174
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 46 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 46 if (O.rows != Lq || O.cols != Dq || O.dtype != Dtype::FP32) {
175 46 O.resize(Lq, Dq, Dtype::FP32);
176 46 }
177
3/6
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 46 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 46 times.
 46 if (Lq == 0 || Dq == 0 || valid_len == 0) return;
178
179 46 const int seq_offset = valid_len - Lq;
180 46 const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(head_dim));
181
182 46 const float* Qp = Q.host_f32();
183 46 const float* Kp = K_cache.host_f32();
184 46 const float* Vp = V_cache.host_f32();
185 46 float* Op = O.host_f32_mut();
186
187 46 std::vector<float> scores;
188
2/2
✓ Branch 0 taken 46 times.
✓ Branch 1 taken 107 times.
 153 for (int q = 0; q < Lq; ++q) {
189 107 const int p_q = seq_offset + q;
190 107 const int klen = p_q + 1; // causal: keys 0..p_q inclusive
191 // Sliding window: keys below `lo` are out of band. window <= 0 keeps the
192 // band unbounded (lo == 0), preserving plain causal exactly.
193
3/4
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 102 times.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 107 const int lo = (window > 0) ? std::max(0, p_q - window + 1) : 0;
194
2/2
✓ Branch 0 taken 763 times.
✓ Branch 1 taken 107 times.
 870 for (int hq = 0; hq < num_q_heads; ++hq) {
195 763 const int hkv = hq / q_per_kv;
196 763 const int q_head_off = hq * head_dim;
197 763 const int kv_head_off = hkv * head_dim;
198 763 const float* qrow = Qp + q * Dq + q_head_off;
199 // The valid causal keys 0..p_q.
200
2/4
✓ Branch 0 taken 763 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 763 times.
✗ Branch 3 not taken.
 1526 attend_head(qrow, Kp, Vp, Dkv, kv_head_off, head_dim, klen, nullptr, lo,
201 763 inv_sqrt, attn_softcap, scores, Op + q * Dq + q_head_off);
202 763 }
203 107 }
204 46 }
205
206 // Masked fixed-capacity single-query decode — reference twin of the CUDA op.
207 // Validity lives entirely in the device-/host-resident FP32 key mask; keys
208 // with mask <= 0.5 are skipped before the dot product and softmax, so with a
209 // valid-prefix mask the result equals flash_attention_decode at L_q == 1.
210 9 void flash_attention_decode_masked(const ::brotensor::Tensor& Q,
211 const ::brotensor::Tensor& K_cache,
212 const ::brotensor::Tensor& V_cache,
213 const float* d_mask,
214 int num_q_heads, int num_kv_heads,
215 ::brotensor::Tensor& O,
216 float attn_softcap, int window) {
217 9 check_fp32(Q, "flash_attention_decode_masked", "Q");
218 9 check_fp32(K_cache, "flash_attention_decode_masked", "K_cache");
219 9 check_fp32(V_cache, "flash_attention_decode_masked", "V_cache");
220
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (d_mask == nullptr) {
221 ✗ throw std::runtime_error("flash_attention_decode_masked: d_mask must not be null");
222 }
223
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (Q.rows != 1) {
224 ✗ throw std::runtime_error("flash_attention_decode_masked: Q must be a single row (L_q == 1)");
225 }
226 9 const int Dq = Q.cols;
227 9 const int Dkv = K_cache.cols;
228 9 const int cap = K_cache.rows;
229
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (V_cache.cols != Dkv) {
230 ✗ throw std::runtime_error("flash_attention_decode_masked: K_cache.cols != V_cache.cols");
231 }
232
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (V_cache.rows != cap) {
233 ✗ throw std::runtime_error("flash_attention_decode_masked: K_cache/V_cache row mismatch");
234 }
235
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (num_q_heads <= 0 || num_kv_heads <= 0) {
236 ✗ throw std::runtime_error("flash_attention_decode_masked: num_q_heads / num_kv_heads must be positive");
237 }
238
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (num_q_heads % num_kv_heads != 0) {
239 ✗ throw std::runtime_error("flash_attention_decode_masked: num_kv_heads must divide num_q_heads");
240 }
241
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (Dq % num_q_heads != 0 || Dkv % num_kv_heads != 0) {
242 ✗ throw std::runtime_error("flash_attention_decode_masked: head_dim does not divide cols cleanly");
243 }
244 9 const int head_dim = Dq / num_q_heads;
245
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (Dkv / num_kv_heads != head_dim) {
246 ✗ throw std::runtime_error("flash_attention_decode_masked: head_dim mismatch between Q and K/V");
247 }
248 9 const int q_per_kv = num_q_heads / num_kv_heads;
249
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 9 if (O.rows != 1 || O.cols != Dq || O.dtype != Dtype::FP32) {
250 9 O.resize(1, Dq, Dtype::FP32);
251 9 }
252
2/4
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 9 times.
 9 if (Dq == 0 || cap == 0) return;
253
254 9 const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(head_dim));
255 9 const float* Qp = Q.host_f32();
256 9 const float* Kp = K_cache.host_f32();
257 9 const float* Vp = V_cache.host_f32();
258 9 float* Op = O.host_f32_mut();
259
260 // Sliding window: the query's absolute position is the highest valid key
261 // index (the last 1 in the mask). window <= 0 keeps the full valid set —
262 // lo == 0, the pre-window behaviour exactly.
263 9 int lo = 0;
264
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 1 time.
 9 if (window > 0) {
265 1 int p_max = -1;
266
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 1 time.
 25 for (int kg = 0; kg < cap; ++kg) {
267
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 14 times.
 24 if (d_mask[kg] > 0.5f) p_max = kg;
268 24 }
269
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 lo = (p_max >= 0) ? std::max(0, p_max - window + 1) : 0;
270 1 }
271
272 9 std::vector<float> scores;
273
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 66 times.
 75 for (int hq = 0; hq < num_q_heads; ++hq) {
274 66 const int hkv = hq / q_per_kv;
275 66 const int q_head_off = hq * head_dim;
276 66 const int kv_head_off = hkv * head_dim;
277
2/4
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 66 times.
✗ Branch 3 not taken.
 132 attend_head(Qp + q_head_off, Kp, Vp, Dkv, kv_head_off, head_dim, cap, d_mask, lo,
278 66 inv_sqrt, attn_softcap, scores, Op + q_head_off);
279 66 }
280 9 }
281
282 // Device-length decode — on the CPU the length pointer is a host pointer, so
283 // this is flash_attention_decode at *d_len (zeros for an empty cache).
284 13 void flash_attention_decode_len(const ::brotensor::Tensor& Q, const ::brotensor::Tensor& K_cache,
285 const ::brotensor::Tensor& V_cache, const int32_t* d_len, int num_q_heads,
286 int num_kv_heads, ::brotensor::Tensor& O) {
287
1/4
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 13 if (d_len == nullptr) throw std::runtime_error("flash_attention_decode_len: d_len must not be null");
288
1/4
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 13 if (*d_len < 0 || *d_len > K_cache.rows) throw std::runtime_error("flash_attention_decode_len: bad length");
289
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (*d_len == 0) {
290 ✗ if (O.rows != Q.rows || O.cols != Q.cols || O.dtype != Q.dtype) O.resize(Q.rows, Q.cols, Q.dtype);
291 ✗ std::fill(O.ptr(), O.ptr() + O.size(), 0.0f);
292 ✗ return;
293 }
294 13 flash_attention_decode(Q, K_cache, V_cache, *d_len, num_q_heads, num_kv_heads, O, 0.0f, 0);
295 13 }
296
297 } // namespace brotensor::detail::cpu
298