GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 91.8% 457 / 0 / 498
Functions: 100.0% 61 / 0 / 61
Branches: 57.2% 87 / 0 / 152
src/ops_attention.cpp
Line Branch Exec Source
1 // brotensor public op wrappers — attention, matmul, LSTM, RoPE, RMSNorm, SwiGLU, reductions.
2
3 #include <brotensor/ops.h>
4 #include <brotensor/tensor.h>
5 #include <brotensor/detail/dispatch.h>
6
7 #include <cstdint>
8 #include <stdexcept>
9
10 namespace brotensor {
11
12 // ─── FP16 linear + GEGLU ───────────────────────────────────────────────────
13
14 49 void linear_forward_batched_fp16(const Tensor& W, const Tensor* bias,
15 const Tensor& X_BD, Tensor& Y_BD) {
16 49 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD});
17
1/2
✓ Branch 0 taken 49 times.
✗ Branch 1 not taken.
 49 if (!v.linear_forward_batched_fp16) detail::throw_not_implemented("linear_forward_batched_fp16", W.device);
18 49 detail::adopt_output(Y_BD, W.device);
19 49 v.linear_forward_batched_fp16(W, bias, X_BD, Y_BD);
20 49 }
21
22 29 void linear_forward_batched_fp16_act(const Tensor& W, const Tensor* bias,
23 const Tensor& X_BD, int act, Tensor& Y_BD) {
24 29 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD});
25
1/2
✓ Branch 0 taken 29 times.
✗ Branch 1 not taken.
 29 if (!v.linear_forward_batched_fp16_act)
26 ✗ detail::throw_not_implemented("linear_forward_batched_fp16_act", W.device);
27 29 detail::adopt_output(Y_BD, W.device);
28 29 v.linear_forward_batched_fp16_act(W, bias, X_BD, act, Y_BD);
29 29 }
30
31 439 void linear_forward_batched_ex(const Tensor& W, const Tensor* bias, const Tensor& X_BD,
32 int act, int epilogue, Tensor* workspace, Tensor& Y_BD) {
33 439 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD, workspace});
34
1/2
✓ Branch 0 taken 439 times.
✗ Branch 1 not taken.
 439 if (!v.linear_forward_batched_ex)
35 ✗ detail::throw_not_implemented("linear_forward_batched_ex", W.device);
36 439 detail::adopt_output(Y_BD, W.device);
37
2/2
✓ Branch 0 taken 285 times.
✓ Branch 1 taken 154 times.
 439 if (workspace) detail::adopt_output(*workspace, W.device);
