GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 90.4% 1772 / 0 / 1961
Functions: 96.6% 253 / 0 / 262
Branches: 49.7% 329 / 0 / 662
src/ops.cpp
Line Branch Exec Source
1 // brotensor public op wrappers.
2 //
3 // One wrapper per public op declared in <brotensor/ops.h>. Each wrapper:
4 // 1. Validates that all operand tensors share a Device (via detail::dispatch
5 // overloads, which throw on mismatch).
6 // 2. Looks up the function pointer in the resolved OpsVTable.
7 // 3. If null, throws "not implemented on <device>".
8 // 4. Otherwise forwards the call.
9 //
10 // Special cases:
11 // * `mse_scalar`, `softmax_xent_segment` have no Tensor operands. They are
12 // host/CPU-only by nature; the wrapper routes directly to the CPU vtable.
13 // Backends other than CPU leave these slots null.
14 // * Ops with optional `const Tensor*` operands dispatch on the first
15 // guaranteed-non-null Tensor (typically the first positional arg).
16 // * Ops taking std::vector<const Tensor*> dispatch on parts[0]->device
17 // (or, for backward variants, on dY).
18 // * `quantize_int8_per_row_host` is a host helper over raw buffers — not in
19 // the vtable. Implemented inline here.
20
21 #include <brotensor/ops.h>
22 #include <brotensor/tensor.h>
23 #include <brotensor/detail/dispatch.h>
24 #include <brotensor/detail/cpu/thread_pool.h>
25
26 #include <algorithm>
27 #include <cmath>
28 #include <cstdint>
29 #include <stdexcept>
30 #include <string>
31 #include <vector>
32
33 namespace brotensor {
34
35 namespace {
36
37 ✗ [[noreturn]] void throw_empty_parts(const char* op) {
38 ✗ std::string m = "brotensor: ";
39 ✗ m += op;
40 ✗ m += ": parts must be non-empty";
41 ✗ throw std::runtime_error(m);
42 ✗ }
43
44 49 const detail::OpsVTable& vtable_from_parts(
45 const std::vector<const Tensor*>& parts, const char* op) {
46
1/2
✓ Branch 0 taken 49 times.
✗ Branch 1 not taken.
 49 if (parts.empty() || parts[0] == nullptr) throw_empty_parts(op);
47 49 Device d = parts[0]->device;
48
2/2
✓ Branch 0 taken 74 times.
✓ Branch 1 taken 49 times.
 123 for (std::size_t i = 1; i < parts.size(); ++i) {
49
2/4
✓ Branch 0 taken 74 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 74 times.
✗ Branch 3 not taken.
 74 if (parts[i] && parts[i]->device != d) {
50 ✗ std::string m = "brotensor: ";
51 ✗ m += op;
52 ✗ m += ": parts[";
53 ✗ m += std::to_string(i);
54 ✗ m += "] on different device";
55 ✗ throw std::runtime_error(m);
56 ✗ }
57 74 }
58 49 return detail::ops_for(d);
59 ✗ }
60
61 const detail::OpsVTable& vtable_from_parts(
62 const std::vector<Tensor*>& parts, const char* op) {
63 if (parts.empty() || parts[0] == nullptr) throw_empty_parts(op);
64 Device d = parts[0]->device;
65 for (std::size_t i = 1; i < parts.size(); ++i) {
66 if (parts[i] && parts[i]->device != d) {
67 std::string m = "brotensor: ";
68 m += op;
69 m += ": parts[";
70 m += std::to_string(i);
71 m += "] on different device";
72 throw std::runtime_error(m);
73 }
74 }
75 return detail::ops_for(d);
76 }
77
78 #define DISPATCH_REQUIRE(opname, vt) \
79 do { \
80 if (!(vt).opname) detail::throw_not_implemented(#opname, _disp_dev); \
81 } while (0)
82
83 } // namespace
84
85 // Helper to extract device for use after dispatch():
86 // const auto& ops = detail::dispatch(...);
87 // Device _disp_dev = X.device; // for DISPATCH_REQUIRE message
88 //
89 // We forgo that pattern and inline the throw with a fresh device argument
90 // where needed.
91
92 // ─── Dense layers + elementwise ────────────────────────────────────────────
93
94 968 void linear_forward(const Tensor& W, const Tensor& b,
95 const Tensor& x, Tensor& y) {
96 968 const auto& v = detail::dispatch(W, b, x, y);
97
1/2
✓ Branch 0 taken 968 times.
✗ Branch 1 not taken.
 968 if (!v.linear_forward) detail::throw_not_implemented("linear_forward", W.device);
98 968 detail::adopt_output(y, W.device);
99 968 v.linear_forward(W, b, x, y);
100 968 }
101
102 30 void linear_backward(const Tensor& W, const Tensor& x, const Tensor& dY,
103 Tensor& dX, Tensor& dW, Tensor& dB) {
104 30 const auto& v = detail::dispatch(W, x, dY, dX, dW, dB);
105
1/2
✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
 30 if (!v.linear_backward) detail::throw_not_implemented("linear_backward", W.device);
106 30 detail::adopt_output(dX, W.device);
107 30 detail::adopt_output(dW, W.device);
108 30 detail::adopt_output(dB, W.device);
109 30 v.linear_backward(W, x, dY, dX, dW, dB);
110 30 }
111
112 93 void relu_forward(const Tensor& x, Tensor& y) {
113 93 const auto& v = detail::dispatch(x, y);
114
1/2
✓ Branch 0 taken 93 times.
✗ Branch 1 not taken.
 93 if (!v.relu_forward) detail::throw_not_implemented("relu_forward", x.device);
115 93 detail::adopt_output(y, x.device);
116 93 v.relu_forward(x, y);
117 93 }
118
119 11 void relu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
120 11 const auto& v = detail::dispatch(x, dY, dX);
121
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.relu_backward) detail::throw_not_implemented("relu_backward", x.device);
122 11 detail::adopt_output(dX, x.device);
123 11 v.relu_backward(x, dY, dX);
124 11 }
125
126 97 void tanh_forward(const Tensor& x, Tensor& y) {
127 97 const auto& v = detail::dispatch(x, y);
128
1/2
✓ Branch 0 taken 97 times.
✗ Branch 1 not taken.
 97 if (!v.tanh_forward) detail::throw_not_implemented("tanh_forward", x.device);
129 97 detail::adopt_output(y, x.device);
130 97 v.tanh_forward(x, y);
131 97 }
132
133 11 void tanh_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
134 11 const auto& v = detail::dispatch(y, dY, dX);
135
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.tanh_backward) detail::throw_not_implemented("tanh_backward", y.device);
136 11 detail::adopt_output(dX, y.device);
137 11 v.tanh_backward(y, dY, dX);
138 11 }
139
140 26 void sigmoid_forward(const Tensor& x, Tensor& y) {
141 26 const auto& v = detail::dispatch(x, y);
142
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.sigmoid_forward) detail::throw_not_implemented("sigmoid_forward", x.device);
143 26 detail::adopt_output(y, x.device);
144 26 v.sigmoid_forward(x, y);
145 26 }
146
147 11 void sigmoid_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
148 11 const auto& v = detail::dispatch(y, dY, dX);
149
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.sigmoid_backward) detail::throw_not_implemented("sigmoid_backward", y.device);
150 11 detail::adopt_output(dX, y.device);
151 11 v.sigmoid_backward(y, dY, dX);
152 11 }
153
154 363 void add_inplace(Tensor& y, const Tensor& x) {
155 363 const auto& v = detail::dispatch(y, x);
156
1/2
✓ Branch 0 taken 363 times.
✗ Branch 1 not taken.
 363 if (!v.add_inplace) detail::throw_not_implemented("add_inplace", y.device);
157 363 detail::adopt_output(y, y.device);
158 363 v.add_inplace(y, x);
159 363 }
160
161 2 void axpby_inplace(Tensor& y, const Tensor& x, float a, float b) {
162 2 const auto& v = detail::dispatch(y, x);
163
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.axpby_inplace) detail::throw_not_implemented("axpby_inplace", y.device);
164 2 detail::adopt_output(y, y.device);
165 2 v.axpby_inplace(y, x, a, b);
166 2 }
167
168 15 void add_scalar_inplace(Tensor& y, float s) {
169 15 const auto& v = detail::dispatch(y);
170
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.add_scalar_inplace) detail::throw_not_implemented("add_scalar_inplace", y.device);
171 15 detail::adopt_output(y, y.device);
172 15 v.add_scalar_inplace(y, s);
173 15 }
174
175 36 void add_channel_bias_inplace(Tensor& y, const Tensor& bias, int C, int L) {
176 36 const auto& v = detail::dispatch(y);
177
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (!v.add_channel_bias_inplace)
178 ✗ detail::throw_not_implemented("add_channel_bias_inplace", y.device);
179 36 detail::adopt_output(y, y.device);
180 36 v.add_channel_bias_inplace(y, bias, C, L);
181 36 }
182
183 287 void scale_inplace(Tensor& y, float s) {
184 287 const auto& v = detail::dispatch(y);
185
1/2
✓ Branch 0 taken 287 times.
✗ Branch 1 not taken.
 287 if (!v.scale_inplace) detail::throw_not_implemented("scale_inplace", y.device);
186 287 detail::adopt_output(y, y.device);
187 287 v.scale_inplace(y, s);
188 287 }
189
190 12 void clamp(Tensor& y, float lo, float hi) {
191 12 const auto& v = detail::dispatch(y);
192
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.clamp) detail::throw_not_implemented("clamp", y.device);
193 12 detail::adopt_output(y, y.device);
194 12 v.clamp(y, lo, hi);
195 12 }
196
197 155 void mul_inplace(Tensor& y, const Tensor& x) {
198 155 const auto& v = detail::dispatch(y, x);
199
1/2
✓ Branch 0 taken 155 times.
✗ Branch 1 not taken.
 155 if (!v.mul_inplace) detail::throw_not_implemented("mul_inplace", y.device);
200 155 detail::adopt_output(y, y.device);
201 155 v.mul_inplace(y, x);
202 155 }
203
204 12 void modulate(const Tensor& X, const Tensor& scale, const Tensor& shift,
205 Tensor& Y) {
206 12 const auto& v = detail::dispatch(X, scale, shift, Y);
207
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.modulate) detail::throw_not_implemented("modulate", X.device);
208 12 detail::adopt_output(Y, X.device);
209 12 v.modulate(X, scale, shift, Y);
210 12 }
211
212 10 void broadcast_mul(const Tensor& X, const Tensor& vv, Tensor& Y) {
213 10 const auto& v = detail::dispatch(X, vv, Y);
214
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.broadcast_mul) detail::throw_not_implemented("broadcast_mul", X.device);
215 10 detail::adopt_output(Y, X.device);
216 10 v.broadcast_mul(X, vv, Y);
217 10 }
218
219 8 void build_slot_mask(const Tensor& x, int offset, int K, int stride, Tensor& mask) {
220 8 const auto& v = detail::dispatch(x, mask);
221
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.build_slot_mask) detail::throw_not_implemented("build_slot_mask", x.device);
222 8 detail::adopt_output(mask, x.device);
223 8 v.build_slot_mask(x, offset, K, stride, mask);
224 8 }
225
226 // ─── Reductions / norm / softmax / attention (training) ────────────────────
227
228 35 void softmax_forward(const Tensor& logits, Tensor& probs, const float* mask) {
229 35 const auto& v = detail::dispatch(logits, probs);
230
1/2
✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
 35 if (!v.softmax_forward) detail::throw_not_implemented("softmax_forward", logits.device);
231 35 detail::adopt_output(probs, logits.device);
232 35 v.softmax_forward(logits, probs, mask);
233 35 }
234
235 44 void softmax_rows_forward(const Tensor& X, Tensor& Y, int rows, int cols) {
236 44 const auto& v = detail::dispatch(X, Y);
237
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.softmax_rows_forward) detail::throw_not_implemented("softmax_rows_forward", X.device);
238 44 detail::adopt_output(Y, X.device);
239 44 v.softmax_rows_forward(X, Y, rows, cols);
240 44 }
241
242 21 void softmax_backward(const Tensor& probs, const Tensor& dProbs, Tensor& dLogits) {
243 21 const auto& v = detail::dispatch(probs, dProbs, dLogits);
244
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (!v.softmax_backward) detail::throw_not_implemented("softmax_backward", probs.device);
245 21 detail::adopt_output(dLogits, probs.device);
246 21 v.softmax_backward(probs, dProbs, dLogits);
247 21 }
248
249 6 void layernorm_forward(const Tensor& x, const Tensor& gamma, const Tensor& beta,
250 Tensor& y, Tensor& xhat,
251 float& mean_out, float& rstd_out, float eps) {
252 6 const auto& v = detail::dispatch(x, gamma, beta, y, xhat);
253
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.layernorm_forward) detail::throw_not_implemented("layernorm_forward", x.device);
254 6 detail::adopt_output(y, x.device);
255 6 detail::adopt_output(xhat, x.device);
256 6 v.layernorm_forward(x, gamma, beta, y, xhat, mean_out, rstd_out, eps);
257 6 }
258
259 12 void layernorm_backward(const Tensor& dY, const Tensor& xhat,
260 const Tensor& gamma, float rstd,
261 Tensor& dX, Tensor& dGamma, Tensor& dBeta) {
262 12 const auto& v = detail::dispatch(dY, xhat, gamma, dX, dGamma, dBeta);
263
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.layernorm_backward) detail::throw_not_implemented("layernorm_backward", dY.device);
264 12 detail::adopt_output(dX, dY.device);
265 12 detail::adopt_output(dGamma, dY.device);
266 12 detail::adopt_output(dBeta, dY.device);
267 12 v.layernorm_backward(dY, xhat, gamma, rstd, dX, dGamma, dBeta);
268 12 }
269
270 14 void attention_forward(const Tensor& X,
271 const Tensor& Wq, const Tensor& Wk,
272 const Tensor& Wv, const Tensor& Wo,
273 const float* d_mask,
274 Tensor& Q, Tensor& K, Tensor& V,
275 Tensor& Attn, Tensor& Y_pre_Wo, Tensor& O) {
276 14 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, Q, K, V);
277
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.attention_forward) detail::throw_not_implemented("attention_forward", X.device);
278 14 detail::adopt_output(Q, X.device);
279 14 detail::adopt_output(K, X.device);
280 14 detail::adopt_output(V, X.device);
281 14 detail::adopt_output(Attn, X.device);
282 14 detail::adopt_output(Y_pre_Wo, X.device);
283 14 detail::adopt_output(O, X.device);
284 14 v.attention_forward(X, Wq, Wk, Wv, Wo, d_mask, Q, K, V, Attn, Y_pre_Wo, O);
285 14 }
286
287 14 void attention_backward(const Tensor& dO, const Tensor& X,
288 const Tensor& Q, const Tensor& K,
289 const Tensor& V, const Tensor& Attn,
290 const Tensor& Y_pre_Wo,
291 const Tensor& Wq, const Tensor& Wk,
292 const Tensor& Wv, const Tensor& Wo,
293 const float* d_mask,
294 Tensor& dX,
295 Tensor& dWq, Tensor& dWk,
296 Tensor& dWv, Tensor& dWo) {
297 14 const auto& v = detail::dispatch(dO, X, Q, K, V, Attn, Y_pre_Wo);
298
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.attention_backward) detail::throw_not_implemented("attention_backward", dO.device);
299 14 detail::adopt_output(dX, dO.device);
300 14 detail::adopt_output(dWq, dO.device);
301 14 detail::adopt_output(dWk, dO.device);
302 14 detail::adopt_output(dWv, dO.device);
303 14 detail::adopt_output(dWo, dO.device);
304 28 v.attention_backward(dO, X, Q, K, V, Attn, Y_pre_Wo,
305 14 Wq, Wk, Wv, Wo, d_mask,
306 14 dX, dWq, dWk, dWv, dWo);
307 14 }
308
309 26 void mha_forward(const Tensor& X,
310 const Tensor& Wq, const Tensor& Wk,
311 const Tensor& Wv, const Tensor& Wo,
312 const Tensor* bq, const Tensor* bk,
313 const Tensor* bv, const Tensor* bo,
314 const float* d_mask, int num_heads,
315 Tensor& Qh, Tensor& Kh, Tensor& Vh,
316 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
317 26 const auto& v = detail::dispatch_with_opts(
318 26 X, Wq, {&Wk, &Wv, &Wo, bq, bk, bv, bo, &Qh, &Kh, &Vh});
319
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.mha_forward) detail::throw_not_implemented("mha_forward", X.device);
320 26 detail::adopt_output(Qh, X.device);
321 26 detail::adopt_output(Kh, X.device);
322 26 detail::adopt_output(Vh, X.device);
323 26 detail::adopt_output(Attnh, X.device);
324 26 detail::adopt_output(Yconcat, X.device);
325 26 detail::adopt_output(O, X.device);
326 52 v.mha_forward(X, Wq, Wk, Wv, Wo, bq, bk, bv, bo, d_mask, num_heads,
327 26 Qh, Kh, Vh, Attnh, Yconcat, O);
328 26 }
329
330 24 void mha_backward(const Tensor& dO, const Tensor& X,
331 const Tensor& Qh, const Tensor& Kh,
332 const Tensor& Vh, const Tensor& Attnh,
333 const Tensor& Yconcat,
334 const Tensor& Wq, const Tensor& Wk,
335 const Tensor& Wv, const Tensor& Wo,
336 const float* d_mask, int num_heads,
337 Tensor& dX,
338 Tensor& dWq, Tensor& dWk,
339 Tensor& dWv, Tensor& dWo,
340 Tensor* dbq, Tensor* dbk,
341 Tensor* dbv, Tensor* dbo) {
342 24 const auto& v = detail::dispatch_with_opts(
343 24 dO, X, {&Qh, &Kh, &Vh, &Attnh, &Yconcat, dbq, dbk, dbv, dbo});
344
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.mha_backward) detail::throw_not_implemented("mha_backward", dO.device);
345 24 detail::adopt_output(dX, dO.device);
346 24 detail::adopt_output(dWq, dO.device);
347 24 detail::adopt_output(dWk, dO.device);
348 24 detail::adopt_output(dWv, dO.device);
349 24 detail::adopt_output(dWo, dO.device);
350
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbq) detail::adopt_output(*dbq, dO.device);
351
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbk) detail::adopt_output(*dbk, dO.device);
352
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbv) detail::adopt_output(*dbv, dO.device);
353
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbo) detail::adopt_output(*dbo, dO.device);
354 48 v.mha_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat,
355 24 Wq, Wk, Wv, Wo, d_mask, num_heads,
356 24 dX, dWq, dWk, dWv, dWo,
357 24 dbq, dbk, dbv, dbo);
358 24 }
359
360 // ─── Pooling / losses / embedding / concat ─────────────────────────────────
361
362 18 void masked_mean_pool_forward(const Tensor& X, const float* d_mask, Tensor& y) {
363 18 const auto& v = detail::dispatch(X, y);
364
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.masked_mean_pool_forward) detail::throw_not_implemented("masked_mean_pool_forward", X.device);
365 18 detail::adopt_output(y, X.device);
366 18 v.masked_mean_pool_forward(X, d_mask, y);
367 18 }
368
369 14 void masked_mean_pool_backward(const Tensor& dY, const float* d_mask,
370 int K, Tensor& dX) {
371 14 const auto& v = detail::dispatch(dY, dX);
372
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.masked_mean_pool_backward) detail::throw_not_implemented("masked_mean_pool_backward", dY.device);
373 14 detail::adopt_output(dX, dY.device);
374 14 v.masked_mean_pool_backward(dY, d_mask, K, dX);
375 14 }
376
377 9 float mse_vec_forward(const Tensor& pred, const Tensor& target) {
378 9 const auto& v = detail::dispatch(pred, target);
379
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.mse_vec_forward) detail::throw_not_implemented("mse_vec_forward", pred.device);
380 9 return v.mse_vec_forward(pred, target);
381 }
382
383 9 void mse_vec_backward(const Tensor& pred, const Tensor& target, Tensor& dPred) {
384 9 const auto& v = detail::dispatch(pred, target, dPred);
385
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.mse_vec_backward) detail::throw_not_implemented("mse_vec_backward", pred.device);
386 9 detail::adopt_output(dPred, pred.device);
387 9 v.mse_vec_backward(pred, target, dPred);
388 9 }
389
390 // CPU-only host helper. No tensors involved — routes through the CPU vtable.
391 3 float mse_scalar(float pred, float target, float& dPred) {
392 3 const auto& v = detail::ops_for(Device::CPU);
393
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.mse_scalar) detail::throw_not_implemented("mse_scalar", Device::CPU);
394 3 return v.mse_scalar(pred, target, dPred);
395 }
396
397 3 float softmax_xent(const Tensor& logits, const Tensor& target,
398 Tensor& probs, Tensor& dLogits, const float* mask) {
399 3 const auto& v = detail::dispatch(logits, target, probs, dLogits);
400
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.softmax_xent) detail::throw_not_implemented("softmax_xent", logits.device);
401 3 detail::adopt_output(probs, logits.device);
402 3 detail::adopt_output(dLogits, logits.device);
403 3 return v.softmax_xent(logits, target, probs, dLogits, mask);
404 }
405
406 // CPU-only host helper. Raw pointers — CPU vtable directly.
407 7 float softmax_xent_segment(const float* logits, const float* target,
408 float* probs, float* dLogits,
409 int n, const float* mask) {
410 7 const auto& v = detail::ops_for(Device::CPU);
411
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.softmax_xent_segment) detail::throw_not_implemented("softmax_xent_segment", Device::CPU);
412 7 return v.softmax_xent_segment(logits, target, probs, dLogits, n, mask);
413 }
414
415 16 float softmax_xent_fused(const Tensor& logits, const Tensor& target,
416 const float* d_mask,
417 Tensor& probs, Tensor& dLogits) {
418 16 const auto& v = detail::dispatch(logits, target, probs, dLogits);
419
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.softmax_xent_fused) detail::throw_not_implemented("softmax_xent_fused", logits.device);
420 16 detail::adopt_output(probs, logits.device);
421 16 detail::adopt_output(dLogits, logits.device);
422 16 return v.softmax_xent_fused(logits, target, d_mask, probs, dLogits);
423 }
424
425 24 void embedding_lookup_forward(const Tensor& table, const int32_t* d_idx,
426 int B, Tensor& out) {
427 24 const auto& v = detail::dispatch(table, out);
428
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.embedding_lookup_forward) detail::throw_not_implemented("embedding_lookup_forward", table.device);
429 24 detail::adopt_output(out, table.device);
430 24 v.embedding_lookup_forward(table, d_idx, B, out);
431 24 }
432
433 25 void embedding_lookup_backward(const Tensor& dOut, const int32_t* d_idx,
434 int B, Tensor& dTable) {
435 25 const auto& v = detail::dispatch(dOut, dTable);
436
1/2
✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
 25 if (!v.embedding_lookup_backward) detail::throw_not_implemented("embedding_lookup_backward", dOut.device);
437 25 detail::adopt_output(dTable, dOut.device);
438 25 v.embedding_lookup_backward(dOut, d_idx, B, dTable);
439 25 }
440
441 9 void concat_rows(const std::vector<const Tensor*>& parts, Tensor& out) {
442 9 const auto& v = vtable_from_parts(parts, "concat_rows");
443
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 9 if (out.data != nullptr && parts[0]->device != out.device) {
444 ✗ throw std::runtime_error("brotensor: concat_rows: out on different device");
445 }
446 9 detail::adopt_output(out, parts[0]->device);
447
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.concat_rows) detail::throw_not_implemented("concat_rows", parts[0]->device);
448 9 v.concat_rows(parts, out);
449 9 }
450
451 9 void split_rows(const Tensor& in, const std::vector<Tensor*>& parts) {
452 9 Device d = in.device;
453
2/2
✓ Branch 0 taken 33 times.
✓ Branch 1 taken 9 times.
 42 for (std::size_t i = 0; i < parts.size(); ++i) {
454
3/6
✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 33 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 33 times.
✗ Branch 5 not taken.
 33 if (parts[i] && parts[i]->data != nullptr && parts[i]->device != d) {
455 ✗ throw std::runtime_error("brotensor: split_rows: parts on different device");
456 }
457 33 }
458 9 const auto& v = detail::ops_for(d);
459
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.split_rows) detail::throw_not_implemented("split_rows", d);
460
3/4
✓ Branch 0 taken 33 times.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 33 times.
 42 for (Tensor* _p : parts) if (_p) detail::adopt_output(*_p, d);
