← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
ckernel_engine.h
Go to the documentation of this file.
1#ifndef CKERNEL_ENGINE_H
2#define CKERNEL_ENGINE_H
3
4#include <stddef.h>
5#include <stdint.h>
6#include "ckernel_audio.h"
7#include "cpu_features.h"
8#include "ckernel_quant.h" /* INT8 block types (block_q8_0, block_q8_K, etc.) */
9#include "mega_fused_attention.h"
10#include "ckernel_rwkv7_decode.h"
11
12#ifdef __cplusplus
13extern "C" {
14#endif
15
16/**
17 * Core math backend interface for C-Kernel-Engine.
18 *
19 * This is intentionally minimal and matches the conventions already used
20 * in C-Transformer for GEMM kernels.
21 *
22 * Layout assumptions (LLM-style shapes):
23 * - A: [M x K], row-major, A(i,k) = A[i*K + k]
24 * - B: [N x K], row-major, B(j,k) = B[j*K + k]
25 * - C: [M x N], row-major, C(i,j) = C[i*N + j]
26 * - bias: optional [N], added per output column j
27 */
28typedef struct {
29 void (*sgemm)(int M, int N, int K,
30 const float *A, int lda,
31 const float *B, int ldb,
32 const float *bias,
33 float *C, int ldc);
34} CKMathBackend;
35
36/**
37 * Obtain the built-in native backend (single-node CPU, C + intrinsics).
38 */
39CKMathBackend ckernel_backend_native(void);
40
41// Enable stricter numeric parity (single-thread + double-accumulation GEMM).
42void ck_set_strict_parity(int enabled);
43int ck_strict_parity_enabled(void);
44void ck_strict_store_next_gemm_a(const float *data, size_t elems);
45const float *ck_strict_consume_next_gemm_a(size_t elems);
46int ck_strict_mtmd_clip_encode_planar_f32(const float *planar,
47 int channels,
48 int height,
49 int width,
50 float *out,
51 size_t out_elems);
52
53// Thread configuration - call once at startup
54// num_threads: 0 = auto-detect physical cores, >0 = use specified count
55void ck_set_num_threads(int num_threads);
56int ck_get_num_threads(void);
57int ck_get_physical_cores(void);
58
59void hyper_stream_expand_bf16(const float *input,
60 float *output,
61 int rows,
62 int streams,
63 int hidden_dim);
64void hyper_stream_expand_f32(const float *input,
65 float *output,
66 int rows,
67 int streams,
68 int hidden_dim);
69void hyper_connection_mix_bf16(const float *hyper_input,
70 const float *norm_weight,
71 const uint16_t *mix_down_weight,
72 const uint16_t *mix_up_weight,
73 const uint16_t *inject_weight,
74 float *mixed_output,
75 float *injection_output,
76 float *normalized_scratch,
77 float *dynamic_scratch,
78 float *mix_scratch,
79 int rows,
80 int streams,
81 int hidden_dim,
82 int dynamic_dim,
83 float eps,
84 int emit_injection);
85void hyper_connection_mix_q4k_q5_0_q4k(const float *hyper_input,
86 const float *norm_weight,
87 const void *mix_down_weight,
88 const void *mix_up_weight,
89 const void *inject_weight,
90 float *mixed_output,
91 float *injection_output,
92 float *normalized_scratch,
93 float *dynamic_scratch,
94 float *mix_scratch,
95 int rows,
96 int streams,
97 int hidden_dim,
98 int dynamic_dim,
99 float eps,
100 int emit_injection);
101void hyper_connection_mix_q6k_q5_0_q4k(const float *hyper_input,
102 const float *norm_weight,
103 const void *mix_down_weight,
104 const void *mix_up_weight,
105 const void *inject_weight,
106 float *mixed_output,
107 float *injection_output,
108 float *normalized_scratch,
109 float *dynamic_scratch,
110 float *mix_scratch,
111 int rows,
112 int streams,
113 int hidden_dim,
114 int dynamic_dim,
115 float eps,
116 int emit_injection);
117void hyper_stream_inject_bf16(const float *hyper_input,
118 const float *block_output,
119 const float *injection_weight,
120 float *output,
121 int rows,
122 int streams,
123 int hidden_dim);
124void hyper_stream_inject_f32(const float *hyper_input,
125 const float *block_output,
126 const float *injection_weight,
127 float *output,
128 int rows,
129 int streams,
130 int hidden_dim);
131
132void qwen4_ple_ngram_embed_bf16(const int32_t *token_ids,
133 const uint16_t *embedding,
134 const int64_t *layer_multipliers,
135 const int64_t *head_offsets,
136 const int64_t *head_vocab_sizes,
137 float *output,
138 const float *token_state_in,
139 float *token_state_out,
140 int rows,
141 int ngram_size,
142 int heads_per_ngram,
143 int head_dim,
144 int eos_token_id,
145 int position);
146void qwen4_ple_ngram_embed_q5_0(const int32_t *token_ids,
147 const void *embedding,
148 const int64_t *layer_multipliers,
149 const int64_t *head_offsets,
150 const int64_t *head_vocab_sizes,
151 float *output,
152 const float *token_state_in,
153 float *token_state_out,
154 int rows,
155 int ngram_size,
156 int heads_per_ngram,
157 int head_dim,
158 int eos_token_id,
159 int position);
160void qwen4_ple_gate_conv_inject_bf16(
161 const float *hyper_input,
162 const float *key_projected,
163 const float *value_projected,
164 const float *norm_key_weight,
165 const float *norm_query_weight,
166 const float *norm_conv_weight,
167 const uint16_t *conv_weight,
168 float *hyper_output,
169 float *key_norm_scratch,
170 float *query_norm_scratch,
171 float *gated_scratch,
172 float *conv_norm_scratch,
173 const float *conv_state_in,
174 float *conv_state_out,
175 int rows,
176 int streams,
177 int hidden_dim,
178 int kernel_size,
179 int dilation,
180 float eps);
181void qwen4_ple_gate_conv_inject_fp16(
182 const float *hyper_input,
183 const float *key_projected,
184 const float *value_projected,
185 const float *norm_key_weight,
186 const float *norm_query_weight,
187 const float *norm_conv_weight,
188 const uint16_t *conv_weight,
189 float *hyper_output,
190 float *key_norm_scratch,
191 float *query_norm_scratch,
192 float *gated_scratch,
193 float *conv_norm_scratch,
194 const float *conv_state_in,
195 float *conv_state_out,
196 int rows,
197 int streams,
198 int hidden_dim,
199 int kernel_size,
200 int dilation,
201 float eps);
202void qwen4_ple_gate_conv_inject_llama_fp16(
203 const float *hyper_input,
204 const float *key_projected,
205 const float *value_projected,
206 const float *norm_key_weight,
207 const float *norm_query_weight,
208 const float *norm_conv_weight,
209 const uint16_t *conv_weight,
210 float *hyper_output,
211 float *key_norm_scratch,
212 float *query_norm_scratch,
213 float *gated_scratch,
214 float *conv_norm_scratch,
215 const float *conv_state_in,
216 float *conv_state_out,
217 int rows,
218 int streams,
219 int hidden_dim,
220 int kernel_size,
221 int dilation,
222 float eps);
223void qwen4_qsa_index_select_bf16(
224 const float *projected_qk, const float *index_key_cache_in,
225 const float *q_norm_weight, const float *k_norm_weight,
226 float *selected_indices, float *index_key_cache_out,
227 float *q_norm_scratch, float *pooled_key_scratch,
228 float *block_score_scratch, int32_t *block_index_scratch,
229 int rows, int query_heads, int index_head_dim, int token_budget,
230 int compress_ratio, int rotary_dim, int context_length,
231 int position, float rope_theta, float eps);
232// Residual kernels used by IR/codegen runtime paths.
233void ck_residual_add_token_major(const float *a,
234 const float *b,
235 float *out,
236 int tokens,
237 int aligned_embed_dim);
238void ck_residual_add_token_major_parallel_dispatch(
239 const float *a,
240 const float *b,
241 float *out,
242 int tokens,
243 int aligned_embed_dim);
244void *ck_memcpy_parallel_dispatch(void *dst, const void *src, size_t size);
245void ck_residual_add_token_major_bf16_storage(const float *a,
246 const float *b,
247 float *out,
248 int tokens,
249 int aligned_embed_dim);
250void ck_residual_add_backward(const float *d_out,
251 float *d_a,
252 float *d_b,
253 int tokens,
254 int aligned_embed_dim);
255
256// Expose the individual GEMM kernels copied from C-Transformer.
257void gemm_naive_parallel(const float *A,
258 const float *B,
259 const float *bias,
260 float *C,
261 int M, int N, int K);
262void gemm_nt_fp32_exact_parallel_dispatch(const float *A,
263 const float *B,
264 const float *bias,
265 float *C,
266 int M, int N, int K);
267
268void gemm_avx512_parallel(const float *A,
269 const float *B,
270 const float *bias,
271 float *C,
272 int M, int N, int K);
273
274void gemm_fine_grained_parallel(const float *A,
275 const float *B,
276 const float *bias,
277 float *C,
278 int M, int N, int K);
279
280 void gemm_blocked_serial(const float *A,
281 const float *B,
282 const float *bias,
283 float *C,
284 int M, int N, int K);
285
286 void gemm_nt_f32_llama_production(const float *A,
287 const float *B,
288 const float *bias,
289 float *C,
290 int M, int N, int K);
291
292 void gemm_nt_f32_llama_production_output_range(const float *A,
293 const float *B,
294 const float *bias,
295 float *C,
296 int M, int N, int K,
297 int output_begin,
298 int output_end);
299
300 /* Training dispatch wrapper: same contract as gemm_blocked_serial.
301 * Uses CK threadpool for row/column partitioning and falls back to serial. */
302 void gemm_blocked_serial_train_parallel_dispatch(const float *A,
303 const float *B,
304 const float *bias,
305 float *C,
306 int M, int N, int K);
307
308 /* Training backward GEMM dispatch wrapper.
309 * Contract matches fc2_backward_kernel/fc1_backward_kernel and uses
310 * CK threadpool for deterministic row/column partitioning. */
311 void gemm_backward_f32_train_parallel_dispatch(const float *d_output,
312 const float *input,
313 const float *W,
314 float *d_input,
315 float *d_W,
316 float *d_b,
317 int T,
318 int aligned_in,
319 int aligned_out,
320 int num_threads);
321
322 /* v2 wrapper: more aggressive threadpool dispatch for training backward GEMM
323 * (especially T=1 microstep shapes). Keeps kernel math unchanged. */
324 void gemm_backward_f32_train_parallel_dispatch_v2(const float *d_output,
325 const float *input,
326 const float *W,
327 float *d_input,
328 float *d_W,
329 float *d_b,
330 int T,
331 int aligned_in,
332 int aligned_out,
333 int num_threads);
334
335/* Optimizer / gradient update kernels used by v7 train runtime. */
336void adamw_update_f32(const float *grad,
337 float *weight,
338 float *m,
339 float *v,
340 size_t numel,
341 float lr,
342 float beta1,
343 float beta2,
344 float eps,
345 float weight_decay,
346 int step);
347
348void adamw_clip_update_multi_f32(float *const *grads,
349 float *const *weights,
350 float *const *m_states,
351 float *const *v_states,
352 const size_t *numels,
353 int tensor_count,
354 float lr,
355 float beta1,
356 float beta2,
357 float eps,
358 float weight_decay,
359 float max_grad_norm,
360 int step);
361
362void gradient_accumulate_f32(float *dst, const float *src, size_t numel);
363void gradient_accumulate_multi_f32(float *const *dsts,
364 const float *const *srcs,
365 const size_t *numels,
366 int tensor_count);
367void gradient_scale_f32(float *grad, size_t numel, float scale);
368float gradient_clip_norm_f32(float *grad, size_t numel, float max_norm);
369float gradient_global_norm_multi_f32(const float *const *grads,
370 const size_t *numels,
371 int tensor_count);
372
373 // Reference BF16 GEMM (A/B/bias in BF16, output BF16).
374void gemm_blocked_serial_bf16(const uint16_t *A,
375 const uint16_t *B,
376 const uint16_t *bias,
377 uint16_t *C,
378 int M, int N, int K);
379
380// BF16 storage, FP32 gradient accumulation for linear backward.
381void gemm_backward_bf16_mixed(const uint16_t *d_output,
382 const uint16_t *input,
383 const uint16_t *weight,
384 float *d_input,
385 float *d_weight,
386 float *d_bias,
387 int tokens,
388 int in_dim,
389 int out_dim);
390
391void gemm_nt_f16(const float *A,
392 const void *B,
393 const float *bias,
394 float *C,
395 int M, int N, int K);
396// SIMD lane count used by the compile-time llama-compatible FP16 dot provider.
397int ck_gemm_nt_f16_simd_lanes(void);
398
399// Diagnostic oracle backed by the dynamically loaded ggml mul_mat graph.
400// Returns zero instead of falling back when the independent oracle is unavailable.
401int ck_gemm_nt_f16_ggml_oracle(const float *A,
402 const void *B,
403 const float *bias,
404 float *C,
405 int M, int N, int K);
406
407void gemm_nt_f16_clipped(const float *A,
408 const void *B,
409 const float *bias,
410 const float *input_min,
411 const float *input_max,
412 const float *output_min,
413 const float *output_max,
414 float *C,
415 int M, int N, int K);
416
417void gemv_bf16(float *y,
418 const void *W,
419 const float *x,
420 int M, int K);
421void gemv_bf16_parallel_dispatch(float *y,
422 const void *W,
423 const float *x,
424 int M, int K);
425void gemv_bf16_bf16_storage(float *y,
426 const void *W,
427 const float *x,
428 int M, int K);
429void gemv_bf16_bf16_storage_parallel_dispatch(float *y,
430 const void *W,
431 const float *x,
432 int M, int K);
433
434void gemm_nt_bf16(const float *A,
435 const void *B,
436 const float *bias,
437 float *C,
438 int M, int N, int K);
439void gemm_nt_bf16_row_range(const float *A,
440 const void *B,
441 const float *bias,
442 float *C,
443 int M, int N, int K,
444 int row_begin, int row_end);
445void gemm_nt_bf16_parallel_dispatch(const float *A,
446 const void *B,
447 const float *bias,
448 float *C,
449 int M, int N, int K);
450void gemm_nt_bf16_bf16_storage(const float *A,
451 const void *B,
452 const float *bias,
453 float *C,
454 int M, int N, int K);
455void gemm_nt_bf16_bf16_storage_row_range(const float *A,
456 const void *B,
457 const float *bias,
458 float *C,
459 int M, int N, int K,
460 int row_begin, int row_end);
461void gemm_nt_bf16_bf16_storage_parallel_dispatch(const float *A,
462 const void *B,
463 const float *bias,
464 float *C,
465 int M, int N, int K);
466void gemm_nt_bf16_accurate_storage(const float *A,
467 const void *B,
468 const float *bias,
469 float *C,
470 int M, int N, int K);
471void gemm_nt_bf16_accurate_storage_row_range(const float *A,
472 const void *B,
473 const float *bias,
474 float *C,
475 int M, int N, int K,
476 int row_begin, int row_end);
477void gemm_nt_bf16_accurate_storage_parallel_dispatch(const float *A,
478 const void *B,
479 const float *bias,
480 float *C,
481 int M, int N, int K);
482
483void gemm_nt_bf16_native_bf16_storage(const float *A,
484 const void *B,
485 const float *bias,
486 float *C,
487 int M, int N, int K);
488void gemm_nt_bf16_amx_bf16_storage(const float *A,
489 const void *B,
490 const float *bias,
491 float *C,
492 int M, int N, int K);
493void gemm_nt_bf16_amx_bf16_storage_workspace(const float *A,
494 const void *B,
495 const float *bias,
496 float *C,
497 int M, int N, int K,
498 uint16_t *a_bf16,
499 size_t a_bf16_bytes);
500void gemm_nt_bf16_prefill_shape_safe_bf16_storage(const float *A,
501 const void *B,
502 const float *bias,
503 float *C,
504 int M, int N, int K);
505void gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace(
506 const float *A, const void *B, const float *bias, float *C,
507 int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes);
508void gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage(const float *A,
509 const void *B,
510 const float *bias,
511 float *C,
512 int M, int N, int K);
513void gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage(
514 const float *A, const void *B, const float *bias, float *C,
515 int M, int N, int K);
516void gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace(
517 const float *A, const void *B, const float *bias, float *C,
518 int M, int N, int K,
519 uint16_t *input_bf16, size_t input_bf16_bytes,
520 uint16_t *output_bf16, size_t output_bf16_bytes,
521 uint16_t *bias_bf16, size_t bias_bf16_bytes);
522void patch_projection_bf16_pytorch_onednn_conv3d_storage(
523 const float *input, const void *weights, const float *bias, float *output,
524 int batch, int out_channels, int in_channels, int temporal,
525 int patch_h, int patch_w);
526void patch_projection_image_bf16_pytorch_onednn_conv3d_storage(
527 const float *image, const void *weights_t0, const void *weights_t1,
528 const float *bias, float *output, int channels, int image_h, int image_w,
529 int patch_size, int out_channels, int merge_size);
530
531void patch_projection_image_bf16_native_storage(
532 const float *image, const void *weights_t0, const void *weights_t1,
533 const float *bias, float *output, int channels, int image_h, int image_w,
534 int patch_size, int out_channels, int merge_size);
535int ck_gemm_bf16_amx_available(void);
536float ck_attention_pytorch_sdpa_scale_f32(int head_dim);
537void attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage(
538 const float *q, const float *k, const float *v, float *output,
539 int num_heads, int num_kv_heads, int num_tokens,
540 int head_dim, int aligned_head_dim, int kv_stride_tokens);
541void attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_token_output(
542 const float *q, const float *k, const float *v, float *output,
543 int num_heads, int num_kv_heads, int num_tokens,
544 int head_dim, int aligned_head_dim, int kv_stride_tokens);
545int ck_gemm_bf16_fp32out_amx_raw(const uint16_t *A,
546 const uint16_t *B,
547 float *C,
548 int M, int N, int K,
549 int accumulate);
550
551// =============================================================================
552// Quantized (GGML-style) GEMM/GEMV helpers
553// =============================================================================
554//
555// These kernels are used for weight-only quantized inference (e.g. Q4_K_M).
556// The "NT" wrapper matches the engine's common layout:
557// A: [M x K] fp32 (token-major)
558// B: [N x K] quantized (row-major by output channel)
559// C: [M x N] fp32
560//
561// NOTE: Q4_K requires K to be a multiple of 256 (QK_K).
562
563void gemv_q4_k(float *y,
564 const void *W,
565 const float *x,
566 int M, int K);
567
568void gemm_q4_k(float *Y,
569 const void *W,
570 const float *X,
571 int M, int N, int K);
572
573void gemm_nt_q4_k(const float *A,
574 const void *B,
575 const float *bias,
576 float *C,
577 int M, int N, int K);
578
579void dequant_q4_k_row(const void *src, float *dst, size_t n_elements);
580
581void gemv_q6_k(float *y,
582 const void *W,
583 const float *x,
584 int M, int K);
585
586void gemm_q6_k(float *Y,
587 const void *W,
588 const float *X,
589 int M, int N, int K);
590
591void gemm_nt_q6_k(const float *A,
592 const void *B,
593 const float *bias,
594 float *C,
595 int M, int N, int K);
596
597// Simple quant GEMM (Q4_0, Q4_1, Q5_0, Q5_1, Q8_0)
598void gemm_nt_q4_0(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
599void gemm_nt_q4_1(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
600void gemm_nt_q5_0(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
601void gemm_nt_q5_1(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
602void gemm_nt_q5_1_q8_1(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
603/* Exactly two rows; each packed Q5_1 weight block is reused for both rows. */
604void gemm_nt_q5_1_q8_1_m2(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
605void gemm_nt_q5_1_q8_1_ref(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K);
606void gemm_nt_q5_k(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
607void gemm_nt_q5_k_q8_k(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K);
608void gemm_nt_q5_k_q8_k_m2(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
609void gemm_nt_q8_0(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
610void gemm_nt_q8_0_q8_0_contract(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
611void gemm_nt_q8_0_q8_0_contract_m2(const float *A, const void *B, const float *bias, float *C, int M, int N, int K);
612
613// GEMV versions (for decode mode - single token)
614void gemv_q4_0(float *y, const void *W, const float *x, int M, int K);
615void gemv_q5_0(float *y, const void *W, const float *x, int M, int K);
616void gemv_q5_1(float *y, const void *W, const float *x, int M, int K);
617void gemv_q5_1_q8_1(float *y, const void *W, const float *x, int M, int K);
618void gemv_q5_1_q8_1_ref(float *y, const void *W, const void *x_q8, int M, int K);
619void gemv_q5_k(float *y, const void *W, const float *x, int M, int K);
620void gemv_q5_k_q8_k(float *y, const void *W, const void *x_q8, int M, int K);
621void gemv_q8_0(float *y, const void *W, const float *x, int M, int K);
622void gemv_q8_0_q8_0_contract(float *y, const void *W, const float *x, int M, int K);
623void gemv_q8_0_q8_0_x4(float *y, const void *W, const void *x_q8, int M, int K);
624
625/* Parallel Q5_0 versions - caller provides ith/nth from OpenMP region */
626void gemv_q5_0_parallel(float *y, const void *W, const float *x,
627 int M, int K, int ith, int nth);
628void gemv_q5_0_parallel_simd(float *y, const void *W, const float *x,
629 int M, int K, int ith, int nth);
630
631void dequant_q6_k_row(const void *src, float *dst, size_t n_elements);
632
633// Simple quant dequantization (Q4_0, Q4_1, Q5_0, Q5_1, Q8_0)
634void dequant_q4_0_row(const void *src, float *dst, size_t n_elements);
635void dequant_q4_1_row(const void *src, float *dst, size_t n_elements);
636void dequant_q5_0_row(const void *src, float *dst, size_t n_elements);
637void dequant_q5_1_row(const void *src, float *dst, size_t n_elements);
638void dequant_q8_0_row(const void *src, float *dst, size_t n_elements);
639
640// ============================================================================
641// INT8 ACTIVATION KERNELS
642// ============================================================================
643
644// Q8_0 quantization (32 elements per block, 34 bytes: 2-byte scale + 32 int8)
645void quantize_row_q8_0(const float *x, void *y, int k);
646
647// Batch Q8_0 quantization (row-major output for GEMM compatibility)
648// Output: each row at offset row * ((k/32) * 34) bytes
649void quantize_batch_q8_0(const float *x, void *y, int num_rows, int k);
650
651// Q5_0 weights x Q8_0 activations
652void gemv_q5_0_q8_0(float *y, const void *W, const void *x_q8, int M, int K);
653
654// Q8_0 weights x Q8_0 activations
655void gemv_q8_0_q8_0(float *y, const void *W, const void *x_q8, int M, int K);
656
657// Fused GEMV: quantize(FP32->Q8_0) + GEMV(Q5_0 weights) + bias add
658void gemv_fused_q5_0_bias_dispatch(float *y, const void *W, const float *x,
659 const float *bias, int M, int K);
660
661// Fused GEMV: quantize(FP32->Q8_0) + GEMV(Q8_0 weights) + bias add
662void gemv_fused_q8_0_bias_dispatch(float *y, const void *W, const float *x,
663 const float *bias, int M, int K);
664
665// Q8_K quantization (256 elements per block, super-block format)
666void quantize_row_q8_k(const float *x, void *y, int k);
667
668// Batch Q8_K quantization (row-major output for GEMM compatibility)
669void quantize_batch_q8_k(const float *x, void *y, int num_rows, int k);
670void quantize_batch_q8_k_4row_nearest_even(const float *x, void *y,
671 int num_rows, int k);
672
673void gemv_q4_k_q8_k(float *y,
674 const void *W,
675 const void *x_q8,
676 int M, int K);
677
678/* Reference implementation (no SIMD) - for testing/comparison */
679void gemv_q4_k_q8_k_ref(float *y,
680 const void *W,
681 const void *x_q8,
682 int M, int K);
683
684/* Parallel version: receives ith (thread index) and nth (total threads).
685 * OpenMP is at orchestration level, kernel processes only rows [r0, r1). */
686void gemv_q4_k_q8_k_parallel(float *y,
687 const void *W,
688 const void *x_q8,
689 int M, int K,
690 int ith, int nth);
691
692/* Parallel SIMD version: combines AVX with parallel row splitting.
693 * Includes row-ahead prefetching to hide memory latency (~50-70ns).
694 * This is the fastest option for multi-threaded decode. */
695void gemv_q4_k_q8_k_parallel_simd(float *y,
696 const void *W,
697 const void *x_q8,
698 int M, int K,
699 int ith, int nth);
700
701void gemm_q4_k_q8_k(float *Y,
702 const void *W,
703 const void *X_q8,
704 int M, int N, int K);
705
706void gemm_nt_q4_k_q8_k(const void *A_q8,
707 const void *B,
708 const float *bias,
709 float *C,
710 int M, int N, int K);
711
712// Q6_K x Q8_K quantized kernels
713void vec_dot_q6_k_q8_k(int n, float *s, const void *vx, const void *vy);
714
715void gemv_q6_k_q8_k(float *y,
716 const void *W,
717 const void *x_q8,
718 int M, int K);
719const char *ck_q6_k_q8_k_provider_name(void);
720const char *ck_q6_k_prepared_provider_name(void);
721
722/* Parallel Q6_K versions - caller provides ith/nth from OpenMP region */
723void gemv_q6_k_q8_k_parallel(float *y, const void *W, const void *x_q8,
724 int M, int K, int ith, int nth);
725void gemv_q6_k_q8_k_parallel_simd(float *y, const void *W, const void *x_q8,
726 int M, int K, int ith, int nth);
727
728void gemm_q6_k_q8_k(float *Y,
729 const void *W,
730 const void *X_q8,
731 int M, int N, int K);
732
733void gemm_nt_q6_k_q8_k(const void *A_q8,
734 const void *B,
735 const float *bias,
736 float *C,
737 int M, int N, int K);
738
739void gemm_nt_q6_k_q8_k_tile(const void *A_q8,
740 const void *B,
741 const float *bias,
742 float *C,
743 int M, int N, int K,
744 int m0, int m1,
745 int n0, int n1);
746
747void gemm_nt_q6_k_q8_k_m4_tile(const void *A_q8,
748 const void *B,
749 const float *bias,
750 float *C,
751 int M, int N, int K,
752 int m0, int m1,
753 int n0, int n1);
754
755void gemm_nt_q6_k_q8_k_m2(const void *A_q8, const void *B,
756 const float *bias, float *C, int M, int N, int K);
757
758void gemm_nt_q6_k_q8_k_tiled(const void *A_q8,
759 const void *B,
760 const float *bias,
761 float *C,
762 int M, int N, int K);
763
764void gemm_nt_q8_0_q8_0(const void *A_q8,
765 const void *B,
766 const float *bias,
767 float *C,
768 int M, int N, int K);
769
770void gemm_nt_q8_0_q8_0_m2n4(const void *A_q8,
771 const void *B,
772 const float *bias,
773 float *C,
774 int M, int N, int K);
775void gemm_nt_q8_0_q8_0_m2n4_tile(const void *A_q8,
776 const void *B,
777 const float *bias,
778 float *C,
779 int M, int N, int K, int ldc);
780
781// GEMM_NN: C[M,N] = A[M,K] @ B[K,N] + bias[N]
782// B is stored row-major as [K,N] (no transpose on B)
783// Used for backward d_input = d_output @ W
784void gemm_nn_parallel(const float *A,
785 const float *B,
786 const float *bias,
787 float *C,
788 int M, int N, int K);
789
790void gemm_nn_avx512(const float *A,
791 const float *B,
792 const float *bias,
793 float *C,
794 int M, int N, int K);
795
796/* ISA-neutral wrapper (dispatches to the best compiled SIMD path). */
797void gemm_nn_simd(const float *A,
798 const float *B,
799 const float *bias,
800 float *C,
801 int M, int N, int K);
802
803void gemm_nn_blocked(const float *A,
804 const float *B,
805 const float *bias,
806 float *C,
807 int M, int N, int K);
808
809// Head-major output projection (reads attention output directly, no flatten)
810// Reads attn_out [num_heads, tokens, head_dim] with strided access
811void ck_gemm_nt_head_major_q5_0(const float *attn_out,
812 const void *wo,
813 const float *bias,
814 float *output,
815 int tokens,
816 int embed_dim,
817 int num_heads,
818 int head_dim);
819
820void ck_gemm_nt_head_major_q8_0(const float *attn_out,
821 const void *wo,
822 const float *bias,
823 float *output,
824 int tokens,
825 int embed_dim,
826 int num_heads,
827 int head_dim);
828
829/* Exact physical-layout providers. Logical values and reduction order are unchanged. */
830void ck_layout_token_to_head_f32(const float *src, float *dst,
831 int tokens, int heads, int head_dim);
832void ck_layout_head_to_token_f32(const float *src, float *dst,
833 int heads, int tokens, int head_dim);
834
835// GEMM_TN: C[M,N] = A[K,M].T @ B[K,N] + bias[N]
836// A is stored row-major as [K,M], B is stored row-major as [K,N]
837// Used for backward d_W = d_output.T @ input
838void gemm_tn_parallel(const float *A,
839 const float *B,
840 const float *bias,
841 float *C,
842 int M, int N, int K);
843
844void gemm_tn_avx512(const float *A,
845 const float *B,
846 const float *bias,
847 float *C,
848 int M, int N, int K);
849
850void gemm_tn_blocked(const float *A,
851 const float *B,
852 const float *bias,
853 float *C,
854 int M, int N, int K);
855
856// Fused GEMM operations (GEMM + bias + activation in one pass)
857void gemm_bias_relu_fused(const float *A,
858 const float *B,
859 const float *bias,
860 float *C,
861 int M, int N, int K);
862
863void gemm_bias_gelu_fused(const float *A,
864 const float *B,
865 const float *bias,
866 float *C,
867 int M, int N, int K);
868
869void gemm_bias_silu_fused(const float *A,
870 const float *B,
871 const float *bias,
872 float *C,
873 int M, int N, int K);
874
875// Fused GEMM + SwiGLU (LLaMA/SmolLM MLP gate+up projection)
876// Computes: output = SiLU(x @ W_gate + b_gate) * (x @ W_up + b_up)
877// Two GEMMs + SwiGLU fused into one pass - intermediates stay in registers
878void gemm_swiglu_fused(const float *x,
879 const float *W_gate,
880 const float *W_up,
881 const float *b_gate, // can be NULL
882 const float *b_up, // can be NULL
883 float *output,
884 int M, int N, int K);
885
886// =============================================================================
887// Fully Fused MLP Decode Kernels (T=1 token generation)
888// =============================================================================
889// These kernels fuse the ENTIRE MLP block: Gate + Up + SwiGLU + Down
890// Key benefit: Intermediate swiglu values stay in L1/L2, never touch DRAM
891// Target: AVX-512 / Intel Xeon 5th Gen (Sapphire/Emerald Rapids)
892
893// Version 1: Tiled fusion with thread-local accumulators
894// Best for: Small number of cores, when critical section overhead is low
895void fused_mlp_swiglu_decode(
896 const float *x, // [D] input
897 const float *W_gate, // [Hff, D] gate projection
898 const float *W_up, // [Hff, D] up projection
899 const float *W_down, // [D, Hff] down projection
900 const float *b_gate, // [Hff] or NULL
901 const float *b_up, // [Hff] or NULL
902 const float *b_down, // [D] or NULL
903 float *output, // [D] output
904 int D, // hidden dimension
905 int Hff); // intermediate dimension
906
907// Version 2: Two-phase (swiglu then down projection)
908// Best for: Many cores (24+), avoids critical section, better parallelism
909void fused_mlp_swiglu_decode_v2(
910 const float *x,
911 const float *W_gate,
912 const float *W_up,
913 const float *W_down,
914 const float *b_gate,
915 const float *b_up,
916 const float *b_down,
917 float *output,
918 int D,
919 int Hff);
920
921// Version 3: Tiled with atomic accumulation
922// Best for: Large L2 cache (2MB+), good cache reuse
923void fused_mlp_swiglu_decode_tiled(
924 const float *x,
925 const float *W_gate,
926 const float *W_up,
927 const float *W_down,
928 const float *b_gate,
929 const float *b_up,
930 const float *b_down,
931 float *output,
932 int D,
933 int Hff);
934
935/* ============================================================================
936 * PREFILL FUSION KERNELS
937 * ============================================================================
938 * These kernels fuse operations for prefill (large batch/sequence) to avoid
939 * writing intermediate activations to DRAM. Fusion helps when activations
940 * exceed L3 cache size.
941 *
942 * For decode (single token), use the non-fused kernels as activations
943 * easily fit in L2 cache anyway.
944 */
945
946/**
947 * @brief Fused RMSNorm + QKV projection for prefill
948 *
949 * Tiles along token dimension to keep intermediate x_norm in L2 cache.
950 * Avoids ~7MB DRAM traffic per layer for seq_len=1024, hidden=896.
951 *
952 * @param scratch Temporary buffer from fused_rmsnorm_qkv_scratch_size()
953 */
954void fused_rmsnorm_qkv_prefill(
955 const float *x, /* [seq_len × hidden] input */
956 const float *gamma, /* [hidden] RMSNorm weights */
957 const float *Wq, /* [q_dim × hidden] Q weights (transposed) */
958 const float *Wk, /* [kv_dim × hidden] K weights (transposed) */
959 const float *Wv, /* [kv_dim × hidden] V weights (transposed) */
960 float *Q, /* [seq_len × q_dim] output */
961 float *K, /* [seq_len × kv_dim] output */
962 float *V, /* [seq_len × kv_dim] output */
963 int seq_len,
964 int hidden,
965 int q_dim,
966 int kv_dim,
967 float eps,
968 float *scratch);
969
970/**
971 * @brief Fused RMSNorm + QKV projection for prefill (head-major outputs)
972 *
973 * Writes Q as [num_heads, seq_len, aligned_head_dim] and K/V with stride
974 * kv_stride_tokens for KV-cache compatibility.
975 */
976void fused_rmsnorm_qkv_prefill_head_major(
977 const float *x,
978 const float *gamma,
979 const float *Wq, const float *Bq,
980 const float *Wk, const float *Bk,
981 const float *Wv, const float *Bv,
982 float *Q,
983 float *K,
984 float *V,
985 int seq_len,
986 int embed_dim,
987 int aligned_embed_dim,
988 int num_heads,
989 int num_kv_heads,
990 int head_dim,
991 int aligned_head_dim,
992 int kv_stride_tokens,
993 float eps,
994 float *scratch);
995
996/**
997 * @brief Fused RMSNorm + QKV projection for prefill (head-major, Q8 activations)
998 *
999 * Supports Q5_0 or Q8_0 weights with Q8_0 activations.
1000 */
1001void fused_rmsnorm_qkv_prefill_head_major_quant(
1002 const float *x,
1003 const float *gamma,
1004 const void *Wq, const float *Bq, CKDataType wq_dt,
1005 const void *Wk, const float *Bk, CKDataType wk_dt,
1006 const void *Wv, const float *Bv, CKDataType wv_dt,
1007 float *Q,
1008 float *K,
1009 float *V,
1010 int seq_len,
1011 int embed_dim,
1012 int aligned_embed_dim,
1013 int num_heads,
1014 int num_kv_heads,
1015 int head_dim,
1016 int aligned_head_dim,
1017 int kv_stride_tokens,
1018 float eps,
1019 void *scratch);
1020
1021/** @brief Unfused version for benchmarking comparison */
1022void unfused_rmsnorm_qkv_prefill(
1023 const float *x,
1024 const float *gamma,
1025 const float *Wq,
1026 const float *Wk,
1027 const float *Wv,
1028 float *x_norm, /* [seq_len × hidden] intermediate buffer */
1029 float *Q,
1030 float *K,
1031 float *V,
1032 int seq_len,
1033 int hidden,
1034 int q_dim,
1035 int kv_dim,
1036 float eps);
1037
1038/** @brief Get scratch buffer size for fused_rmsnorm_qkv_prefill */
1039size_t fused_rmsnorm_qkv_scratch_size(int hidden);
1040
1041/** @brief Get scratch buffer size for fused_rmsnorm_qkv_prefill_head_major_quant */
1042size_t fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(int aligned_embed_dim);
1043
1044/**
1045 * @brief Fused MLP (Gate + Up + SwiGLU + Down) for prefill
1046 *
1047 * Tiles along token dimension to keep gate/up/hidden in L3 cache.
1048 *
1049 * @param scratch Temporary buffer from fused_mlp_swiglu_scratch_size()
1050 */
1051void fused_mlp_swiglu_prefill(
1052 const float *x, /* [seq_len × hidden] input */
1053 const float *W_gate, /* [intermediate × hidden] (transposed) */
1054 const float *W_up, /* [intermediate × hidden] (transposed) */
1055 const float *W_down, /* [hidden × intermediate] (transposed) */
1056 float *output, /* [seq_len × hidden] output */
1057 int seq_len,
1058 int hidden,
1059 int intermediate,
1060 float *scratch);
1061
1062/**
1063 * @brief Fused MLP (Gate + Up + SwiGLU + Down) for prefill with biases
1064 */
1065void fused_mlp_swiglu_prefill_bias(
1066 const float *x,
1067 const float *W_gate,
1068 const float *W_up,
1069 const float *W_down,
1070 const float *B_gate,
1071 const float *B_up,
1072 const float *B_down,
1073 float *output,
1074 int seq_len,
1075 int hidden,
1076 int intermediate,
1077 float *scratch);
1078
1079/** @brief Get scratch buffer size for fused_mlp_swiglu_prefill */
1080size_t fused_mlp_swiglu_scratch_size(int intermediate);
1081
1082/**
1083 * @brief Quantized fused MLP for prefill (W1=gate+up, W2=down)
1084 *
1085 * W1 uses Q8_0 activations (Q5_0/Q8_0 weights), W2 uses Q8_K activations
1086 * (Q4_K/Q6_K weights).
1087 */
1088void fused_mlp_swiglu_prefill_w1w2_quant(
1089 const float *x,
1090 const void *W1,
1091 const float *B1,
1092 CKDataType w1_dt,
1093 const void *W2,
1094 const float *B2,
1095 CKDataType w2_dt,
1096 float *output,
1097 int seq_len,
1098 int embed_dim,
1099 int aligned_embed_dim,
1100 int intermediate_dim,
1101 int aligned_intermediate_dim,
1102 void *scratch);
1103
1104/** @brief Get scratch buffer size for fused_mlp_swiglu_prefill_w1w2_quant */
1105size_t fused_mlp_swiglu_prefill_w1w2_quant_scratch_size(int aligned_embed_dim,
1106 int aligned_intermediate_dim);
1107
1108// High-performance GEMM microkernel with 8x8 register blocking
1109// Inspired by oneDNN/BLIS - keeps all 64 accumulator values in registers
1110// C[M,N] = A[M,K] @ B[K,N] or C[M,N] = A[M,K] @ B[N,K].T
1111// B_transposed: 0 = B is [K,N], 1 = B is [N,K] (transposed, common in NN weights)
1112void gemm_microkernel(const float *A,
1113 const float *B,
1114 float *C,
1115 int M, int N, int K,
1116 int B_transposed);
1117
1118// Cache-blocked GEMM using 8x8 microkernels (B not transposed)
1119void gemm_microkernel_blocked(const float *A,
1120 const float *B,
1121 float *C,
1122 int M, int N, int K);
1123
1124// Cache-blocked GEMM for B transposed (common in NN)
1125void gemm_microkernel_blocked_bt(const float *A,
1126 const float *B,
1127 float *C,
1128 int M, int N, int K);
1129
1130// Optimized GEMM with matrix packing (best for large matrices)
1131// Packs A and B into contiguous layouts for optimal cache access
1132void gemm_microkernel_packed(const float *A,
1133 const float *B,
1134 float *C,
1135 int M, int N, int K);
1136
1137// LayerNorm forward kernels, copied from C-Transformer.
1138void layernorm_naive_serial(const float *input,
1139 const float *gamma,
1140 const float *beta,
1141 float *output,
1142 float *mean_cache,
1143 float *rstd_cache,
1144 int tokens, int d_model, int aligned_embed_dim,
1145 float eps);
1146
1147void layernorm_forward_rolled_slice(const float *__restrict input_slice_base,
1148 const float *__restrict gamma,
1149 const float *__restrict beta,
1150 float *__restrict output_slice_base,
1151 float *__restrict mean_cache_slice,
1152 float *__restrict rstd_cache_slice,
1153 int num_tokens_in_slice,
1154 int d_model,
1155 int aligned_embed_dim,
1156 float eps);
1157
1158/* BF16 LayerNorm forward (rolled) - caller provides scratch buffers */
1159void layernorm_forward_rolled_slice_bf16(const uint16_t *__restrict input_slice_base,
1160 const float *__restrict gamma,
1161 const float *__restrict beta,
1162 uint16_t *__restrict output_slice_base,
1163 float *__restrict mean_cache_slice,
1164 float *__restrict rstd_cache_slice,
1165 int num_tokens_in_slice,
1166 int d_model,
1167 int aligned_embed_dim,
1168 float eps,
1169 float *scratch_input, /* [num_tokens * aligned_embed_dim] */
1170 float *scratch_output); /* [num_tokens * aligned_embed_dim] */
1171
1172void layernorm_forward_unrolled_slice(const float *__restrict input_slice_base,
1173 const float *__restrict gamma,
1174 const float *__restrict beta,
1175 float *__restrict output_slice_base,
1176 float *__restrict mean_cache_slice,
1177 float *__restrict rstd_cache_slice,
1178 int num_tokens_in_slice,
1179 int d_model,
1180 float eps);
1181
1182/* BF16 LayerNorm forward (unrolled) - caller provides scratch buffers */
1183void layernorm_forward_unrolled_slice_bf16(const uint16_t *__restrict input_slice_base,
1184 const float *__restrict gamma,
1185 const float *__restrict beta,
1186 uint16_t *__restrict output_slice_base,
1187 float *__restrict mean_cache_slice,
1188 float *__restrict rstd_cache_slice,
1189 int num_tokens_in_slice,
1190 int d_model,
1191 float eps,
1192 float *scratch_input, /* [num_tokens * d_model] */
1193 float *scratch_output); /* [num_tokens * d_model] */
1194
1195void layernorm_naive_serial_matched_precision(const float *input,
1196 const float *gamma,
1197 const float *beta,
1198 float *output,
1199 float *mean_cache,
1200 float *rstd_cache,
1201 int tokens, int d_model, float eps);
1202void layernorm_naive_serial_bf16_storage(const float *input,
1203 const float *gamma,
1204 const float *beta,
1205 float *output,
1206 float *mean_cache,
1207 float *rstd_cache,
1208 int tokens, int d_model, float eps);
1209void layernorm_pytorch_welford_bf16_storage(const float *input,
1210 const float *gamma,
1211 const float *beta,
1212 float *output,
1213 float *mean_cache,
1214 float *rstd_cache,
1215 int tokens, int d_model, float eps);
1216
1217void layernorm_backward_kernel(const float *d_output,
1218 const float *input,
1219 const float *gamma,
1220 const float *mean,
1221 const float *rstd,
1222 float *d_input,
1223 float *d_gamma,
1224 float *d_beta,
1225 int tokens, int d_model, int aligned_embed_dim);
1226
1227/* BF16 LayerNorm backward - caller provides scratch buffers */
1228void layernorm_backward_kernel_bf16(const uint16_t *d_output,
1229 const uint16_t *input,
1230 const float *gamma,
1231 const float *mean,
1232 const float *rstd,
1233 uint16_t *d_input,
1234 float *d_gamma,
1235 float *d_beta,
1236 int tokens, int d_model, int aligned_embed_dim,
1237 float *scratch_d_output, /* [tokens * aligned_embed_dim] */
1238 float *scratch_input, /* [tokens * aligned_embed_dim] */
1239 float *scratch_d_input); /* [tokens * aligned_embed_dim] */
1240
1241void feature_slice_copy(const float *src,
1242 float *dst,
1243 int rows,
1244 int src_dim,
1245 int dst_dim,
1246 int dst_feature_offset);
1247
1248void feature_concat(const float *main_input,
1249 const float *branch_input,
1250 float *output,
1251 int rows,
1252 int main_dim,
1253 int branch_slice_dim,
1254 int num_branch_slices);
1255
1256// RMSNorm forward/backward kernels.
1257void rmsnorm_forward(const float *input,
1258 const float *gamma,
1259 float *output,
1260 float *rstd_cache,
1261 int tokens,
1262 int d_model,
1263 int aligned_embed_dim,
1264 float eps);
1265void rmsnorm_forward_parallel_dispatch(
1266 const float *input,
1267 const float *gamma,
1268 float *output,
1269 float *rstd_cache,
1270 int tokens,
1271 int d_model,
1272 int aligned_embed_dim,
1273 float eps);
1274void rmsnorm_forward_fp64_sum(const float *input,
1275 const float *gamma,
1276 float *output,
1277 float *rstd_cache,
1278 int tokens,
1279 int d_model,
1280 int aligned_embed_dim,
1281 float eps);
1282void rmsnorm_forward_llama_production(const float *input,
1283 const float *gamma,
1284 float *output,
1285 float *rstd_cache,
1286 int tokens,
1287 int d_model,
1288 int aligned_embed_dim,
1289 float eps);
1290void rmsnorm_forward_pytorch_bf16_storage(const float *input,
1291 const float *gamma,
1292 float *output,
1293 float *rstd_cache,
1294 int tokens,
1295 int d_model,
1296 int aligned_embed_dim,
1297 float eps);
1298void rmsnorm_forward_strided_pytorch_bf16_storage(const float *input,
1299 const float *gamma,
1300 float *output,
1301 float *rstd_cache,
1302 int tokens,
1303 int d_model,
1304 int input_stride,
1305 int output_stride,
1306 float eps);
1307void rmsnorm_forward_qwen3next_pytorch_bf16_storage(const float *input,
1308 const float *gamma,
1309 float *output,
1310 float *rstd_cache,
1311 int tokens,
1312 int d_model,
1313 int aligned_embed_dim,
1314 float eps);
1315void rmsnorm_forward_muse_centered_pytorch_bf16_storage(
1316 const float *input,
1317 const float *zero_centered_weight,
1318 float *output,
1319 float *rstd_cache,
1320 int tokens,
1321 int d_model,
1322 int aligned_embed_dim,
1323 float eps);
1324void rmsnorm_forward_muse_weighted_pytorch_bf16_storage(
1325 const float *input,
1326 const float *weight,
1327 float *output,
1328 float *rstd_cache,
1329 int tokens,
1330 int d_model,
1331 int aligned_embed_dim,
1332 float eps);
1333void rmsnorm_forward_muse_unweighted_pytorch_bf16_storage(
1334 const float *input,
1335 float *output,
1336 float *rstd_cache,
1337 int tokens,
1338 int d_model,
1339 int aligned_embed_dim,
1340 float eps);
1341void rmsnorm_forward_fp32_square_fp64_sum(const float *input,
1342 const float *gamma,
1343 float *output,
1344 float *rstd_cache,
1345 int tokens,
1346 int d_model,
1347 int aligned_embed_dim,
1348 float eps);
1349void rmsnorm_forward_strided_f32(const float *input,
1350 const float *gamma,
1351 float *output,
1352 float *rstd_cache,
1353 int tokens,
1354 int d_model,
1355 int input_stride,
1356 int output_stride,
1357 float eps);
1358void rmsnorm_forward_kv_lora(const float *input,
1359 const float *gamma,
1360 float *output,
1361 float *rstd_cache,
1362 int tokens,
1363 int d_model,
1364 int aligned_embed_dim,
1365 float eps);
1366void rmsnorm_forward_no_weight(const float *input,
1367 float *output,
1368 float *rstd_cache,
1369 int tokens,
1370 int d_model,
1371 int aligned_embed_dim,
1372 float eps);
1373void rmsnorm_forward_no_weight_llama_production(const float *input,
1374 float *output,
1375 float *rstd_cache,
1376 int tokens,
1377 int d_model,
1378 int aligned_embed_dim,
1379 float eps);
1380void gemma4_v_norm_forward(const float *input,
1381 float *output,
1382 float *rstd_cache,
1383 int tokens,
1384 int num_kv_heads,
1385 int head_dim,
1386 float eps);
1387void gemma4_v_norm_forward_parallel_dispatch(const float *input,
1388 float *output,
1389 float *rstd_cache,
1390 int tokens,
1391 int num_kv_heads,
1392 int head_dim,
1393 float eps);
1394
1395void rmsnorm_backward(const float *d_output,
1396 const float *input,
1397 const float *gamma,
1398 const float *rstd_cache,
1399 float *d_input,
1400 float *d_gamma,
1401 int tokens,
1402 int d_model,
1403 int aligned_embed_dim);
1404
1405/* QK norm kernels (Qwen-style per-head RMSNorm before RoPE). */
1406void qk_norm_forward(float *q,
1407 float *k,
1408 const float *q_gamma,
1409 const float *k_gamma,
1410 int num_heads,
1411 int num_kv_heads,
1412 int num_tokens,
1413 int head_dim,
1414 float eps);
1415void qk_norm_forward_parallel_dispatch(float *q,
1416 float *k,
1417 const float *q_gamma,
1418 const float *k_gamma,
1419 int num_heads,
1420 int num_kv_heads,
1421 int num_tokens,
1422 int head_dim,
1423 float eps);
1424void qk_norm_forward_llama_production_parallel_dispatch(
1425 float *q,
1426 float *k,
1427 const float *q_gamma,
1428 const float *k_gamma,
1429 int num_heads,
1430 int num_kv_heads,
1431 int num_tokens,
1432 int head_dim,
1433 float eps);
1434void qk_norm_forward_fp64_sum(float *q,
1435 float *k,
1436 const float *q_gamma,
1437 const float *k_gamma,
1438 int num_heads,
1439 int num_kv_heads,
1440 int num_tokens,
1441 int head_dim,
1442 float eps);
1443void qk_norm_forward_llama_production(float *q,
1444 float *k,
1445 const float *q_gamma,
1446 const float *k_gamma,
1447 int num_heads,
1448 int num_kv_heads,
1449 int num_tokens,
1450 int head_dim,
1451 float eps);
1452void qk_norm_forward_pytorch_bf16_storage(float *q,
1453 float *k,
1454 const float *q_gamma,
1455 const float *k_gamma,
1456 int num_heads,
1457 int num_kv_heads,
1458 int num_tokens,
1459 int head_dim,
1460 float eps);
1461void qk_norm_forward_qwen4_pytorch_bf16_storage(float *q,
1462 float *k,
1463 const float *q_gamma,
1464 const float *k_gamma,
1465 int num_heads,
1466 int num_kv_heads,
1467 int num_tokens,
1468 int head_dim,
1469 float eps);
1470void qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage(
1471 float *q,
1472 float *k,
1473 int num_heads,
1474 int num_kv_heads,
1475 int num_tokens,
1476 int head_dim,
1477 float eps,
1478 float q_scale);
1479void qk_norm_forward_prefill_exact(float *q, float *k,
1480 const float *q_gamma, const float *k_gamma,
1481 int num_heads, int num_kv_heads,
1482 int num_tokens, int head_dim, float eps);
1483void qk_norm_forward_decode_exact(float *q, float *k,
1484 const float *q_gamma, const float *k_gamma,
1485 int num_heads, int num_kv_heads,
1486 int num_tokens, int head_dim, float eps);
1487void q_norm_forward(float *q,
1488 const float *q_gamma,
1489 int num_heads,
1490 int num_tokens,
1491 int head_dim,
1492 float eps);
1493
1494
1495void gemma4_per_layer_prepare_forward(float *per_layer_input,
1496 const float *hidden,
1497 const int32_t *token_ids,
1498 const void *per_layer_token_emb,
1499 const uint16_t *per_layer_model_proj,
1500 const float *per_layer_proj_norm,
1501 int tokens,
1502 int num_layers,
1503 int embed_dim,
1504 int per_layer_dim,
1505 int vocab_size,
1506 float eps);
1507
1508void gemma4_per_layer_prepare_bf16_forward(float *per_layer_input,
1509 const float *hidden,
1510 const int32_t *token_ids,
1511 const uint16_t *per_layer_token_emb,
1512 const uint16_t *per_layer_model_proj,
1513 const float *per_layer_proj_norm,
1514 int tokens,
1515 int num_layers,
1516 int embed_dim,
1517 int per_layer_dim,
1518 int vocab_size,
1519 float eps);
1520
1521void gemma4_per_layer_embed_forward(float *hidden,
1522 const float *per_layer_input,
1523 const float *inp_gate,
1524 const float *proj,
1525 const float *post_norm,
1526 const float *out_scale,
1527 int tokens,
1528 int layer,
1529 int num_layers,
1530 int embed_dim,
1531 int per_layer_dim,
1532 float eps,
1533 const int *segment_lengths,
1534 int num_segments,
1535 int singleton_tail);
1536
1537void assistant_layer_scale_forward(float *hidden,
1538 const float *scale,
1539 int tokens,
1540 int embed_dim);
1541
1542void gemma4_final_logit_softcap_forward(float *logits,
1543 int tokens,
1544 int vocab_size,
1545 float cap);
1546void final_logit_scale_f32(float *logits,
1547 int tokens,
1548 int vocab_size,
1549 float scale);
1550void final_logit_scale_muse_pytorch_bf16_storage(float *logits,
1551 int tokens,
1552 int vocab_size,
1553 float scale);
1554void final_logit_softcap_muse_pytorch_bf16_storage(float *logits,
1555 int tokens,
1556 int vocab_size,
1557 float cap);
1558void qk_norm_backward(const float *d_q_out,
1559 const float *d_k_out,
1560 const float *q_in,
1561 const float *k_in,
1562 const float *q_gamma,
1563 const float *k_gamma,
1564 float *d_q_in,
1565 float *d_k_in,
1566 float *d_q_gamma,
1567 float *d_k_gamma,
1568 int num_heads,
1569 int num_kv_heads,
1570 int num_tokens,
1571 int head_dim,
1572 float eps);
1573
1574/* Last selected qk_norm_backward ISA: scalar=0, avx=1, avx2=2, avx_vnni=3. */
1575int qk_norm_backward_last_isa(void);
1576
1577void rmsnorm_forward_bf16(const uint16_t *input,
1578 const float *gamma,
1579 uint16_t *output,
1580 float *rstd_cache,
1581 int tokens,
1582 int d_model,
1583 int aligned_embed_dim,
1584 float eps);
1585
1586void rmsnorm_backward_bf16(const uint16_t *d_output,
1587 const uint16_t *input,
1588 const float *gamma,
1589 const float *rstd_cache,
1590 uint16_t *d_input,
1591 float *d_gamma,
1592 int tokens,
1593 int d_model,
1594 int aligned_embed_dim);
1595
1596/* INT8 RMSNorm forward - caller provides scratch buffers */
1597void rmsnorm_forward_int8(const int8_t *input,
1598 const float *gamma,
1599 int8_t *output,
1600 float *rstd_cache,
1601 int tokens,
1602 int d_model,
1603 int aligned_embed_dim,
1604 float eps,
1605 float *scratch_input, /* [tokens * aligned_embed_dim] */
1606 float *scratch_output); /* [tokens * aligned_embed_dim] */
1607
1608/* INT8 RMSNorm backward - caller provides scratch buffers */
1609void rmsnorm_backward_int8(const int8_t *d_output,
1610 const int8_t *input,
1611 const float *gamma,
1612 const float *rstd_cache,
1613 int8_t *d_input,
1614 float *d_gamma,
1615 int tokens,
1616 int d_model,
1617 int aligned_embed_dim,
1618 float *scratch_d_output, /* [tokens * aligned_embed_dim] */
1619 float *scratch_input, /* [tokens * aligned_embed_dim] */
1620 float *scratch_d_input); /* [tokens * aligned_embed_dim] */
1621
1622/* INT4 RMSNorm forward - caller provides scratch buffers */
1623void rmsnorm_forward_int4(const uint8_t *input,
1624 const float *gamma,
1625 uint8_t *output,
1626 float *rstd_cache,
1627 int tokens,
1628 int d_model,
1629 int aligned_embed_dim,
1630 float eps,
1631 float *scratch_input, /* [tokens * aligned_embed_dim] */
1632 float *scratch_output); /* [tokens * aligned_embed_dim] */
1633
1634/* INT4 RMSNorm backward - caller provides scratch buffers */
1635void rmsnorm_backward_int4(const uint8_t *d_output,
1636 const uint8_t *input,
1637 const float *gamma,
1638 const float *rstd_cache,
1639 uint8_t *d_input,
1640 float *d_gamma,
1641 int tokens,
1642 int d_model,
1643 int aligned_embed_dim,
1644 float *scratch_d_output, /* [tokens * aligned_embed_dim] */
1645 float *scratch_input, /* [tokens * aligned_embed_dim] */
1646 float *scratch_d_input); /* [tokens * aligned_embed_dim] */
1647
1648// GELU forward kernel (fast approximation), copied from C-Transformer.
1649void gelu_fast_inplace(float *data, size_t n);
1650
1651// Scalar-only exact GELU forward using standard library tanhf.
1652// Slower but provides maximum accuracy. Used by BF16 wrapper.
1653void gelu_exact_inplace(float *data, size_t n);
1654void gelu_erf_fp64_f32_inplace(float *data, size_t n);
1655void gelu_erf_fp64_f32_parallel_dispatch(float *data, size_t n);
1656void gelu_pytorch_erf_f32_inplace(float *data, size_t n);
1657
1658// GGML-compatible GELU forward matching llama.cpp's FP16 table semantics.
1659void gelu_ggml_inplace(float *data, size_t n);
1660void gelu_pytorch_tanh_bf16_storage(float *data, size_t n);
1661void gelu_erf_bf16_storage(float *data, size_t n);
1662void gelu_pytorch_erf_sleef_bf16_storage(float *data, size_t n);
1663
1664// GELU backward using tanh-based derivative (vectorized, uses fast tanh approx).
1665void gelu_backward_exact(const float *input,
1666 const float *d_output,
1667 float *d_input,
1668 size_t n);
1669
1670// Scalar-only exact GELU backward using standard library tanhf.
1671// Slower but provides maximum accuracy. Used by BF16 wrapper.
1672void gelu_backward_scalar(const float *input,
1673 const float *d_output,
1674 float *d_input,
1675 size_t n);
1676
1677void gelu_backward_fast(const float *input,
1678 const float *d_output,
1679 float *d_input,
1680 size_t n);
1681
1682// BF16 variants relying on the same floating-point logic.
1683/* BF16 GELU - caller provides scratch buffer [n] floats */
1684void gelu_fast_inplace_bf16(uint16_t *data, size_t n, float *scratch);
1685void gelu_backward_exact_bf16(const uint16_t *input,
1686 const uint16_t *d_output,
1687 uint16_t *d_input,
1688 size_t n,
1689 float *scratch_input,
1690 float *scratch_d_output,
1691 float *scratch_d_input);
1692void gelu_backward_fast_bf16(const uint16_t *input,
1693 const uint16_t *d_output,
1694 uint16_t *d_input,
1695 size_t n,
1696 float *scratch_input,
1697 float *scratch_d_output,
1698 float *scratch_d_input);
1699
1700// GeGLU: out = GELU(a) * b where x = [a, b] along last dimension
1701// Input shape: [tokens, 2 * dim], Output shape: [tokens, dim]
1702void geglu_forward_fp32(const float *x, float *out, int tokens, int dim);
1703void geglu_forward_exact(const float *x, float *out, int tokens, int dim);
1704void geglu_forward_quick(const float *x, float *out, int tokens, int dim);
1705void geglu_forward_ggml_native(const float *x, float *out, int tokens, int dim);
1706void gelu_ggml_native_inplace(float *data, size_t n);
1707void geglu_forward_bf16(const uint16_t *x, uint16_t *out, int tokens, int dim, float *scratch);
1708void geglu_backward_fp32(const float *x,
1709 const float *d_out,
1710 float *d_x,
1711 int tokens,
1712 int dim);
1713void geglu_backward_bf16_mixed(const uint16_t *x,
1714 const uint16_t *d_out,
1715 float *d_x,
1716 int tokens,
1717 int dim);
1718
1719 // ReLU kernels.
1720 void relu_forward(const float *input, float *output, size_t n);
1721 void relu_forward_inplace(float *data, size_t n);
1722 void relu_backward(const float *input,
1723 const float *d_output,
1724 float *d_input,
1725 size_t n);
1726 void relu2_forward(const float *input, float *output, size_t n);
1727 void relu2_backward(const float *input,
1728 const float *d_output,
1729 float *d_input,
1730 size_t n);
1731
1732 void relu_forward_bf16(const uint16_t *input, uint16_t *output, size_t n);
1733 void relu_forward_inplace_bf16(uint16_t *data, size_t n);
1734 void relu_backward_bf16(const uint16_t *input,
1735 const uint16_t *d_output,
1736 uint16_t *d_input,
1737 size_t n);
1738
1739 // Causal softmax kernel on head-major attention scores, copied from C-Transformer.
1740 void causal_softmax_head_major(float *scores,
1741 int num_heads,
1742 int num_tokens,
1743 int aligned_context_window);
1744
1745 // Scalar-only exact causal softmax using standard library expf.
1746 // Slower but provides maximum accuracy. Used by BF16 attention wrapper.
1747 void causal_softmax_head_major_exact(float *scores,
1748 int num_heads,
1749 int num_tokens,
1750 int aligned_context_window);
1751
1752 void backward_causal_softmax_head_major(float *d_scores,
1753 const float *weights,
1754 int num_heads,
1755 int num_tokens,
1756 int aligned_context_window);
1757
1758 /* BF16 causal softmax - caller provides scratch buffer */
1759 void causal_softmax_head_major_bf16(uint16_t *scores,
1760 int num_heads,
1761 int num_tokens,
1762 int aligned_context_window,
1763 float *scratch); /* [num_heads * aligned_context_window * aligned_context_window] */
1764
1765 /* BF16 backward causal softmax - caller provides scratch buffers */
1766 void backward_causal_softmax_head_major_bf16(uint16_t *d_scores,
1767 const uint16_t *weights,
1768 int num_heads,
1769 int num_tokens,
1770 int aligned_context_window,
1771 float *scratch_d_scores, /* [num_heads * aligned_context_window * aligned_context_window] */
1772 float *scratch_weights); /* [num_heads * aligned_context_window * aligned_context_window] */
1773
1774// Scaled dot-product attention (causal) in head-major layout.
1775// Q/K/V layout: [head][token][head_dim] with stride aligned_head_dim.
1776// scores: [head][query_token][key_token] with stride aligned_context_window.
1777// output: same layout as Q/V.
1778void attention_forward_causal_head_major(const float *q,
1779 const float *k,
1780 const float *v,
1781 float *scores,
1782 float *output,
1783 int num_heads,
1784 int num_tokens,
1785 int head_dim,
1786 int aligned_head_dim,
1787 int aligned_context_window);
1788
1789// Exact version using standard library expf (slower but accurate).
1790void attention_forward_causal_head_major_exact(const float *q,
1791 const float *k,
1792 const float *v,
1793 float *scores,
1794 float *output,
1795 int num_heads,
1796 int num_tokens,
1797 int head_dim,
1798 int aligned_head_dim,
1799 int aligned_context_window);
1800
1801// GQA-aware attention: Q has num_heads, K/V have num_kv_heads.
1802void attention_forward_causal_head_major_gqa(const float *q,
1803 const float *k,
1804 const float *v,
1805 float *scores,
1806 float *output,
1807 int num_heads,
1808 int num_kv_heads,
1809 int num_tokens,
1810 int head_dim,
1811 int aligned_head_dim,
1812 int aligned_context_window);
1813
1814// Exact GQA version using standard library expf (slower but accurate).
1815void attention_forward_causal_head_major_gqa_exact(const float *q,
1816 const float *k,
1817 const float *v,
1818 float *scores,
1819 float *output,
1820 int num_heads,
1821 int num_kv_heads,
1822 int num_tokens,
1823 int head_dim,
1824 int aligned_head_dim,
1825 int aligned_context_window);
1826
1827/* BF16 attention forward - caller provides scratch buffers (no internal malloc) */
1828void attention_forward_causal_head_major_gqa_bf16(const uint16_t *q,
1829 const uint16_t *k,
1830 const uint16_t *v,
1831 float *scores,
1832 float *output,
1833 int num_heads,
1834 int num_kv_heads,
1835 int num_tokens,
1836 int head_dim,
1837 int aligned_head_dim,
1838 int aligned_context_window,
1839 float *scratch_q, /* [num_heads * num_tokens * aligned_head_dim] */
1840 float *scratch_k, /* [num_kv_heads * num_tokens * aligned_head_dim] */
1841 float *scratch_v); /* [num_kv_heads * num_tokens * aligned_head_dim] */
1842
1843// Flash-style causal attention forward (no score/weight matrix materialization).
1844// Head-major layout:
1845// Q: [num_heads, num_tokens, aligned_head_dim]
1846// K/V: [num_kv_heads, num_tokens, aligned_head_dim]
1847// out: [num_heads, num_tokens, aligned_head_dim]
1848void attention_forward_causal_head_major_gqa_flash(const float *q,
1849 const float *k,
1850 const float *v,
1851 float *output,
1852 int num_heads,
1853 int num_kv_heads,
1854 int num_tokens,
1855 int head_dim,
1856 int aligned_head_dim);
1857
1858// Flash-style full / bidirectional attention for encoder-style prefill.
1859void attention_forward_full_head_major_gqa_flash(const float *q,
1860 const float *k,
1861 const float *v,
1862 float *output,
1863 int num_heads,
1864 int num_kv_heads,
1865 int num_tokens,
1866 int head_dim,
1867 int aligned_head_dim);
1868
1869
1870void deepseek_mla_attention_f32(const float *q,
1871 const float *k,
1872 const float *v,
1873 float *output,
1874 int num_heads,
1875 int num_kv_heads,
1876 int num_tokens,
1877 int qk_head_dim,
1878 int v_head_dim);
1879void deepseek_mla_attention_f32_workspace(const float *q,
1880 const float *k,
1881 const float *v,
1882 float *output,
1883 int num_heads,
1884 int num_kv_heads,
1885 int num_tokens,
1886 int qk_head_dim,
1887 int v_head_dim,
1888 float scale,
1889 float *scores,
1890 size_t scores_bytes);
1891void deepseek_mla_attention_f32_parallel_dispatch(const float *q,
1892 const float *k,
1893 const float *v,
1894 float *output,
1895 int num_heads,
1896 int num_kv_heads,
1897 int num_tokens,
1898 int qk_head_dim,
1899 int v_head_dim,
1900 float scale,
1901 float *scores,
1902 size_t scores_bytes);
1903/* Publisher eager BF16 arithmetic phases over FP32 storage buffers. */
1904void deepseek_mla_attention_bf16_phases_workspace(const float *q,
1905 const float *k,
1906 const float *v,
1907 float *output,
1908 int num_heads,
1909 int num_kv_heads,
1910 int num_tokens,
1911 int qk_head_dim,
1912 int v_head_dim,
1913 float scale,
1914 float *scores,
1915 size_t scores_bytes);
1916void deepseek_mla_attention_bf16_phases_parallel_dispatch(const float *q,
1917 const float *k,
1918 const float *v,
1919 float *output,
1920 int num_heads,
1921 int num_kv_heads,
1922 int num_tokens,
1923 int qk_head_dim,
1924 int v_head_dim,
1925 float scale,
1926 float *scores,
1927 size_t scores_bytes);
1928
1929void deepseek_mla_kv_cache_batch_store_f32(float *k_cache,
1930 float *v_cache,
1931 const float *k,
1932 const float *v,
1933 int num_tokens,
1934 int num_kv_heads,
1935 int qk_head_dim,
1936 int v_head_dim,
1937 int max_seq_len,
1938 int cache_stride);
1939
1940void deepseek_mla_kv_cache_store_f32(float *k_cache,
1941 float *v_cache,
1942 const float *k,
1943 const float *v,
1944 int pos,
1945 int num_kv_heads,
1946 int qk_head_dim,
1947 int v_head_dim,
1948 int max_seq_len,
1949 int cache_stride);
1950
1951void deepseek_mla_attention_decode_f32(const float *q,
1952 const float *k_cache,
1953 const float *v_cache,
1954 float *output,
1955 int num_heads,
1956 int num_kv_heads,
1957 int cache_len,
1958 int qk_head_dim,
1959 int v_head_dim,
1960 int max_seq_len,
1961 int cache_stride);
1962void deepseek_mla_attention_decode_f32_workspace(const float *q,
1963 const float *k_cache,
1964 const float *v_cache,
1965 float *output,
1966 int num_heads,
1967 int num_kv_heads,
1968 int cache_len,
1969 int qk_head_dim,
1970 int v_head_dim,
1971 int max_seq_len,
1972 int cache_stride,
1973 float scale,
1974 float *scores,
1975 size_t scores_bytes);
1976void deepseek_mla_attention_decode_bf16_phases_workspace(const float *q,
1977 const float *k_cache,
1978 const float *v_cache,
1979 float *output,
1980 int num_heads,
1981 int num_kv_heads,
1982 int cache_len,
1983 int qk_head_dim,
1984 int v_head_dim,
1985 int max_seq_len,
1986 int cache_stride,
1987 float scale,
1988 float *scores,
1989 size_t scores_bytes);
1990
1991void attention_forward_causal_head_major_gqa_flash_strided(const float *q,
1992 const float *k,
1993 const float *v,
1994 float *output,
1995 int num_heads,
1996 int num_kv_heads,
1997 int num_tokens,
1998 int head_dim,
1999 int aligned_head_dim,
2000 int kv_stride_tokens);
2001void attention_forward_causal_head_major_gqa_muse_eager_bf16_storage(
2002 const float *q, const float *k, const float *v, float *output,
2003 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2004 int aligned_head_dim, int kv_stride_tokens,
2005 float *scores, size_t scores_bytes,
2006 uint16_t *key_rows, size_t key_rows_bytes,
2007 uint16_t *value_columns, size_t value_columns_bytes,
2008 uint16_t *gemm_rows, size_t gemm_rows_bytes,
2009 uint16_t *gemm_scores, size_t gemm_scores_bytes);
2010void attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_sliding(
2011 const float *q, const float *k, const float *v, float *output,
2012 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2013 int aligned_head_dim, int kv_stride_tokens, int sliding_window,
2014 float *scores, size_t scores_bytes,
2015 uint16_t *key_rows, size_t key_rows_bytes,
2016 uint16_t *value_columns, size_t value_columns_bytes,
2017 uint16_t *gemm_rows, size_t gemm_rows_bytes,
2018 uint16_t *gemm_scores, size_t gemm_scores_bytes);
2019
2020void attention_forward_causal_head_major_gqa_flash_strided_token_output(
2021 const float *q,
2022 const float *k,
2023 const float *v,
2024 float *output,
2025 int num_heads,
2026 int num_kv_heads,
2027 int num_tokens,
2028 int head_dim,
2029 int aligned_head_dim,
2030 int kv_stride_tokens);
2031
2032void attention_forward_full_head_major_gqa_flash_strided(const float *q,
2033 const float *k,
2034 const float *v,
2035 float *output,
2036 int num_heads,
2037 int num_kv_heads,
2038 int num_tokens,
2039 int head_dim,
2040 int aligned_head_dim,
2041 int kv_stride_tokens);
2042
2043void attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(
2044 const float *q,
2045 const float *k,
2046 const float *v,
2047 float *output,
2048 int num_heads,
2049 int num_kv_heads,
2050 int num_tokens,
2051 int head_dim,
2052 int aligned_head_dim,
2053 int kv_stride_tokens);
2054
2055void attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided(
2056 const float *q,
2057 const float *k,
2058 const float *v,
2059 float *output,
2060 int num_heads,
2061 int num_kv_heads,
2062 int num_tokens,
2063 int head_dim,
2064 int aligned_head_dim,
2065 int kv_stride_tokens);
2066
2067void attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided(
2068 const float *q,
2069 const float *k,
2070 const float *v,
2071 float *output,
2072 int num_heads,
2073 int num_kv_heads,
2074 int num_tokens,
2075 int head_dim,
2076 int aligned_head_dim,
2077 int kv_stride_tokens);
2078
2079void attention_forward_full_head_major_gqa_flash_strided_bf16_storage(
2080 const float *q, const float *k, const float *v, float *output,
2081 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2082 int aligned_head_dim, int kv_stride_tokens);
2083
2084void attention_forward_full_head_major_gqa_sdpa_bf16_storage(
2085 const float *q, const float *k, const float *v, float *output,
2086 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2087 int aligned_head_dim, int kv_stride_tokens);
2088
2089void attention_forward_causal_head_major_gqa_flash_strided_gemma4(const float *q,
2090 const float *k,
2091 const float *v,
2092 float *output,
2093 int num_heads,
2094 int num_kv_heads,
2095 int num_tokens,
2096 int head_dim,
2097 int aligned_head_dim,
2098 int kv_stride_tokens);
2099void attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output(
2100 const float *q, const float *k, const float *v, float *output,
2101 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2102 int aligned_head_dim, int kv_stride_tokens);
2103void attention_forward_causal_head_major_shared_kv_gemma4(const float *q,
2104 float *output,
2105 int num_heads,
2106 int num_tokens,
2107 int head_dim,
2108 int aligned_head_dim,
2109 int kv_stride_tokens);
2110
2111void attention_forward_full_head_major_gqa_flash_strided_gemma4(const float *q,
2112 const float *k,
2113 const float *v,
2114 float *output,
2115 int num_heads,
2116 int num_kv_heads,
2117 int num_tokens,
2118 int head_dim,
2119 int aligned_head_dim,
2120 int kv_stride_tokens);
2121
2122void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4(const float *q,
2123 const float *k,
2124 const float *v,
2125 float *output,
2126 int num_heads,
2127 int num_kv_heads,
2128 int num_tokens,
2129 int head_dim,
2130 int aligned_head_dim,
2131 int kv_stride_tokens,
2132 int visual_start,
2133 int visual_tokens);
2134void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_token_output(
2135 const float *q, const float *k, const float *v, float *output,
2136 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2137 int aligned_head_dim, int kv_stride_tokens, int visual_start,
2138 int visual_tokens);
2139
2140// Regular exact full / bidirectional attention for encoder-style prefill.
2141// This matches the non-flash CPU reference path more closely than the online
2142// flash reduction used by the causal kernels.
2143void attention_forward_full_head_major_gqa_exact_strided(const float *q,
2144 const float *k,
2145 const float *v,
2146 float *output,
2147 int num_heads,
2148 int num_kv_heads,
2149 int num_tokens,
2150 int head_dim,
2151 int aligned_head_dim,
2152 int kv_stride_tokens);
2153
2154// GGML-compatible non-flash full / bidirectional attention for encoder-style
2155// prefill. Mirrors the CPU mul_mat -> softmax -> mul_mat path more closely than
2156// the fused exact variant, and is intended for parity-sensitive vision models.
2157void attention_forward_full_head_major_gqa_ggml_strided(const float *q,
2158 const float *k,
2159 const float *v,
2160 float *output,
2161 int num_heads,
2162 int num_kv_heads,
2163 int num_tokens,
2164 int head_dim,
2165 int aligned_head_dim,
2166 int kv_stride_tokens);
2167void attention_forward_full_head_major_gqa_ggml_strided_workspace(
2168 const float *q, const float *k, const float *v, float *output,
2169 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2170 int aligned_head_dim, int kv_stride_tokens,
2171 float *score_rows, size_t score_rows_bytes,
2172 float *v_columns, size_t v_columns_bytes,
2173 float *probability_row, size_t probability_row_bytes);
2174
2175// Llama-parity flash attention variant that rounds K/V through FP16 before use.
2176void attention_forward_causal_head_major_gqa_flash_strided_f16kv(const float *q,
2177 const float *k,
2178 const float *v,
2179 float *output,
2180 int num_heads,
2181 int num_kv_heads,
2182 int num_tokens,
2183 int head_dim,
2184 int aligned_head_dim,
2185 int kv_stride_tokens);
2186void attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace(
2187 const float *q, const float *k, const float *v, float *output,
2188 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2189 int aligned_head_dim, int kv_stride_tokens,
2190 float *rounded_kv, size_t rounded_kv_bytes);
2191
2192// Decode attention for a single token using a KV cache (flash-style).
2193// q_token: [num_heads, aligned_head_dim]
2194// k_cache/v_cache: [num_kv_heads, cache_capacity, aligned_head_dim]
2195// out_token: [num_heads, aligned_head_dim]
2196void attention_forward_decode_head_major_gqa_flash(const float *q_token,
2197 const float *k_cache,
2198 const float *v_cache,
2199 float *out_token,
2200 int num_heads,
2201 int num_kv_heads,
2202 int kv_tokens,
2203 int cache_capacity,
2204 int head_dim,
2205 int aligned_head_dim);
2206void attention_forward_decode_head_major_gqa_muse_eager_bf16_storage(
2207 const float *q_token, const float *k_cache, const float *v_cache,
2208 float *out_token, int num_heads, int num_kv_heads, int kv_tokens,
2209 int cache_capacity, int head_dim, int aligned_head_dim,
2210 float *scores, size_t scores_bytes,
2211 uint16_t *key_rows, size_t key_rows_bytes,
2212 uint16_t *value_columns, size_t value_columns_bytes,
2213 uint16_t *gemm_rows, size_t gemm_rows_bytes,
2214 uint16_t *gemm_scores, size_t gemm_scores_bytes);
2215void attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_sliding(
2216 const float *q_token, const float *k_cache, const float *v_cache,
2217 float *out_token, int num_heads, int num_kv_heads, int kv_tokens,
2218 int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window,
2219 float *scores, size_t scores_bytes,
2220 uint16_t *key_rows, size_t key_rows_bytes,
2221 uint16_t *value_columns, size_t value_columns_bytes,
2222 uint16_t *gemm_rows, size_t gemm_rows_bytes,
2223 uint16_t *gemm_scores, size_t gemm_scores_bytes);
2224void attention_forward_decode_head_major_gqa_flash_gemma4(const float *q_token,
2225 const float *k_cache,
2226 const float *v_cache,
2227 float *out_token,
2228 int num_heads,
2229 int num_kv_heads,
2230 int kv_tokens,
2231 int cache_capacity,
2232 int head_dim,
2233 int aligned_head_dim);
2234void attention_forward_decode_head_major_shared_kv_gemma4(const float *q_token,
2235 const float *k_cache,
2236 const float *v_cache,
2237 float *out_token,
2238 int num_heads,
2239 int kv_tokens,
2240 int cache_capacity,
2241 int head_dim,
2242 int aligned_head_dim);
2243
2244// Chunk decode attention: q_chunk/out_chunk use compact head-major
2245// [num_heads, q_tokens, aligned_head_dim] layout while K/V are stored in the
2246// persistent cache [num_kv_heads, cache_capacity, aligned_head_dim].
2247void attention_forward_chunk_head_major_gqa_flash_gemma4(const float *q_chunk,
2248 const float *k_cache,
2249 const float *v_cache,
2250 float *out_chunk,
2251 int num_heads,
2252 int num_kv_heads,
2253 int q_tokens,
2254 int kv_tokens,
2255 int cache_capacity,
2256 int head_dim,
2257 int aligned_head_dim);
2258
2259// Llama-parity decode flash attention variant that rounds K/V through FP16 before use.
2260void attention_forward_decode_head_major_gqa_flash_f16kv(const float *q_token,
2261 const float *k_cache,
2262 const float *v_cache,
2263 float *out_token,
2264 int num_heads,
2265 int num_kv_heads,
2266 int kv_tokens,
2267 int cache_capacity,
2268 int head_dim,
2269 int aligned_head_dim);
2270
2271void attention_forward_decode_head_major_gqa_flash_f16cache(const float *q_token,
2272 const uint16_t *k_cache,
2273 const uint16_t *v_cache,
2274 float *out_token,
2275 int num_heads,
2276 int num_kv_heads,
2277 int kv_tokens,
2278 int cache_capacity,
2279 int head_dim,
2280 int aligned_head_dim);
2281
2282// Complete attention reduction contracts. These IDs describe numerical
2283// semantics, not an ISA or performance mode.
2293
2294// Research selectors for partitioning independent prefill attention outputs.
2295// They do not alter the numerical reduction performed for any output row.
2303
2310
2311// v8.5 explicit numerical-contract entry point. Unlike the legacy wrapper,
2312// this function never infers reduction semantics from strict/debug state.
2313ck_attention_status_t attention_forward_decode_head_major_gqa_flash_f16cache_contract(
2314 const float *q_token,
2315 const uint16_t *k_cache,
2316 const uint16_t *v_cache,
2317 float *out_token,
2318 int num_heads,
2319 int num_kv_heads,
2320 int kv_tokens,
2321 int cache_capacity,
2322 int head_dim,
2323 int aligned_head_dim,
2324 ck_attention_reduction_t reduction);
2325ck_attention_status_t attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract(
2326 const float *q_token,
2327 const uint16_t *k_cache,
2328 const uint16_t *v_cache,
2329 float *out_token,
2330 int num_heads,
2331 int num_kv_heads,
2332 int kv_tokens,
2333 int cache_capacity,
2334 int head_dim,
2335 int aligned_head_dim,
2336 ck_attention_reduction_t reduction);
2337void attention_forward_sparse_token_major_gqa_bf16cache_pytorch_cpu_flash_contract(
2338 const float *query,
2339 const uint16_t *key_cache,
2340 const uint16_t *value_cache,
2341 const float *selected_indices,
2342 float *output,
2343 float *score_scratch,
2344 int rows,
2345 int query_heads,
2346 int kv_heads,
2347 int head_dim,
2348 int selection_width,
2349 int context_length,
2350 int position);
2351// Returns whether the sparse QSA BF16 CPU-flash provider is present in this
2352// build. Unlike ordinary BF16 GQA, this provider requires AVX-512F.
2353int ck_attention_sparse_bf16_pytorch_gqa_available(void);
2354// Returns whether the exact native PyTorch BF16 GQA provider can execute in
2355// this build and process. This probes required MKL and SLEEF symbols.
2356int ck_attention_bf16_pytorch_gqa_available(void);
2357
2358// Causal prefill over a cache-preserving segment. Current-segment K/V rows
2359// must already be appended at [past_tokens, past_tokens + q_tokens).
2360ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract(
2361 const float *q,
2362 const uint16_t *k_cache,
2363 const uint16_t *v_cache,
2364 float *output,
2365 int num_heads,
2366 int num_kv_heads,
2367 int q_tokens,
2368 int past_tokens,
2369 int cache_capacity,
2370 int head_dim,
2371 int aligned_head_dim,
2372 ck_attention_reduction_t reduction);
2373ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(
2374 const float *q,
2375 const uint16_t *k_cache,
2376 const uint16_t *v_cache,
2377 float *output,
2378 int num_heads,
2379 int num_kv_heads,
2380 int q_tokens,
2381 int past_tokens,
2382 int cache_capacity,
2383 int head_dim,
2384 int aligned_head_dim,
2385 ck_attention_reduction_t reduction,
2386 float *token_workspace,
2387 size_t token_workspace_bytes);
2388ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace(
2389 const float *q, const uint16_t *k_cache, const uint16_t *v_cache,
2390 float *output, int num_heads, int num_kv_heads, int q_tokens,
2391 int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim,
2392 ck_attention_reduction_t reduction, float *token_workspace,
2393 size_t token_workspace_bytes);
2394ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace(
2395 const float *q, const uint16_t *k_cache, const uint16_t *v_cache,
2396 float *output, int num_heads, int num_kv_heads, int q_tokens,
2397 int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim,
2398 int sliding_window, ck_attention_reduction_t reduction,
2399 float *token_workspace, size_t token_workspace_bytes);
2400ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace(
2401 const float *q, const uint16_t *k_cache, const uint16_t *v_cache,
2402 float *output, int num_heads, int num_kv_heads, int q_tokens,
2403 int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim,
2404 ck_attention_reduction_t reduction,
2405 float *token_workspace, size_t token_workspace_bytes,
2406 const int *segment_lengths, int num_segments);
2407ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace(
2408 const float *q, const uint16_t *k_cache, const uint16_t *v_cache,
2409 float *output, int num_heads, int num_kv_heads, int q_tokens,
2410 int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim,
2411 int sliding_window, ck_attention_reduction_t reduction,
2412 float *token_workspace, size_t token_workspace_bytes,
2413 const int *segment_lengths, int num_segments);
2414ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_gemma4_contract(
2415 const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache,
2416 float *out_token, int num_heads, int num_kv_heads, int kv_tokens,
2417 int cache_capacity, int head_dim, int aligned_head_dim,
2418 ck_attention_reduction_t reduction);
2419ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_regular_noscale(
2420 const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache,
2421 float *out_token, int num_heads, int num_kv_heads, int kv_tokens,
2422 int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window);
2423ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract(
2424 const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache,
2425 float *out_token, int num_heads, int num_kv_heads, int kv_tokens,
2426 int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window,
2427 ck_attention_reduction_t reduction);
2428ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace(
2429 const float *q,
2430 const uint16_t *k_cache,
2431 const uint16_t *v_cache,
2432 float *output,
2433 int num_heads,
2434 int num_kv_heads,
2435 int q_tokens,
2436 int past_tokens,
2437 int cache_capacity,
2438 int head_dim,
2439 int aligned_head_dim,
2440 ck_attention_reduction_t reduction,
2441 float *token_workspace,
2442 size_t token_workspace_bytes,
2443 void *gqa_workspace,
2444 size_t gqa_workspace_bytes,
2445 int route_num_heads,
2446 int route_num_kv_heads,
2447 int route_head_dim,
2448 int route_query_tokens,
2449 int route_min_kv_tokens,
2450 int route_workers,
2451 int route_query_tile_size,
2452 int route_concurrent_query_tiles);
2453
2454// Isolated scheduling research entry point for the qtile64 numerical contract.
2455// Production circuits continue to use the contract entry point above.
2456ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule(
2457 const float *q,
2458 const uint16_t *k_cache,
2459 const uint16_t *v_cache,
2460 float *output,
2461 int num_heads,
2462 int num_kv_heads,
2463 int q_tokens,
2464 int past_tokens,
2465 int cache_capacity,
2466 int head_dim,
2467 int aligned_head_dim,
2468 ck_attention_prefill_schedule_t schedule);
2469
2470// Configurable GQA data-reuse entry point. A fixed 64-key block preserves the
2471// qtile64 numerical contract while the kernel map controls query-row tiling and
2472// concurrent query tiles without changing any output row's key reduction order.
2473size_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes(
2474 int num_heads,
2475 int num_kv_heads,
2476 int head_dim,
2477 int workers,
2478 int query_tile_size,
2479 int concurrent_query_tiles);
2480ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(
2481 const float *q,
2482 const uint16_t *k_cache,
2483 const uint16_t *v_cache,
2484 float *output,
2485 int num_heads,
2486 int num_kv_heads,
2487 int q_tokens,
2488 int past_tokens,
2489 int cache_capacity,
2490 int head_dim,
2491 int aligned_head_dim,
2492 int query_tile_size,
2493 int concurrent_query_tiles,
2494 void *workspace,
2495 size_t workspace_bytes);
2496
2497ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace(
2498 const float *q, const uint16_t *k_cache, const uint16_t *v_cache,
2499 float *output, int num_heads, int num_kv_heads, int q_tokens,
2500 int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim,
2501 ck_attention_reduction_t reduction, float *token_workspace,
2502 size_t token_workspace_bytes, const int *segment_lengths,
2503 int num_segments);
2504ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract(
2505 const float *q,
2506 const uint16_t *k_cache,
2507 const uint16_t *v_cache,
2508 float *output,
2509 int num_heads,
2510 int num_kv_heads,
2511 int q_tokens,
2512 int past_tokens,
2513 int cache_capacity,
2514 int head_dim,
2515 int aligned_head_dim,
2516 ck_attention_reduction_t reduction);
2517ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(
2518 const float *q,
2519 const uint16_t *k_cache,
2520 const uint16_t *v_cache,
2521 float *output,
2522 int num_heads,
2523 int num_kv_heads,
2524 int q_tokens,
2525 int past_tokens,
2526 int cache_capacity,
2527 int head_dim,
2528 int aligned_head_dim,
2529 ck_attention_reduction_t reduction,
2530 float *token_workspace,
2531 size_t token_workspace_bytes);
2532
2533// Full-matrix PyTorch math-SDPA prefill. Unlike the row-at-a-time segmented
2534// provider above, this entry point materializes [head, query, key] scores and
2535// applies the causal mask over the complete key width before FP32 softmax.
2536ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract(
2537 const float *q,
2538 const uint16_t *k_cache,
2539 const uint16_t *v_cache,
2540 float *output,
2541 int num_heads,
2542 int num_kv_heads,
2543 int q_tokens,
2544 int past_tokens,
2545 int cache_capacity,
2546 int head_dim,
2547 int aligned_head_dim,
2548 ck_attention_reduction_t reduction);
2549
2550// Deterministic llama.cpp-style FP16 split-KV oracle. The explicit chunk count
2551// makes diagnostics independent of the host's available core count. At and
2552// above a 512-row scheduling extent, worker boundaries use valid KV rounded
2553// to 256 rows while storage reads remain bounded by kv_tokens. This means the
2554// first split decode shape is valid KV=257 -> scheduling extent=512.
2555void attention_forward_decode_head_major_gqa_flash_f16cache_split(const float *q_token,
2556 const uint16_t *k_cache,
2557 const uint16_t *v_cache,
2558 float *out_token,
2559 int num_heads,
2560 int num_kv_heads,
2561 int kv_tokens,
2562 int cache_capacity,
2563 int head_dim,
2564 int aligned_head_dim,
2565 int split_chunks);
2566
2567// Decode attention for a single token using a KV cache (REGULAR - NOT flash).
2568// q_token: [num_heads, aligned_head_dim]
2569// k_cache/v_cache: [num_kv_heads, cache_capacity, aligned_head_dim]
2570// out_token: [num_heads, aligned_head_dim]
2571// WARNING: This is O(n) complexity, not true flash attention!
2572void attention_forward_decode_head_major_gqa_regular(const float *q_token,
2573 const float *k_cache,
2574 const float *v_cache,
2575 float *out_token,
2576 int num_heads,
2577 int num_kv_heads,
2578 int kv_tokens,
2579 int cache_capacity,
2580 int head_dim,
2581 int aligned_head_dim);
2582
2583// Sliding-window attention forward (prefill, flash-style)
2584// Each token attends to the last `sliding_window` tokens.
2585// sliding_window: window size (0 or negative = no limit, like regular causal)
2586void attention_forward_causal_head_major_gqa_flash_strided_sliding(
2587 const float *q,
2588 const float *k,
2589 const float *v,
2590 float *output,
2591 int num_heads,
2592 int num_kv_heads,
2593 int num_tokens,
2594 int head_dim,
2595 int aligned_head_dim,
2596 int kv_stride_tokens,
2597 int sliding_window);
2598void attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace(
2599 const float *q, const float *k, const float *v, float *output,
2600 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
2601 int aligned_head_dim, int kv_stride_tokens, int sliding_window,
2602 float *scores, size_t scores_bytes,
2603 float *value_columns, size_t value_columns_bytes,
2604 float *scaled_scores, size_t scaled_scores_bytes);
2605void attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace(
2606 const float *q, const float *k, const float *v, float *output,
2607 int num_heads, int num_kv_heads, int live_tokens, int kv_stride_tokens,
2608 int head_dim, int aligned_head_dim, int sliding_window,
2609 float *scores, size_t scores_bytes,
2610 float *value_columns, size_t value_columns_bytes,
2611 float *scaled_scores, size_t scaled_scores_bytes);
2612
2613void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4(
2614 const float *q,
2615 const float *k,
2616 const float *v,
2617 float *output,
2618 int num_heads,
2619 int num_kv_heads,
2620 int num_tokens,
2621 int head_dim,
2622 int aligned_head_dim,
2623 int kv_stride_tokens,
2624 int sliding_window);
2625void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output(
2626 const float *q,
2627 const float *k,
2628 const float *v,
2629 float *output,
2630 int num_heads,
2631 int num_kv_heads,
2632 int num_tokens,
2633 int head_dim,
2634 int aligned_head_dim,
2635 int kv_stride_tokens,
2636 int sliding_window);
2637void attention_forward_causal_head_major_shared_kv_sliding_gemma4(
2638 const float *q,
2639 float *output,
2640 int num_heads,
2641 int num_tokens,
2642 int head_dim,
2643 int aligned_head_dim,
2644 int kv_stride_tokens,
2645 int sliding_window);
2646
2647// Gated DeltaNet recurrent update used by qwen3next/Qwen3.5 linear attention.
2648// Layout:
2649// q, k, v : [num_heads, state_dim]
2650// g, beta : [num_heads]
2651// state_* : [num_heads, state_dim, state_dim] row-major per head
2652// out : [num_heads, state_dim]
2653//
2654// Qwen3.5/qwen3next also uses an SSM causal depthwise convolution before the
2655// recurrent DeltaNet update. That kernel consumes:
2656// conv_x : [seqs, channels, kernel_size - 1 + num_tokens]
2657// kernel : [channels, kernel_size]
2658// out : [seqs, num_tokens, channels]
2659// and matches ggml's GGML_OP_SSM_CONV layout/semantics.
2660void ssm_conv1d_forward(const float *conv_x,
2661 const float *kernel,
2662 float *out,
2663 int kernel_size,
2664 int num_channels,
2665 int num_tokens,
2666 int num_seqs);
2667void ssm_conv1d_forward_llama_production(const float *conv_x,
2668 const float *kernel,
2669 float *out,
2670 int kernel_size,
2671 int num_channels,
2672 int num_tokens,
2673 int num_seqs);
2674void ssm_conv1d_forward_llama_production_serial(const float *conv_x,
2675 const float *kernel,
2676 float *out,
2677 int kernel_size,
2678 int num_channels,
2679 int num_tokens,
2680 int num_seqs);
2681void ssm_conv1d_forward_llama_fma(const float *conv_x,
2682 const float *kernel,
2683 float *out,
2684 int kernel_size,
2685 int num_channels,
2686 int num_tokens,
2687 int num_seqs);
2688void ssm_conv1d_forward_pytorch_bf16_storage(const float *conv_x,
2689 const float *kernel,
2690 float *out,
2691 int kernel_size,
2692 int num_channels,
2693 int num_tokens,
2694 int num_seqs);
2695
2696// Backward for the causal depthwise SSM convolution used in qwen3next/Qwen3.5.
2697// Layout:
2698// d_out : [seqs, num_tokens, channels]
2699// conv_x : [seqs, channels, kernel_size - 1 + num_tokens]
2700// kernel : [channels, kernel_size]
2701// d_conv_x : same as conv_x
2702// d_kernel : same as kernel
2703void ssm_conv1d_backward(const float *d_out,
2704 const float *conv_x,
2705 const float *kernel,
2706 float *d_conv_x,
2707 float *d_kernel,
2708 int kernel_size,
2709 int num_channels,
2710 int num_tokens,
2711 int num_seqs);
2712
2713// Split a packed recurrent QKV projection into explicit Q, K, and V rows.
2714// Layout:
2715// packed_qkv : [rows, q_dim + k_dim + v_dim] row-major
2716// q : [rows, q_dim]
2717// k : [rows, k_dim]
2718// v : [rows, v_dim]
2719//
2720// For qwen3.5 recurrent layers, rows usually means num_tokens * num_seqs after
2721// the packed projection has been flattened into a token-major matrix.
2722//
2723// Full-attention qwen3.5 layers also use a joint Q+gate projection before the
2724// standard attention block:
2725// packed_qg : [rows, interleaved (q_group, gate_group) segments]
2726// q : [rows, q_dim]
2727// gate : [rows, gate_dim]
2728// group_dim : per-group width; if q_dim == gate_dim == group_dim the layout
2729// degenerates to a single [Q | G] split, otherwise it is treated
2730// as [Q0 | G0 | Q1 | G1 | ...].
2731void split_q_gate_forward(const float *packed_qg,
2732 float *q,
2733 float *gate,
2734 int rows,
2735 int q_dim,
2736 int gate_dim,
2737 int group_dim);
2738
2739void split_q_gate_backward(const float *d_q,
2740 const float *d_gate,
2741 float *d_packed_qg,
2742 int rows,
2743 int q_dim,
2744 int gate_dim,
2745 int group_dim);
2746
2747void recurrent_split_qkv_forward(const float *packed_qkv,
2748 float *q,
2749 float *k,
2750 float *v,
2751 int rows,
2752 int q_dim,
2753 int k_dim,
2754 int v_dim);
2755
2756// Split packed QKV rows into head-major Q/K/V buffers used by transformer
2757// attention kernels.
2758// Layout:
2759// packed_qkv : [rows, q_dim + k_dim + v_dim] token-major
2760// q : [num_heads, rows, q_dim / num_heads] head-major
2761// k : [num_kv_heads, rows, k_dim / num_kv_heads] head-major
2762// v : [num_kv_heads, rows, v_dim / num_kv_heads] head-major
2763void split_qkv_packed_head_major_forward(const float *packed_qkv,
2764 float *q,
2765 float *k,
2766 float *v,
2767 int rows,
2768 int q_dim,
2769 int k_dim,
2770 int v_dim,
2771 int num_heads,
2772 int num_kv_heads);
2773
2774// Backward for the packed recurrent QKV split.
2775// Layout:
2776// d_q : [rows, q_dim]
2777// d_k : [rows, k_dim]
2778// d_v : [rows, v_dim]
2779// d_packed_qkv : [rows, q_dim + k_dim + v_dim]
2780void recurrent_split_qkv_backward(const float *d_q,
2781 const float *d_k,
2782 const float *d_v,
2783 float *d_packed_qkv,
2784 int rows,
2785 int q_dim,
2786 int k_dim,
2787 int v_dim);
2788
2789// Transform recurrent alpha rows into the DeltaNet gate using per-dimension
2790// dt bias and scale weights.
2791// Layout:
2792// alpha : [rows, dim]
2793// dt_bias : [dim]
2794// a : [dim]
2795// gate : [rows, dim]
2796// Formula:
2797// gate[row, col] = softplus(alpha[row, col] + dt_bias[col]) * a[col]
2798void recurrent_dt_gate_forward(const float *alpha,
2799 const float *dt_bias,
2800 const float *a,
2801 float *gate,
2802 int rows,
2803 int num_heads,
2804 int state_dim);
2805
2806// Bit-exact PyTorch FP32 softplus contract for Qwen4-Exp recurrent gates.
2807// The validated AVX-512 provider follows PyTorch's vector exp/log1p path.
2808void recurrent_dt_gate_forward_pytorch_fp32(const float *alpha,
2809 const float *dt_bias,
2810 const float *a,
2811 float *gate,
2812 int rows,
2813 int num_heads,
2814 int state_dim);
2815
2816// Backward for recurrent_dt_gate_forward.
2817// Layout:
2818// d_gate : [rows, dim]
2819// alpha : [rows, dim]
2820// dt_bias : [dim]
2821// a : [dim]
2822// d_alpha : [rows, dim]
2823// d_dt_bias : [dim]
2824// d_a : [dim]
2825// Expanded Qwen3.5/KDA-style dt gate.
2826// Layout:
2827// alpha : [rows, num_heads]
2828// dt_bias : [num_heads]
2829// a : [num_heads, state_dim] (already converted to -exp(A_log))
2830// gate : [rows, num_heads * state_dim]
2831void recurrent_dt_gate_expanded_forward(const float *alpha,
2832 const float *dt_bias,
2833 const float *a,
2834 float *gate,
2835 int rows,
2836 int num_heads,
2837 int state_dim);
2838
2839void recurrent_dt_gate_backward(const float *d_gate,
2840 const float *alpha,
2841 const float *dt_bias,
2842 const float *a,
2843 float *d_alpha,
2844 float *d_dt_bias,
2845 float *d_a,
2846 int rows,
2847 int dim);
2848
2849// Update the causal convolution history for a recurrent SSM block.
2850// Layout:
2851// state_in : [num_seqs, channels, history_len]
2852// q : [num_seqs * num_tokens, q_dim]
2853// k : [num_seqs * num_tokens, k_dim]
2854// v : [num_seqs * num_tokens, v_dim]
2855// conv_x : [num_seqs, channels, history_len + num_tokens]
2856// state_out : [num_seqs, channels, history_len]
2857void recurrent_conv_state_update_forward(const float *state_in,
2858 const float *q,
2859 const float *k,
2860 const float *v,
2861 float *conv_x,
2862 float *state_out,
2863 int history_len,
2864 int num_seqs,
2865 int num_tokens,
2866 int q_dim,
2867 int k_dim,
2868 int v_dim);
2869
2870// Backward for recurrent_conv_state_update_forward.
2871void recurrent_conv_state_update_backward(const float *d_conv_x,
2872 const float *d_state_out,
2873 float *d_state_in,
2874 float *d_q,
2875 float *d_k,
2876 float *d_v,
2877 int history_len,
2878 int num_seqs,
2879 int num_tokens,
2880 int q_dim,
2881 int k_dim,
2882 int v_dim);
2883void recurrent_conv_state_update_backward_workspace(const float *d_conv_x,
2884 const float *d_state_out,
2885 float *d_state_in,
2886 float *d_q,
2887 float *d_k,
2888 float *d_v,
2889 float *d_conv_total,
2890 int history_len,
2891 int num_seqs,
2892 int num_tokens,
2893 int q_dim,
2894 int k_dim,
2895 int v_dim);
2896
2897// Apply SiLU elementwise to recurrent convolution output rows.
2898// Layout:
2899// x : [rows, dim]
2900// out : [rows, dim]
2901void recurrent_silu_forward(const float *x,
2902 float *out,
2903 int rows,
2904 int dim);
2905void recurrent_silu_forward_pytorch_bf16_storage(const float *x,
2906 float *out,
2907 int rows,
2908 int dim);
2909void recurrent_silu_forward_pytorch_bf16_input_fp32_output(const float *x,
2910 float *out,
2911 int rows,
2912 int dim);
2913void recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(const float *x,
2914 float *out,
2915 int rows,
2916 int dim);
2917void recurrent_silu_forward_ggml(const float *x,
2918 float *out,
2919 int rows,
2920 int dim);
2921void recurrent_sigmoid_forward_ggml(const float *x,
2922 float *out,
2923 int rows,
2924 int dim);
2925
2926// Backward for recurrent_silu_forward.
2927void recurrent_silu_backward(const float *d_out,
2928 const float *x,
2929 float *d_x,
2930 int rows,
2931 int dim);
2932
2933// Split the post-convolution recurrent packed QKV rows into explicit Q, K, and V.
2934// Layout matches recurrent_split_qkv_forward.
2935void recurrent_split_conv_qkv_forward(const float *packed_qkv,
2936 float *q,
2937 float *k,
2938 float *v,
2939 int rows,
2940 int q_dim,
2941 int k_dim,
2942 int v_dim);
2943
2944// Backward for recurrent_split_conv_qkv_forward.
2945void recurrent_split_conv_qkv_backward(const float *d_q,
2946 const float *d_k,
2947 const float *d_v,
2948 float *d_packed_qkv,
2949 int rows,
2950 int q_dim,
2951 int k_dim,
2952 int v_dim);
2953
2954// Apply per-head L2 normalization to recurrent Q and K rows in-place.
2955// Layout:
2956// q : [rows, q_dim]
2957// k : [rows, k_dim]
2958// where q_dim and k_dim are multiples of head_dim.
2959void recurrent_qk_l2_norm_forward(float *q,
2960 float *k,
2961 int rows,
2962 int q_dim,
2963 int k_dim,
2964 int head_dim,
2965 float eps);
2966void recurrent_qk_l2_norm_pytorch_bf16_storage(float *q,
2967 float *k,
2968 int rows,
2969 int q_dim,
2970 int k_dim,
2971 int expanded_heads,
2972 int head_dim,
2973 float eps);
2974void recurrent_qk_l2_norm_pytorch_fp32_output(float *q,
2975 float *k,
2976 int rows,
2977 int q_dim,
2978 int k_dim,
2979 int head_dim,
2980 float eps);
2981
2982// Backward for recurrent_qk_l2_norm_forward.
2983void recurrent_qk_l2_norm_backward(const float *d_q_out,
2984 const float *d_k_out,
2985 const float *q,
2986 const float *k,
2987 float *d_q,
2988 float *d_k,
2989 int rows,
2990 int q_dim,
2991 int k_dim,
2992 int head_dim,
2993 float eps);
2994
2995// Per-head RMSNorm followed by SiLU(z) gating for recurrent outputs.
2996// Layout:
2997// x : [rows, num_heads * head_dim]
2998// gate : [rows, num_heads * head_dim]
2999// weight : [head_dim]
3000// out : [rows, num_heads * head_dim]
3001void recurrent_norm_gate_forward(const float *x,
3002 const float *gate,
3003 const float *weight,
3004 float *out,
3005 int rows,
3006 int num_heads,
3007 int head_dim,
3008 float eps);
3009void recurrent_norm_sigmoid_gate_llama_avx2_forward(const float *x,
3010 const float *gate,
3011 const float *weight,
3012 float *out,
3013 int rows,
3014 int num_heads,
3015 int head_dim,
3016 float eps);
3017
3018// Backward for recurrent_norm_gate_forward.
3019void recurrent_norm_gate_backward(const float *d_out,
3020 const float *x,
3021 const float *gate,
3022 const float *weight,
3023 float *d_x,
3024 float *d_gate,
3025 float *d_weight,
3026 int rows,
3027 int num_heads,
3028 int head_dim,
3029 float eps);
3030
3031// Nemotron-H/Mamba2 scalar reference kernels.
3032// These cover the decode-state contract first; chunked prefill scan is a
3033// separate scheduling contract and should not be hidden inside this API.
3034void mamba2_in_proj_split_f32(const float *projected,
3035 float *gate,
3036 float *hidden_bc,
3037 float *dt,
3038 int rows,
3039 int d_mlp,
3040 int intermediate_dim,
3041 int conv_dim,
3042 int num_heads);
3043
3044void mamba2_conv1d_decode_f32(const float *state_in,
3045 const float *x,
3046 const float *weight,
3047 const float *bias,
3048 float *conv_out,
3049 float *state_out,
3050 int rows,
3051 int conv_dim,
3052 int kernel_size);
3053
3054void mamba2_conv1d_f32_channel_range(const float *state_in,
3055 const float *x,
3056 const float *weight,
3057 const float *bias,
3058 float *conv_out,
3059 float *state_out,
3060 int rows,
3061 int conv_dim,
3062 int kernel_size,
3063 int channel_begin,
3064 int channel_end);
3065
3066void mamba2_conv1d_f32_parallel_dispatch(const float *state_in,
3067 const float *x,
3068 const float *weight,
3069 const float *bias,
3070 float *conv_out,
3071 float *state_out,
3072 int rows,
3073 int conv_dim,
3074 int kernel_size);
3075
3076void mamba2_dt_softplus_f32(const float *dt,
3077 const float *dt_bias,
3078 float *dt_out,
3079 int rows,
3080 int num_heads,
3081 float dt_min,
3082 float dt_max);
3083
3084void mamba2_selective_state_update_decode_f32(const float *state_in,
3085 const float *x,
3086 const float *dt,
3087 const float *a,
3088 const float *b,
3089 const float *c,
3090 const float *d,
3091 float *state_out,
3092 float *y,
3093 int rows,
3094 int num_heads,
3095 int head_dim,
3096 int state_dim,
3097 int num_groups);
3098
3099void mamba2_selective_scan_f32(const float *state_init,
3100 const float *x,
3101 const float *dt,
3102 const float *a,
3103 const float *b,
3104 const float *c,
3105 const float *d,
3106 float *state_out,
3107 float *y,
3108 int batch,
3109 int seq_len,
3110 int num_heads,
3111 int head_dim,
3112 int state_dim,
3113 int num_groups);
3114
3115void mamba2_selective_scan_f32_head_range(const float *state_init,
3116 const float *x,
3117 const float *dt,
3118 const float *a,
3119 const float *b,
3120 const float *c,
3121 const float *d,
3122 float *state_out,
3123 float *y,
3124 int batch,
3125 int seq_len,
3126 int num_heads,
3127 int head_dim,
3128 int state_dim,
3129 int num_groups,
3130 int head_begin,
3131 int head_end);
3132
3133void mamba2_selective_scan_f32_parallel_dispatch(const float *state_init,
3134 const float *x,
3135 const float *dt,
3136 const float *a,
3137 const float *b,
3138 const float *c,
3139 const float *d,
3140 float *state_out,
3141 float *y,
3142 int batch,
3143 int seq_len,
3144 int num_heads,
3145 int head_dim,
3146 int state_dim,
3147 int num_groups);
3148
3149void mamba2_rmsnorm_gate_f32(const float *x,
3150 const float *gate,
3151 const float *weight,
3152 float *out,
3153 int rows,
3154 int inner_dim,
3155 int group_size,
3156 float eps);
3157
3158// Apply sigmoid(gate) to the full-attention qwen3.5 gate path and multiply it
3159// elementwise with the attention output before the output projection.
3160// Layout:
3161// x : [rows, num_heads * state_dim]
3162// gate : [rows, num_heads * state_dim]
3163// out : [rows, num_heads * state_dim]
3164void attn_gate_sigmoid_mul_forward(const float *x,
3165 const float *gate,
3166 float *out,
3167 int rows,
3168 int num_heads,
3169 int state_dim);
3170void attn_gate_sigmoid_mul_pytorch_bf16_storage(const float *x,
3171 const float *gate,
3172 float *out,
3173 int rows,
3174 int num_heads,
3175 int state_dim);
3176
3177void attn_gate_softplus_mul_forward(const float *x,
3178 const float *gate,
3179 float *out,
3180 int rows,
3181 int num_heads,
3182 int state_dim);
3183
3184void attn_gate_sigmoid_mul_backward(const float *d_out,
3185 const float *x,
3186 const float *gate,
3187 float *d_x,
3188 float *d_gate,
3189 int rows,
3190 int num_heads,
3191 int state_dim);
3192
3193void gated_deltanet_autoregressive_forward(const float *q,
3194 const float *k,
3195 const float *v,
3196 const float *g,
3197 const float *beta,
3198 const float *state_in,
3199 float *state_out,
3200 float *out,
3201 int num_heads,
3202 int state_dim,
3203 float norm_eps);
3204
3205void gated_deltanet_prefill_forward(const float *q,
3206 const float *k,
3207 const float *v,
3208 const float *g,
3209 const float *beta,
3210 const float *state_in,
3211 float *state_out,
3212 float *out,
3213 int rows,
3214 int num_heads,
3215 int state_dim,
3216 float norm_eps);
3217
3218void gated_deltanet_llama_avx2_forward(const float *q,
3219 const float *k,
3220 const float *v,
3221 const float *g,
3222 const float *beta,
3223 const float *state_in,
3224 float *state_out,
3225 float *out,
3226 int num_heads,
3227 int group_count,
3228 int state_dim,
3229 float norm_eps);
3230
3231void gated_deltanet_llama_avx2_prefill_forward(const float *q,
3232 const float *k,
3233 const float *v,
3234 const float *g,
3235 const float *beta,
3236 const float *state_in,
3237 float *state_out,
3238 float *out,
3239 int rows,
3240 int num_heads,
3241 int group_count,
3242 int state_dim,
3243 float norm_eps);
3244
3245void gated_deltanet_llama_chunk64_prefill_forward(const float *q,
3246 const float *k,
3247 const float *v,
3248 const float *g,
3249 const float *beta,
3250 const float *state_in,
3251 float *state_out,
3252 float *out,
3253 int rows,
3254 int num_heads,
3255 int group_count,
3256 int state_dim,
3257 float norm_eps);
3258
3259void gated_deltanet_llama_chunk64_head_forward(const float *q,
3260 const float *k,
3261 const float *v,
3262 const float *g,
3263 const float *beta,
3264 const float *state_in,
3265 float *state_out,
3266 float *out,
3267 int rows,
3268 int num_heads,
3269 int group_count,
3270 int head,
3271 int state_dim);
3272
3273void gated_deltanet_pytorch_grouped_bf16_forward(const float *q,
3274 const float *k,
3275 const float *v,
3276 const float *g,
3277 const float *beta,
3278 const float *state_in,
3279 float *state_out,
3280 float *out,
3281 int num_heads,
3282 int group_count,
3283 int state_dim,
3284 float norm_eps);
3285
3286void gated_deltanet_pytorch_grouped_bf16_prefill_forward(
3287 const float *q,
3288 const float *k,
3289 const float *v,
3290 const float *g,
3291 const float *beta,
3292 const float *state_in,
3293 float *state_out,
3294 float *out,
3295 int rows,
3296 int num_heads,
3297 int group_count,
3298 int state_dim,
3299 float norm_eps);
3300
3301void recurrent_norm_gate_llama_avx2_forward(const float *x,
3302 const float *gate,
3303 const float *weight,
3304 float *out,
3305 int rows,
3306 int num_heads,
3307 int head_dim,
3308 float eps);
3309void recurrent_norm_gate_pytorch_bf16_storage(const float *x,
3310 const float *gate,
3311 const float *weight,
3312 float *out,
3313 int rows,
3314 int num_heads,
3315 int head_dim,
3316 float eps);
3317void recurrent_norm_sigmoid_gate_pytorch_bf16_storage(const float *x,
3318 const float *gate,
3319 const float *weight,
3320 float *out,
3321 int rows,
3322 int num_heads,
3323 int head_dim,
3324 float eps);
3325
3326// Gated DeltaNet recurrent backward used by qwen3next/Qwen3.5 linear attention.
3327// Layout:
3328// d_out : [num_heads, state_dim]
3329// d_state_out : [num_heads, state_dim, state_dim] row-major per head
3330// q, k, v : [num_heads, state_dim]
3331// g, beta : [num_heads]
3332// state_* : [num_heads, state_dim, state_dim] row-major per head
3333// d_q/d_k/d_v : [num_heads, state_dim]
3334// d_g/d_beta : [num_heads]
3335void gated_deltanet_autoregressive_backward(const float *d_out,
3336 const float *d_state_out,
3337 const float *q,
3338 const float *k,
3339 const float *v,
3340 const float *g,
3341 const float *beta,
3342 const float *state_in,
3343 const float *state_out,
3344 float *d_q,
3345 float *d_k,
3346 float *d_v,
3347 float *d_g,
3348 float *d_beta,
3349 float *d_state_in,
3350 int num_heads,
3351 int state_dim,
3352 float norm_eps);
3353
3354// Sliding-window attention forward (decode, flash-style)
3355// Single query token attends to the last `sliding_window` tokens in KV cache.
3356void attention_forward_decode_head_major_gqa_flash_sliding(
3357 const float *q_token,
3358 const float *k_cache,
3359 const float *v_cache,
3360 float *out_token,
3361 int num_heads,
3362 int num_kv_heads,
3363 int kv_tokens,
3364 int cache_capacity,
3365 int head_dim,
3366 int aligned_head_dim,
3367 int sliding_window);
3368
3369void attention_forward_decode_head_major_gqa_flash_sliding_gemma4(
3370 const float *q_token,
3371 const float *k_cache,
3372 const float *v_cache,
3373 float *out_token,
3374 int num_heads,
3375 int num_kv_heads,
3376 int kv_tokens,
3377 int cache_capacity,
3378 int head_dim,
3379 int aligned_head_dim,
3380 int sliding_window);
3381void attention_forward_decode_head_major_shared_kv_sliding_gemma4(
3382 const float *q_token,
3383 const float *k_cache,
3384 const float *v_cache,
3385 float *out_token,
3386 int num_heads,
3387 int kv_tokens,
3388 int cache_capacity,
3389 int head_dim,
3390 int aligned_head_dim,
3391 int sliding_window);
3392
3393// TRUE Flash Attention (O(1) for decode) - Tri Dao's algorithm
3394// out: [T_q, H, D_h]
3395// q: [T_q, H, D_h]
3396// k: [T_k, H, D_h]
3397// v: [T_k, H, D_h]
3398// T_q: Query tokens (1 for decode)
3399// T_k: Context length
3400// H: Number of heads
3401// D_h: Head dimension
3402// scale: 1/sqrt(D_h)
3403void attention_flash_decode(float *out,
3404 const float *q,
3405 const float *k,
3406 const float *v,
3407 int T_q,
3408 int T_k,
3409 int H,
3410 int D_h,
3411 float scale);
3412
3413// Diagnostics for flash attention tuning (used by unit tests).
3414int ck_flash_attn_choose_tile_k(int D_h);
3415int ck_flash_attn_fast_exp_kind(void);
3416
3417// Orchestration wrapper for TRUE flash attention
3418void ck_attention_flash_decode_wrapper(const float *q_token,
3419 const float *k_cache,
3420 const float *v_cache,
3421 float *out_token,
3422 int num_heads,
3423 int num_kv_heads,
3424 int kv_tokens,
3425 int cache_capacity,
3426 int head_dim,
3427 int aligned_head_dim);
3428
3429// KV cache helper (write one token for all KV heads).
3430void kv_cache_write_head_major(const float *__restrict k_token,
3431 const float *__restrict v_token,
3432 float *__restrict k_cache,
3433 float *__restrict v_cache,
3434 int num_kv_heads,
3435 int token_index,
3436 int cache_capacity,
3437 int head_dim,
3438 int aligned_head_dim);
3439
3440void kv_cache_store(float *__restrict kv_cache_k,
3441 float *__restrict kv_cache_v,
3442 const float *__restrict k,
3443 const float *__restrict v,
3444 int layer,
3445 int pos,
3446 int num_kv_heads,
3447 int head_dim,
3448 int max_seq_len);
3449void kv_cache_store_shared_q(float *__restrict kv_cache_k,
3450 float *__restrict kv_cache_v,
3451 const float *__restrict q,
3452 int layer,
3453 int pos,
3454 int num_heads,
3455 int head_dim,
3456 int max_seq_len);
3457
3458void kv_cache_store_f16(uint16_t *__restrict kv_cache_k,
3459 uint16_t *__restrict kv_cache_v,
3460 const float *__restrict k,
3461 const float *__restrict v,
3462 int layer,
3463 int pos,
3464 int num_kv_heads,
3465 int head_dim,
3466 int max_seq_len);
3467void kv_cache_store_bf16(uint16_t *__restrict kv_cache_k,
3468 uint16_t *__restrict kv_cache_v,
3469 const float *__restrict k,
3470 const float *__restrict v,
3471 int layer,
3472 int pos,
3473 int num_kv_heads,
3474 int head_dim,
3475 int max_seq_len);
3476void kv_cache_store_batch_f32(float *__restrict kv_cache_k,
3477 float *__restrict kv_cache_v,
3478 const float *__restrict k,
3479 const float *__restrict v,
3480 int start_pos,
3481 int num_tokens,
3482 int num_kv_heads,
3483 int head_dim,
3484 int max_seq_len);
3485void kv_cache_store_batch_bf16(uint16_t *__restrict kv_cache_k,
3486 uint16_t *__restrict kv_cache_v,
3487 const float *__restrict k,
3488 const float *__restrict v,
3489 int start_pos,
3490 int num_tokens,
3491 int num_kv_heads,
3492 int head_dim,
3493 int max_seq_len);
3494
3495void kv_cache_store_batch_f16(uint16_t *__restrict kv_cache_k,
3496 uint16_t *__restrict kv_cache_v,
3497 const float *__restrict k,
3498 const float *__restrict v,
3499 int start_pos,
3500 int num_tokens,
3501 int num_kv_heads,
3502 int head_dim,
3503 int max_seq_len);
3504
3505// Repack a head-major tensor from a packed `[head, tokens, aligned_head_dim]`
3506// layout into a KV-cache-compatible layout `[head, cache_capacity, aligned_head_dim]`
3507// in-place. This is used after prefill when forward kernels write head slices
3508// back-to-back using `tokens` as the head stride, but decode expects a fixed
3509// `cache_capacity` stride.
3510void kv_cache_repack_head_major_inplace(float *buf,
3511 int num_heads,
3512 int tokens,
3513 int cache_capacity,
3514 int aligned_head_dim);
3515
3516// MLP forward kernel (FC1 -> GELU -> FC2), generic token-parallel version.
3517void mlp_token_parallel(const float *input,
3518 const float *W_fc1,
3519 const float *b_fc1,
3520 const float *W_fc2,
3521 const float *b_fc2,
3522 float *fc1_output,
3523 float *output,
3524 int T,
3525 int aligned_dim,
3526 int num_threads);
3527
3528// Exact version using scalar GELU with standard library tanhf.
3529// Slower but provides maximum accuracy. Used for correctness testing.
3530void mlp_token_parallel_exact(const float *input,
3531 const float *W_fc1,
3532 const float *b_fc1,
3533 const float *W_fc2,
3534 const float *b_fc2,
3535 float *fc1_output,
3536 float *output,
3537 int T,
3538 int aligned_dim,
3539 int num_threads);
3540
3541/* BF16 MLP forward - caller provides scratch buffers */
3542void mlp_token_parallel_bf16(const uint16_t *input,
3543 const uint16_t *W_fc1,
3544 const uint16_t *b_fc1,
3545 const uint16_t *W_fc2,
3546 const uint16_t *b_fc2,
3547 float *fc1_output,
3548 float *output,
3549 int T,
3550 int aligned_dim,
3551 int num_threads,
3552 float *scratch_bias1_f, /* [4*D] */
3553 float *scratch_bias2_f, /* [D] */
3554 uint16_t *scratch_fc1_bf16); /* [T * 4*D] */
3555
3556/* BF16 MLP forward with FP32 activations - caller provides scratch buffers */
3557void mlp_token_parallel_bf16_fp32act(const uint16_t *input,
3558 const uint16_t *W_fc1,
3559 const uint16_t *b_fc1,
3560 const uint16_t *W_fc2,
3561 const uint16_t *b_fc2,
3562 float *fc1_output,
3563 float *output,
3564 int T,
3565 int aligned_dim,
3566 int num_threads,
3567 float *scratch_input_f, /* [T * D] */
3568 float *scratch_bias1_f, /* [4*D] */
3569 float *scratch_bias2_f, /* [D] */
3570 uint16_t *scratch_fc1_bf16); /* [T * 4*D] */
3571
3572/* BF16 MLP backward: BF16 storage, FP32 gradient accumulation. */
3573void mlp_token_parallel_bf16_backward_mixed(const uint16_t *input,
3574 const uint16_t *W_fc1,
3575 const uint16_t *b_fc1,
3576 const uint16_t *W_fc2,
3577 const uint16_t *d_output,
3578 float *d_input,
3579 float *d_W_fc1,
3580 float *d_b_fc1,
3581 float *d_W_fc2,
3582 float *d_b_fc2,
3583 int T,
3584 int aligned_dim,
3585 int num_threads,
3586 float *scratch_fc1_pre,
3587 uint16_t *scratch_fc1_act_bf16,
3588 float *scratch_d_fc1);
3589
3590// MLP FC1/FC2 backward kernels (generic), adapted from C-Transformer.
3591void fc2_backward_kernel(const float *d_output,
3592 const float *fc2_input,
3593 const float *W_fc2,
3594 float *d_input,
3595 float *d_W_fc2,
3596 float *d_b_fc2,
3597 int T,
3598 int aligned_in,
3599 int aligned_out,
3600 int num_threads);
3601
3602void fc1_backward_kernel(const float *d_output,
3603 const float *fc1_input,
3604 const float *W_fc1,
3605 float *d_input,
3606 float *d_W_fc1,
3607 float *d_b_fc1,
3608 int T,
3609 int aligned_in,
3610 int aligned_out,
3611 int num_threads);
3612
3613// Sigmoid kernels (scalar + vector forms).
3614float sigmoid_scalar(float x);
3615
3616void sigmoid_forward(const float *input,
3617 float *output,
3618 size_t n);
3619
3620void sigmoid_backward(const float *input,
3621 const float *d_output,
3622 float *d_input,
3623 size_t n);
3624
3625/* BF16 sigmoid - caller provides scratch buffers [n] floats each */
3626void sigmoid_forward_bf16(const uint16_t *input,
3627 uint16_t *output,
3628 size_t n,
3629 float *scratch_input,
3630 float *scratch_output);
3631
3632void sigmoid_backward_bf16(const uint16_t *input,
3633 const uint16_t *d_output,
3634 uint16_t *d_input,
3635 size_t n,
3636 float *scratch_input,
3637 float *scratch_d_output,
3638 float *scratch_d_input);
3639
3640// SwiGLU activation kernels (forward + backward).
3641// Input layout per token: [gate[0..D-1], value[0..D-1]], size 2*D.
3642// Output: [D].
3643 void swiglu_forward(const float *input,
3644 float *output,
3645 int tokens,
3646 int dim);
3647
3648 void swiglu_forward_q8_k(const float *input,
3649 void *output_q8,
3650 int tokens,
3651 int dim);
3652
3653 void swiglu_backward(const float *input,
3654 const float *d_output,
3655 float *d_input,
3656 int tokens,
3657 int dim);
3658
3659 // Exact versions using standard library expf (slower but accurate)
3660 void swiglu_forward_exact(const float *input,
3661 float *output,
3662 int tokens,
3663 int dim);
3664
3665void swiglu_forward_ggml(const float *input,
3666 float *output,
3667 int tokens,
3668 int dim);
3669
3670void swiglu_forward_ggml_split(const float *gate,
3671 const float *up,
3672 float *output,
3673 int tokens,
3674 int dim);
3675
3676 void swiglu_forward_pytorch_bf16_storage(const float *input,
3677 float *output,
3678 int tokens,
3679 int dim);
3680
3681 void swiglu_backward_exact(const float *input,
3682 const float *d_output,
3683 float *d_input,
3684 int tokens,
3685 int dim);
3686
3687 void swiglu_forward_bf16(const uint16_t *input,
3688 uint16_t *output,
3689 int tokens,
3690 int dim);
3691
3692 void swiglu_backward_bf16(const uint16_t *input,
3693 const uint16_t *d_output,
3694 uint16_t *d_input,
3695 int tokens,
3696 int dim);
3697
3698// =============================================================================
3699// Element-wise addition kernels (for residual connections)
3700// =============================================================================
3701
3702// Forward: y = a + b
3703void add_forward_bf16(const uint16_t *a,
3704 const uint16_t *b,
3705 uint16_t *y,
3706 size_t n);
3707
3708// Forward with scale: y = a + alpha * b
3709void add_scaled_forward_bf16(const uint16_t *a,
3710 const uint16_t *b,
3711 uint16_t *y,
3712 float alpha,
3713 size_t n);
3714
3715// In-place: a += b
3716void add_inplace_bf16(uint16_t *a,
3717 const uint16_t *b,
3718 size_t n);
3719
3720// In-place scaled: a += alpha * b
3721void add_scaled_inplace_bf16(uint16_t *a,
3722 const uint16_t *b,
3723 float alpha,
3724 size_t n);
3725
3726// Backward: d_a = d_y, d_b = d_y (gradient passthrough)
3727void add_backward_bf16(const uint16_t *d_y,
3728 uint16_t *d_a,
3729 uint16_t *d_b,
3730 size_t n);
3731
3732// 2D version for [tokens, dim] shaped tensors
3733void add_forward_2d_bf16(const uint16_t *a,
3734 const uint16_t *b,
3735 uint16_t *y,
3736 int tokens,
3737 int dim,
3738 int aligned_dim);
3739
3740// FP32 versions
3741void add_forward_f32(const float *a,
3742 const float *b,
3743 float *y,
3744 size_t n);
3745
3746void add_inplace_f32(float *a,
3747 const float *b,
3748 size_t n);
3749
3750// =============================================================================
3751// AXPY kernels (for MoE expert accumulation)
3752// =============================================================================
3753
3754// In-place AXPY: y += alpha * x
3755void axpy_f32(float *y,
3756 const float *x,
3757 float alpha,
3758 int n);
3759
3760// Scaled copy: y = alpha * x
3761void scal_copy_f32(float *y,
3762 const float *x,
3763 float alpha,
3764 int n);
3765
3766// Weighted sum: y = sum_i(weights[i] * vectors[i])
3767void weighted_sum_f32(float *y,
3768 const float **vectors,
3769 const float *weights,
3770 int k,
3771 int n);
3772
3773// Zero-then-accumulate: y = 0; y += alpha * x
3774void axpy_zero_f32(float *y,
3775 const float *x,
3776 float alpha,
3777 int n);
3778
3779// Batched 2D AXPY: Y[t,:] += alpha * X[t,:]
3780void axpy_2d_f32(float *Y,
3781 const float *X,
3782 float alpha,
3783 int num_tokens,
3784 int dim,
3785 int y_stride,
3786 int x_stride);
3787
3788// MoE expert accumulation: output += routing_weight * expert_output
3789void moe_accumulate_expert_f32(float *output,
3790 const float *expert_output,
3791 float routing_weight,
3792 int hidden_dim);
3793
3794// Routed MoE expert MLP: output += weight * down(relu2(up(hidden))).
3795void moe_relu2_expert_forward_f32(const float *hidden,
3796 const int *indices,
3797 const float *routing_weights,
3798 const float *expert_up,
3799 const float *expert_down,
3800 float *output,
3801 int rows,
3802 int hidden_dim,
3803 int intermediate_dim,
3804 int n_experts,
3805 int top_k);
3806
3807void moe_relu2_expert_forward_q5_0_q8_0(const float *hidden,
3808 const int *indices,
3809 const float *routing_weights,
3810 const void *expert_up,
3811 const void *expert_down,
3812 float *output,
3813 int rows,
3814 int hidden_dim,
3815 int intermediate_dim,
3816 int n_experts,
3817 int top_k);
3818
3819void moe_relu2_expert_forward_q5_0_q5_0(const float *hidden,
3820 const int *indices,
3821 const float *routing_weights,
3822 const void *expert_up,
3823 const void *expert_down,
3824 float *output,
3825 int rows,
3826 int hidden_dim,
3827 int intermediate_dim,
3828 int n_experts,
3829 int top_k);
3830
3831void moe_relu2_shared_forward_q5_1_q8_0(const float *hidden,
3832 const float *routed,
3833 const void *shared_up,
3834 const void *shared_down,
3835 float *output,
3836 int rows,
3837 int hidden_dim,
3838 int intermediate_dim);
3839
3840// Routed MoE SwiGLU expert MLP: output += weight * down(silu(gate(hidden)) * up(hidden)).
3841void moe_swiglu_expert_forward_f32(const float *hidden,
3842 const int *indices,
3843 const float *routing_weights,
3844 const float *expert_gate,
3845 const float *expert_up,
3846 const float *expert_down,
3847 float *output,
3848 int rows,
3849 int hidden_dim,
3850 int intermediate_dim,
3851 int n_experts,
3852 int top_k);
3853
3854void moe_swiglu_shared_forward_f32(const float *hidden,
3855 const float *routed,
3856 const float *shared_gate,
3857 const float *shared_up,
3858 const float *shared_down,
3859 float *output,
3860 int rows,
3861 int hidden_dim,
3862 int intermediate_dim);
3863
3864void moe_swiglu_expert_forward_bf16(const float *hidden,
3865 const int *indices,
3866 const float *routing_weights,
3867 const uint16_t *expert_gate,
3868 const uint16_t *expert_up,
3869 const uint16_t *expert_down,
3870 float *output,
3871 int rows,
3872 int hidden_dim,
3873 int intermediate_dim,
3874 int n_experts,
3875 int top_k);
3876
3877void moe_swiglu_expert_forward_bf16_row_range(
3878 const float *hidden, const int *indices, const float *routing_weights,
3879 const uint16_t *expert_gate, const uint16_t *expert_up,
3880 const uint16_t *expert_down, float *output, int rows, int hidden_dim,
3881 int intermediate_dim, int n_experts, int top_k,
3882 int row_begin, int row_end);
3883
3884void moe_swiglu_expert_forward_bf16_parallel_dispatch(
3885 const float *hidden, const int *indices, const float *routing_weights,
3886 const uint16_t *expert_gate, const uint16_t *expert_up,
3887 const uint16_t *expert_down, float *output, int rows, int hidden_dim,
3888 int intermediate_dim, int n_experts, int top_k);
3889
3890void moe_swiglu_packed_expert_forward_bf16(
3891 const float *hidden, const int *indices, const float *routing_weights,
3892 const uint16_t *expert_gate_up, const uint16_t *expert_down,
3893 float *output, int rows, int hidden_dim, int intermediate_dim,
3894 int n_experts, int top_k);
3895
3896size_t moe_swiglu_expert_q4k_q5k_workspace_bytes(int hidden_dim,
3897 int intermediate_dim);
3898size_t moe_swiglu_expert_q4k_q8_0_workspace_bytes(int hidden_dim,
3899 int intermediate_dim);
3900size_t moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes(
3901 int hidden_dim, int intermediate_dim);
3902
3903int moe_swiglu_expert_forward_q4k_q5k_workspace(
3904 const float *hidden,
3905 const int *indices,
3906 const float *routing_weights,
3907 const void *expert_gate,
3908 const void *expert_up,
3909 const void *expert_down,
3910 float *output,
3911 int rows,
3912 int hidden_dim,
3913 int intermediate_dim,
3914 int n_experts,
3915 int top_k,
3916 void *workspace,
3917 size_t workspace_bytes);
3918
3919int moe_swiglu_expert_forward_q4k_q6k_workspace(
3920 const float *hidden,
3921 const int *indices,
3922 const float *routing_weights,
3923 const void *expert_gate,
3924 const void *expert_up,
3925 const void *expert_down,
3926 float *output,
3927 int rows,
3928 int hidden_dim,
3929 int intermediate_dim,
3930 int n_experts,
3931 int top_k,
3932 void *workspace,
3933 size_t workspace_bytes);
3934
3935int moe_swiglu_expert_forward_q4k_q5_0_workspace(
3936 const float *hidden,
3937 const int *indices,
3938 const float *routing_weights,
3939 const void *expert_gate,
3940 const void *expert_up,
3941 const void *expert_down,
3942 float *output,
3943 int rows,
3944 int hidden_dim,
3945 int intermediate_dim,
3946 int n_experts,
3947 int top_k,
3948 void *workspace,
3949 size_t workspace_bytes);
3950
3951int moe_swiglu_expert_forward_q4k_q4k_workspace(
3952 const float *hidden,
3953 const int *indices,
3954 const float *routing_weights,
3955 const void *expert_gate,
3956 const void *expert_up,
3957 const void *expert_down,
3958 float *output,
3959 int rows,
3960 int hidden_dim,
3961 int intermediate_dim,
3962 int n_experts,
3963 int top_k,
3964 void *workspace,
3965 size_t workspace_bytes);
3966
3967int moe_swiglu_expert_forward_q4k_q6k_parallel_workspace(
3968 const float *hidden,
3969 const int *indices,
3970 const float *routing_weights,
3971 const void *expert_gate,
3972 const void *expert_up,
3973 const void *expert_down,
3974 float *output,
3975 int rows,
3976 int hidden_dim,
3977 int intermediate_dim,
3978 int n_experts,
3979 int top_k,
3980 void *workspace,
3981 size_t workspace_bytes);
3982
3983int moe_swiglu_expert_forward_q4k_q4k_parallel_workspace(
3984 const float *hidden,
3985 const int *indices,
3986 const float *routing_weights,
3987 const void *expert_gate,
3988 const void *expert_up,
3989 const void *expert_down,
3990 float *output,
3991 int rows,
3992 int hidden_dim,
3993 int intermediate_dim,
3994 int n_experts,
3995 int top_k,
3996 void *workspace,
3997 size_t workspace_bytes);
3998
3999int moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace(
4000 const float *hidden,
4001 const int *indices,
4002 const float *routing_weights,
4003 const void *expert_gate,
4004 const void *expert_up,
4005 const void *expert_down,
4006 float *output,
4007 int rows,
4008 int hidden_dim,
4009 int intermediate_dim,
4010 int n_experts,
4011 int top_k,
4012 void *workspace,
4013 size_t workspace_bytes);
4014
4015int moe_swiglu_expert_forward_q4k_q8_0_workspace(
4016 const float *hidden, const int *indices, const float *routing_weights,
4017 const void *expert_gate, const void *expert_up, const void *expert_down,
4018 float *output, int rows, int hidden_dim, int intermediate_dim,
4019 int n_experts, int top_k, void *workspace, size_t workspace_bytes);
4020
4021int moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace(
4022 const float *hidden, const int *indices, const float *routing_weights,
4023 const void *expert_gate, const void *expert_up, const void *expert_down,
4024 float *output, int rows, int hidden_dim, int intermediate_dim,
4025 int n_experts, int top_k, void *workspace, size_t workspace_bytes);
4026
4027int moe_swiglu_shared_forward_q4k_q6k_workspace(
4028 const float *hidden,
4029 const float *routed,
4030 const void *shared_gate,
4031 const void *shared_up,
4032 const void *shared_down,
4033 float *output,
4034 int rows,
4035 int hidden_dim,
4036 int intermediate_dim,
4037 void *workspace,
4038 size_t workspace_bytes);
4039
4040int moe_swiglu_shared_forward_q4k_q4k_workspace(
4041 const float *hidden,
4042 const float *routed,
4043 const void *shared_gate,
4044 const void *shared_up,
4045 const void *shared_down,
4046 float *output,
4047 int rows,
4048 int hidden_dim,
4049 int intermediate_dim,
4050 void *workspace,
4051 size_t workspace_bytes);
4052
4053int moe_swiglu_shared_forward_q4k_q6k_parallel_workspace(
4054 const float *hidden,
4055 const float *routed,
4056 const void *shared_gate,
4057 const void *shared_up,
4058 const void *shared_down,
4059 float *output,
4060 int rows,
4061 int hidden_dim,
4062 int intermediate_dim,
4063 void *workspace,
4064 size_t workspace_bytes);
4065
4066int moe_swiglu_shared_forward_q4k_q4k_parallel_workspace(
4067 const float *hidden,
4068 const float *routed,
4069 const void *shared_gate,
4070 const void *shared_up,
4071 const void *shared_down,
4072 float *output,
4073 int rows,
4074 int hidden_dim,
4075 int intermediate_dim,
4076 void *workspace,
4077 size_t workspace_bytes);
4078
4079int moe_swiglu_expert_forward_q4k_q5k_parallel_workspace(
4080 const float *hidden,
4081 const int *indices,
4082 const float *routing_weights,
4083 const void *expert_gate,
4084 const void *expert_up,
4085 const void *expert_down,
4086 float *output,
4087 int rows,
4088 int hidden_dim,
4089 int intermediate_dim,
4090 int n_experts,
4091 int top_k,
4092 void *workspace,
4093 size_t workspace_bytes);
4094
4095int moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace(
4096 const float *hidden,
4097 const int *indices,
4098 const float *routing_weights,
4099 const void *expert_gate,
4100 const void *expert_up,
4101 const void *expert_down,
4102 float *output,
4103 int rows,
4104 int hidden_dim,
4105 int intermediate_dim,
4106 int n_experts,
4107 int top_k,
4108 void *workspace,
4109 size_t workspace_bytes);
4110
4111int moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace(
4112 const float *hidden,
4113 const int *indices,
4114 const float *routing_weights,
4115 const void *expert_gate,
4116 const void *expert_up,
4117 const void *expert_down,
4118 const void *expert_gate_packed,
4119 const void *expert_up_packed,
4120 float *output,
4121 int rows,
4122 int hidden_dim,
4123 int intermediate_dim,
4124 int n_experts,
4125 int top_k,
4126 void *workspace,
4127 size_t workspace_bytes);
4128
4129int moe_swiglu_expert_forward_q4k_q5k_auto_workspace(
4130 const float *hidden,
4131 const int *indices,
4132 const float *routing_weights,
4133 const void *expert_gate,
4134 const void *expert_up,
4135 const void *expert_down,
4136 float *output,
4137 int rows,
4138 int hidden_dim,
4139 int intermediate_dim,
4140 int n_experts,
4141 int top_k,
4142 void *workspace,
4143 size_t workspace_bytes);
4144
4145int moe_swiglu_expert_forward_q4k_q5k_auto_prepared_workspace(
4146 const float *hidden,
4147 const int *indices,
4148 const float *routing_weights,
4149 const void *expert_gate,
4150 const void *expert_up,
4151 const void *expert_down,
4152 float *output,
4153 int rows,
4154 int hidden_dim,
4155 int intermediate_dim,
4156 int n_experts,
4157 int top_k,
4158 void *workspace,
4159 size_t workspace_bytes);
4160
4161size_t moe_swiglu_shared_q8_0_gated_workspace_bytes(int hidden_dim,
4162 int intermediate_dim);
4163
4164int moe_swiglu_shared_forward_q8_0_gated_workspace(
4165 const float *hidden,
4166 const float *routed,
4167 const void *shared_gate,
4168 const void *shared_up,
4169 const void *shared_down,
4170 const float *shared_gate_input,
4171 float *output,
4172 int rows,
4173 int hidden_dim,
4174 int intermediate_dim,
4175 void *workspace,
4176 size_t workspace_bytes);
4177
4178int moe_swiglu_shared_forward_q8_0_gated_parallel_workspace(
4179 const float *hidden,
4180 const float *routed,
4181 const void *shared_gate,
4182 const void *shared_up,
4183 const void *shared_down,
4184 const float *shared_gate_input,
4185 float *output,
4186 int rows,
4187 int hidden_dim,
4188 int intermediate_dim,
4189 void *workspace,
4190 size_t workspace_bytes);
4191
4192int moe_swiglu_shared_forward_q4k_q8_0_gated_workspace(
4193 const float *hidden, const float *routed, const void *shared_gate,
4194 const void *shared_up, const void *shared_down,
4195 const float *shared_gate_input, float *output, int rows, int hidden_dim,
4196 int intermediate_dim, void *workspace, size_t workspace_bytes);
4197
4198int moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace(
4199 const float *hidden, const float *routed, const void *shared_gate,
4200 const void *shared_up, const void *shared_down,
4201 const float *shared_gate_input, float *output, int rows, int hidden_dim,
4202 int intermediate_dim, void *workspace, size_t workspace_bytes);
4203
4204int moe_swiglu_shared_forward_q4k_q5_0_gated_workspace(
4205 const float *hidden, const float *routed, const void *shared_gate,
4206 const void *shared_up, const void *shared_down,
4207 const float *shared_gate_input, float *output, int rows, int hidden_dim,
4208 int intermediate_dim, void *workspace, size_t workspace_bytes);
4209
4210int moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace(
4211 const float *hidden, const float *routed, const void *shared_gate,
4212 const void *shared_up, const void *shared_down,
4213 const float *shared_gate_input, float *output, int rows, int hidden_dim,
4214 int intermediate_dim, void *workspace, size_t workspace_bytes);
4215
4216void moe_swiglu_shared_forward_bf16(const float *hidden,
4217 const float *routed,
4218 const uint16_t *shared_gate,
4219 const uint16_t *shared_up,
4220 const uint16_t *shared_down,
4221 float *output,
4222 int rows,
4223 int hidden_dim,
4224 int intermediate_dim);
4225void moe_swiglu_shared_forward_bf16_row_range(
4226 const float *hidden, const float *routed, const uint16_t *shared_gate,
4227 const uint16_t *shared_up, const uint16_t *shared_down, float *output,
4228 int rows, int hidden_dim, int intermediate_dim,
4229 int row_begin, int row_end);
4230void moe_swiglu_shared_forward_bf16_parallel_dispatch(
4231 const float *hidden, const float *routed, const uint16_t *shared_gate,
4232 const uint16_t *shared_up, const uint16_t *shared_down, float *output,
4233 int rows, int hidden_dim, int intermediate_dim);
4234void moe_swiglu_shared_forward_bf16_gated(
4235 const float *hidden, const float *routed, const uint16_t *shared_gate,
4236 const uint16_t *shared_up, const uint16_t *shared_down,
4237 const uint16_t *shared_router, float *output, int rows,
4238 int hidden_dim, int intermediate_dim);
4239void moe_swiglu_shared_forward_bf16_gated_row_range(
4240 const float *hidden, const float *routed, const uint16_t *shared_gate,
4241 const uint16_t *shared_up, const uint16_t *shared_down,
4242 const uint16_t *shared_router, float *output, int rows,
4243 int hidden_dim, int intermediate_dim, int row_begin, int row_end);
4244void moe_swiglu_shared_forward_bf16_gated_parallel_dispatch(
4245 const float *hidden, const float *routed, const uint16_t *shared_gate,
4246 const uint16_t *shared_up, const uint16_t *shared_down,
4247 const uint16_t *shared_router, float *output, int rows,
4248 int hidden_dim, int intermediate_dim);
4249void farskip_swiglu_shared_combine_bf16(const float *hidden,
4250 const float *routed,
4251 const float *post_attn_residual,
4252 const uint16_t *shared_gate,
4253 const uint16_t *shared_up,
4254 const uint16_t *shared_down,
4255 float *main_output,
4256 float *routed_free_output,
4257 int rows,
4258 int hidden_dim,
4259 int intermediate_dim);
4260void farskip_swiglu_shared_combine_bf16_row_range(
4261 const float *hidden, const float *routed,
4262 const float *post_attn_residual, const uint16_t *shared_gate,
4263 const uint16_t *shared_up, const uint16_t *shared_down,
4264 float *main_output, float *routed_free_output, int rows,
4265 int hidden_dim, int intermediate_dim, int row_begin, int row_end);
4266void farskip_swiglu_shared_combine_bf16_parallel_dispatch(
4267 const float *hidden, const float *routed,
4268 const float *post_attn_residual, const uint16_t *shared_gate,
4269 const uint16_t *shared_up, const uint16_t *shared_down,
4270 float *main_output, float *routed_free_output, int rows,
4271 int hidden_dim, int intermediate_dim);
4272
4273void group_limited_topk_router_sigmoid_f32(const float *logits,
4274 const float *correction_bias,
4275 int *indices,
4276 float *weights,
4277 int rows,
4278 int n_experts,
4279 int top_k,
4280 int n_group,
4281 int topk_group,
4282 int norm_topk_prob,
4283 float routed_scaling_factor);
4284
4285void moe_relu2_expert_backward_f32(const float *d_output,
4286 const float *hidden,
4287 const int *indices,
4288 const float *routing_weights,
4289 const float *expert_up,
4290 const float *expert_down,
4291 float *d_hidden,
4292 float *d_routing_weights,
4293 float *d_expert_up,
4294 float *d_expert_down,
4295 int rows,
4296 int hidden_dim,
4297 int intermediate_dim,
4298 int n_experts,
4299 int top_k);
4300
4301// =============================================================================
4302// Top-K selection kernels (for MoE router dispatch)
4303// =============================================================================
4304
4305// Find top-K indices and values from scores
4306void topk_f32(const float *scores,
4307 int n,
4308 int k,
4309 int *indices,
4310 float *values);
4311
4312// Top-K with softmax-normalized weights
4313void topk_softmax_f32(const float *scores,
4314 int n,
4315 int k,
4316 int *indices,
4317 float *weights);
4318
4319size_t moe_softmax_topk_router_workspace_bytes(int n_experts);
4320
4321int moe_softmax_topk_router_llama_f32_workspace(
4322 const float *logits,
4323 int *indices,
4324 float *weights,
4325 int rows,
4326 int n_experts,
4327 int top_k,
4328 float routed_scaling_factor,
4329 void *workspace,
4330 size_t workspace_bytes);
4331int moe_softmax_topk_router_pytorch_bf16_workspace(
4332 const float *logits,
4333 int *indices,
4334 float *weights,
4335 int rows,
4336 int n_experts,
4337 int top_k,
4338 float routed_scaling_factor,
4339 void *workspace,
4340 size_t workspace_bytes);
4341
4342// Backward for hard top-k followed by softmax over selected values.
4343void topk_softmax_backward_f32(const int *indices,
4344 const float *weights,
4345 const float *d_weights,
4346 float *d_scores,
4347 int num_tokens,
4348 int n_experts_or_keys,
4349 int k);
4350
4351// Batched top-K for multiple tokens
4352void topk_batched_f32(const float *scores,
4353 int num_tokens,
4354 int n_experts,
4355 int k,
4356 int *indices,
4357 float *weights);
4358
4359// Nemotron-H/DeepSeek-style group-limited MoE router over sigmoid scores.
4360void nemotron_group_limited_topk_router_f32(const float *scores,
4361 const float *correction_bias,
4362 int *indices,
4363 float *weights,
4364 int rows,
4365 int n_experts,
4366 int top_k,
4367 int n_group,
4368 int topk_group,
4369 int norm_topk_prob,
4370 float routed_scaling_factor);
4371
4372// Argmax (top-1)
4373int argmax_f32(const float *scores, int n);
4374
4375// Greedy speculative verification for one candidate token.
4376// accepted=1 when draft_token equals argmax(target_logits), otherwise 0.
4377void speculative_verify_greedy_f32(const float *target_logits,
4378 int vocab_size,
4379 int draft_token,
4380 int *accepted,
4381 int *verified_token);
4382
4383// Greedy one-token speculative commit state update.
4384// Appends verified_token when capacity allows, advances target_position, and
4385// keeps draft_position synchronized for the one-token verifier milestone.
4386void speculative_commit_one_i32(int accepted,
4387 int verified_token,
4388 int *token_buffer,
4389 int *token_count,
4390 int max_tokens,
4391 int *target_position,
4392 int *draft_position,
4393 int *accepted_count,
4394 int *rejected_count);
4395
4396// =============================================================================
4397// DeepSeek-style scalar reference kernels
4398// =============================================================================
4399
4400void deepseek_mhc_mix_f32(const float *streams,
4401 const float *mix,
4402 float *out,
4403 int tokens,
4404 int n_streams,
4405 int dim);
4406
4407void deepseek_mhc_mix_backward_f32(const float *d_out,
4408 const float *streams,
4409 const float *mix,
4410 float *d_streams,
4411 float *d_mix,
4412 int tokens,
4413 int n_streams,
4414 int dim);
4415
4416void deepseek_dsa_topk_softmax_f32(const float *scores,
4417 int *indices,
4418 float *weights,
4419 int tokens,
4420 int heads,
4421 int key_count,
4422 int top_k);
4423
4424void deepseek_dsa_topk_softmax_backward_f32(const int *indices,
4425 const float *weights,
4426 const float *d_weights,
4427 float *d_scores,
4428 int tokens,
4429 int heads,
4430 int key_count,
4431 int top_k);
4432
4433void deepseek_csa_attention_f32(const float *q,
4434 const float *k,
4435 const float *v,
4436 const int *indices,
4437 float *out,
4438 float *attn,
4439 int query_tokens,
4440 int key_tokens,
4441 int heads,
4442 int dim,
4443 int top_k,
4444 float scale);
4445
4446int attention_forward_query_key_head_major_f32(
4447 const float *query,
4448 const float *key,
4449 const float *value,
4450 float *output,
4451 float *score_scratch,
4452 int num_heads,
4453 int query_tokens,
4454 int key_tokens,
4455 int head_dim,
4456 float scale);
4457
4458int attention_forward_query_key_head_major_f32_decode_heads(
4459 const float *query,
4460 const float *key,
4461 const float *value,
4462 float *output,
4463 float *score_scratch,
4464 int num_heads,
4465 int query_tokens,
4466 int key_tokens,
4467 int head_dim,
4468 float scale);
4469
4470int attention_forward_query_key_head_major_f32_packed_k(
4471 const float *query,
4472 const float *key,
4473 const float *value,
4474 float *output,
4475 float *score_scratch,
4476 float *key_transpose_scratch,
4477 int num_heads,
4478 int query_tokens,
4479 int key_tokens,
4480 int head_dim,
4481 float scale);
4482
4483int attention_forward_query_key_head_major_tiled_f16kv_fp32(
4484 const float *query,
4485 const float *key,
4486 const float *value,
4487 float *output,
4488 int num_heads,
4489 int query_tokens,
4490 int key_tokens,
4491 int head_dim,
4492 float scale);
4493
4494void deepseek_csa_attention_backward_f32(const float *d_out,
4495 const float *q,
4496 const float *k,
4497 const float *v,
4498 const int *indices,
4499 const float *attn,
4500 float *d_q,
4501 float *d_k,
4502 float *d_v,
4503 int query_tokens,
4504 int key_tokens,
4505 int heads,
4506 int dim,
4507 int top_k,
4508 float scale);
4509
4510void deepseek_hybrid_attention_f32(const float *q,
4511 const float *k,
4512 const float *v,
4513 const int *indices,
4514 float *out,
4515 float *attn,
4516 int query_tokens,
4517 int key_tokens,
4518 int heads,
4519 int dim,
4520 int top_k,
4521 float scale,
4522 int mode);
4523
4524void deepseek_mla_kv_decompress_f32(const float *compressed_kv,
4525 const float *kv_b_proj,
4526 float *k_nope,
4527 float *value,
4528 int tokens,
4529 int heads,
4530 int kv_lora_rank,
4531 int qk_nope_dim,
4532 int v_dim);
4533
4534void deepseek_mla_kv_decompress_bf16(const float *compressed_kv,
4535 const uint16_t *kv_b_proj,
4536 float *k_nope,
4537 float *value,
4538 int tokens,
4539 int heads,
4540 int kv_lora_rank,
4541 int qk_nope_dim,
4542 int v_dim);
4543void deepseek_mla_kv_decompress_bf16_token_range(const float *compressed_kv,
4544 const uint16_t *kv_b_proj,
4545 float *k_nope,
4546 float *value,
4547 int tokens,
4548 int heads,
4549 int kv_lora_rank,
4550 int qk_nope_dim,
4551 int v_dim,
4552 int token_begin,
4553 int token_end);
4554void deepseek_mla_kv_decompress_bf16_parallel_dispatch(
4555 const float *compressed_kv,
4556 const uint16_t *kv_b_proj,
4557 float *k_nope,
4558 float *value,
4559 int tokens,
4560 int heads,
4561 int kv_lora_rank,
4562 int qk_nope_dim,
4563 int v_dim);
4564
4565void deepseek_mla_partial_rope_concat_f32(const float *q_nope,
4566 const float *q_pe,
4567 const float *k_nope,
4568 const float *k_pe,
4569 const float *cos,
4570 const float *sin,
4571 float *query,
4572 float *key,
4573 int tokens,
4574 int heads,
4575 int qk_nope_dim,
4576 int qk_rope_dim);
4577
4578void deepseek_mla_partial_rope_concat_packed_f32(const float *q_packed,
4579 const float *k_nope,
4580 const float *kv_a_packed,
4581 const float *cos,
4582 const float *sin,
4583 float *query,
4584 float *key,
4585 int tokens,
4586 int heads,
4587 int kv_lora_rank,
4588 int qk_nope_dim,
4589 int qk_rope_dim);
4590
4591void deepseek_mla_partial_rope_concat_packed_bf16_storage(
4592 const float *q_packed,
4593 const float *k_nope,
4594 const float *kv_a_packed,
4595 const float *cos,
4596 const float *sin,
4597 float *query,
4598 float *key,
4599 int tokens,
4600 int heads,
4601 int kv_lora_rank,
4602 int qk_nope_dim,
4603 int qk_rope_dim);
4604
4605// Attention backward (GQA-aware): computes d_q, d_k, d_v.
4606void attention_backward_causal_head_major_gqa(
4607 const float *d_output,
4608 const float *q,
4609 const float *k,
4610 const float *v,
4611 const float *attn_weights,
4612 float *d_q,
4613 float *d_k,
4614 float *d_v,
4615 float *d_scores,
4616 int num_heads,
4617 int num_kv_heads,
4618 int num_tokens,
4619 int head_dim,
4620 int aligned_head_dim,
4621 int aligned_context_window);
4622
4623// Attention backward (non-GQA): num_kv_heads == num_heads.
4624void attention_backward_causal_head_major(
4625 const float *d_output,
4626 const float *q,
4627 const float *k,
4628 const float *v,
4629 const float *attn_weights,
4630 float *d_q,
4631 float *d_k,
4632 float *d_v,
4633 float *d_scores,
4634 int num_heads,
4635 int num_tokens,
4636 int head_dim,
4637 int aligned_head_dim,
4638 int aligned_context_window);
4639
4640/* BF16 attention backward - caller provides scratch buffers (no internal malloc) */
4641void attention_backward_causal_head_major_gqa_bf16(
4642 const uint16_t *d_output,
4643 float *d_x,
4644 const uint16_t *q,
4645 const uint16_t *k,
4646 const uint16_t *v,
4647 const float *attn_weights,
4648 float *d_q,
4649 float *d_k,
4650 float *d_v,
4651 float *d_scores,
4652 int num_heads,
4653 int num_kv_heads,
4654 int num_tokens,
4655 int head_dim,
4656 int aligned_head_dim,
4657 int aligned_context_window,
4658 float *scratch_d_output, /* [num_heads * num_tokens * aligned_head_dim] */
4659 float *scratch_q, /* [num_heads * num_tokens * aligned_head_dim] */
4660 float *scratch_k, /* [num_kv_heads * num_tokens * aligned_head_dim] */
4661 float *scratch_v); /* [num_kv_heads * num_tokens * aligned_head_dim] */
4662
4663// RoPE (Rotary Position Embedding) kernels.
4664// Precompute cos/sin cache: [max_seq_len, rotary_dim/2].
4665void rope_precompute_cache_split(float *cos_cache,
4666 float *sin_cache,
4667 int max_seq_len,
4668 int head_dim,
4669 float base);
4670
4671void rope_precompute_cache(float *cos_cache,
4672 float *sin_cache,
4673 int max_seq_len,
4674 int head_dim,
4675 float base,
4676 int rotary_dim,
4677 const char *scaling_type,
4678 float scaling_factor);
4679
4680/* Parameterized YaRN cache with explicit positions. The BF16 variant stores
4681 * the exact RNE BF16 cos/sin values consumed by PyTorch BF16 RoPE. */
4682void yarn_rope_cache_explicit_positions_f32(float *cos_cache,
4683 float *sin_cache,
4684 const int32_t *positions,
4685 int num_tokens,
4686 int rotary_dim,
4687 float freq_base,
4688 float factor,
4689 int original_context,
4690 float beta_fast,
4691 float beta_slow,
4692 float mscale,
4693 float mscale_all_dim);
4694void yarn_rope_cache_contiguous_positions_f32(float *cos_cache,
4695 float *sin_cache,
4696 int num_tokens,
4697 int rotary_dim,
4698 float freq_base,
4699 float factor,
4700 int original_context,
4701 float beta_fast,
4702 float beta_slow,
4703 float mscale,
4704 float mscale_all_dim);
4705void yarn_rope_cache_explicit_positions_bf16(uint16_t *cos_cache,
4706 uint16_t *sin_cache,
4707 const int32_t *positions,
4708 int num_tokens,
4709 int rotary_dim,
4710 float freq_base,
4711 float factor,
4712 int original_context,
4713 float beta_fast,
4714 float beta_slow,
4715 float mscale,
4716 float mscale_all_dim);
4717
4718void rope_precompute_cache_llama_cpu(float *cos_cache,
4719 float *sin_cache,
4720 int max_seq_len,
4721 int head_dim,
4722 float base,
4723 int rotary_dim,
4724 const char *scaling_type,
4725 float scaling_factor);
4726
4727// Apply RoPE forward in-place: x[num_heads, num_tokens, aligned_head_dim].
4728void rope_forward(float *x,
4729 const float *cos_cache,
4730 const float *sin_cache,
4731 int num_heads,
4732 int num_tokens,
4733 int head_dim,
4734 int aligned_head_dim,
4735 int pos_offset);
4736
4737void rope_forward_with_rotary_dim(float *x,
4738 const float *cos_cache,
4739 const float *sin_cache,
4740 int num_heads,
4741 int num_tokens,
4742 int head_dim,
4743 int aligned_head_dim,
4744 int pos_offset,
4745 int rotary_dim);
4746
4747// RoPE backward: inverse rotation.
4748void rope_backward(const float *d_out,
4749 float *d_x,
4750 const float *cos_cache,
4751 const float *sin_cache,
4752 int num_heads,
4753 int num_tokens,
4754 int head_dim,
4755 int aligned_head_dim,
4756 int pos_offset);
4757
4758/* BF16 RoPE forward - caller provides scratch buffer */
4759void rope_forward_bf16(uint16_t *x,
4760 const float *cos_cache,
4761 const float *sin_cache,
4762 int num_heads,
4763 int num_tokens,
4764 int head_dim,
4765 int aligned_head_dim,
4766 int pos_offset,
4767 float *scratch); /* [num_heads * num_tokens * aligned_head_dim] */
4768
4769/* BF16 RoPE backward - caller provides scratch buffers */
4770void rope_backward_bf16(const uint16_t *d_out,
4771 uint16_t *d_x,
4772 const float *cos_cache,
4773 const float *sin_cache,
4774 int num_heads,
4775 int num_tokens,
4776 int head_dim,
4777 int aligned_head_dim,
4778 int pos_offset,
4779 float *scratch_d_out, /* [num_heads * num_tokens * aligned_head_dim] */
4780 float *scratch_d_x); /* [num_heads * num_tokens * aligned_head_dim] */
4781
4782// RoPE backward in-place.
4783void rope_backward_inplace(float *d_x,
4784 const float *cos_cache,
4785 const float *sin_cache,
4786 int num_heads,
4787 int num_tokens,
4788 int head_dim,
4789 int aligned_head_dim,
4790 int pos_offset);
4791
4792void rope_forward_strided(float *x,
4793 const float *cos_cache,
4794 const float *sin_cache,
4795 int num_heads,
4796 int num_tokens,
4797 int head_dim,
4798 int aligned_head_dim,
4799 int pos_offset,
4800 int head_stride_tokens);
4801
4802void rope_forward_strided_with_rotary_dim(float *x,
4803 const float *cos_cache,
4804 const float *sin_cache,
4805 int num_heads,
4806 int num_tokens,
4807 int head_dim,
4808 int aligned_head_dim,
4809 int pos_offset,
4810 int head_stride_tokens,
4811 int rotary_dim);
4812
4813// Combined RoPE for Q and K.
4814void rope_forward_qk(float *q,
4815 float *k,
4816 const float *cos_cache,
4817 const float *sin_cache,
4818 int num_heads,
4819 int num_kv_heads,
4820 int num_tokens,
4821 int head_dim,
4822 int aligned_head_dim,
4823 int pos_offset);
4824
4825void rope_forward_qk_with_rotary_dim(float *q,
4826 float *k,
4827 const float *cos_cache,
4828 const float *sin_cache,
4829 int num_heads,
4830 int num_kv_heads,
4831 int num_tokens,
4832 int head_dim,
4833 int aligned_head_dim,
4834 int pos_offset,
4835 int rotary_dim);
4836
4837void rope_forward_qk_split_llama_token_range_f32(
4838 float *q, float *k, const float *freq_factors, int use_freq_factors,
4839 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
4840 int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base,
4841 int token_begin, int token_end);
4842
4843void rope_forward_qk_split_direct_f32(float *q,
4844 float *k,
4845 const float *freq_factors,
4846 int use_freq_factors,
4847 int num_heads,
4848 int num_kv_heads,
4849 int num_tokens,
4850 int head_dim,
4851 int aligned_head_dim,
4852 int pos_offset,
4853 int rotary_dim,
4854 float freq_base);
4855void rope_forward_qk_split_direct_muse_pytorch_bf16_storage(
4856 float *q, float *k, const float *freq_factors, int use_freq_factors,
4857 int num_heads, int num_kv_heads, int num_tokens, int head_dim,
4858 int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base);
4859void rope_forward_qk_split_direct_token_range_f32(
4860 float *q,
4861 float *k,
4862 const float *freq_factors,
4863 int use_freq_factors,
4864 int num_heads,
4865 int num_kv_heads,
4866 int num_tokens,
4867 int head_dim,
4868 int aligned_head_dim,
4869 int pos_offset,
4870 int rotary_dim,
4871 float freq_base,
4872 int token_begin,
4873 int token_end);
4874void rope_forward_q_split_direct_f32(float *q,
4875 const float *freq_factors,
4876 int use_freq_factors,
4877 int num_heads,
4878 int num_tokens,
4879 int head_dim,
4880 int aligned_head_dim,
4881 int pos_offset,
4882 int rotary_dim,
4883 float freq_base);
4884void rope_forward_q_split_llama_f32(float *q,
4885 const float *freq_factors,
4886 int use_freq_factors,
4887 int num_heads,
4888 int num_tokens,
4889 int head_dim,
4890 int aligned_head_dim,
4891 int pos_offset,
4892 int rotary_dim,
4893 float freq_base);
4894
4895void rope_forward_qk_gemma4_direct(float *q,
4896 float *k,
4897 const float *freq_factors,
4898 int use_freq_factors,
4899 int num_heads,
4900 int num_kv_heads,
4901 int num_tokens,
4902 int head_dim,
4903 int aligned_head_dim,
4904 int pos_offset,
4905 int rotary_dim,
4906 float freq_base);
4907void rope_forward_qk_with_rotary_dim_cache_stride(float *q,
4908 float *k,
4909 const float *cos_cache,
4910 const float *sin_cache,
4911 int num_heads,
4912 int num_kv_heads,
4913 int num_tokens,
4914 int head_dim,
4915 int aligned_head_dim,
4916 int pos_offset,
4917 int rotary_dim,
4918 int cache_rotary_dim);
4919
4920void rope_forward_qk_pairwise_with_rotary_dim(float *q,
4921 float *k,
4922 const float *cos_cache,
4923 const float *sin_cache,
4924 int num_heads,
4925 int num_kv_heads,
4926 int num_tokens,
4927 int head_dim,
4928 int aligned_head_dim,
4929 int pos_offset,
4930 int rotary_dim);
4931
4932void rope_forward_qk_pairwise_llama_cpu(float *q,
4933 float *k,
4934 const float *cos_cache,
4935 const float *sin_cache,
4936 int num_heads,
4937 int num_kv_heads,
4938 int num_tokens,
4939 int head_dim,
4940 int aligned_head_dim,
4941 int pos_offset,
4942 int rotary_dim);
4943
4944void mrope_qk_text(float *q,
4945 float *k,
4946 int num_heads,
4947 int num_kv_heads,
4948 int num_tokens,
4949 int head_dim,
4950 int aligned_head_dim,
4951 int pos_offset,
4952 int n_dims,
4953 int section_0,
4954 int section_1,
4955 int section_2,
4956 int section_3,
4957 int n_ctx_orig,
4958 float freq_base,
4959 float freq_scale,
4960 float ext_factor,
4961 float attn_factor,
4962 float beta_fast,
4963 float beta_slow);
4964
4965void mrope_qk_vision(float *q,
4966 float *k,
4967 const int32_t *positions,
4968 int num_heads,
4969 int num_kv_heads,
4970 int num_tokens,
4971 int head_dim,
4972 int aligned_head_dim,
4973 int n_dims,
4974 int section_0,
4975 int section_1,
4976 int section_2,
4977 int section_3,
4978 int n_ctx_orig,
4979 float freq_base,
4980 float freq_scale,
4981 float ext_factor,
4982 float attn_factor,
4983 float beta_fast,
4984 float beta_slow);
4985
4986
4987void mrope_qk_vision_bf16_storage(float *q, float *k, const int32_t *positions,
4988 int num_heads, int num_kv_heads, int num_tokens,
4989 int head_dim, int aligned_head_dim, int n_dims,
4990 int section_0, int section_1, int section_2, int section_3,
4991 int n_ctx_orig, float freq_base, float freq_scale,
4992 float ext_factor, float attn_factor,
4993 float beta_fast, float beta_slow);
4994void mrope_qk_vision_bf16_pytorch_storage(float *q, float *k, const int32_t *positions,
4995 int num_heads, int num_kv_heads, int num_tokens,
4996 int head_dim, int aligned_head_dim, int n_dims,
4997 int section_0, int section_1, int section_2, int section_3,
4998 int n_ctx_orig, float freq_base, float freq_scale,
4999 float ext_factor, float attn_factor,
5000 float beta_fast, float beta_slow);
5001
5002void mrope_qk_text_imrope_positions_bf16_pytorch_storage(
5003 float *q, float *k, const int32_t *positions,
5004 int num_heads, int num_kv_heads, int num_tokens,
5005 int head_dim, int aligned_head_dim, int n_dims,
5006 int section_0, int section_1, int section_2, int section_3,
5007 int n_ctx_orig, float freq_base, float freq_scale,
5008 float ext_factor, float attn_factor, float beta_fast, float beta_slow);
5009
5010void mrope_qk_vision_fp16_storage(float *q, float *k, const int32_t *positions,
5011 int num_heads, int num_kv_heads, int num_tokens,
5012 int head_dim, int aligned_head_dim, int n_dims,
5013 int section_0, int section_1, int section_2, int section_3,
5014 int n_ctx_orig, float freq_base, float freq_scale,
5015 float ext_factor, float attn_factor,
5016 float beta_fast, float beta_slow);
5017
5018void rope_forward_qk_gemma4v_vision_xy(float *q,
5019 float *k,
5020 int num_heads,
5021 int num_kv_heads,
5022 int num_tokens,
5023 int head_dim,
5024 int aligned_head_dim,
5025 int grid_w,
5026 int rotary_dim,
5027 float freq_base);
5028
5029void mrope_qk_imrope_positions(float *q,
5030 float *k,
5031 const int32_t *positions,
5032 int num_heads,
5033 int num_kv_heads,
5034 int num_tokens,
5035 int head_dim,
5036 int aligned_head_dim,
5037 int n_dims,
5038 int section_0,
5039 int section_1,
5040 int section_2,
5041 int section_3,
5042 int n_ctx_orig,
5043 float freq_base,
5044 float freq_scale,
5045 float ext_factor,
5046 float attn_factor,
5047 float beta_fast,
5048 float beta_slow);
5049
5050void mrope_qk_text_imrope(float *q,
5051 float *k,
5052 int num_heads,
5053 int num_kv_heads,
5054 int num_tokens,
5055 int head_dim,
5056 int aligned_head_dim,
5057 int pos_offset,
5058 int n_dims,
5059 int section_0,
5060 int section_1,
5061 int section_2,
5062 int section_3,
5063 int n_ctx_orig,
5064 float freq_base,
5065 float freq_scale,
5066 float ext_factor,
5067 float attn_factor,
5068 float beta_fast,
5069 float beta_slow);
5070
5071void mrope_qk_text_imrope_bf16_pytorch_storage(float *q,
5072 float *k,
5073 int num_heads,
5074 int num_kv_heads,
5075 int num_tokens,
5076 int head_dim,
5077 int aligned_head_dim,
5078 int pos_offset,
5079 int n_dims,
5080 int section_0,
5081 int section_1,
5082 int section_2,
5083 int section_3,
5084 int n_ctx_orig,
5085 float freq_base,
5086 float freq_scale,
5087 float ext_factor,
5088 float attn_factor,
5089 float beta_fast,
5090 float beta_slow);
5091
5092void rope_forward_qk_strided(float *q,
5093 float *k,
5094 const float *cos_cache,
5095 const float *sin_cache,
5096 int num_heads,
5097 int num_kv_heads,
5098 int num_tokens,
5099 int head_dim,
5100 int aligned_head_dim,
5101 int pos_offset,
5102 int q_stride_tokens,
5103 int k_stride_tokens);
5104
5105void rope_forward_qk_strided_with_rotary_dim(float *q,
5106 float *k,
5107 const float *cos_cache,
5108 const float *sin_cache,
5109 int num_heads,
5110 int num_kv_heads,
5111 int num_tokens,
5112 int head_dim,
5113 int aligned_head_dim,
5114 int pos_offset,
5115 int q_stride_tokens,
5116 int k_stride_tokens,
5117 int rotary_dim);
5118
5119void rope_backward_qk(const float *d_q_out,
5120 const float *d_k_out,
5121 float *d_q,
5122 float *d_k,
5123 const float *cos_cache,
5124 const float *sin_cache,
5125 int num_heads,
5126 int num_kv_heads,
5127 int num_tokens,
5128 int head_dim,
5129 int aligned_head_dim,
5130 int pos_offset);
5131
5132void rope_backward_qk_pairwise_with_rotary_dim(const float *d_q_out,
5133 const float *d_k_out,
5134 float *d_q,
5135 float *d_k,
5136 const float *cos_cache,
5137 const float *sin_cache,
5138 int num_heads,
5139 int num_kv_heads,
5140 int num_tokens,
5141 int head_dim,
5142 int aligned_head_dim,
5143 int pos_offset,
5144 int rotary_dim);
5145
5146/* BF16 RoPE forward for Q and K - caller provides scratch buffers */
5147void rope_forward_qk_bf16(uint16_t *q,
5148 uint16_t *k,
5149 const float *cos_cache,
5150 const float *sin_cache,
5151 int num_heads,
5152 int num_kv_heads,
5153 int num_tokens,
5154 int head_dim,
5155 int aligned_head_dim,
5156 int pos_offset,
5157 float *scratch_q, /* [num_heads * num_tokens * aligned_head_dim] */
5158 float *scratch_k); /* [num_kv_heads * num_tokens * aligned_head_dim] */
5159
5160void rope_forward_qk_bf16_with_rotary_dim(uint16_t *q,
5161 uint16_t *k,
5162 const float *cos_cache,
5163 const float *sin_cache,
5164 int num_heads,
5165 int num_kv_heads,
5166 int num_tokens,
5167 int head_dim,
5168 int aligned_head_dim,
5169 int pos_offset,
5170 int rotary_dim,
5171 float *scratch_q,
5172 float *scratch_k);
5173
5174void rope_forward_bf16_with_rotary_dim(uint16_t *x,
5175 const float *cos_cache,
5176 const float *sin_cache,
5177 int num_heads,
5178 int num_tokens,
5179 int head_dim,
5180 int aligned_head_dim,
5181 int pos_offset,
5182 int rotary_dim,
5183 float *scratch);
5184
5185/* BF16 RoPE backward for Q and K - caller provides scratch buffers */
5186void rope_backward_qk_bf16(const uint16_t *d_q_out,
5187 const uint16_t *d_k_out,
5188 uint16_t *d_q,
5189 uint16_t *d_k,
5190 const float *cos_cache,
5191 const float *sin_cache,
5192 int num_heads,
5193 int num_kv_heads,
5194 int num_tokens,
5195 int head_dim,
5196 int aligned_head_dim,
5197 int pos_offset,
5198 float *scratch_dq_out, /* [num_heads * num_tokens * aligned_head_dim] */
5199 float *scratch_dq, /* [num_heads * num_tokens * aligned_head_dim] */
5200 float *scratch_dk_out, /* [num_kv_heads * num_tokens * aligned_head_dim] */
5201 float *scratch_dk); /* [num_kv_heads * num_tokens * aligned_head_dim] */
5202
5203// Token embedding lookup (optionally adds positional embeddings).
5204// token_embeddings: [vocab_size x aligned_embed_dim]
5205// pos_embeddings: [context_window x aligned_embed_dim] or NULL if add_pos == 0.
5206// output: [context_window x aligned_embed_dim]
5207void embedding_forward(const int32_t *token_ids,
5208 int token_count,
5209 int vocab_size,
5210 const float *token_embeddings,
5211 const float *pos_embeddings,
5212 float *output,
5213 int embed_dim,
5214 int aligned_embed_dim,
5215 int context_window,
5216 int add_pos);
5217
5218void embedding_forward_q4_k(const int32_t *token_ids,
5219 int token_count,
5220 int vocab_size,
5221 const void *token_embeddings,
5222 const float *pos_embeddings,
5223 float *output,
5224 int embed_dim,
5225 int aligned_embed_dim,
5226 int context_window,
5227 int add_pos);
5228
5229void embedding_forward_q5_0(const int32_t *token_ids,
5230 int token_count,
5231 int vocab_size,
5232 const void *token_embeddings,
5233 const float *pos_embeddings,
5234 float *output,
5235 int embed_dim,
5236 int aligned_embed_dim,
5237 int context_window,
5238 int add_pos);
5239
5240void embedding_forward_q6_k(const int32_t *token_ids,
5241 int token_count,
5242 int vocab_size,
5243 const void *token_embeddings,
5244 const float *pos_embeddings,
5245 float *output,
5246 int embed_dim,
5247 int aligned_embed_dim,
5248 int context_window,
5249 int add_pos);
5250
5251void embedding_forward_q8_0(const int32_t *token_ids,
5252 int token_count,
5253 int vocab_size,
5254 const void *token_embeddings,
5255 const float *pos_embeddings,
5256 float *output,
5257 int embed_dim,
5258 int aligned_embed_dim,
5259 int context_window,
5260 int add_pos);
5261
5262 void embedding_forward_bf16(const int32_t *token_ids,
5263 int token_count,
5264 int vocab_size,
5265 const uint16_t *token_embeddings,
5266 const uint16_t *pos_embeddings,
5267 uint16_t *output,
5268 int embed_dim,
5269 int aligned_embed_dim,
5270 int context_window,
5271 int add_pos);
5272
5273void embedding_forward_bf16_fp32(const int32_t *token_ids,
5274 int token_count,
5275 int vocab_size,
5276 const uint16_t *token_embeddings,
5277 const float *pos_embeddings,
5278 float *output,
5279 int embed_dim,
5280 int aligned_embed_dim,
5281 int context_window,
5282 int add_pos);
5283
5284// Embedding backward: accumulates into d_token_embeddings and d_pos_embeddings.
5285// d_output: [context_window x aligned_embed_dim]
5286// d_token_embeddings: [vocab_size x aligned_embed_dim]
5287// d_pos_embeddings: [context_window x aligned_embed_dim] (optional)
5288 void embedding_backward(const int32_t *token_ids,
5289 int token_count,
5290 const float *d_output,
5291 float *d_token_embeddings,
5292 float *d_pos_embeddings,
5293 int vocab_size,
5294 int embed_dim,
5295 int aligned_embed_dim,
5296 int context_window,
5297 int add_pos);
5298
5299 void embedding_backward_bf16(const int32_t *token_ids,
5300 int token_count,
5301 const uint16_t *d_output,
5302 uint16_t *d_token_embeddings,
5303 uint16_t *d_pos_embeddings,
5304 int vocab_size,
5305 int embed_dim,
5306 int aligned_embed_dim,
5307 int context_window,
5308 int add_pos);
5309
5310 void embedding_backward_bf16_mixed(const int32_t *token_ids,
5311 int token_count,
5312 const uint16_t *d_output,
5313 float *d_token_embeddings,
5314 float *d_pos_embeddings,
5315 int vocab_size,
5316 int embed_dim,
5317 int aligned_embed_dim,
5318 int context_window,
5319 int add_pos);
5320
5321// Softmax cross-entropy loss + gradient w.r.t logits.
5322// logits: [tokens x vocab_size], targets: [tokens], d_logits: [tokens x vocab_size]
5323 void softmax_cross_entropy_loss(const float *logits,
5324 const int32_t *targets,
5325 int tokens,
5326 int vocab_size,
5327 float *d_logits,
5328 float *loss_out);
5329 void softmax_cross_entropy_loss_ptref(const float *logits,
5330 const int32_t *targets,
5331 int tokens,
5332 int vocab_size,
5333 float *d_logits,
5334 float *loss_out);
5335
5336 /* BF16 softmax cross-entropy loss - caller provides scratch buffers */
5337 void softmax_cross_entropy_loss_bf16(const uint16_t *logits,
5338 const int32_t *targets,
5339 int tokens,
5340 int vocab_size,
5341 uint16_t *d_logits,
5342 float *loss_out,
5343 float *scratch_logits, /* [tokens * vocab_size] */
5344 float *scratch_d_logits); /* [tokens * vocab_size] */
5345
5346 // Vision helpers (patchify/unpatchify).
5347 void im2patch(const float *image,
5348 float *patches,
5349 int C, int H, int W, int P);
5350 void patch2im(const float *d_patches,
5351 float *d_image,
5352 int C, int H, int W, int P);
5353 void position_embeddings_add(float *x,
5354 const float *position_embd,
5355 int num_tokens,
5356 int embed_dim,
5357 int num_positions);
5358 void position_embeddings_add_at_offset(float *x,
5359 const float *position_embd,
5360 int num_tokens,
5361 int embed_dim,
5362 int num_positions,
5363 int start_position);
5364 void position_embeddings_add_tiled_2d(float *x,
5365 const float *position_embd,
5366 int grid_h,
5367 int grid_w,
5368 int embed_dim,
5369 int merge_size,
5370 int source_grid_size);
5371 void position_embeddings_add_tiled_2d_align_corners(float *x,
5372 const float *position_embd,
5373 int grid_h,
5374 int grid_w,
5375 int embed_dim,
5376 int merge_size,
5377 int source_grid_size);
5378 void position_embeddings_add_tiled_2d_align_corners_bf16(float *x,
5379 const float *position_embd,
5380 int grid_h,
5381 int grid_w,
5382 int embed_dim,
5383 int merge_size,
5384 int source_grid_size);
5385 void position_embeddings_add_tiled_2d_align_corners_fp32_interp_bf16(
5386 float *x,
5387 const float *position_embd,
5388 int grid_h,
5389 int grid_w,
5390 int embed_dim,
5391 int merge_size,
5392 int source_grid_size);
5393 void position_embeddings_add_gemma4v_xy(float *x,
5394 const float *position_embd,
5395 int grid_h,
5396 int grid_w,
5397 int embed_dim,
5398 int source_grid_size);
5399 void vision_position_ids_2d_merge(int32_t *positions,
5400 int grid_h,
5401 int grid_w,
5402 int merge_size);
5403 void rowwise_bias_add(float *x,
5404 const float *bias,
5405 int rows,
5406 int dim);
5407 void add_stream_inplace(float *a,
5408 const float *b,
5409 size_t n);
5410 void add_stream_reorder_2d(float *main_inout,
5411 float *aux_scratch,
5412 int grid_h,
5413 int grid_w,
5414 int embed_dim,
5415 int merge_size);
5416 void spatial_merge_2x2(const float *input,
5417 float *output,
5418 int grid_h,
5419 int grid_w,
5420 int embed_dim);
5421 void spatial_merge_contiguous_tiled(const float *input,
5422 float *output,
5423 int grid_h,
5424 int grid_w,
5425 int embed_dim,
5426 int merge_size);
5427 void spatial_average_pool_contiguous(const float *input,
5428 float *output,
5429 int grid_h,
5430 int grid_w,
5431 int embed_dim,
5432 int merge_size);
5433 void feature_concat_2way(const float *main_input,
5434 const float *branch_input,
5435 float *output,
5436 int rows,
5437 int main_dim,
5438 int branch_slice_dim,
5439 int num_branch_slices);
5440void gemma4_vision_projector_prep_forward(const float *input,
5441 float *output,
5442 int tokens,
5443 int dim,
5444 float scale,
5445 float eps);
5446
5447int ck_multimodal_prefix_insert_f32(const float *source_rows,
5448 int32_t *token_ids,
5449 float *decoder_rows,
5450 int row_count,
5451 int source_row_stride,
5452 int decoder_row_stride,
5453 int copy_dim,
5454 int start_row,
5455 int decoder_capacity);
5456int ck_multimodal_mrope_positions_2d(int32_t *positions,
5457 int total_tokens,
5458 int prefix_start,
5459 int position_base,
5460 int prefix_tokens,
5461 int grid_x,
5462 int grid_y,
5463 int text_pos);
5464
5465
5466 void im2patch_bf16(const uint16_t *image,
5467 uint16_t *patches,
5468 int C, int H, int W, int P);
5469 void patch2im_bf16(const uint16_t *d_patches,
5470 uint16_t *d_image,
5471 int C, int H, int W, int P);
5472
5473#ifdef __cplusplus
5474} // extern "C"
5475#endif
5476
5477#endif // CKERNEL_ENGINE_H
ckernel_audio.h
CKDataType
CKDataType
Supported data types in C-Kernel-Engine.
Definition ckernel_dtype.h:28
hyper_connection_mix_q4k_q5_0_q4k
void hyper_connection_mix_q4k_q5_0_q4k(const float *hyper_input, const float *norm_weight, const void *mix_down_weight, const void *mix_up_weight, const void *inject_weight, float *mixed_output, float *injection_output, float *normalized_scratch, float *dynamic_scratch, float *mix_scratch, int rows, int streams, int hidden_dim, int dynamic_dim, float eps, int emit_injection)
Definition hyper_connection_kernels.c:350
gemm_nt_q4_0
void gemm_nt_q4_0(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Matrix-matrix multiply: C[M,N] = A[M,K] @ B[N,K]^T + bias.
Definition gemm_kernels_q4_0.c:176
gemma4_per_layer_prepare_forward
void gemma4_per_layer_prepare_forward(float *per_layer_input, const float *hidden, const int32_t *token_ids, const void *per_layer_token_emb, const uint16_t *per_layer_model_proj, const float *per_layer_proj_norm, int tokens, int num_layers, int embed_dim, int per_layer_dim, int vocab_size, float eps)
Definition gemma4_per_layer_embed.c:272
add_stream_reorder_2d
void add_stream_reorder_2d(float *main_inout, float *aux_scratch, int grid_h, int grid_w, int embed_dim, int merge_size)
Definition vision_kernels.c:726
attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes, const int *segment_lengths, int num_segments)
Definition attention_kernels.c:7988
ck_attention_prefill_schedule_t
ck_attention_prefill_schedule_t
Definition ckernel_engine.h:2296
CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES
@ CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES
Definition ckernel_engine.h:2299
CK_ATTN_PREFILL_SCHEDULE_KV_HEADS
@ CK_ATTN_PREFILL_SCHEDULE_KV_HEADS
Definition ckernel_engine.h:2297
CK_ATTN_PREFILL_SCHEDULE_QUERY_HEADS
@ CK_ATTN_PREFILL_SCHEDULE_QUERY_HEADS
Definition ckernel_engine.h:2298
CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES
@ CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES
Definition ckernel_engine.h:2301
CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES
@ CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES
Definition ckernel_engine.h:2300
rope_backward_qk_pairwise_with_rotary_dim
void rope_backward_qk_pairwise_with_rotary_dim(const float *d_q_out, const float *d_k_out, float *d_q, float *d_k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)
Definition rope_kernels.c:2933
ck_attention_pytorch_sdpa_scale_f32
float ck_attention_pytorch_sdpa_scale_f32(int head_dim)
Definition attention_kernels.c:4592
attention_forward_causal_head_major_gqa_bf16
void attention_forward_causal_head_major_gqa_bf16(const uint16_t *q, const uint16_t *k, const uint16_t *v, float *scores, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window, float *scratch_q, float *scratch_k, float *scratch_v)
Definition attention_kernels.c:1651
dequant_q4_0_row
void dequant_q4_0_row(const void *src, float *dst, size_t n_elements)
Dequantize Q4_0 row (multiple blocks)
Definition dequant_kernels.c:63
hyper_connection_mix_bf16
void hyper_connection_mix_bf16(const float *hyper_input, const float *norm_weight, const uint16_t *mix_down_weight, const uint16_t *mix_up_weight, const uint16_t *inject_weight, float *mixed_output, float *injection_output, float *normalized_scratch, float *dynamic_scratch, float *mix_scratch, int rows, int streams, int hidden_dim, int dynamic_dim, float eps, int emit_injection)
Definition hyper_connection_kernels.c:73
gemm_nt_q6_k_q8_k_tile
void gemm_nt_q6_k_q8_k_tile(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K, int m0, int m1, int n0, int n1)
Compute one C[m0:m1, n0:n1] tile for Q8_K activations x Q6_K weights.
Definition gemm_kernels_q6k_q8k.c:1760
attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)
Definition attention_kernels.c:7975
embedding_forward_q6_k
void embedding_forward_q6_k(const int32_t *token_ids, int token_count, int vocab_size, const void *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels.c:241
attention_forward_causal_head_major_gqa_exact
void attention_forward_causal_head_major_gqa_exact(const float *q, const float *k, const float *v, float *scores, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)
Definition attention_kernels.c:1579
axpy_f32
void axpy_f32(float *y, const float *x, float alpha, int n)
In-place AXPY: y += alpha * x.
Definition axpy_kernels.c:105
rmsnorm_forward_int8
void rmsnorm_forward_int8(const int8_t *input, const float *gamma, int8_t *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps, float *scratch_input, float *scratch_output)
Definition rmsnorm_kernels_int8.c:58
position_embeddings_add_gemma4v_xy
void position_embeddings_add_gemma4v_xy(float *x, const float *position_embd, int grid_h, int grid_w, int embed_dim, int source_grid_size)
Definition vision_kernels.c:209
moe_swiglu_expert_q4k_q5k_workspace_bytes
size_t moe_swiglu_expert_q4k_q5k_workspace_bytes(int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:727
gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage
void gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:1601
deepseek_mla_partial_rope_concat_f32
void deepseek_mla_partial_rope_concat_f32(const float *q_nope, const float *q_pe, const float *k_nope, const float *k_pe, const float *cos, const float *sin, float *query, float *key, int tokens, int heads, int qk_nope_dim, int qk_rope_dim)
Definition deepseek_kernels.c:367
softmax_cross_entropy_loss_ptref
void softmax_cross_entropy_loss_ptref(const float *logits, const int32_t *targets, int tokens, int vocab_size, float *d_logits, float *loss_out)
Definition loss_kernels.c:292
attention_forward_decode_head_major_gqa_flash_f16cache
void attention_forward_decode_head_major_gqa_flash_f16cache(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:9506
attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_prefill_schedule_t schedule)
Definition attention_kernels.c:7743
gemm_q6_k
void gemm_q6_k(float *Y, const void *W, const float *X, int M, int N, int K)
Definition gemm_kernels_q6k.c:195
ck_gemm_nt_head_major_q8_0
void ck_gemm_nt_head_major_q8_0(const float *attn_out, const void *wo, const float *bias, float *output, int tokens, int embed_dim, int num_heads, int head_dim)
Output projection from head-major attention (Q8_0 weights)
Definition gemm_head_major_output.c:315
attention_forward_causal_head_major_gqa_flash_strided
void attention_forward_causal_head_major_gqa_flash_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:4496
ck_flash_attn_choose_tile_k
int ck_flash_attn_choose_tile_k(int D_h)
Definition attention_flash_true.c:108
kv_cache_store_bf16
void kv_cache_store_bf16(uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:195
rope_forward_qk_split_llama_token_range_f32
void rope_forward_qk_split_llama_token_range_f32(float *q, float *k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base, int token_begin, int token_end)
Definition rope_kernels.c:1338
ck_memcpy_parallel_dispatch
void * ck_memcpy_parallel_dispatch(void *dst, const void *src, size_t size)
kv_cache_store_f16
void kv_cache_store_f16(uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:161
attention_forward_decode_head_major_shared_kv_gemma4
void attention_forward_decode_head_major_shared_kv_gemma4(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:6345
rope_forward_qk_strided
void rope_forward_qk_strided(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int q_stride_tokens, int k_stride_tokens)
Definition rope_kernels.c:2868
moe_accumulate_expert_f32
void moe_accumulate_expert_f32(float *output, const float *expert_output, float routing_weight, int hidden_dim)
Accumulate expert output: output += routing_weight * expert_output.
Definition axpy_kernels.c:307
attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4
void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
Definition attention_kernels_sliding.c:763
swiglu_forward_exact
void swiglu_forward_exact(const float *input, float *output, int tokens, int dim)
Definition swiglu_kernels.c:652
ck_gemm_nt_f16_ggml_oracle
int ck_gemm_nt_f16_ggml_oracle(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_f16.c:262
mamba2_selective_state_update_decode_f32
void mamba2_selective_state_update_decode_f32(const float *state_in, const float *x, const float *dt, const float *a, const float *b, const float *c, const float *d, float *state_out, float *y, int rows, int num_heads, int head_dim, int state_dim, int num_groups)
Definition mamba2_kernels.c:204
spatial_merge_2x2
void spatial_merge_2x2(const float *input, float *output, int grid_h, int grid_w, int embed_dim)
Definition vision_kernels.c:577
gated_deltanet_autoregressive_backward
void gated_deltanet_autoregressive_backward(const float *d_out, const float *d_state_out, const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, const float *state_out, float *d_q, float *d_k, float *d_v, float *d_g, float *d_beta, float *d_state_in, int num_heads, int state_dim, float norm_eps)
Definition deltanet_kernels.c:2038
qwen4_ple_gate_conv_inject_bf16
void qwen4_ple_gate_conv_inject_bf16(const float *hyper_input, const float *key_projected, const float *value_projected, const float *norm_key_weight, const float *norm_query_weight, const float *norm_conv_weight, const uint16_t *conv_weight, float *hyper_output, float *key_norm_scratch, float *query_norm_scratch, float *gated_scratch, float *conv_norm_scratch, const float *conv_state_in, float *conv_state_out, int rows, int streams, int hidden_dim, int kernel_size, int dilation, float eps)
Definition qwen4_exp_kernels.c:402
recurrent_split_conv_qkv_backward
void recurrent_split_conv_qkv_backward(const float *d_q, const float *d_k, const float *d_v, float *d_packed_qkv, int rows, int q_dim, int k_dim, int v_dim)
Definition recurrent_split_kernels.c:108
rmsnorm_backward_bf16
void rmsnorm_backward_bf16(const uint16_t *d_output, const uint16_t *input, const float *gamma, const float *rstd_cache, uint16_t *d_input, float *d_gamma, int tokens, int d_model, int aligned_embed_dim)
Definition rmsnorm_kernels_bf16.c:113
fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size
size_t fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(int aligned_embed_dim)
Get scratch buffer size for fused_rmsnorm_qkv_prefill_head_major_quant.
Definition prefill_fused_gemm.c:651
rope_forward_qk_gemma4v_vision_xy
void rope_forward_qk_gemma4v_vision_xy(float *q, float *k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int grid_w, int rotary_dim, float freq_base)
Definition rope_kernels.c:1494
recurrent_qk_l2_norm_pytorch_fp32_output
void recurrent_qk_l2_norm_pytorch_fp32_output(float *q, float *k, int rows, int q_dim, int k_dim, int head_dim, float eps)
Definition recurrent_qk_norm_kernels.c:335
rmsnorm_forward_no_weight
void rmsnorm_forward_no_weight(const float *input, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:769
gemm_nt_fp32_exact_parallel_dispatch
void gemm_nt_fp32_exact_parallel_dispatch(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:179
attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided
void attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:3638
attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)
Definition attention_kernels.c:9429
swiglu_forward_pytorch_bf16_storage
void swiglu_forward_pytorch_bf16_storage(const float *input, float *output, int tokens, int dim)
Definition swiglu_kernels.c:861
position_embeddings_add_tiled_2d_align_corners
void position_embeddings_add_tiled_2d_align_corners(float *x, const float *position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)
Definition vision_kernels.c:342
gemv_q5_1
void gemv_q5_1(float *y, const void *W, const float *x, int M, int K)
Auto-dispatch GEMV.
Definition gemm_kernels_q5_1.c:187
rmsnorm_forward_fp64_sum
void rmsnorm_forward_fp64_sum(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:154
qwen4_ple_gate_conv_inject_llama_fp16
void qwen4_ple_gate_conv_inject_llama_fp16(const float *hyper_input, const float *key_projected, const float *value_projected, const float *norm_key_weight, const float *norm_query_weight, const float *norm_conv_weight, const uint16_t *conv_weight, float *hyper_output, float *key_norm_scratch, float *query_norm_scratch, float *gated_scratch, float *conv_norm_scratch, const float *conv_state_in, float *conv_state_out, int rows, int streams, int hidden_dim, int kernel_size, int dilation, float eps)
Definition qwen4_exp_kernels.c:436
gemm_naive_parallel
void gemm_naive_parallel(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:206
qk_norm_forward_qwen4_pytorch_bf16_storage
void qk_norm_forward_qwen4_pytorch_bf16_storage(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:390
rmsnorm_forward_qwen3next_pytorch_bf16_storage
void rmsnorm_forward_qwen3next_pytorch_bf16_storage(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:455
recurrent_norm_sigmoid_gate_llama_avx2_forward
void recurrent_norm_sigmoid_gate_llama_avx2_forward(const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
Definition recurrent_norm_kernels.c:82
recurrent_qk_l2_norm_pytorch_bf16_storage
void recurrent_qk_l2_norm_pytorch_bf16_storage(float *q, float *k, int rows, int q_dim, int k_dim, int expanded_heads, int head_dim, float eps)
Definition recurrent_qk_norm_kernels.c:457
attention_forward_decode_head_major_gqa_f16cache_regular_noscale
ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_regular_noscale(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
Definition attention_kernels.c:6844
rmsnorm_forward_llama_production
void rmsnorm_forward_llama_production(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:213
gemm_nt_f32_llama_production_output_range
void gemm_nt_f32_llama_production_output_range(const float *A, const float *B, const float *bias, float *C, int M, int N, int K, int output_begin, int output_end)
Definition gemm_kernels.c:1143
attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace
void attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, float *rounded_kv, size_t rounded_kv_bytes)
Definition attention_kernels.c:5377
moe_swiglu_shared_forward_bf16
void moe_swiglu_shared_forward_bf16(const float *hidden, const float *routed, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:3229
gemm_nt_q5_k_q8_k_m2
void gemm_nt_q5_k_q8_k_m2(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_k.c:985
moe_relu2_expert_forward_q5_0_q8_0
void moe_relu2_expert_forward_q5_0_q8_0(const float *hidden, const int *indices, const float *routing_weights, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:3499
attention_forward_full_head_major_gqa_flash_strided_bf16_storage
void attention_forward_full_head_major_gqa_flash_strided_bf16_storage(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:4994
gemm_nt_q8_0_q8_0_contract_m2
void gemm_nt_q8_0_q8_0_contract_m2(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q8_0_q8_0_contract.c:418
gemv_fused_q8_0_bias_dispatch
void gemv_fused_q8_0_bias_dispatch(float *y, const void *W, const float *x, const float *bias, int M, int K)
Definition gemv_fused_quant_bias.c:525
gelu_erf_fp64_f32_inplace
void gelu_erf_fp64_f32_inplace(float *data, size_t n)
Definition gelu_kernels.c:567
recurrent_qk_l2_norm_forward
void recurrent_qk_l2_norm_forward(float *q, float *k, int rows, int q_dim, int k_dim, int head_dim, float eps)
Definition recurrent_qk_norm_kernels.c:124
rope_forward_qk_split_direct_muse_pytorch_bf16_storage
void rope_forward_qk_split_direct_muse_pytorch_bf16_storage(float *q, float *k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)
Definition rope_kernels.c:1250
ck_residual_add_token_major_bf16_storage
void ck_residual_add_token_major_bf16_storage(const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)
Definition add_kernels_bf16.c:67
swiglu_forward
void swiglu_forward(const float *input, float *output, int tokens, int dim)
Definition swiglu_kernels.c:283
gemm_bias_silu_fused
void gemm_bias_silu_fused(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_fused_kernels.c:177
gemm_nt_bf16_accurate_storage_row_range
void gemm_nt_bf16_accurate_storage_row_range(const float *A, const void *B, const float *bias, float *C, int M, int N, int K, int row_begin, int row_end)
Definition gemm_kernels_bf16.c:2070
geglu_backward_bf16_mixed
void geglu_backward_bf16_mixed(const uint16_t *x, const uint16_t *d_out, float *d_x, int tokens, int dim)
Definition geglu_kernels.c:154
backward_causal_softmax_head_major_bf16
void backward_causal_softmax_head_major_bf16(uint16_t *d_scores, const uint16_t *weights, int num_heads, int num_tokens, int aligned_context_window, float *scratch_d_scores, float *scratch_weights)
Definition softmax_kernels_bf16.c:53
attention_forward_decode_head_major_gqa_flash_gemma4
void attention_forward_decode_head_major_gqa_flash_gemma4(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:6256
add_scaled_forward_bf16
void add_scaled_forward_bf16(const uint16_t *a, const uint16_t *b, uint16_t *y, float alpha, size_t n)
Definition add_kernels_bf16.c:86
attention_forward_decode_head_major_gqa_flash_f16kv
void attention_forward_decode_head_major_gqa_flash_f16kv(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:6361
mrope_qk_vision
void mrope_qk_vision(float *q, float *k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2624
mamba2_conv1d_f32_parallel_dispatch
void mamba2_conv1d_f32_parallel_dispatch(const float *state_in, const float *x, const float *weight, const float *bias, float *conv_out, float *state_out, int rows, int conv_dim, int kernel_size)
geglu_forward_exact
void geglu_forward_exact(const float *x, float *out, int tokens, int dim)
Definition geglu_kernels.c:42
qk_norm_forward_decode_exact
void qk_norm_forward_decode_exact(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
gemm_swiglu_fused
void gemm_swiglu_fused(const float *x, const float *W_gate, const float *W_up, const float *b_gate, const float *b_up, float *output, int M, int N, int K)
Definition gemm_fused_kernels.c:241
split_q_gate_backward
void split_q_gate_backward(const float *d_q, const float *d_gate, float *d_packed_qg, int rows, int q_dim, int gate_dim, int group_dim)
Definition hybrid_attention_kernels.c:80
ck_set_num_threads
void ck_set_num_threads(int num_threads)
Definition ckernel_strict.c:280
attention_forward_query_key_head_major_f32_decode_heads
int attention_forward_query_key_head_major_f32_decode_heads(const float *query, const float *key, const float *value, float *output, float *score_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)
Definition attention_kernels.c:260
attention_forward_decode_head_major_gqa_flash_sliding_gemma4
void attention_forward_decode_head_major_gqa_flash_sliding_gemma4(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
Definition attention_kernels_sliding.c:903
attn_gate_sigmoid_mul_pytorch_bf16_storage
void attn_gate_sigmoid_mul_pytorch_bf16_storage(const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
Definition hybrid_attention_kernels.c:199
moe_swiglu_expert_forward_q4k_q8_0_workspace
int moe_swiglu_expert_forward_q4k_q8_0_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:946
swiglu_backward
void swiglu_backward(const float *input, const float *d_output, float *d_input, int tokens, int dim)
Definition swiglu_kernels.c:432
attention_flash_decode
void attention_flash_decode(float *out, const float *q, const float *k, const float *v, int T_q, int T_k, int H, int D_h, float scale)
Main flash attention function with SIMD dispatch.
Definition attention_flash_true.c:696
moe_swiglu_shared_forward_f32
void moe_swiglu_shared_forward_f32(const float *hidden, const float *routed, const float *shared_gate, const float *shared_up, const float *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:3140
mrope_qk_imrope_positions
void mrope_qk_imrope_positions(float *q, float *k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2776
hyper_stream_expand_f32
void hyper_stream_expand_f32(const float *input, float *output, int rows, int streams, int hidden_dim)
Definition hyper_connection_kernels.c:34
gelu_ggml_inplace
void gelu_ggml_inplace(float *data, size_t n)
Definition gelu_kernels.c:805
gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace
void gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace(const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *input_bf16, size_t input_bf16_bytes, uint16_t *output_bf16, size_t output_bf16_bytes, uint16_t *bias_bf16, size_t bias_bf16_bytes)
Definition gemm_kernels_bf16.c:1614
gelu_backward_exact
void gelu_backward_exact(const float *input, const float *d_output, float *d_input, size_t n)
Definition gelu_kernels.c:369
ck_gemm_nt_head_major_q5_0
void ck_gemm_nt_head_major_q5_0(const float *attn_out, const void *wo, const float *bias, float *output, int tokens, int embed_dim, int num_heads, int head_dim)
Output projection from head-major attention (auto-dispatch)
Definition gemm_head_major_output.c:290
axpy_zero_f32
void axpy_zero_f32(float *y, const float *x, float alpha, int n)
Zero output then accumulate: y = 0; y += alpha * x.
Definition axpy_kernels.c:239
moe_swiglu_shared_forward_q8_0_gated_workspace
int moe_swiglu_shared_forward_q8_0_gated_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, const float *shared_gate_input, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2666
fused_mlp_swiglu_prefill
void fused_mlp_swiglu_prefill(const float *x, const float *W_gate, const float *W_up, const float *W_down, float *output, int seq_len, int hidden, int intermediate, float *scratch)
Fused MLP (Gate + Up + SwiGLU + Down) for prefill.
Definition prefill_fused_gemm.c:878
gelu_exact_inplace
void gelu_exact_inplace(float *data, size_t n)
Definition gelu_kernels.c:558
gemm_nt_q4_k
void gemm_nt_q4_k(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q4k.c:708
moe_swiglu_shared_forward_bf16_row_range
void moe_swiglu_shared_forward_bf16_row_range(const float *hidden, const float *routed, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim, int row_begin, int row_end)
Definition axpy_kernels.c:3181
recurrent_norm_gate_backward
void recurrent_norm_gate_backward(const float *d_out, const float *x, const float *gate, const float *weight, float *d_x, float *d_gate, float *d_weight, int rows, int num_heads, int head_dim, float eps)
Definition recurrent_norm_kernels.c:182
recurrent_sigmoid_forward_ggml
void recurrent_sigmoid_forward_ggml(const float *x, float *out, int rows, int dim)
Definition recurrent_gate_kernels.c:543
hyper_stream_inject_bf16
void hyper_stream_inject_bf16(const float *hyper_input, const float *block_output, const float *injection_weight, float *output, int rows, int streams, int hidden_dim)
Definition hyper_connection_kernels.c:398
add_inplace_bf16
void add_inplace_bf16(uint16_t *a, const uint16_t *b, size_t n)
Definition add_kernels_bf16.c:119
attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes, void *gqa_workspace, size_t gqa_workspace_bytes, int route_num_heads, int route_num_kv_heads, int route_head_dim, int route_query_tokens, int route_min_kv_tokens, int route_workers, int route_query_tile_size, int route_concurrent_query_tiles)
Definition attention_kernels.c:7686
gemm_blocked_serial_train_parallel_dispatch
void gemm_blocked_serial_train_parallel_dispatch(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition ck_parallel_train.c:167
mamba2_in_proj_split_f32
void mamba2_in_proj_split_f32(const float *projected, float *gate, float *hidden_bc, float *dt, int rows, int d_mlp, int intermediate_dim, int conv_dim, int num_heads)
Definition mamba2_kernels.c:65
attention_backward_causal_head_major_gqa_bf16
void attention_backward_causal_head_major_gqa_bf16(const uint16_t *d_output, float *d_x, const uint16_t *q, const uint16_t *k, const uint16_t *v, const float *attn_weights, float *d_q, float *d_k, float *d_v, float *d_scores, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window, float *scratch_d_output, float *scratch_q, float *scratch_k, float *scratch_v)
Definition attention_kernels.c:9739
ck_residual_add_backward
void ck_residual_add_backward(const float *d_out, float *d_a, float *d_b, int tokens, int aligned_embed_dim)
Definition ckernel_orchestration.c:151
gemm_nn_parallel
void gemm_nn_parallel(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:398
gemm_nt_f16
void gemm_nt_f16(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
NT GEMM wrapper for FP16 weights with the engine's standard ABI.
Definition gemm_kernels_f16.c:970
deepseek_mla_attention_f32_workspace
void deepseek_mla_attention_f32_workspace(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, size_t scores_bytes)
Definition deepseek_kernels.c:800
fused_mlp_swiglu_decode_tiled
void fused_mlp_swiglu_decode_tiled(const float *x, const float *W_gate, const float *W_up, const float *W_down, const float *b_gate, const float *b_up, const float *b_down, float *output, int D, int Hff)
Definition mlp_fused_decode.c:429
gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage
void gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:1586
rope_forward
void rope_forward(float *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)
Definition rope_kernels.c:796
kv_cache_repack_head_major_inplace
void kv_cache_repack_head_major_inplace(float *buf, int num_heads, int tokens, int cache_capacity, int aligned_head_dim)
Definition kv_cache_kernels.c:49
gelu_pytorch_erf_f32_inplace
void gelu_pytorch_erf_f32_inplace(float *data, size_t n)
Definition gelu_kernels.c:632
fused_mlp_swiglu_decode
void fused_mlp_swiglu_decode(const float *x, const float *W_gate, const float *W_up, const float *W_down, const float *b_gate, const float *b_up, const float *b_down, float *output, int D, int Hff)
Definition mlp_fused_decode.c:154
gemv_q5_1_q8_1_ref
void gemv_q5_1_q8_1_ref(float *y, const void *W, const void *x_q8, int M, int K)
Definition gemm_kernels_q5_1_q8_1.c:250
position_embeddings_add_tiled_2d
void position_embeddings_add_tiled_2d(float *x, const float *position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)
Definition vision_kernels.c:248
gemm_backward_f32_train_parallel_dispatch_v2
void gemm_backward_f32_train_parallel_dispatch_v2(const float *d_output, const float *input, const float *W, float *d_input, float *d_W, float *d_b, int T, int aligned_in, int aligned_out, int num_threads)
Definition ck_parallel_train.c:699
gemma4_v_norm_forward_parallel_dispatch
void gemma4_v_norm_forward_parallel_dispatch(const float *input, float *output, float *rstd_cache, int tokens, int num_kv_heads, int head_dim, float eps)
yarn_rope_cache_explicit_positions_bf16
void yarn_rope_cache_explicit_positions_bf16(uint16_t *cos_cache, uint16_t *sin_cache, const int32_t *positions, int num_tokens, int rotary_dim, float freq_base, float factor, int original_context, float beta_fast, float beta_slow, float mscale, float mscale_all_dim)
Definition rope_kernels.c:535
gemm_bias_relu_fused
void gemm_bias_relu_fused(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_fused_kernels.c:84
gemm_nt_q5_1_q8_1_ref
void gemm_nt_q5_1_q8_1_ref(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_1_q8_1.c:273
moe_swiglu_expert_forward_bf16
void moe_swiglu_expert_forward_bf16(const float *hidden, const int *indices, const float *routing_weights, const uint16_t *expert_gate, const uint16_t *expert_up, const uint16_t *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:563
attention_forward_full_head_major_gqa_exact_strided
void attention_forward_full_head_major_gqa_exact_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5446
layernorm_naive_serial_matched_precision
void layernorm_naive_serial_matched_precision(const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, float eps)
Definition layernorm_kernels.c:764
gemv_q6_k
void gemv_q6_k(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_q6k.c:169
topk_batched_f32
void topk_batched_f32(const float *scores, int num_tokens, int n_experts, int k, int *indices, float *weights)
Batched top-K selection for multiple tokens.
Definition topk_kernels.c:465
moe_relu2_expert_backward_f32
void moe_relu2_expert_backward_f32(const float *d_output, const float *hidden, const int *indices, const float *routing_weights, const float *expert_up, const float *expert_down, float *d_hidden, float *d_routing_weights, float *d_expert_up, float *d_expert_down, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:3730
backward_causal_softmax_head_major
void backward_causal_softmax_head_major(float *d_scores, const float *weights, int num_heads, int num_tokens, int aligned_context_window)
Definition softmax_kernels.c:382
gemv_bf16_bf16_storage
void gemv_bf16_bf16_storage(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_bf16.c:719
moe_swiglu_expert_forward_q4k_q4k_workspace
int moe_swiglu_expert_forward_q4k_q4k_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1108
fused_mlp_swiglu_prefill_w1w2_quant_scratch_size
size_t fused_mlp_swiglu_prefill_w1w2_quant_scratch_size(int aligned_embed_dim, int aligned_intermediate_dim)
Get scratch buffer size for fused_mlp_swiglu_prefill_w1w2_quant.
Definition prefill_fused_gemm.c:1062
attention_forward_causal_head_major
void attention_forward_causal_head_major(const float *q, const float *k, const float *v, float *scores, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)
Definition attention_kernels.c:1355
attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4
void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int visual_start, int visual_tokens)
Definition attention_kernels.c:5290
gemv_q8_0
void gemv_q8_0(float *y, const void *W, const float *x, int M, int K)
Auto-dispatch GEMV for Q8_0 weights based on CPU features.
Definition gemm_kernels_q8_0.c:694
rope_forward_strided_with_rotary_dim
void rope_forward_strided_with_rotary_dim(float *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int head_stride_tokens, int rotary_dim)
Definition rope_kernels.c:852
attention_forward_causal_head_major_gqa
void attention_forward_causal_head_major_gqa(const float *q, const float *k, const float *v, float *scores, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)
Definition attention_kernels.c:1509
rmsnorm_forward_muse_weighted_pytorch_bf16_storage
void rmsnorm_forward_muse_weighted_pytorch_bf16_storage(const float *input, const float *weight, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:485
recurrent_split_qkv_backward
void recurrent_split_qkv_backward(const float *d_q, const float *d_k, const float *d_v, float *d_packed_qkv, int rows, int q_dim, int k_dim, int v_dim)
Definition recurrent_split_kernels.c:77
patch2im
void patch2im(const float *d_patches, float *d_image, int C, int H, int W, int P)
Definition vision_kernels.c:123
fused_rmsnorm_qkv_prefill_head_major
void fused_rmsnorm_qkv_prefill_head_major(const float *x, const float *gamma, const float *Wq, const float *Bq, const float *Wk, const float *Bk, const float *Wv, const float *Bv, float *Q, float *K, float *V, int seq_len, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int kv_stride_tokens, float eps, float *scratch)
Fused RMSNorm + QKV projection for prefill (head-major outputs)
Definition prefill_fused_gemm.c:441
ck_multimodal_prefix_insert_f32
int ck_multimodal_prefix_insert_f32(const float *source_rows, int32_t *token_ids, float *decoder_rows, int row_count, int source_row_stride, int decoder_row_stride, int copy_dim, int start_row, int decoder_capacity)
Definition embedding_kernels.c:342
attention_forward_decode_head_major_gqa_f16cache_gemma4_contract
ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_gemma4_contract(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)
Definition attention_kernels.c:6910
swiglu_forward_bf16
void swiglu_forward_bf16(const uint16_t *input, uint16_t *output, int tokens, int dim)
Definition swiglu_kernels_bf16.c:66
ck_gemm_bf16_amx_available
int ck_gemm_bf16_amx_available(void)
Definition gemm_kernels_bf16.c:1203
qk_norm_forward_parallel_dispatch
void qk_norm_forward_parallel_dispatch(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
moe_swiglu_shared_forward_q4k_q6k_workspace
int moe_swiglu_shared_forward_q4k_q6k_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1185
rope_backward_bf16
void rope_backward_bf16(const uint16_t *d_out, uint16_t *d_x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float *scratch_d_out, float *scratch_d_x)
Definition rope_kernels_bf16.c:67
attention_forward_full_head_major_gqa_flash_strided
void attention_forward_full_head_major_gqa_flash_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:4541
rope_backward_qk
void rope_backward_qk(const float *d_q_out, const float *d_k_out, float *d_q, float *d_k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)
Definition rope_kernels.c:2916
relu_backward
void relu_backward(const float *input, const float *d_output, float *d_input, size_t n)
Definition relu_kernels.c:84
mlp_token_parallel_bf16
void mlp_token_parallel_bf16(const uint16_t *input, const uint16_t *W_fc1, const uint16_t *b_fc1, const uint16_t *W_fc2, const uint16_t *b_fc2, float *fc1_output, float *output, int T, int aligned_dim, int num_threads, float *scratch_bias1_f, float *scratch_bias2_f, uint16_t *scratch_fc1_bf16)
Definition mlp_kernels_bf16.c:104
attention_forward_causal_head_major_gqa_flash_strided_sliding
void attention_forward_causal_head_major_gqa_flash_strided_sliding(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
Definition attention_kernels_sliding.c:566
geglu_backward_fp32
void geglu_backward_fp32(const float *x, const float *d_out, float *d_x, int tokens, int dim)
Definition geglu_kernels.c:114
ssm_conv1d_forward_llama_production_serial
void ssm_conv1d_forward_llama_production_serial(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:189
layernorm_pytorch_welford_bf16_storage
void layernorm_pytorch_welford_bf16_storage(const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, float eps)
Definition layernorm_kernels.c:952
rope_forward_qk_strided_with_rotary_dim
void rope_forward_qk_strided_with_rotary_dim(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int q_stride_tokens, int k_stride_tokens, int rotary_dim)
Definition rope_kernels.c:2886
embedding_forward_q4_k
void embedding_forward_q4_k(const int32_t *token_ids, int token_count, int vocab_size, const void *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels.c:76
gated_deltanet_llama_avx2_prefill_forward
void gated_deltanet_llama_avx2_prefill_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)
Definition deltanet_kernels.c:1181
deepseek_csa_attention_f32
void deepseek_csa_attention_f32(const float *q, const float *k, const float *v, const int *indices, float *out, float *attn, int query_tokens, int key_tokens, int heads, int dim, int top_k, float scale)
Definition deepseek_kernels.c:562
attention_forward_query_key_head_major_f32_packed_k
int attention_forward_query_key_head_major_f32_packed_k(const float *query, const float *key, const float *value, float *output, float *score_scratch, float *key_transpose_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)
Definition attention_kernels.c:301
topk_softmax_backward_f32
void topk_softmax_backward_f32(const int *indices, const float *weights, const float *d_weights, float *d_scores, int num_tokens, int n_experts_or_keys, int k)
Backward for hard top-k followed by softmax over selected values.
Definition topk_kernels.c:408
moe_swiglu_expert_forward_bf16_parallel_dispatch
void moe_swiglu_expert_forward_bf16_parallel_dispatch(const float *hidden, const int *indices, const float *routing_weights, const uint16_t *expert_gate, const uint16_t *expert_up, const uint16_t *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)
Definition attention_kernels.c:9352
position_embeddings_add_tiled_2d_align_corners_bf16
void position_embeddings_add_tiled_2d_align_corners_bf16(float *x, const float *position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)
Definition vision_kernels.c:397
gemm_nt_q4_1
void gemm_nt_q4_1(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
GEMM with transposed Q4_1 weights: C = A @ B^T.
Definition gemm_kernels_q4_1.c:256
ck_q6_k_q8_k_provider_name
const char * ck_q6_k_q8_k_provider_name(void)
Definition gemm_kernels_q6k_q8k.c:1389
dequant_q5_0_row
void dequant_q5_0_row(const void *src, float *dst, size_t n_elements)
Dequantize Q5_0 row (multiple blocks)
Definition dequant_kernels.c:198
add_forward_f32
void add_forward_f32(const float *a, const float *b, float *y, size_t n)
Definition add_kernels_bf16.c:280
mrope_qk_text
void mrope_qk_text(float *q, float *k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2293
hyper_connection_mix_q6k_q5_0_q4k
void hyper_connection_mix_q6k_q5_0_q4k(const float *hyper_input, const float *norm_weight, const void *mix_down_weight, const void *mix_up_weight, const void *inject_weight, float *mixed_output, float *injection_output, float *normalized_scratch, float *dynamic_scratch, float *mix_scratch, int rows, int streams, int hidden_dim, int dynamic_dim, float eps, int emit_injection)
Definition hyper_connection_kernels.c:374
quantize_batch_q8_k_4row_nearest_even
void quantize_batch_q8_k_4row_nearest_even(const float *x, void *y, int num_rows, int k)
Definition gemm_kernels_q4k_q8k.c:270
rope_precompute_cache_llama_cpu
void rope_precompute_cache_llama_cpu(float *cos_cache, float *sin_cache, int max_seq_len, int head_dim, float base, int rotary_dim, const char *scaling_type, float scaling_factor)
Definition rope_kernels.c:560
rmsnorm_forward_muse_unweighted_pytorch_bf16_storage
void rmsnorm_forward_muse_unweighted_pytorch_bf16_storage(const float *input, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:500
deepseek_mla_attention_bf16_phases_parallel_dispatch
void deepseek_mla_attention_bf16_phases_parallel_dispatch(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, size_t scores_bytes)
Definition deepseek_kernels.c:942
kv_cache_store_shared_q
void kv_cache_store_shared_q(float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict q, int layer, int pos, int num_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:142
gradient_clip_norm_f32
float gradient_clip_norm_f32(float *grad, size_t numel, float max_norm)
Clip gradient norm (fp32)
Definition optimizer_kernels.c:1169
gemma4_per_layer_embed_forward
void gemma4_per_layer_embed_forward(float *hidden, const float *per_layer_input, const float *inp_gate, const float *proj, const float *post_norm, const float *out_scale, int tokens, int layer, int num_layers, int embed_dim, int per_layer_dim, float eps, const int *segment_lengths, int num_segments, int singleton_tail)
Definition gemma4_per_layer_embed.c:417
qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage
void qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage(float *q, float *k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps, float q_scale)
Definition qk_norm_kernels.c:407
rmsnorm_forward_bf16
void rmsnorm_forward_bf16(const uint16_t *input, const float *gamma, uint16_t *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels_bf16.c:24
gemv_q8_0_q8_0_contract
void gemv_q8_0_q8_0_contract(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_q8_0_q8_0_contract.c:385
gemv_q5_0_q8_0
void gemv_q5_0_q8_0(float *y, const void *W, const void *x_q8, int M, int K)
Matrix-vector multiply with Q5_0 weights and Q8_0 input.
Definition gemm_kernels_q5_0.c:1640
recurrent_silu_forward_pytorch_bf16_input_fp32_output
void recurrent_silu_forward_pytorch_bf16_input_fp32_output(const float *x, float *out, int rows, int dim)
Definition recurrent_gate_kernels.c:300
attn_gate_sigmoid_mul_forward
void attn_gate_sigmoid_mul_forward(const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
Definition hybrid_attention_kernels.c:154
moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace
int moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1904
recurrent_dt_gate_expanded_forward
void recurrent_dt_gate_expanded_forward(const float *alpha, const float *dt_bias, const float *a, float *gate, int rows, int num_heads, int state_dim)
Definition recurrent_gate_kernels.c:84
add_inplace_f32
void add_inplace_f32(float *a, const float *b, size_t n)
Definition add_kernels_bf16.c:314
qwen4_qsa_index_select_bf16
void qwen4_qsa_index_select_bf16(const float *projected_qk, const float *index_key_cache_in, const float *q_norm_weight, const float *k_norm_weight, float *selected_indices, float *index_key_cache_out, float *q_norm_scratch, float *pooled_key_scratch, float *block_score_scratch, int32_t *block_index_scratch, int rows, int query_heads, int index_head_dim, int token_budget, int compress_ratio, int rotary_dim, int context_length, int position, float rope_theta, float eps)
Definition qwen4_exp_kernels.c:488
gemv_q5_k_q8_k
void gemv_q5_k_q8_k(float *y, const void *W, const void *x_q8, int M, int K)
Definition gemm_kernels_q5_k.c:892
attention_forward_decode_head_major_gqa_flash_sliding
void attention_forward_decode_head_major_gqa_flash_sliding(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
Definition attention_kernels_sliding.c:817
ck_strict_store_next_gemm_a
void ck_strict_store_next_gemm_a(const float *data, size_t elems)
Definition ckernel_strict.c:48
kv_cache_store_batch_f16
void kv_cache_store_batch_f16(uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:255
recurrent_split_conv_qkv_forward
void recurrent_split_conv_qkv_forward(const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
Definition recurrent_split_kernels.c:97
rope_forward_qk_with_rotary_dim_cache_stride
void rope_forward_qk_with_rotary_dim_cache_stride(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, int cache_rotary_dim)
Definition rope_kernels.c:1509
attention_forward_causal_head_major_gqa_flash_strided_gemma4
void attention_forward_causal_head_major_gqa_flash_strided_gemma4(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5098
attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output
void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
Definition attention_kernels_sliding.c:782
moe_swiglu_expert_forward_q4k_q5_0_workspace
int moe_swiglu_expert_forward_q4k_q5_0_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1027
attention_forward_causal_head_major_gqa_muse_eager_bf16_storage
void attention_forward_causal_head_major_gqa_muse_eager_bf16_storage(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, float *scores, size_t scores_bytes, uint16_t *key_rows, size_t key_rows_bytes, uint16_t *value_columns, size_t value_columns_bytes, uint16_t *gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)
Definition attention_kernels.c:9206
attention_forward_full_head_major_gqa_flash
void attention_forward_full_head_major_gqa_flash(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:4465
relu_forward_inplace_bf16
void relu_forward_inplace_bf16(uint16_t *data, size_t n)
Definition relu_kernels_bf16.c:48
attention_forward_chunk_head_major_gqa_flash_gemma4
void attention_forward_chunk_head_major_gqa_flash_gemma4(const float *q_chunk, const float *k_cache, const float *v_cache, float *out_chunk, int num_heads, int num_kv_heads, int q_tokens, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:6299
recurrent_conv_state_update_forward
void recurrent_conv_state_update_forward(const float *state_in, const float *q, const float *k, const float *v, float *conv_x, float *state_out, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
Definition recurrent_state_kernels.c:100
rope_forward_qk_pairwise_llama_cpu
void rope_forward_qk_pairwise_llama_cpu(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)
Definition rope_kernels.c:1630
deepseek_mla_attention_decode_f32_workspace
void deepseek_mla_attention_decode_f32_workspace(const float *q, const float *k_cache, const float *v_cache, float *output, int num_heads, int num_kv_heads, int cache_len, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride, float scale, float *scores, size_t scores_bytes)
Definition deepseek_kernels.c:1133
gemv_q5_0
void gemv_q5_0(float *y, const void *W, const float *x, int M, int K)
Auto-dispatch GEMV for Q5_0 weights based on CPU features.
Definition gemm_kernels_q5_0.c:550
im2patch_bf16
void im2patch_bf16(const uint16_t *image, uint16_t *patches, int C, int H, int W, int P)
Definition vision_kernels_bf16.c:22
attention_forward_causal_head_major_exact
void attention_forward_causal_head_major_exact(const float *q, const float *k, const float *v, float *scores, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)
Definition attention_kernels.c:1431
position_embeddings_add_tiled_2d_align_corners_fp32_interp_bf16
void position_embeddings_add_tiled_2d_align_corners_fp32_interp_bf16(float *x, const float *position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)
Definition vision_kernels.c:458
gemm_nt_bf16_bf16_storage_parallel_dispatch
void gemm_nt_bf16_bf16_storage_parallel_dispatch(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:2128
layernorm_naive_serial_bf16_storage
void layernorm_naive_serial_bf16_storage(const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, float eps)
Definition layernorm_kernels.c:786
attn_gate_softplus_mul_forward
void attn_gate_softplus_mul_forward(const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
Definition hybrid_attention_kernels.c:228
gemm_microkernel
void gemm_microkernel(const float *A, const float *B, float *C, int M, int N, int K, int B_transposed)
Definition gemm_microkernel.c:1134
moe_swiglu_shared_forward_bf16_parallel_dispatch
void moe_swiglu_shared_forward_bf16_parallel_dispatch(const float *hidden, const float *routed, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim)
attention_forward_query_key_head_major_f32
int attention_forward_query_key_head_major_f32(const float *query, const float *key, const float *value, float *output, float *score_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)
Definition attention_kernels.c:210
mlp_token_parallel
void mlp_token_parallel(const float *input, const float *W_fc1, const float *b_fc1, const float *W_fc2, const float *b_fc2, float *fc1_output, float *output, int T, int aligned_dim, int num_threads)
Definition mlp_kernels.c:49
ck_layout_token_to_head_f32
void ck_layout_token_to_head_f32(const float *src, float *dst, int tokens, int heads, int head_dim)
Definition layout_kernels.c:6
recurrent_conv_state_update_backward_workspace
void recurrent_conv_state_update_backward_workspace(const float *d_conv_x, const float *d_state_out, float *d_state_in, float *d_q, float *d_k, float *d_v, float *d_conv_total, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
Definition recurrent_state_kernels.c:188
deepseek_mhc_mix_f32
void deepseek_mhc_mix_f32(const float *streams, const float *mix, float *out, int tokens, int n_streams, int dim)
Definition deepseek_kernels.c:30
gemm_nt_q8_0_q8_0_m2n4_tile
void gemm_nt_q8_0_q8_0_m2n4_tile(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K, int ldc)
Definition gemm_batch_int8.c:603
group_limited_topk_router_sigmoid_f32
void group_limited_topk_router_sigmoid_f32(const float *logits, const float *correction_bias, int *indices, float *weights, int rows, int n_experts, int top_k, int n_group, int topk_group, int norm_topk_prob, float routed_scaling_factor)
Definition topk_kernels.c:798
kv_cache_store_batch_f32
void kv_cache_store_batch_f32(float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:223
moe_swiglu_expert_forward_q4k_q5k_parallel_workspace
int moe_swiglu_expert_forward_q4k_q5k_parallel_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1470
moe_swiglu_shared_forward_bf16_gated
void moe_swiglu_shared_forward_bf16_gated(const float *hidden, const float *routed, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, const uint16_t *shared_router, float *output, int rows, int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:3322
gemm_nt_q6_k_q8_k_m4_tile
void gemm_nt_q6_k_q8_k_m4_tile(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K, int m0, int m1, int n0, int n1)
Definition gemm_kernels_q6k_q8k.c:1797
moe_softmax_topk_router_llama_f32_workspace
int moe_softmax_topk_router_llama_f32_workspace(const float *logits, int *indices, float *weights, int rows, int n_experts, int top_k, float routed_scaling_factor, void *workspace, size_t workspace_bytes)
Definition topk_kernels.c:316
deepseek_mla_kv_decompress_bf16
void deepseek_mla_kv_decompress_bf16(const float *compressed_kv, const uint16_t *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim)
Definition deepseek_kernels.c:270
softmax_cross_entropy_loss_bf16
void softmax_cross_entropy_loss_bf16(const uint16_t *logits, const int32_t *targets, int tokens, int vocab_size, uint16_t *d_logits, float *loss_out, float *scratch_logits, float *scratch_d_logits)
Definition loss_kernels_bf16.c:25
gemm_nt_q5_1_q8_1
void gemm_nt_q5_1_q8_1(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_1_q8_1.c:320
deepseek_mla_kv_decompress_f32
void deepseek_mla_kv_decompress_f32(const float *compressed_kv, const float *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim)
Definition deepseek_kernels.c:185
moe_relu2_expert_forward_f32
void moe_relu2_expert_forward_f32(const float *hidden, const int *indices, const float *routing_weights, const float *expert_up, const float *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:374
feature_slice_copy
void feature_slice_copy(const float *src, float *dst, int rows, int src_dim, int dst_dim, int dst_feature_offset)
Definition vision_kernels.c:783
rope_backward_qk_bf16
void rope_backward_qk_bf16(const uint16_t *d_q_out, const uint16_t *d_k_out, uint16_t *d_q, uint16_t *d_k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float *scratch_dq_out, float *scratch_dq, float *scratch_dk_out, float *scratch_dk)
Definition rope_kernels_bf16.c:139
gemm_nt_bf16_native_bf16_storage
void gemm_nt_bf16_native_bf16_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:1078
gemm_nt_q4_k_q8_k
void gemm_nt_q4_k_q8_k(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q4k_q8k.c:648
attention_forward_causal_head_major_shared_kv_gemma4
void attention_forward_causal_head_major_shared_kv_gemma4(const float *q, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5144
qk_norm_backward
void qk_norm_backward(const float *d_q_out, const float *d_k_out, const float *q_in, const float *k_in, const float *q_gamma, const float *k_gamma, float *d_q_in, float *d_k_in, float *d_q_gamma, float *d_k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:456
recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output
void recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(const float *x, float *out, int rows, int dim)
Definition recurrent_gate_kernels.c:340
attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided
void attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:3601
attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)
Definition attention_kernels.c:8155
deepseek_mla_kv_decompress_bf16_token_range
void deepseek_mla_kv_decompress_bf16_token_range(const float *compressed_kv, const uint16_t *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim, int token_begin, int token_end)
Definition deepseek_kernels.c:225
ck_layout_head_to_token_f32
void ck_layout_head_to_token_f32(const float *src, float *dst, int heads, int tokens, int head_dim)
Definition layout_kernels.c:19
recurrent_silu_forward_ggml
void recurrent_silu_forward_ggml(const float *x, float *out, int rows, int dim)
Definition recurrent_gate_kernels.c:510
weighted_sum_f32
void weighted_sum_f32(float *y, const float **vectors, const float *weights, int k, int n)
Weighted sum of k vectors: y = sum_i(weights[i] * vectors[i])
Definition axpy_kernels.c:206
spatial_merge_contiguous_tiled
void spatial_merge_contiguous_tiled(const float *input, float *output, int grid_h, int grid_w, int embed_dim, int merge_size)
Definition vision_kernels.c:616
gradient_accumulate_multi_f32
void gradient_accumulate_multi_f32(float *const *dsts, const float *const *srcs, const size_t *numels, int tensor_count)
Definition optimizer_kernels.c:951
rmsnorm_backward_int4
void rmsnorm_backward_int4(const uint8_t *d_output, const uint8_t *input, const float *gamma, const float *rstd_cache, uint8_t *d_input, float *d_gamma, int tokens, int d_model, int aligned_embed_dim, float *scratch_d_output, float *scratch_input, float *scratch_d_input)
Definition rmsnorm_kernels_int4.c:104
ck_gemm_nt_f16_simd_lanes
int ck_gemm_nt_f16_simd_lanes(void)
Definition gemm_kernels_f16.c:583
kv_cache_write_head_major
void kv_cache_write_head_major(const float *__restrict k_token, const float *__restrict v_token, float *__restrict k_cache, float *__restrict v_cache, int num_kv_heads, int token_index, int cache_capacity, int head_dim, int aligned_head_dim)
Definition kv_cache_kernels.c:81
qk_norm_forward_llama_production
void qk_norm_forward_llama_production(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:362
attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes, const int *segment_lengths, int num_segments)
Definition attention_kernels.c:8003
ckernel_backend_native
CKMathBackend ckernel_backend_native(void)
Definition backend_native.c:39
embedding_forward_bf16
void embedding_forward_bf16(const int32_t *token_ids, int token_count, int vocab_size, const uint16_t *token_embeddings, const uint16_t *pos_embeddings, uint16_t *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels_bf16.c:21
gelu_backward_fast_bf16
void gelu_backward_fast_bf16(const uint16_t *input, const uint16_t *d_output, uint16_t *d_input, size_t n, float *scratch_input, float *scratch_d_output, float *scratch_d_input)
Definition gelu_kernels_bf16.c:69
attention_forward_full_head_major_gqa_ggml_strided
void attention_forward_full_head_major_gqa_ggml_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5764
gemm_nt_bf16_accurate_storage
void gemm_nt_bf16_accurate_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:2185
qk_norm_forward_pytorch_bf16_storage
void qk_norm_forward_pytorch_bf16_storage(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:375
gemm_nt_q8_0_q8_0_m2n4
void gemm_nt_q8_0_q8_0_m2n4(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_batch_int8.c:586
causal_softmax_head_major_exact
void causal_softmax_head_major_exact(float *scores, int num_heads, int num_tokens, int aligned_context_window)
Definition softmax_kernels.c:339
recurrent_norm_sigmoid_gate_pytorch_bf16_storage
void recurrent_norm_sigmoid_gate_pytorch_bf16_storage(const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
Definition recurrent_norm_kernels.c:149
gemm_q6_k_q8_k
void gemm_q6_k_q8_k(float *Y, const void *W, const void *X_q8, int M, int N, int K)
GEMM: Y = W @ X^T where W is Q6_K and X is Q8_K.
Definition gemm_kernels_q6k_q8k.c:1508
attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)
Definition attention_kernels.c:9469
gemv_q4_k_q8_k_parallel
void gemv_q4_k_q8_k_parallel(float *y, const void *W, const void *x_q8, int M, int K, int ith, int nth)
Definition gemm_kernels_q4k_q8k.c:464
gemm_nt_q5_0
void gemm_nt_q5_0(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_0.c:833
attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage
void attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5047
attention_forward_decode_head_major_gqa_muse_eager_bf16_storage
void attention_forward_decode_head_major_gqa_muse_eager_bf16_storage(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, float *scores, size_t scores_bytes, uint16_t *key_rows, size_t key_rows_bytes, uint16_t *value_columns, size_t value_columns_bytes, uint16_t *gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)
Definition attention_kernels.c:9242
rope_backward_inplace
void rope_backward_inplace(float *d_x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)
Definition rope_kernels.c:991
qwen4_ple_ngram_embed_bf16
void qwen4_ple_ngram_embed_bf16(const int32_t *token_ids, const uint16_t *embedding, const int64_t *layer_multipliers, const int64_t *head_offsets, const int64_t *head_vocab_sizes, float *output, const float *token_state_in, float *token_state_out, int rows, int ngram_size, int heads_per_ngram, int head_dim, int eos_token_id, int position)
Definition qwen4_exp_kernels.c:121
layernorm_naive_serial
void layernorm_naive_serial(const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition layernorm_kernels.c:175
fc1_backward_kernel
void fc1_backward_kernel(const float *d_output, const float *fc1_input, const float *W_fc1, float *d_input, float *d_W_fc1, float *d_b_fc1, int T, int aligned_in, int aligned_out, int num_threads)
Definition mlp_kernels.c:174
qk_norm_forward
void qk_norm_forward(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:335
attention_forward_full_head_major_gqa_flash_strided_gemma4
void attention_forward_full_head_major_gqa_flash_strided_gemma4(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5158
gemm_nt_q6_k
void gemm_nt_q6_k(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q6k.c:212
ck_strict_consume_next_gemm_a
const float * ck_strict_consume_next_gemm_a(size_t elems)
Definition ckernel_strict.c:70
gemm_nt_bf16_bf16_storage_row_range
void gemm_nt_bf16_bf16_storage_row_range(const float *A, const void *B, const float *bias, float *C, int M, int N, int K, int row_begin, int row_end)
Definition gemm_kernels_bf16.c:2054
swiglu_backward_exact
void swiglu_backward_exact(const float *input, const float *d_output, float *d_input, int tokens, int dim)
Definition swiglu_kernels.c:922
rmsnorm_forward_parallel_dispatch
void rmsnorm_forward_parallel_dispatch(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
qwen4_ple_gate_conv_inject_fp16
void qwen4_ple_gate_conv_inject_fp16(const float *hyper_input, const float *key_projected, const float *value_projected, const float *norm_key_weight, const float *norm_query_weight, const float *norm_conv_weight, const uint16_t *conv_weight, float *hyper_output, float *key_norm_scratch, float *query_norm_scratch, float *gated_scratch, float *conv_norm_scratch, const float *conv_state_in, float *conv_state_out, int rows, int streams, int hidden_dim, int kernel_size, int dilation, float eps)
Definition qwen4_exp_kernels.c:419
swiglu_backward_bf16
void swiglu_backward_bf16(const uint16_t *input, const uint16_t *d_output, uint16_t *d_input, int tokens, int dim)
Definition swiglu_kernels_bf16.c:108
qk_norm_forward_fp64_sum
void qk_norm_forward_fp64_sum(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:351
mlp_token_parallel_bf16_fp32act
void mlp_token_parallel_bf16_fp32act(const uint16_t *input, const uint16_t *W_fc1, const uint16_t *b_fc1, const uint16_t *W_fc2, const uint16_t *b_fc2, float *fc1_output, float *output, int T, int aligned_dim, int num_threads, float *scratch_input_f, float *scratch_bias1_f, float *scratch_bias2_f, uint16_t *scratch_fc1_bf16)
Definition mlp_kernels_bf16.c:199
attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace
void attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window, float *scores, size_t scores_bytes, float *value_columns, size_t value_columns_bytes, float *scaled_scores, size_t scaled_scores_bytes)
Definition attention_kernels.c:6086
mrope_qk_vision_bf16_pytorch_storage
void mrope_qk_vision_bf16_pytorch_storage(float *q, float *k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
moe_swiglu_expert_forward_q4k_q4k_parallel_workspace
int moe_swiglu_expert_forward_q4k_q4k_parallel_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1881
relu2_backward
void relu2_backward(const float *input, const float *d_output, float *d_input, size_t n)
Definition relu_kernels.c:147
attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int query_tile_size, int concurrent_query_tiles, void *workspace, size_t workspace_bytes)
Definition attention_kernels.c:7585
gelu_backward_exact_bf16
void gelu_backward_exact_bf16(const uint16_t *input, const uint16_t *d_output, uint16_t *d_input, size_t n, float *scratch_input, float *scratch_d_output, float *scratch_d_input)
Definition gelu_kernels_bf16.c:46
kv_cache_store_batch_bf16
void kv_cache_store_batch_bf16(uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:291
gemv_q6_k_q8_k
void gemv_q6_k_q8_k(float *y, const void *W, const void *x_q8, int M, int K)
GEMV: y = W @ x where W is Q6_K and x is Q8_K.
Definition gemm_kernels_q6k_q8k.c:1356
embedding_backward_bf16
void embedding_backward_bf16(const int32_t *token_ids, int token_count, const uint16_t *d_output, uint16_t *d_token_embeddings, uint16_t *d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels_bf16.c:123
gemv_q5_k
void gemv_q5_k(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_q5_k.c:1009
moe_swiglu_shared_forward_bf16_gated_row_range
void moe_swiglu_shared_forward_bf16_gated_row_range(const float *hidden, const float *routed, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, const uint16_t *shared_router, float *output, int rows, int hidden_dim, int intermediate_dim, int row_begin, int row_end)
Definition axpy_kernels.c:3245
gemv_q8_0_q8_0_x4
void gemv_q8_0_q8_0_x4(float *y, const void *W, const void *x_q8, int M, int K)
Definition gemm_kernels_q8_0.c:1428
mamba2_conv1d_f32_channel_range
void mamba2_conv1d_f32_channel_range(const float *state_in, const float *x, const float *weight, const float *bias, float *conv_out, float *state_out, int rows, int conv_dim, int kernel_size, int channel_begin, int channel_end)
Definition mamba2_kernels.c:102
embedding_backward
void embedding_backward(const int32_t *token_ids, int token_count, const float *d_output, float *d_token_embeddings, float *d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels.c:296
ck_get_physical_cores
int ck_get_physical_cores(void)
Definition ckernel_strict.c:170
mlp_token_parallel_exact
void mlp_token_parallel_exact(const float *input, const float *W_fc1, const float *b_fc1, const float *W_fc2, const float *b_fc2, float *fc1_output, float *output, int T, int aligned_dim, int num_threads)
Definition mlp_kernels.c:84
gated_deltanet_llama_avx2_forward
void gated_deltanet_llama_avx2_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps)
Definition deltanet_kernels.c:817
geglu_forward_fp32
void geglu_forward_fp32(const float *x, float *out, int tokens, int dim)
Definition geglu_kernels.c:70
deepseek_mla_attention_f32_parallel_dispatch
void deepseek_mla_attention_f32_parallel_dispatch(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, size_t scores_bytes)
Definition deepseek_kernels.c:931
moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace
int moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1927
relu_forward_bf16
void relu_forward_bf16(const uint16_t *input, uint16_t *output, size_t n)
Definition relu_kernels_bf16.c:27
attention_forward_decode_head_major_gqa_flash
void attention_forward_decode_head_major_gqa_flash(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:6212
gemm_nt_bf16_amx_bf16_storage
void gemm_nt_bf16_amx_bf16_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:1318
patch_projection_bf16_pytorch_onednn_conv3d_storage
void patch_projection_bf16_pytorch_onednn_conv3d_storage(const float *input, const void *weights, const float *bias, float *output, int batch, int out_channels, int in_channels, int temporal, int patch_h, int patch_w)
Definition gemm_kernels_bf16.c:1633
gemm_tn_parallel
void gemm_tn_parallel(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:687
rmsnorm_forward_pytorch_bf16_storage
void rmsnorm_forward_pytorch_bf16_storage(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:426
gemm_nt_q8_0
void gemm_nt_q8_0(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q8_0.c:879
relu_forward
void relu_forward(const float *input, float *output, size_t n)
Definition relu_kernels.c:26
im2patch
void im2patch(const float *image, float *patches, int C, int H, int W, int P)
Definition vision_kernels.c:82
fused_rmsnorm_qkv_prefill
void fused_rmsnorm_qkv_prefill(const float *x, const float *gamma, const float *Wq, const float *Wk, const float *Wv, float *Q, float *K, float *V, int seq_len, int hidden, int q_dim, int kv_dim, float eps, float *scratch)
Fused RMSNorm + QKV projection for prefill.
Definition prefill_fused_gemm.c:393
quantize_batch_q8_k
void quantize_batch_q8_k(const float *x, void *y, int num_rows, int k)
Batch quantize FP32 to Q8_K format (row-major output)
Definition gemm_kernels_q8_0.c:283
vec_dot_q6_k_q8_k
void vec_dot_q6_k_q8_k(int n, float *s, const void *vx, const void *vy)
Q6_K x Q8_K dot product (single row)
Definition gemm_kernels_q6k_q8k.c:1339
fc2_backward_kernel
void fc2_backward_kernel(const float *d_output, const float *fc2_input, const float *W_fc2, float *d_input, float *d_W_fc2, float *d_b_fc2, int T, int aligned_in, int aligned_out, int num_threads)
Definition mlp_kernels.c:126
attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes
size_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes(int num_heads, int num_kv_heads, int head_dim, int workers, int query_tile_size, int concurrent_query_tiles)
Definition attention_kernels.c:7064
quantize_row_q8_k
void quantize_row_q8_k(const float *x, void *y, int k)
Definition gemm_kernels_q4k_q8k.c:229
qk_norm_forward_prefill_exact
void qk_norm_forward_prefill_exact(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
axpy_2d_f32
void axpy_2d_f32(float *Y, const float *X, float alpha, int num_tokens, int dim, int y_stride, int x_stride)
Batched AXPY for 2D tensors: Y[t,:] += alpha * X[t,:].
Definition axpy_kernels.c:272
ck_set_strict_parity
void ck_set_strict_parity(int enabled)
Definition ckernel_strict.c:28
rmsnorm_backward_int8
void rmsnorm_backward_int8(const int8_t *d_output, const int8_t *input, const float *gamma, const float *rstd_cache, int8_t *d_input, float *d_gamma, int tokens, int d_model, int aligned_embed_dim, float *scratch_d_output, float *scratch_input, float *scratch_d_input)
Definition rmsnorm_kernels_int8.c:84
recurrent_silu_forward
void recurrent_silu_forward(const float *x, float *out, int rows, int dim)
Definition recurrent_gate_kernels.c:138
mlp_token_parallel_bf16_backward_mixed
void mlp_token_parallel_bf16_backward_mixed(const uint16_t *input, const uint16_t *W_fc1, const uint16_t *b_fc1, const uint16_t *W_fc2, const uint16_t *d_output, float *d_input, float *d_W_fc1, float *d_b_fc1, float *d_W_fc2, float *d_b_fc2, int T, int aligned_dim, int num_threads, float *scratch_fc1_pre, uint16_t *scratch_fc1_act_bf16, float *scratch_d_fc1)
Definition mlp_kernels_bf16.c:268
gemv_bf16_bf16_storage_parallel_dispatch
void gemv_bf16_bf16_storage_parallel_dispatch(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_bf16.c:695
layernorm_backward_kernel
void layernorm_backward_kernel(const float *d_output, const float *input, const float *gamma, const float *mean, const float *rstd, float *d_input, float *d_gamma, float *d_beta, int tokens, int d_model, int aligned_embed_dim)
Definition layernorm_kernels.c:1007
ssm_conv1d_forward_llama_fma
void ssm_conv1d_forward_llama_fma(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:272
moe_swiglu_expert_forward_bf16_row_range
void moe_swiglu_expert_forward_bf16_row_range(const float *hidden, const int *indices, const float *routing_weights, const uint16_t *expert_gate, const uint16_t *expert_up, const uint16_t *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, int row_begin, int row_end)
Definition axpy_kernels.c:505
moe_swiglu_shared_q8_0_gated_workspace_bytes
size_t moe_swiglu_shared_q8_0_gated_workspace_bytes(int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:2649
moe_relu2_shared_forward_q5_1_q8_0
void moe_relu2_shared_forward_q5_1_q8_0(const float *hidden, const float *routed, const void *shared_up, const void *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:3662
sigmoid_forward_bf16
void sigmoid_forward_bf16(const uint16_t *input, uint16_t *output, size_t n, float *scratch_input, float *scratch_output)
Definition sigmoid_kernels_bf16.c:27
gemm_blocked_serial_bf16
void gemm_blocked_serial_bf16(const uint16_t *A, const uint16_t *B, const uint16_t *bias, uint16_t *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:431
sigmoid_backward_bf16
void sigmoid_backward_bf16(const uint16_t *input, const uint16_t *d_output, uint16_t *d_input, size_t n, float *scratch_input, float *scratch_d_output, float *scratch_d_input)
Definition sigmoid_kernels_bf16.c:45
geglu_forward_ggml_native
void geglu_forward_ggml_native(const float *x, float *out, int tokens, int dim)
Definition geglu_kernels.c:17
rope_forward_qk_pairwise_with_rotary_dim
void rope_forward_qk_pairwise_with_rotary_dim(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)
Definition rope_kernels.c:1573
gemm_nn_avx512
void gemm_nn_avx512(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:420
causal_softmax_head_major
void causal_softmax_head_major(float *scores, int num_heads, int num_tokens, int aligned_context_window)
Definition softmax_kernels.c:144
patch_projection_image_bf16_native_storage
void patch_projection_image_bf16_native_storage(const float *image, const void *weights_t0, const void *weights_t1, const float *bias, float *output, int channels, int image_h, int image_w, int patch_size, int out_channels, int merge_size)
Definition gemm_kernels_bf16.c:2009
moe_relu2_expert_forward_q5_0_q5_0
void moe_relu2_expert_forward_q5_0_q5_0(const float *hidden, const int *indices, const float *routing_weights, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:3581
gemm_nt_q6_k_q8_k_tiled
void gemm_nt_q6_k_q8_k_tiled(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Experimental single-thread tiled NT GEMM wrapper.
Definition gemm_kernels_q6k_q8k.c:1870
attention_forward_decode_head_major_gqa_regular
void attention_forward_decode_head_major_gqa_regular(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
WARNING: This is NOT true flash attention!
Definition attention_kernels.c:9629
argmax_f32
int argmax_f32(const float *scores, int n)
Find index of maximum value.
Definition topk_kernels.c:500
gemm_nt_f16_clipped
void gemm_nt_f16_clipped(const float *A, const void *B, const float *bias, const float *input_min, const float *input_max, const float *output_min, const float *output_max, float *C, int M, int N, int K)
Definition gemm_kernels_f16.c:998
assistant_layer_scale_forward
void assistant_layer_scale_forward(float *hidden, const float *scale, int tokens, int embed_dim)
Definition gemma4_per_layer_embed.c:475
attention_forward_causal_head_major_gqa_flash_strided_f16kv
void attention_forward_causal_head_major_gqa_flash_strided_f16kv(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5417
gemv_q5_1_q8_1
void gemv_q5_1_q8_1(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_q5_1_q8_1.c:300
split_qkv_packed_head_major_forward
void split_qkv_packed_head_major_forward(const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim, int num_heads, int num_kv_heads)
Definition recurrent_split_kernels.c:25
ck_gemm_bf16_fp32out_amx_raw
int ck_gemm_bf16_fp32out_amx_raw(const uint16_t *A, const uint16_t *B, float *C, int M, int N, int K, int accumulate)
Definition gemm_kernels_bf16.c:1212
gemma4_vision_projector_prep_forward
void gemma4_vision_projector_prep_forward(const float *input, float *output, int tokens, int dim, float scale, float eps)
Definition vision_kernels.c:633
unfused_rmsnorm_qkv_prefill
void unfused_rmsnorm_qkv_prefill(const float *x, const float *gamma, const float *Wq, const float *Wk, const float *Wv, float *x_norm, float *Q, float *K, float *V, int seq_len, int hidden, int q_dim, int kv_dim, float eps)
Unfused version for benchmarking comparison.
Definition prefill_fused_gemm.c:667
add_forward_bf16
void add_forward_bf16(const uint16_t *a, const uint16_t *b, uint16_t *y, size_t n)
Definition add_kernels_bf16.c:38
embedding_forward_bf16_fp32
void embedding_forward_bf16_fp32(const int32_t *token_ids, int token_count, int vocab_size, const uint16_t *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels_bf16.c:73
add_forward_2d_bf16
void add_forward_2d_bf16(const uint16_t *a, const uint16_t *b, uint16_t *y, int tokens, int dim, int aligned_dim)
Definition add_kernels_bf16.c:231
yarn_rope_cache_contiguous_positions_f32
void yarn_rope_cache_contiguous_positions_f32(float *cos_cache, float *sin_cache, int num_tokens, int rotary_dim, float freq_base, float factor, int original_context, float beta_fast, float beta_slow, float mscale, float mscale_all_dim)
Definition rope_kernels.c:517
gemma4_per_layer_prepare_bf16_forward
void gemma4_per_layer_prepare_bf16_forward(float *per_layer_input, const float *hidden, const int32_t *token_ids, const uint16_t *per_layer_token_emb, const uint16_t *per_layer_model_proj, const float *per_layer_proj_norm, int tokens, int num_layers, int embed_dim, int per_layer_dim, int vocab_size, float eps)
Definition gemma4_per_layer_embed.c:309
rmsnorm_forward
void rmsnorm_forward(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:735
gemma4_final_logit_softcap_forward
void gemma4_final_logit_softcap_forward(float *logits, int tokens, int vocab_size, float cap)
Definition gemma4_per_layer_embed.c:491
feature_concat_2way
void feature_concat_2way(const float *main_input, const float *branch_input, float *output, int rows, int main_dim, int branch_slice_dim, int num_branch_slices)
Definition vision_kernels.c:756
gemm_avx512_parallel
void gemm_avx512_parallel(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:230
gemv_q4_k
void gemv_q4_k(float *y, const void *W, const float *x, int M, int K)
Auto-dispatch GEMV based on available SIMD.
Definition gemm_kernels_q4k.c:301
attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_sliding
void attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_sliding(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, float *scores, size_t scores_bytes, uint16_t *key_rows, size_t key_rows_bytes, uint16_t *value_columns, size_t value_columns_bytes, uint16_t *gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)
Definition attention_kernels.c:9260
farskip_swiglu_shared_combine_bf16_parallel_dispatch
void farskip_swiglu_shared_combine_bf16_parallel_dispatch(const float *hidden, const float *routed, const float *post_attn_residual, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, float *main_output, float *routed_free_output, int rows, int hidden_dim, int intermediate_dim)
gemm_nt_bf16
void gemm_nt_bf16(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:755
gated_deltanet_pytorch_grouped_bf16_forward
void gated_deltanet_pytorch_grouped_bf16_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps)
Definition deltanet_kernels.c:1139
ssm_conv1d_forward_pytorch_bf16_storage
void ssm_conv1d_forward_pytorch_bf16_storage(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:299
gelu_fast_inplace_bf16
void gelu_fast_inplace_bf16(uint16_t *data, size_t n, float *scratch)
Definition gelu_kernels_bf16.c:31
swiglu_forward_q8_k
void swiglu_forward_q8_k(const float *input, void *output_q8, int tokens, int dim)
Definition swiglu_kernels.c:365
dequant_q8_0_row
void dequant_q8_0_row(const void *src, float *dst, size_t n_elements)
Dequantize Q8_0 row (multiple blocks)
Definition dequant_kernels.c:288
gradient_global_norm_multi_f32
float gradient_global_norm_multi_f32(const float *const *grads, const size_t *numels, int tensor_count)
Definition optimizer_kernels.c:1207
rope_forward_bf16
void rope_forward_bf16(uint16_t *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float *scratch)
Definition rope_kernels_bf16.c:28
gemm_nt_bf16_parallel_dispatch
void gemm_nt_bf16_parallel_dispatch(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:784
mrope_qk_text_imrope
void mrope_qk_text_imrope(float *q, float *k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2363
attention_forward_decode_head_major_shared_kv_sliding_gemma4
void attention_forward_decode_head_major_shared_kv_sliding_gemma4(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
Definition attention_kernels_sliding.c:976
gated_deltanet_pytorch_grouped_bf16_prefill_forward
void gated_deltanet_pytorch_grouped_bf16_prefill_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)
Definition deltanet_kernels.c:1293
q_norm_forward
void q_norm_forward(float *q, const float *q_gamma, int num_heads, int num_tokens, int head_dim, float eps)
Definition qk_norm_kernels.c:434
attention_forward_sparse_token_major_gqa_bf16cache_pytorch_cpu_flash_contract
void attention_forward_sparse_token_major_gqa_bf16cache_pytorch_cpu_flash_contract(const float *query, const uint16_t *key_cache, const uint16_t *value_cache, const float *selected_indices, float *output, float *score_scratch, int rows, int query_heads, int kv_heads, int head_dim, int selection_width, int context_length, int position)
Definition attention_kernels.c:8966
gemm_nn_simd
void gemm_nn_simd(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:577
moe_softmax_topk_router_workspace_bytes
size_t moe_softmax_topk_router_workspace_bytes(int n_experts)
Definition topk_kernels.c:177
gemm_nt_q5_k_q8_k
void gemm_nt_q5_k_q8_k(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_k.c:914
rope_forward_qk_split_direct_token_range_f32
void rope_forward_qk_split_direct_token_range_f32(float *q, float *k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base, int token_begin, int token_end)
Definition rope_kernels.c:1274
attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace
void attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int live_tokens, int kv_stride_tokens, int head_dim, int aligned_head_dim, int sliding_window, float *scores, size_t scores_bytes, float *value_columns, size_t value_columns_bytes, float *scaled_scores, size_t scaled_scores_bytes)
Definition attention_kernels.c:6185
attention_forward_query_key_head_major_tiled_f16kv_fp32
int attention_forward_query_key_head_major_tiled_f16kv_fp32(const float *query, const float *key, const float *value, float *output, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)
Definition attention_kernels.c:3655
moe_swiglu_expert_forward_q4k_q5k_workspace
int moe_swiglu_expert_forward_q4k_q5k_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:781
rope_precompute_cache
void rope_precompute_cache(float *cos_cache, float *sin_cache, int max_seq_len, int head_dim, float base, int rotary_dim, const char *scaling_type, float scaling_factor)
Definition rope_kernels.c:384
gemm_nt_q5_1
void gemm_nt_q5_1(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
GEMM with transposed Q5_1 weights: C = A @ B^T.
Definition gemm_kernels_q5_1.c:314
ck_strict_mtmd_clip_encode_planar_f32
int ck_strict_mtmd_clip_encode_planar_f32(const float *planar, int channels, int height, int width, float *out, size_t out_elems)
Definition ckernel_strict.c:84
layernorm_backward_kernel_bf16
void layernorm_backward_kernel_bf16(const uint16_t *d_output, const uint16_t *input, const float *gamma, const float *mean, const float *rstd, uint16_t *d_input, float *d_gamma, float *d_beta, int tokens, int d_model, int aligned_embed_dim, float *scratch_d_output, float *scratch_input, float *scratch_d_input)
Definition layernorm_kernels_bf16.c:84
rope_forward_bf16_with_rotary_dim
void rope_forward_bf16_with_rotary_dim(uint16_t *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float *scratch)
Definition rope_kernels_bf16.c:42
moe_swiglu_shared_forward_q8_0_gated_parallel_workspace
int moe_swiglu_shared_forward_q8_0_gated_parallel_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, const float *shared_gate_input, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2981
gemv_q6_k_q8_k_parallel
void gemv_q6_k_q8_k_parallel(float *y, const void *W, const void *x_q8, int M, int K, int ith, int nth)
Parallel reference GEMV for Q6_K × Q8_K.
Definition gemm_kernels_q6k_q8k.c:1422
recurrent_silu_forward_pytorch_bf16_storage
void recurrent_silu_forward_pytorch_bf16_storage(const float *x, float *out, int rows, int dim)
Definition recurrent_gate_kernels.c:253
sigmoid_scalar
float sigmoid_scalar(float x)
Definition sigmoid_kernels.c:26
deepseek_csa_attention_backward_f32
void deepseek_csa_attention_backward_f32(const float *d_out, const float *q, const float *k, const float *v, const int *indices, const float *attn, float *d_q, float *d_k, float *d_v, int query_tokens, int key_tokens, int heads, int dim, int top_k, float scale)
Definition deepseek_kernels.c:621
moe_swiglu_shared_forward_q4k_q4k_parallel_workspace
int moe_swiglu_shared_forward_q4k_q4k_parallel_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2071
ck_residual_add_token_major_parallel_dispatch
void ck_residual_add_token_major_parallel_dispatch(const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)
attention_forward_causal_head_major_gqa_flash
void attention_forward_causal_head_major_gqa_flash(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim)
Definition attention_kernels.c:4444
ck_attention_flash_decode_wrapper
void ck_attention_flash_decode_wrapper(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
Wrapper to call TRUE flash attention from orchestration layer.
Definition ckernel_orchestration.c:72
deepseek_mla_kv_cache_batch_store_f32
void deepseek_mla_kv_cache_batch_store_f32(float *k_cache, float *v_cache, const float *k, const float *v, int num_tokens, int num_kv_heads, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride)
Definition deepseek_kernels.c:974
mamba2_selective_scan_f32_parallel_dispatch
void mamba2_selective_scan_f32_parallel_dispatch(const float *state_init, const float *x, const float *dt, const float *a, const float *b, const float *c, const float *d, float *state_out, float *y, int batch, int seq_len, int num_heads, int head_dim, int state_dim, int num_groups)
sigmoid_backward
void sigmoid_backward(const float *input, const float *d_output, float *d_input, size_t n)
Definition sigmoid_kernels.c:138
rmsnorm_forward_strided_f32
void rmsnorm_forward_strided_f32(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)
Definition rmsnorm_kernels.c:567
hyper_stream_inject_f32
void hyper_stream_inject_f32(const float *hyper_input, const float *block_output, const float *injection_weight, float *output, int rows, int streams, int hidden_dim)
Definition hyper_connection_kernels.c:429
quantize_row_q8_0
void quantize_row_q8_0(const float *x, void *y, int k)
Quantize FP32 to Q8_0 format (scalar reference)
Definition gemm_kernels_q8_0.c:125
gemv_bf16
void gemv_bf16(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_bf16.c:629
deepseek_mla_attention_bf16_phases_workspace
void deepseek_mla_attention_bf16_phases_workspace(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, size_t scores_bytes)
Definition deepseek_kernels.c:823
gelu_pytorch_erf_sleef_bf16_storage
void gelu_pytorch_erf_sleef_bf16_storage(float *data, size_t n)
Definition gelu_kernels.c:734
gated_deltanet_autoregressive_forward
void gated_deltanet_autoregressive_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)
Definition deltanet_kernels.c:1956
ck_residual_add_token_major
void ck_residual_add_token_major(const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)
Definition ckernel_orchestration.c:139
attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract
ck_attention_status_t attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)
Definition attention_kernels.c:9278
mrope_qk_text_imrope_positions_bf16_pytorch_storage
void mrope_qk_text_imrope_positions_bf16_pytorch_storage(float *q, float *k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2550
rope_forward_qk
void rope_forward_qk(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)
Definition rope_kernels.c:1094
attention_forward_full_head_major_gqa_ggml_strided_workspace
void attention_forward_full_head_major_gqa_ggml_strided_workspace(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, float *score_rows, size_t score_rows_bytes, float *v_columns, size_t v_columns_bytes, float *probability_row, size_t probability_row_bytes)
Definition attention_kernels.c:5523
moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace
int moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2570
attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)
Definition attention_kernels.c:7962
gated_deltanet_llama_chunk64_head_forward
void gated_deltanet_llama_chunk64_head_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int head, int state_dim)
Definition deltanet_kernels.c:1252
recurrent_norm_gate_pytorch_bf16_storage
void recurrent_norm_gate_pytorch_bf16_storage(const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
Definition recurrent_norm_kernels.c:111
fused_rmsnorm_qkv_scratch_size
size_t fused_rmsnorm_qkv_scratch_size(int hidden)
Get scratch buffer size for fused_rmsnorm_qkv_prefill.
Definition prefill_fused_gemm.c:739
gelu_ggml_native_inplace
void gelu_ggml_native_inplace(float *data, size_t n)
Definition gelu_kernels.c:793
gemv_q8_0_q8_0
void gemv_q8_0_q8_0(float *y, const void *W, const void *x_q8, int M, int K)
Matrix-vector multiply with Q8_0 weights and Q8_0 input.
Definition gemm_kernels_q8_0.c:1405
adamw_clip_update_multi_f32
void adamw_clip_update_multi_f32(float *const *grads, float *const *weights, float *const *m_states, float *const *v_states, const size_t *numels, int tensor_count, float lr, float beta1, float beta2, float eps, float weight_decay, float max_grad_norm, int step)
Definition optimizer_kernels.c:638
attn_gate_sigmoid_mul_backward
void attn_gate_sigmoid_mul_backward(const float *d_out, const float *x, const float *gate, float *d_x, float *d_gate, int rows, int num_heads, int state_dim)
Definition hybrid_attention_kernels.c:253
gradient_scale_f32
void gradient_scale_f32(float *grad, size_t numel, float scale)
Definition optimizer_kernels.c:1073
mrope_qk_vision_bf16_storage
void mrope_qk_vision_bf16_storage(float *q, float *k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2735
gemm_nt_q8_0_q8_0
void gemm_nt_q8_0_q8_0(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
gemm_nt_q8_0_q8_0 with optional bias (matches header signature)
Definition gemm_batch_int8.c:552
deepseek_mla_attention_f32
void deepseek_mla_attention_f32(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim)
Definition deepseek_kernels.c:953
attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_token_output
void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_token_output(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int visual_start, int visual_tokens)
Definition attention_kernels.c:5302
relu2_forward
void relu2_forward(const float *input, float *output, size_t n)
Definition relu_kernels.c:120
rope_forward_with_rotary_dim
void rope_forward_with_rotary_dim(float *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)
Definition rope_kernels.c:809
rope_forward_q_split_direct_f32
void rope_forward_q_split_direct_f32(float *q, const float *freq_factors, int use_freq_factors, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)
Definition rope_kernels.c:1374
rmsnorm_forward_kv_lora
void rmsnorm_forward_kv_lora(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:757
deepseek_mhc_mix_backward_f32
void deepseek_mhc_mix_backward_f32(const float *d_out, const float *streams, const float *mix, float *d_streams, float *d_mix, int tokens, int n_streams, int dim)
Definition deepseek_kernels.c:53
nemotron_group_limited_topk_router_f32
void nemotron_group_limited_topk_router_f32(const float *scores, const float *correction_bias, int *indices, float *weights, int rows, int n_experts, int top_k, int n_group, int topk_group, int norm_topk_prob, float routed_scaling_factor)
Definition topk_kernels.c:780
moe_swiglu_expert_forward_q4k_q5k_auto_workspace
int moe_swiglu_expert_forward_q4k_q5k_auto_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2619
gemm_fine_grained_parallel
void gemm_fine_grained_parallel(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:286
moe_swiglu_expert_q4k_q8_0_workspace_bytes
size_t moe_swiglu_expert_q4k_q8_0_workspace_bytes(int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:742
moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes
size_t moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes(int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:761
spatial_average_pool_contiguous
void spatial_average_pool_contiguous(const float *input, float *output, int grid_h, int grid_w, int embed_dim, int merge_size)
Definition vision_kernels.c:657
rope_forward_q_split_llama_f32
void rope_forward_q_split_llama_f32(float *q, const float *freq_factors, int use_freq_factors, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)
Definition rope_kernels.c:1390
ck_q6_k_prepared_provider_name
const char * ck_q6_k_prepared_provider_name(void)
Definition gemm_kernels_q6k_q8k.c:62
relu_forward_inplace
void relu_forward_inplace(float *data, size_t n)
Definition relu_kernels.c:54
gemm_nt_bf16_accurate_storage_parallel_dispatch
void gemm_nt_bf16_accurate_storage_parallel_dispatch(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:2161
gelu_fast_inplace
void gelu_fast_inplace(float *data, size_t n)
Definition gelu_kernels.c:244
gemv_q4_k_q8_k
void gemv_q4_k_q8_k(float *y, const void *W, const void *x_q8, int M, int K)
Definition gemm_kernels_q4k_q8k.c:497
ck_multimodal_mrope_positions_2d
int ck_multimodal_mrope_positions_2d(int32_t *positions, int total_tokens, int prefix_start, int position_base, int prefix_tokens, int grid_x, int grid_y, int text_pos)
Definition rope_kernels.c:2982
gemv_q5_0_parallel
void gemv_q5_0_parallel(float *y, const void *W, const float *x, int M, int K, int ith, int nth)
Parallel reference GEMV for Q5_0 × FP32.
Definition gemm_kernels_q5_0.c:579
yarn_rope_cache_explicit_positions_f32
void yarn_rope_cache_explicit_positions_f32(float *cos_cache, float *sin_cache, const int32_t *positions, int num_tokens, int rotary_dim, float freq_base, float factor, int original_context, float beta_fast, float beta_slow, float mscale, float mscale_all_dim)
Definition rope_kernels.c:498
deepseek_dsa_topk_softmax_f32
void deepseek_dsa_topk_softmax_f32(const float *scores, int *indices, float *weights, int tokens, int heads, int key_count, int top_k)
Definition deepseek_kernels.c:86
gemma4_v_norm_forward
void gemma4_v_norm_forward(const float *input, float *output, float *rstd_cache, int tokens, int num_kv_heads, int head_dim, float eps)
Definition rmsnorm_kernels.c:839
attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output
void attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5121
gemv_q4_k_q8_k_ref
void gemv_q4_k_q8_k_ref(float *y, const void *W, const void *x_q8, int M, int K)
Definition gemm_kernels_q4k_q8k.c:425
qk_norm_backward_last_isa
int qk_norm_backward_last_isa(void)
Definition qk_norm_kernels.c:120
farskip_swiglu_shared_combine_bf16
void farskip_swiglu_shared_combine_bf16(const float *hidden, const float *routed, const float *post_attn_residual, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, float *main_output, float *routed_free_output, int rows, int hidden_dim, int intermediate_dim)
Definition axpy_kernels.c:3405
adamw_update_f32
void adamw_update_f32(const float *grad, float *weight, float *m, float *v, size_t numel, float lr, float beta1, float beta2, float eps, float weight_decay, int step)
Definition optimizer_kernels.c:592
geglu_forward_bf16
void geglu_forward_bf16(const uint16_t *x, uint16_t *out, int tokens, int dim, float *scratch)
Definition geglu_kernels.c:98
farskip_swiglu_shared_combine_bf16_row_range
void farskip_swiglu_shared_combine_bf16_row_range(const float *hidden, const float *routed, const float *post_attn_residual, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, float *main_output, float *routed_free_output, int rows, int hidden_dim, int intermediate_dim, int row_begin, int row_end)
Definition axpy_kernels.c:3352
gemm_q4_k
void gemm_q4_k(float *Y, const void *W, const float *X, int M, int N, int K)
Auto-dispatch GEMM based on available SIMD.
Definition gemm_kernels_q4k.c:486
ssm_conv1d_forward
void ssm_conv1d_forward(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:128
feature_concat
void feature_concat(const float *main_input, const float *branch_input, float *output, int rows, int main_dim, int branch_slice_dim, int num_branch_slices)
Definition vision_kernels.c:804
gradient_accumulate_f32
void gradient_accumulate_f32(float *dst, const float *src, size_t numel)
Definition optimizer_kernels.c:925
deepseek_dsa_topk_softmax_backward_f32
void deepseek_dsa_topk_softmax_backward_f32(const int *indices, const float *weights, const float *d_weights, float *d_scores, int tokens, int heads, int key_count, int top_k)
Definition deepseek_kernels.c:153
recurrent_qk_l2_norm_backward
void recurrent_qk_l2_norm_backward(const float *d_q_out, const float *d_k_out, const float *q, const float *k, float *d_q, float *d_k, int rows, int q_dim, int k_dim, int head_dim, float eps)
Definition recurrent_qk_norm_kernels.c:472
deepseek_mla_kv_decompress_bf16_parallel_dispatch
void deepseek_mla_kv_decompress_bf16_parallel_dispatch(const float *compressed_kv, const uint16_t *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim)
Definition deepseek_kernels.c:308
gelu_backward_scalar
void gelu_backward_scalar(const float *input, const float *d_output, float *d_input, size_t n)
Definition gelu_kernels.c:846
rmsnorm_forward_int4
void rmsnorm_forward_int4(const uint8_t *input, const float *gamma, uint8_t *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps, float *scratch_input, float *scratch_output)
Definition rmsnorm_kernels_int4.c:78
ck_attention_reduction_t
ck_attention_reduction_t
Definition ckernel_engine.h:2284
CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA
@ CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA
Definition ckernel_engine.h:2289
CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64_REGULAR_TAIL_NOSCALE
@ CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64_REGULAR_TAIL_NOSCALE
Definition ckernel_engine.h:2291
CK_ATTN_REDUCTION_F16_REGULAR_NOSCALE
@ CK_ATTN_REDUCTION_F16_REGULAR_NOSCALE
Definition ckernel_engine.h:2290
CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE
@ CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE
Definition ckernel_engine.h:2286
CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64
@ CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64
Definition ckernel_engine.h:2288
CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE
@ CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE
Definition ckernel_engine.h:2287
CK_ATTN_REDUCTION_FP32_ONLINE
@ CK_ATTN_REDUCTION_FP32_ONLINE
Definition ckernel_engine.h:2285
gemm_q4_k_q8_k
void gemm_q4_k_q8_k(float *Y, const void *W, const void *X_q8, int M, int N, int K)
Definition gemm_kernels_q4k_q8k.c:605
gemm_backward_bf16_mixed
void gemm_backward_bf16_mixed(const uint16_t *d_output, const uint16_t *input, const uint16_t *weight, float *d_input, float *d_weight, float *d_bias, int tokens, int in_dim, int out_dim)
Definition gemm_kernels_bf16.c:2194
embedding_forward_q5_0
void embedding_forward_q5_0(const int32_t *token_ids, int token_count, int vocab_size, const void *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels.c:131
geglu_forward_quick
void geglu_forward_quick(const float *x, float *out, int tokens, int dim)
Definition geglu_kernels.c:55
mrope_qk_text_imrope_bf16_pytorch_storage
void mrope_qk_text_imrope_bf16_pytorch_storage(float *q, float *k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2586
gemm_microkernel_blocked
void gemm_microkernel_blocked(const float *A, const float *B, float *C, int M, int N, int K)
Definition gemm_microkernel.c:934
dequant_q5_1_row
void dequant_q5_1_row(const void *src, float *dst, size_t n_elements)
Dequantize Q5_1 row (multiple blocks)
Definition dequant_kernels.c:257
fused_mlp_swiglu_decode_v2
void fused_mlp_swiglu_decode_v2(const float *x, const float *W_gate, const float *W_up, const float *W_down, const float *b_gate, const float *b_up, const float *b_down, float *output, int D, int Hff)
Definition mlp_fused_decode.c:318
rmsnorm_forward_strided_pytorch_bf16_storage
void rmsnorm_forward_strided_pytorch_bf16_storage(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)
Definition rmsnorm_kernels.c:440
gemm_nt_q6_k_q8_k
void gemm_nt_q6_k_q8_k(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q6k_q8k.c:1548
moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace
int moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, const float *shared_gate_input, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:3023
attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract
ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, ck_attention_reduction_t reduction)
Definition attention_kernels.c:6921
rmsnorm_backward
void rmsnorm_backward(const float *d_output, const float *input, const float *gamma, const float *rstd_cache, float *d_input, float *d_gamma, int tokens, int d_model, int aligned_embed_dim)
Definition rmsnorm_kernels.c:888
dequant_q4_1_row
void dequant_q4_1_row(const void *src, float *dst, size_t n_elements)
Dequantize Q4_1 row (multiple blocks)
Definition dequant_kernels.c:141
swiglu_forward_ggml
void swiglu_forward_ggml(const float *input, float *output, int tokens, int dim)
Definition swiglu_kernels.c:727
moe_swiglu_shared_forward_q4k_q5_0_gated_workspace
int moe_swiglu_shared_forward_q4k_q5_0_gated_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, const float *shared_gate_input, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2852
final_logit_scale_muse_pytorch_bf16_storage
void final_logit_scale_muse_pytorch_bf16_storage(float *logits, int tokens, int vocab_size, float scale)
Definition logit_kernels.c:24
deepseek_mla_partial_rope_concat_packed_f32
void deepseek_mla_partial_rope_concat_packed_f32(const float *q_packed, const float *k_nope, const float *kv_a_packed, const float *cos, const float *sin, float *query, float *key, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int qk_rope_dim)
Definition deepseek_kernels.c:406
gemm_nt_q6_k_q8_k_m2
void gemm_nt_q6_k_q8_k_m2(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q6k_q8k.c:1854
recurrent_dt_gate_forward_pytorch_fp32
void recurrent_dt_gate_forward_pytorch_fp32(const float *alpha, const float *dt_bias, const float *a, float *gate, int rows, int num_heads, int state_dim)
Definition recurrent_gate_kernels.c:180
gemv_bf16_parallel_dispatch
void gemv_bf16_parallel_dispatch(float *y, const void *W, const float *x, int M, int K)
Definition gemm_kernels_bf16.c:653
gemv_fused_q5_0_bias_dispatch
void gemv_fused_q5_0_bias_dispatch(float *y, const void *W, const float *x, const float *bias, int M, int K)
Definition gemv_fused_quant_bias.c:510
gated_deltanet_prefill_forward
void gated_deltanet_prefill_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int state_dim, float norm_eps)
Definition deltanet_kernels.c:1999
fused_mlp_swiglu_prefill_bias
void fused_mlp_swiglu_prefill_bias(const float *x, const float *W_gate, const float *W_up, const float *W_down, const float *B_gate, const float *B_up, const float *B_down, float *output, int seq_len, int hidden, int intermediate, float *scratch)
Fused MLP (Gate + Up + SwiGLU + Down) for prefill with biases.
Definition prefill_fused_gemm.c:746
layernorm_forward_rolled_slice
void layernorm_forward_rolled_slice(const float *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, float *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, int aligned_embed_dim, float eps)
Definition layernorm_kernels.c:398
position_embeddings_add_at_offset
void position_embeddings_add_at_offset(float *x, const float *position_embd, int num_tokens, int embed_dim, int num_positions, int start_position)
Definition vision_kernels.c:187
embedding_forward_q8_0
void embedding_forward_q8_0(const int32_t *token_ids, int token_count, int vocab_size, const void *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels.c:186
attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided
void attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:3621
attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_sliding
void attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_sliding(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window, float *scores, size_t scores_bytes, uint16_t *key_rows, size_t key_rows_bytes, uint16_t *value_columns, size_t value_columns_bytes, uint16_t *gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)
Definition attention_kernels.c:9224
swiglu_forward_ggml_split
void swiglu_forward_ggml_split(const float *gate, const float *up, float *output, int tokens, int dim)
Definition swiglu_kernels.c:743
rope_forward_qk_with_rotary_dim
void rope_forward_qk_with_rotary_dim(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)
Definition rope_kernels.c:1109
attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)
Definition attention_kernels.c:7770
attention_forward_causal_head_major_shared_kv_sliding_gemma4
void attention_forward_causal_head_major_shared_kv_sliding_gemma4(const float *q, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
Definition attention_kernels_sliding.c:801
ck_flash_attn_fast_exp_kind
int ck_flash_attn_fast_exp_kind(void)
Definition attention_flash_true.c:112
gated_deltanet_llama_chunk64_prefill_forward
void gated_deltanet_llama_chunk64_prefill_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)
Definition deltanet_kernels.c:1217
mrope_qk_vision_fp16_storage
void mrope_qk_vision_fp16_storage(float *q, float *k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)
Definition rope_kernels.c:2736
moe_swiglu_shared_forward_q4k_q4k_workspace
int moe_swiglu_shared_forward_q4k_q4k_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1242
attention_backward_causal_head_major
void attention_backward_causal_head_major(const float *d_output, const float *q, const float *k, const float *v, const float *attn_weights, float *d_q, float *d_k, float *d_v, float *d_scores, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)
Definition attention_kernels.c:9931
deepseek_mla_attention_decode_bf16_phases_workspace
void deepseek_mla_attention_decode_bf16_phases_workspace(const float *q, const float *k_cache, const float *v_cache, float *output, int num_heads, int num_kv_heads, int cache_len, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride, float scale, float *scores, size_t scores_bytes)
Definition deepseek_kernels.c:1144
rope_forward_qk_split_direct_f32
void rope_forward_qk_split_direct_f32(float *q, float *k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)
Definition rope_kernels.c:1180
embedding_forward
void embedding_forward(const int32_t *token_ids, int token_count, int vocab_size, const float *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels.c:22
mamba2_dt_softplus_f32
void mamba2_dt_softplus_f32(const float *dt, const float *dt_bias, float *dt_out, int rows, int num_heads, float dt_min, float dt_max)
Definition mamba2_kernels.c:171
attention_forward_decode_head_major_gqa_flash_f16cache_split
void attention_forward_decode_head_major_gqa_flash_f16cache_split(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int split_chunks)
Definition attention_kernels.c:6784
ssm_conv1d_forward_llama_production
void ssm_conv1d_forward_llama_production(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:207
moe_swiglu_expert_forward_f32
void moe_swiglu_expert_forward_f32(const float *hidden, const int *indices, const float *routing_weights, const float *expert_gate, const float *expert_up, const float *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:450
gelu_erf_bf16_storage
void gelu_erf_bf16_storage(float *data, size_t n)
Definition gelu_kernels.c:661
layernorm_forward_unrolled_slice_bf16
void layernorm_forward_unrolled_slice_bf16(const uint16_t *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, uint16_t *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, float eps, float *scratch_input, float *scratch_output)
Definition layernorm_kernels_bf16.c:57
moe_swiglu_packed_expert_forward_bf16
void moe_swiglu_packed_expert_forward_bf16(const float *hidden, const int *indices, const float *routing_weights, const uint16_t *expert_gate_up, const uint16_t *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)
Definition axpy_kernels.c:582
gemm_bias_gelu_fused
void gemm_bias_gelu_fused(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_fused_kernels.c:131
fused_mlp_swiglu_prefill_w1w2_quant
void fused_mlp_swiglu_prefill_w1w2_quant(const float *x, const void *W1, const float *B1, CKDataType w1_dt, const void *W2, const float *B2, CKDataType w2_dt, float *output, int seq_len, int embed_dim, int aligned_embed_dim, int intermediate_dim, int aligned_intermediate_dim, void *scratch)
Quantized fused MLP for prefill (W1=gate+up, W2=down)
Definition prefill_fused_gemm.c:964
rope_forward_qk_bf16
void rope_forward_qk_bf16(uint16_t *q, uint16_t *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float *scratch_q, float *scratch_k)
Definition rope_kernels_bf16.c:94
ck_attention_sparse_bf16_pytorch_gqa_available
int ck_attention_sparse_bf16_pytorch_gqa_available(void)
Definition attention_kernels.c:9016
attention_backward_causal_head_major_gqa
void attention_backward_causal_head_major_gqa(const float *d_output, const float *q, const float *k, const float *v, const float *attn_weights, float *d_q, float *d_k, float *d_v, float *d_scores, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)
Definition attention_kernels.c:9792
recurrent_split_qkv_forward
void recurrent_split_qkv_forward(const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
Definition recurrent_split_kernels.c:5
qk_norm_forward_llama_production_parallel_dispatch
void qk_norm_forward_llama_production_parallel_dispatch(float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
embedding_backward_bf16_mixed
void embedding_backward_bf16_mixed(const int32_t *token_ids, int token_count, const uint16_t *d_output, float *d_token_embeddings, float *d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
Definition embedding_kernels_bf16.c:166
ck_attention_status_t
ck_attention_status_t
Definition ckernel_engine.h:2304
CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT
@ CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT
Definition ckernel_engine.h:2307
CK_ATTENTION_STATUS_OK
@ CK_ATTENTION_STATUS_OK
Definition ckernel_engine.h:2305
CK_ATTENTION_STATUS_INVALID_ARGUMENT
@ CK_ATTENTION_STATUS_INVALID_ARGUMENT
Definition ckernel_engine.h:2306
CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE
@ CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE
Definition ckernel_engine.h:2308
topk_f32
void topk_f32(const float *scores, int n, int k, int *indices, float *values)
Find top-K indices and values from a score vector.
Definition topk_kernels.c:51
speculative_verify_greedy_f32
void speculative_verify_greedy_f32(const float *target_logits, int vocab_size, int draft_token, int *accepted, int *verified_token)
Greedy one-token speculative verification.
Definition topk_kernels.c:585
attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace
ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace(const float *q, const uint16_t *k_cache, const uint16_t *v_cache, float *output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes, const int *segment_lengths, int num_segments)
Definition attention_kernels.c:8130
gemm_microkernel_blocked_bt
void gemm_microkernel_blocked_bt(const float *A, const float *B, float *C, int M, int N, int K)
Definition gemm_microkernel.c:1058
dequant_q6_k_row
void dequant_q6_k_row(const void *src, float *dst, size_t n_elements)
Dequantize Q6_K row (multiple blocks)
Definition dequant_kernels.c:422
attention_forward_full_head_major_gqa_sdpa_bf16_storage
void attention_forward_full_head_major_gqa_sdpa_bf16_storage(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5022
mamba2_selective_scan_f32
void mamba2_selective_scan_f32(const float *state_init, const float *x, const float *dt, const float *a, const float *b, const float *c, const float *d, float *state_out, float *y, int batch, int seq_len, int num_heads, int head_dim, int state_dim, int num_groups)
Definition mamba2_kernels.c:339
mamba2_selective_scan_f32_head_range
void mamba2_selective_scan_f32_head_range(const float *state_init, const float *x, const float *dt, const float *a, const float *b, const float *c, const float *d, float *state_out, float *y, int batch, int seq_len, int num_heads, int head_dim, int state_dim, int num_groups, int head_begin, int head_end)
Definition mamba2_kernels.c:264
deepseek_mla_attention_decode_f32
void deepseek_mla_attention_decode_f32(const float *q, const float *k_cache, const float *v_cache, float *output, int num_heads, int num_kv_heads, int cache_len, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride)
Definition deepseek_kernels.c:1155
rmsnorm_forward_fp32_square_fp64_sum
void rmsnorm_forward_fp32_square_fp64_sum(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
patch_projection_image_bf16_pytorch_onednn_conv3d_storage
void patch_projection_image_bf16_pytorch_onednn_conv3d_storage(const float *image, const void *weights_t0, const void *weights_t1, const float *bias, float *output, int channels, int image_h, int image_w, int patch_size, int out_channels, int merge_size)
Definition gemm_kernels_bf16.c:1819
moe_swiglu_expert_forward_q4k_q5k_auto_prepared_workspace
int moe_swiglu_expert_forward_q4k_q5k_auto_prepared_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
gemv_q6_k_q8_k_parallel_simd
void gemv_q6_k_q8_k_parallel_simd(float *y, const void *W, const void *x_q8, int M, int K, int ith, int nth)
Parallel SIMD GEMV for Q6_K × Q8_K.
Definition gemm_kernels_q6k_q8k.c:1454
gemm_nt_bf16_bf16_storage
void gemm_nt_bf16_bf16_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:2152
add_stream_inplace
void add_stream_inplace(float *a, const float *b, size_t n)
Definition vision_kernels.c:713
rmsnorm_forward_muse_centered_pytorch_bf16_storage
void rmsnorm_forward_muse_centered_pytorch_bf16_storage(const float *input, const float *zero_centered_weight, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:470
qwen4_ple_ngram_embed_q5_0
void qwen4_ple_ngram_embed_q5_0(const int32_t *token_ids, const void *embedding, const int64_t *layer_multipliers, const int64_t *head_offsets, const int64_t *head_vocab_sizes, float *output, const float *token_state_in, float *token_state_out, int rows, int ngram_size, int heads_per_ngram, int head_dim, int eos_token_id, int position)
Definition qwen4_exp_kernels.c:141
recurrent_norm_gate_llama_avx2_forward
void recurrent_norm_gate_llama_avx2_forward(const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
Definition recurrent_norm_kernels.c:53
final_logit_softcap_muse_pytorch_bf16_storage
void final_logit_softcap_muse_pytorch_bf16_storage(float *logits, int tokens, int vocab_size, float cap)
Definition logit_kernels.c:40
split_q_gate_forward
void split_q_gate_forward(const float *packed_qg, float *q, float *gate, int rows, int q_dim, int gate_dim, int group_dim)
Definition hybrid_attention_kernels.c:41
gemm_tn_blocked
void gemm_tn_blocked(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:769
ssm_conv1d_backward
void ssm_conv1d_backward(const float *d_out, const float *conv_x, const float *kernel, float *d_conv_x, float *d_kernel, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:316
rope_forward_qk_bf16_with_rotary_dim
void rope_forward_qk_bf16_with_rotary_dim(uint16_t *q, uint16_t *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float *scratch_q, float *scratch_k)
Definition rope_kernels_bf16.c:112
recurrent_silu_backward
void recurrent_silu_backward(const float *d_out, const float *x, float *d_x, int rows, int dim)
Definition recurrent_gate_kernels.c:557
attention_forward_decode_head_major_gqa_flash_f16cache_contract
ck_attention_status_t attention_forward_decode_head_major_gqa_flash_f16cache_contract(const float *q_token, const uint16_t *k_cache, const uint16_t *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)
Definition attention_kernels.c:6933
rmsnorm_forward_no_weight_llama_production
void rmsnorm_forward_no_weight_llama_production(const float *input, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:806
rope_forward_qk_gemma4_direct
void rope_forward_qk_gemma4_direct(float *q, float *k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)
Definition rope_kernels.c:1407
moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace
int moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, const float *shared_gate_input, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:3002
recurrent_conv_state_update_backward
void recurrent_conv_state_update_backward(const float *d_conv_x, const float *d_state_out, float *d_state_in, float *d_q, float *d_k, float *d_v, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
Definition recurrent_state_kernels.c:257
gemm_nt_f32_llama_production
void gemm_nt_f32_llama_production(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:1182
topk_softmax_f32
void topk_softmax_f32(const float *scores, int n, int k, int *indices, float *weights)
Find top-K indices with softmax-normalized weights.
Definition topk_kernels.c:136
gemm_backward_f32_train_parallel_dispatch
void gemm_backward_f32_train_parallel_dispatch(const float *d_output, const float *input, const float *W, float *d_input, float *d_W, float *d_b, int T, int aligned_in, int aligned_out, int num_threads)
Definition ck_parallel_train.c:610
gemm_nt_bf16_prefill_shape_safe_bf16_storage
void gemm_nt_bf16_prefill_shape_safe_bf16_storage(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_bf16.c:1354
gemv_q5_0_parallel_simd
void gemv_q5_0_parallel_simd(float *y, const void *W, const float *x, int M, int K, int ith, int nth)
Parallel SIMD GEMV for Q5_0 × FP32 with prefetching.
Definition gemm_kernels_q5_0.c:625
speculative_commit_one_i32
void speculative_commit_one_i32(int accepted, int verified_token, int *token_buffer, int *token_count, int max_tokens, int *target_position, int *draft_position, int *accepted_count, int *rejected_count)
Commit one verified speculative token and update decode counters.
Definition topk_kernels.c:611
rowwise_bias_add
void rowwise_bias_add(float *x, const float *bias, int rows, int dim)
Definition vision_kernels.c:696
moe_swiglu_shared_forward_bf16_gated_parallel_dispatch
void moe_swiglu_shared_forward_bf16_gated_parallel_dispatch(const float *hidden, const float *routed, const uint16_t *shared_gate, const uint16_t *shared_up, const uint16_t *shared_down, const uint16_t *shared_router, float *output, int rows, int hidden_dim, int intermediate_dim)
kv_cache_store
void kv_cache_store(float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)
Definition kv_cache_kernels.c:122
gelu_backward_fast
void gelu_backward_fast(const float *input, const float *d_output, float *d_input, size_t n)
Definition gelu_kernels.c:870
moe_swiglu_expert_forward_q4k_q6k_parallel_workspace
int moe_swiglu_expert_forward_q4k_q6k_parallel_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:1858
deepseek_hybrid_attention_f32
void deepseek_hybrid_attention_f32(const float *q, const float *k, const float *v, const int *indices, float *out, float *attn, int query_tokens, int key_tokens, int heads, int dim, int top_k, float scale, int mode)
Definition deepseek_kernels.c:683
softmax_cross_entropy_loss
void softmax_cross_entropy_loss(const float *logits, const int32_t *targets, int tokens, int vocab_size, float *d_logits, float *loss_out)
Definition loss_kernels.c:259
attention_forward_causal_head_major_gqa_flash_strided_token_output
void attention_forward_causal_head_major_gqa_flash_strided_token_output(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:4518
causal_softmax_head_major_bf16
void causal_softmax_head_major_bf16(uint16_t *scores, int num_heads, int num_tokens, int aligned_context_window, float *scratch)
Definition softmax_kernels_bf16.c:31
position_embeddings_add
void position_embeddings_add(float *x, const float *position_embd, int num_tokens, int embed_dim, int num_positions)
Definition vision_kernels.c:168
ck_get_num_threads
int ck_get_num_threads(void)
Definition ckernel_strict.c:310
patch2im_bf16
void patch2im_bf16(const uint16_t *d_patches, uint16_t *d_image, int C, int H, int W, int P)
Definition vision_kernels_bf16.c:57
gelu_erf_fp64_f32_parallel_dispatch
void gelu_erf_fp64_f32_parallel_dispatch(float *data, size_t n)
Definition gelu_kernels.c:606
moe_softmax_topk_router_pytorch_bf16_workspace
int moe_softmax_topk_router_pytorch_bf16_workspace(const float *logits, int *indices, float *weights, int rows, int n_experts, int top_k, float routed_scaling_factor, void *workspace, size_t workspace_bytes)
Definition topk_kernels.c:375
ck_strict_parity_enabled
int ck_strict_parity_enabled(void)
Definition ckernel_strict.c:43
vision_position_ids_2d_merge
void vision_position_ids_2d_merge(int32_t *positions, int grid_h, int grid_w, int merge_size)
Definition vision_kernels.c:529
deepseek_mla_kv_cache_store_f32
void deepseek_mla_kv_cache_store_f32(float *k_cache, float *v_cache, const float *k, const float *v, int pos, int num_kv_heads, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride)
Definition deepseek_kernels.c:1011
moe_swiglu_shared_forward_q4k_q6k_parallel_workspace
int moe_swiglu_shared_forward_q4k_q6k_parallel_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2052
hyper_stream_expand_bf16
void hyper_stream_expand_bf16(const float *input, float *output, int rows, int streams, int hidden_dim)
Definition hyper_connection_kernels.c:53
moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace
int moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, const void *expert_gate_packed, const void *expert_up_packed, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2593
rope_precompute_cache_split
void rope_precompute_cache_split(float *cos_cache, float *sin_cache, int max_seq_len, int head_dim, float base)
Definition rope_kernels.c:340
add_backward_bf16
void add_backward_bf16(const uint16_t *d_y, uint16_t *d_a, uint16_t *d_b, size_t n)
Definition add_kernels_bf16.c:187
gemm_nt_q5_k
void gemm_nt_q5_k(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_k.c:1028
attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_token_output
void attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_token_output(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
Definition attention_kernels.c:5072
gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace
void gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace(const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes)
Definition gemm_kernels_bf16.c:1340
rope_backward
void rope_backward(const float *d_out, float *d_x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)
Definition rope_kernels.c:884
gemm_nt_bf16_row_range
void gemm_nt_bf16_row_range(const float *A, const void *B, const float *bias, float *C, int M, int N, int K, int row_begin, int row_end)
Definition gemm_kernels_bf16.c:727
dequant_q4_k_row
void dequant_q4_k_row(const void *src, float *dst, size_t n_elements)
Dequantize Q4_K row (multiple blocks)
Definition dequant_kernels.c:372
gemm_nt_bf16_amx_bf16_storage_workspace
void gemm_nt_bf16_amx_bf16_storage_workspace(const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes)
Definition gemm_kernels_bf16.c:1265
recurrent_dt_gate_backward
void recurrent_dt_gate_backward(const float *d_gate, const float *alpha, const float *dt_bias, const float *a, float *d_alpha, float *d_dt_bias, float *d_a, int rows, int dim)
Definition recurrent_gate_kernels.c:105
scal_copy_f32
void scal_copy_f32(float *y, const float *x, float alpha, int n)
Scaled copy: y = alpha * x.
Definition axpy_kernels.c:156
gemm_microkernel_packed
void gemm_microkernel_packed(const float *A, const float *B, float *C, int M, int N, int K)
Definition gemm_microkernel.c:840
gemm_blocked_serial
void gemm_blocked_serial(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:849
deepseek_mla_partial_rope_concat_packed_bf16_storage
void deepseek_mla_partial_rope_concat_packed_bf16_storage(const float *q_packed, const float *k_nope, const float *kv_a_packed, const float *cos, const float *sin, float *query, float *key, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int qk_rope_dim)
Definition deepseek_kernels.c:473
rope_forward_strided
void rope_forward_strided(float *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int head_stride_tokens)
Definition rope_kernels.c:837
moe_swiglu_expert_forward_q4k_q6k_workspace
int moe_swiglu_expert_forward_q4k_q6k_workspace(const float *hidden, const int *indices, const float *routing_weights, const void *expert_gate, const void *expert_up, const void *expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:869
recurrent_norm_gate_forward
void recurrent_norm_gate_forward(const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
Definition recurrent_norm_kernels.c:18
fused_mlp_swiglu_scratch_size
size_t fused_mlp_swiglu_scratch_size(int intermediate)
Get scratch buffer size for fused_mlp_swiglu_prefill.
Definition prefill_fused_gemm.c:898
sigmoid_forward
void sigmoid_forward(const float *input, float *output, size_t n)
Definition sigmoid_kernels.c:122
layernorm_forward_rolled_slice_bf16
void layernorm_forward_rolled_slice_bf16(const uint16_t *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, uint16_t *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, int aligned_embed_dim, float eps, float *scratch_input, float *scratch_output)
Definition layernorm_kernels_bf16.c:30
relu_backward_bf16
void relu_backward_bf16(const uint16_t *input, const uint16_t *d_output, uint16_t *d_input, size_t n)
Definition relu_kernels_bf16.c:69
final_logit_scale_f32
void final_logit_scale_f32(float *logits, int tokens, int vocab_size, float scale)
Definition logit_kernels.c:9
recurrent_dt_gate_forward
void recurrent_dt_gate_forward(const float *alpha, const float *dt_bias, const float *a, float *gate, int rows, int num_heads, int state_dim)
Definition recurrent_gate_kernels.c:66
quantize_batch_q8_0
void quantize_batch_q8_0(const float *x, void *y, int num_rows, int k)
Batch quantize FP32 to Q8_0 format (row-major output)
Definition gemm_kernels_q8_0.c:256
gelu_pytorch_tanh_bf16_storage
void gelu_pytorch_tanh_bf16_storage(float *data, size_t n)
Definition gelu_kernels.c:642
ck_attention_bf16_pytorch_gqa_available
int ck_attention_bf16_pytorch_gqa_available(void)
Definition attention_kernels.c:9025
gemm_nt_q5_1_q8_1_m2
void gemm_nt_q5_1_q8_1_m2(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q5_1_q8_1.c:396
gemv_q4_0
void gemv_q4_0(float *y, const void *W, const float *x, int M, int K)
Auto-dispatch GEMV.
Definition gemm_kernels_q4_0.c:132
moe_swiglu_shared_forward_q4k_q8_0_gated_workspace
int moe_swiglu_shared_forward_q4k_q8_0_gated_workspace(const float *hidden, const float *routed, const void *shared_gate, const void *shared_up, const void *shared_down, const float *shared_gate_input, float *output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)
Definition axpy_kernels.c:2839
gemv_q4_k_q8_k_parallel_simd
void gemv_q4_k_q8_k_parallel_simd(float *y, const void *W, const void *x_q8, int M, int K, int ith, int nth)
Definition gemm_kernels_q4k_avx.c:263
layernorm_forward_unrolled_slice
void layernorm_forward_unrolled_slice(const float *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, float *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, float eps)
Definition layernorm_kernels.c:730
gemm_nn_blocked
void gemm_nn_blocked(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:590
mamba2_conv1d_decode_f32
void mamba2_conv1d_decode_f32(const float *state_in, const float *x, const float *weight, const float *bias, float *conv_out, float *state_out, int rows, int conv_dim, int kernel_size)
Definition mamba2_kernels.c:148
add_scaled_inplace_bf16
void add_scaled_inplace_bf16(uint16_t *a, const uint16_t *b, float alpha, size_t n)
Definition add_kernels_bf16.c:149
fused_rmsnorm_qkv_prefill_head_major_quant
void fused_rmsnorm_qkv_prefill_head_major_quant(const float *x, const float *gamma, const void *Wq, const float *Bq, CKDataType wq_dt, const void *Wk, const float *Bk, CKDataType wk_dt, const void *Wv, const float *Bv, CKDataType wv_dt, float *Q, float *K, float *V, int seq_len, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int kv_stride_tokens, float eps, void *scratch)
Fused RMSNorm + QKV projection for prefill (head-major, Q8 activations)
Definition prefill_fused_gemm.c:519
mamba2_rmsnorm_gate_f32
void mamba2_rmsnorm_gate_f32(const float *x, const float *gate, const float *weight, float *out, int rows, int inner_dim, int group_size, float eps)
Definition mamba2_kernels.c:373
gemm_tn_avx512
void gemm_tn_avx512(const float *A, const float *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels.c:709
gemm_nt_q8_0_q8_0_contract
void gemm_nt_q8_0_q8_0_contract(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition gemm_kernels_q8_0_q8_0_contract.c:447
ckernel_quant.h
Quantization block structures for weight-only quantization.
ckernel_rwkv7_decode.h
cpu_features.h
mega_fused_attention.h
Mega-Fused Attention Kernel.
C
#define C(color)
Definition show_config.c:39
vocab_size
int vocab_size
Definition true_bpe.h:197