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