API Reference

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl(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, int output_token_major)

Flash attention decode with sliding window Testtest_attention.py::TestAttentionForward::test_sliding_window_decode

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl(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, int output_token_major)

Forward pass computation

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)

Forward pass computation

void ck_gemm_add_bias(float * C, const float * bias, int M, int N)

int ck_gemm_bf16_amx_available(void)

void ck_gemm_bf16_amx_work(int ith, int nth, void * opaque)

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)

void ck_gemm_bf16_native_work(int ith, int nth, void * opaque)

int ck_gemm_dynamic_schedule_enabled(void)

Return non-zero when independent GEMM tiles should use dynamic claiming.

void ck_gemm_f16_input_fp16_work(int ith, int nth, void * opaque)

int ck_gemm_f16_pick_active_threads(const ck_threadpool_t * pool, int M, int N, int K)

int ck_gemm_f16_threadpool_enabled(int M, int N, int K)

int ck_gemm_nn_impl_probe_enabled(void)

void ck_gemm_nt_bf16_exact_rows(int begin, int end, void * opaque)

void ck_gemm_nt_bf16_storage_exact_rows(int begin, int end, void * opaque)

int ck_gemm_nt_f16_ggml_oracle(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

int ck_gemm_nt_f16_simd_lanes(void)

void ck_gemm_nt_f32_llama_production_output(const float * A, const float * B, const float * bias, float * C, int M, int N, int K, int index)

void ck_gemm_nt_fp32_exact_rows(int begin, int end, void * opaque)

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)

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)

void ck_gemm_nt_quant(const float * A, const void * B, const float * bias, float * C, int M, int N, int K, CKDataType dtype)

void ck_gemma4_dequant_q5_k_block(const ck_gemma4_block_q5_K * block, float * out)

void ck_gemma4_embed_range(int begin, int end, void * opaque)

float ck_gemma4_gelu(float x)

void ck_gemma4_prepare_bf16_range(int begin, int end, void * opaque)

void ck_gemma4_prepare_parallel(int tokens, ck_range_fn_t fn, ck_gemma4_prepare_args_t * args)

void ck_gemma4_prepare_q5_range(int begin, int end, void * opaque)

uint8_t ck_gemma4_q5_k_value(const ck_gemma4_block_q5_K * block, int subblock, int i)

void ck_gemma4_rmsnorm_tmp(const float * x, const float * gamma, float * out, int n, float eps)

void ck_gemma4_unpack_q5_k_scales(const uint8_t * scales, uint8_t * sc, uint8_t * m)

int ck_get_gemm_schedule(void)

Return the configured process-wide GEMM scheduling policy.

float ck_llama_regular_gemm_f16(const float * probability, const float * value_column, int count)

int ck_set_gemm_schedule(int policy)

Set the process-wide GEMM tile scheduling policy.

const float * ck_strict_consume_next_gemm_a(size_t elems)

void ck_strict_store_next_gemm_a(const float * data, size_t elems)

void ck_test_gemm_q4_k(const void * weight_q4k, const float * input_f32, float * output, int rows, int cols, int n_tokens)

Q4_K GEMM - batched matrix multiply with quantized weights.

void ck_test_gemm_q5_0(const void * weight_q5_0, const float * input_f32, float * output, int rows, int cols, int n_tokens)

Test Q5_0 x Q8_0 GEMM (batch matrix multiply)

void ck_test_gemm_q6_k(const void * weight_q6k, const float * input_f32, float * output, int rows, int cols, int n_tokens)

Test Q6_K x Q8_K GEMM (batch matrix multiply)

void ck_test_gemm_q8_0(const void * weight_q8_0, const float * input_f32, float * output, int rows, int cols, int n_tokens)

Test Q8_0 x Q8_0 GEMM (batch matrix multiply)

void ck_train_gemm_backward_serial(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)

Backward pass / gradient computation

void ck_train_gemm_backward_work(int ith, int nth, void * argp)

Backward pass / gradient computation

void ck_train_gemm_nn_compute_rows(const float * A, const float * B, const float * bias, float * C, int row_start, int row_end, int N, int K)

void ck_train_gemm_nn_work(int ith, int nth, void * argp)

void ck_train_gemm_nt_compute_rows(const float * A, const float * B, const float * bias, float * C, int row_start, int row_end, int N, int K)

void ck_train_gemm_tn_compute_rows(const float * A, const float * B, const float * bias, float * C, int row_start, int row_end, int M, int N, int K)

void ck_train_gemm_work(int ith, int nth, void * argp)

void ckernel_sgemm_native(int M, int N, int K, const float * A, int lda, const float * B, int ldb, const float * bias, float * C, int ldc)

Native GEMM backend that directly reuses the C-Transformer GEMM kernel.

void compute_gemm_params(const CPUInfo * cpu, GEMMParams * params)

void fused_rmsnorm_gemm_2d_tiled(const float * x, const float * gamma, const float * W, float * output, int seq_len, int hidden, int out_dim, float eps, float * x_norm_scratch)

Fused RMSNorm + single GEMM with 2D tiling (weight reuse)

