GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 89.2% 166 / 0 / 186
Functions: 100.0% 5 / 0 / 5
Branches: 46.7% 86 / 0 / 184
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 <cmath>
30 #include <cstring>
31 #include <stdexcept>
32 #include <string>
33 #include <vector>
34
35 namespace brotensor::detail::cpu {
36
37 namespace {
38
39 170 inline void check_fp32(const ::brotensor::Tensor& t,
40 const char* op, const char* name) {
41
1/2
✓ Branch 0 taken 170 times.
✗ Branch 1 not taken.
 170 if (t.dtype != Dtype::FP32) {
42 ✗ throw std::runtime_error(std::string(op) + ": " + name +
43 " must be FP32 (CPU backend is FP32-only)");
44 }
45 170 }
46
47 // Gemma-2 tanh logit soft-capping, applied to a raw (already 1/sqrt(hd)-scaled)
48 // score. softcap <= 0 disables it — returns s unchanged (bit-identical path).
49 4262 inline float apply_softcap(float s, float softcap) {
50
2/2
✓ Branch 0 taken 167 times.
✓ Branch 1 taken 4095 times.
 4262 if (softcap > 0.0f) return softcap * std::tanh(s / softcap);
51 4095 return s;
52 4262 }
53
54 } // namespace
55
56 14 void kv_cache_append(const ::brotensor::Tensor& K_new,
57 const ::brotensor::Tensor& V_new,
58 int cur_len,
59 ::brotensor::Tensor& K_cache,
60 ::brotensor::Tensor& V_cache) {
61 14 check_fp32(K_new, "kv_cache_append", "K_new");
62 14 check_fp32(V_new, "kv_cache_append", "V_new");
63 14 check_fp32(K_cache, "kv_cache_append", "K_cache");
64 14 check_fp32(V_cache, "kv_cache_append", "V_cache");
65
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 ||
66 14 K_cache.cols != V_cache.cols) {
67 ✗ throw std::runtime_error("kv_cache_append: column mismatch");
68 }
69
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (K_new.rows != V_new.rows) {
70 ✗ throw std::runtime_error("kv_cache_append: K_new/V_new row mismatch");
71 }
72
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (K_cache.rows != V_cache.rows) {
73 ✗ throw std::runtime_error("kv_cache_append: K_cache/V_cache row mismatch");
74 }
75 14 const int L_new = K_new.rows;
76 14 const int L_max = K_cache.rows;
77 14 const int D = K_new.cols;
78
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (cur_len < 0 || cur_len + L_new > L_max) {
79 ✗ throw std::runtime_error("kv_cache_append: cur_len + L_new exceeds cache capacity");
80 }
81
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;
82
83 14 const std::size_t n = static_cast<std::size_t>(L_new) * D * sizeof(float);
84 14 const std::size_t dst_off = static_cast<std::size_t>(cur_len) * D;
85 14 std::memcpy(K_cache.host_f32_mut() + dst_off, K_new.host_f32(), n);
86 14 std::memcpy(V_cache.host_f32_mut() + dst_off, V_new.host_f32(), n);
87 14 }
88
89 31 void flash_attention_decode(const ::brotensor::Tensor& Q,
90 const ::brotensor::Tensor& K_cache,
91 const ::brotensor::Tensor& V_cache,
92 int valid_len,
93 int num_q_heads, int num_kv_heads,
94 ::brotensor::Tensor& O,
95 float attn_softcap, int window) {
96 31 check_fp32(Q, "flash_attention_decode", "Q");
97 31 check_fp32(K_cache, "flash_attention_decode", "K_cache");
98 31 check_fp32(V_cache, "flash_attention_decode", "V_cache");
99 31 const int Lq = Q.rows;
100 31 const int Dq = Q.cols;
101 31 const int Dkv = K_cache.cols;
102
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (V_cache.cols != Dkv) {
103 ✗ throw std::runtime_error("flash_attention_decode: K_cache.cols != V_cache.cols");
104 }
105
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (valid_len < 0 || valid_len > K_cache.rows || valid_len > V_cache.rows) {
106 ✗ throw std::runtime_error("flash_attention_decode: invalid valid_len");
107 }
108
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (valid_len < Lq) {
109 ✗ throw std::runtime_error("flash_attention_decode: valid_len must be >= Lq");
110 }
111
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (num_q_heads <= 0 || num_kv_heads <= 0) {
112 ✗ throw std::runtime_error("flash_attention_decode: num_q_heads / num_kv_heads must be positive");
113 }
114
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (num_q_heads % num_kv_heads != 0) {
115 ✗ throw std::runtime_error("flash_attention_decode: num_kv_heads must divide num_q_heads");
116 }
117
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (Dq % num_q_heads != 0 || Dkv % num_kv_heads != 0) {
118 ✗ throw std::runtime_error("flash_attention_decode: head_dim does not divide cols cleanly");
119 }
120 31 const int head_dim = Dq / num_q_heads;
121
1/2
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
 31 if (Dkv / num_kv_heads != head_dim) {
122 ✗ throw std::runtime_error("flash_attention_decode: head_dim mismatch between Q and K/V");
123 }
124 31 const int q_per_kv = num_q_heads / num_kv_heads;
125
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 31 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 31 if (O.rows != Lq || O.cols != Dq || O.dtype != Dtype::FP32) {
126 31 O.resize(Lq, Dq, Dtype::FP32);
127 31 }
128
3/6
✓ Branch 0 taken 31 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 31 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 31 times.
 31 if (Lq == 0 || Dq == 0 || valid_len == 0) return;
129
130 31 const int seq_offset = valid_len - Lq;
131 31 const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(head_dim));
132
133 31 const float* Qp = Q.host_f32();
134 31 const float* Kp = K_cache.host_f32();
135 31 const float* Vp = V_cache.host_f32();
136 31 float* Op = O.host_f32_mut();
137
138 31 std::vector<float> scores;
139
2/2
✓ Branch 0 taken 69 times.
✓ Branch 1 taken 31 times.
 100 for (int q = 0; q < Lq; ++q) {
140 69 const int p_q = seq_offset + q;
141 69 const int klen = p_q + 1; // causal: keys 0..p_q inclusive
142 // Sliding window: keys below `lo` are out of band. window <= 0 keeps the
143 // band unbounded (lo == 0), preserving plain causal exactly.
144
3/4
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 64 times.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 69 const int lo = (window > 0) ? std::max(0, p_q - window + 1) : 0;
145
2/2
✓ Branch 0 taken 245 times.
✓ Branch 1 taken 69 times.
 314 for (int hq = 0; hq < num_q_heads; ++hq) {
146 245 const int hkv = hq / q_per_kv;
147 245 const int q_head_off = hq * head_dim;
148 245 const int kv_head_off = hkv * head_dim;
149 245 const float* qrow = Qp + q * Dq + q_head_off;
150
151 // Scores against the valid causal keys.
152
1/2
✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
 245 scores.assign(klen, 0.0f);
153 245 float run_max = -1e30f;
154
2/2
✓ Branch 0 taken 3352 times.
✓ Branch 1 taken 245 times.
 3597 for (int kg = 0; kg < klen; ++kg) {
155
2/2
✓ Branch 0 taken 296 times.
✓ Branch 1 taken 3056 times.
 3352 if (kg < lo) { scores[kg] = -1e30f; continue; } // out of window
156 3056 const float* krow = Kp + kg * Dkv + kv_head_off;
157 3056 float dot = 0.0f;
158
2/2
✓ Branch 0 taken 49856 times.
✓ Branch 1 taken 3056 times.
 52912 for (int d = 0; d < head_dim; ++d) dot += qrow[d] * krow[d];
159 3056 float s = dot * inv_sqrt;
160
1/2
✓ Branch 0 taken 3056 times.
✗ Branch 1 not taken.
 3056 s = apply_softcap(s, attn_softcap); // Gemma-2 tanh soft-cap
161 3056 scores[kg] = s;
162
2/2
✓ Branch 0 taken 2414 times.
✓ Branch 1 taken 642 times.
 3056 if (s > run_max) run_max = s;
163 3056 }
164 // Stable softmax.
165 245 float sum = 0.0f;
166
2/2
✓ Branch 0 taken 3352 times.
✓ Branch 1 taken 245 times.
 3597 for (int kg = 0; kg < klen; ++kg) {
167 3352 const float e = std::exp(scores[kg] - run_max);
168 3352 scores[kg] = e;
169 3352 sum += e;
170 3352 }
171
1/2
✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
 245 const float inv = (sum > 0.0f) ? (1.0f / sum) : 0.0f;
172 // Weighted sum of V (KV head's V).
173 245 float* orow = Op + q * Dq + q_head_off;
174
2/2
✓ Branch 0 taken 3128 times.
✓ Branch 1 taken 245 times.
 3373 for (int d = 0; d < head_dim; ++d) {
175 3128 float acc = 0.0f;
176
2/2
✓ Branch 0 taken 54336 times.
✓ Branch 1 taken 3128 times.
 57464 for (int kg = 0; kg < klen; ++kg) {
177 54336 acc += scores[kg] * Vp[kg * Dkv + kv_head_off + d];
178 54336 }
179 3128 orow[d] = acc * inv;
180 3128 }
181 245 }
182 69 }
183 31 }
184
185 // Masked fixed-capacity single-query decode — reference twin of the CUDA op.
186 // Validity lives entirely in the device-/host-resident FP32 key mask; keys
187 // with mask <= 0.5 are skipped before the dot product and softmax, so with a
188 // valid-prefix mask the result equals flash_attention_decode at L_q == 1.
189 7 void flash_attention_decode_masked(const ::brotensor::Tensor& Q,
190 const ::brotensor::Tensor& K_cache,
191 const ::brotensor::Tensor& V_cache,
192 const float* d_mask,
193 int num_q_heads, int num_kv_heads,
194 ::brotensor::Tensor& O,
195 float attn_softcap, int window) {
196 7 check_fp32(Q, "flash_attention_decode_masked", "Q");
197 7 check_fp32(K_cache, "flash_attention_decode_masked", "K_cache");
198 7 check_fp32(V_cache, "flash_attention_decode_masked", "V_cache");
199
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (d_mask == nullptr) {
200 ✗ throw std::runtime_error("flash_attention_decode_masked: d_mask must not be null");
201 }
202
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (Q.rows != 1) {
203 ✗ throw std::runtime_error("flash_attention_decode_masked: Q must be a single row (L_q == 1)");
204 }
205 7 const int Dq = Q.cols;
206 7 const int Dkv = K_cache.cols;
207 7 const int cap = K_cache.rows;
208
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (V_cache.cols != Dkv) {
209 ✗ throw std::runtime_error("flash_attention_decode_masked: K_cache.cols != V_cache.cols");
210 }
211
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (V_cache.rows != cap) {
212 ✗ throw std::runtime_error("flash_attention_decode_masked: K_cache/V_cache row mismatch");
213 }
214
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (num_q_heads <= 0 || num_kv_heads <= 0) {
215 ✗ throw std::runtime_error("flash_attention_decode_masked: num_q_heads / num_kv_heads must be positive");
216 }
217
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (num_q_heads % num_kv_heads != 0) {
218 ✗ throw std::runtime_error("flash_attention_decode_masked: num_kv_heads must divide num_q_heads");
219 }
220
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (Dq % num_q_heads != 0 || Dkv % num_kv_heads != 0) {
221 ✗ throw std::runtime_error("flash_attention_decode_masked: head_dim does not divide cols cleanly");
222 }
223 7 const int head_dim = Dq / num_q_heads;
224
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (Dkv / num_kv_heads != head_dim) {
225 ✗ throw std::runtime_error("flash_attention_decode_masked: head_dim mismatch between Q and K/V");
226 }
227 7 const int q_per_kv = num_q_heads / num_kv_heads;
228
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 if (O.rows != 1 || O.cols != Dq || O.dtype != Dtype::FP32) {
229 7 O.resize(1, Dq, Dtype::FP32);
230 7 }
231
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 7 times.
 7 if (Dq == 0 || cap == 0) return;
232
233 7 const float inv_sqrt = 1.0f / std::sqrt(static_cast<float>(head_dim));
234 7 const float* Qp = Q.host_f32();
235 7 const float* Kp = K_cache.host_f32();
236 7 const float* Vp = V_cache.host_f32();
237 7 float* Op = O.host_f32_mut();
238
239 // Sliding window: the query's absolute position is the highest valid key
240 // index (the last 1 in the mask). window <= 0 keeps the full valid set —
241 // lo == 0, the pre-window behaviour exactly.
242 7 int lo = 0;
243
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
 7 if (window > 0) {
244 1 int p_max = -1;
245
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 1 time.
 25 for (int kg = 0; kg < cap; ++kg) {
246
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 14 times.
 24 if (d_mask[kg] > 0.5f) p_max = kg;
247 24 }
248
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 lo = (p_max >= 0) ? std::max(0, p_max - window + 1) : 0;
249 1 }
250
251 7 std::vector<float> scores;
252
2/2
✓ Branch 0 taken 34 times.
✓ Branch 1 taken 7 times.
 41 for (int hq = 0; hq < num_q_heads; ++hq) {
253 34 const int hkv = hq / q_per_kv;
254 34 const int q_head_off = hq * head_dim;
255 34 const int kv_head_off = hkv * head_dim;
256 34 const float* qrow = Qp + q_head_off;
257
258
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 scores.assign(static_cast<std::size_t>(cap), 0.0f);
259 34 float run_max = -1e30f;
260
2/2
✓ Branch 0 taken 3808 times.
✓ Branch 1 taken 34 times.
 3842 for (int kg = 0; kg < cap; ++kg) {
261
4/4
✓ Branch 0 taken 1242 times.
✓ Branch 1 taken 2566 times.
✓ Branch 2 taken 36 times.
✓ Branch 3 taken 1206 times.
 3808 if (d_mask[kg] <= 0.5f || kg < lo) {
262 2602 scores[static_cast<std::size_t>(kg)] = -1e30f;
263 2602 continue;
264 }
265 1206 const float* krow = Kp + static_cast<std::size_t>(kg) * Dkv + kv_head_off;
266 1206 float dot = 0.0f;
267
2/2
✓ Branch 0 taken 27248 times.
✓ Branch 1 taken 1206 times.
 28454 for (int d = 0; d < head_dim; ++d) dot += qrow[d] * krow[d];
268 1206 float s = dot * inv_sqrt;
269
1/2
✓ Branch 0 taken 1206 times.
✗ Branch 1 not taken.
 1206 s = apply_softcap(s, attn_softcap); // Gemma-2 tanh soft-cap
270 1206 scores[static_cast<std::size_t>(kg)] = s;
271
2/2
✓ Branch 0 taken 1096 times.
✓ Branch 1 taken 110 times.
 1206 if (s > run_max) run_max = s;
272 1206 }
273 34 float sum = 0.0f;
274
2/2
✓ Branch 0 taken 3808 times.
✓ Branch 1 taken 34 times.
 3842 for (int kg = 0; kg < cap; ++kg) {
275
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 3808 times.
 3808 const float e = (run_max <= -1e29f)
276 ? 0.0f
277 3808 : std::exp(scores[static_cast<std::size_t>(kg)] - run_max);
278 3808 scores[static_cast<std::size_t>(kg)] = e;
279 3808 sum += e;
280 3808 }
281
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 const float inv = (sum > 0.0f) ? (1.0f / sum) : 0.0f;
282 34 float* orow = Op + q_head_off;
283
2/2
✓ Branch 0 taken 912 times.
✓ Branch 1 taken 34 times.
 946 for (int d = 0; d < head_dim; ++d) {
284 912 float acc = 0.0f;
285
2/2
✓ Branch 0 taken 124672 times.
✓ Branch 1 taken 912 times.
 125584 for (int kg = 0; kg < cap; ++kg) {
286 124672 const float s = scores[static_cast<std::size_t>(kg)];
287 // Skip zero-weight keys: identical math for any finite V, and
288 // masked rows may hold garbage the multiply would propagate.
289
2/2
✓ Branch 0 taken 97424 times.
✓ Branch 1 taken 27248 times.
 124672 if (s == 0.0f) continue;
290 27248 acc += s * Vp[static_cast<std::size_t>(kg) * Dkv + kv_head_off + d];
291 27248 }
292 912 orow[d] = acc * inv;
293 912 }
294 34 }
295 7 }
296
297 } // namespace brotensor::detail::cpu
298