GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 90.3% 84 / 0 / 93
Functions: 100.0% 6 / 0 / 6
Branches: 22.1% 27 / 0 / 122
src/cpu/gated_delta_rule.cpp
Line Branch Exec Source
1 // ─── CPU Gated Delta Rule ──────────────────────────────────────────────────
2 //
3 // FP32-only host implementation of the Gated DeltaNet matrix-valued recurrence
4 // used by hybrid linear-attention text decoders (the linear-attention layers
5 // alternate with standard gated attention, handled via flash_attention_decode).
6 //
7 // Per token t, per head h (FLA / HF Qwen3.5 ordering — decay applied BEFORE
8 // the delta read, so u_t is computed against the decayed state):
9 // alpha_t = exp(-softplus(a_raw_t) * exp(log_A_h)) (decay gate, in (0,1])
10 // beta_t = sigmoid(beta_raw_t) (write strength)
11 // S_pre_t = alpha_t * S_{t-1} (decayed state)
12 // u_t = S_pre_t k_t (predicted v)
13 // S_t = S_pre_t + beta_t * (v_t - u_t) k_t^T
14 // o_t = S_t q_t
15 // per-head state S has shape (d_v, d_k); o_t in R^{d_v}, q_t/k_t in R^{d_k},
16 // v_t in R^{d_v}.
17 //
18 // The chunked WY/UT-transform exists for GPU throughput; on CPU a plain
19 // sequential scan is the same complexity and clearer. gated_delta_rule_chunked
20 // and gated_delta_rule_step therefore share the same inner loop; the only
21 // reason for two ops is the GPU split (where chunked prefill matters).
22
23 #include <brotensor/tensor.h>
24
25 #include <cmath>
26 #include <stdexcept>
27 #include <string>
28
29 namespace brotensor::detail::cpu {
30
31 namespace {
32
33 112 inline void check_fp32(const ::brotensor::Tensor& t,
34 const char* op, const char* name) {
35
1/2
✓ Branch 0 taken 112 times.
✗ Branch 1 not taken.
 112 if (t.dtype != Dtype::FP32) {
36 ✗ throw std::runtime_error(std::string(op) + ": " + name +
37 " must be FP32 (CPU backend is FP32-only)");
38 }
39 112 }
40
41 // log(1 + exp(x)) — numerically stable.
42 96 inline float softplus(float x) {
43 96 return std::max(x, 0.0f) + std::log1p(std::exp(-std::abs(x)));
44 }
45
46 // 1 / (1 + exp(-x)).
47 96 inline float sigmoid(float x) {
48
2/2
✓ Branch 0 taken 47 times.
✓ Branch 1 taken 49 times.
 96 if (x >= 0.0f) {
49 47 const float e = std::exp(-x);
50 47 return 1.0f / (1.0f + e);
51 } else {
52 49 const float e = std::exp(x);
53 49 return e / (1.0f + e);
54 }
55 96 }
56
57 // Shared per-call validation and inner scan. Both chunked and step share these
58 // rules; only the op name differs in error messages.
59 16 void run_scan(const ::brotensor::Tensor& Q,
60 const ::brotensor::Tensor& K,
61 const ::brotensor::Tensor& V,
62 const ::brotensor::Tensor& a_raw,
63 const ::brotensor::Tensor& beta,
64 const ::brotensor::Tensor& log_A,
65 int num_heads, int d_k, int d_v,
66 ::brotensor::Tensor& state,
67 ::brotensor::Tensor& O,
68 const char* op) {
69 16 check_fp32(Q, op, "Q");
70 16 check_fp32(K, op, "K");
71 16 check_fp32(V, op, "V");
72 16 check_fp32(a_raw, op, "a_raw");
73 16 check_fp32(beta, op, "beta");
74 16 check_fp32(log_A, op, "log_A");
75 16 check_fp32(state, op, "state");
76
77
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (num_heads <= 0 || d_k <= 0 || d_v <= 0) {
78 ✗ throw std::runtime_error(std::string(op) +
79 ": num_heads, d_k, d_v must be positive");
80 }
81
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (Q.cols != num_heads * d_k || K.cols != num_heads * d_k) {
82 ✗ throw std::runtime_error(std::string(op) +
83 ": Q/K cols must equal num_heads * d_k");
84 }
85
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (V.cols != num_heads * d_v) {
86 ✗ throw std::runtime_error(std::string(op) +
87 ": V.cols must equal num_heads * d_v");
88 }
89
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (K.rows != Q.rows || V.rows != Q.rows) {
90 ✗ throw std::runtime_error(std::string(op) +
91 ": Q/K/V row count mismatch");
92 }
93
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (a_raw.rows != Q.rows || a_raw.cols != num_heads) {
94 ✗ throw std::runtime_error(std::string(op) +
95 ": a_raw must be (L, num_heads)");
96 }
97
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (beta.rows != Q.rows || beta.cols != num_heads) {
98 ✗ throw std::runtime_error(std::string(op) +
99 ": beta must be (L, num_heads)");
100 }
101
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (log_A.rows != num_heads || log_A.cols != 1) {
102 ✗ throw std::runtime_error(std::string(op) +
103 ": log_A must be (num_heads, 1)");
104 }
105
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (state.rows != num_heads || state.cols != d_v * d_k) {
106 ✗ throw std::runtime_error(std::string(op) +
107 ": state must be (num_heads, d_v*d_k)");
108 }
109
110 16 const int L = Q.rows;
111 16 const int Dq = Q.cols;
112 16 const int Dv = V.cols;
113
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 16 if (O.rows != L || O.cols != Dv || O.dtype != Dtype::FP32) {
114 16 O.resize(L, Dv, Dtype::FP32);
115 16 }
116
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (L == 0) return;
117
118 16 const float* Qp = Q.host_f32();
119 16 const float* Kp = K.host_f32();
120 16 const float* Vp = V.host_f32();
121 16 const float* Ap = a_raw.host_f32(); // (L, num_heads)
122 16 const float* Bp = beta.host_f32(); // (L, num_heads)
123 16 const float* logA = log_A.host_f32(); // (num_heads, 1)
124 16 float* Sp = state.host_f32_mut(); // (num_heads, d_v * d_k)
125 16 float* Op = O.host_f32_mut();
126
127 16 const int Sh_stride = d_v * d_k; // per-head state stride
128
129 // Per head independently — each head has its own S, and there is no
130 // cross-head interaction. We loop heads-outer / tokens-inner so the
131 // d_v * d_k state stays hot in cache for the whole sequence.
132
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 34 times.
 50 for (int h = 0; h < num_heads; ++h) {
133 34 float* S = Sp + h * Sh_stride; // S[v, k] = S[v*d_k + k]
134 34 const float exp_A = std::exp(logA[h]);
135
136
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 34 times.
 130 for (int t = 0; t < L; ++t) {
137 96 const float* qt = Qp + t * Dq + h * d_k;
138 96 const float* kt = Kp + t * Dq + h * d_k;
139 96 const float* vt = Vp + t * Dv + h * d_v;
140
141 96 const float a_raw_t = Ap[t * num_heads + h];
142 96 const float beta_t = sigmoid(Bp[t * num_heads + h]);
143 // alpha = exp(-softplus(a_raw) * exp(log_A)) ∈ (0, 1]
144 96 const float alpha = std::exp(-softplus(a_raw_t) * exp_A);
145
146 // u = S * k (shape d_v). S row v stride d_k.
147 // The output row (orow) doubles as scratch for delta_v between
148 // the update pass and the output pass below — no separate
149 // allocation needed.
150 //
151 // The recurrence per head per token is:
152 // 1) u_v = sum_k S[v,k] * k[k] (against decayed S)
153 // 2) delta_v = v[v] - u_v
154 // 3) S[v,k] = alpha * S[v,k] + beta * delta_v * k[k]
155 // 4) o_v = sum_k S[v,k] * q[k]
156 // computed in two passes per head per token:
157 // pass A: compute delta_v (needs u_v against decayed S), then
158 // update S in place.
159 // pass B: compute o_v from the updated S.
160
161 // FLA / HF Qwen3.5 ordering: decay S FIRST, then compute u against
162 // the decayed S, then add the delta-write. Without this ordering,
163 // brotensor's recurrence diverges from HF's `torch_recurrent_gated
164 // _delta_rule` at every token after the first by a factor that
165 // depends on alpha, producing percent-level logit drift in
166 // Qwen3.5-VL.
167 //
168 // u_v is computed against the decayed state S_decayed = alpha *
169 // S_old, i.e. u_v = dot(S_decayed[v], k) = alpha * dot(S_old[v], k).
170 // That lets us read S_old[v,:] once, derive u_v directly from it
171 // (skipping the separately-materialized decayed-state pass), and
172 // then overwrite the row in place with the final
173 // S_new[v,k] = alpha * S_old[v,k] + beta * delta_v * k[k] —
174 // safe in place since each element's write only consumes its own
175 // pre-overwrite value. This folds what was a decay-only pass plus
176 // a re-read-and-rewrite pass into a single pass (2 sweeps of the
177 // state matrix total, counting the output pass below, instead of
178 // 3).
179 96 float* orow = Op + t * Dv + h * d_v;
180
2/2
✓ Branch 0 taken 1277 times.
✓ Branch 1 taken 96 times.
 1373 for (int v = 0; v < d_v; ++v) {
181 1277 float* Sv = S + v * d_k;
182 1277 float u_v = 0.0f;
183
2/2
✓ Branch 0 taken 25640 times.
✓ Branch 1 taken 1277 times.
 26917 for (int k = 0; k < d_k; ++k) u_v += Sv[k] * kt[k];
184 1277 u_v *= alpha;
185 1277 const float delta_v = vt[v] - u_v;
186 1277 orow[v] = delta_v; // stash delta in orow
187 1277 const float scale_v = beta_t * delta_v;
188
2/2
✓ Branch 0 taken 25640 times.
✓ Branch 1 taken 1277 times.
 26917 for (int k = 0; k < d_k; ++k) {
189 25640 Sv[k] = alpha * Sv[k] + scale_v * kt[k];
190 25640 }
191 1277 }
192 // pass B — o = S @ q (overwrites the stashed delta)
193
2/2
✓ Branch 0 taken 1277 times.
✓ Branch 1 taken 96 times.
 1373 for (int v = 0; v < d_v; ++v) {
194 1277 const float* Sv = S + v * d_k;
195 1277 float o_v = 0.0f;
196
2/2
✓ Branch 0 taken 25640 times.
✓ Branch 1 taken 1277 times.
 26917 for (int k = 0; k < d_k; ++k) o_v += Sv[k] * qt[k];
197 1277 orow[v] = o_v;
198 1277 }
199 96 }
200 34 }
201 16 }
202
203 } // namespace
204
205 7 void gated_delta_rule_chunked(const ::brotensor::Tensor& Q,
206 const ::brotensor::Tensor& K,
207 const ::brotensor::Tensor& V,
208 const ::brotensor::Tensor& a_raw,
209 const ::brotensor::Tensor& beta,
210 const ::brotensor::Tensor& log_A,
211 int num_heads, int d_k, int d_v,
212 ::brotensor::Tensor& state,
213 ::brotensor::Tensor& O) {
214 14 run_scan(Q, K, V, a_raw, beta, log_A,
215 7 num_heads, d_k, d_v, state, O, "gated_delta_rule_chunked");
216 7 }
217
218 9 void gated_delta_rule_step(const ::brotensor::Tensor& Q,
219 const ::brotensor::Tensor& K,
220 const ::brotensor::Tensor& V,
221 const ::brotensor::Tensor& a_raw,
222 const ::brotensor::Tensor& beta,
223 const ::brotensor::Tensor& log_A,
224 int num_heads, int d_k, int d_v,
225 ::brotensor::Tensor& state,
226 ::brotensor::Tensor& O) {
227 18 run_scan(Q, K, V, a_raw, beta, log_A,
228 9 num_heads, d_k, d_v, state, O, "gated_delta_rule_step");
229 9 }
230
231 } // namespace brotensor::detail::cpu
232