void gemm_avx512_parallel(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

const char * gemm_batch_int8_impl_name(void)

Get the best implementation name for logging/debugging.

void gemm_bf16_fp32out(const uint16_t * A, const uint16_t * B, const float * bias, float * C, int M, int N, int K)

void gemm_bias_gelu_fused(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_bias_relu_fused(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_bias_silu_fused(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_blocked_serial(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

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)

void gemm_blocked_serial_train_parallel_dispatch(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_f16(float * Y, const uint16_t * W, const float * X, int M, int N, int K)

Auto-dispatch GEMM based on available SIMD.

void gemm_f16_backward(float * dX, const uint16_t * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_f16_input_fp16_ref(float * Y, const uint16_t * W, const float * X, int M, int N, int K)

int gemm_f16_input_fp16_threadpool(float * Y, const uint16_t * W, const float * X, int M, int N, int K)

void gemm_f16_ref(float * Y, const uint16_t * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with FP16 weights (scalar reference)

void gemm_fine_grained_parallel(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

const char * gemm_get_backend(void)

void gemm_init_threads(void)

void gemm_microkernel(const float * A, const float * B, float * C, int M, int N, int K, int B_transposed)

void gemm_microkernel_blocked(const float * A, const float * B, float * C, int M, int N, int K)

void gemm_microkernel_blocked_bt(const float * A, const float * B, float * C, int M, int N, int K)

void gemm_microkernel_edge(int m, int n, int K, const float * A, int lda, const float * B, int ldb, float * C, int ldc, int first_k)

void gemm_microkernel_packed(const float * A, const float * B, float * C, int M, int N, int K)

void gemm_microkernel_sequential(const float * A, const float * B, float * C, int M, int N, int K)

void gemm_naive_parallel(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_naive_serial_double(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_naive_serial_float(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nn_avx512(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nn_avx512_probe(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nn_bf16(const uint16_t * A, const uint16_t * B, const uint16_t * bias, uint16_t * C, int M, int N, int K)

void gemm_nn_blocked(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nn_parallel(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nn_serial_double(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nn_simd(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt(const float * input, const float * weight, float * output, int rows, int cols, int common)

void gemm_nt_bf16(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_bf16_amx_bf16_storage(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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)

void gemm_nt_bf16_bf16_storage(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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)

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)

void gemm_nt_bf16_native_bf16_storage(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_bf16_parallel_dispatch(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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)

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)

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)

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)

void gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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)

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.

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)

int gemm_nt_f16_ggml_strict(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_f32_llama_production(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

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)

void gemm_nt_fp32_exact_parallel_dispatch(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_matvec_parallel(const float * A, const float * B, const float * bias, float * C, int N, int K)

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.

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.

void gemm_nt_q4_k(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_meta_q8_k(const void * A_q8, const void * B_packed, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_meta_q8_k_threaded(const void * A_q8, const void * B_packed, const float * bias, float * C, int M, int N, int K, int active_threads)

void gemm_nt_q4_k_packed_meta_q8_k_threaded_nsplit(const void * A_q8, const void * B_packed, const float * bias, float * C, int M, int N, int K, int active_threads)

void gemm_nt_q4_k_packed_meta_q8_k_tile(const void * A_q8, const void * B_packed, const float * bias, float * C, int M, int N, int K, int m0, int m1, int n0, int n1)

void gemm_nt_q4_k_packed_meta_x16_gateup_swiglu_fused_vnni(const void * A_q8, const void * B_packed_x16, const float * bias, float * C, int M, int D, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_meta_x16_q8_k_llama_order(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_meta_x16_q8_k_threaded_mreuse(const void * A_q8, const void * B_packed_x16, const float * bias, float * C, int M, int N, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_meta_x16_q8_k_threaded_mtile(const void * A_q8, const void * B_packed_x16, const float * bias, float * C, int M, int N, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_meta_x8_q8_k(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_meta_x8_q8_k_split_min_threaded_4m(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K, int active_threads)

void gemm_nt_q4_k_packed_meta_x8_q8_k_split_min_threaded_8m(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K, int active_threads)

void gemm_nt_q4_k_packed_meta_x8_q8_k_split_min_threaded_mreuse(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_meta_x8_q8_k_superblock_order(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_meta_x8_q8_k_threaded_mreuse(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_meta_x8_q8_k_threaded_mtile(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_meta_x8_q8_k_threaded_nsplit(const void * A_q8, const void * B_packed_x8, const float * bias, float * C, int M, int N, int K, int active_threads)

void gemm_nt_q4_k_packed_u8_q8_k(const void * A_q8, const void * B_packed, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_u8_x16_q8_k_threaded_mtile(const void * A_q8, const void * B_packed_u8_x16, const float * bias, float * C, int M, int N, int K, int tile_m, int active_threads)

void gemm_nt_q4_k_packed_vnni_x16_q8_k_gemv_order(const void * A_q8, const void * B_packed_x16, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q4_k_packed_vnni_x16_q8_k_split_min_threaded_16m(const void * A_q8, const void * B_packed_vnni_x16, const float * bias, float * C, int M, int N, int K, int active_threads)

void gemm_nt_q4_k_packed_vnni_x8_q8_k_split_min_threaded_4m(const void * A_q8, const void * B_packed_vnni_x8, const float * bias, float * C, int M, int N, int K, int active_threads)

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)

void gemm_nt_q4_k_q8_k_gateup_swiglu_fused_vnni(const void * A_q8, const void * B_gate_up, const float * bias, float * C, int M, int D, int K, int threads)

void gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch(const void * input, const void * weight, const float * bias, float * output, int rows, int output_dim, int input_dim)

void gemm_nt_q5_0(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_0_q8_0(const void * A_q8, const void * B_q5, const float * bias, float * C, int M, int N, int K)

Batch GEMM with Q5_0 weights and Q8_0 activations for prefill.

void gemm_nt_q5_0_q8_0_m2n4(const void * A_q8, const void * B_q5, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_0_q8_0_m2n4_tile(const void * A_q8, const void * B_q5, const float * bias, float * C, int M, int N, int K, int ldc)

void gemm_nt_q5_0_q8_0_m4n2(const void * A_q8, const void * B_q5, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_0_q8_0_m4n2_tile(const void * A_q8, const void * B_q5, const float * bias, float * C, int M, int N, int K, int ldc)

void gemm_nt_q5_0_q8_0_parallel_dispatch(const void * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_0_q8_0_ref(const void * A, const void * B, float * C, int M, int N, int K)

Dispatcher for gemm_nt_q8_0_q8_0.

void gemm_nt_q5_0_q8_0_unroll_avx(const void * A_q8, const void * B_q5, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_0_ref(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

GEMM with transposed Q5_0 weights: C = A @ B^T.

void gemm_nt_q5_0_sse(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_0_sse_v2(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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.

void gemm_nt_q5_1_q8_1(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_1_q8_1_m4(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_1_q8_1_m8(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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)

void gemm_nt_q5_k(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_k_prepared(const float * A, const void * B_prepared, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_k_prepared_m4(const float * A, const void * B_prepared, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_k_prepared_q8_m4_nrange(const void * A_q8, const void * B_prepared, const float * bias, float * C, int M, int N, int K, int n_begin, int n_end)

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)

void gemm_nt_q5_k_q8_k_ref(const void * A_q8, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_k_ref(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q5_k_ref_fp32(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q6_k(const float * A, const void * B, const float * bias, float * C, int M, int N, int 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.

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)

void gemm_nt_q6_k_q8_k_parallel_dispatch(const void * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q6_k_q8_k_prepared(const void * A_q8, const void * B_prepared, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q6_k_q8_k_prepared_avx512_vnni(const void * A_q8, const void * B_prepared, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q6_k_q8_k_prepared_tile(const void * A_q8, const void * B_prepared, const float * bias, float * C, int M, int N, int K, int m0, int m1, int n0, int n1)

void gemm_nt_q6_k_q8_k_prepared_tile_impl(const void * A_q8, const void * B_prepared, const float * bias, float * C, int M, int N, int K, int m0, int m1, int n0, int n1, int use_avx512_vnni)

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.

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.

void gemm_nt_q6_k_ref(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q6_k_sse(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q8_0(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

void gemm_nt_q8_0_dispatch(const void * A_q8, const void * B, const float * bias, float * C, int M, int N, int K, CKDataType dt)

void gemm_nt_q8_0_mlp_dispatch(const void * A_q8, const void * B, const float * bias, float * C, int M, int N, int K, CKDataType dt)

void gemm_nt_q8_0_q8_0(const void * A_q8, const void * B_q8, const float * bias, float * C, int M, int N, int K)

gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

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)

int gemm_nt_q8_0_q8_0_ggml_strict(const float * A, const void * B, const float * bias, float * C, int M, int N, int K)

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)

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)

void gemm_nt_q8_0_q8_0_ref(const void * A, const void * B, float * C, int M, int N, int K)

Scalar reference: gemm_nt_q8_0_q8_0.

void gemm_nt_q8_0_rowloop(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.

void gemm_nt_q8_k_mlp_dispatch(const void * A_q8, const void * B, const float * bias, float * C, int M, int N, int K, CKDataType dt)

void gemm_nt_q8_k_qkv_dispatch(const void * A_q8k, const void * B, const float * bias, float * C, int M, int N, int K, CKDataType dt)

void gemm_q4_0(float * Y, const void * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with Q4_0 weights.

void gemm_q4_0_backward(float * dX, const void * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_q4_1(float * Y, const void * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with Q4_1 weights.

void gemm_q4_1_backward(float * dX, const void * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_q4_gateup_swiglu_thread_fn(int ith, int nth, void * args)

void gemm_q4_gateup_swiglu_x16_thread_fn(int ith, int nth, void * args)

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.

void gemm_q4_k_backward(float * dX, const void * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_q4_k_q8_k(float * Y, const void * W, const void * X_q8, int M, int N, int K)

void gemm_q4_k_q8_k_compact_rows4(float * output, int output_stride, const void * weights, const void *const input_rows, int rows, int output_dim, int input_dim)

void gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4(float * output, const void * weights_packed, const void * input_q8, int rows, int output_dim, int input_dim)

void gemm_q4_k_q8_k_ref(float * Y, const void * W, const void * X_q8, int M, int N, int K)

void gemm_q4_k_q8_k_thread_fn(int ith, int nth, void * args)

void gemm_q4_k_ref(float * Y, const void * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with Q4_K weights (scalar reference)

void gemm_q4_packed_meta_nsplit_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x16_mreuse_process_job(const gemm_q4_packed_meta_x16_thread_work_t * a, int job, int mt, int tile_m)

void gemm_q4_packed_meta_x16_mreuse_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x16_mtile_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x8_mreuse_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x8_mtile_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x8_nsplit_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x8_split_min_4m_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x8_split_min_8m_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_meta_x8_split_min_mreuse_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_u8_x16_mtile_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_vnni_x16_q8k_16m_thread_fn(int ith, int nth, void * args)

void gemm_q4_packed_vnni_x8_q8k_4m_job(gemm_q4_packed_vnni_x8_thread_work_t * a, int job, int row_tiles)

void gemm_q4_packed_vnni_x8_q8k_4m_range_fn(int begin, int end, void * args)

void gemm_q4_packed_vnni_x8_q8k_4m_thread_fn(int ith, int nth, void * args)

void gemm_q5_0(float * Y, const void * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with Q5_0 weights.

void gemm_q5_0_backward(float * dX, const void * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_q5_1(float * Y, const void * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with Q5_1 weights.

void gemm_q5_1_backward(float * dX, const void * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_q5_k_q8_k_compact_rows4(float * output, int output_stride, const void * weights, const void *const input_rows, int rows, int output_dim, int input_dim)

void gemm_q6_k(float * Y, const void * W, const float * X, int M, int N, int 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.

void gemm_q8_0(float * Y, const void * W, const float * X, int M, int N, int K)

Matrix-matrix multiply with Q8_0 weights.

void gemm_q8_0_backward(float * dX, const void * W, const float * dY, int M, int N, int K)

Batched backward pass.

void gemm_q8_0_q8_0_m2n4(float * C, const void * W, const void * A_q8, int M, int N, int K)

void gemm_q8_0_q8_0_m2n4_strided(float * C, int ldc, const void * W, const void * A_q8, int M, int N, int K)

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)

void gemm_tile_nt_strided(const float * A, const float * B_tile, float * C, int tile_m, int tile_n, int K, int C_stride)

GEMM tile with N-dimension tiling (weight reuse)

void gemm_tn_avx512(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_tn_bf16(const uint16_t * A, const uint16_t * B, const uint16_t * bias, uint16_t * C, int M, int N, int K)

void gemm_tn_blocked(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_tn_parallel(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemm_tn_serial_double(const float * A, const float * B, const float * bias, float * C, int M, int N, int K)

void gemma4_final_logit_softcap_forward(float * logits, int tokens, int vocab_size, float cap)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void gemma4_v_norm_forward(const float * input, float * output, float * rstd_cache, int tokens, int num_kv_heads, int head_dim, float eps)

Forward pass computation

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)

Forward pass computation

void gemma4_vision_projector_prep_forward(const float * input, float * output, int tokens, int dim, float scale, float eps)

Forward pass computation

const GEMMParams * get_gemm_params(void)

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)

void rope_forward_gemma4v_vision_xy_one(float * x, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int grid_w, int rotary_dim, float freq_base)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

void layernorm_forward_ggml_exact(const float * input, const float * gamma, const float * beta, float * output, float * mean_cache, float * rstd_cache, int tokens, int d_model, int input_stride, int output_stride, int aligned_embed_dim, float eps)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void layernorm_forward_unrolled_slice_scalar(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)

Forward pass computation

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)

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)

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)

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)

void zero_layernorm_padding(float * out_ptr, int d_model, int aligned_embed_dim)

void ck_layer_backward_rmsnorm_swiglu(const CKLayerBackwardParams * p)

Backward pass / gradient computation

void ck_layer_forward_rmsnorm_swiglu(const CKLayerForwardParams * p)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode(const CKLayerForwardParams * p, int token_index, int cache_capacity)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode_fused(const CKLayerForwardParams * p, int token_index, int cache_capacity)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode_fused_attn(const CKLayerForwardParams * p, int token_index, int cache_capacity)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl(const CKLayerForwardParams * p, int token_index, int cache_capacity, int fuse_mlp)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_mlp(const CKLayerForwardParams * p, int token_index, int cache_capacity)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode_q4_k(const CKLayerForwardParamsQ4K * p, int token_index, int cache_capacity)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_decode_quant(const CKLayerForwardParamsQ4K * p, int token_index, int cache_capacity)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_q4_k(const CKLayerForwardParamsQ4K * p)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_quant(const CKLayerForwardParamsQ4K * p)

Forward pass computation

void ck_layer_forward_rmsnorm_swiglu_ref(const CKLayerForwardParams * p)

Forward pass computation

void ck_test_rmsnorm(const float * input, const float * weight, float * output, int n_tokens, int dim, float eps)

RMSNorm.

void fused_rmsnorm(const float * input, const float * gamma, const float * beta, float * output, int hidden, float eps)

Fused RMSNorm - writes to pre-allocated buffer.

void fused_rmsnorm_linear_q4k(float * y, const float * x, const float * gamma, const void * W_q4k, int M, int K, float eps)

Fused RMSNorm + Q4_K Linear projection.

void fused_rmsnorm_qkv(const float * input, const float * gamma, const float * W_qkv, const float * b_qkv, float * q_out, float * k_out, float * v_out, int hidden, int num_heads, int num_kv_heads, int head_dim, float eps)

Fused RMSNorm with fused QKV projection.

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.

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)

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)

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.

size_t fused_rmsnorm_qkv_scratch_size(int hidden)

Get scratch buffer size for fused_rmsnorm_qkv_prefill.

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)

void mega_fuse_rmsnorm_qkv(float * q_out, float * k_out, float * v_out, const float * input, const float * gamma, const float * W_qkv, const float * b_qkv, int hidden, int num_heads, int num_kv_heads, int head_dim, float eps)

Phase 1: Fused RMSNorm + QKV (intermediates in registers)

void mega_fuse_rmsnorm_qkv_avx(float * q_out, float * k_out, float * v_out, const float * input, const float * gamma, const float * wq, const float * bq, const float * wk, const float * bk, const float * wv, const float * bv, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, float eps)

Fused RMSNorm + QKV for decode (single token)

void mega_fuse_rmsnorm_qkv_rope(float * q_out, float * k_out, float * v_out, const float * input, const float * gamma, const float * W_qkv, const float * b_qkv, const float * rope_cos, const float * rope_sin, int pos, int hidden, int num_heads, int num_kv_heads, int head_dim, int max_seq, float eps)

Phase 2: Fused RMSNorm + QKV + RoPE.

void qwen4_group_rmsnorm_llama(const float * input, const float * weight, float * output, int groups, int hidden_dim, float eps)

void qwen4_group_rmsnorm_pytorch_bf16(const float * input, const float * weight, float * output, int groups, int hidden_dim, float eps)

void qwen4_shared_head_rmsnorm(const float * input, const float * weight, float * output, int heads, int head_dim, float eps)

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)

RMSNorm backward pass Testtest_rmsnorm.py::TestRMSNormBackward::test_backward_tokens test_rmsnorm.py::TestRMSNormBackward::test_backward_single test_parity.py::test_rmsnorm_backward_parity

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

void rmsnorm_backward_strict_scalar(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)

Backward pass / gradient computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void rmsnorm_forward_no_weight(const float * input, float * output, float * rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void rmsnorm_forward_pytorch_bf16_storage_impl(const float * input, const float * gamma, float * output, float * rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps, int qwen3next_weight_order)

Forward pass computation

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)

Forward pass computation

void rmsnorm_forward_strict_scalar(const float * input, const float * gamma, float * output, float * rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

Forward pass computation

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)

RMSNorm forward pass Testtest_rmsnorm.py::TestRMSNormForward::test_fp32_tokens test_rmsnorm.py::TestRMSNormForward::test_fp32_single test_rmsnorm.py::TestRMSNormForward::test_perf_rolled test_layernorm.py::TestLayerNormForward::test_rmsnorm_compat test_parity.py::test_rmsnorm_parity

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)

Forward pass computation

float rmsnorm_llama_production_rstd(float mean_eps)

void rmsnorm_q8_k_fused(const float * input, const float * gamma, void * vy, int tokens, int d_model, int aligned_embed_dim, float eps)

void rmsnorm_qkv_fp32_fused(const float * x, const float * rms_weight, const float * wq, const float * wk, const float * wv, float * q_out, float * k_out, float * v_out, int embed_dim, int q_dim, int kv_dim, float eps)

void rmsnorm_qkv_fp32_fused_v2(const float * x, const float * rms_weight, const float * wq, const float * wk, const float * wv, float * q_out, float * k_out, float * v_out, int embed_dim, int q_dim, int kv_dim, float eps)

void rmsnorm_qkv_fp32_fused_v3(const float * x, const float * rms_weight, const float * wq, const float * wk, const float * wv, float * q_out, float * k_out, float * v_out, int embed_dim, int q_dim, int kv_dim, float eps)

void rmsnorm_qkv_q4k_fused(const float * x, const float * rms_weight, const void * wq, const void * wk, const void * wv, float * q_out, float * k_out, float * v_out, int embed_dim, int q_dim, int kv_dim, float eps)

void rmsnorm_qkv_separate_fp32(const float * x, const float * rms_weight, const float * wq, const float * wk, const float * wv, float * normed, float * q_out, float * k_out, float * v_out, int embed_dim, int q_dim, int kv_dim, float eps)

void rmsnorm_tile(const float * input, const float * gamma, float * output, int tile_m, int embed_dim, int aligned_embed_dim, float eps)

Compute RMSNorm for a tile of tokens.

void simple_rmsnorm(const float * input, const float * gamma, float * output, int tokens, int d_model, float eps)

void unfused_rmsnorm_linear_q4k_ref(float * y, const float * x, const float * gamma, const void * W_q4k, int M, int K, float eps)

Reference (unfused) implementation for correctness testing.

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.

void ck_gelu_ggml_runtime_init(void)

void ck_gelu_ggml_table_init(void)

void ck_gelu_reference_math_init(void)

ck_gelu_math_f64_fn ck_gelu_system_erf(void)

ck_gelu_math_f32_fn ck_gelu_system_tanhf(void)

float ck_gelu_tanh_f32(float x)

float ck_gelu_tanh_ggml_reference_f32(float x)

float ck_gelu_tanh_parity_f32(float x)

void ck_gelu_try_bind_runtime(void * handle)

float fast_gelu_scalar(float x)

void gelu_backward_exact(const float * input, const float * d_output, float * d_input, size_t n)

Backward pass / gradient computation

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)

Backward pass / gradient computation

void gelu_backward_fast(const float * input, const float * d_output, float * d_input, size_t n)

Backward pass / gradient computation

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)

Backward pass / gradient computation

void gelu_backward_scalar(const float * input, const float * d_output, float * d_input, size_t n)

Backward pass / gradient computation

float gelu_derivative_scalar(float x)

void gelu_erf_bf16_storage(float * data, size_t n)

void gelu_erf_fp64_f32_inplace(float * data, size_t n)

void gelu_exact_inplace(float * data, size_t n)

void gelu_fast_inplace(float * data, size_t n)

GELU activation forward (fast approximation, in-place) Testtest_gelu.py::TestGELUForward::test_gelu_fast_inplace test_gelu.py::TestGELUForward::test_gelu_vs_exact test_parity.py::test_gelu_parity

void gelu_fast_inplace_bf16(uint16_t * data, size_t n, float * scratch)

void gelu_ggml_inplace(float * data, size_t n)

void gelu_ggml_native_inplace(float * data, size_t n)

void gelu_pytorch_erf_f32_inplace(float * data, size_t n)

void gelu_pytorch_erf_sleef_bf16_storage(float * data, size_t n)

void gelu_pytorch_tanh_bf16_storage(float * data, size_t n)

float gelu_scalar(float x)

void backward_causal_softmax_head_major(float * d_scores, const float * weights, int num_heads, int num_tokens, int aligned_context_window)

Backward pass / gradient computation

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)

Backward pass / gradient computation

void causal_softmax_head_major(float * scores, int num_heads, int num_tokens, int aligned_context_window)

Causal softmax (in-place, row-wise) Testtest_softmax.py::TestSoftmaxForward::test_causal_softmax test_softmax.py::TestSoftmaxForward::test_causal_vs_softmax test_attention.py::TestAttentionForward::test_softmax_correctness

void causal_softmax_head_major_bf16(uint16_t * scores, int num_heads, int num_tokens, int aligned_context_window, float * scratch)

void causal_softmax_head_major_exact(float * scores, int num_heads, int num_tokens, int aligned_context_window)

Causal softmax (exact version using stdlib expf) Testtest_softmax.py::TestSoftmaxForward::test_causal_softmax_exact test_softmax.py::TestSoftmaxForward::test_exact_vs_fast

double ck_moe_llama_softmax_row(float * probabilities, const float * logits, int n_experts, float max_value)

void ck_test_softmax(const float * input, float * output, int n)

Softmax (simple, non-causal)

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)

Backward pass / gradient computation

void deepseek_dsa_topk_softmax_f32(const float * scores, int * indices, float * weights, int tokens, int heads, int key_count, int top_k)

void ds_softmax(float * x, int n)

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)

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)

size_t moe_softmax_topk_router_workspace_bytes(int n_experts)

void softmax(float * x, int n)

void softmax_cross_entropy_loss(const float * logits, const int32_t * targets, int tokens, int vocab_size, float * d_logits, float * loss_out)

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)

void softmax_cross_entropy_loss_index_mean_impl(const float * logits, const int32_t * targets, int tokens, int vocab_size, float * d_logits, float * loss_out, int force_strict_math)

void softmax_cross_entropy_loss_legacy_mean_impl(const float * logits, const int32_t * targets, int tokens, int vocab_size, float * d_logits, float * loss_out)

void softmax_cross_entropy_loss_ptref(const float * logits, const int32_t * targets, int tokens, int vocab_size, float * d_logits, float * loss_out)

void softmax_inplace(float * x, int n)

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.

void topk_softmax_f32(const float * scores, int n, int k, int * indices, float * weights)

Find top-K indices with softmax-normalized weights.

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)

Causal attention backward (non-GQA version) Testtest_attention_backward.py::TestAttentionBackward::test_backward test_attention_backward.py::TestAttentionBackward::test_backward_vs_separate test_parity.py::test_attention_backward_parity

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)

GQA causal attention backward (score-matrix version) Testtest_attention_backward.py::TestAttentionBackwardGQA::test_gqa_backward test_attention_backward.py::TestAttentionBackwardGQA::test_gqa_vs_separate test_parity.py::test_attention_backward_parity

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)

BF16 attention backward with caller-provided scratch buffers Testbf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_backward

void attention_flash_cleanup(void)

Clean up flash attention resources.

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.

void attention_flash_decode_scalar(float * out, const float * q, const float * k, const float * v, int T_q, int T_k, int H, int D_h, float scale)

Scalar flash-style attention (online softmax)

void attention_flash_init(int max_context, int max_heads, int max_head_dim)

Initialize flash attention buffers.

void attention_flash_query_causal(const float * q_vec, const float * k_head, const float * v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * out_vec)

void attention_flash_query_causal_exact(const float * q_vec, const float * k_head, const float * v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * out_vec)

void attention_flash_query_causal_exact_f16kv(const float * q_vec, const float * k_head, const float * v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * out_vec)

void attention_flash_query_causal_exact_prerounded_f16kv(const float * q_vec, const float * k_head, const float * v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * out_vec)

void attention_flash_query_sliding(const float * q_vec, const float * k_head, const float * v_head, int query_pos, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * out_vec, int sliding_window)

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)

Causal attention forward (score-matrix version) Testtest_attention.py::TestAttentionForward::test_causal_forward test_attention.py::TestAttentionForward::test_gqa_broadcast test_attention.py::TestAttentionForward::test_exact_vs_fast test_parity.py::test_attention_parity

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)

Causal attention forward (exact version using stdlib expf) Testtest_attention.py::TestAttentionForward::test_exact_single test_attention.py::TestAttentionForward::test_exact_vs_fast

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)

GQA causal attention forward (score-matrix version) Testtest_attention.py::TestAttentionForward::test_gqa_forward test_attention.py::TestAttentionForward::test_gqa_broadcast test_attention_backward.py::TestAttentionBackwardGQA::test_gqa_backward test_parity.py::test_attention_gqa_parity

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)

BF16 GQA causal attention forward Testbf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_forward bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_gqa bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_flash

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)

GQA causal attention forward (exact version using stdlib expf) Testtest_attention.py::TestAttentionForward::test_gqa_exact bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_gqa

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)

Forward pass computation

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)

Flash attention forward with custom KV stride (for KV cache) Testtest_flash_attention.py::TestFlashAttention::test_flash_strided test_kv_cache_attention.py::TestKVCacheAttention::test_flash_attention

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)

Forward pass computation

void attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial(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)

Forward pass computation

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)

Forward pass computation

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)

Flash attention forward with sliding window (prefill) Testtest_attention.py::TestAttentionForward::test_sliding_window_prefill

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Flash attention decode (single token attends to KV cache) Testtest_flash_attention.py::TestFlashAttention::test_flash_decode test_kv_cache_attention.py::TestKVCacheAttention::test_flash_decode test_fused_attention_decode.py::TestFusedAttentionDecode::test_flash_decode test_attention.py::TestAttentionForward::test_flash_decode

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned(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, int partition_tokens)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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!

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void attention_forward_head_major_gqa_flash_impl(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 causal, int round_full_kv_fp16, int output_token_major, float scale)

Forward pass computation

int attention_forward_head_major_gqa_unfused_f16_strict(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 causal, int output_token_major, float scale, int debug_layer_id)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void attention_mlp_fused_fp32(const float * q, const float * k_cache, const float * v_cache, int seq_len, int num_heads, int num_kv_heads, int head_dim, float attn_scale, const float * wo, const float * residual_1, const float * rms_weight, float eps, const float * w_gate, const float * w_up, const float * w_down, int embed_dim, int intermediate_dim, float * hidden_out)

void attention_mlp_fused_q4k(const float * q, const float * k_cache, const float * v_cache, int seq_len, int num_heads, int num_kv_heads, int head_dim, float attn_scale, const void * wo, const float * residual_1, const float * rms_weight, float eps, const void * w_gate, const void * w_up, const void * w_down, int embed_dim, int intermediate_dim, float * hidden_out)

void attention_mlp_separate_fp32(const float * q, const float * k_cache, const float * v_cache, int seq_len, int num_heads, int num_kv_heads, int head_dim, float attn_scale, const float * wo, const float * residual_1, const float * rms_weight, float eps, const float * w_gate, const float * w_up, const float * w_down, int embed_dim, int intermediate_dim, float * attn_out_buf, float * hidden_after_attn_buf, float * normed_buf, float * gate_buf, float * up_buf, float * mlp_out_buf, float * hidden_out)

size_t attention_output_index(int h, int token, int num_heads, int num_tokens, int aligned_head_dim, int token_major)

void attention_query_full_exact_regular(const float * q_vec, const float * k_head, const float * v_cols, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * score_row, float * out_vec, int layer_id, int head_id, int query_id)

void attention_query_full_ggml_regular(const float * q_vec, const float * k_head, const float * v_cols, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * score_row, float * out_vec, int layer_id, int head_id, int query_id)

size_t ck_attention_align64_size(size_t value)

void ck_attention_bf16_pytorch_flash_work(int ith, int nth, void * opaque)

int ck_attention_bf16_pytorch_gqa_available(void)

void ck_attention_bf16_sdpa_work(int ith, int nth, void * opaque)

void ck_attention_causal_f16kv_work(int ith, int nth, void * opaque)

float ck_attention_dot_f16_llama(const uint16_t * x, const uint16_t * y, int n)

float ck_attention_dot_f16_unfused_llama(const uint16_t * x, const uint16_t * y, int n)

void ck_attention_f16_prefill_gqa_reuse_work(int ith, int nth, void * opaque)

size_t ck_attention_f16_prefill_gqa_reuse_worker_bytes(int num_heads, int num_kv_heads, int head_dim, int workers, int query_tile_size, int concurrent_query_tiles)

ck_attention_status_t ck_attention_f16_prefill_qtile64_dispatch(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 cache_is_bf16, size_t q_head_stride, size_t output_head_stride, ck_attention_prefill_schedule_t schedule)

void ck_attention_f16_prefill_qtile64_work(int ith, int nth, void * opaque)

float ck_attention_f16_reduce_expf(float value)

void ck_attention_f16_split_work(int ith, int nth, void * opaque)

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.

void ck_attention_flash_query_auto(const float * q_vec, const float * k_head, const float * v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float * out_vec)

ck_attention_status_t ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_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, ck_attention_prefill_schedule_t schedule)

Forward pass computation

void ck_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, int query_tile_size)

Forward pass computation

int ck_attention_forward_query_key_head_major_f32_run(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)

Forward pass computation

int ck_attention_full_bf16_pytorch_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, int kv_stride_tokens, int output_token_major)

int ck_attention_full_bf16_sdpa_amx_range(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 head_begin, int head_step, int output_token_major)

int ck_attention_full_bf16_sdpa_tiled(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 ck_attention_full_bf16_sdpa_tiled_range(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 head_begin, int head_step)

int ck_attention_full_ggml_graph_oracle_multihead(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 scale)

void ck_attention_full_grid_work(int ith, int nth, void * opaque)

void ck_attention_full_tiled_f16kv_fp32_range(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 query_tile_size, int ith, int nth)

void ck_attention_full_tiled_f16kv_fp32_work(int ith, int nth, void * opaque)

int ck_attention_ggml_out_graph_enabled(void)

void ck_attention_gqa_team_barrier_wait(ck_attention_gqa_team_barrier_t * barrier)

int ck_attention_head_full_ggml_graph_oracle_regular(const float * q_head, const float * k_head, const float * v_head, float * out_head, int num_tokens, int head_dim, int aligned_head_dim, float scale)

void ck_attention_llama_regular_impl(const float * q, const float * k, const float * v, float * output, int num_heads, int num_kv_heads, int query_tokens, int live_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, int batched_prefill)

void ck_attention_llama_regular_query(const float * query, const float * key_head, const float * value_columns, float * output, float * scores, float * scaled_scores, int live_tokens, int padded_tokens, int query_position, int head_dim, int aligned_head_dim, int sliding_window, int batched_prefill)

void ck_attention_mad_f16_llama(uint16_t * y, const uint16_t * x, float scale, int n)

void ck_attention_matmul_f32_accum(float * c, const float * a, const float * b, int m, int k, int n)

int ck_attention_oracle_dump_enabled(void)

int ck_attention_oracle_dump_layer_id(void)

void ck_attention_oracle_dump_meta(const char * name, int layer_id, const struct ggml_tensor * t)

void ck_attention_oracle_dump_tensor(const char * name, int layer_id, const struct ggml_tensor * t)

int ck_attention_oracle_exact_dump_layer(int * layer_id_out)

int ck_attention_oracle_meta_dump_enabled(void)

size_t ck_attention_oracle_qkv_index(int h, int t, int d, int num_tokens, int aligned_head_dim)

int ck_attention_oracle_should_dump_layer(int layer_id)

float ck_attention_oracle_tensor_f32_at(const struct ggml_tensor * t, size_t i0, size_t i1, size_t i2, size_t i3)

int ck_attention_parallel_enabled(int total_queries, int num_tokens, int head_dim)

int ck_attention_pick_active_threads(const ck_threadpool_t * pool, int total_queries, int num_tokens)

void ck_attention_project_head_major(const float * attn_out, const float * wo, const float * bo, float * out, float * scratch, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void ck_attention_project_head_major_backward(const float * d_out, const float * attn_out, const float * wo, float * d_attn_out, float * d_wo, float * d_bo, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

Backward pass / gradient computation

void ck_attention_project_head_major_decode_token(const float * attn_token, const float * wo, const float * bo, float * out_token, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void ck_attention_project_head_major_decode_token_residual(const float * attn_token, const float * wo, const float * bo, const float * residual_in, float * proj_out, float * residual_out, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void ck_attention_project_head_major_q4_k(const float * attn_out, const void * wo, const float * bo, float * out, float * scratch, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void ck_attention_project_head_major_q4_k_q8_k(const float * attn_out, const void * wo, const float * bo, float * out, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void ck_attention_project_head_major_quant(const float * attn_out, const void * wo, const float * bo, float * out, float * scratch, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim, CKDataType wo_dtype)

void ck_attention_project_head_major_ref(const float * attn_out, const float * wo, const float * bo, float * out, float * scratch, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

float ck_attention_pytorch_sdpa_scale_f32(int head_dim)

void ck_attention_query_key_f32_transpose_work(int ith, int nth, void * opaque)

void ck_attention_query_key_f32_work(int ith, int nth, void * opaque)

float ck_attention_reference_expf(float value)

int ck_attention_reverse_out_dot_enabled(void)

void ck_attention_scale_f16_llama(uint16_t * y, float scale, int n)

int ck_attention_sparse_bf16_pytorch_gqa_available(void)

float ck_attention_strict_scale_f32(int head_dim)

int ck_attention_strict_unfused_f16_enabled(void)

void ck_attention_trace(const char * branch, int layer_id, int head_id)

void ck_attention_trace_float(const char * tag, int layer_id, int head_id, float value)

void ck_attention_trace_query(const char * tag, int layer_id, int head_id, int query_id, int value)

float ck_attention_u16_cache_to_f32(uint16_t value, int cache_is_bf16)

int ck_attention_vec_dump_enabled(void)

void ck_attention_vec_dump_exact_query(const float * q_vec, const float * k_head, const float * out_vec, int kv_tokens, int head_dim, int aligned_head_dim, float scale, int layer_id, int head_id, int query_id)

int ck_attention_vec_dump_next_layer_id(void)

int ck_attention_vec_dump_parse_env_int(const char * name, int * out)

void ck_attention_vec_dump_selected_query(const float * raw_scores, const float * probs, const float * out_vec, const float * v_cols, int kv_tokens, int head_dim, int layer_id, int head_id, int query_id)

int ck_attention_vec_dump_should_emit(int layer_id, int head_id, int query_id)

void ck_attention_vec_dump_tensor(const char * name, int layer_id, int query_id, const float * data, size_t elem_count)

int ck_attention_vec_dump_vcols_enabled(void)

void ck_sliding_attention_compute_one(const ck_sliding_attention_args_t * a, int job)

int ck_sliding_attention_parallel_disabled(void)

int ck_sliding_attention_pick_threads(ck_threadpool_t * pool, int total_jobs, int num_tokens, int head_dim)

void ck_sliding_attention_work_fn(int ith, int nth, void * args)

void ck_test_attention_causal(const float * q, const float * k, const float * v, float * out, int num_heads, int num_kv_heads, int tokens, int seq_len, int head_dim)

Multi-head causal attention for prefill (head-major layout)

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)

Backward pass / gradient computation

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)

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)

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)

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)

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)

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)

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)

void ds_mla_attention_f32_query_range(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, int query_begin, int query_end, int query_step)

void ds_mla_attention_f32_work(int ith, int nth, void * opaque)

void fused_flash_attention_all_heads(float * o_out, const float * q_all, const float * kv_cache_k, const float * kv_cache_v, int num_heads, int num_kv_heads, int head_dim, int seq_len, int kv_tile_size)

Fused Flash Attention for all heads (parallel dispatch)

void fused_flash_attention_head(float * o_out, const float * q, const float * kv_cache_k, const float * kv_cache_v, int kv_head_idx, int seq_len, int head_dim, int kv_tile_size)

Fused Flash Attention for single head.

void mega_fuse_flash_attention_avx(float * o_out, const float * q, const float * kv_cache_k, const float * kv_cache_v, int num_heads, int num_kv_heads, int seq_len, int cache_capacity, int head_dim, int aligned_head_dim)

Flash attention with online softmax (AVX version)

void mega_fused_attention(float * output, const float * input, const float * residual, const float * W_qkv, const float * b_qkv, const float * W_o, const float * b_o, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int pos, int seq_len, int hidden, int num_heads, int num_kv_heads, int head_dim, int max_seq, float eps)

Complete mega-fused attention block.

void mega_fused_attention_decode(float * output, const float * input, const float * residual, const float * ln1_gamma, const float * wq, const float * bq, const float * wk, const float * bk, const float * wv, const float * bv, const float * wo, const float * bo, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int pos, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int cache_capacity, float eps)

Mega-fused attention for decode mode (single token)

void mega_fused_attention_decode_q5_0(float * output, const float * input, const float * residual, const void * wq_q5_0, const void * wk_q5_0, const void * wv_q8_0, const void * wo_q5_0, const float * ln_gamma, const float * bq, const float * bk, const float * bv, const float * bo, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int pos, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int cache_capacity, float eps, void * scratch)

Serial mega-fused attention decode kernel.

void mega_fused_attention_decode_q5_0_parallel_simd(float * output, const float * input, const float * residual, const void * wq_q5_0, const void * wk_q5_0, const void * wv_q8_0, const void * wo_q5_0, const float * ln_gamma, const float * bq, const float * bk, const float * bv, const float * bo, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int pos, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int cache_capacity, float eps, void * scratch, int ith, int nth)

Parallel SIMD mega-fused attention decode kernel (threadpool-aware)

int mega_fused_attention_decode_scratch_size(int AE, int H, int KV, int AD)

Calculate scratch buffer size needed for the kernel.

void mega_fused_attention_decode_workspace(float * output, const float * input, const float * residual, const float * ln1_gamma, const float * wq, const float * bq, const float * wk, const float * bk, const float * wv, const float * bv, const float * wo, const float * bo, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int pos, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int cache_capacity, float eps, float * q_output_workspace, size_t q_output_workspace_bytes, float * kv_workspace, size_t kv_workspace_bytes)

Full mega-fused attention for decode.

void mega_fused_attention_prefill(float * output, const float * input, const float * residual, const float * ln1_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, const void * wo, const float * bo, CKDataType wo_dt, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int start_pos, int tokens, int cache_capacity, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, float eps, void * scratch)

Mega-fused attention for prefill mode (multiple tokens)

void mega_fused_attention_prefill_q8_0(float * output, const float * input, const float * residual, const float * ln1_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, const void * wo, const float * bo, CKDataType wo_dt, float * kv_cache_k, float * kv_cache_v, const float * rope_cos, const float * rope_sin, int start_pos, int tokens, int cache_capacity, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, float eps, void * scratch)

Mega-fused prefill attention kernel (Q8_0 out-proj)

size_t mega_fused_attention_prefill_q8_0_scratch_size(int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

Get scratch buffer size for mega_fused_attention_prefill_q8_0.

size_t mega_fused_attention_prefill_scratch_size(int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

Get scratch buffer size for mega_fused_attention_prefill.

void simple_attention(const float * q, const float * k, const float * v, float * output, int num_heads, int num_kv_heads, int seq_len, int head_dim)

void ck_mlp_swiglu_forward(const float * input, const float * w1, const float * b1, const float * w2, const float * b2, float * fc1_out, float * swiglu_out, float * output, int tokens, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_fully_fused_token(const float * input_row, const float * w1, const float * b1, const float * w2, const float * b2, float * output_row, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_fused_token(const float * input_row, const float * w1, const float * b1, const float * w2, const float * b2, float * swiglu_row, float * output_row, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_q4_k(const float * input, const void * w1, const float * b1, const void * w2, const float * b2, float * fc1_out, float * swiglu_out, float * output, int tokens, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_q4_k_q8_k(const float * input, const void * w1, const float * b1, const void * w2, const float * b2, float * fc1_out, float * swiglu_out, float * output, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_q4_k_q8_k_prefill(const float * input, const void * w1, const float * b1, const void * w2, const float * b2, float * fc1_out, float * swiglu_out, float * output, int tokens, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_quant(const float * input, const void * w1, const float * b1, CKDataType w1_dtype, const void * w2, const float * b2, CKDataType w2_dtype, float * fc1_out, float * swiglu_out, float * output, int tokens, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_mlp_swiglu_forward_ref(const float * input, const float * w1, const float * b1, const float * w2, const float * b2, float * fc1_out, float * swiglu_out, float * output, int tokens, int aligned_embed_dim, int aligned_intermediate_dim)

Forward pass computation

void ck_test_outproj_mlp_fused_q5_0(const float * attn_out, const float * residual, const float * ln2_gamma, const void * wo, const void * w1, const void * w2, float * output, int tokens, int num_heads, int head_dim, int embed_dim, int intermediate, float eps, int w2_is_q6k)

Test mega-fused OutProj + MLP kernel (Q5_0 weights)

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)

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)

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)

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.

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.

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)

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.

size_t fused_mlp_swiglu_scratch_size(int intermediate)

Get scratch buffer size for fused_mlp_swiglu_prefill.

void layer_fused_attn_mlp_qkv_q4k(const float * q, const float * k_cache, const float * v_cache, int seq_len, float attn_scale, const void * wo, const float * rms_weight_mlp, const void * w_gate, const void * w_up, const void * w_down, const float * rms_weight_attn, const void * wq_next, const void * wk_next, const void * wv_next, const float * residual_in, int embed_dim, int intermediate_dim, int num_heads, int num_kv_heads, int head_dim, float eps, float * q_next, float * k_next, float * v_next, float * hidden_out)

void mega_fused_outproj_mlp_prefill(float * output, const float * attn_out, const float * residual, const float * ln2_gamma, const void * wo, const float * bo, int wo_dt, const void * w1, const float * b1, int w1_dt, const void * w2, const float * b2, int w2_dt, int tokens, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim, int intermediate_dim, int aligned_intermediate_dim, float eps, void * scratch)

size_t mega_fused_outproj_mlp_prefill_scratch_size(int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim, int aligned_intermediate_dim)

Get scratch buffer size for mega_fused_outproj_mlp_prefill.

void mlp_fused_fp32_v2(const float * hidden_in, const float * rms_weight, float eps, const float * w_gate, const float * w_up, const float * w_down, int embed_dim, int intermediate_dim, float * hidden_out)

void mlp_fused_fp32_v3(const float * hidden_in, const float * rms_weight, float eps, const float * w_gate, const float * w_up, const float * w_down, int embed_dim, int intermediate_dim, float * hidden_out)

void mlp_parallel(const void * ln2_q8, const void * W_gate, const void * W_up, const void * W_down, float * gate_buf, float * up_buf, float * swiglu_buf, void * down_q8, float * mlp_out, int intermediate, int embed_dim, int num_threads)

Parallel MLP (gate/up + SwiGLU + down projection).

int mlp_q8_0_dtype_supported(CKDataType dt)

int mlp_q8_k_dtype_supported(CKDataType dt)

void mlp_separate_fp32(const float * hidden_in, const float * rms_weight, float eps, const float * w_gate, const float * w_up, const float * w_down, float * normed_buf, float * gate_buf, float * up_buf, int embed_dim, int intermediate_dim, float * hidden_out)

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)

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)

Optimized MLP Forward (BF16 weights, FP32 activations)

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)

BF16 MLP backward with FP32 gradient accumulation.

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)

Alternative: Fully FP32 activations throughout

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)

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)

Backward pass / gradient computation

void attn_gate_sigmoid_mul_forward(const float * x, const float * gate, float * out, int rows, int num_heads, int state_dim)

Forward pass computation

void attn_gate_sigmoid_mul_pytorch_bf16_storage(const float * x, const float * gate, float * out, int rows, int num_heads, int state_dim)

float ck_deltanet_llama_sigmoidf(float x)

float ck_deltanet_sigmoidf(float x)

float ck_moe_sigmoid_f32(float x)

float ck_sigmoid_bf16(float value)

void ck_test_attn_gate_sigmoid_mul(const float * x, const float * gate, float * out, int rows, int dim)

Multiply attention output rows by sigmoid(gate) elementwise.

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)

float hybrid_sigmoid(float x)

float mamba2_sigmoid_f32(float x)

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)

Forward pass computation

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)

float recurrent_sigmoid(float x)

void recurrent_sigmoid_forward_ggml(const float * x, float * out, int rows, int dim)

Forward pass computation

void recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(const float * x, float * out, int rows, int dim)

Forward pass computation

float recurrent_sigmoid_local(float x)

void sigmoid_backward(const float * input, const float * d_output, float * d_input, size_t n)

Backward pass / gradient computation

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)

Backward pass / gradient computation

void sigmoid_forward(const float * input, float * output, size_t n)

Forward pass computation

void sigmoid_forward_bf16(const uint16_t * input, uint16_t * output, size_t n, float * scratch_input, float * scratch_output)

Forward pass computation

float sigmoid_scalar(float x)

float sigmoid_scalar_parity(float x)

int ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl(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)

Forward pass computation

int ck_moe_swiglu_nvfp4_projection(const float * hidden, const void * gate, float gate_scale, const void * up, float up_scale, const void * down, float down_scale, float * result, int hidden_dim, int intermediate_dim, void * workspace)

void ck_test_swiglu(const float * gate_up, float * output, int n_tokens, int intermediate_dim)

SwiGLU activation.

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)

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)

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)

void moe_swiglu_expert_backward_f32(const float * d_output, const float * hidden, const int * indices, const float * routing_weights, const float * expert_gate, const float * expert_up, const float * expert_down, float * d_hidden, float * d_routing_weights, float * d_expert_gate, float * d_expert_up, float * d_expert_down, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

Backward pass / gradient computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

int moe_swiglu_expert_forward_nvfp4_workspace(const float * hidden, const int * indices, const float * routing_weights, const void * expert_gate, const float * expert_gate_scales, const void * expert_up, const float * expert_up_scales, const void * expert_down, const float * expert_down_scales, float * output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void * workspace, size_t workspace_bytes)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

size_t moe_swiglu_expert_q4k_q5k_bucketed_workspace_bytes(int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

size_t moe_swiglu_expert_q4k_q5k_workspace_bytes(int hidden_dim, int intermediate_dim)

size_t moe_swiglu_expert_q4k_q8_0_workspace_bytes(int hidden_dim, int intermediate_dim)

size_t moe_swiglu_nvfp4_workspace_bytes(int hidden_dim, int intermediate_dim)

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)

Forward pass computation

void moe_swiglu_shared_backward_f32(const float * d_output, const float * hidden, const float * shared_gate, const float * shared_up, const float * shared_down, float * d_hidden, float * d_routed, float * d_shared_gate, float * d_shared_up, float * d_shared_down, int rows, int hidden_dim, int intermediate_dim)

Backward pass / gradient computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

int moe_swiglu_shared_forward_nvfp4_workspace(const float * hidden, const float * routed, const void * shared_gate, const float * shared_gate_scale, const void * shared_up, const float * shared_up_scale, const void * shared_down, const float * shared_down_scale, float * output, int rows, int hidden_dim, int intermediate_dim, float combination_scale, void * workspace, size_t workspace_bytes)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

size_t moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes(int hidden_dim, int intermediate_dim)

size_t moe_swiglu_shared_q8_0_gated_workspace_bytes(int hidden_dim, int intermediate_dim)

void swiglu_backward(const float * input, const float * d_output, float * d_input, int tokens, int dim)

SwiGLU backward pass Testtest_swiglu.py::TestSwiGLUBackward::test_backward_tokens test_swiglu.py::TestSwiGLUBackward::test_backward_single test_parity.py::test_swiglu_backward_parity

void swiglu_backward_bf16(const uint16_t * input, const uint16_t * d_output, uint16_t * d_input, int tokens, int dim)

Backward pass / gradient computation

void swiglu_backward_exact(const float * input, const float * d_output, float * d_input, int tokens, int dim)

SwiGLU backward pass (exact version using stdlib sigmoid) Testtest_swiglu.py::TestSwiGLUBackward::test_exact_vs_fast test_swiglu.py::TestSwiGLUBackward::test_exact_single

void swiglu_forward(const float * input, float * output, int tokens, int dim)

SwiGLU forward pass Testtest_swiglu.py::TestSwiGLUForward::test_forward_tokens test_swiglu.py::TestSwiGLUForward::test_forward_single test_mlp.py::TestMLPForward::test_swiglu_mlp test_fused_swiglu_decode.py::TestFusedSwiGLUDecode::test_fused_swiglu_decode test_parity.py::test_swiglu_parity

void swiglu_forward_bf16(const uint16_t * input, uint16_t * output, int tokens, int dim)

Forward pass computation

void swiglu_forward_exact(const float * input, float * output, int tokens, int dim)

SwiGLU forward pass (exact version using stdlib sigmoid) Testtest_swiglu.py::TestSwiGLUForward::test_exact_vs_fast test_swiglu.py::TestSwiGLUForward::test_exact_single

void swiglu_forward_ggml(const float * input, float * output, int tokens, int dim)

Forward pass computation

void swiglu_forward_ggml_split(const float * gate, const float * up, float * output, int tokens, int dim)

Forward pass computation

void swiglu_forward_pytorch_bf16_storage(const float * input, float * output, int tokens, int dim)

Forward pass computation

void swiglu_forward_q8_k(const float * input, void * output_q8, int tokens, int dim)

Forward pass computation

void apply_rope(float * x, int seq_len, int head_dim)

void apply_rope_inline(float * q, float * k, const float * rope_cos, const float * rope_sin, int pos, int H, int KV, int AD)

void ck_model_precompute_rope(void * model)

Precompute RoPE cos/sin caches. Call once after allocation, before inference.

void ck_mrope_round_storage(float * data, size_t count, int storage_kind)

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)

ck_ggml_rope_multi_inplace_fn ck_resolve_ggml_rope_multi_inplace(void)

void ck_rope_ensure_ggml_loaded(void)

float ck_rope_reference_cosf(float value)

float ck_rope_reference_powf(float base, float exponent)

float ck_rope_reference_sinf(float value)

void * ck_rope_resolve_ggml_symbol(const char * name)

ck_rope_math_f32_fn ck_rope_resolve_system_math_f32(const char * name)

void ck_test_rope(float * q, float * k, int n_tokens, int n_heads, int n_heads_kv, int head_dim, int pos_offset, float theta)

RoPE (Rotary Position Embedding)

void ck_test_rope_interleaved(float * q, float * k, int n_tokens, int n_heads, int n_heads_kv, int head_dim, int pos_offset, float theta)

RoPE with interleaved format (for llama.cpp compatibility)

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)

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)

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)

void ds_mla_apply_kimi_rope(const float * src, float * dst, const float * cos_row, const float * sin_row, int dim)

int explicit_mrope_apply_ggml_exact(float * x, const int32_t * positions, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, const int sections, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, int rope_type)

void explicit_mrope_apply_head(float * x, const int32_t * positions, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, const int sections, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, int is_imrope)

void fused_rope_inplace(float * q, float * k, const float * rope_cos, const float * rope_sin, int pos, int num_heads, int num_kv_heads, int head_dim, int max_seq)

Fused RoPE application (in-place on pre-allocated buffers)

void mega_fuse_rope_inplace_avx(float * q, float * k, const float * rope_cos, const float * rope_sin, int pos, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim)

Apply RoPE to Q and K (in-place, from L1)

void model_precompute_rope(MODELModel * model)

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)

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)

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)

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)

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)

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)

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)

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)

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)

void mrope_rotate_pair(float x0, float x1, float cos_theta, float sin_theta, float * out0, float * out1)

void qwen2_0_5b_decode_precompute_rope(QWEN2_0_5B_DECODEModel * model)

void qwen4_rope_split_inplace(float * vector, int rotary_dim, int position, float theta)

void rope_apply_decode_pairwise_llama_cpu(float * rows, const float * cos_row, const float * sin_row, int num_heads, int aligned_head_dim, int rotary_dim)

void rope_apply_head(float * x, const float * cos_cache, const float * sin_cache, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

void rope_apply_head_pairwise(float * x, const float * cos_cache, const float * sin_cache, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

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)

RoPE backward (inverse rotation) Testtest_rope.py::TestRoPEBackward::test_rope_backward test_rope.py::TestRoPEBackward::test_rope_backward_vs_separate

void rope_backward_apply_head_pairwise(const float * d_out, float * d_x, const float * cos_cache, const float * sin_cache, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

Backward pass / gradient computation

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)

Backward pass / gradient computation

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)

RoPE backward in-place (overwrite with inverse rotation) Testtest_rope.py::TestRoPEBackward::test_rope_backward_inplace

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)

RoPE backward for both dQ and dK Testtest_rope.py::TestRoPEBackward::test_rope_backward_qk

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

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)

RoPE forward (head-major layout, in-place) Testtest_rope.py::TestRoPEForward::test_rope_forward test_rope.py::TestRoPEForward::test_rope_vs_separate test_parity.py::test_rope_parity

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

RoPE forward for both Q and K (common inference pattern) Testtest_rope.py::TestRoPEForward::test_rope_forward_qk test_fused_attention_decode.py::TestFusedAttentionDecode::test_qk_rope test_parity.py::test_rope_qk_parity

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

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)

RoPE forward for both Q and K with custom strides (KV cache layouts) Testtest_rope.py::TestRoPEForward::test_rope_forward_qk_strided test_kv_cache_attention.py::TestKVCacheAttention::test_qk_rope_strided

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)

Forward pass computation

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)