461 9 v.split_rows(in, parts);
462 9 }
463
464 16 void concat_batched_rows(const std::vector<const Tensor*>& parts, Tensor& out) {
465 16 const auto& v = vtable_from_parts(parts, "concat_batched_rows");
466
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 16 if (out.data != nullptr && parts[0]->device != out.device) {
467 ✗ throw std::runtime_error("brotensor: concat_batched_rows: out on different device");
468 }
469 16 detail::adopt_output(out, parts[0]->device);
470
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.concat_batched_rows) detail::throw_not_implemented("concat_batched_rows", parts[0]->device);
471 16 v.concat_batched_rows(parts, out);
472 16 }
473
474 24 void concat_nchw_channels(const std::vector<const Tensor*>& parts,
475 int N, int H, int W,
476 const std::vector<int>& C_per_part,
477 Tensor& out) {
478 24 const auto& v = vtable_from_parts(parts, "concat_nchw_channels");
479
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 24 if (out.data != nullptr && parts[0]->device != out.device) {
480 ✗ throw std::runtime_error("brotensor: concat_nchw_channels: out on different device");
481 }
482 24 detail::adopt_output(out, parts[0]->device);
483
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.concat_nchw_channels) detail::throw_not_implemented("concat_nchw_channels", parts[0]->device);
484 24 v.concat_nchw_channels(parts, N, H, W, C_per_part, out);
485 24 }
486
487 20 void concat_nchw_channels_backward(const Tensor& dY, int N, int H, int W,
488 const std::vector<int>& C_per_part,
489 const std::vector<Tensor*>& parts) {
490 20 Device d = dY.device;
491
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
 60 for (std::size_t i = 0; i < parts.size(); ++i) {
492
2/6
✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 40 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 40 if (parts[i] && parts[i]->data != nullptr && parts[i]->device != d) {
493 ✗ throw std::runtime_error("brotensor: concat_nchw_channels_backward: parts on different device");
494 }
495 40 }
496 20 const auto& v = detail::ops_for(d);
497
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 if (!v.concat_nchw_channels_backward) detail::throw_not_implemented("concat_nchw_channels_backward", d);
498
3/4
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 40 times.
 60 for (Tensor* _p : parts) if (_p) detail::adopt_output(*_p, d);
499 20 v.concat_nchw_channels_backward(dY, N, H, W, C_per_part, parts);
500 20 }
501
502 8 void copy_d2d(const Tensor& src, int src_off, Tensor& dst, int dst_off, int n) {
503 8 const auto& v = detail::dispatch(src, dst);
504
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.copy_d2d) detail::throw_not_implemented("copy_d2d", src.device);
505 8 detail::adopt_output(dst, src.device);
506 8 v.copy_d2d(src, src_off, dst, dst_off, n);
507 8 }
508
509 8 void copy_d2d_strided(const Tensor& src, int src_off, int src_pitch,
510 Tensor& dst, int dst_off, int dst_pitch,
511 int width, int height) {
512
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (width > src_pitch || width > dst_pitch) {
513 ✗ throw std::runtime_error("brotensor: copy_d2d_strided: width exceeds pitch");
514 }
515 8 const auto& v = detail::dispatch(src, dst);
516
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.copy_d2d_strided) detail::throw_not_implemented("copy_d2d_strided", src.device);
517 8 detail::adopt_output(dst, src.device);
518 8 v.copy_d2d_strided(src, src_off, src_pitch, dst, dst_off, dst_pitch, width, height);
519 8 }
520
521 24 void cast(const Tensor& src, Tensor& dst, Dtype out_dtype) {
522 24 const auto& v = detail::dispatch(src, dst);
523
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.cast) detail::throw_not_implemented("cast", src.device);
524 24 detail::adopt_output(dst, src.device);
525 24 v.cast(src, dst, out_dtype);
526 24 }
527
528 // ─── Inference batched + optim ─────────────────────────────────────────────
529
530 26 void layernorm_forward_inference_batched(const Tensor& X_RD,
531 const Tensor& gamma,
532 const Tensor& beta,
533 Tensor& Y_RD, float eps) {
534 26 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD);
535
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.layernorm_forward_inference_batched)
536 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched", X_RD.device);
537 26 detail::adopt_output(Y_RD, X_RD.device);
538 26 v.layernorm_forward_inference_batched(X_RD, gamma, beta, Y_RD, eps);
539 26 }
540
541 16 void layernorm_forward_batched_with_caches(const Tensor& X_RD,
542 const Tensor& gamma,
543 const Tensor& beta,
544 Tensor& Y_RD, Tensor& Xhat_RD,
545 Tensor& Mean_R, Tensor& Rstd_R,
546 float eps) {
547 16 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD, Xhat_RD, Mean_R, Rstd_R);
548
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.layernorm_forward_batched_with_caches)
549 ✗ detail::throw_not_implemented("layernorm_forward_batched_with_caches", X_RD.device);
550 16 detail::adopt_output(Y_RD, X_RD.device);
551 16 detail::adopt_output(Xhat_RD, X_RD.device);
552 16 detail::adopt_output(Mean_R, X_RD.device);
553 16 detail::adopt_output(Rstd_R, X_RD.device);
554 32 v.layernorm_forward_batched_with_caches(X_RD, gamma, beta, Y_RD, Xhat_RD,
555 16 Mean_R, Rstd_R, eps);
556 16 }
557
558 16 void layernorm_backward_batched_with_caches(const Tensor& dY_RD,
559 const Tensor& Xhat_RD,
560 const Tensor& gamma,
561 const Tensor& Rstd_R,
562 Tensor& dX_RD,
563 Tensor& dGamma, Tensor& dBeta) {
564 16 const auto& v = detail::dispatch(dY_RD, Xhat_RD, gamma, Rstd_R, dX_RD, dGamma, dBeta);
565
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.layernorm_backward_batched_with_caches)
566 ✗ detail::throw_not_implemented("layernorm_backward_batched_with_caches", dY_RD.device);
567 16 detail::adopt_output(dX_RD, dY_RD.device);
568 16 detail::adopt_output(dGamma, dY_RD.device);
569 16 detail::adopt_output(dBeta, dY_RD.device);
570 32 v.layernorm_backward_batched_with_caches(dY_RD, Xhat_RD, gamma, Rstd_R,
571 16 dX_RD, dGamma, dBeta);
572 16 }
573
574 5 void sgd_step(Tensor& param, Tensor& grad, Tensor& velocity,
575 float lr, float momentum) {
576 5 const auto& v = detail::dispatch(param, grad, velocity);
577
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.sgd_step) detail::throw_not_implemented("sgd_step", param.device);
578 5 detail::adopt_output(param, param.device);
579 5 detail::adopt_output(grad, param.device);
580 5 detail::adopt_output(velocity, param.device);
581 5 v.sgd_step(param, grad, velocity, lr, momentum);
582 5 }
583
584 50 void adam_step(Tensor& param, const Tensor& grad,
585 Tensor& m, Tensor& v_buf,
586 float lr, float beta1, float beta2, float eps, int step) {
587 50 const auto& v = detail::dispatch(param, grad, m, v_buf);
588
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 if (!v.adam_step) detail::throw_not_implemented("adam_step", param.device);
589 50 detail::adopt_output(param, param.device);
590 50 detail::adopt_output(m, param.device);
591 50 detail::adopt_output(v_buf, param.device);
592 50 v.adam_step(param, grad, m, v_buf, lr, beta1, beta2, eps, step);
593 50 }
594
595 3 void xavier_init(Tensor& W, uint64_t& rng_state) {
596 3 const auto& v = detail::dispatch(W);
597
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.xavier_init) detail::throw_not_implemented("xavier_init", W.device);
598 3 detail::adopt_output(W, W.device);
599 3 v.xavier_init(W, rng_state);
600 3 }
601
602 // ─── Batched inference variants ────────────────────────────────────────────
603
604 22 void linear_forward_batched(const Tensor& W, const Tensor& bias,
605 const Tensor& X_BD, Tensor& Y_BD) {
606 22 const auto& v = detail::dispatch(W, bias, X_BD, Y_BD);
607
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 22 if (!v.linear_forward_batched) detail::throw_not_implemented("linear_forward_batched", W.device);
608 22 detail::adopt_output(Y_BD, W.device);
609 22 v.linear_forward_batched(W, bias, X_BD, Y_BD);
610 22 }
611
612 3 void relu_forward_batched(const Tensor& X_BD, Tensor& Y_BD) {
613 3 const auto& v = detail::dispatch(X_BD, Y_BD);
614
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.relu_forward_batched) detail::throw_not_implemented("relu_forward_batched", X_BD.device);
615 3 detail::adopt_output(Y_BD, X_BD.device);
616 3 v.relu_forward_batched(X_BD, Y_BD);
617 3 }
618
619 9 void tanh_forward_batched(const Tensor& X_BD, Tensor& Y_BD) {
620 9 const auto& v = detail::dispatch(X_BD, Y_BD);
621
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.tanh_forward_batched) detail::throw_not_implemented("tanh_forward_batched", X_BD.device);
622 9 detail::adopt_output(Y_BD, X_BD.device);
623 9 v.tanh_forward_batched(X_BD, Y_BD);
624 9 }
625
626 3 void add_inplace_batched(Tensor& Y_BD, const Tensor& X_BD) {
627 3 const auto& v = detail::dispatch(Y_BD, X_BD);
628
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.add_inplace_batched) detail::throw_not_implemented("add_inplace_batched", Y_BD.device);
629 3 detail::adopt_output(Y_BD, Y_BD.device);
630 3 v.add_inplace_batched(Y_BD, X_BD);
631 3 }
632
633 // ─── Batched backward variants ─────────────────────────────────────────────
634
635 9 void linear_backward_batched(const Tensor& W, const Tensor& X_BD,
636 const Tensor& dY_BD,
637 Tensor& dX_BD, Tensor& dW, Tensor& dB) {
638 9 const auto& v = detail::dispatch(W, X_BD, dY_BD, dX_BD, dW, dB);
639
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.linear_backward_batched) detail::throw_not_implemented("linear_backward_batched", W.device);
640 9 detail::adopt_output(dX_BD, W.device);
641 9 detail::adopt_output(dW, W.device);
642 9 detail::adopt_output(dB, W.device);
643 9 v.linear_backward_batched(W, X_BD, dY_BD, dX_BD, dW, dB);
644 9 }
645
646 6 void relu_backward_batched(const Tensor& X_BD, const Tensor& dY_BD, Tensor& dX_BD) {
647 6 const auto& v = detail::dispatch(X_BD, dY_BD, dX_BD);
648
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.relu_backward_batched) detail::throw_not_implemented("relu_backward_batched", X_BD.device);
649 6 detail::adopt_output(dX_BD, X_BD.device);
650 6 v.relu_backward_batched(X_BD, dY_BD, dX_BD);
651 6 }
652
653 6 void tanh_backward_batched(const Tensor& Y_BD, const Tensor& dY_BD, Tensor& dX_BD) {
654 6 const auto& v = detail::dispatch(Y_BD, dY_BD, dX_BD);
655
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.tanh_backward_batched) detail::throw_not_implemented("tanh_backward_batched", Y_BD.device);
656 6 detail::adopt_output(dX_BD, Y_BD.device);
657 6 v.tanh_backward_batched(Y_BD, dY_BD, dX_BD);
658 6 }
659
660 // ─── Batched per-sample losses ─────────────────────────────────────────────
661
662 6 void mse_vec_per_sample(const Tensor& pred, const Tensor& target,
663 Tensor& dPred, Tensor& loss_per_sample) {
664 6 const auto& v = detail::dispatch(pred, target, dPred, loss_per_sample);
665
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.mse_vec_per_sample) detail::throw_not_implemented("mse_vec_per_sample", pred.device);
666 6 detail::adopt_output(dPred, pred.device);
667 6 detail::adopt_output(loss_per_sample, pred.device);
668 6 v.mse_vec_per_sample(pred, target, dPred, loss_per_sample);
669 6 }
670
671 12 void softmax_xent_fused_batched(const Tensor& logits_BL, const Tensor& target_BL,
672 const float* d_mask_BL,
673 const int* d_head_offsets, int n_heads,
674 Tensor& probs_BL, Tensor& dLogits_BL,
675 Tensor& loss_per_sample) {
676 12 const auto& v = detail::dispatch(logits_BL, target_BL, probs_BL, dLogits_BL, loss_per_sample);
677
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.softmax_xent_fused_batched)
678 ✗ detail::throw_not_implemented("softmax_xent_fused_batched", logits_BL.device);
679 12 detail::adopt_output(probs_BL, logits_BL.device);
680 12 detail::adopt_output(dLogits_BL, logits_BL.device);
681 12 detail::adopt_output(loss_per_sample, logits_BL.device);
682 24 v.softmax_xent_fused_batched(logits_BL, target_BL, d_mask_BL,
683 12 d_head_offsets, n_heads,
684 12 probs_BL, dLogits_BL, loss_per_sample);
685 12 }
686
687 11 void bce_with_logits_fused_batched(const Tensor& logits_BL, const Tensor& target_BL,
688 const float* d_mask_BL,
689 float pos_weight,
690 Tensor& probs_BL, Tensor& dLogits_BL,
691 Tensor& loss_per_sample) {
692 11 const auto& v = detail::dispatch(logits_BL, target_BL, probs_BL, dLogits_BL, loss_per_sample);
693
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.bce_with_logits_fused_batched)
694 ✗ detail::throw_not_implemented("bce_with_logits_fused_batched", logits_BL.device);
695 11 detail::adopt_output(probs_BL, logits_BL.device);
696 11 detail::adopt_output(dLogits_BL, logits_BL.device);
697 11 detail::adopt_output(loss_per_sample, logits_BL.device);
698 22 v.bce_with_logits_fused_batched(logits_BL, target_BL, d_mask_BL,
699 11 pos_weight,
700 11 probs_BL, dLogits_BL, loss_per_sample);
701 11 }
702
703 // ─── Conv2d ────────────────────────────────────────────────────────────────
704
705 44 void conv2d_forward(const Tensor& X, const Tensor& Wt, const Tensor* bias,
706 int N, int C_in, int H, int W,
707 int C_out, int kH, int kW,
708 int stride_h, int stride_w,
709 int pad_h, int pad_w,
710 int dil_h, int dil_w,
711 int groups, Tensor& Y) {
712 44 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
713
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.conv2d_forward) detail::throw_not_implemented("conv2d_forward", X.device);
714 44 detail::adopt_output(Y, X.device);
715 88 v.conv2d_forward(X, Wt, bias, N, C_in, H, W, C_out, kH, kW,
716 44 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, Y);
717 44 }
718
719 28 void conv2d_backward_input(const Tensor& Wt, const Tensor& dY,
720 int N, int C_in, int H, int W,
721 int C_out, int kH, int kW,
722 int stride_h, int stride_w,
723 int pad_h, int pad_w,
724 int dil_h, int dil_w,
725 int groups, Tensor& dX) {
726 28 const auto& v = detail::dispatch(Wt, dY, dX);
727
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 if (!v.conv2d_backward_input) detail::throw_not_implemented("conv2d_backward_input", Wt.device);
728 28 detail::adopt_output(dX, Wt.device);
729 56 v.conv2d_backward_input(Wt, dY, N, C_in, H, W, C_out, kH, kW,
730 28 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, dX);
731 28 }
732
733 26 void conv2d_backward_weight(const Tensor& X, const Tensor& dY,
734 int N, int C_in, int H, int W,
735 int C_out, int kH, int kW,
736 int stride_h, int stride_w,
737 int pad_h, int pad_w,
738 int dil_h, int dil_w,
739 int groups, Tensor& dWt) {
740 26 const auto& v = detail::dispatch(X, dY, dWt);
741
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.conv2d_backward_weight) detail::throw_not_implemented("conv2d_backward_weight", X.device);
742 26 detail::adopt_output(dWt, X.device);
743 52 v.conv2d_backward_weight(X, dY, N, C_in, H, W, C_out, kH, kW,
744 26 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, dWt);
745 26 }
746
747 16 void conv2d_backward_bias(const Tensor& dY,
748 int N, int C_out, int H_out, int W_out,
749 Tensor& dB) {
750 16 const auto& v = detail::dispatch(dY, dB);
751
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.conv2d_backward_bias) detail::throw_not_implemented("conv2d_backward_bias", dY.device);
752 16 detail::adopt_output(dB, dY.device);
753 16 v.conv2d_backward_bias(dY, N, C_out, H_out, W_out, dB);
754 16 }
755
756 12 void deform_conv2d_forward(const Tensor& X, const Tensor& offset,
757 const Tensor* mask, const Tensor& Wt,
758 const Tensor* bias,
759 int N, int C_in, int H, int W,
760 int C_out, int kH, int kW,
761 int stride_h, int stride_w,
762 int pad_h, int pad_w,
763 int dil_h, int dil_w,
764 int groups, int deform_groups, Tensor& Y) {
765 12 const auto& v = detail::dispatch_with_opts(X, Wt, {&offset, mask, bias, &Y});
766
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.deform_conv2d_forward) detail::throw_not_implemented("deform_conv2d_forward", X.device);
767 12 detail::adopt_output(Y, X.device);
768 24 v.deform_conv2d_forward(X, offset, mask, Wt, bias, N, C_in, H, W,
769 12 C_out, kH, kW, stride_h, stride_w, pad_h, pad_w,
770 12 dil_h, dil_w, groups, deform_groups, Y);
771 12 }
772
773 // ─── Conv3d ────────────────────────────────────────────────────────────────
774
775 37 void conv3d_forward(const Tensor& X, const Tensor& Wt, const Tensor* bias,
776 int N, int C_in, int T, int H, int W,
777 int C_out, int kT, int kH, int kW,
778 int stride_t, int stride_h, int stride_w,
779 int pad_t, int pad_h, int pad_w,
780 int dil_t, int dil_h, int dil_w,
781 int groups, Tensor& Y) {
782 37 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
783
1/2
✓ Branch 0 taken 37 times.
✗ Branch 1 not taken.
 37 if (!v.conv3d_forward) detail::throw_not_implemented("conv3d_forward", X.device);
784 37 detail::adopt_output(Y, X.device);
785 74 v.conv3d_forward(X, Wt, bias, N, C_in, T, H, W, C_out, kT, kH, kW,
786 37 stride_t, stride_h, stride_w,
787 37 pad_t, pad_h, pad_w,
788 37 dil_t, dil_h, dil_w, groups, Y);
789 37 }
790
791 6 void conv3d_int8w_fp16_forward(const Tensor& X, const Tensor& W_int8,
792 const Tensor& scales, const Tensor* bias,
793 int N, int C_in, int T, int H, int W,
794 int C_out, int kT, int kH, int kW,
795 int stride_t, int stride_h, int stride_w,
796 int pad_t, int pad_h, int pad_w,
797 int dil_t, int dil_h, int dil_w,
798 int groups, Tensor& Y) {
799 6 const auto& v = detail::dispatch_with_opts(X, W_int8, {&scales, bias, &Y});
800
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
 6 if (!v.conv3d_int8w_fp16_forward)
801 1 detail::throw_not_implemented("conv3d_int8w_fp16_forward", X.device);
802 5 detail::adopt_output(Y, X.device);
803 10 v.conv3d_int8w_fp16_forward(X, W_int8, scales, bias,
804 5 N, C_in, T, H, W, C_out, kT, kH, kW,
805 5 stride_t, stride_h, stride_w,
806 5 pad_t, pad_h, pad_w,
807 5 dil_t, dil_h, dil_w, groups, Y);
808 5 }
809
810 // ─── GroupNorm ─────────────────────────────────────────────────────────────
811
812 18 void group_norm_forward(const Tensor& X, const Tensor& gamma, const Tensor& beta,
813 int N, int C, int H, int W, int num_groups,
814 float eps, Tensor& Y) {
815 18 const auto& v = detail::dispatch(X, gamma, beta, Y);
816
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.group_norm_forward) detail::throw_not_implemented("group_norm_forward", X.device);
817 18 detail::adopt_output(Y, X.device);
818 18 v.group_norm_forward(X, gamma, beta, N, C, H, W, num_groups, eps, Y);
819 18 }
820
821 18 void group_norm_backward(const Tensor& X, const Tensor& gamma, const Tensor& dY,
822 int N, int C, int H, int W, int num_groups, float eps,
823 Tensor& dX, Tensor& dGamma, Tensor& dBeta) {
824 18 const auto& v = detail::dispatch(X, gamma, dY, dX, dGamma, dBeta);
825
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.group_norm_backward) detail::throw_not_implemented("group_norm_backward", X.device);
826 18 detail::adopt_output(dX, X.device);
827 18 detail::adopt_output(dGamma, X.device);
828 18 detail::adopt_output(dBeta, X.device);
829 36 v.group_norm_backward(X, gamma, dY, N, C, H, W, num_groups, eps,
830 18 dX, dGamma, dBeta);
831 18 }
832
833 // ─── Activations ───────────────────────────────────────────────────────────
834
835 10 void silu_forward(const Tensor& x, Tensor& y) {
836 10 const auto& v = detail::dispatch(x, y);
837
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.silu_forward) detail::throw_not_implemented("silu_forward", x.device);
838 10 detail::adopt_output(y, x.device);
839 10 v.silu_forward(x, y);
840 10 }
841 6 void silu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
842 6 const auto& v = detail::dispatch(x, dY, dX);
843
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.silu_backward) detail::throw_not_implemented("silu_backward", x.device);
844 6 detail::adopt_output(dX, x.device);
845 6 v.silu_backward(x, dY, dX);
846 6 }
847 10 void gelu_forward(const Tensor& x, Tensor& y) {
848 10 const auto& v = detail::dispatch(x, y);
849
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.gelu_forward) detail::throw_not_implemented("gelu_forward", x.device);
850 10 detail::adopt_output(y, x.device);
851 10 v.gelu_forward(x, y);
852 10 }
853 6 void gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
854 6 const auto& v = detail::dispatch(x, dY, dX);
855
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.gelu_backward) detail::throw_not_implemented("gelu_backward", x.device);
856 6 detail::adopt_output(dX, x.device);
857 6 v.gelu_backward(x, dY, dX);
858 6 }
859 8 void gelu_exact_forward(const Tensor& x, Tensor& y) {
860 8 const auto& v = detail::dispatch(x, y);
861
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.gelu_exact_forward) detail::throw_not_implemented("gelu_exact_forward", x.device);
862 8 detail::adopt_output(y, x.device);
863 8 v.gelu_exact_forward(x, y);
864 8 }
865 6 void gelu_exact_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
866 6 const auto& v = detail::dispatch(x, dY, dX);
867
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.gelu_exact_backward) detail::throw_not_implemented("gelu_exact_backward", x.device);
868 6 detail::adopt_output(dX, x.device);
869 6 v.gelu_exact_backward(x, dY, dX);
870 6 }
871 8 void quick_gelu_forward(const Tensor& x, Tensor& y) {
872 8 const auto& v = detail::dispatch(x, y);
873
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.quick_gelu_forward) detail::throw_not_implemented("quick_gelu_forward", x.device);
874 8 detail::adopt_output(y, x.device);
875 8 v.quick_gelu_forward(x, y);
876 8 }
877 6 void quick_gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
878 6 const auto& v = detail::dispatch(x, dY, dX);
879
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.quick_gelu_backward) detail::throw_not_implemented("quick_gelu_backward", x.device);
880 6 detail::adopt_output(dX, x.device);
881 6 v.quick_gelu_backward(x, dY, dX);
882 6 }
883
884 // ─── Resample ──────────────────────────────────────────────────────────────
885
886 15 void upsample_nearest_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
887 15 const auto& v = detail::dispatch(X, Y);
888
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.upsample_nearest_2x) detail::throw_not_implemented("upsample_nearest_2x", X.device);
889 15 detail::adopt_output(Y, X.device);
890 15 v.upsample_nearest_2x(X, N, C, H, W, Y);
891 15 }
892 15 void upsample_bilinear_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
893 15 const auto& v = detail::dispatch(X, Y);
894
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.upsample_bilinear_2x) detail::throw_not_implemented("upsample_bilinear_2x", X.device);
895 15 detail::adopt_output(Y, X.device);
896 15 v.upsample_bilinear_2x(X, N, C, H, W, Y);
897 15 }
898 14 void downsample_avg_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
899 14 const auto& v = detail::dispatch(X, Y);
900
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.downsample_avg_2x) detail::throw_not_implemented("downsample_avg_2x", X.device);
901 14 detail::adopt_output(Y, X.device);
902 14 v.downsample_avg_2x(X, N, C, H, W, Y);
903 14 }
904 12 void upsample_nearest_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
905 12 const auto& v = detail::dispatch(dY, dX);
906
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.upsample_nearest_2x_backward) detail::throw_not_implemented("upsample_nearest_2x_backward", dY.device);
907 12 detail::adopt_output(dX, dY.device);
908 12 v.upsample_nearest_2x_backward(dY, N, C, H, W, dX);
909 12 }
910 12 void upsample_bilinear_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
911 12 const auto& v = detail::dispatch(dY, dX);
912
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.upsample_bilinear_2x_backward) detail::throw_not_implemented("upsample_bilinear_2x_backward", dY.device);
913 12 detail::adopt_output(dX, dY.device);
914 12 v.upsample_bilinear_2x_backward(dY, N, C, H, W, dX);
915 12 }
916 12 void downsample_avg_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
917 12 const auto& v = detail::dispatch(dY, dX);
918
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.downsample_avg_2x_backward) detail::throw_not_implemented("downsample_avg_2x_backward", dY.device);
919 12 detail::adopt_output(dX, dY.device);
920 12 v.downsample_avg_2x_backward(dY, N, C, H, W, dX);
921 12 }
922
923 76 void interp2d_forward(const Tensor& X,
924 int N, int C, int H_in, int W_in, int H_out, int W_out,
925 int mode, Tensor& Y) {
926 76 const auto& v = detail::dispatch(X, Y);
927
1/2
✓ Branch 0 taken 76 times.
✗ Branch 1 not taken.
 76 if (!v.interp2d_forward) detail::throw_not_implemented("interp2d_forward", X.device);
928 76 detail::adopt_output(Y, X.device);
929 76 v.interp2d_forward(X, N, C, H_in, W_in, H_out, W_out, mode, Y);
930 76 }
931 13 void interp2d_backward(const Tensor& dY,
932 int N, int C, int H_in, int W_in, int H_out, int W_out,
933 int mode, Tensor& dX) {
934 13 const auto& v = detail::dispatch(dY, dX);
935
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (!v.interp2d_backward) detail::throw_not_implemented("interp2d_backward", dY.device);
936 13 detail::adopt_output(dX, dY.device);
937 13 v.interp2d_backward(dY, N, C, H_in, W_in, H_out, W_out, mode, dX);
938 13 }
939 18 void interp2d_align_corners_forward(const Tensor& X,
940 int N, int C, int H_in, int W_in,
941 int H_out, int W_out, int mode, Tensor& Y) {
942 18 const auto& v = detail::dispatch(X, Y);
943
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.interp2d_align_corners_forward)
944 ✗ detail::throw_not_implemented("interp2d_align_corners_forward", X.device);
945 18 detail::adopt_output(Y, X.device);
946 18 v.interp2d_align_corners_forward(X, N, C, H_in, W_in, H_out, W_out, mode, Y);
947 18 }
948
949 61 void pad2d_forward(const Tensor& X, int N, int C, int H, int W,
950 int pad_top, int pad_bottom, int pad_left, int pad_right,
951 int mode, Tensor& Y) {
952 61 const auto& v = detail::dispatch(X, Y);
953
1/2
✓ Branch 0 taken 61 times.
✗ Branch 1 not taken.
 61 if (!v.pad2d_forward) detail::throw_not_implemented("pad2d_forward", X.device);
954 61 detail::adopt_output(Y, X.device);
955 61 v.pad2d_forward(X, N, C, H, W, pad_top, pad_bottom, pad_left, pad_right, mode, Y);
956 61 }
957 11 void pad2d_backward(const Tensor& dY, int N, int C, int H, int W,
958 int pad_top, int pad_bottom, int pad_left, int pad_right,
959 int mode, Tensor& dX) {
960 11 const auto& v = detail::dispatch(dY, dX);
961
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.pad2d_backward) detail::throw_not_implemented("pad2d_backward", dY.device);
962 11 detail::adopt_output(dX, dY.device);
963 11 v.pad2d_backward(dY, N, C, H, W, pad_top, pad_bottom, pad_left, pad_right, mode, dX);
964 11 }
965
966 12 void slice2d_forward(const Tensor& X, int N, int C, int H, int W,
967 int h0, int w0, int H_out, int W_out, Tensor& Y) {
968 12 const auto& v = detail::dispatch(X, Y);
969
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.slice2d_forward) detail::throw_not_implemented("slice2d_forward", X.device);
970 12 detail::adopt_output(Y, X.device);
971 12 v.slice2d_forward(X, N, C, H, W, h0, w0, H_out, W_out, Y);
972 12 }
973 6 void slice2d_backward(const Tensor& dY, int N, int C, int H, int W,
974 int h0, int w0, int H_out, int W_out, Tensor& dX) {
975 6 const auto& v = detail::dispatch(dY, dX);
976
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.slice2d_backward) detail::throw_not_implemented("slice2d_backward", dY.device);
977 6 detail::adopt_output(dX, dY.device);
978 6 v.slice2d_backward(dY, N, C, H, W, h0, w0, H_out, W_out, dX);
979 6 }
980
981 10 void unfold2d_forward(const Tensor& X, int N, int C, int H, int W,
982 int kH, int kW, int stride_h, int stride_w,
983 int pad_top, int pad_bottom, int pad_left, int pad_right,
984 int mode, Tensor& Y) {
985 10 const auto& v = detail::dispatch(X, Y);
986
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.unfold2d_forward) detail::throw_not_implemented("unfold2d_forward", X.device);
987 10 detail::adopt_output(Y, X.device);
988 20 v.unfold2d_forward(X, N, C, H, W, kH, kW, stride_h, stride_w,
989 10 pad_top, pad_bottom, pad_left, pad_right, mode, Y);
990 10 }
991
992 10 void l2_normalize_nchw_forward(const Tensor& X, int N, int C, int H, int W,
993 float eps, Tensor& Y) {
994 10 const auto& v = detail::dispatch(X, Y);
995
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.l2_normalize_nchw_forward)
996 ✗ detail::throw_not_implemented("l2_normalize_nchw_forward", X.device);
997 10 detail::adopt_output(Y, X.device);
998 10 v.l2_normalize_nchw_forward(X, N, C, H, W, eps, Y);
999 10 }
1000
1001 10 void convex_upsample_forward(const Tensor& X, const Tensor& Mask,
1002 int N, int C, int H, int W, int scale, Tensor& Y) {
1003 10 const auto& v = detail::dispatch(X, Mask, Y);
1004
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.convex_upsample_forward)
1005 ✗ detail::throw_not_implemented("convex_upsample_forward", X.device);
1006 10 detail::adopt_output(Y, X.device);
1007 10 v.convex_upsample_forward(X, Mask, N, C, H, W, scale, Y);
1008 10 }
1009
1010 17 void top_k_rows(const Tensor& X, int k, Tensor& Vals, Tensor& Idx) {
1011 17 const auto& v = detail::dispatch(X, Vals, Idx);
1012
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.top_k_rows) detail::throw_not_implemented("top_k_rows", X.device);
1013 17 detail::adopt_output(Vals, X.device);
1014 17 detail::adopt_output(Idx, X.device);
1015 17 v.top_k_rows(X, k, Vals, Idx);
1016 17 }
1017
1018 44 void adaptive_avg_pool2d_forward(const Tensor& X, int N, int C, int H, int W,
1019 int H_out, int W_out, Tensor& Y) {
1020 44 const auto& v = detail::dispatch(X, Y);
1021
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.adaptive_avg_pool2d_forward)
1022 ✗ detail::throw_not_implemented("adaptive_avg_pool2d_forward", X.device);
1023 44 detail::adopt_output(Y, X.device);
1024 44 v.adaptive_avg_pool2d_forward(X, N, C, H, W, H_out, W_out, Y);
1025 44 }
1026 6 void adaptive_avg_pool2d_backward(const Tensor& dY, int N, int C, int H, int W,
1027 int H_out, int W_out, Tensor& dX) {
1028 6 const auto& v = detail::dispatch(dY, dX);
1029
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.adaptive_avg_pool2d_backward)
1030 ✗ detail::throw_not_implemented("adaptive_avg_pool2d_backward", dY.device);
1031 6 detail::adopt_output(dX, dY.device);
1032 6 v.adaptive_avg_pool2d_backward(dY, N, C, H, W, H_out, W_out, dX);
1033 6 }
1034
1035 19 void max_pool2d_forward(const Tensor& X, int N, int C, int H, int W,
1036 int kH, int kW, int stride_h, int stride_w,
1037 int pad_h, int pad_w, Tensor& Y, Tensor& Idx) {
1038 19 const auto& v = detail::dispatch(X, Y, Idx);
1039
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.max_pool2d_forward)
1040 ✗ detail::throw_not_implemented("max_pool2d_forward", X.device);
1041 19 detail::adopt_output(Y, X.device);
1042 19 detail::adopt_output(Idx, X.device);
1043 38 v.max_pool2d_forward(X, N, C, H, W, kH, kW, stride_h, stride_w,
1044 19 pad_h, pad_w, Y, Idx);
1045 19 }
1046 8 void max_pool2d_backward(const Tensor& dY, const Tensor& Idx,
1047 int N, int C, int H, int W, int H_out, int W_out,
1048 Tensor& dX) {
1049 8 const auto& v = detail::dispatch(dY, Idx, dX);
1050
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.max_pool2d_backward)
1051 ✗ detail::throw_not_implemented("max_pool2d_backward", dY.device);
1052 8 detail::adopt_output(dX, dY.device);
1053 8 v.max_pool2d_backward(dY, Idx, N, C, H, W, H_out, W_out, dX);
1054 8 }
1055
1056 12 void gather_rows(const Tensor& X, const Tensor& Idx, Tensor& Y) {
1057 12 const auto& v = detail::dispatch(X, Idx, Y);
1058
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.gather_rows) detail::throw_not_implemented("gather_rows", X.device);
1059 12 detail::adopt_output(Y, X.device);
1060 12 v.gather_rows(X, Idx, Y);
1061 12 }
1062 9 void scatter_rows_add(const Tensor& dY, const Tensor& Idx, int R, Tensor& dX) {
1063 9 const auto& v = detail::dispatch(dY, Idx, dX);
1064
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.scatter_rows_add)
1065 ✗ detail::throw_not_implemented("scatter_rows_add", dY.device);
1066 9 detail::adopt_output(dX, dY.device);
1067 9 v.scatter_rows_add(dY, Idx, R, dX);
1068 9 }
1069 4 void scatter_rows(const Tensor& Y, const Tensor& Idx, Tensor& X) {
1070 4 const auto& v = detail::dispatch(Y, Idx, X);
1071
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.scatter_rows) detail::throw_not_implemented("scatter_rows", Y.device);
1072 // X is an in-place destination (existing rows are preserved), so it is
1073 // never adopted/resized here — the backend validates its shape instead.
1074 4 v.scatter_rows(Y, Idx, X);
1075 4 }
1076
1077 82 void conv_transpose2d_forward(const Tensor& X, const Tensor& Wt,
1078 const Tensor* bias,
1079 int N, int C_in, int H, int W,
1080 int C_out, int kH, int kW,
1081 int stride_h, int stride_w,
1082 int pad_h, int pad_w,
1083 int output_padding_h, int output_padding_w,
1084 int dil_h, int dil_w, int groups,
1085 Tensor& Y) {
1086 82 const auto& v = detail::dispatch(X, Wt, Y);
1087
1/2
✓ Branch 0 taken 82 times.
✗ Branch 1 not taken.
 82 if (!v.conv_transpose2d_forward)
1088 ✗ detail::throw_not_implemented("conv_transpose2d_forward", X.device);
1089 82 detail::adopt_output(Y, X.device);
1090 164 v.conv_transpose2d_forward(X, Wt, bias, N, C_in, H, W, C_out, kH, kW,
1091 82 stride_h, stride_w, pad_h, pad_w,
1092 82 output_padding_h, output_padding_w,
1093 82 dil_h, dil_w, groups, Y);
1094 82 }
1095 8 void conv_transpose2d_backward_input(const Tensor& Wt, const Tensor& dY,
1096 int N, int C_in, int H, int W,
1097 int C_out, int kH, int kW,
1098 int stride_h, int stride_w,
1099 int pad_h, int pad_w,
1100 int output_padding_h, int output_padding_w,
1101 int dil_h, int dil_w, int groups,
1102 Tensor& dX) {
1103 8 const auto& v = detail::dispatch(Wt, dY, dX);
1104
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.conv_transpose2d_backward_input)
1105 ✗ detail::throw_not_implemented("conv_transpose2d_backward_input", Wt.device);
1106 8 detail::adopt_output(dX, Wt.device);
1107 16 v.conv_transpose2d_backward_input(Wt, dY, N, C_in, H, W, C_out, kH, kW,
1108 8 stride_h, stride_w, pad_h, pad_w,
1109 8 output_padding_h, output_padding_w,
1110 8 dil_h, dil_w, groups, dX);
1111 8 }
1112 7 void conv_transpose2d_backward_weight(const Tensor& X, const Tensor& dY,
1113 int N, int C_in, int H, int W,
1114 int C_out, int kH, int kW,
1115 int stride_h, int stride_w,
1116 int pad_h, int pad_w,
1117 int output_padding_h, int output_padding_w,
1118 int dil_h, int dil_w, int groups,
1119 Tensor& dWt) {
1120 7 const auto& v = detail::dispatch(X, dY, dWt);
1121
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.conv_transpose2d_backward_weight)
1122 ✗ detail::throw_not_implemented("conv_transpose2d_backward_weight", X.device);
1123 7 detail::adopt_output(dWt, X.device);
1124 14 v.conv_transpose2d_backward_weight(X, dY, N, C_in, H, W, C_out, kH, kW,
1125 7 stride_h, stride_w, pad_h, pad_w,
1126 7 output_padding_h, output_padding_w,
1127 7 dil_h, dil_w, groups, dWt);
1128 7 }
1129 3 void conv_transpose2d_backward_bias(const Tensor& dY, int N, int C_out,
1130 int H_out, int W_out, Tensor& dB) {
1131 3 const auto& v = detail::dispatch(dY, dB);
1132
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.conv_transpose2d_backward_bias)
1133 ✗ detail::throw_not_implemented("conv_transpose2d_backward_bias", dY.device);
1134 3 detail::adopt_output(dB, dY.device);
1135 3 v.conv_transpose2d_backward_bias(dY, N, C_out, H_out, W_out, dB);
1136 3 }
1137
1138 14 void window_partition_forward(const Tensor& X, int N, int C, int H, int W,
1139 int window, Tensor& Y) {
1140 14 const auto& v = detail::dispatch(X, Y);
1141
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.window_partition_forward)
1142 ✗ detail::throw_not_implemented("window_partition_forward", X.device);
1143 14 detail::adopt_output(Y, X.device);
1144 14 v.window_partition_forward(X, N, C, H, W, window, Y);
1145 14 }
1146 8 void window_reverse_forward(const Tensor& X, int N, int C, int H, int W,
1147 int window, Tensor& Y) {
1148 8 const auto& v = detail::dispatch(X, Y);
1149
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.window_reverse_forward)
1150 ✗ detail::throw_not_implemented("window_reverse_forward", X.device);
1151 8 detail::adopt_output(Y, X.device);
1152 8 v.window_reverse_forward(X, N, C, H, W, window, Y);
1153 8 }
1154
1155 // ─── FP16 linear + GEGLU ───────────────────────────────────────────────────
1156
1157 6 void linear_forward_batched_fp16(const Tensor& W, const Tensor* bias,
1158 const Tensor& X_BD, Tensor& Y_BD) {
1159 6 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD});
1160
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.linear_forward_batched_fp16) detail::throw_not_implemented("linear_forward_batched_fp16", W.device);
1161 6 detail::adopt_output(Y_BD, W.device);
1162 6 v.linear_forward_batched_fp16(W, bias, X_BD, Y_BD);
1163 6 }
1164
1165 8 void linear_forward_batched_fp16_act(const Tensor& W, const Tensor* bias,
1166 const Tensor& X_BD, int act, Tensor& Y_BD) {
1167 8 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD});
1168
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.linear_forward_batched_fp16_act)
1169 ✗ detail::throw_not_implemented("linear_forward_batched_fp16_act", W.device);
1170 8 detail::adopt_output(Y_BD, W.device);
1171 8 v.linear_forward_batched_fp16_act(W, bias, X_BD, act, Y_BD);
1172 8 }
1173
1174 10 void geglu_forward(const Tensor& X, Tensor& Y) {
1175 10 const auto& v = detail::dispatch(X, Y);
1176
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.geglu_forward) detail::throw_not_implemented("geglu_forward", X.device);
1177 10 detail::adopt_output(Y, X.device);
1178 10 v.geglu_forward(X, Y);
1179 10 }
1180 6 void geglu_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
1181 6 const auto& v = detail::dispatch(X, dY, dX);
1182
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.geglu_backward) detail::throw_not_implemented("geglu_backward", X.device);
1183 6 detail::adopt_output(dX, X.device);
1184 6 v.geglu_backward(X, dY, dX);
1185 6 }
1186 8 void geglu_exact_forward(const Tensor& X, Tensor& Y) {
1187 8 const auto& v = detail::dispatch(X, Y);
1188
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.geglu_exact_forward) detail::throw_not_implemented("geglu_exact_forward", X.device);
1189 8 detail::adopt_output(Y, X.device);
1190 8 v.geglu_exact_forward(X, Y);
1191 8 }
1192 6 void geglu_exact_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
1193 6 const auto& v = detail::dispatch(X, dY, dX);
1194
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.geglu_exact_backward) detail::throw_not_implemented("geglu_exact_backward", X.device);
1195 6 detail::adopt_output(dX, X.device);
1196 6 v.geglu_exact_backward(X, dY, dX);
1197 6 }
1198
1199 // ─── Causal mask helper ────────────────────────────────────────────────────
1200
1201 10 void build_causal_mask_row(int L, int q, Tensor& mask) {
1202 10 const auto& v = detail::dispatch(mask);
1203
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.build_causal_mask_row) detail::throw_not_implemented("build_causal_mask_row", mask.device);
1204 10 detail::adopt_output(mask, mask.device);
1205 10 v.build_causal_mask_row(L, q, mask);
1206 10 }
1207
1208 // ─── Cross-attention family ────────────────────────────────────────────────
1209
1210 16 void cross_attention_forward(const Tensor& X, const Tensor& Ctx,
1211 const Tensor& Wq, const Tensor& Wk,
1212 const Tensor& Wv, const Tensor& Wo,
1213 const float* d_mask, int num_heads, Tensor& O) {
1214 16 const auto& v = detail::dispatch(X, Ctx, Wq, Wk, Wv, Wo, O);
1215
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.cross_attention_forward) detail::throw_not_implemented("cross_attention_forward", X.device);
1216 16 detail::adopt_output(O, X.device);
1217 16 v.cross_attention_forward(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads, O);
1218 16 }
1219
1220 10 void cross_attention_forward_with_attn(const Tensor& X, const Tensor& Ctx,
1221 const Tensor& Wq, const Tensor& Wk,
1222 const Tensor& Wv, const Tensor& Wo,
1223 const float* d_mask,
1224 const Tensor* attn_logit_bias,
1225 int num_heads, Tensor& O, Tensor& AttnAvg) {
1226 10 const auto& v = detail::dispatch_with_opts(X, Ctx, {&Wq, &Wk, &Wv, &Wo, attn_logit_bias, &O, &AttnAvg});
1227
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.cross_attention_forward_with_attn)
1228 ✗ detail::throw_not_implemented("cross_attention_forward_with_attn", X.device);
1229 10 detail::adopt_output(O, X.device);
1230 10 detail::adopt_output(AttnAvg, X.device);
1231 20 v.cross_attention_forward_with_attn(X, Ctx, Wq, Wk, Wv, Wo, d_mask,
1232 10 attn_logit_bias, num_heads, O, AttnAvg);
1233 10 }
1234
1235 16 void self_attention_forward_train(const Tensor& X,
1236 const Tensor& Wq, const Tensor& Wk,
1237 const Tensor& Wv, const Tensor& Wo,
1238 const float* d_mask, int num_heads,
1239 Tensor& Qh, Tensor& Kh, Tensor& Vh,
1240 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
1241 16 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, Qh, Kh, Vh);
1242
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.self_attention_forward_train)
1243 ✗ detail::throw_not_implemented("self_attention_forward_train", X.device);
1244 16 detail::adopt_output(Qh, X.device);
1245 16 detail::adopt_output(Kh, X.device);
1246 16 detail::adopt_output(Vh, X.device);
1247 16 detail::adopt_output(Attnh, X.device);
1248 16 detail::adopt_output(Yconcat, X.device);
1249 16 detail::adopt_output(O, X.device);
1250 32 v.self_attention_forward_train(X, Wq, Wk, Wv, Wo, d_mask, num_heads,
1251 16 Qh, Kh, Vh, Attnh, Yconcat, O);
1252 16 }
1253
1254 14 void self_attention_backward(const Tensor& dO, const Tensor& X,
1255 const Tensor& Qh, const Tensor& Kh,
1256 const Tensor& Vh, const Tensor& Attnh,
1257 const Tensor& Yconcat,
1258 const Tensor& Wq, const Tensor& Wk,
1259 const Tensor& Wv, const Tensor& Wo,
1260 const float* d_mask, int num_heads,
1261 Tensor& dX,
1262 Tensor& dWq, Tensor& dWk,
1263 Tensor& dWv, Tensor& dWo) {
1264 14 const auto& v = detail::dispatch(dO, X, Qh, Kh, Vh, Attnh, Yconcat);
1265
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.self_attention_backward)
1266 ✗ detail::throw_not_implemented("self_attention_backward", dO.device);
1267 14 detail::adopt_output(dX, dO.device);
1268 14 detail::adopt_output(dWq, dO.device);
1269 14 detail::adopt_output(dWk, dO.device);
1270 14 detail::adopt_output(dWv, dO.device);
1271 14 detail::adopt_output(dWo, dO.device);
1272 28 v.self_attention_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat,
1273 14 Wq, Wk, Wv, Wo, d_mask, num_heads,
1274 14 dX, dWq, dWk, dWv, dWo);
1275 14 }
1276
1277 10 void attention_token_moments(const Tensor& Attn, int h_lat, int w_lat,
1278 Tensor& mass, Tensor& centroid) {
1279 10 const auto& v = detail::dispatch(Attn, mass, centroid);
1280
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.attention_token_moments) detail::throw_not_implemented("attention_token_moments", Attn.device);
1281 10 detail::adopt_output(mass, Attn.device);
1282 10 detail::adopt_output(centroid, Attn.device);
1283 10 v.attention_token_moments(Attn, h_lat, w_lat, mass, centroid);
1284 10 }
1285
1286 10 void cross_attention_forward_train(const Tensor& X, const Tensor& Ctx,
1287 const Tensor& Wq, const Tensor& Wk,
1288 const Tensor& Wv, const Tensor& Wo,
1289 const float* d_mask, int num_heads,
1290 Tensor& Qh, Tensor& Kh, Tensor& Vh,
1291 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
1292 10 const auto& v = detail::dispatch(X, Ctx, Wq, Wk, Wv, Wo, Qh, Kh);
1293
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.cross_attention_forward_train)
1294 ✗ detail::throw_not_implemented("cross_attention_forward_train", X.device);
1295 10 detail::adopt_output(Qh, X.device);
1296 10 detail::adopt_output(Kh, X.device);
1297 10 detail::adopt_output(Vh, X.device);
1298 10 detail::adopt_output(Attnh, X.device);
1299 10 detail::adopt_output(Yconcat, X.device);
1300 10 detail::adopt_output(O, X.device);
1301 20 v.cross_attention_forward_train(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads,
1302 10 Qh, Kh, Vh, Attnh, Yconcat, O);
1303 10 }
1304
1305 10 void cross_attention_backward(const Tensor& dO, const Tensor& X, const Tensor& Ctx,
1306 const Tensor& Qh, const Tensor& Kh,
1307 const Tensor& Vh, const Tensor& Attnh,
1308 const Tensor& Yconcat,
1309 const Tensor& Wq, const Tensor& Wk,
1310 const Tensor& Wv, const Tensor& Wo,
1311 const float* d_mask, int num_heads,
1312 Tensor& dX, Tensor& dCtx,
1313 Tensor& dWq, Tensor& dWk,
1314 Tensor& dWv, Tensor& dWo) {
1315 10 const auto& v = detail::dispatch(dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat);
1316
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.cross_attention_backward)
1317 ✗ detail::throw_not_implemented("cross_attention_backward", dO.device);
1318 10 detail::adopt_output(dX, dO.device);
1319 10 detail::adopt_output(dCtx, dO.device);
1320 10 detail::adopt_output(dWq, dO.device);
1321 10 detail::adopt_output(dWk, dO.device);
1322 10 detail::adopt_output(dWv, dO.device);
1323 10 detail::adopt_output(dWo, dO.device);
1324 20 v.cross_attention_backward(dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat,
1325 10 Wq, Wk, Wv, Wo, d_mask, num_heads,
1326 10 dX, dCtx, dWq, dWk, dWv, dWo);
1327 10 }
1328
1329 // ─── FP16 LN inference + FP16 self-attention ───────────────────────────────
1330
1331 ✗ void layernorm_forward_inference_batched_fp16(const Tensor& X_RD,
1332 const Tensor& gamma,
1333 const Tensor& beta,
1334 Tensor& Y_RD, float eps) {
1335 ✗ const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD);
1336 ✗ if (!v.layernorm_forward_inference_batched_fp16)
1337 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched_fp16", X_RD.device);
1338 ✗ detail::adopt_output(Y_RD, X_RD.device);
1339 ✗ v.layernorm_forward_inference_batched_fp16(X_RD, gamma, beta, Y_RD, eps);
1340 ✗ }
1341
1342 12 void self_attention_forward(const Tensor& X,
1343 const Tensor& Wq, const Tensor& Wk,
1344 const Tensor& Wv, const Tensor& Wo,
1345 const float* d_mask, int num_heads, Tensor& O) {
1346 12 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, O);
1347
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.self_attention_forward) detail::throw_not_implemented("self_attention_forward", X.device);
1348 12 detail::adopt_output(O, X.device);
1349 12 v.self_attention_forward(X, Wq, Wk, Wv, Wo, d_mask, num_heads, O);
1350 12 }
1351
1352 39 void self_attention_bias_forward(const Tensor& X,
1353 const Tensor& Wq, const Tensor& Wk,
1354 const Tensor& Wv, const Tensor& Wo,
1355 const Tensor* bq, const Tensor* bk,
1356 const Tensor* bv, const Tensor* bo,
1357 const float* d_mask,
1358 const Tensor* attn_bias,
1359 int num_heads, float scale, Tensor& O) {
1360 39 const auto& v = detail::dispatch_with_opts(X, Wq, {&Wk, &Wv, &Wo, bq, bk, bv, bo, attn_bias, &O});
1361
1/2
✓ Branch 0 taken 39 times.
✗ Branch 1 not taken.
 39 if (!v.self_attention_bias_forward)
1362 ✗ detail::throw_not_implemented("self_attention_bias_forward", X.device);
1363 39 detail::adopt_output(O, X.device);
1364 78 v.self_attention_bias_forward(X, Wq, Wk, Wv, Wo, bq, bk, bv, bo, d_mask, attn_bias,
1365 39 num_heads, scale, O);
1366 39 }
1367
1368 23 void self_attention_decomposed_rel_pos_forward(
1369 const Tensor& X,
1370 const Tensor& Wq, const Tensor* bq,
1371 const Tensor& Wk, const Tensor* bk,
1372 const Tensor& Wv, const Tensor* bv,
1373 const Tensor& Wo, const Tensor* bo,
1374 const Tensor& rel_pos_h, const Tensor& rel_pos_w,
1375 int num_heads, int grid_h, int grid_w, float scale, Tensor& O) {
1376 23 const auto& v = detail::dispatch_with_opts(
1377 23 X, Wq, {bq, &Wk, bk, &Wv, bv, &Wo, bo, &rel_pos_h, &rel_pos_w, &O});
1378
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.self_attention_decomposed_rel_pos_forward)
1379 ✗ detail::throw_not_implemented("self_attention_decomposed_rel_pos_forward",
1380 ✗ X.device);
1381 23 detail::adopt_output(O, X.device);
1382 46 v.self_attention_decomposed_rel_pos_forward(
1383 23 X, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
1384 23 num_heads, grid_h, grid_w, scale, O);
1385 23 }
1386
1387 15 void self_attention_decomposed_rel_pos_windowed_forward(
1388 const Tensor& X,
1389 const Tensor& Wq, const Tensor* bq,
1390 const Tensor& Wk, const Tensor* bk,
1391 const Tensor& Wv, const Tensor* bv,
1392 const Tensor& Wo, const Tensor* bo,
1393 const Tensor& rel_pos_h, const Tensor& rel_pos_w,
1394 int num_heads, int grid_h, int grid_w, int window, float scale,
1395 Tensor& O) {
1396 15 const auto& v = detail::dispatch_with_opts(
1397 15 X, Wq, {bq, &Wk, bk, &Wv, bv, &Wo, bo, &rel_pos_h, &rel_pos_w, &O});
1398
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.self_attention_decomposed_rel_pos_windowed_forward)
1399 ✗ detail::throw_not_implemented(
1400 ✗ "self_attention_decomposed_rel_pos_windowed_forward", X.device);
1401 15 detail::adopt_output(O, X.device);
1402 30 v.self_attention_decomposed_rel_pos_windowed_forward(
1403 15 X, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
1404 15 num_heads, grid_h, grid_w, window, scale, O);
1405 15 }
1406
1407 5 void self_attention_bias_int8w_fp16(const Tensor& X,
1408 const Tensor& Wq_int8, const Tensor& sq,
1409 const Tensor& Wk_int8, const Tensor& sk,
1410 const Tensor& Wv_int8, const Tensor& sv,
1411 const Tensor& Wo_int8, const Tensor& so,
1412 const float* d_mask,
1413 const Tensor* attn_bias,
1414 int num_heads, float scale, Tensor& O) {
1415 5 const auto& v = detail::dispatch_with_opts(
1416 15 X, Wq_int8, {&sq, &Wk_int8, &sk, &Wv_int8, &sv, &Wo_int8, &so,
1417 10 attn_bias, &O});
1418
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.self_attention_bias_int8w_fp16)
1419 ✗ detail::throw_not_implemented("self_attention_bias_int8w_fp16", X.device);
1420 5 detail::adopt_output(O, X.device);
1421 10 v.self_attention_bias_int8w_fp16(X, Wq_int8, sq, Wk_int8, sk,
1422 5 Wv_int8, sv, Wo_int8, so,
1423 5 d_mask, attn_bias, num_heads, scale, O);
1424 5 }
1425
1426 // ─── Flash attention family ────────────────────────────────────────────────
1427
1428 29 void flash_attention_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1429 const float* d_mask, int num_heads, bool causal,
1430 Tensor& O) {
1431 29 const auto& v = detail::dispatch(Q, K, V, O);
1432
1/2
✓ Branch 0 taken 29 times.
✗ Branch 1 not taken.
 29 if (!v.flash_attention_forward) detail::throw_not_implemented("flash_attention_forward", Q.device);
1433 29 detail::adopt_output(O, Q.device);
1434 29 v.flash_attention_forward(Q, K, V, d_mask, num_heads, causal, O);
1435 29 }
1436
1437 15 void flash_attention_gqa_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1438 const float* d_mask, int num_q_heads, int num_kv_heads,
1439 bool causal, Tensor& O) {
1440 15 const auto& v = detail::dispatch(Q, K, V, O);
1441
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.flash_attention_gqa_forward)
1442 ✗ detail::throw_not_implemented("flash_attention_gqa_forward", Q.device);
1443 15 detail::adopt_output(O, Q.device);
1444 15 v.flash_attention_gqa_forward(Q, K, V, d_mask, num_q_heads, num_kv_heads, causal, O);
1445 15 }
1446
1447 ✗ void flash_attention_windowed_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1448 const float* d_mask, int num_heads, int window,
1449 Tensor& O) {
1450 ✗ const auto& v = detail::dispatch(Q, K, V, O);
1451 ✗ if (!v.flash_attention_windowed_forward)
1452 ✗ detail::throw_not_implemented("flash_attention_windowed_forward", Q.device);
1453 ✗ detail::adopt_output(O, Q.device);
1454 ✗ v.flash_attention_windowed_forward(Q, K, V, d_mask, num_heads, window, O);
1455 ✗ }
1456
1457 27 void flash_attention_varlen_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1458 const int32_t* cu_seqlens_q,
1459 const int32_t* cu_seqlens_k,
1460 int batch_size, int max_seqlen_q, int max_seqlen_k,
1461 int num_heads, int head_dim, bool causal,
1462 Tensor& O) {
1463 27 const auto& v = detail::dispatch(Q, K, V, O);
1464
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.flash_attention_varlen_forward)
1465 ✗ detail::throw_not_implemented("flash_attention_varlen_forward", Q.device);
1466 27 detail::adopt_output(O, Q.device);
1467 54 v.flash_attention_varlen_forward(Q, K, V, cu_seqlens_q, cu_seqlens_k,
1468 27 batch_size, max_seqlen_q, max_seqlen_k,
1469 27 num_heads, head_dim, causal, O);
1470 27 }
1471
1472 27 void flash_attention_varlen_backward(const Tensor& Q, const Tensor& K, const Tensor& V,
1473 const Tensor& O, const Tensor& dO,
1474 const int32_t* cu_seqlens_q,
1475 const int32_t* cu_seqlens_k,
1476 int batch_size, int max_seqlen_q, int max_seqlen_k,
1477 int num_heads, int head_dim, bool causal,
1478 Tensor& dQ, Tensor& dK, Tensor& dV) {
1479 27 const auto& v = detail::dispatch(Q, K, V, O, dO, dQ, dK, dV);
1480
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.flash_attention_varlen_backward)
1481 ✗ detail::throw_not_implemented("flash_attention_varlen_backward", Q.device);
1482 27 detail::adopt_output(dQ, Q.device);
1483 27 detail::adopt_output(dK, Q.device);
1484 27 detail::adopt_output(dV, Q.device);
1485 54 v.flash_attention_varlen_backward(Q, K, V, O, dO,
1486 27 cu_seqlens_q, cu_seqlens_k,
1487 27 batch_size, max_seqlen_q, max_seqlen_k,
1488 27 num_heads, head_dim, causal,
1489 27 dQ, dK, dV);
1490 27 }
1491
1492 24 void flash_attention_qkvo_forward(const Tensor& X, const Tensor* Ctx,
1493 const Tensor& Wq, const Tensor* bq,
1494 const Tensor& Wk, const Tensor* bk,
1495 const Tensor& Wv, const Tensor* bv,
1496 const Tensor& Wo, const Tensor* bo,
1497 const float* d_mask, int num_heads,
1498 bool causal, Tensor& O) {
1499 24 const auto& v = detail::dispatch_with_opts(
1500 24 X, Wq, {Ctx, bq, &Wk, bk, &Wv, bv, &Wo, bo, &O});
1501
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.flash_attention_qkvo_forward)
1502 ✗ detail::throw_not_implemented("flash_attention_qkvo_forward", X.device);
1503 24 detail::adopt_output(O, X.device);
1504 48 v.flash_attention_qkvo_forward(X, Ctx, Wq, bq, Wk, bk, Wv, bv, Wo, bo,
1505 24 d_mask, num_heads, causal, O);
1506 24 }
1507
1508 18 void flash_attention_qkvo_backward(
1509 const Tensor& X, const Tensor* Ctx,
1510 const Tensor& Wq, const Tensor* bq,
1511 const Tensor& Wk, const Tensor* bk,
1512 const Tensor& Wv, const Tensor* bv,
1513 const Tensor& Wo, const Tensor* bo,
1514 const float* d_mask, int num_heads, bool causal,
1515 const Tensor& dO,
1516 Tensor& dX, Tensor* dCtx,
1517 Tensor& dWq, Tensor* dbq,
1518 Tensor& dWk, Tensor* dbk,
1519 Tensor& dWv, Tensor* dbv,
1520 Tensor& dWo, Tensor* dbo) {
1521 18 const auto& v = detail::dispatch_with_opts(
1522 198 X, Wq, {Ctx, bq, &Wk, bk, &Wv, bv, &Wo, bo, &dO,
1523 180 &dX, dCtx, &dWq, dbq, &dWk, dbk, &dWv, dbv, &dWo, dbo});
1524
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.flash_attention_qkvo_backward)
1525 ✗ detail::throw_not_implemented("flash_attention_qkvo_backward", X.device);
1526 18 detail::adopt_output(dX, X.device);
1527
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 if (dCtx) detail::adopt_output(*dCtx, X.device);
1528 18 detail::adopt_output(dWq, X.device);
1529
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbq) detail::adopt_output(*dbq, X.device);
1530 18 detail::adopt_output(dWk, X.device);
1531
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbk) detail::adopt_output(*dbk, X.device);
1532 18 detail::adopt_output(dWv, X.device);
1533
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbv) detail::adopt_output(*dbv, X.device);
1534 18 detail::adopt_output(dWo, X.device);
1535
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbo) detail::adopt_output(*dbo, X.device);
1536 36 v.flash_attention_qkvo_backward(X, Ctx, Wq, bq, Wk, bk, Wv, bv, Wo, bo,
1537 18 d_mask, num_heads, causal, dO,
1538 18 dX, dCtx, dWq, dbq, dWk, dbk,
1539 18 dWv, dbv, dWo, dbo);
1540 18 }
1541
1542 16 void flash_attention_backward(const Tensor& Q, const Tensor& K, const Tensor& V,
1543 const Tensor& O, const Tensor& dO,
1544 const float* d_mask, int num_heads, bool causal,
1545 Tensor& dQ, Tensor& dK, Tensor& dV) {
1546 16 const auto& v = detail::dispatch(Q, K, V, O, dO, dQ, dK, dV);
1547
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.flash_attention_backward) detail::throw_not_implemented("flash_attention_backward", Q.device);
1548 16 detail::adopt_output(dQ, Q.device);
1549 16 detail::adopt_output(dK, Q.device);
1550 16 detail::adopt_output(dV, Q.device);
1551 32 v.flash_attention_backward(Q, K, V, O, dO, d_mask, num_heads, causal,
1552 16 dQ, dK, dV);
1553 16 }
1554
1555 8 void flash_attention_project_kv(const Tensor& ctx,
1556 const Tensor& Wk, const Tensor* bk,
1557 const Tensor& Wv, const Tensor* bv,
1558 Tensor& K_out, Tensor& V_out) {
1559 8 const auto& v = detail::dispatch_with_opts(ctx, Wk, {bk, &Wv, bv, &K_out, &V_out});
1560
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.flash_attention_project_kv)
1561 ✗ detail::throw_not_implemented("flash_attention_project_kv", ctx.device);
1562 8 detail::adopt_output(K_out, ctx.device);
1563 8 detail::adopt_output(V_out, ctx.device);
1564 8 v.flash_attention_project_kv(ctx, Wk, bk, Wv, bv, K_out, V_out);
1565 8 }
1566
1567 12 void flash_attention_q_with_kv_cached_forward(const Tensor& X,
1568 const Tensor& K, const Tensor& V,
1569 const Tensor& Wq, const Tensor* bq,
1570 const Tensor& Wo, const Tensor* bo,
1571 const float* d_mask, int num_heads,
1572 bool causal, Tensor& O) {
1573 12 const auto& v = detail::dispatch_with_opts(X, K, {&V, &Wq, bq, &Wo, bo, &O});
1574
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.flash_attention_q_with_kv_cached_forward)
1575 ✗ detail::throw_not_implemented("flash_attention_q_with_kv_cached_forward", X.device);
1576 12 detail::adopt_output(O, X.device);
1577 24 v.flash_attention_q_with_kv_cached_forward(X, K, V, Wq, bq, Wo, bo,
1578 12 d_mask, num_heads, causal, O);
1579 12 }
1580
1581 // ─── NCHW <-> sequence ─────────────────────────────────────────────────────
1582
1583 19 void nchw_to_sequence(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
1584 19 const auto& v = detail::dispatch(X, Y);
1585
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.nchw_to_sequence) detail::throw_not_implemented("nchw_to_sequence", X.device);
1586 19 detail::adopt_output(Y, X.device);
1587 19 v.nchw_to_sequence(X, N, C, H, W, Y);
1588 19 }
1589
1590 19 void sequence_to_nchw(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
1591 19 const auto& v = detail::dispatch(X, Y);
1592
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.sequence_to_nchw) detail::throw_not_implemented("sequence_to_nchw", X.device);
1593 19 detail::adopt_output(Y, X.device);
1594 19 v.sequence_to_nchw(X, N, C, H, W, Y);
1595 19 }
1596
1597 34 void spatial_merge_2x2_forward(const Tensor& X, int N, int C, int H, int W,
1598 bool channel_major, Tensor& Y) {
1599 34 const auto& v = detail::dispatch(X, Y);
1600
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!v.spatial_merge_2x2_forward)
1601 ✗ detail::throw_not_implemented("spatial_merge_2x2_forward", X.device);
1602 34 detail::adopt_output(Y, X.device);
1603 34 v.spatial_merge_2x2_forward(X, N, C, H, W, channel_major, Y);
1604 34 }
1605
1606 47 void pixel_shuffle_upsample_2x_forward(const Tensor& X, int N, int C_in,
1607 int H, int W, int C_out, Tensor& Y) {
1608 47 const auto& v = detail::dispatch(X, Y);
1609
1/2
✓ Branch 0 taken 47 times.
✗ Branch 1 not taken.
 47 if (!v.pixel_shuffle_upsample_2x_forward)
1610 ✗ detail::throw_not_implemented("pixel_shuffle_upsample_2x_forward", X.device);
1611 47 detail::adopt_output(Y, X.device);
1612 47 v.pixel_shuffle_upsample_2x_forward(X, N, C_in, H, W, C_out, Y);
1613 47 }
1614
1615 12 void patch_unpack_forward(const Tensor& tokens, int hp, int wp, int P,
1616 int C_total, int C_keep, bool channel_major,
1617 Tensor& Y) {
1618 12 const auto& v = detail::dispatch(tokens, Y);
1619
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.patch_unpack_forward)
1620 ✗ detail::throw_not_implemented("patch_unpack_forward", tokens.device);
1621 12 detail::adopt_output(Y, tokens.device);
1622 12 v.patch_unpack_forward(tokens, hp, wp, P, C_total, C_keep, channel_major, Y);
1623 12 }
1624
1625 // ─── ResBlock ──────────────────────────────────────────────────────────────
1626
1627 16 void resblock_forward(const Tensor& X,
1628 const Tensor& gamma1, const Tensor& beta1,
1629 const Tensor& W1, const Tensor* b1,
1630 const Tensor* t_emb_shift,
1631 const Tensor& gamma2, const Tensor& beta2,
1632 const Tensor& W2, const Tensor* b2,
1633 const Tensor* Wskip, const Tensor* bskip,
1634 int N, int C_in, int C_out, int H, int W,
1635 int num_groups, float eps, Tensor& Y) {
1636 16 const auto& v = detail::dispatch_with_opts(
1637 64 X, gamma1, {&beta1, &W1, b1, t_emb_shift, &gamma2, &beta2, &W2, b2,
1638 48 Wskip, bskip, &Y});
1639
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.resblock_forward) detail::throw_not_implemented("resblock_forward", X.device);
1640 16 detail::adopt_output(Y, X.device);
1641 32 v.resblock_forward(X, gamma1, beta1, W1, b1, t_emb_shift,
1642 16 gamma2, beta2, W2, b2, Wskip, bskip,
1643 16 N, C_in, C_out, H, W, num_groups, eps, Y);
1644 16 }
1645
1646 ✗ void resblock_forward_int8w_fp16(const Tensor& X,
1647 const Tensor& gamma1, const Tensor& beta1,
1648 const Tensor& W1_int8, const Tensor& s1,
1649 const Tensor* b1,
1650 const Tensor* t_emb_shift,
1651 const Tensor& gamma2, const Tensor& beta2,
1652 const Tensor& W2_int8, const Tensor& s2,
1653 const Tensor* b2,
1654 const Tensor* Wskip_int8, const Tensor* sskip,
1655 const Tensor* bskip,
1656 int N, int C_in, int C_out, int H, int W,
1657 int num_groups, float eps, Tensor& Y) {
1658 ✗ const auto& v = detail::dispatch_with_opts(
1659 ✗ X, gamma1, {&beta1, &W1_int8, &s1, b1, t_emb_shift,
1660 ✗ &gamma2, &beta2, &W2_int8, &s2, b2,
1661 ✗ Wskip_int8, sskip, bskip, &Y});
1662 ✗ if (!v.resblock_forward_int8w_fp16)
1663 ✗ detail::throw_not_implemented("resblock_forward_int8w_fp16", X.device);
1664 ✗ detail::adopt_output(Y, X.device);
1665 ✗ v.resblock_forward_int8w_fp16(X, gamma1, beta1, W1_int8, s1, b1, t_emb_shift,
1666 ✗ gamma2, beta2, W2_int8, s2, b2,
1667 ✗ Wskip_int8, sskip, bskip,
1668 ✗ N, C_in, C_out, H, W, num_groups, eps, Y);
1669 ✗ }
1670
1671 16 void resblock_backward(const Tensor& X,
1672 const Tensor& gamma1, const Tensor& beta1,
1673 const Tensor& W1, const Tensor* b1,
1674 const Tensor* t_emb_shift,
1675 const Tensor& gamma2, const Tensor& beta2,
1676 const Tensor& W2, const Tensor* b2,
1677 const Tensor* Wskip, const Tensor* bskip,
1678 int N, int C_in, int C_out, int H, int W,
1679 int num_groups, float eps,
1680 const Tensor& dY,
1681 Tensor& dX,
1682 Tensor& dGamma1, Tensor& dBeta1,
1683 Tensor& dW1, Tensor* db1,
1684 Tensor* dt_emb_shift,
1685 Tensor& dGamma2, Tensor& dBeta2,
1686 Tensor& dW2, Tensor* db2,
1687 Tensor* dWskip, Tensor* dbskip) {
1688 16 const auto& v = detail::dispatch_with_opts(
1689 256 X, gamma1, {&beta1, &W1, b1, t_emb_shift, &gamma2, &beta2, &W2, b2,
1690 128 Wskip, bskip, &dY, &dX, &dGamma1, &dBeta1, &dW1, db1,
1691 112 dt_emb_shift, &dGamma2, &dBeta2, &dW2, db2, dWskip, dbskip});
1692
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.resblock_backward) detail::throw_not_implemented("resblock_backward", X.device);
1693 16 detail::adopt_output(dX, X.device);
1694 16 detail::adopt_output(dGamma1, X.device);
1695 16 detail::adopt_output(dBeta1, X.device);
1696 16 detail::adopt_output(dW1, X.device);
1697
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
 16 if (db1) detail::adopt_output(*db1, X.device);