38 439 v.linear_forward_batched_ex(W, bias, X_BD, act, epilogue, workspace, Y_BD);
39 439 }
40
41 12 void geglu_forward(const Tensor& X, Tensor& Y) {
42 12 const auto& v = detail::dispatch(X, Y);
43
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.geglu_forward) detail::throw_not_implemented("geglu_forward", X.device);
44 12 detail::adopt_output(Y, X.device);
45 12 v.geglu_forward(X, Y);
46 12 }
47 8 void geglu_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
48 8 const auto& v = detail::dispatch(X, dY, dX);
49
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.geglu_backward) detail::throw_not_implemented("geglu_backward", X.device);
50 8 detail::adopt_output(dX, X.device);
51 8 v.geglu_backward(X, dY, dX);
52 8 }
53 11 void geglu_exact_forward(const Tensor& X, Tensor& Y) {
54 11 const auto& v = detail::dispatch(X, Y);
55
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.geglu_exact_forward) detail::throw_not_implemented("geglu_exact_forward", X.device);
56 11 detail::adopt_output(Y, X.device);
57 11 v.geglu_exact_forward(X, Y);
58 11 }
59 9 void geglu_exact_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
60 9 const auto& v = detail::dispatch(X, dY, dX);
61
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.geglu_exact_backward) detail::throw_not_implemented("geglu_exact_backward", X.device);
62 9 detail::adopt_output(dX, X.device);
63 9 v.geglu_exact_backward(X, dY, dX);
64 9 }
65
66 // ─── Causal mask helper ────────────────────────────────────────────────────
67
68 23 void build_causal_mask_row(int L, int q, Tensor& mask) {
69 23 const auto& v = detail::dispatch(mask);
70
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.build_causal_mask_row) detail::throw_not_implemented("build_causal_mask_row", mask.device);
71 23 detail::adopt_output(mask, mask.device);
72 23 v.build_causal_mask_row(L, q, mask);
73 23 }
74
75 // ─── Cross-attention family ────────────────────────────────────────────────
76
77 27 void cross_attention_forward(const Tensor& X, const Tensor& Ctx,
78 const Tensor& Wq, const Tensor& Wk,
79 const Tensor& Wv, const Tensor& Wo,
80 const float* d_mask, int num_heads, Tensor& O) {
81 27 const auto& v = detail::dispatch(X, Ctx, Wq, Wk, Wv, Wo, O);
82
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.cross_attention_forward) detail::throw_not_implemented("cross_attention_forward", X.device);
83 27 detail::adopt_output(O, X.device);
84 27 v.cross_attention_forward(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads, O);
85 27 }
86
87 17 void cross_attention_forward_with_attn(const Tensor& X, const Tensor& Ctx,
88 const Tensor& Wq, const Tensor& Wk,
89 const Tensor& Wv, const Tensor& Wo,
90 const float* d_mask,
91 const Tensor* attn_logit_bias,
92 int num_heads, Tensor& O, Tensor& AttnAvg) {
93 17 const auto& v = detail::dispatch_with_opts(X, Ctx, {&Wq, &Wk, &Wv, &Wo, attn_logit_bias, &O, &AttnAvg});
94
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.cross_attention_forward_with_attn)
95 ✗ detail::throw_not_implemented("cross_attention_forward_with_attn", X.device);
96 17 detail::adopt_output(O, X.device);
97 17 detail::adopt_output(AttnAvg, X.device);
98 34 v.cross_attention_forward_with_attn(X, Ctx, Wq, Wk, Wv, Wo, d_mask,
99 17 attn_logit_bias, num_heads, O, AttnAvg);
100 17 }
101
102 21 void self_attention_forward_train(const Tensor& X,
103 const Tensor& Wq, const Tensor& Wk,
104 const Tensor& Wv, const Tensor& Wo,
105 const float* d_mask, int num_heads,
106 Tensor& Qh, Tensor& Kh, Tensor& Vh,
107 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
108 21 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, Qh, Kh, Vh);
109
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (!v.self_attention_forward_train)
110 ✗ detail::throw_not_implemented("self_attention_forward_train", X.device);
111 21 detail::adopt_output(Qh, X.device);
112 21 detail::adopt_output(Kh, X.device);
113 21 detail::adopt_output(Vh, X.device);
114 21 detail::adopt_output(Attnh, X.device);
115 21 detail::adopt_output(Yconcat, X.device);
116 21 detail::adopt_output(O, X.device);
117 42 v.self_attention_forward_train(X, Wq, Wk, Wv, Wo, d_mask, num_heads,
118 21 Qh, Kh, Vh, Attnh, Yconcat, O);
119 21 }
120
121 19 void self_attention_backward(const Tensor& dO, const Tensor& X,
122 const Tensor& Qh, const Tensor& Kh,
123 const Tensor& Vh, const Tensor& Attnh,
124 const Tensor& Yconcat,
125 const Tensor& Wq, const Tensor& Wk,
126 const Tensor& Wv, const Tensor& Wo,
127 const float* d_mask, int num_heads,
128 Tensor& dX,
129 Tensor& dWq, Tensor& dWk,
130 Tensor& dWv, Tensor& dWo) {
131 19 const auto& v = detail::dispatch(dO, X, Qh, Kh, Vh, Attnh, Yconcat);
132
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.self_attention_backward)
133 ✗ detail::throw_not_implemented("self_attention_backward", dO.device);
134 19 detail::adopt_output(dX, dO.device);
135 19 detail::adopt_output(dWq, dO.device);
136 19 detail::adopt_output(dWk, dO.device);
137 19 detail::adopt_output(dWv, dO.device);
138 19 detail::adopt_output(dWo, dO.device);
139 38 v.self_attention_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat,
140 19 Wq, Wk, Wv, Wo, d_mask, num_heads,
141 19 dX, dWq, dWk, dWv, dWo);
142 19 }
143
144 16 void attention_token_moments(const Tensor& Attn, int h_lat, int w_lat,
145 Tensor& mass, Tensor& centroid) {
146
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (h_lat <= 0 || w_lat <= 0 || Attn.rows != h_lat * w_lat) {
147 ✗ detail::throw_spatial_2d_error("attention_token_moments", "Attn",
148 ✗ h_lat * w_lat, Attn.cols, h_lat, w_lat,
149 ✗ Attn.rows, Attn.cols);
150 }
151 16 const auto& v = detail::dispatch(Attn, mass, centroid);
152
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.attention_token_moments) detail::throw_not_implemented("attention_token_moments", Attn.device);
153 16 detail::adopt_output(mass, Attn.device);
154 16 detail::adopt_output(centroid, Attn.device);
155 16 v.attention_token_moments(Attn, h_lat, w_lat, mass, centroid);
156 16 }
157
158 19 void cross_attention_forward_train(const Tensor& X, const Tensor& Ctx,
159 const Tensor& Wq, const Tensor& Wk,
160 const Tensor& Wv, const Tensor& Wo,
161 const float* d_mask, int num_heads,
162 Tensor& Qh, Tensor& Kh, Tensor& Vh,
163 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
164 19 const auto& v = detail::dispatch(X, Ctx, Wq, Wk, Wv, Wo, Qh, Kh);
165
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.cross_attention_forward_train)
166 ✗ detail::throw_not_implemented("cross_attention_forward_train", X.device);
167 19 detail::adopt_output(Qh, X.device);
168 19 detail::adopt_output(Kh, X.device);
169 19 detail::adopt_output(Vh, X.device);
170 19 detail::adopt_output(Attnh, X.device);
171 19 detail::adopt_output(Yconcat, X.device);
172 19 detail::adopt_output(O, X.device);
173 38 v.cross_attention_forward_train(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads,
174 19 Qh, Kh, Vh, Attnh, Yconcat, O);
175 19 }
176
177 19 void cross_attention_backward(const Tensor& dO, const Tensor& X, const Tensor& Ctx,
178 const Tensor& Qh, const Tensor& Kh,
179 const Tensor& Vh, const Tensor& Attnh,
180 const Tensor& Yconcat,
181 const Tensor& Wq, const Tensor& Wk,
182 const Tensor& Wv, const Tensor& Wo,
183 const float* d_mask, int num_heads,
184 Tensor& dX, Tensor& dCtx,
185 Tensor& dWq, Tensor& dWk,
186 Tensor& dWv, Tensor& dWo) {
187 19 const auto& v = detail::dispatch(dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat);
188
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.cross_attention_backward)
189 ✗ detail::throw_not_implemented("cross_attention_backward", dO.device);
190 19 detail::adopt_output(dX, dO.device);
191 19 detail::adopt_output(dCtx, dO.device);
192 19 detail::adopt_output(dWq, dO.device);
193 19 detail::adopt_output(dWk, dO.device);
194 19 detail::adopt_output(dWv, dO.device);
195 19 detail::adopt_output(dWo, dO.device);
196 38 v.cross_attention_backward(dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat,
197 19 Wq, Wk, Wv, Wo, d_mask, num_heads,
198 19 dX, dCtx, dWq, dWk, dWv, dWo);
199 19 }
200
201 // ─── FP16 LN inference + FP16 self-attention ───────────────────────────────
202
203 3 void layernorm_forward_inference_batched_fp16(const Tensor& X_RD,
204 const Tensor& gamma,
205 const Tensor& beta,
206 Tensor& Y_RD, float eps) {
207 3 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD);
208
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.layernorm_forward_inference_batched_fp16)
209 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched_fp16", X_RD.device);
210 3 detail::adopt_output(Y_RD, X_RD.device);
211 3 v.layernorm_forward_inference_batched_fp16(X_RD, gamma, &beta, Y_RD, eps);
212 3 }
213
214 2 void layernorm_forward_inference_batched_fp16(const Tensor& X_RD,
215 const Tensor& gamma,
216 Tensor& Y_RD, float eps) {
217 2 const auto& v = detail::dispatch(X_RD, gamma, Y_RD);
218
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.layernorm_forward_inference_batched_fp16)
219 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched_fp16", X_RD.device);
220 2 detail::adopt_output(Y_RD, X_RD.device);
221 2 v.layernorm_forward_inference_batched_fp16(X_RD, gamma, nullptr, Y_RD, eps);
222 2 }
223
224 14 void self_attention_forward(const Tensor& X,
225 const Tensor& Wq, const Tensor& Wk,
226 const Tensor& Wv, const Tensor& Wo,
227 const float* d_mask, int num_heads, Tensor& O) {
228 14 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, O);
229
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.self_attention_forward) detail::throw_not_implemented("self_attention_forward", X.device);
230 14 detail::adopt_output(O, X.device);
231 14 v.self_attention_forward(X, Wq, Wk, Wv, Wo, d_mask, num_heads, O);
232 14 }
233
234 41 void self_attention_bias_forward(const Tensor& X,
235 const Tensor& Wq, const Tensor& Wk,
236 const Tensor& Wv, const Tensor& Wo,
237 const Tensor* bq, const Tensor* bk,
238 const Tensor* bv, const Tensor* bo,
239 const float* d_mask,
240 const Tensor* attn_bias,
241 int num_heads, float scale, Tensor& O) {
242 41 const auto& v = detail::dispatch_with_opts(X, Wq, {&Wk, &Wv, &Wo, bq, bk, bv, bo, attn_bias, &O});
243
1/2
✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
 41 if (!v.self_attention_bias_forward)
244 ✗ detail::throw_not_implemented("self_attention_bias_forward", X.device);
245 41 detail::adopt_output(O, X.device);
246 82 v.self_attention_bias_forward(X, Wq, Wk, Wv, Wo, bq, bk, bv, bo, d_mask, attn_bias,
247 41 num_heads, scale, O);
248 41 }
249
250 8 void rel_pos_bias_xl_forward(const Tensor& Qv, const Tensor& Pk,
251 int num_heads, int head_dim, Tensor& Bias) {
252 8 const auto& v = detail::dispatch_with_opts(Qv, Pk, {&Bias});
253
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.rel_pos_bias_xl_forward)
254 ✗ detail::throw_not_implemented("rel_pos_bias_xl_forward", Qv.device);
255 8 detail::adopt_output(Bias, Qv.device);
256 8 v.rel_pos_bias_xl_forward(Qv, Pk, num_heads, head_dim, Bias);
257 8 }
258
259 27 void self_attention_decomposed_rel_pos_forward(
260 const Tensor& X,
261 const Tensor& Wq, const Tensor* bq,
262 const Tensor& Wk, const Tensor* bk,
263 const Tensor& Wv, const Tensor* bv,
264 const Tensor& Wo, const Tensor* bo,
265 const Tensor& rel_pos_h, const Tensor& rel_pos_w,
266 int num_heads, int grid_h, int grid_w, float scale, Tensor& O) {
267
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (grid_h <= 0 || grid_w <= 0 || X.rows != grid_h * grid_w) {
268 ✗ detail::throw_spatial_2d_error("self_attention_decomposed_rel_pos_forward", "X",
269 ✗ grid_h * grid_w, X.cols, grid_h, grid_w,
270 ✗ X.rows, X.cols);
271 }
272 27 const auto& v = detail::dispatch_with_opts(
273 27 X, Wq, {bq, &Wk, bk, &Wv, bv, &Wo, bo, &rel_pos_h, &rel_pos_w, &O});
274
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.self_attention_decomposed_rel_pos_forward)
275 ✗ detail::throw_not_implemented("self_attention_decomposed_rel_pos_forward",
276 ✗ X.device);
277 27 detail::adopt_output(O, X.device);
278 54 v.self_attention_decomposed_rel_pos_forward(
279 27 X, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
280 27 num_heads, grid_h, grid_w, scale, O);
281 27 }
282
283 19 void self_attention_decomposed_rel_pos_windowed_forward(
284 const Tensor& X,
285 const Tensor& Wq, const Tensor* bq,
286 const Tensor& Wk, const Tensor* bk,
287 const Tensor& Wv, const Tensor* bv,
288 const Tensor& Wo, const Tensor* bo,
289 const Tensor& rel_pos_h, const Tensor& rel_pos_w,
290 int num_heads, int grid_h, int grid_w, int window, float scale,
291 Tensor& O) {
292
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (grid_h <= 0 || grid_w <= 0 || X.rows != grid_h * grid_w) {
293 ✗ detail::throw_spatial_2d_error(
294 "self_attention_decomposed_rel_pos_windowed_forward", "X",
295 ✗ grid_h * grid_w, X.cols, grid_h, grid_w, X.rows, X.cols);
296 }
297 19 const auto& v = detail::dispatch_with_opts(
298 19 X, Wq, {bq, &Wk, bk, &Wv, bv, &Wo, bo, &rel_pos_h, &rel_pos_w, &O});
299
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.self_attention_decomposed_rel_pos_windowed_forward)
300 ✗ detail::throw_not_implemented(
301 ✗ "self_attention_decomposed_rel_pos_windowed_forward", X.device);
302 19 detail::adopt_output(O, X.device);
303 38 v.self_attention_decomposed_rel_pos_windowed_forward(
304 19 X, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
305 19 num_heads, grid_h, grid_w, window, scale, O);
306 19 }
307
308 9 void self_attention_bias_int8w_fp16(const Tensor& X,
309 const Tensor& Wq_int8, const Tensor& sq,
310 const Tensor& Wk_int8, const Tensor& sk,
311 const Tensor& Wv_int8, const Tensor& sv,
312 const Tensor& Wo_int8, const Tensor& so,
313 const float* d_mask,
314 const Tensor* attn_bias,
315 int num_heads, float scale, Tensor& O) {
316 9 const auto& v = detail::dispatch_with_opts(
317 27 X, Wq_int8, {&sq, &Wk_int8, &sk, &Wv_int8, &sv, &Wo_int8, &so,
318 18 attn_bias, &O});
319
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.self_attention_bias_int8w_fp16)
320 ✗ detail::throw_not_implemented("self_attention_bias_int8w_fp16", X.device);
321 9 detail::adopt_output(O, X.device);
322 18 v.self_attention_bias_int8w_fp16(X, Wq_int8, sq, Wk_int8, sk,
323 9 Wv_int8, sv, Wo_int8, so,
324 9 d_mask, attn_bias, num_heads, scale, O);
325 9 }
326
327 // ─── Flash attention family ────────────────────────────────────────────────
328
329 165 void flash_attention_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
330 const float* d_mask, int num_heads, bool causal,
331 Tensor& O) {
332 165 const auto& v = detail::dispatch(Q, K, V, O);
333
1/2
✓ Branch 0 taken 165 times.
✗ Branch 1 not taken.
 165 if (!v.flash_attention_forward) detail::throw_not_implemented("flash_attention_forward", Q.device);
334 165 detail::adopt_output(O, Q.device);
335 165 v.flash_attention_forward(Q, K, V, d_mask, num_heads, causal, O);
336 165 }
337
338 28 void flash_attention_gqa_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
339 const float* d_mask, int num_q_heads, int num_kv_heads,
340 bool causal, Tensor& O) {
341 28 const auto& v = detail::dispatch(Q, K, V, O);
342
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 if (!v.flash_attention_gqa_forward)
343 ✗ detail::throw_not_implemented("flash_attention_gqa_forward", Q.device);
344 28 detail::adopt_output(O, Q.device);
345 28 v.flash_attention_gqa_forward(Q, K, V, d_mask, num_q_heads, num_kv_heads, causal, O);
346 28 }
347
348 51 void flash_attention_windowed_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
349 const float* d_mask, int num_heads, int window,
350 Tensor& O, bool causal) {
351 51 const auto& v = detail::dispatch(Q, K, V, O);
352
1/2
✓ Branch 0 taken 51 times.
✗ Branch 1 not taken.
 51 if (!v.flash_attention_windowed_forward)
353 ✗ detail::throw_not_implemented("flash_attention_windowed_forward", Q.device);
354 51 detail::adopt_output(O, Q.device);
355 51 v.flash_attention_windowed_forward(Q, K, V, d_mask, num_heads, window, O, causal);
356 51 }
357
358 57 void flash_attention_varlen_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
359 const int32_t* cu_seqlens_q,
360 const int32_t* cu_seqlens_k,
361 int batch_size, int max_seqlen_q, int max_seqlen_k,
362 int num_heads, int head_dim, bool causal,
363 Tensor& O) {
364 57 const auto& v = detail::dispatch(Q, K, V, O);
365
1/2
✓ Branch 0 taken 57 times.
✗ Branch 1 not taken.
 57 if (!v.flash_attention_varlen_forward)
366 ✗ detail::throw_not_implemented("flash_attention_varlen_forward", Q.device);
367 57 detail::adopt_output(O, Q.device);
368 114 v.flash_attention_varlen_forward(Q, K, V, cu_seqlens_q, cu_seqlens_k,
369 57 batch_size, max_seqlen_q, max_seqlen_k,
370 57 num_heads, head_dim, causal, O);
371 57 }
372
373 47 void flash_attention_varlen_backward(const Tensor& Q, const Tensor& K, const Tensor& V,
374 const Tensor& O, const Tensor& dO,
375 const int32_t* cu_seqlens_q,
376 const int32_t* cu_seqlens_k,
377 int batch_size, int max_seqlen_q, int max_seqlen_k,
378 int num_heads, int head_dim, bool causal,
379 Tensor& dQ, Tensor& dK, Tensor& dV) {
380 47 const auto& v = detail::dispatch(Q, K, V, O, dO, dQ, dK, dV);
381
1/2
✓ Branch 0 taken 47 times.
✗ Branch 1 not taken.
 47 if (!v.flash_attention_varlen_backward)
382 ✗ detail::throw_not_implemented("flash_attention_varlen_backward", Q.device);
383 47 detail::adopt_output(dQ, Q.device);
384 47 detail::adopt_output(dK, Q.device);
385 47 detail::adopt_output(dV, Q.device);
386 94 v.flash_attention_varlen_backward(Q, K, V, O, dO,
387 47 cu_seqlens_q, cu_seqlens_k,
388 47 batch_size, max_seqlen_q, max_seqlen_k,
389 47 num_heads, head_dim, causal,
390 47 dQ, dK, dV);
391 47 }
392
393 39 void flash_attention_qkvo_forward(const Tensor& X, const Tensor* Ctx,
394 const Tensor& Wq, const Tensor* bq,
395 const Tensor& Wk, const Tensor* bk,
396 const Tensor& Wv, const Tensor* bv,
397 const Tensor& Wo, const Tensor* bo,
398 const float* d_mask, int num_heads,
399 bool causal, Tensor& O) {
400 39 const auto& v = detail::dispatch_with_opts(
401 39 X, Wq, {Ctx, bq, &Wk, bk, &Wv, bv, &Wo, bo, &O});
402
1/2
✓ Branch 0 taken 39 times.
✗ Branch 1 not taken.
 39 if (!v.flash_attention_qkvo_forward)
403 ✗ detail::throw_not_implemented("flash_attention_qkvo_forward", X.device);
404 39 detail::adopt_output(O, X.device);
405 78 v.flash_attention_qkvo_forward(X, Ctx, Wq, bq, Wk, bk, Wv, bv, Wo, bo,
406 39 d_mask, num_heads, causal, O);
407 39 }
408
409 39 void flash_attention_qkvo_backward(
410 const Tensor& X, const Tensor* Ctx,
411 const Tensor& Wq, const Tensor* bq,
412 const Tensor& Wk, const Tensor* bk,
413 const Tensor& Wv, const Tensor* bv,
414 const Tensor& Wo, const Tensor* bo,
415 const float* d_mask, int num_heads, bool causal,
416 const Tensor& dO,
417 Tensor& dX, Tensor* dCtx,
418 Tensor& dWq, Tensor* dbq,
419 Tensor& dWk, Tensor* dbk,
420 Tensor& dWv, Tensor* dbv,
421 Tensor& dWo, Tensor* dbo) {
422 39 const auto& v = detail::dispatch_with_opts(
423 429 X, Wq, {Ctx, bq, &Wk, bk, &Wv, bv, &Wo, bo, &dO,
424 390 &dX, dCtx, &dWq, dbq, &dWk, dbk, &dWv, dbv, &dWo, dbo});
425
1/2
✓ Branch 0 taken 39 times.
✗ Branch 1 not taken.
 39 if (!v.flash_attention_qkvo_backward)
426 ✗ detail::throw_not_implemented("flash_attention_qkvo_backward", X.device);
427 39 detail::adopt_output(dX, X.device);
428
2/2
✓ Branch 0 taken 25 times.
✓ Branch 1 taken 14 times.
 39 if (dCtx) detail::adopt_output(*dCtx, X.device);
429 39 detail::adopt_output(dWq, X.device);
430
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 26 times.
 39 if (dbq) detail::adopt_output(*dbq, X.device);
431 39 detail::adopt_output(dWk, X.device);
432
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 26 times.
 39 if (dbk) detail::adopt_output(*dbk, X.device);
433 39 detail::adopt_output(dWv, X.device);
434
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 26 times.
 39 if (dbv) detail::adopt_output(*dbv, X.device);
435 39 detail::adopt_output(dWo, X.device);
436
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 26 times.
 39 if (dbo) detail::adopt_output(*dbo, X.device);
437 78 v.flash_attention_qkvo_backward(X, Ctx, Wq, bq, Wk, bk, Wv, bv, Wo, bo,
438 39 d_mask, num_heads, causal, dO,
439 39 dX, dCtx, dWq, dbq, dWk, dbk,
440 39 dWv, dbv, dWo, dbo);
441 39 }
442
443 74 void flash_attention_backward(const Tensor& Q, const Tensor& K, const Tensor& V,
444 const Tensor& O, const Tensor& dO,
445 const float* d_mask, int num_heads, bool causal,
446 Tensor& dQ, Tensor& dK, Tensor& dV) {
447 74 const auto& v = detail::dispatch(Q, K, V, O, dO, dQ, dK, dV);
448
1/2
✓ Branch 0 taken 74 times.
✗ Branch 1 not taken.
 74 if (!v.flash_attention_backward) detail::throw_not_implemented("flash_attention_backward", Q.device);
449 74 detail::adopt_output(dQ, Q.device);
450 74 detail::adopt_output(dK, Q.device);
451 74 detail::adopt_output(dV, Q.device);
452 148 v.flash_attention_backward(Q, K, V, O, dO, d_mask, num_heads, causal,
453 74 dQ, dK, dV);
454 74 }
455
456 11 void flash_attention_project_kv(const Tensor& ctx,
457 const Tensor& Wk, const Tensor* bk,
458 const Tensor& Wv, const Tensor* bv,
459 Tensor& K_out, Tensor& V_out) {
460 11 const auto& v = detail::dispatch_with_opts(ctx, Wk, {bk, &Wv, bv, &K_out, &V_out});
461
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.flash_attention_project_kv)
462 ✗ detail::throw_not_implemented("flash_attention_project_kv", ctx.device);
463 11 detail::adopt_output(K_out, ctx.device);
464 11 detail::adopt_output(V_out, ctx.device);
465 11 v.flash_attention_project_kv(ctx, Wk, bk, Wv, bv, K_out, V_out);
466 11 }
467
468 15 void flash_attention_q_with_kv_cached_forward(const Tensor& X,
469 const Tensor& K, const Tensor& V,
470 const Tensor& Wq, const Tensor* bq,
471 const Tensor& Wo, const Tensor* bo,
472 const float* d_mask, int num_heads,
473 bool causal, Tensor& O) {
474 15 const auto& v = detail::dispatch_with_opts(X, K, {&V, &Wq, bq, &Wo, bo, &O});
475
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.flash_attention_q_with_kv_cached_forward)
476 ✗ detail::throw_not_implemented("flash_attention_q_with_kv_cached_forward", X.device);
477 15 detail::adopt_output(O, X.device);
478 30 v.flash_attention_q_with_kv_cached_forward(X, K, V, Wq, bq, Wo, bo,
479 15 d_mask, num_heads, causal, O);
480 15 }
481
482 // ─── Matmul + RoPE + RMSNorm + SwiGLU + KV-cache + Llama ───────────────────
483
484 114 void matmul(const Tensor& A, const Tensor& B, Tensor& C) {
485 114 const auto& v = detail::dispatch(A, B, C);
486
1/2
✓ Branch 0 taken 114 times.
✗ Branch 1 not taken.
 114 if (!v.matmul) detail::throw_not_implemented("matmul", A.device);
487 114 detail::adopt_output(C, A.device);
488 114 v.matmul(A, B, C);
489 114 }
490
491 80 void matmul_abt(const Tensor& A, const Tensor& B, Tensor& C,
492 int batch, int M, int N, int K,
493 long long strideA, long long strideB, long long strideC,
494 const Tensor* bias, int act) {
495 80 const auto& v = detail::dispatch(A, B, C);
496
1/2
✓ Branch 0 taken 80 times.
✗ Branch 1 not taken.
 80 if (!v.matmul_abt) detail::throw_not_implemented("matmul_abt", A.device);
497 80 detail::adopt_output(C, A.device);
498 80 v.matmul_abt(A, B, C, batch, M, N, K, strideA, strideB, strideC, bias, act);
499 80 }
500
501 62 void matmul_backward(const Tensor& A, const Tensor& B, const Tensor& dC,
502 Tensor& dA, Tensor& dB) {
503 62 const auto& v = detail::dispatch(A, B, dC, dA, dB);
504
1/2
✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
 62 if (!v.matmul_backward) detail::throw_not_implemented("matmul_backward", A.device);
505 62 detail::adopt_output(dA, A.device);
506 62 detail::adopt_output(dB, A.device);
507 62 v.matmul_backward(A, B, dC, dA, dB);
508 62 }
509
510 909 void lstm_forward_train(const Tensor& X, const Tensor& W_ih, const Tensor& W_hh,
511 const Tensor* b_ih, const Tensor* b_hh,
512 const Tensor* h0, const Tensor* c0, int T, int B,
513 Tensor& Y, Tensor& gates, Tensor& C,
514 Tensor* hT, Tensor* cT) {
515 909 const auto& v = detail::dispatch(X);
516
1/2
✓ Branch 0 taken 909 times.
✗ Branch 1 not taken.
 909 if (!v.lstm_forward_train) detail::throw_not_implemented("lstm_forward_train", X.device);
517 909 detail::adopt_output(Y, X.device);
518 909 detail::adopt_output(gates, X.device);
519 909 detail::adopt_output(C, X.device);
520
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 14 times.
 909 if (hT) detail::adopt_output(*hT, X.device);
521
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 14 times.
 909 if (cT) detail::adopt_output(*cT, X.device);
522 909 v.lstm_forward_train(X, W_ih, W_hh, b_ih, b_hh, h0, c0, T, B, Y, gates, C, hT, cT);
523 909 }
524
525 415 void lstm_backward(const Tensor& X, const Tensor& W_ih, const Tensor& W_hh,
526 const Tensor* h0, const Tensor* c0,
527 const Tensor& Y, const Tensor& gates, const Tensor& C,
528 const Tensor& dY, int T, int B,
529 Tensor& dX, Tensor& dW_ih, Tensor& dW_hh,
530 Tensor* db_ih, Tensor* db_hh,
531 Tensor* dh0, Tensor* dc0) {
532 415 const auto& v = detail::dispatch(X);
533
1/2
✓ Branch 0 taken 415 times.
✗ Branch 1 not taken.
 415 if (!v.lstm_backward) detail::throw_not_implemented("lstm_backward", X.device);
534 415 detail::adopt_output(dX, X.device);
535 415 detail::adopt_output(dW_ih, X.device);
536 415 detail::adopt_output(dW_hh, X.device);
537
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 415 times.
 415 if (db_ih) detail::adopt_output(*db_ih, X.device);
538
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 413 times.
 415 if (db_hh) detail::adopt_output(*db_hh, X.device);
539
2/2
✓ Branch 0 taken 402 times.
✓ Branch 1 taken 13 times.
 415 if (dh0) detail::adopt_output(*dh0, X.device);
540
2/2
✓ Branch 0 taken 402 times.
✓ Branch 1 taken 13 times.
 415 if (dc0) detail::adopt_output(*dc0, X.device);
541 830 v.lstm_backward(X, W_ih, W_hh, h0, c0, Y, gates, C, dY, T, B,
542 415 dX, dW_ih, dW_hh, db_ih, db_hh, dh0, dc0);
543 415 }
544
545 16 void rope_forward(const Tensor& X, int head_dim, int num_heads,
546 int seq_offset, float theta_base, Tensor& Y) {
547 16 const auto& v = detail::dispatch(X, Y);
548
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.rope_forward) detail::throw_not_implemented("rope_forward", X.device);
549 16 detail::adopt_output(Y, X.device);
550 16 v.rope_forward(X, head_dim, num_heads, seq_offset, theta_base, Y);
551 16 }
552
553 10 void rope_backward(const Tensor& dY, int head_dim, int num_heads,
554 int seq_offset, float theta_base, Tensor& dX) {
555 10 const auto& v = detail::dispatch(dY, dX);
556
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.rope_backward) detail::throw_not_implemented("rope_backward", dY.device);
557 10 detail::adopt_output(dX, dY.device);
558 10 v.rope_backward(dY, head_dim, num_heads, seq_offset, theta_base, dX);
559 10 }
560
561 38 void rope_apply(const Tensor& X, const Tensor& cos_tbl, const Tensor& sin_tbl,
562 int head_dim, int num_heads, Tensor& Y) {
563 38 const auto& v = detail::dispatch(X, cos_tbl, sin_tbl, Y);
564
1/2
✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
 38 if (!v.rope_apply) detail::throw_not_implemented("rope_apply", X.device);
565 38 detail::adopt_output(Y, X.device);
566 38 v.rope_apply(X, cos_tbl, sin_tbl, head_dim, num_heads, Y);
567 38 }
568
569 6 void rope_apply_perhead(const Tensor& X, const Tensor& cos_tbl,
570 const Tensor& sin_tbl, int head_dim, int num_heads,
571 Tensor& Y) {
572 6 const auto& v = detail::dispatch(X, cos_tbl, sin_tbl, Y);
573
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.rope_apply_perhead)
574 ✗ detail::throw_not_implemented("rope_apply_perhead", X.device);
575 6 detail::adopt_output(Y, X.device);
576 6 v.rope_apply_perhead(X, cos_tbl, sin_tbl, head_dim, num_heads, Y);
577 6 }
578
579 4 void rope_apply_backward(const Tensor& dY, const Tensor& cos_tbl,
580 const Tensor& sin_tbl, int head_dim, int num_heads,
581 Tensor& dX) {
582 4 const auto& v = detail::dispatch(dY, cos_tbl, sin_tbl, dX);
583
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rope_apply_backward)
584 ✗ detail::throw_not_implemented("rope_apply_backward", dY.device);
585 4 detail::adopt_output(dX, dY.device);
586 4 v.rope_apply_backward(dY, cos_tbl, sin_tbl, head_dim, num_heads, dX);
587 4 }
588
589 15 void rope_apply_mrope(const Tensor& X,
590 const Tensor& cos_t, const Tensor& sin_t,
591 const Tensor& cos_h, const Tensor& sin_h,
592 const Tensor& cos_w, const Tensor& sin_w,
593 const int32_t* pos_t, const int32_t* pos_h,
594 const int32_t* pos_w,
595 int head_dim, int num_heads,
596 int d_t, int d_h, int d_w,
597 Tensor& Y) {
598 30 const auto& v = detail::dispatch(X, cos_t, sin_t, cos_h, sin_h, cos_w,
599 15 sin_w, Y);
600
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.rope_apply_mrope)
601 ✗ detail::throw_not_implemented("rope_apply_mrope", X.device);
602 15 detail::adopt_output(Y, X.device);
603 30 v.rope_apply_mrope(X, cos_t, sin_t, cos_h, sin_h, cos_w, sin_w,
604 15 pos_t, pos_h, pos_w, head_dim, num_heads,
605 15 d_t, d_h, d_w, Y);
606 15 }
607
608 754 void rms_norm_forward(const Tensor& X, const Tensor& gamma, float eps, Tensor& Y) {
609 754 const auto& v = detail::dispatch(X, gamma, Y);
610
1/2
✓ Branch 0 taken 754 times.
✗ Branch 1 not taken.
 754 if (!v.rms_norm_forward) detail::throw_not_implemented("rms_norm_forward", X.device);
611 754 detail::adopt_output(Y, X.device);
612 754 v.rms_norm_forward(X, gamma, eps, Y);
613 754 }
614
615 16 void rms_norm_backward(const Tensor& X, const Tensor& gamma, const Tensor& dY,
616 float eps, Tensor& dX, Tensor& dGamma) {
617 16 const auto& v = detail::dispatch(X, gamma, dY, dX, dGamma);
618
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.rms_norm_backward) detail::throw_not_implemented("rms_norm_backward", X.device);
619 16 detail::adopt_output(dX, X.device);
620 16 detail::adopt_output(dGamma, X.device);
621 16 v.rms_norm_backward(X, gamma, dY, eps, dX, dGamma);
622 16 }
623
624 496 void swiglu_forward(const Tensor& X, Tensor& Y) {
625 496 const auto& v = detail::dispatch(X, Y);
626
1/2
✓ Branch 0 taken 496 times.
✗ Branch 1 not taken.
 496 if (!v.swiglu_forward) detail::throw_not_implemented("swiglu_forward", X.device);
627 496 detail::adopt_output(Y, X.device);
628 496 v.swiglu_forward(X, Y);
629 496 }
630 12 void swiglu_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
631 12 const auto& v = detail::dispatch(X, dY, dX);
632
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.swiglu_backward) detail::throw_not_implemented("swiglu_backward", X.device);
633 12 detail::adopt_output(dX, X.device);
634 12 v.swiglu_backward(X, dY, dX);
635 12 }
636
637 36 void kv_cache_append(const Tensor& K_new, const Tensor& V_new, int cur_len,
638 Tensor& K_cache, Tensor& V_cache) {
639 36 const auto& v = detail::dispatch(K_new, V_new, K_cache, V_cache);
640
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (!v.kv_cache_append) detail::throw_not_implemented("kv_cache_append", K_new.device);
641 36 detail::adopt_output(K_cache, K_new.device);
642 36 detail::adopt_output(V_cache, K_new.device);
643 36 v.kv_cache_append(K_new, V_new, cur_len, K_cache, V_cache);
644 36 }
645
646 61 void flash_attention_decode(const Tensor& Q,
647 const Tensor& K_cache, const Tensor& V_cache,
648 int valid_len, int num_q_heads, int num_kv_heads,
649 Tensor& O, float attn_softcap, int window) {
650 61 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
651
1/2
✓ Branch 0 taken 61 times.
✗ Branch 1 not taken.
 61 if (!v.flash_attention_decode) detail::throw_not_implemented("flash_attention_decode", Q.device);
652 61 detail::adopt_output(O, Q.device);
653 122 v.flash_attention_decode(Q, K_cache, V_cache, valid_len,
654 61 num_q_heads, num_kv_heads, O, attn_softcap, window);
655 61 }
656
657 16 void flash_attention_decode_masked(const Tensor& Q,
658 const Tensor& K_cache,
659 const Tensor& V_cache,
660 const float* d_mask,
661 int num_q_heads, int num_kv_heads,
662 Tensor& O, float attn_softcap, int window) {
663 16 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
664
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.flash_attention_decode_masked) {
665 ✗ detail::throw_not_implemented("flash_attention_decode_masked", Q.device);
666 }
667 16 detail::adopt_output(O, Q.device);
668 32 v.flash_attention_decode_masked(Q, K_cache, V_cache, d_mask,
669 16 num_q_heads, num_kv_heads, O,
670 16 attn_softcap, window);
671 16 }
672
673 52 void flash_attention_decode_len(const Tensor& Q, const Tensor& K_cache, const Tensor& V_cache,
674 const int32_t* d_len, int num_q_heads, int num_kv_heads, Tensor& O) {
675 52 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
676
1/2
✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
 52 if (!v.flash_attention_decode_len) detail::throw_not_implemented("flash_attention_decode_len", Q.device);
677 52 detail::adopt_output(O, Q.device);
678 52 v.flash_attention_decode_len(Q, K_cache, V_cache, d_len, num_q_heads, num_kv_heads, O);
679 52 }
680
681 // ─── Public reductions ─────────────────────────────────────────────────────
682
683 15 void sum_rows(const Tensor& X, Tensor& Y) {
684 15 const auto& v = detail::dispatch(X, Y);
685
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.sum_rows) detail::throw_not_implemented("sum_rows", X.device);
686 15 detail::adopt_output(Y, X.device);
687 15 v.sum_rows(X, Y);
688 15 }
689 14 void sum_cols(const Tensor& X, Tensor& Y) {
690 14 const auto& v = detail::dispatch(X, Y);
691
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.sum_cols) detail::throw_not_implemented("sum_cols", X.device);
692 14 detail::adopt_output(Y, X.device);
693 14 v.sum_cols(X, Y);
694 14 }
695 23 void argmax_rows(const Tensor& X, Tensor& Idx) {
696 23 const auto& v = detail::dispatch(X, Idx);
697
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.argmax_rows) detail::throw_not_implemented("argmax_rows", X.device);
698 23 detail::adopt_output(Idx, X.device);
699 23 v.argmax_rows(X, Idx);
700 23 }
701 5 void segment_softmax_stats(const Tensor& logits, const Tensor& seg_offsets, Tensor& out) {
702 5 const auto& v = detail::dispatch(logits, seg_offsets, out);
703
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.segment_softmax_stats)
704 ✗ detail::throw_not_implemented("segment_softmax_stats", logits.device);
705 5 detail::adopt_output(out, logits.device);
706 5 v.segment_softmax_stats(logits, seg_offsets, out);
707 5 }
708 1203 void flash_attention_packed_qkv_forward(const Tensor& QKV, const Tensor& seq_bounds,
709 int num_heads, int window, Tensor& O) {
710 1203 const auto& v = detail::dispatch(QKV, seq_bounds, O);
711
1/2
✓ Branch 0 taken 1203 times.
✗ Branch 1 not taken.
 1203 if (!v.flash_attention_packed_qkv_forward)
712 ✗ detail::throw_not_implemented("flash_attention_packed_qkv_forward", QKV.device);
713 1203 detail::adopt_output(O, QKV.device);
714 1203 v.flash_attention_packed_qkv_forward(QKV, seq_bounds, num_heads, window, O);
715 1203 }
716 22 void flash_attention_packed_qkv_backward(const Tensor& QKV, const Tensor& dO, const Tensor& seq_bounds,
717 int num_heads, int window, Tensor& dQKV) {
718 22 const auto& v = detail::dispatch(QKV, dO, seq_bounds, dQKV);
719
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 22 if (!v.flash_attention_packed_qkv_backward)
720 ✗ detail::throw_not_implemented("flash_attention_packed_qkv_backward", QKV.device);
721 22 detail::adopt_output(dQKV, QKV.device);
722 22 v.flash_attention_packed_qkv_backward(QKV, dO, seq_bounds, num_heads, window, dQKV);
723 22 }
724 8 void rope_qkv_packed_inplace(Tensor& QKV, const Tensor& cos_tbl, const Tensor& sin_tbl,
725 const Tensor& pos, int num_heads, int head_dim) {
726 8 const auto& v = detail::dispatch(QKV, cos_tbl, sin_tbl, pos);
727
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.rope_qkv_packed_inplace)
728 ✗ detail::throw_not_implemented("rope_qkv_packed_inplace", QKV.device);
729 8 v.rope_qkv_packed_inplace(QKV, cos_tbl, sin_tbl, pos, num_heads, head_dim);
730 8 }
731 13 void rows_count_above(const Tensor& X, float t_lo, float t_hi, Tensor& counts) {
732 13 const auto& v = detail::dispatch(X, counts);
733
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (!v.rows_count_above)
734 ✗ detail::throw_not_implemented("rows_count_above", X.device);
735 13 detail::adopt_output(counts, X.device);
736 13 v.rows_count_above(X, t_lo, t_hi, counts);
737 13 }
738 11 void threshold_u8(const Tensor& X, float t, Tensor& Y) {
739 11 const auto& v = detail::dispatch(X, Y);
740
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.threshold_u8) detail::throw_not_implemented("threshold_u8", X.device);
741 11 detail::adopt_output(Y, X.device);
742 11 v.threshold_u8(X, t, Y);
743 11 }
744
745 } // namespace brotensor
746