Forward pass computation

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)

Forward pass computation

void rope_forward_split_direct_one(float * x, 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)

Forward pass computation

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)

RoPE forward with custom head stride (for KV cache layouts) Testtest_rope.py::TestRoPEForward::test_rope_strided test_kv_cache_attention.py::TestKVCacheAttention::test_rope_decode

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)

Forward pass computation

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)

Forward pass computation

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)

Precompute RoPE cos/sin cache with rotary_dim and scaling support Testtest_rope.py::TestRoPECache::test_cache_computation test_rope.py::TestRoPECache::test_cache_values

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)

void rope_precompute_cache_split(float * cos_cache, float * sin_cache, int max_seq_len, int head_dim, float base)

Precompute RoPE cos/sin cache (split layout: head_dim/2) Legacy layout used before rotary_dim/scaling support.

void text_mrope_apply_head(float * x, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, const int sections, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, int is_imrope)

void text_mrope_apply_positions_pytorch_bf16_storage(float * x, const int32_t * positions, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, const int sections, float freq_base, float freq_scale)

void text_mrope_apply_pytorch_bf16_storage(float * x, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, float freq_base, float freq_scale)

void text_mrope_yarn(float theta_extrap, float freq_scale, const float corr_dims, int chan, float ext_factor, float attn_factor, float * cos_theta, float * sin_theta)