1698
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 if (dt_emb_shift) detail::adopt_output(*dt_emb_shift, X.device);
1699 16 detail::adopt_output(dGamma2, X.device);
1700 16 detail::adopt_output(dBeta2, X.device);
1701 16 detail::adopt_output(dW2, X.device);
1702
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
 16 if (db2) detail::adopt_output(*db2, X.device);
1703
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 if (dWskip) detail::adopt_output(*dWskip, X.device);
1704
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 if (dbskip) detail::adopt_output(*dbskip, X.device);
1705 32 v.resblock_backward(X, gamma1, beta1, W1, b1, t_emb_shift,
1706 16 gamma2, beta2, W2, b2, Wskip, bskip,
1707 16 N, C_in, C_out, H, W, num_groups, eps,
1708 16 dY, dX, dGamma1, dBeta1, dW1, db1, dt_emb_shift,
1709 16 dGamma2, dBeta2, dW2, db2, dWskip, dbskip);
1710 16 }
1711
1712 // ─── Matmul + RoPE + RMSNorm + SwiGLU + KV-cache + Llama ───────────────────
1713
1714 60 void matmul(const Tensor& A, const Tensor& B, Tensor& C) {
1715 60 const auto& v = detail::dispatch(A, B, C);
1716
1/2
✓ Branch 0 taken 60 times.
✗ Branch 1 not taken.
 60 if (!v.matmul) detail::throw_not_implemented("matmul", A.device);
1717 60 detail::adopt_output(C, A.device);
1718 60 v.matmul(A, B, C);
1719 60 }
1720
1721 51 void matmul_abt(const Tensor& A, const Tensor& B, Tensor& C,
1722 int batch, int M, int N, int K,
1723 long long strideA, long long strideB, long long strideC,
1724 const Tensor* bias, int act) {
1725 51 const auto& v = detail::dispatch(A, B, C);
1726
1/2
✓ Branch 0 taken 51 times.
✗ Branch 1 not taken.
 51 if (!v.matmul_abt) detail::throw_not_implemented("matmul_abt", A.device);
1727 51 detail::adopt_output(C, A.device);
1728 51 v.matmul_abt(A, B, C, batch, M, N, K, strideA, strideB, strideC, bias, act);
1729 51 }
1730
1731 16 void matmul_backward(const Tensor& A, const Tensor& B, const Tensor& dC,
1732 Tensor& dA, Tensor& dB) {
1733 16 const auto& v = detail::dispatch(A, B, dC, dA, dB);
1734
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.matmul_backward) detail::throw_not_implemented("matmul_backward", A.device);
1735 16 detail::adopt_output(dA, A.device);
1736 16 detail::adopt_output(dB, A.device);
1737 16 v.matmul_backward(A, B, dC, dA, dB);
1738 16 }
1739
1740 907 void lstm_forward_train(const Tensor& X, const Tensor& W_ih, const Tensor& W_hh,
1741 const Tensor* b_ih, const Tensor* b_hh,
1742 const Tensor* h0, const Tensor* c0, int T, int B,
1743 Tensor& Y, Tensor& gates, Tensor& C,
1744 Tensor* hT, Tensor* cT) {
1745 907 const auto& v = detail::dispatch(X);
1746
1/2
✓ Branch 0 taken 907 times.
✗ Branch 1 not taken.
 907 if (!v.lstm_forward_train) detail::throw_not_implemented("lstm_forward_train", X.device);
1747 907 detail::adopt_output(Y, X.device);
1748 907 detail::adopt_output(gates, X.device);
1749 907 detail::adopt_output(C, X.device);
1750
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 12 times.
 907 if (hT) detail::adopt_output(*hT, X.device);
1751
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 12 times.
 907 if (cT) detail::adopt_output(*cT, X.device);
1752 907 v.lstm_forward_train(X, W_ih, W_hh, b_ih, b_hh, h0, c0, T, B, Y, gates, C, hT, cT);
1753 907 }
1754
1755 413 void lstm_backward(const Tensor& X, const Tensor& W_ih, const Tensor& W_hh,
1756 const Tensor* h0, const Tensor* c0,
1757 const Tensor& Y, const Tensor& gates, const Tensor& C,
1758 const Tensor& dY, int T, int B,
1759 Tensor& dX, Tensor& dW_ih, Tensor& dW_hh,
1760 Tensor* db_ih, Tensor* db_hh,
1761 Tensor* dh0, Tensor* dc0) {
1762 413 const auto& v = detail::dispatch(X);
1763
1/2
✓ Branch 0 taken 413 times.
✗ Branch 1 not taken.
 413 if (!v.lstm_backward) detail::throw_not_implemented("lstm_backward", X.device);
1764 413 detail::adopt_output(dX, X.device);
1765 413 detail::adopt_output(dW_ih, X.device);
1766 413 detail::adopt_output(dW_hh, X.device);
1767
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 413 times.
 413 if (db_ih) detail::adopt_output(*db_ih, X.device);
1768
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 413 times.
 413 if (db_hh) detail::adopt_output(*db_hh, X.device);
1769
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 13 times.
 413 if (dh0) detail::adopt_output(*dh0, X.device);
1770
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 13 times.
 413 if (dc0) detail::adopt_output(*dc0, X.device);
1771 826 v.lstm_backward(X, W_ih, W_hh, h0, c0, Y, gates, C, dY, T, B,
1772 413 dX, dW_ih, dW_hh, db_ih, db_hh, dh0, dc0);
1773 413 }
1774
1775 14 void rope_forward(const Tensor& X, int head_dim, int num_heads,
1776 int seq_offset, float theta_base, Tensor& Y) {
1777 14 const auto& v = detail::dispatch(X, Y);
1778
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.rope_forward) detail::throw_not_implemented("rope_forward", X.device);
1779 14 detail::adopt_output(Y, X.device);
1780 14 v.rope_forward(X, head_dim, num_heads, seq_offset, theta_base, Y);
1781 14 }
1782
1783 8 void rope_backward(const Tensor& dY, int head_dim, int num_heads,
1784 int seq_offset, float theta_base, Tensor& dX) {
1785 8 const auto& v = detail::dispatch(dY, dX);
1786
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.rope_backward) detail::throw_not_implemented("rope_backward", dY.device);
1787 8 detail::adopt_output(dX, dY.device);
1788 8 v.rope_backward(dY, head_dim, num_heads, seq_offset, theta_base, dX);
1789 8 }
1790
1791 26 void rope_apply(const Tensor& X, const Tensor& cos_tbl, const Tensor& sin_tbl,
1792 int head_dim, int num_heads, Tensor& Y) {
1793 26 const auto& v = detail::dispatch(X, cos_tbl, sin_tbl, Y);
1794
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.rope_apply) detail::throw_not_implemented("rope_apply", X.device);
1795 26 detail::adopt_output(Y, X.device);
1796 26 v.rope_apply(X, cos_tbl, sin_tbl, head_dim, num_heads, Y);
1797 26 }
1798
1799 2 void rope_apply_perhead(const Tensor& X, const Tensor& cos_tbl,
1800 const Tensor& sin_tbl, int head_dim, int num_heads,
1801 Tensor& Y) {
1802 2 const auto& v = detail::dispatch(X, cos_tbl, sin_tbl, Y);
1803
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.rope_apply_perhead)
1804 ✗ detail::throw_not_implemented("rope_apply_perhead", X.device);
1805 2 detail::adopt_output(Y, X.device);
1806 2 v.rope_apply_perhead(X, cos_tbl, sin_tbl, head_dim, num_heads, Y);
1807 2 }
1808
1809 4 void rope_apply_backward(const Tensor& dY, const Tensor& cos_tbl,
1810 const Tensor& sin_tbl, int head_dim, int num_heads,
1811 Tensor& dX) {
1812 4 const auto& v = detail::dispatch(dY, cos_tbl, sin_tbl, dX);
1813
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rope_apply_backward)
1814 ✗ detail::throw_not_implemented("rope_apply_backward", dY.device);
1815 4 detail::adopt_output(dX, dY.device);
1816 4 v.rope_apply_backward(dY, cos_tbl, sin_tbl, head_dim, num_heads, dX);
1817 4 }
1818
1819 11 void rope_apply_mrope(const Tensor& X,
1820 const Tensor& cos_t, const Tensor& sin_t,
1821 const Tensor& cos_h, const Tensor& sin_h,
1822 const Tensor& cos_w, const Tensor& sin_w,
1823 const int32_t* pos_t, const int32_t* pos_h,
1824 const int32_t* pos_w,
1825 int head_dim, int num_heads,
1826 int d_t, int d_h, int d_w,
1827 Tensor& Y) {
1828 22 const auto& v = detail::dispatch(X, cos_t, sin_t, cos_h, sin_h, cos_w,
1829 11 sin_w, Y);
1830
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.rope_apply_mrope)
1831 ✗ detail::throw_not_implemented("rope_apply_mrope", X.device);
1832 11 detail::adopt_output(Y, X.device);
1833 22 v.rope_apply_mrope(X, cos_t, sin_t, cos_h, sin_h, cos_w, sin_w,
1834 11 pos_t, pos_h, pos_w, head_dim, num_heads,
1835 11 d_t, d_h, d_w, Y);
1836 11 }
1837
1838 18 void rms_norm_forward(const Tensor& X, const Tensor& gamma, float eps, Tensor& Y) {
1839 18 const auto& v = detail::dispatch(X, gamma, Y);
1840
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.rms_norm_forward) detail::throw_not_implemented("rms_norm_forward", X.device);
1841 18 detail::adopt_output(Y, X.device);
1842 18 v.rms_norm_forward(X, gamma, eps, Y);
1843 18 }
1844
1845 14 void rms_norm_backward(const Tensor& X, const Tensor& gamma, const Tensor& dY,
1846 float eps, Tensor& dX, Tensor& dGamma) {
1847 14 const auto& v = detail::dispatch(X, gamma, dY, dX, dGamma);
1848
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.rms_norm_backward) detail::throw_not_implemented("rms_norm_backward", X.device);
1849 14 detail::adopt_output(dX, X.device);
1850 14 detail::adopt_output(dGamma, X.device);
1851 14 v.rms_norm_backward(X, gamma, dY, eps, dX, dGamma);
1852 14 }
1853
1854 10 void swiglu_forward(const Tensor& X, Tensor& Y) {
1855 10 const auto& v = detail::dispatch(X, Y);
1856
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.swiglu_forward) detail::throw_not_implemented("swiglu_forward", X.device);
1857 10 detail::adopt_output(Y, X.device);
1858 10 v.swiglu_forward(X, Y);
1859 10 }
1860 10 void swiglu_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
1861 10 const auto& v = detail::dispatch(X, dY, dX);
1862
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.swiglu_backward) detail::throw_not_implemented("swiglu_backward", X.device);
1863 10 detail::adopt_output(dX, X.device);
1864 10 v.swiglu_backward(X, dY, dX);
1865 10 }
1866
1867 28 void kv_cache_append(const Tensor& K_new, const Tensor& V_new, int cur_len,
1868 Tensor& K_cache, Tensor& V_cache) {
1869 28 const auto& v = detail::dispatch(K_new, V_new, K_cache, V_cache);
1870
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 if (!v.kv_cache_append) detail::throw_not_implemented("kv_cache_append", K_new.device);
1871 28 detail::adopt_output(K_cache, K_new.device);
1872 28 detail::adopt_output(V_cache, K_new.device);
1873 28 v.kv_cache_append(K_new, V_new, cur_len, K_cache, V_cache);
1874 28 }
1875
1876 51 void flash_attention_decode(const Tensor& Q,
1877 const Tensor& K_cache, const Tensor& V_cache,
1878 int valid_len, int num_q_heads, int num_kv_heads,
1879 Tensor& O, float attn_softcap, int window) {
1880 51 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
1881
1/2
✓ Branch 0 taken 51 times.
✗ Branch 1 not taken.
 51 if (!v.flash_attention_decode) detail::throw_not_implemented("flash_attention_decode", Q.device);
1882 51 detail::adopt_output(O, Q.device);
1883 102 v.flash_attention_decode(Q, K_cache, V_cache, valid_len,
1884 51 num_q_heads, num_kv_heads, O, attn_softcap, window);
1885 51 }
1886
1887 12 void flash_attention_decode_masked(const Tensor& Q,
1888 const Tensor& K_cache,
1889 const Tensor& V_cache,
1890 const float* d_mask,
1891 int num_q_heads, int num_kv_heads,
1892 Tensor& O, float attn_softcap, int window) {
1893 12 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
1894
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.flash_attention_decode_masked) {
1895 ✗ detail::throw_not_implemented("flash_attention_decode_masked", Q.device);
1896 }
1897 12 detail::adopt_output(O, Q.device);
1898 24 v.flash_attention_decode_masked(Q, K_cache, V_cache, d_mask,
1899 12 num_q_heads, num_kv_heads, O,
1900 12 attn_softcap, window);
1901 12 }
1902
1903 // ─── Public reductions ─────────────────────────────────────────────────────
1904
1905 15 void sum_rows(const Tensor& X, Tensor& Y) {
1906 15 const auto& v = detail::dispatch(X, Y);
1907
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.sum_rows) detail::throw_not_implemented("sum_rows", X.device);
1908 15 detail::adopt_output(Y, X.device);
1909 15 v.sum_rows(X, Y);
1910 15 }
1911 14 void sum_cols(const Tensor& X, Tensor& Y) {
1912 14 const auto& v = detail::dispatch(X, Y);
1913
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.sum_cols) detail::throw_not_implemented("sum_cols", X.device);
1914 14 detail::adopt_output(Y, X.device);
1915 14 v.sum_cols(X, Y);
1916 14 }
1917 17 void argmax_rows(const Tensor& X, Tensor& Idx) {
1918 17 const auto& v = detail::dispatch(X, Idx);
1919
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.argmax_rows) detail::throw_not_implemented("argmax_rows", X.device);
1920 17 detail::adopt_output(Idx, X.device);
1921 17 v.argmax_rows(X, Idx);
1922 17 }
1923 11 void rows_count_above(const Tensor& X, float t_lo, float t_hi, Tensor& counts) {
1924 11 const auto& v = detail::dispatch(X, counts);
1925
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.rows_count_above)
1926 ✗ detail::throw_not_implemented("rows_count_above", X.device);
1927 11 detail::adopt_output(counts, X.device);
1928 11 v.rows_count_above(X, t_lo, t_hi, counts);
1929 11 }
1930 9 void threshold_u8(const Tensor& X, float t, Tensor& Y) {
1931 9 const auto& v = detail::dispatch(X, Y);
1932
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.threshold_u8) detail::throw_not_implemented("threshold_u8", X.device);
1933 9 detail::adopt_output(Y, X.device);
1934 9 v.threshold_u8(X, t, Y);
1935 9 }
1936
1937 // ─── Diffusion sampler steps + timestep embedding ──────────────────────────
1938
1939 14 void ddim_step(const Tensor& x_t, const Tensor& eps_pred,
1940 float alpha_t, float alpha_prev, float sigma_t, Tensor& x_prev) {
1941 14 const auto& v = detail::dispatch(x_t, eps_pred, x_prev);
1942
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.ddim_step) detail::throw_not_implemented("ddim_step", x_t.device);
1943 14 detail::adopt_output(x_prev, x_t.device);
1944 14 v.ddim_step(x_t, eps_pred, alpha_t, alpha_prev, sigma_t, x_prev);
1945 14 }
1946
1947 12 void euler_step(const Tensor& x_t, const Tensor& eps_pred,
1948 float sigma_t, float sigma_prev, Tensor& x_prev) {
1949 12 const auto& v = detail::dispatch(x_t, eps_pred, x_prev);
1950
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.euler_step) detail::throw_not_implemented("euler_step", x_t.device);
1951 12 detail::adopt_output(x_prev, x_t.device);
1952 12 v.euler_step(x_t, eps_pred, sigma_t, sigma_prev, x_prev);
1953 12 }
1954
1955 12 void dpmpp_2m_step(const Tensor& x_t, const Tensor& eps_pred,
1956 const Tensor& x0_prev, float sigma_t,
1957 float c_xt, float c_x0t, float c_x0prev,
1958 Tensor& x_prev, Tensor& x0_out) {
1959 12 const auto& v = detail::dispatch(x_t, eps_pred, x0_prev, x_prev, x0_out);
1960
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.dpmpp_2m_step) detail::throw_not_implemented("dpmpp_2m_step", x_t.device);
1961 12 detail::adopt_output(x_prev, x_t.device);
1962 12 detail::adopt_output(x0_out, x_t.device);
1963 24 v.dpmpp_2m_step(x_t, eps_pred, x0_prev, sigma_t, c_xt, c_x0t, c_x0prev,
1964 12 x_prev, x0_out);
1965 12 }
1966
1967 8 void timestep_embedding(const Tensor& timesteps, int dim, float max_period,
1968 Tensor& Y) {
1969 8 const auto& v = detail::dispatch(timesteps, Y);
1970
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.timestep_embedding) detail::throw_not_implemented("timestep_embedding", timesteps.device);
1971 8 detail::adopt_output(Y, timesteps.device);
1972 8 v.timestep_embedding(timesteps, dim, max_period, Y);
1973 8 }
1974
1975 // ─── INT8 weight-only quantisation (W8A16) ─────────────────────────────────
1976
1977 // Host helper — pure host buffers, no device dispatch. Rows are independent,
1978 // so they fan out over the CPU thread pool: the multi-GB checkpoint quantise
1979 // during a quantized model load is this function's dominant caller, and the
1980 // serial version left it minutes-long on 10B+-parameter models.
1981 20 void quantize_int8_per_row_host(const uint16_t* W_fp16,
1982 int out, int in,
1983 int8_t* W_int8_out,
1984 float* scales_out) {
1985
2/4
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
 20 if (out <= 0 || in <= 0) {
1986 ✗ for (int r = 0; r < out; ++r) scales_out[r] = 0.0f;
1987 ✗ return;
1988 }
1989
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 848 detail::cpu::parallel_for(static_cast<std::size_t>(out), [&](std::size_t r) {
1990 828 const uint16_t* row = W_fp16 + r * static_cast<std::size_t>(in);
1991 828 float amax = 0.0f;
1992
2/2
✓ Branch 0 taken 32568 times.
✓ Branch 1 taken 828 times.
 33396 for (int c = 0; c < in; ++c) {
1993 32568 const float v = fp16_bits_to_fp32(row[c]);
1994 32568 const float a = std::fabs(v);
1995
2/2
✓ Branch 0 taken 29051 times.
✓ Branch 1 taken 3517 times.
 32568 if (a > amax) amax = a;
1996 32568 }
1997
1/2
✓ Branch 0 taken 828 times.
✗ Branch 1 not taken.
 828 const float scale = (amax > 0.0f) ? (amax / 127.0f) : 0.0f;
1998
1/2
✓ Branch 0 taken 828 times.
✗ Branch 1 not taken.
 828 const float inv = (scale > 0.0f) ? (1.0f / scale) : 0.0f;
1999 828 scales_out[r] = scale;
2000 828 int8_t* dst = W_int8_out + r * static_cast<std::size_t>(in);
2001
2/2
✓ Branch 0 taken 33792 times.
✓ Branch 1 taken 828 times.
 34620 for (int c = 0; c < in; ++c) {
2002 33792 const float v = fp16_bits_to_fp32(row[c]);
2003 33792 int q = static_cast<int>(std::lrint(v * inv));
2004
1/2
✓ Branch 0 taken 33792 times.
✗ Branch 1 not taken.
 33792 if (q < -127) q = -127;
2005
1/2
✓ Branch 0 taken 33792 times.
✗ Branch 1 not taken.
 33792 if (q > 127) q = 127;
2006 33792 dst[c] = static_cast<int8_t>(q);
2007 33792 }
2008 828 });
2009 20 }
2010
2011 1 void matmul_int8w_fp16(const Tensor& W_int8, const Tensor& scales,
2012 const Tensor& X, Tensor& Y) {
2013 1 const auto& v = detail::dispatch(W_int8, scales, X, Y);
2014
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 if (!v.matmul_int8w_fp16) detail::throw_not_implemented("matmul_int8w_fp16", W_int8.device);
2015 ✗ detail::adopt_output(Y, W_int8.device);
2016 ✗ v.matmul_int8w_fp16(W_int8, scales, X, Y);
2017 ✗ }
2018
2019 ✗ void conv2d_int8w_fp16_forward(const Tensor& X,
2020 const Tensor& W_int8, const Tensor& scales,
2021 const Tensor* bias,
2022 int N, int C_in, int H, int W,
2023 int C_out, int kH, int kW,
2024 int stride_h, int stride_w,
2025 int pad_h, int pad_w,
2026 int dil_h, int dil_w, int groups,
2027 Tensor& Y) {
2028 ✗ const auto& v = detail::dispatch_with_opts(X, W_int8, {&scales, bias, &Y});
2029 ✗ if (!v.conv2d_int8w_fp16_forward)
2030 ✗ detail::throw_not_implemented("conv2d_int8w_fp16_forward", X.device);
2031 ✗ detail::adopt_output(Y, X.device);
2032 ✗ v.conv2d_int8w_fp16_forward(X, W_int8, scales, bias,
2033 ✗ N, C_in, H, W, C_out, kH, kW,
2034 ✗ stride_h, stride_w, pad_h, pad_w,
2035 ✗ dil_h, dil_w, groups, Y);
2036 ✗ }
2037
2038 ✗ void linear_forward_batched_int8w_fp16(const Tensor& W_int8,
2039 const Tensor& scales,
2040 const Tensor* bias,
2041 const Tensor& X_BD, Tensor& Y_BD) {
2042 ✗ const auto& v = detail::dispatch_with_opts(W_int8, scales, {bias, &X_BD, &Y_BD});
2043 ✗ if (!v.linear_forward_batched_int8w_fp16)
2044 ✗ detail::throw_not_implemented("linear_forward_batched_int8w_fp16", W_int8.device);
2045 ✗ detail::adopt_output(Y_BD, W_int8.device);
2046 ✗ v.linear_forward_batched_int8w_fp16(W_int8, scales, bias, X_BD, Y_BD);
2047 ✗ }
2048
2049 2 void dequant_q4k_to_fp16(const Tensor& W_q4k, Tensor& W_fp16) {
2050 2 const auto& v = detail::dispatch(W_q4k, W_fp16);
2051
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 2 if (!v.dequant_q4k_to_fp16)
2052 1 detail::throw_not_implemented("dequant_q4k_to_fp16", W_q4k.device);
2053 1 detail::adopt_output(W_fp16, W_q4k.device);
2054 1 v.dequant_q4k_to_fp16(W_q4k, W_fp16);
2055 1 }
2056
2057 2 void linear_forward_q4k_fp16(const Tensor& W_q4k, const Tensor* bias,
2058 const Tensor& x, Tensor& y) {
2059 2 const auto& v = detail::dispatch_with_opts(W_q4k, {bias, &x, &y});
2060
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.linear_forward_q4k_fp16)
2061 ✗ detail::throw_not_implemented("linear_forward_q4k_fp16", W_q4k.device);
2062 2 detail::adopt_output(y, W_q4k.device);
2063 2 v.linear_forward_q4k_fp16(W_q4k, bias, x, y);
2064 2 }
2065
2066 1 void linear_forward_batched_q4k_fp16(const Tensor& W_q4k, const Tensor* bias,
2067 const Tensor& X_BD, Tensor& Y_BD) {
2068 1 const auto& v = detail::dispatch_with_opts(W_q4k, {bias, &X_BD, &Y_BD});
2069
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (!v.linear_forward_batched_q4k_fp16)
2070 ✗ detail::throw_not_implemented("linear_forward_batched_q4k_fp16", W_q4k.device);
2071 1 detail::adopt_output(Y_BD, W_q4k.device);
2072 1 v.linear_forward_batched_q4k_fp16(W_q4k, bias, X_BD, Y_BD);
2073 1 }
2074
2075 4 void dequant_q8_0_to_fp16(const Tensor& W_q8, Tensor& W_fp16) {
2076 4 const auto& v = detail::dispatch(W_q8, W_fp16);
2077
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.dequant_q8_0_to_fp16)
2078 ✗ detail::throw_not_implemented("dequant_q8_0_to_fp16", W_q8.device);
2079 4 detail::adopt_output(W_fp16, W_q8.device);
2080 4 v.dequant_q8_0_to_fp16(W_q8, W_fp16);
2081 4 }
2082
2083 2 void linear_forward_q8_0_fp16(const Tensor& W_q8, const Tensor* bias,
2084 const Tensor& x, Tensor& y) {
2085 2 const auto& v = detail::dispatch_with_opts(W_q8, {bias, &x, &y});
2086
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.linear_forward_q8_0_fp16)
2087 ✗ detail::throw_not_implemented("linear_forward_q8_0_fp16", W_q8.device);
2088 2 detail::adopt_output(y, W_q8.device);
2089 2 v.linear_forward_q8_0_fp16(W_q8, bias, x, y);
2090 2 }
2091
2092 1 void linear_forward_batched_q8_0_fp16(const Tensor& W_q8, const Tensor* bias,
2093 const Tensor& X_BD, Tensor& Y_BD) {
2094 1 const auto& v = detail::dispatch_with_opts(W_q8, {bias, &X_BD, &Y_BD});
2095
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (!v.linear_forward_batched_q8_0_fp16)
2096 ✗ detail::throw_not_implemented("linear_forward_batched_q8_0_fp16", W_q8.device);
2097 1 detail::adopt_output(Y_BD, W_q8.device);
2098 1 v.linear_forward_batched_q8_0_fp16(W_q8, bias, X_BD, Y_BD);
2099 1 }
2100
2101 1 void dequant_q6k_to_fp16(const Tensor& W_q6k, Tensor& W_fp16) {
2102 1 const auto& v = detail::dispatch(W_q6k, W_fp16);
2103
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (!v.dequant_q6k_to_fp16)
2104 ✗ detail::throw_not_implemented("dequant_q6k_to_fp16", W_q6k.device);
2105 1 detail::adopt_output(W_fp16, W_q6k.device);
2106 1 v.dequant_q6k_to_fp16(W_q6k, W_fp16);
2107 1 }
2108
2109 2 void linear_forward_q6k_fp16(const Tensor& W_q6k, const Tensor* bias,
2110 const Tensor& x, Tensor& y) {
2111 2 const auto& v = detail::dispatch_with_opts(W_q6k, {bias, &x, &y});
2112
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.linear_forward_q6k_fp16)
2113 ✗ detail::throw_not_implemented("linear_forward_q6k_fp16", W_q6k.device);
2114 2 detail::adopt_output(y, W_q6k.device);
2115 2 v.linear_forward_q6k_fp16(W_q6k, bias, x, y);
2116 2 }
2117
2118 1 void linear_forward_batched_q6k_fp16(const Tensor& W_q6k, const Tensor* bias,
2119 const Tensor& X_BD, Tensor& Y_BD) {
2120 1 const auto& v = detail::dispatch_with_opts(W_q6k, {bias, &X_BD, &Y_BD});
2121
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (!v.linear_forward_batched_q6k_fp16)
2122 ✗ detail::throw_not_implemented("linear_forward_batched_q6k_fp16", W_q6k.device);
2123 1 detail::adopt_output(Y_BD, W_q6k.device);
2124 1 v.linear_forward_batched_q6k_fp16(W_q6k, bias, X_BD, Y_BD);
2125 1 }
2126
2127 ✗ void flash_attention_project_kv_int8w_fp16(const Tensor& ctx,
2128 const Tensor& Wk_int8,
2129 const Tensor& sk,
2130 const Tensor* bk,
2131 const Tensor& Wv_int8,
2132 const Tensor& sv,
2133 const Tensor* bv,
2134 Tensor& K_out, Tensor& V_out) {
2135 ✗ const auto& v = detail::dispatch_with_opts(
2136 ✗ ctx, Wk_int8, {&sk, bk, &Wv_int8, &sv, bv, &K_out, &V_out});
2137 ✗ if (!v.flash_attention_project_kv_int8w_fp16)
2138 ✗ detail::throw_not_implemented("flash_attention_project_kv_int8w_fp16", ctx.device);
2139 ✗ detail::adopt_output(K_out, ctx.device);
2140 ✗ detail::adopt_output(V_out, ctx.device);
2141 ✗ v.flash_attention_project_kv_int8w_fp16(ctx, Wk_int8, sk, bk,
2142 ✗ Wv_int8, sv, bv, K_out, V_out);
2143 ✗ }
2144
2145 ✗ void flash_attention_q_with_kv_cached_int8w_fp16(const Tensor& X,
2146 const Tensor& K, const Tensor& V,
2147 const Tensor& Wq_int8,
2148 const Tensor& sq,
2149 const Tensor* bq,
2150 const Tensor& Wo_int8,
2151 const Tensor& so,
2152 const Tensor* bo,
2153 const float* d_mask,
2154 int num_heads, bool causal,
2155 Tensor& O) {
2156 ✗ const auto& v = detail::dispatch_with_opts(
2157 ✗ X, K, {&V, &Wq_int8, &sq, bq, &Wo_int8, &so, bo, &O});
2158 ✗ if (!v.flash_attention_q_with_kv_cached_int8w_fp16)
2159 ✗ detail::throw_not_implemented("flash_attention_q_with_kv_cached_int8w_fp16", X.device);
2160 ✗ detail::adopt_output(O, X.device);
2161 ✗ v.flash_attention_q_with_kv_cached_int8w_fp16(X, K, V, Wq_int8, sq, bq,
2162 ✗ Wo_int8, so, bo,
2163 ✗ d_mask, num_heads, causal, O);
2164 ✗ }
2165
2166 ✗ void flash_attention_qkvo_int8w_fp16(const Tensor& X, const Tensor* Ctx,
2167 const Tensor& Wq_int8, const Tensor& sq, const Tensor* bq,
2168 const Tensor& Wk_int8, const Tensor& sk, const Tensor* bk,
2169 const Tensor& Wv_int8, const Tensor& sv, const Tensor* bv,
2170 const Tensor& Wo_int8, const Tensor& so, const Tensor* bo,
2171 const float* d_mask, int num_heads, bool causal,
2172 Tensor& O) {
2173 ✗ const auto& v = detail::dispatch_with_opts(
2174 ✗ X, Wq_int8, {Ctx, &sq, bq, &Wk_int8, &sk, bk, &Wv_int8, &sv, bv,
2175 ✗ &Wo_int8, &so, bo, &O});
2176 ✗ if (!v.flash_attention_qkvo_int8w_fp16)
2177 ✗ detail::throw_not_implemented("flash_attention_qkvo_int8w_fp16", X.device);
2178 ✗ detail::adopt_output(O, X.device);
2179 ✗ v.flash_attention_qkvo_int8w_fp16(X, Ctx, Wq_int8, sq, bq,
2180 ✗ Wk_int8, sk, bk, Wv_int8, sv, bv,
2181 ✗ Wo_int8, so, bo,
2182 ✗ d_mask, num_heads, causal, O);
2183 ✗ }
2184
2185 // ─── Spectral / FFT core (brosoundml) ──────────────────────────────────────
2186
2187 54 void complex_mul(const Tensor& a, const Tensor& b, Tensor& y) {
2188 54 const auto& v = detail::dispatch(a, b, y);
2189
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 if (!v.complex_mul) detail::throw_not_implemented("complex_mul", a.device);
2190 54 detail::adopt_output(y, a.device);
2191 54 v.complex_mul(a, b, y);
2192 54 }
2193 5 void complex_mul_backward(const Tensor& a, const Tensor& b, const Tensor& dY,
2194 Tensor& dA, Tensor& dB) {
2195 5 const auto& v = detail::dispatch(a, b, dY, dA, dB);
2196
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.complex_mul_backward)
2197 ✗ detail::throw_not_implemented("complex_mul_backward", a.device);
2198 5 detail::adopt_output(dA, a.device);
2199 5 detail::adopt_output(dB, a.device);
2200 5 v.complex_mul_backward(a, b, dY, dA, dB);
2201 5 }
2202 36 void complex_abs(const Tensor& z, Tensor& y) {
2203 36 const auto& v = detail::dispatch(z, y);
2204
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (!v.complex_abs) detail::throw_not_implemented("complex_abs", z.device);
2205 36 detail::adopt_output(y, z.device);
2206 36 v.complex_abs(z, y);
2207 36 }
2208 3 void complex_abs_backward(const Tensor& z, const Tensor& dY, Tensor& dZ) {
2209 3 const auto& v = detail::dispatch(z, dY, dZ);
2210
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.complex_abs_backward)
2211 ✗ detail::throw_not_implemented("complex_abs_backward", z.device);
2212 3 detail::adopt_output(dZ, z.device);
2213 3 v.complex_abs_backward(z, dY, dZ);
2214 3 }
2215 3 void complex_angle(const Tensor& z, Tensor& y) {
2216 3 const auto& v = detail::dispatch(z, y);
2217
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.complex_angle) detail::throw_not_implemented("complex_angle", z.device);
2218 3 detail::adopt_output(y, z.device);
2219 3 v.complex_angle(z, y);
2220 3 }
2221 3 void complex_from_polar(const Tensor& mag, const Tensor& phase, Tensor& y) {
2222 3 const auto& v = detail::dispatch(mag, phase, y);
2223
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.complex_from_polar)
2224 ✗ detail::throw_not_implemented("complex_from_polar", mag.device);
2225 3 detail::adopt_output(y, mag.device);
2226 3 v.complex_from_polar(mag, phase, y);
2227 3 }
2228 339 void fft(const Tensor& x, Tensor& y) {
2229 339 const auto& v = detail::dispatch(x, y);
2230
1/2
✓ Branch 0 taken 339 times.
✗ Branch 1 not taken.
 339 if (!v.fft) detail::throw_not_implemented("fft", x.device);
2231 339 detail::adopt_output(y, x.device);
2232 339 v.fft(x, y);
2233 339 }
2234 23 void ifft(const Tensor& x, Tensor& y) {
2235 23 const auto& v = detail::dispatch(x, y);
2236
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.ifft) detail::throw_not_implemented("ifft", x.device);
2237 23 detail::adopt_output(y, x.device);
2238 23 v.ifft(x, y);
2239 23 }
2240 335 void rfft(const Tensor& x, Tensor& y) {
2241 335 const auto& v = detail::dispatch(x, y);
2242
1/2
✓ Branch 0 taken 335 times.
✗ Branch 1 not taken.
 335 if (!v.rfft) detail::throw_not_implemented("rfft", x.device);
2243 335 detail::adopt_output(y, x.device);
2244 335 v.rfft(x, y);
2245 335 }
2246 367 void irfft(const Tensor& x, int L, Tensor& y) {
2247 367 const auto& v = detail::dispatch(x, y);
2248
1/2
✓ Branch 0 taken 367 times.
✗ Branch 1 not taken.
 367 if (!v.irfft) detail::throw_not_implemented("irfft", x.device);
2249 367 detail::adopt_output(y, x.device);
2250 367 v.irfft(x, L, y);
2251 367 }
2252 9 void rfft_backward(const Tensor& dY, int L, Tensor& dX) {
2253 9 const auto& v = detail::dispatch(dY, dX);
2254
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.rfft_backward) detail::throw_not_implemented("rfft_backward", dY.device);
2255 9 detail::adopt_output(dX, dY.device);
2256 9 v.rfft_backward(dY, L, dX);
2257 9 }
2258 9 void irfft_backward(const Tensor& dY, Tensor& dX) {
2259 9 const auto& v = detail::dispatch(dY, dX);
2260
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.irfft_backward)
2261 ✗ detail::throw_not_implemented("irfft_backward", dY.device);
2262 9 detail::adopt_output(dX, dY.device);
2263 9 v.irfft_backward(dY, dX);
2264 9 }
2265
2266 1290 void stft(const Tensor& signal, const Tensor& window,
2267 int N, int n_fft, int hop_length, int win_length,
2268 bool center, bool normalized, Tensor& spec) {
2269 1290 const auto& v = detail::dispatch(signal, window, spec);
2270
1/2
✓ Branch 0 taken 1290 times.
✗ Branch 1 not taken.
 1290 if (!v.stft) detail::throw_not_implemented("stft", signal.device);
2271 1290 detail::adopt_output(spec, signal.device);
2272 2580 v.stft(signal, window, N, n_fft, hop_length, win_length,
2273 1290 center, normalized, spec);
2274 1290 }
2275 14 void stft_backward(const Tensor& dSpec, const Tensor& window,
2276 int N, int signal_len, int n_fft, int hop_length,
2277 int win_length, bool center, bool normalized,
2278 Tensor& dSignal) {
2279 14 const auto& v = detail::dispatch(dSpec, window, dSignal);
2280
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.stft_backward)
2281 ✗ detail::throw_not_implemented("stft_backward", dSpec.device);
2282 14 detail::adopt_output(dSignal, dSpec.device);
2283 28 v.stft_backward(dSpec, window, N, signal_len, n_fft, hop_length,
2284 14 win_length, center, normalized, dSignal);
2285 14 }
2286 5114 void istft(const Tensor& spec, const Tensor& window,
2287 int N, int signal_len, int n_fft, int hop_length, int win_length,
2288 bool center, bool normalized, Tensor& signal) {
2289 5114 const auto& v = detail::dispatch(spec, window, signal);
2290
1/2
✓ Branch 0 taken 5114 times.
✗ Branch 1 not taken.
 5114 if (!v.istft) detail::throw_not_implemented("istft", spec.device);
2291 5114 detail::adopt_output(signal, spec.device);
2292 10228 v.istft(spec, window, N, signal_len, n_fft, hop_length, win_length,
2293 5114 center, normalized, signal);
2294 5114 }
2295 14 void istft_backward(const Tensor& dSignal, const Tensor& window,
2296 int N, int signal_len, int n_fft, int hop_length,
2297 int win_length, bool center, bool normalized,
2298 Tensor& dSpec) {
2299 14 const auto& v = detail::dispatch(dSignal, window, dSpec);
2300
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.istft_backward)
2301 ✗ detail::throw_not_implemented("istft_backward", dSignal.device);
2302 14 detail::adopt_output(dSpec, dSignal.device);
2303 28 v.istft_backward(dSignal, window, N, signal_len, n_fft, hop_length,
2304 14 win_length, center, normalized, dSpec);
2305 14 }
2306
2307 // ─── 1D convolution family (brosoundml) ────────────────────────────────────
2308 //
2309 // conv1d / conv1d backward / conv1d_int8w_fp16 / causal_conv1d are header-only
2310 // inline wrappers in ops.h (they forward to the conv2d ops); only the three
2311 // genuinely new ops below have dispatcher wrappers.
2312
2313 144 void conv_transpose1d_forward(const Tensor& X, const Tensor& Wt,
2314 const Tensor* bias,
2315 int N, int C_in, int L, int C_out, int kL,
2316 int stride, int padding, int output_padding,
2317 int dilation, int groups, Tensor& Y) {
2318 144 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
2319
1/2
✓ Branch 0 taken 144 times.
✗ Branch 1 not taken.
 144 if (!v.conv_transpose1d_forward)
2320 ✗ detail::throw_not_implemented("conv_transpose1d_forward", X.device);
2321 144 detail::adopt_output(Y, X.device);
2322 288 v.conv_transpose1d_forward(X, Wt, bias, N, C_in, L, C_out, kL, stride,
2323 144 padding, output_padding, dilation, groups, Y);
2324 144 }
2325
2326 12 void conv_transpose1d_backward_input(const Tensor& Wt, const Tensor& dY,
2327 int N, int C_in, int L, int C_out, int kL,
2328 int stride, int padding,
2329 int output_padding, int dilation,
2330 int groups, Tensor& dX) {
2331 12 const auto& v = detail::dispatch(Wt, dY, dX);
2332
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.conv_transpose1d_backward_input)
2333 ✗ detail::throw_not_implemented("conv_transpose1d_backward_input", Wt.device);
2334 12 detail::adopt_output(dX, Wt.device);
2335 24 v.conv_transpose1d_backward_input(Wt, dY, N, C_in, L, C_out, kL, stride,
2336 12 padding, output_padding, dilation,
2337 12 groups, dX);
2338 12 }
2339
2340 10 void conv_transpose1d_backward_weight(const Tensor& X, const Tensor& dY,
2341 int N, int C_in, int L, int C_out, int kL,
2342 int stride, int padding,
2343 int output_padding, int dilation,
2344 int groups, Tensor& dWt) {
2345 10 const auto& v = detail::dispatch(X, dY, dWt);
2346
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.conv_transpose1d_backward_weight)
2347 ✗ detail::throw_not_implemented("conv_transpose1d_backward_weight", X.device);
2348 10 detail::adopt_output(dWt, X.device);
2349 20 v.conv_transpose1d_backward_weight(X, dY, N, C_in, L, C_out, kL, stride,
2350 10 padding, output_padding, dilation,
2351 10 groups, dWt);
2352 10 }
2353
2354 8 void conv_transpose1d_backward_bias(const Tensor& dY, int N, int C_out,
2355 int L_out, Tensor& dB) {
2356 8 const auto& v = detail::dispatch(dY, dB);
2357
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.conv_transpose1d_backward_bias)
2358 ✗ detail::throw_not_implemented("conv_transpose1d_backward_bias", dY.device);
2359 8 detail::adopt_output(dB, dY.device);
2360 8 v.conv_transpose1d_backward_bias(dY, N, C_out, L_out, dB);
2361 8 }
2362
2363 26 void causal_conv1d_update(const Tensor& X, const Tensor& Wt, const Tensor* bias,
2364 int N, int C, int L_step, int kL, int dilation,
2365 Tensor& state, Tensor& Y) {
2366 26 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &state, &Y});
2367
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.causal_conv1d_update)
2368 ✗ detail::throw_not_implemented("causal_conv1d_update", X.device);
2369 26 detail::adopt_output(state, X.device);
2370 26 detail::adopt_output(Y, X.device);
2371 26 v.causal_conv1d_update(X, Wt, bias, N, C, L_step, kL, dilation, state, Y);
2372 26 }
2373
2374 157 void pad1d_forward(const Tensor& X, int N, int C, int L,
2375 int pad_left, int pad_right, int mode, Tensor& Y) {
2376 157 const auto& v = detail::dispatch(X, Y);
2377
1/2
✓ Branch 0 taken 157 times.
✗ Branch 1 not taken.
 157 if (!v.pad1d_forward)
2378 ✗ detail::throw_not_implemented("pad1d_forward", X.device);
2379 157 detail::adopt_output(Y, X.device);
2380 157 v.pad1d_forward(X, N, C, L, pad_left, pad_right, mode, Y);
2381 157 }
2382
2383 9 void pad1d_backward(const Tensor& dY, int N, int C, int L,
2384 int pad_left, int pad_right, int mode, Tensor& dX) {
2385 9 const auto& v = detail::dispatch(dY, dX);
2386
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.pad1d_backward)
2387 ✗ detail::throw_not_implemented("pad1d_backward", dY.device);
2388 9 detail::adopt_output(dX, dY.device);
2389 9 v.pad1d_backward(dY, N, C, L, pad_left, pad_right, mode, dX);
2390 9 }
2391
2392 // ─── Vocoder / codec activations (brosoundml) ──────────────────────────────
2393
2394 124 void snake_forward(const Tensor& X, const Tensor& alpha, const Tensor* beta,
2395 int N, int C, int L, Tensor& Y) {
2396 124 const auto& v = detail::dispatch_with_opts(X, alpha, {beta, &Y});
2397
1/2
✓ Branch 0 taken 124 times.
✗ Branch 1 not taken.
 124 if (!v.snake_forward)
2398 ✗ detail::throw_not_implemented("snake_forward", X.device);
2399 124 detail::adopt_output(Y, X.device);
2400 124 v.snake_forward(X, alpha, beta, N, C, L, Y);
2401 124 }
2402
2403 11 void snake_backward(const Tensor& X, const Tensor& alpha, const Tensor* beta,
2404 const Tensor& dY, int N, int C, int L,
2405 Tensor& dX, Tensor& dAlpha, Tensor* dBeta) {
2406 11 const auto& v = detail::dispatch_with_opts(
2407 11 X, alpha, {beta, &dY, &dX, &dAlpha, dBeta});
2408
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.snake_backward)
2409 ✗ detail::throw_not_implemented("snake_backward", X.device);
2410 11 detail::adopt_output(dX, X.device);
2411 11 detail::adopt_output(dAlpha, X.device);
2412
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 8 times.
 11 if (dBeta) detail::adopt_output(*dBeta, X.device);