void vision_mrope_apply_head(float * x, const int32_t * positions, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, const int sections, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void vision_mrope_yarn(float theta_extrap, float freq_scale, const float corr_dims, int chan, float ext_factor, float attn_factor, float * cos_theta, float * sin_theta)

float vision_mrope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base)

void vision_mrope_yarn_corr_dims(int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims)

float vision_mrope_yarn_ramp(float low, float high, int chan)

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)

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)

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)

void yarn_rope_cache_explicit_positions_impl(float * cos_f32, float * sin_f32, uint16_t * cos_bf16, uint16_t * sin_bf16, 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)

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)

Backward pass / gradient computation

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)

Backward pass / gradient computation

 _Static_assert(sizeof(MagicHeader), "MagicHeader must be 64 bytes")

 __attribute__((unused))

void accum_q4_k_packed_meta_x16_q8_k_block(float acc, const block_q4_K_packed_meta_x16 * w, int active, const block_q8_K * x)

void accum_q4_k_packed_meta_x16_q8_k_block_mreuse(float acc, const block_q4_K_packed_meta_x16 * w, int active, const block_q8_K * A, int blocks_per_vec, int block_index, int m0, int m_count)

void accum_q4_k_packed_meta_x16_q8_k_block_mreuse_chunk4(float acc, const block_q4_K_packed_meta_x16 * w, int active, const block_q8_K * A, int blocks_per_vec, int block_index, int m0, int m_count)

void accum_q4_k_packed_meta_x8_q8_k_block(float acc, const block_q4_K_packed_meta_x8 * w, int active, const block_q8_K * x)

void accum_q4_k_packed_meta_x8_q8_k_block_mreuse(float acc, const block_q4_K_packed_meta_x8 * w, int active, const block_q8_K * A, int blocks_per_vec, int block_index, int m0, int m_count)

void accum_q4_k_packed_meta_x8_q8_k_gemv_block(float acc, float acc_min, const block_q4_K_packed_meta_x8 * w, int active, const block_q8_K * x)

void accum_q4_k_packed_meta_x8_q8_k_superblock(float acc, float acc_min, const block_q4_K_packed_meta_x8 * w, int active, const block_q8_K * x)

void accum_q4_k_packed_meta_x8_q8_k_superblock_rows(float acc, float acc_min, const block_q4_K_packed_meta_x8 * w, int active, const block_q8_K * x, int rows)

void accum_q4_k_packed_u8_x16_q8_k_block(float acc, const block_q4_K_packed_u8_x16 * w, int active, const block_q8_K * x)

void accum_q4_k_packed_vnni_x16_q8_k_16m_superblock(float acc, float acc_min, const block_q4_K_packed_vnni_x16 * w, const block_q8_K * x, int rows)

void accum_q4_k_packed_vnni_x16_q8_k_gemv_block(float acc, float acc_min, const block_q4_K_packed_vnni_x16 * w, const block_q8_K * x)

void accum_q4_k_packed_vnni_x8_q8_k_4m_superblock(float acc, float acc_min, const block_q4_K_packed_vnni_x8 * w, const block_q8_K * x, int rows)

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)

void adamw_update_bf16(const uint16_t * grad, uint16_t * weight, float * m, float * v, size_t numel, float lr, float beta1, float beta2, float eps, float weight_decay, int step)

AdamW optimizer update (bf16 weights/gradients, fp32 optimizer state)

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)

void adamw_update_f32_impl(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)

AdamW optimizer update (fp32 version)

void add_backward_bf16(const uint16_t * d_y, uint16_t * d_a, uint16_t * d_b, size_t n)

Backward pass / gradient computation

void add_bias_tile(float * out, const float * bias, int tile_m, int out_dim)

void add_forward_2d_bf16(const uint16_t * a, const uint16_t * b, uint16_t * y, int tokens, int dim, int aligned_dim)

Forward pass computation

void add_forward_bf16(const uint16_t * a, const uint16_t * b, uint16_t * y, size_t n)

Forward pass computation

void add_forward_f32(const float * a, const float * b, float * y, size_t n)

Element-wise add: y = a + b Testtest_add.py::TestAddForward::test_add_forward_f32 test_add.py::TestAddForward::test_add_inplace_f32 test_multi_layer_parity.py::TestMultiLayerParity::test_residual_add

void add_inplace_bf16(uint16_t * a, const uint16_t * b, size_t n)

void add_inplace_f32(float * a, const float * b, size_t n)

void add_scaled_forward_bf16(const uint16_t * a, const uint16_t * b, uint16_t * y, float alpha, size_t n)

Forward pass computation

void add_scaled_inplace_bf16(uint16_t * a, const uint16_t * b, float alpha, size_t n)

void add_stream_inplace(float * a, const float * b, size_t n)

void add_stream_reorder_2d(float * main_inout, float * aux_scratch, int grid_h, int grid_w, int embed_dim, int merge_size)

size_t align_up(size_t n, size_t align)

size_t align_up_bytes(size_t n, size_t align)

size_t align_up_elems(size_t elems, size_t elem_bytes, size_t align_bytes)

size_t align_up_size(size_t value, size_t align)

bool amx_available(void)

int apply_bpe_merges(CKTrueBPE * bpe, CKBPETokenList * list)

char * apply_chat_template(const ChatTemplate * tmpl, const char * system, const char * user)

CKMemArenaKind arena_for_role(CKBufferRole role)

int argmax_f32(const float * scores, int n)

Find index of maximum value.

void assistant_layer_scale_forward(float * hidden, const float * scale, int tokens, int embed_dim)

Forward pass computation

void attn_gate_softplus_mul_forward(const float * x, const float * gate, float * out, int rows, int num_heads, int state_dim)

Forward pass computation

int audio_conv1d_channel_major_f32(const float * input, const float * weight, const float * bias, float * output, int input_channels, int output_channels, int input_frames, int kernel_size, int stride, int padding, int output_frames)

int audio_conv2d_whc_grouped_f32(const float * input, const float * weight, const float * bias, float * output, int input_width, int input_height, int input_channels, int output_channels, int kernel_width, int kernel_height, int stride_width, int stride_height, int padding_width, int padding_height, int groups, int output_width, int output_height)

int audio_feature_normalize_per_feature_f32(const float * input, float * output, int channels, int frames, float epsilon)

int audio_glu_split_channel_major_f32(const float * input, float * output, int channels, int frames)

double audio_hz_to_mel_slaney(double hz)

int audio_log_mel_time_major_f32(const float * power, const float * mel_filters, float * log_mel, int frames, int bins, int channels, float epsilon)

double audio_mel_to_hz_slaney(double mel)

int audio_pad_or_truncate_f32(const float * input, int input_frames, float * output, int output_frames)

int audio_pcm_s16_to_mono_f32(const int16_t * interleaved, int n_frames, int n_channels, float * mono)

int audio_preemphasis_f32(const float * input, float * output, int frames, float coefficient)

int audio_relative_shift_f32(const float * raw_scores, float * scores, int heads, int query_frames)

int audio_resample_linear_f32(const float * input, int input_frames, int input_rate, float * output, int output_frames, int output_rate)

int audio_resample_windowed_sinc_f32(const float * input, int input_frames, int input_rate, float * output, int output_frames, int output_rate, int radius)

int audio_resampled_frame_count(int input_frames, int input_rate, int output_rate)

int audio_stft_power_centered_window_f32(const float * samples, int n_samples, const float * window, int window_length, const float * cos_table, const float * sin_table, int n_fft, int hop_length, int reflect_padding, float * power, int n_frames)

int audio_stft_power_fft400_f32(const float * samples, int n_samples, const float * window, const float * cos_table, const float * sin_table, int hop_length, float * power, int n_frames, float * fft_scratch)

void audio_stft_power_fft400_frame_f32(const float * samples, int n_samples, int frame, const float * window, const float * cos_table, const float * sin_table, float * power, float * fft_scratch)

int audio_stft_power_precomputed_f32(const float * samples, int n_samples, const float * window, const float * cos_table, const float * sin_table, int n_fft, int hop_length, float * power, int n_frames)

int audio_stft_precompute_tables_f32(int n_fft, float * window, float * cos_table, float * sin_table)

int audio_transpose_channel_to_token_f32(const float * input, float * output, int channels, int frames)

int audio_wav_decode_memory_pcm16_mono_f32(const uint8_t * bytes, size_t byte_count, float * mono, int mono_capacity, CKAudioWavInfo * info)

int audio_wav_decode_memory_pcm16_mono_window_f32(const uint8_t * bytes, size_t byte_count, int start_frame, float * mono, int mono_capacity, CKAudioWavInfo * info)

int audio_wav_decode_pcm16_mono_f32(const uint8_t * bytes, size_t byte_count, const CKAudioWavInfo * info, float * mono, int mono_capacity)

int audio_wav_parse_memory(const uint8_t * bytes, size_t byte_count, CKAudioWavInfo * info)

int audio_whisper_log_mel_from_power_reference_f32(const float * power, const float * mel_filters, int n_mels, int n_frames, float * log_mel)

int audio_whisper_log_mel_reference_f32(const float * samples, int n_samples, const float * mel_filters, int n_mels, float * power_scratch, float * log_mel, int n_frames)

int audio_whisper_log_mel_window_wav_pcm16_f32(const uint8_t * bytes, size_t byte_count, int start_frame, int target_sample_rate, const float * window, const float * cos_table, const float * sin_table, const float * mel_filters, int n_mels, int output_frames, float * log_mel)

int audio_whisper_mel_filters_slaney_f32(int sample_rate, int n_fft, int n_mels, float * mel_filters)

int audio_whisper_stft_power_reference_f32(const float * samples, int n_samples, float * power, int n_frames)

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,:].

void axpy_f32(float * y, const float * x, float alpha, int n)

In-place AXPY: y += alpha * x.

void axpy_zero_f32(float * y, const float * x, float alpha, int n)

Zero output then accumulate: y = 0; y += alpha * x.

void barrier_init(ck_barrier_t * b, int n_threads)

void barrier_wait(ck_barrier_t * b)

Spin-wait barrier. All threads must call this. Uses phase counter to allow re-use without reset.

void bf16_tensor_to_float(const uint16_t * src, float * dst, size_t count)

float bf16_to_float(uint16_t v)

size_t buffer_bytes(const CKIRV2Buffer * buf, const CKModelConfig * cfg, const CKV2AlignInfo * align)

int buffer_enabled(const CKIRV2Graph * graph, const CKIRV2Buffer * buf, int training_enabled)

int build_plan(const CKIRV2Graph * graph, CKMemPlan * plan, size_t alignment_bytes, int training_enabled, int tokens_override)

size_t bump_bytes(size_t * off, size_t bytes, size_t align)

int byte_to_gpt2(unsigned char byte, char * out)

int cache_compare(const void * a, const void * b)

int ce_legacy_mode_enabled(void)

int ce_targets_all_valid_no_ignore(const int32_t * targets, int tokens, int vocab_size)

void ck_accum_multi_parallel_work(int ith, int nth, void * argp)

void ck_accum_parallel_work(int ith, int nth, void * argp)

void ck_adamw_multi_parallel_work(int ith, int nth, void * argp)

void ck_adamw_parallel_work(int ith, int nth, void * argp)

void ck_add_inplace(float * dst, const float * src, int tokens, int aligned_embed_dim)

void ck_audio_conv1d_channel_major_f32_work(int ith, int nth, void * opaque)

void ck_audio_conv2d_whc_grouped_f32_range(int begin, int end, void * opaque)

void ck_audio_glu_split_f32_range(int begin, int end, void * opaque)

void ck_audio_relative_shift_f32_range(int begin, int end, void * opaque)

int ck_available_logical_cpus(void)

void ck_bf16_convert_work(int ith, int nth, void * opaque)

float ck_bf16_dot_contract(const float * a, const float * b, int count)

float ck_bf16_round(float value)

void ck_bf16_round_work(int ith, int nth, void * opaque)

float ck_bf16_to_f32(uint16_t v)

void ck_bind_deltanet_llama_libm(void)

void ck_bind_deltanet_pytorch_primitives(void)

void ck_bind_hybrid_llama_libm(void)

void ck_bind_recurrent_llama_libm(void)

int ck_buffer_should_alloc(const CKBufferSpec * spec)

int ck_buffer_uses_weight_dtype(const CKBufferSpec * spec)

int ck_build_decoder_backward_ir(const CKIRGraph * forward, CKIRGraph * backward)

Build a naive backward IR graph from a forward decoder IR.

int ck_build_decoder_ir(const CKModelConfig * cfg, CKIRGraph * graph)

Build a simple decoder-only IR graph for the given config.

void ck_bump_alloc_free(ck_bump_alloc_t * alloc)

int ck_bump_alloc_init(ck_bump_alloc_t * alloc, const char * weights_path, size_t total_size, size_t weights_base, size_t activations_base)

int ck_bump_alloc_needs_weight_materialization(const ck_bump_alloc_t * alloc)

void ck_bump_alloc_reset(ck_bump_alloc_t * alloc)

void ck_codegen_c_skeleton(const CKIRGraph * forward, const CKIRGraph * backward, FILE * out)

Emit a C skeleton for forward + backward execution based on the IR.

int ck_codegen_emit_runtime(const CKIRGraph * forward, const char * path, CKEmitMode mode)

Emit a C runtime file that stitches kernels for the given forward IR.

const char * ck_codegen_v2_dtype_name(CKDataType dtype)

void ck_codegen_v2_emit_dispatch(FILE * out, const CKIRV2Graph * graph)

int ck_codegen_v2_emit_preamble(FILE * out)

int ck_codegen_v2_emit_runtime(const CKIRV2Graph * graph, const char * path, CKEmitMode mode)

Emit a C runtime file from a CKIRV2Graph.

void ck_codegen_v2_emit_schedule(FILE * out, const CKIRV2Graph * graph, const char * prefill_runtime, const char * decode_runtime, const char * backward_runtime)

void ck_codegen_v2_emit_sections(FILE * out, const CKIRV2Graph * graph, const CKMemPlan * prefill_plan, const CKMemPlan * decode_plan, const CKMemPlan * backward_plan)

void ck_codegen_v2_emit_struct(FILE * out, const CKIRV2Graph * graph, const CKMemPlan * plan, const char * tag)

void ck_debug_check_buffer(const char * stage, const float * buf, int size)

void ck_debug_check_q4k_weights(const char * stage, const void * q4_buf, int num_blocks)

void ck_debug_check_q8k(const char * stage, const void * q8_buf, int num_blocks)

int ck_deltanet_ceil_log2(int value)

int ck_deltanet_force_ref(void)

void ck_deltanet_pytorch_gate_values(const float * g, const float * beta, float * gate_values, float * beta_values, int num_heads)

void ck_deltanet_pytorch_outer_sum(const float * matrix, const float * row_weights, float * output, int state_dim)

float ck_dot_f16_f16_local(const uint16_t * w, const uint16_t * x, int k)

float ck_dot_f32(const float * a, const float * b, int len)

float ck_dot_q6_k_q8_k_fast_or_ref(const block_q6_K * w, const block_q8_K * x, int K)

size_t ck_dtype_block_bytes(CKDataType dt)

Get bytes per block for quantized types.

size_t ck_dtype_block_size(CKDataType dt)

Get the number of elements per quantization block.

size_t ck_dtype_bytes(CKDataType dt)

Get bytes per element for non-quantized types.

int ck_dtype_is_quantized(CKDataType dt)

Check if a data type is block-quantized (GGML-style)

size_t ck_dtype_row_bytes(CKDataType dt, size_t n_elements)