2413 11 v.snake_backward(X, alpha, beta, dY, N, C, L, dX, dAlpha, dBeta);
2414 11 }
2415
2416 137 void elu_forward(const Tensor& x, float alpha, Tensor& y) {
2417 137 const auto& v = detail::dispatch(x, y);
2418
1/2
✓ Branch 0 taken 137 times.
✗ Branch 1 not taken.
 137 if (!v.elu_forward) detail::throw_not_implemented("elu_forward", x.device);
2419 137 detail::adopt_output(y, x.device);
2420 137 v.elu_forward(x, alpha, y);
2421 137 }
2422
2423 9 void elu_backward(const Tensor& x, const Tensor& dY, float alpha, Tensor& dX) {
2424 9 const auto& v = detail::dispatch(x, dY, dX);
2425
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.elu_backward) detail::throw_not_implemented("elu_backward", x.device);
2426 9 detail::adopt_output(dX, x.device);
2427 9 v.elu_backward(x, dY, alpha, dX);
2428 9 }
2429
2430 94 void leaky_relu_forward(const Tensor& x, float negative_slope, Tensor& y) {
2431 94 const auto& v = detail::dispatch(x, y);
2432
1/2
✓ Branch 0 taken 94 times.
✗ Branch 1 not taken.
 94 if (!v.leaky_relu_forward)
2433 ✗ detail::throw_not_implemented("leaky_relu_forward", x.device);
2434 94 detail::adopt_output(y, x.device);
2435 94 v.leaky_relu_forward(x, negative_slope, y);
2436 94 }
2437
2438 6 void leaky_relu_backward(const Tensor& x, const Tensor& dY,
2439 float negative_slope, Tensor& dX) {
2440 6 const auto& v = detail::dispatch(x, dY, dX);
2441
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.leaky_relu_backward)
2442 ✗ detail::throw_not_implemented("leaky_relu_backward", x.device);
2443 6 detail::adopt_output(dX, x.device);
2444 6 v.leaky_relu_backward(x, dY, negative_slope, dX);
2445 6 }
2446
2447 // ─── Codec quantization (brosoundml) ───────────────────────────────────────
2448
2449 9 void vq_encode_forward(const Tensor& x, const Tensor& codebook,
2450 Tensor& indices, Tensor& quantized) {
2451 9 const auto& v = detail::dispatch(x, codebook, indices, quantized);
2452
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.vq_encode_forward)
2453 ✗ detail::throw_not_implemented("vq_encode_forward", x.device);
2454 9 detail::adopt_output(indices, x.device);
2455 9 detail::adopt_output(quantized, x.device);
2456 9 v.vq_encode_forward(x, codebook, indices, quantized);
2457 9 }
2458
2459 6 void vq_encode_backward(const Tensor& dQuantized, Tensor& dX) {
2460 6 const auto& v = detail::dispatch(dQuantized, dX);
2461
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.vq_encode_backward)
2462 ✗ detail::throw_not_implemented("vq_encode_backward", dQuantized.device);
2463 6 detail::adopt_output(dX, dQuantized.device);
2464 6 v.vq_encode_backward(dQuantized, dX);
2465 6 }
2466
2467 8 void fsq_quantize_forward(const Tensor& x, const Tensor& levels,
2468 Tensor& quantized, Tensor& packed_indices) {
2469 8 const auto& v = detail::dispatch(x, levels, quantized, packed_indices);
2470
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.fsq_quantize_forward)
2471 ✗ detail::throw_not_implemented("fsq_quantize_forward", x.device);
2472 8 detail::adopt_output(quantized, x.device);
2473 8 detail::adopt_output(packed_indices, x.device);
2474 8 v.fsq_quantize_forward(x, levels, quantized, packed_indices);
2475 8 }
2476
2477 6 void fsq_quantize_backward(const Tensor& dQuantized, Tensor& dX) {
2478 6 const auto& v = detail::dispatch(dQuantized, dX);
2479
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.fsq_quantize_backward)
2480 ✗ detail::throw_not_implemented("fsq_quantize_backward", dQuantized.device);
2481 6 detail::adopt_output(dX, dQuantized.device);
2482 6 v.fsq_quantize_backward(dQuantized, dX);
2483 6 }
2484
2485 // ─── 1D resampling (brosoundml) ────────────────────────────────────────────
2486
2487 152 void resample1d_forward(const Tensor& X, int N, int C, int L_in, int L_out,
2488 int mode, Tensor& Y) {
2489 152 const auto& v = detail::dispatch(X, Y);
2490
1/2
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
 152 if (!v.resample1d_forward)
2491 ✗ detail::throw_not_implemented("resample1d_forward", X.device);
2492 152 detail::adopt_output(Y, X.device);
2493 152 v.resample1d_forward(X, N, C, L_in, L_out, mode, Y);
2494 152 }
2495
2496 16 void resample1d_backward(const Tensor& dY, int N, int C, int L_in, int L_out,
2497 int mode, Tensor& dX) {
2498 16 const auto& v = detail::dispatch(dY, dX);
2499
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.resample1d_backward)
2500 ✗ detail::throw_not_implemented("resample1d_backward", dY.device);
2501 16 detail::adopt_output(dX, dY.device);
2502 16 v.resample1d_backward(dY, N, C, L_in, L_out, mode, dX);
2503 16 }
2504
2505 // ─── log / exp / round elementwise (brosoundml) ────────────────────────────
2506
2507 52 void log_forward(const Tensor& x, Tensor& y) {
2508 52 const auto& v = detail::dispatch(x, y);
2509
1/2
✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
 52 if (!v.log_forward) detail::throw_not_implemented("log_forward", x.device);
2510 52 detail::adopt_output(y, x.device);
2511 52 v.log_forward(x, y);
2512 52 }
2513
2514 5 void log_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2515 5 const auto& v = detail::dispatch(x, dY, dX);
2516
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.log_backward) detail::throw_not_implemented("log_backward", x.device);
2517 5 detail::adopt_output(dX, x.device);
2518 5 v.log_backward(x, dY, dX);
2519 5 }
2520
2521 48 void exp_forward(const Tensor& x, Tensor& y) {
2522 48 const auto& v = detail::dispatch(x, y);
2523
1/2
✓ Branch 0 taken 48 times.
✗ Branch 1 not taken.
 48 if (!v.exp_forward) detail::throw_not_implemented("exp_forward", x.device);
2524 48 detail::adopt_output(y, x.device);
2525 48 v.exp_forward(x, y);
2526 48 }
2527
2528 7 void exp_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2529 7 const auto& v = detail::dispatch(x, dY, dX);
2530
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.exp_backward) detail::throw_not_implemented("exp_backward", x.device);
2531 7 detail::adopt_output(dX, x.device);
2532 7 v.exp_backward(x, dY, dX);
2533 7 }
2534
2535 7 void round_forward(const Tensor& x, Tensor& y) {
2536 7 const auto& v = detail::dispatch(x, y);
2537
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.round_forward) detail::throw_not_implemented("round_forward", x.device);
2538 7 detail::adopt_output(y, x.device);
2539 7 v.round_forward(x, y);
2540 7 }
2541
2542 3 void round_backward(const Tensor& dY, Tensor& dX) {
2543 3 const auto& v = detail::dispatch(dY, dX);
2544
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.round_backward) detail::throw_not_implemented("round_backward", dY.device);
2545 3 detail::adopt_output(dX, dY.device);
2546 3 v.round_backward(dY, dX);
2547 3 }
2548
2549 // ─── Autoregressive logit sampling (brosoundml CHUNK 7, family F) ───────────
2550
2551 50266 void sample_logits(const Tensor& logits, float temperature, int top_k,
2552 float top_p, uint64_t key, uint64_t counter,
2553 Tensor& indices) {
2554 50266 const auto& v = detail::dispatch(logits, indices);
2555
1/2
✓ Branch 0 taken 50266 times.
✗ Branch 1 not taken.
 50266 if (!v.sample_logits) detail::throw_not_implemented("sample_logits", logits.device);
2556 50266 detail::adopt_output(indices, logits.device);
2557 50266 v.sample_logits(logits, temperature, top_k, top_p, key, counter, indices);
2558 50266 }
2559
2560 13 void sample_logits_into(const Tensor& logits, float temperature, int top_k,
2561 float top_p, uint64_t key, Tensor& counter,
2562 Tensor& scratch, Tensor& indices) {
2563 13 const auto& v = detail::dispatch(logits, indices);
2564
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (!v.sample_logits_into)
2565 ✗ detail::throw_not_implemented("sample_logits_into", logits.device);
2566 13 detail::adopt_output(indices, logits.device);
2567 26 v.sample_logits_into(logits, temperature, top_k, top_p, key, counter,
2568 13 scratch, indices);
2569 13 }
2570
2571 // ─── L2 norm + Gated Delta Rule (linear-attention text path) ───────────────
2572
2573 73 void l2_norm_forward(const Tensor& X, int head_dim, int num_heads,
2574 float eps, Tensor& Y) {
2575 73 const auto& v = detail::dispatch(X, Y);
2576
1/2
✓ Branch 0 taken 73 times.
✗ Branch 1 not taken.
 73 if (!v.l2_norm_forward) detail::throw_not_implemented("l2_norm_forward", X.device);
2577 73 detail::adopt_output(Y, X.device);
2578 73 v.l2_norm_forward(X, head_dim, num_heads, eps, Y);
2579 73 }
2580
2581 11 void l2_norm_backward(const Tensor& X, int head_dim, int num_heads,
2582 float eps, const Tensor& dY, Tensor& dX) {
2583 11 const auto& v = detail::dispatch(X, dY, dX);
2584
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.l2_norm_backward) detail::throw_not_implemented("l2_norm_backward", X.device);
2585 11 detail::adopt_output(dX, X.device);
2586 11 v.l2_norm_backward(X, head_dim, num_heads, eps, dY, dX);
2587 11 }
2588
2589 11 void gated_delta_rule_chunked(const Tensor& Q, const Tensor& K, const Tensor& V,
2590 const Tensor& a_raw, const Tensor& beta,
2591 const Tensor& log_A,
2592 int num_heads, int d_k, int d_v,
2593 Tensor& state, Tensor& O) {
2594 11 const auto& v = detail::dispatch(Q, K, V, a_raw, beta, log_A, state, O);
2595
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.gated_delta_rule_chunked)
2596 ✗ detail::throw_not_implemented("gated_delta_rule_chunked", Q.device);
2597 11 detail::adopt_output(state, Q.device);
2598 11 detail::adopt_output(O, Q.device);
2599 22 v.gated_delta_rule_chunked(Q, K, V, a_raw, beta, log_A,
2600 11 num_heads, d_k, d_v, state, O);
2601 11 }
2602
2603 12 void gated_delta_rule_step(const Tensor& Q, const Tensor& K, const Tensor& V,
2604 const Tensor& a_raw, const Tensor& beta,
2605 const Tensor& log_A,
2606 int num_heads, int d_k, int d_v,
2607 Tensor& state, Tensor& O) {
2608 12 const auto& v = detail::dispatch(Q, K, V, a_raw, beta, log_A, state, O);
2609
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.gated_delta_rule_step)
2610 ✗ detail::throw_not_implemented("gated_delta_rule_step", Q.device);
2611 12 detail::adopt_output(state, Q.device);
2612 12 detail::adopt_output(O, Q.device);
2613 24 v.gated_delta_rule_step(Q, K, V, a_raw, beta, log_A,
2614 12 num_heads, d_k, d_v, state, O);
2615 12 }
2616
2617 // ─── BatchNorm ─────────────────────────────────────────────────────────────
2618
2619 9 void batch_norm_forward(const Tensor& X,
2620 const Tensor& gamma, const Tensor& beta,
2621 Tensor& running_mean, Tensor& running_var,
2622 int N, int C, int H, int W,
2623 float eps, float momentum,
2624 Tensor& Y,
2625 Tensor& saved_mean, Tensor& saved_rstd) {
2626 18 const auto& v = detail::dispatch(X, gamma, beta, running_mean, running_var,
2627 9 Y, saved_mean, saved_rstd);
2628
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.batch_norm_forward)
2629 ✗ detail::throw_not_implemented("batch_norm_forward", X.device);
2630 9 detail::adopt_output(running_mean, X.device);
2631 9 detail::adopt_output(running_var, X.device);
2632 9 detail::adopt_output(Y, X.device);
2633 9 detail::adopt_output(saved_mean, X.device);
2634 9 detail::adopt_output(saved_rstd, X.device);
2635 18 v.batch_norm_forward(X, gamma, beta, running_mean, running_var,
2636 9 N, C, H, W, eps, momentum,
2637 9 Y, saved_mean, saved_rstd);
2638 9 }
2639
2640 11 void batch_norm_inference(const Tensor& X,
2641 const Tensor& gamma, const Tensor& beta,
2642 const Tensor& running_mean,
2643 const Tensor& running_var,
2644 int N, int C, int H, int W,
2645 float eps,
2646 Tensor& Y) {
2647 11 const auto& v = detail::dispatch(X, gamma, beta, running_mean, running_var, Y);
2648
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.batch_norm_inference)
2649 ✗ detail::throw_not_implemented("batch_norm_inference", X.device);
2650 11 detail::adopt_output(Y, X.device);
2651 22 v.batch_norm_inference(X, gamma, beta, running_mean, running_var,
2652 11 N, C, H, W, eps, Y);
2653 11 }
2654
2655 5 void batch_norm_backward(const Tensor& X,
2656 const Tensor& gamma,
2657 const Tensor& saved_mean,
2658 const Tensor& saved_rstd,
2659 const Tensor& dY,
2660 int N, int C, int H, int W,
2661 Tensor& dX,
2662 Tensor& dGamma, Tensor& dBeta) {
2663 10 const auto& v = detail::dispatch(X, gamma, saved_mean, saved_rstd, dY,
2664 5 dX, dGamma, dBeta);
2665
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.batch_norm_backward)
2666 ✗ detail::throw_not_implemented("batch_norm_backward", X.device);
2667 5 detail::adopt_output(dX, X.device);
2668 5 detail::adopt_output(dGamma, X.device);
2669 5 detail::adopt_output(dBeta, X.device);
2670 10 v.batch_norm_backward(X, gamma, saved_mean, saved_rstd, dY,
2671 5 N, C, H, W, dX, dGamma, dBeta);
2672 5 }
2673
2674 // ─── Image preprocessing helpers ───────────────────────────────────────────
2675
2676 6 void image_normalize(const Tensor& X,
2677 const Tensor& mean, const Tensor& std_,
2678 int N, int C, int H, int W,
2679 Tensor& Y) {
2680 6 const auto& v = detail::dispatch(X, mean, std_, Y);
2681
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.image_normalize)
2682 ✗ detail::throw_not_implemented("image_normalize", X.device);
2683 6 detail::adopt_output(Y, X.device);
2684 6 v.image_normalize(X, mean, std_, N, C, H, W, Y);
2685 6 }
2686
2687 9 void image_u8_to_f32_nhwc_to_nchw(const uint8_t* src,
2688 int N, int H, int W, int C,
2689 float scale, float bias,
2690 Tensor& Y) {
2691 // No tensor inputs to dispatch on — Y is the only one. adopt_output
2692 // pins an uncommitted Y to the default device before lookup.
2693 9 const auto& v = detail::dispatch(Y);
2694
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.image_u8_to_f32_nhwc_to_nchw)
2695 ✗ detail::throw_not_implemented("image_u8_to_f32_nhwc_to_nchw", Y.device);
2696 9 detail::adopt_output(Y, Y.device);
2697 9 v.image_u8_to_f32_nhwc_to_nchw(src, N, H, W, C, scale, bias, Y);
2698 9 }
2699
2700 // ─── Counter-based noise generation (Philox 4x32-10) ───────────────────────
2701
2702 6 void randn(uint64_t key, uint64_t counter, Tensor& Y) {
2703 6 const auto& v = detail::dispatch(Y);
2704
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.randn) detail::throw_not_implemented("randn", Y.device);
2705 6 detail::adopt_output(Y, Y.device);
2706 6 v.randn(key, counter, Y);
2707 6 }
2708
2709 4 void rand_uniform(uint64_t key, uint64_t counter, Tensor& Y) {
2710 4 const auto& v = detail::dispatch(Y);
2711
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rand_uniform) detail::throw_not_implemented("rand_uniform", Y.device);
2712 4 detail::adopt_output(Y, Y.device);
2713 4 v.rand_uniform(key, counter, Y);
2714 4 }
2715
2716 4 void rand_bernoulli(float p, uint64_t key, uint64_t counter, Tensor& Y) {
2717 4 const auto& v = detail::dispatch(Y);
2718
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rand_bernoulli) detail::throw_not_implemented("rand_bernoulli", Y.device);
2719 4 detail::adopt_output(Y, Y.device);
2720 4 v.rand_bernoulli(p, key, counter, Y);
2721 4 }
2722
2723 2 void randn_truncated(float lo, float hi, uint64_t key, uint64_t counter,
2724 Tensor& Y) {
2725 2 const auto& v = detail::dispatch(Y);
2726
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.randn_truncated) detail::throw_not_implemented("randn_truncated", Y.device);
2727 2 detail::adopt_output(Y, Y.device);
2728 2 v.randn_truncated(lo, hi, key, counter, Y);
2729 2 }
2730
2731 // ─── StyleGAN3 synthesis-input primitives — sin/cos/rsqrt + pixel_norm ──────
2732
2733 32 void sin_forward(const Tensor& x, Tensor& y) {
2734 32 const auto& v = detail::dispatch(x, y);
2735
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.sin_forward) detail::throw_not_implemented("sin_forward", x.device);
2736 32 detail::adopt_output(y, x.device);
2737 32 v.sin_forward(x, y);
2738 32 }
2739
2740 3 void sin_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2741 3 const auto& v = detail::dispatch(x, dY, dX);
2742
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.sin_backward) detail::throw_not_implemented("sin_backward", x.device);
2743 3 detail::adopt_output(dX, x.device);
2744 3 v.sin_backward(x, dY, dX);
2745 3 }
2746
2747 32 void cos_forward(const Tensor& x, Tensor& y) {
2748 32 const auto& v = detail::dispatch(x, y);
2749
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.cos_forward) detail::throw_not_implemented("cos_forward", x.device);
2750 32 detail::adopt_output(y, x.device);
2751 32 v.cos_forward(x, y);
2752 32 }
2753
2754 3 void cos_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2755 3 const auto& v = detail::dispatch(x, dY, dX);
2756
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.cos_backward) detail::throw_not_implemented("cos_backward", x.device);
2757 3 detail::adopt_output(dX, x.device);
2758 3 v.cos_backward(x, dY, dX);
2759 3 }
2760
2761 32 void rsqrt_forward(const Tensor& x, Tensor& y) {
2762 32 const auto& v = detail::dispatch(x, y);
2763
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.rsqrt_forward) detail::throw_not_implemented("rsqrt_forward", x.device);
2764 32 detail::adopt_output(y, x.device);
2765 32 v.rsqrt_forward(x, y);
2766 32 }
2767
2768 7 void rsqrt_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
2769 7 const auto& v = detail::dispatch(y, dY, dX);
2770
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.rsqrt_backward) detail::throw_not_implemented("rsqrt_backward", y.device);
2771 7 detail::adopt_output(dX, y.device);
2772 7 v.rsqrt_backward(y, dY, dX);
2773 7 }
2774
2775 55 void pixel_norm_forward(const Tensor& X, float eps, Tensor& Y) {
2776 55 const auto& v = detail::dispatch(X, Y);
2777
1/2
✓ Branch 0 taken 55 times.
✗ Branch 1 not taken.
 55 if (!v.pixel_norm_forward) detail::throw_not_implemented("pixel_norm_forward", X.device);
2778 55 detail::adopt_output(Y, X.device);
2779 55 v.pixel_norm_forward(X, eps, Y);
2780 55 }
2781
2782 7 void pixel_norm_backward(const Tensor& X, const Tensor& dY, float eps, Tensor& dX) {
2783 7 const auto& v = detail::dispatch(X, dY, dX);
2784
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.pixel_norm_backward) detail::throw_not_implemented("pixel_norm_backward", X.device);
2785 7 detail::adopt_output(dX, X.device);
2786 7 v.pixel_norm_backward(X, dY, eps, dX);
2787 7 }
2788
2789 // ─── StyleGAN3 bias_act ─────────────────────────────────────────────────────
2790
2791 42 void bias_act_forward(const Tensor& X, const Tensor* b,
2792 int N, int C, int HW, int act, float alpha,
2793 float gain, float clamp, Tensor& Y) {
2794 42 const auto& v = detail::dispatch_with_opts(X, {b, &Y});
2795
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 if (!v.bias_act_forward) detail::throw_not_implemented("bias_act_forward", X.device);
2796 42 detail::adopt_output(Y, X.device);
2797 42 v.bias_act_forward(X, b, N, C, HW, act, alpha, gain, clamp, Y);
2798 42 }
2799
2800 42 void bias_act_backward(const Tensor& dY, const Tensor& X, const Tensor* b,
2801 int N, int C, int HW, int act, float alpha,
2802 float gain, float clamp, Tensor& dX, Tensor* dB) {
2803 42 const auto& v = detail::dispatch_with_opts(dY, X, {b, dB, &dX});
2804
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 if (!v.bias_act_backward) detail::throw_not_implemented("bias_act_backward", dY.device);
2805 42 detail::adopt_output(dX, dY.device);
2806
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 22 times.
 42 if (dB) detail::adopt_output(*dB, dY.device);
2807 42 v.bias_act_backward(dY, X, b, N, C, HW, act, alpha, gain, clamp, dX, dB);
2808 42 }
2809
2810 // ─── StyleGAN3 upfirdn2d ────────────────────────────────────────────────────
2811
2812 25 void upfirdn2d_forward(const Tensor& X, const Tensor& f,
2813 int N, int C, int H, int Wd, int fH, int fW,
2814 int up_x, int up_y, int down_x, int down_y,
2815 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
2816 bool flip_filter, float gain, Tensor& Y) {
2817 25 const auto& v = detail::dispatch(X, f, Y);
2818
1/2
✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
 25 if (!v.upfirdn2d_forward) detail::throw_not_implemented("upfirdn2d_forward", X.device);
2819 25 detail::adopt_output(Y, X.device);
2820 50 v.upfirdn2d_forward(X, f, N, C, H, Wd, fH, fW, up_x, up_y, down_x, down_y,
2821 25 pad_x0, pad_x1, pad_y0, pad_y1, flip_filter, gain, Y);
2822 25 }
2823
2824 21 void upfirdn2d_backward(const Tensor& dY, const Tensor& f,
2825 int N, int C, int H, int Wd, int fH, int fW,
2826 int up_x, int up_y, int down_x, int down_y,
2827 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
2828 bool flip_filter, float gain, Tensor& dX) {
2829 21 const auto& v = detail::dispatch(dY, f, dX);
2830
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (!v.upfirdn2d_backward) detail::throw_not_implemented("upfirdn2d_backward", dY.device);
2831 21 detail::adopt_output(dX, dY.device);
2832 42 v.upfirdn2d_backward(dY, f, N, C, H, Wd, fH, fW, up_x, up_y, down_x, down_y,
2833 21 pad_x0, pad_x1, pad_y0, pad_y1, flip_filter, gain, dX);
2834 21 }
2835
2836 // ─── StyleGAN3 modulated_conv2d ─────────────────────────────────────────────
2837
2838 12 void modulated_conv2d_forward(const Tensor& X, const Tensor& W, const Tensor& s,
2839 int N, int C_in, int H, int Wd,
2840 int C_out, int kH, int kW,
2841 int pad_h, int pad_w,
2842 bool demodulate, float eps,
2843 Tensor& dcoef, Tensor& Y) {
2844 12 const auto& v = detail::dispatch(X, W, s, dcoef, Y);
2845
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.modulated_conv2d_forward)
2846 ✗ detail::throw_not_implemented("modulated_conv2d_forward", X.device);
2847 12 detail::adopt_output(dcoef, X.device);
2848 12 detail::adopt_output(Y, X.device);
2849 24 v.modulated_conv2d_forward(X, W, s, N, C_in, H, Wd, C_out, kH, kW,
2850 12 pad_h, pad_w, demodulate, eps, dcoef, Y);
2851 12 }
2852
2853 16 void modulated_conv2d_backward(const Tensor& X, const Tensor& W, const Tensor& s,
2854 const Tensor& dcoef, const Tensor& dY,
2855 int N, int C_in, int H, int Wd,
2856 int C_out, int kH, int kW,
2857 int pad_h, int pad_w, bool demodulate, float eps,
2858 Tensor& dX, Tensor& dW, Tensor& ds) {
2859 16 const auto& v = detail::dispatch(X, W, s, dcoef, dY, dX, dW, ds);
2860
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.modulated_conv2d_backward)
2861 ✗ detail::throw_not_implemented("modulated_conv2d_backward", X.device);
2862 16 detail::adopt_output(dX, X.device);
2863 16 detail::adopt_output(dW, X.device);
2864 16 detail::adopt_output(ds, X.device);
2865 32 v.modulated_conv2d_backward(X, W, s, dcoef, dY, N, C_in, H, Wd, C_out, kH, kW,
2866 16 pad_h, pad_w, demodulate, eps, dX, dW, ds);
2867 16 }
2868
2869 } // namespace brotensor
2870