Calculate total bytes for n_elements of given dtype.

int ck_dtype_supported(CKDataTypeMask mask, CKDataType dt)

int ck_env_int_default(const char * name, int fallback)

int ck_env_truthy_or_qwen3vl_ocr_profile(const char * name)

int ck_env_value_truthy(const char * v)

float ck_expf(float x)

ck_f16_ggml_build_forward_expand_fn ck_f16_resolve_ggml_build_forward_expand(void)

Forward pass computation

ck_f16_ggml_cpu_init_fn ck_f16_resolve_ggml_cpu_init(void)

ck_f16_ggml_free_fn ck_f16_resolve_ggml_free(void)

ck_f16_ggml_get_data_fn ck_f16_resolve_ggml_get_data(void)

ck_f16_ggml_get_data_f32_fn ck_f16_resolve_ggml_get_data_f32(void)

ck_f16_ggml_graph_compute_with_ctx_fn ck_f16_resolve_ggml_graph_compute_with_ctx(void)

ck_f16_ggml_init_fn ck_f16_resolve_ggml_init(void)

ck_f16_ggml_mul_mat_fn ck_f16_resolve_ggml_mul_mat(void)

ck_f16_ggml_new_graph_fn ck_f16_resolve_ggml_new_graph(void)

ck_f16_ggml_new_tensor_2d_fn ck_f16_resolve_ggml_new_tensor_2d(void)

void ck_f32_to_f16_row_local(uint16_t * dst, const float * src, int n)

float ck_fast_expf(float x)

const CKBufferSpec * ck_find_buffer_spec(const char * name)

const CKKernelSpec * ck_find_kernel_spec(const char * name)

const char * ck_first_layer_buffer_name(void)

int ck_flash_attn_choose_tile_k(int D_h)

int ck_flash_attn_fast_exp_kind(void)

int ck_flash_attn_tile_k(int D_h)

void ck_fma_f32_to_f16(const float * a, const float * b, const float * c, uint16_t * dst, int n)

FMA in FP32, store result as FP16: dst = a * b + c.

float ck_fp16_to_fp32(ck_half h)

void ck_fp16_to_fp32_2d(const uint16_t * src, float * dst, int rows, int cols, int src_stride, int dst_stride)

Convert 2D FP16 matrix to FP32 with strided access.

void ck_fp16_to_fp32_row(const uint16_t * src, float * dst, int n)

Convert FP16 row to FP32 (auto-select best implementation)

float ck_fp16_to_fp32_scalar(uint16_t h)

float ck_fp16_to_fp32_soft(ck_half h)

Convert FP16 (ck_half) to FP32 — software implementation.

float ck_fp32_from_bits(uint32_t u32)

uint32_t ck_fp32_to_bits(float f)

ck_half ck_fp32_to_fp16(float f)

void ck_fp32_to_fp16_2d(const float * src, uint16_t * dst, int rows, int cols, int src_stride, int dst_stride)

Convert 2D FP32 matrix to FP16 with strided access.

void ck_fp32_to_fp16_inplace(float * data, void * scratch, int n)

Convert FP32 to FP16 in-place using scratch buffer.

void ck_fp32_to_fp16_row(const float * src, uint16_t * dst, int n)

Convert FP32 row to FP16 (auto-select best implementation)

uint16_t ck_fp32_to_fp16_scalar(float f)

ck_half ck_fp32_to_fp16_soft(float f)

Convert FP32 to FP16 (ck_half) — software implementation.

void ck_gemv_bf16_rows(int begin, int end, void * opaque)

void ck_gemv_bf16_storage_rows(int begin, int end, void * opaque)

int ck_get_block_q4_k_size(void)

Get Q4_K block size in bytes.

int ck_get_block_q6_k_size(void)

Get Q6_K block size in bytes.

int ck_get_block_q8_k_size(void)

Get Q8_K block size in bytes.

ck_capability_t ck_get_capabilities(void)

Get current platform capabilities.

int ck_get_num_threads(void)

int ck_get_physical_cores(void)

int ck_get_qk_k(void)

Get QK_K (elements per super-block)

ck_threadpool_t * ck_get_threadpool(void)

Get the global thread pool handle for dispatch. Convenience wrapper — initializes on first call.

float ck_ggml_vec_dot_f32_contig(const float * x, const float * y, int n)

double ck_ggml_vec_soft_max_row(int n, float * y, const float * x, float max)

void * ck_huge_alloc(size_t bytes)

Allocate a large, contiguous memory region for model weights/activations.

void ck_huge_free(void * ptr, size_t bytes)

Free memory allocated by ck_huge_alloc.

void ck_ir_dump(const CKIRGraph * graph, FILE * out)

Dump a human-readable view of the IR to the given stream.

void ck_ir_free(CKIRGraph * graph)

Free any heap-allocated memory owned by the graph.

int ck_ir_parse_json(const char * path, CKIRGraph * graph)

Parse a JSON IR map file (as produced by ck_ir_serialize_json) back into a CKIRGraph. This enables a two-stage pipeline:

int ck_ir_serialize_json(const CKIRGraph * graph, const char * path)

Serialize a CKIRGraph to a simple JSON IR map file.

size_t ck_ir_v2_align_up_bytes(size_t n, size_t align)

size_t ck_ir_v2_align_up_elems(size_t elems, size_t elem_bytes, size_t align_bytes)

int ck_ir_v2_apply_meta(const char * path, CKIRV2Graph * graph)

int ck_ir_v2_apply_weight_dtypes(const char * json, const char * end, CKIRV2Graph * graph)

int ck_ir_v2_build_decoder(const CKModelConfig * cfg, CKIRV2Graph * graph)

int ck_ir_v2_build_decoder_backward(const CKIRV2Graph * forward, CKIRV2Graph * backward)

Backward pass / gradient computation

int ck_ir_v2_copy_buffer_spec(const CKBufferSpec * spec, CKIRV2Buffer * out)

void ck_ir_v2_copy_shape(CKDimToken * dst, const CKDimToken * src)

CKDimKind ck_ir_v2_dim_kind_from_name(const char * name)

const char * ck_ir_v2_dim_name(CKDimKind dim)

const char * ck_ir_v2_dtype_name(CKDataType dtype)

void ck_ir_v2_emit_dimensions(FILE * out, const CKModelConfig * cfg, const CKIRV2AlignInfo * align, int tokens_override)

void ck_ir_v2_emit_memory_plan(FILE * out, const CKIRV2Graph * graph, const CKMemPlan * plan)

void ck_ir_v2_emit_resolved_shape(FILE * out, const CKModelConfig * cfg, const CKIRV2AlignInfo * align, const CKDimToken * shape, int tokens_override)

int ck_ir_v2_emit_shape(FILE * out, const CKDimToken * shape)

const char * ck_ir_v2_find_array_end(const char * open, const char * end)

int ck_ir_v2_find_buffer_index(const CKIRV2Graph * graph, const char * name)

const CKBufferSpec * ck_ir_v2_find_buffer_spec(const char * name)

const CKKernelSpec * ck_ir_v2_find_kernel_spec(const char * name)

const char * ck_ir_v2_find_key(const char * json, const char * key, const char * end)

void ck_ir_v2_free(CKIRV2Graph * graph)

void ck_ir_v2_free_buffer(CKIRV2Buffer * buf)

void ck_ir_v2_free_node(CKIRV2Node * node)

int ck_ir_v2_lower_copy_buffers(const CKIRV2Graph * input, CKIRV2Graph * output)

int ck_ir_v2_lower_copy_nodes(const CKIRV2Graph * input, CKIRV2LowerMode mode, CKIRV2Graph * output)

int ck_ir_v2_lower_emit_json(const CKIRV2Graph * input, CKIRV2LowerMode mode, const char * path)

int ck_ir_v2_lower_graph(const CKIRV2Graph * input, CKIRV2LowerMode mode, CKIRV2Graph * output, CKMemPlan * plan)

int ck_ir_v2_lower_mode_from_string(const char * name, CKIRV2LowerMode * out_mode)

const char * ck_ir_v2_lower_mode_name(CKIRV2LowerMode mode)

int ck_ir_v2_lower_node_enabled(const CKIRV2Node * node, CKIRV2LowerMode mode)

char * ck_ir_v2_lower_strdup(const char * s)

const char * ck_ir_v2_mem_arena_name(CKMemArenaKind arena)

const char * ck_ir_v2_next_object(const char * cur, const char * end, const char ** obj_start, const char ** obj_end)

int ck_ir_v2_parse_bindings(const char * obj_start, const char * obj_end, CKIRV2Graph * graph, CKIRV2Node * node)

int ck_ir_v2_parse_bool(const char * json, const char * key, const char * end, int * out_val)

int ck_ir_v2_parse_buffers(const char * json, const char * end, CKIRV2Graph * graph)

CKDimKind ck_ir_v2_parse_dim_kind(const char * obj_start, const char * obj_end)

CKDataType ck_ir_v2_parse_dtype(const char * s)

int ck_ir_v2_parse_float(const char * json, const char * key, const char * end, float * out_val)

int ck_ir_v2_parse_int(const char * json, const char * key, const char * end, int * out_val)

int ck_ir_v2_parse_json(const char * path, CKIRV2Graph * graph)

int ck_ir_v2_parse_nodes(const char * json, const char * end, CKIRV2Graph * graph)

CKBufferRole ck_ir_v2_parse_role(const char * s)

CKBufferScope ck_ir_v2_parse_scope(const char * s)

int ck_ir_v2_parse_shape(const char * obj_start, const char * obj_end, CKDimToken * shape_out)

int ck_ir_v2_parse_string(const char * start, const char * end, char ** out_str)

int ck_ir_v2_parse_string_field(const char * json, const char * key, const char * end, char ** out_str)

void ck_ir_v2_resolve_align(const CKModelConfig * cfg, size_t alignment_bytes, CKIRV2AlignInfo * align)

size_t ck_ir_v2_resolve_dim_value(const CKModelConfig * cfg, const CKIRV2AlignInfo * align, CKDimKind dim, int tokens_override)

const char * ck_ir_v2_role_name(CKBufferRole role)

const char * ck_ir_v2_scope_name(CKBufferScope scope)

const char * ck_ir_v2_select_kernel(const CKKernelSpec * spec, CKDataType dtype, int backward)

int ck_ir_v2_serialize_json(const CKIRV2Graph * graph, const char * path)

int ck_ir_v2_serialize_json_internal(const CKIRV2Graph * graph, const CKMemPlan * plan, const char * mode, int tokens_override, int base_context_window, const char * path)

int ck_ir_v2_serialize_json_with_plan(const CKIRV2Graph * graph, const struct CKMemPlan * plan, const char * mode, int tokens_override, int base_context_window, const char * path)

const char * ck_ir_v2_skip_string(const char * cur, const char * end)

const char * ck_ir_v2_skip_ws(const char * cur, const char * end)

char * ck_ir_v2_strdup(const char * s)

int ck_ir_validate_supported(const CKIRGraph * graph)

int ck_layer_debug_enabled(void)

void ck_layout_head_to_token_f32(const float * src, float * dst, int heads, int tokens, int head_dim)

void ck_layout_token_to_head_f32(const float * src, float * dst, int tokens, int heads, int head_dim)

int ck_llama_kv_pad_256(int live_tokens, int capacity)

float ck_llama_regular_dot_f16(const float * a, const float * b, int count)

int ck_load_weights_manifest_v4(void * base, const char * weights_path, const char * manifest_path)

Load BUMPWGT4 weights into a v4 model buffer using a manifest map.

void ck_local_fp16_to_fp32_2d(const uint16_t * src, float * dst, int rows, int cols, int src_stride, int dst_stride)

void ck_local_fp16_to_fp32_row(const uint16_t * src, float * dst, int n)

void ck_local_fp32_to_bf16_row(const float * src, uint16_t * dst, int n)

void ck_local_fp32_to_fp16_row(const float * src, uint16_t * dst, int n)

int ck_mamba_debug_enabled(void)

void ck_mamba_debug_finite(const char * name, const float * x, size_t n)

int ck_mem_plan_build_inference(const CKIRV2Graph * graph, CKMemPlan * plan, size_t alignment_bytes)

int ck_mem_plan_build_inference_with_tokens(const CKIRV2Graph * graph, CKMemPlan * plan, size_t alignment_bytes, int tokens_override)

int ck_mem_plan_build_training(const CKIRV2Graph * graph, CKMemPlan * plan, size_t alignment_bytes)

int ck_mem_plan_build_training_with_tokens(const CKIRV2Graph * graph, CKMemPlan * plan, size_t alignment_bytes, int tokens_override)

void ck_mem_plan_free(CKMemPlan * plan)

void * ck_memcpy_parallel_dispatch(void * dst, const void * src, size_t size)

int ck_memory_allocate(CKModel * model, int use_hugepages)

Allocate the planned memory.

void ck_memory_free(CKModel * model)

Free the model memory.

size_t ck_memory_plan(const CKSectionConfig * sections, int num_sections, int mode, uint32_t fusion_flags, CKModel * out_model)

Plan memory layout for a model.

void ck_metrics_cleanup(void)

Cleanup and free resources

CKMetricsContext * ck_metrics_create_context(void)

void ck_metrics_ctx_end(CKMetricsContext * ctx, const char * status)

bool ck_metrics_ctx_init(CKMetricsContext * ctx, const char * run_id, const char * endpoint, CKMetricsMode mode)

void ck_metrics_ctx_log_f(CKMetricsContext * ctx, const char * name, double value)

void ck_metrics_ctx_log_i(CKMetricsContext * ctx, const char * name, int64_t value)

void ck_metrics_ctx_step(CKMetricsContext * ctx, int64_t step)

void ck_metrics_destroy_context(CKMetricsContext * ctx)

void ck_metrics_end(const char * status)

End the training run

void ck_metrics_generate_run_id(char * buffer, size_t size)

Generate a unique run ID based on current timestamp

int64_t ck_metrics_get_memory_mb(void)

Get current memory usage in MB (platform-specific)

bool ck_metrics_init(const char * run_id, const char * endpoint, CKMetricsMode mode)

Initialize metrics logging

bool ck_metrics_init_full(const char * run_id, const char * endpoint, CKMetricsMode mode, const char * model, const char * dataset, int batch_size, double lr, int max_steps)

Initialize with full configuration

void ck_metrics_log_f(const char * name, double value)

Log a float metric (e.g., loss, learning rate)

void ck_metrics_log_i(const char * name, int64_t value)

Log an integer metric (e.g., step, tokens_per_sec)

void ck_metrics_log_s(const char * name, const char * value)

Log a string metric (e.g., phase, status)

void ck_metrics_step(int64_t step)

Flush metrics for the current step and advance to next step Call this at the end of each training step

double ck_metrics_timestamp(void)

Get current timestamp in seconds with microsecond precision

int ck_min(int a, int b)

int ck_min_i(int a, int b)

int ck_model_allocate(CKModel * model, int hugepage_mode)

Allocate the planned memory. hugepage_mode 0=normal, 1=2MB hugepages, 2=1GB hugepages 0 on success, -1 on failure

int ck_model_config_from_hf_json(const char * path, CKModelConfig * cfg)

Parse a HuggingFace-style config.json into CKModelConfig.

void * ck_model_create(void)

Create and allocate model memory. Returns opaque model pointer, or NULL on failure.

void ck_model_decode(void * model, const int * token, int token_index)

Decode single token at position token_index. Used for autoregressive generation.

void ck_model_forward(void * model, const int * tokens, int num_tokens)

Forward pass (prefill) - process multiple tokens. Used for initial prompt processing.

void ck_model_free(void * model)

Free model memory.

void * ck_model_get_base(void * model)

Get model base pointer (for weight loading).

const CKModelConfig * ck_model_get_config(void)

Get model configuration (dimensions, sizes, etc.) This is available before allocation.

float * ck_model_get_logits(void * model)

Get pointer to output logits buffer. Size is vocab_size floats.

size_t ck_model_get_total_bytes(void * model)

Get total model size in bytes.

int ck_model_load_weights(void * model, const char * bump_path)

Load weights from BUMP file into model. Returns 0 on success, -1 on failure.

int ck_model_load_weights_flat(TransformerModel * m, const char * path)

Load weights from a single flat binary file into model->memory_base.

size_t ck_model_plan(CKModel * model, const CKSectionConfig * configs, int num_sections, int training_enabled, uint32_t fusion_flags)

Plan memory layout for complete model. Returns total bytes needed.

int ck_model_verify_canaries(void * model)

Verify memory canaries (debug). Returns number of corrupted canaries (0 = OK).

size_t ck_moe_align64(size_t value)

float ck_moe_bf16_round(float x)

int ck_moe_bucket_expert_for_position(const int * offsets, int n_experts, int position)

int ck_moe_bucket_layout(int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, ck_moe_q4k_q5k_bucket_layout_t * layout)

int ck_moe_debug_enabled(void)

void ck_moe_debug_finite(const char * name, const float * x, size_t n)

size_t ck_moe_down_idx(int e, int h, int i, int hidden_dim, int intermediate_dim)

float ck_moe_dsilu_f32(float x)

void ck_moe_llama_weighted_accumulate(float * output, const float * expert_output, float route_weight, int n)

void ck_moe_q4k_llama_projection(float * output, const void * weights, const void * input_q8, int output_dim, int input_dim, void * scratch)

size_t ck_moe_q4k_llama_projection_scratch_bytes(int output_dim, int input_dim)

void ck_moe_q4k_mixed_parallel_work(int ith, int nth, void * opaque)

int ck_moe_q4k_mixed_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, ck_moe_expert_workspace_fn serial_fn, ck_moe_down_kind_t down_kind)

int ck_moe_q4k_mixed_route_parallel(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 hidden_dim, int intermediate_dim, int n_experts, int top_k, void * workspace, size_t workspace_bytes, size_t workspace_stride, ck_threadpool_t * pool, ck_moe_down_kind_t down_kind)

void ck_moe_q4k_mixed_route_work(int ith, int nth, void * opaque)

void ck_moe_q4k_q5k_bucket_work(int ith, int nth, void * opaque)

void ck_moe_q4k_q5k_parallel_work(int ith, int nth, void * opaque)

void ck_moe_q4k_q5k_quantize_work(int ith, int nth, void * opaque)

int ck_moe_q4k_q5k_route_parallel(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 hidden_dim, int intermediate_dim, int n_experts, int top_k, void * workspace, size_t workspace_bytes, size_t workspace_stride, ck_threadpool_t * pool)

void ck_moe_q4k_q5k_route_work(int ith, int nth, void * opaque)

int ck_moe_shared_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, size_t stride, ck_moe_shared_gated_workspace_fn serial_fn)

int ck_moe_shared_q4k_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, void(*)(float *, const void *, const void *, int, int) down_projection)

void ck_moe_shared_q4k_parallel_work(int ith, int nth, void * opaque)

int ck_moe_shared_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, ck_moe_shared_workspace_fn serial_fn)

void ck_moe_shared_q8_0_parallel_work(int ith, int nth, void * opaque)

float ck_moe_silu_f32(float x)

int ck_moe_size_add(size_t a, size_t b, size_t * result)

int ck_moe_size_mul(size_t a, size_t b, size_t * result)

size_t ck_moe_up_idx(int e, int i, int h, int intermediate_dim, int hidden_dim)

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)

uint32_t ck_murmurhash3(const char * key, uint32_t len, uint32_t seed)

MurmurHash3-32bit hash function (original HPC_Embeddings version).

void ck_murmurhash3_128(const void * key, size_t len, uint32_t seed, uint64_t * out1, uint64_t * out2)

MurmurHash3-128bit hash function (produces two 64-bit values).

uint32_t ck_murmurhash3_32(const void * key, size_t len, uint32_t seed)

MurmurHash3-32bit hash function (alternative name).

uint32_t ck_murmurhash3_str(const char * key, uint32_t seed)

uint32_t ck_murmurhash3_strn(const char * str, size_t len, uint32_t seed)

MurmurHash3-32bit for string with length.

int ck_nearest_int(float fval)

int ck_nearest_int_fused(float fval)

int ck_nearest_int_q8_0(float fval)

int ck_nearest_int_q8_0_ref(float fval)

size_t ck_nvfp4_align64(size_t value)

void ck_nvfp4_gemv_rows(int begin, int end, void * opaque)

const char * ck_op_name(CKOpType op)

int ck_op_supported(CKOpType op)

int ck_opt_pick_active_threads(int nth, size_t work_items, size_t min_chunk)

void ck_parallel_for_worker(int ith, int nth, void * opaque)

int ck_parse_env_int(const char * name)

void ck_patch_projection_bf16_native_work(int ith, int nth, void * opaque)

int ck_plan_step_enabled(const CKPlanStep * step, const CKIRGraph * cfg)

void * ck_pool_alloc(CKMemPool * pool, size_t size)

void ck_pool_free(CKMemPool * pool)

void ck_pool_init(CKMemPool * pool)

char * ck_pool_strdup(CKMemPool * pool, const char * s, int len)

int ck_prefill_forward(const void * weights, const int32_t * tokens, int n_tokens, float * hidden_out, void * kv_cache, int kv_pos)

Forward pass computation

int ck_q4k_debug_q8_contract(void)

int ck_q4k_packed_vnni_x16_available(void)

int ck_q4k_packed_vnni_x8_available(void)

int ck_q4k_packed_vnni_x8_compact_order_available(void)

int ck_q4k_q8k_force_ref(void)

float ck_q4k_silu_f32(float x)

int ck_q4k_x16_chunk4_enabled(void)

void ck_q5_k_prepare_weight(const void * src, void * dst, int N, int K)

size_t ck_q5_k_prepared_block_size(void)

int ck_q5k_debug_fp32_fallback(void)

int ck_q5k_debug_generic_dot(void)

void ck_q6_k_prepare_weight(const void * src, void * dst, int N, int K)

size_t ck_q6_k_prepared_block_size(void)

const char * ck_q6_k_prepared_provider_name(void)

const char * ck_q6_k_q8_k_provider_name(void)

int ck_q6k_q8k_force_ref(void)

int ck_q80_contract_cached_input_enabled(void)

int ck_q80_contract_dump_enabled(void)

void ck_q80_contract_dump_tensor(const char * name, int layer_id, const float * data, size_t elem_count)

ck_q80_ggml_build_forward_expand_fn ck_q80_resolve_ggml_build_forward_expand(void)

Forward pass computation

ck_q80_ggml_cpu_init_fn ck_q80_resolve_ggml_cpu_init(void)

ck_q80_ggml_free_fn ck_q80_resolve_ggml_free(void)

ck_q80_ggml_get_data_fn ck_q80_resolve_ggml_get_data(void)

ck_q80_ggml_get_data_f32_fn ck_q80_resolve_ggml_get_data_f32(void)

ck_q80_ggml_graph_compute_with_ctx_fn ck_q80_resolve_ggml_graph_compute_with_ctx(void)

ck_q80_ggml_init_fn ck_q80_resolve_ggml_init(void)

ck_q80_ggml_mul_mat_fn ck_q80_resolve_ggml_mul_mat(void)

ck_q80_ggml_nbytes_fn ck_q80_resolve_ggml_nbytes(void)

ck_q80_ggml_new_graph_fn ck_q80_resolve_ggml_new_graph(void)

ck_q80_ggml_new_tensor_2d_fn ck_q80_resolve_ggml_new_tensor_2d(void)

int ck_q8_0_debug_ref(void)

int ck_q8_0_fp32_m4n4_enabled(void)

int ck_q8_0_outproj_enabled(void)

int ck_q8_0_q8_0_debug_ref(void)

int ck_q8k_activations_enabled(void)

void ck_qkv_project_head_major(const float * input, 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 tokens, int kv_stride_tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_backward(const float * d_q, const float * d_k, const float * d_v, const float * input, const float * wq, const float * bq, const float * wk, const float * bk, const float * wv, const float * bv, float * d_input, float * d_wq, float * d_bq, float * d_wk, float * d_bk, float * d_wv, float * d_bv, float * scratch, int tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim, int num_threads)

Backward pass / gradient computation

void ck_qkv_project_head_major_q4_k(const float * input, const void * wq, const float * bq, const void * wk, const float * bk, const void * wv, const float * bv, float * q, float * k, float * v, int tokens, int kv_stride_tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_q4_k_q8_k(const float * input, const void * wq, const float * bq, const void * wk, const float * bk, const void * wv, const float * bv, float * q, float * k, float * v, int tokens, int kv_stride_tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_quant(const float * input, const void * wq, const float * bq, CKDataType wq_dtype, const void * wk, const float * bk, CKDataType wk_dtype, const void * wv, const float * bv, CKDataType wv_dtype, float * q, float * k, float * v, int tokens, int kv_stride_tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_ref(const float * input, 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 tokens, int kv_stride_tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_token(const float * input_row, const float * wq, const float * bq, const float * wk, const float * bk, const float * wv, const float * bv, float * q_token, float * k_token, float * v_token, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_token_q4_k(const float * input_row, const void * wq, const float * bq, const void * wk, const float * bk, const void * wv, const float * bv, float * q_token, float * k_token, float * v_token, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_token_q4_k_q8_k(const block_q8_K * input_q8, const void * wq, const float * bq, const void * wk, const float * bk, const void * wv, const float * bv, float * q_token, float * k_token, float * v_token, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

void ck_qkv_project_head_major_token_quant(const float * input_row, const void * wq, const float * bq, CKDataType wq_dtype, const void * wk, const float * bk, CKDataType wk_dtype, const void * wv, const float * bv, CKDataType wv_dtype, float * q_token, float * k_token, float * v_token, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)

size_t ck_quant_block_size(int type)

Get the block size (number of weights per block) for a quant type.

size_t ck_quant_row_size(int type, int64_t n_elements)

Calculate total bytes needed for n_elements with given quant type.

size_t ck_quant_type_size(int type)

Get the byte size per block for a quant type.

void ck_residual_add_backward(const float * d_out, float * d_a, float * d_b, int tokens, int aligned_embed_dim)

Backward pass / gradient computation

void ck_residual_add_token_major(const float * a, const float * b, float * out, int tokens, int aligned_embed_dim)

void ck_residual_add_token_major_bf16_storage(const float * a, const float * b, float * out, int tokens, int aligned_embed_dim)

void ck_residual_add_token_major_parallel_dispatch(const float * a, const float * b, float * out, int tokens, int aligned_embed_dim)

ck_ggml_backend_cpu_set_n_threads_fn ck_resolve_ggml_backend_cpu_set_n_threads(void)

ck_ggml_backend_free_fn ck_resolve_ggml_backend_free(void)

ck_ggml_backend_get_default_buffer_type_fn ck_resolve_ggml_backend_get_default_buffer_type(void)

ck_ggml_backend_init_by_type_fn ck_resolve_ggml_backend_init_by_type(void)

ck_ggml_backend_sched_alloc_graph_fn ck_resolve_ggml_backend_sched_alloc_graph(void)

ck_ggml_backend_sched_free_fn ck_resolve_ggml_backend_sched_free(void)

ck_ggml_backend_sched_graph_compute_fn ck_resolve_ggml_backend_sched_graph_compute(void)

ck_ggml_backend_sched_new_fn ck_resolve_ggml_backend_sched_new(void)

ck_ggml_backend_sched_reset_fn ck_resolve_ggml_backend_sched_reset(void)

ck_ggml_backend_tensor_get_fn ck_resolve_ggml_backend_tensor_get(void)

ck_ggml_backend_tensor_set_fn ck_resolve_ggml_backend_tensor_set(void)

ck_ggml_build_forward_expand_fn ck_resolve_ggml_build_forward_expand(void)

Forward pass computation

ck_ggml_cont_fn ck_resolve_ggml_cont(void)

ck_ggml_cont_2d_fn ck_resolve_ggml_cont_2d(void)

ck_ggml_cpu_init_fn ck_resolve_ggml_cpu_init(void)

ck_ggml_free_fn ck_resolve_ggml_free(void)

ck_ggml_get_data_fn ck_resolve_ggml_get_data(void)

ck_ggml_graph_compute_with_ctx_fn ck_resolve_ggml_graph_compute_with_ctx(void)

ck_ggml_init_fn ck_resolve_ggml_init(void)

ck_ggml_mul_mat_graph_fn ck_resolve_ggml_mul_mat_graph(void)

ck_ggml_new_graph_fn ck_resolve_ggml_new_graph(void)

ck_ggml_new_tensor_1d_fn ck_resolve_ggml_new_tensor_1d(void)

ck_ggml_new_tensor_2d_fn ck_resolve_ggml_new_tensor_2d(void)

ck_ggml_permute_fn ck_resolve_ggml_permute(void)

ck_ggml_set_input_fn ck_resolve_ggml_set_input(void)

ck_ggml_soft_max_ext_fn ck_resolve_ggml_soft_max_ext(void)

ck_ggml_view_3d_fn ck_resolve_ggml_view_3d(void)

void ck_round_fp16_buffer(const float * src, float * dst, size_t count)

float ck_round_fp16_scalar(float x)

int ck_round_nearest(float v)

Round to nearest int, half away from zero (matches quantize_row_q8_0)

int ck_sample_top_p_v8(float * logits, int vocab_size, float temperature, float top_p, float random_value)

void ck_scale_f32_to_f16(const float * src, float scale, uint16_t * dst, int n)

Scale FP32 array and store as FP16: dst = scale * src.

void ck_scale_parallel_work(int ith, int nth, void * argp)

void ck_section_config_init(CKSectionConfig * config, size_t simd_align)

Initialize section config with computed alignments.

size_t ck_section_plan(CKSection * section, const CKSectionConfig * config, int training_enabled, size_t base_offset)

Plan memory layout for a single section. Returns bytes needed for this section.

void ck_session_v8_cancel(CKSessionV8 * session)

void ck_session_v8_close(CKSessionV8 * session)

int ck_session_v8_decode(CKSessionV8 * session, const int32_t * tokens, int32_t token_count, char * output, int32_t capacity)

int ck_session_v8_encode(CKSessionV8 * session, const char * text, int32_t * output, int32_t capacity)

int ck_session_v8_format_chat(CKSessionV8 * session, const char * system_text, const char * user_text, char * output, int32_t capacity)

int ck_session_v8_generate(CKSessionV8 * session, const CKSessionGenerateRequestV8 * request, ck_session_token_callback_v8 callback, void * user_data, CKSessionGenerateResultV8 * result)

uint32_t ck_session_v8_get_abi_version(void)

int ck_session_v8_get_model_descriptor(const CKSessionV8 * session, CKModelRuntimeDescriptorV8 * descriptor, size_t descriptor_size)

const char * ck_session_v8_last_error(const CKSessionV8 * session)

int ck_session_v8_open(const CKSessionConfigV8 * config, CKSessionV8 ** session_out)

int ck_session_v8_reset(CKSessionV8 * session)

void ck_set_num_threads(int num_threads)

void ck_set_strict_parity(int enabled)

int ck_speed_profile_qwen3vl_ocr_fast(void)

void ck_ssm_conv1d_llama_channel_range(int begin, int end, void * opaque)

void ck_ssm_conv1d_llama_fma_channel_range(int begin, int end, void * opaque)

int ck_strict_mtmd_clip_encode_planar_f32(const float * planar, int channels, int height, int width, float * out, size_t out_elems)

int ck_strict_parity_enabled(void)

void ck_sum_sq_multi_parallel_work(int ith, int nth, void * argp)

void ck_sum_sq_parallel_work(int ith, int nth, void * argp)

void ck_test_dequant_q4_0(const void * src, float * dst, int n)

Dequantize Q4_0 data to FP32.

void ck_test_dequant_q4_k(const void * src, float * dst, int n)

Dequantize Q4_K data to FP32.

void ck_test_dequant_q6_k(const void * src, float * dst, int n)

Dequantize Q6_K data to FP32.

void ck_test_gated_deltanet_autoregressive(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)

Gated DeltaNet autoregressive update.

void ck_test_gemv_q4_k(const void * weight_q4k, const float * input_f32, float * output, int cols)

Q4_K GEMV - dot product of quantized weights and FP32 input.

void ck_test_gemv_q5_0(const void * weight_q5_0, const float * input_f32, float * output, int rows, int cols)

Q5_0 GEMV - matrix-vector multiply with Q5_0 weights.

void ck_test_gemv_q5_0_q8_0(const void * weight_q5_0, const float * input_f32, float * output, int rows, int cols)

Q5_0 x Q8_0 quantized GEMV - matches llama.cpp's approach.

void ck_test_gemv_q6_k(const void * weight_q6k, const float * input_f32, float * output, int cols)

Q6_K GEMV.

void ck_test_gemv_q8_0(const void * weight_q8_0, const float * input_f32, float * output, int rows, int cols)

Q8_0 GEMV - matrix-vector multiply with Q8_0 weights.

void ck_test_gemv_q8_0_q8_0(const void * weight_q8_0, const float * input_f32, float * output, int rows, int cols)

Q8_0 x Q8_0 quantized GEMV - matches llama.cpp's approach.

void ck_test_quantize_q8_k(const float * src, void * dst, int n)

Quantize FP32 to Q8_K (for activations)

void ck_test_recurrent_conv_state_update(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)

Build the recurrent convolution input history window.

void ck_test_recurrent_dt_gate(const float * alpha, const float * dt_bias, const float * a, float * gate, int rows, int dim)

Transform recurrent alpha rows into the DeltaNet gate.

void ck_test_recurrent_norm_gate(const float * x, const float * gate, const float * weight, float * out, int rows, int num_heads, int head_dim, float eps)

Per-head RMSNorm followed by SiLU(z) gating for recurrent outputs.

void ck_test_recurrent_qk_l2_norm(float * q, float * k, int rows, int q_dim, int k_dim, int head_dim, float eps)

Apply per-head L2 normalization to recurrent Q/K rows in-place.

void ck_test_recurrent_silu(const float * x, float * out, int rows, int dim)

Apply SiLU elementwise to recurrent rows.

void ck_test_recurrent_split_conv_qkv(const float * packed_qkv, float * q, float * k, float * v, int rows, int q_dim, int k_dim, int v_dim)

Split the post-convolution recurrent packed QKV rows.

void ck_test_recurrent_split_qkv(const float * packed_qkv, float * q, float * k, float * v, int rows, int q_dim, int k_dim, int v_dim)

Split a packed recurrent QKV matrix into explicit Q, K, and V outputs.

void ck_test_split_q_gate(const float * packed_qg, float * q, float * gate, int rows, int q_dim, int gate_dim, int group_dim)

Split a packed full-attention Q+gate matrix into Q rows and gate rows.

void ck_test_ssm_conv1d(const float * conv_x, const float * kernel, float * out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

qwen3next/Qwen3.5 SSM causal depthwise convolution.

void ck_test_vec_dot_q4_k_q8_k(const void * weight_q4_k, const void * input_q8_k, float * output, int cols)

Direct Q4_K x Q8_K dot product using identical pre-quantized bytes.

void ck_test_vec_dot_q5_0_q8_0(const void * weight_q5_0, const void * input_q8_0, float * output, int cols)

Direct Q5_0 x Q8_0 dot product test (takes pre-quantized Q8_0 input)

void ck_test_vec_dot_q6_k_q8_k(const void * weight_q6_k, const void * input_q8_k, float * output, int cols)

Direct Q6_K x Q8_K dot product using identical pre-quantized bytes.

void ck_test_vec_dot_q8_0_q8_0(const void * weight_q8_0, const void * input_q8_0, float * output, int cols)

Direct Q8_0 x Q8_0 dot product test (takes pre-quantized Q8_0 input)

void ck_threadpool_barrier(ck_threadpool_t * pool)

Barrier synchronization within a dispatched work function.

int ck_threadpool_bounded_capacity(int default_threads, int logical_threads)

Compute the bounded capacity for an SMT-safe provider. The default width is preserved and at most half of the additional logical CPUs are reserved.

int ck_threadpool_capacity(const ck_threadpool_t * pool)

Get the maximum worker capacity available to explicit dispatch_n calls.

ck_threadpool_t * ck_threadpool_create(int n_threads)

Create a thread pool with n_threads total threads. Thread 0 is the calling (main) thread; n_threads-1 workers are spawned.

ck_threadpool_t * ck_threadpool_create_capacity(int default_threads, int capacity_threads)

Create a pool whose ordinary dispatch width is smaller than its worker capacity. Exact providers may opt into the additional workers with ck_threadpool_dispatch_n(); ordinary dispatch remains at default_threads.

void ck_threadpool_destroy(ck_threadpool_t * pool)

Destroy the thread pool. Signals all workers to exit and joins them. Safe to call with NULL.

void ck_threadpool_dispatch(ck_threadpool_t * pool, ck_work_fn_t fn, void * args)

Dispatch work to all threads and wait for completion.

void ck_threadpool_dispatch_n(ck_threadpool_t * pool, int active_threads, ck_work_fn_t fn, void * args)

Dispatch work to a subset of the pool and wait for completion.

ck_threadpool_t * ck_threadpool_global(void)

Get or create the global thread pool. Thread-safe (uses pthread_once internally). Uses ck_get_num_threads() for the default width. In automatic mode the pool may reserve a bounded subset of SMT siblings as explicit provider capacity; ordinary dispatch remains at the default width.

void ck_threadpool_global_destroy(void)

Destroy the global thread pool. Called during engine shutdown.

void ck_threadpool_init(void)

Initialize the global thread pool. Called once during engine startup (e.g., from ck_model_init). Uses ck_get_num_threads() for thread count (respects CK_NUM_THREADS env).

int ck_threadpool_n_threads(const ck_threadpool_t * pool)

Get the ordinary/default dispatch width (including the main thread).

void ck_threadpool_parallel_for_n(ck_threadpool_t * pool, int active_threads, int begin, int end, int grain_size, ck_range_fn_t fn, void * args)

Dynamically distribute independent ranges through the persistent pool.

void ck_threadpool_pause(ck_threadpool_t * pool)

Pause workers — they sleep on condvar (0% CPU). Call between batches or during interactive waiting. Workers wake on next dispatch or resume.

void ck_threadpool_profile_reset(ck_threadpool_t * pool)

Enable profiling and reset cumulative dispatch timing counters.

void ck_threadpool_profile_snapshot(const ck_threadpool_t * pool, ck_threadpool_profile_t * profile)

Snapshot cumulative dispatch timing counters without stopping workers.

void ck_threadpool_resume(ck_threadpool_t * pool)

Resume workers — transition from sleep to spin-wait. Call before starting a new batch of work.

void ck_threadpool_shutdown(void)

Shut down the global thread pool. Called during engine teardown. Workers are joined and freed.

int ck_threadpool_thread_id(const ck_threadpool_t * pool)

Get thread index for current thread (0 = main, -1 if not in pool)

int ck_tokenizer_add_merge(CKTokenizer * tok, int32_t left, int32_t right, int32_t merged)

int ck_tokenizer_add_special_token(CKTokenizer * tok, const char * name, int32_t id)

Add special token (UNK, BOS, EOS, PAD, MASK).

int32_t ck_tokenizer_add_token(CKTokenizer * tok, const char * token, int len)

CKTokenizer * ck_tokenizer_create(CKTokenizerType type)

CKTokenizer * ck_tokenizer_create_bpe(void)

Create tokenizer with default BPE config.

CKTokenizer * ck_tokenizer_create_spm(void)

Create tokenizer with default SPM config.

CKTokenizer * ck_tokenizer_create_wordpiece(void)

Create tokenizer with default WordPiece config.

int ck_tokenizer_decode(const CKTokenizer * tok, const int32_t * ids, int num_ids, char * text, int max_len)

Decode token IDs to text.

CKSpacePrefixStyle ck_tokenizer_detect_space_prefix_style(CKTokenizer * tok)

int ck_tokenizer_encode(const CKTokenizer * tok, const char * text, int text_len, int32_t * ids, int max_ids)

Encode text to token IDs using greedy longest-match.

int ck_tokenizer_encode_spm_dispatch(const CKTokenizer * tok, const char * text, int text_len, int32_t * ids, int max_ids)

int ck_tokenizer_encode_spm_impl(const CKTokenizer * tok, const char * text, int text_len, int32_t * ids, int max_ids)

int ck_tokenizer_encode_spm_llama_impl(const CKTokenizer * tok, const char * text, int text_len, int32_t * ids, int max_ids)

int ck_tokenizer_encode_spm_plain_segment(const CKTokenizer * tok, const char * text, int text_len, int32_t * ids, int max_ids)

int ck_tokenizer_encode_tokens(const CKTokenizer * tok, const char * text, int text_len, const char ** out_tokens, int max_tokens)

Encode and return tokens as array of strings.

int ck_tokenizer_encode_with_special(CKTokenizer * tok, const char * text, int text_len, int32_t * ids, int max_ids, bool add_special)

Encode with special token handling.

void ck_tokenizer_free(CKTokenizer * tok)

uint32_t ck_tokenizer_hash(const char * key, size_t len)

uint32_t ck_tokenizer_hash_str(const char * key)

void ck_tokenizer_hash_table_clear(CKTokenizerHashTable * table, bool free_values)

Clear all entries (but keep bucket array).

bool ck_tokenizer_hash_table_contains(CKTokenizerHashTable * table, const char * key)

Check if key exists.

size_t ck_tokenizer_hash_table_count(CKTokenizerHashTable * table)

Get the number of entries.

CKTokenizerHashTable * ck_tokenizer_hash_table_create(size_t bucket_count)

Create a hash table.

int ck_tokenizer_hash_table_delete(CKTokenizerHashTable * table, const char * key, bool free_value)

Delete a key.

void ck_tokenizer_hash_table_free(CKTokenizerHashTable * table, bool free_values)

Free a hash table.

int ck_tokenizer_hash_table_insert(CKTokenizerHashTable * table, const char * key, void * value)

Insert a key-value pair.

int ck_tokenizer_hash_table_iterate(CKTokenizerHashTable * table, CKTokenizerHashCallback callback, void * user_data)

size_t ck_tokenizer_hash_table_keys(CKTokenizerHashTable * table, const char ** out_keys, size_t max_keys)

Get all keys as an array.

void * ck_tokenizer_hash_table_lookup(CKTokenizerHashTable * table, const char * key)

Look up a key.

void * ck_tokenizer_hash_table_lookup_avx(CKTokenizerHashTable * table, const char * key)

const char * ck_tokenizer_id_to_token(const CKTokenizer * tok, int32_t id)

int ck_tokenizer_init(CKTokenizer * tok)

int ck_tokenizer_load(CKTokenizer * tok, const char * path)

int ck_tokenizer_load_binary(CKTokenizer * tok, int vocab_size, const int32_t * offsets, const char * strings, int num_merges, const int32_t * merges)

Load vocabulary from memory-mapped binary data.

int ck_tokenizer_load_binary_with_scores(CKTokenizer * tok, int vocab_size, const int32_t * offsets, const char * strings, const float * scores, const uint8_t * types, int num_merges, const int32_t * merges)

Load vocabulary from memory-mapped binary data with scores and types.

int ck_tokenizer_load_gguf(CKTokenizer * tok, const char * path)

Load vocabulary from GGUF file.

int ck_tokenizer_load_json(CKTokenizer * tok, const char * path)

Load vocabulary from JSON file (HuggingFace format).

int ck_tokenizer_load_merges(CKTokenizer * tok, const char * path)

Load BPE merges from text file.

int ck_tokenizer_load_text(CKTokenizer * tok, const char * path)

Load vocabulary from text file (one token per line).

int32_t ck_tokenizer_lookup(const CKTokenizer * tok, const char * token, int len)

int32_t ck_tokenizer_lookup_exact(const CKTokenizer * tok, const char * token)

int32_t ck_tokenizer_lookup_exact_n(const CKTokenizer * tok, const char * text, int text_len)

int ck_tokenizer_lookup_merge(const CKTokenizer * tok, int32_t left, int32_t right)

void * ck_tokenizer_mempool_alloc(CKTokenizerMemPool * pool, size_t size)

Allocate from pool.

void * ck_tokenizer_mempool_alloc_aligned(CKTokenizerMemPool * pool, size_t size, size_t align)

Allocate aligned memory from pool.

size_t ck_tokenizer_mempool_alloc_count(CKTokenizerMemPool * pool)

Get allocation count.

size_t ck_tokenizer_mempool_available(CKTokenizerMemPool * pool)

Get available bytes in pool.

void ck_tokenizer_mempool_free(CKTokenizerMemPool * pool)

Free a memory pool.

int ck_tokenizer_mempool_init(CKTokenizerMemPool * pool, size_t size)

Initialize a memory pool.

void ck_tokenizer_mempool_reset(CKTokenizerMemPool * pool)

Reset pool (mark all memory as free).

char * ck_tokenizer_mempool_strdup(CKTokenizerMemPool * pool, const char * str)

Allocate and copy string (strdup equivalent).

char * ck_tokenizer_mempool_strndup(CKTokenizerMemPool * pool, const char * str, int len)

Allocate and copy string with length.

size_t ck_tokenizer_mempool_used(CKTokenizerMemPool * pool)

Get used bytes in pool.

void ck_tokenizer_reset(CKTokenizer * tok)

void ck_tokenizer_set_add_bos_eos(CKTokenizer * tok, bool add_bos, bool add_eos)

void ck_tokenizer_set_add_space_prefix(CKTokenizer * tok, bool add_space_prefix)

void ck_tokenizer_set_space_prefix_style(CKTokenizer * tok, CKSpacePrefixStyle style)

void ck_tokenizer_set_special_ids(CKTokenizer * tok, int32_t unk, int32_t bos, int32_t eos, int32_t pad, int32_t mask)

void ck_tokenizer_set_spm_mode(CKTokenizer * tok, CKSpmMode spm_mode)

void ck_tokenizer_set_use_trie(CKTokenizer * tok, bool use_trie)

size_t ck_tokenizer_utf8_normalize_nfc(const char * src, size_t src_len, char * dst, size_t dst_size)

size_t ck_tokenizer_vocab_size(const CKTokenizer * tok)

Get vocabulary size.

void ck_topk_insert_desc(int idx, float val, int * indices, float * values, int k)

void ck_train_bias_reduce_compute_range(const float * d_output, float * d_bias, int T, int out_start, int out_end, int aligned_out)