API Reference

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

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_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_q5_1(const void * weight_q5_1, const float * input_f32, float * output, int rows, int cols, int n_tokens)

Test Q5_1 x Q8_0 GEMM (batch matrix multiply)

void ck_test_gemm_q5_k(const void * weight_q5_k, const float * input_f32, float * output, int rows, int cols, int n_tokens)

Test Q5_K x Q8_K 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 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)

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_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_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_nn_avx512(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_nt(const float * input, const float * weight, float * output, int rows, int cols, int common)

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_q8_k(const void * A_q8, const void * B, const float * bias, float * C, int M, int N, int K)

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_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_k(const float * A, 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_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_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)

Matrix-matrix multiply: C[M,N] = A[M,K] @ B[N,K]^T + bias.

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, const void * B, const float * bias, float * C, int M, int N, int K)

gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

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_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_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_ref(float * Y, const void * W, const void * X_q8, int M, int N, int K)

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_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_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_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)

const GEMMParams * get_gemm_params(void)

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_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_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 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 (v3 optimized)

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 size for fused prefill.

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 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_forward(const float * input, const float * gamma, float * output, float * rstd_cache, int tokens, int d_model, int aligned_embed_dim, 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_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_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_q8_k_fused(const float * input, const float * gamma, void * vy, int tokens, int d_model, int aligned_embed_dim, float eps)

Fused RMSNorm + Q8_K Quantization

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 comparison.

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

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)

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

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

Softmax (simple, non-causal)

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_inplace(float * x, int n)

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_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)

Flash attention forward for GQA (prefill, no score materialization) Testtest_flash_attention.py::TestFlashAttention::test_flash_forward test_flash_attention.py::TestFlashAttention::test_flash_vs_score_matrix test_flash_attention.py::TestFlashAttention::test_flash_gqa test_attention.py::TestAttentionForward::test_flash_forward

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_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_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_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)

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

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_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)

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_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)

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 ck_test_attention_decode_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 sliding_window)

Test sliding-window attention (decode mode)

void ck_test_attention_sliding_window(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, int sliding_window)

Test sliding-window attention (prefill)

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_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 for prefill with proper tiling.

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 size for fused MLP.

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, CKDataType wo_dt, const void * w1, const float * b1, CKDataType w1_dt, const void * w2, const float * b2, CKDataType 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)

Mega-fused post-attention block (out-proj + RMSNorm2 + MLP) for prefill.

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_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 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)

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

SwiGLU activation.

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 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_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 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 qwen2_0_5b_decode_precompute_rope(QWEN2_0_5B_DECODEModel * model)

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)

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_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_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_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_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_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_precompute_cache(float * cos_cache, float * sin_cache, int max_seq_len, int head_dim, float base)

Precompute RoPE cos/sin cache Testtest_rope.py::TestRoPECache::test_cache_computation test_rope.py::TestRoPECache::test_cache_values

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")

struct __attribute__((packed))

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)

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)

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 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)

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

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_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)

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

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)

float ck_expf(float x)

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.

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.

int ck_get_block_q4_k_size(void)

Get Q4_K block size in bytes.

int ck_get_block_q5_1_size(void)

Get Q5_1 block size in bytes (24 bytes per 32 weights)

int ck_get_block_q5_k_size(void)

Get Q5_K block size in bytes (176 bytes per 256 weights)

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

Get QK5_1 (elements per Q5_1 block)

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.

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 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)

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.

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)

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).

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)

const char * ck_op_name(CKOpType op)

int ck_op_supported(CKOpType op)

int ck_parse_env_int(const char * name)

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_q8_0_outproj_enabled(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)

int ck_round_nearest(float v)

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

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_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_set_num_threads(int num_threads)

void ck_set_strict_parity(int enabled)

int ck_strict_parity_enabled(void)

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_q5_1(const void * src, float * dst, int n)

Dequantize Q5_1 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_geglu(const float * x, float * out, int n_tokens, int dim)

Test GeGLU activation.

void ck_test_geglu_backward(const float * x, const float * d_out, float * d_x, int n_tokens, int dim)

Test GeGLU backward.

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_q5_1(const void * weight_q5_1, const float * input_f32, float * output, int rows, int cols)

Q5_1 GEMV - matrix-vector multiply with Q5_1 weights (32-element blocks)

void ck_test_gemv_q5_k(const void * weight_q5_k, const float * input_f32, float * output, int rows, int cols)

Q5_K GEMV - matrix-vector multiply with Q5_K weights (256-element super-blocks)

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_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_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.

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.

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.

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 auto-detection.

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 total thread count (including main thread)

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_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)

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)

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_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)

int ck_tokenizer_vocab_size(const CKTokenizer * tok)

void ck_trie_clear(CKTrie * trie)

CKTrie * ck_trie_create(size_t max_nodes)

int32_t ck_trie_find_longest(const CKTrie * trie, const char * text, size_t text_len, size_t start_pos, size_t * match_len)

void ck_trie_free(CKTrie * trie)

bool ck_trie_has_prefix(const CKTrie * trie, const char * text, size_t text_len, size_t pos)

int ck_trie_insert(CKTrie * trie, const char * token, int32_t token_id, bool is_special, int32_t priority)

size_t ck_trie_node_count(const CKTrie * trie)

int ck_true_bpe_add_merge(CKTrueBPE * bpe, int32_t left_id, int32_t right_id, int32_t merged_id, int32_t priority)

int ck_true_bpe_add_merge_by_tokens(CKTrueBPE * bpe, const char * left, const char * right, int32_t priority)

int ck_true_bpe_add_special_token(CKTrueBPE * bpe, const char * token, int32_t id)

int ck_true_bpe_add_token(CKTrueBPE * bpe, const char * token, int32_t id, float score)

CKTrueBPE * ck_true_bpe_create(void)

int ck_true_bpe_decode(const CKTrueBPE * bpe, const int32_t * ids, int num_ids, char * text, int max_len)

CKSpacePrefixStyle ck_true_bpe_detect_space_style(CKTrueBPE * bpe)

int ck_true_bpe_encode(CKTrueBPE * bpe, const char * text, int text_len, int32_t * ids, int max_ids)

void ck_true_bpe_free(CKTrueBPE * bpe)

const char * ck_true_bpe_id_to_token(const CKTrueBPE * bpe, int32_t id)

int ck_true_bpe_load_binary(CKTrueBPE * bpe, int vocab_size, const int32_t * offsets, const char * strings, int num_merges, const int32_t * merges)

int32_t ck_true_bpe_lookup(const CKTrueBPE * bpe, const char * token)

int32_t ck_true_bpe_num_merges(const CKTrueBPE * bpe)

void ck_true_bpe_set_config(CKTrueBPE * bpe, const CKBPEConfig * config)

void ck_true_bpe_set_special_ids(CKTrueBPE * bpe, int32_t unk, int32_t bos, int32_t eos, int32_t pad)

size_t ck_true_bpe_vocab_size(const CKTrueBPE * bpe)

size_t ck_utf8_byte_to_offset(const char * str, size_t len, size_t byte_offset)

Get the character index from byte offset.

int ck_utf8_char_length(unsigned char c)

Get the length of a UTF-8 character from its first byte.

size_t ck_utf8_count_chars(const char * str, size_t len)

Count UTF-8 characters in a string.

uint32_t ck_utf8_decode_2(const char * s)

Get 2-byte UTF-8 sequence value.

uint32_t ck_utf8_decode_3(const char * s)

Get 3-byte UTF-8 sequence value.

uint32_t ck_utf8_decode_4(const char * s)

Get 4-byte UTF-8 sequence value.

unsigned char ck_utf8_first_byte(const char * s)

Get the first byte of a UTF-8 character.

int ck_utf8_from_cp(uint32_t cp, char * out)

Write a Unicode code point as UTF-8.

int ck_utf8_is_continuation(unsigned char c)

Get the continuation byte mask and value.

bool ck_utf8_is_valid(const char * str, size_t len)

Check if a byte sequence is valid UTF-8.

bool ck_utf8_is_whitespace(uint32_t cp)

Check if character is whitespace (Unicode White_Space property).

int32_t ck_utf8_next_char(const char ** str, int * out_len)

Get next UTF-8 character, return its code point.

size_t ck_utf8_normalize_nfc(const char * src, size_t src_len, char * dst, size_t dst_size)

Normalize UTF-8 string (Unicode normalization form NFC).

size_t ck_utf8_offset_to_byte(const char * str, size_t len, size_t n)

Get the byte offset of the N-th character.

size_t ck_utf8_validate(const char * str, size_t len)

Validate a UTF-8 string.

const char * ck_weight_dtype_expr(const CKBufferSpec * spec)

CKMathBackend ckernel_backend_native(void)

Obtain the built-in native backend (single-node CPU, C + intrinsics).

int8_t clamp_int8(float value)

CKV2AlignInfo compute_align(const CKModelConfig * cfg)

float compute_rms_scale(const float * x, int n, float eps)

float compute_rms_scale_internal(const float * x, int n, float eps)

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

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

Convert FP16 tensor to FP32.

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

Convert FP32 tensor to FP16.

void convert_float_to_int4(const float * src, uint8_t * dst, size_t count)

void convert_float_to_int8(const float * src, int8_t * dst, size_t count)

void convert_int4_to_float(const uint8_t * src, float * dst, size_t count)

void convert_int8_to_float(const int8_t * src, float * dst, size_t count)

int count_set_bits(const char * hex_mask)

void cpu_features_init(void)

CKTokenizerHashEntry * create_entry(const char * key, const void * value, size_t value_size)

CKTrieNode * create_node(void)

int decode_bpe_token(const char * token, char * out, int max)

Decode GPT-2 byte-level BPE representation back to actual bytes.

int8_t decode_int4(uint8_t packed, int index)

void decode_layer_parallel(float * hidden, const void * ln1_weight, const void * ln2_weight, const void * WQ, const void * WK, const void * WV, const void * WO, const void * W_gate, const void * W_up, const void * W_down, float * k_cache, float * v_cache, int token_index, float * scratch, int embed_dim, int intermediate, int H, int H_kv, int head_dim, int max_seq, float eps, int num_threads)

Process one transformer layer in parallel.

void dequant_q4_0_block(const block_q4_0 * block, float * output)

Dequantize a single Q4_0 block to FP32.

void dequant_q4_0_row(const void * src, float * dst, size_t n_elements)

Dequantize Q4_0 row (multiple blocks)

void dequant_q4_1_block(const block_q4_1 * block, float * output)

Dequantize a single Q4_1 block to FP32.

void dequant_q4_1_row(const void * src, float * dst, size_t n_elements)

Dequantize Q4_1 row (multiple blocks)

void dequant_q4_k_block(const block_q4_K * block, float * output)

Dequantize a single Q4_K block to FP32.

void dequant_q4_k_row(const void * src, float * dst, size_t n_elements)

Dequantize Q4_K row (multiple blocks)

void dequant_q5_0_block(const block_q5_0 * block, float * output)

Dequantize a single Q5_0 block to FP32.

void dequant_q5_0_row(const void * src, float * dst, size_t n_elements)

Dequantize Q5_0 row (multiple blocks)

void dequant_q5_1_block(const block_q5_1 * block, float * output)

Dequantize a single Q5_1 block to FP32.

void dequant_q5_1_row(const void * src, float * dst, size_t n_elements)

Dequantize Q5_1 row (multiple blocks)

void dequant_q6_k_block(const block_q6_K * block, float * output)

Dequantize a single Q6_K block to FP32.

void dequant_q6_k_row(const void * src, float * dst, size_t n_elements)

Dequantize Q6_K row (multiple blocks)

void dequant_q8_0_block(const block_q8_0 * block, float * output)

Dequantize a single Q8_0 block to FP32.

void dequant_q8_0_row(const void * src, float * dst, size_t n_elements)

Dequantize Q8_0 row (multiple blocks)

void dequant_row(CKDataType dtype, const void * src, float * dst, size_t n_elements)

Dequantize a row of quantized data to FP32.

ck_huge_alloc_entry_t * detach_allocation(void * ptr)

ChatTemplateType detect_chat_template(const char * model_name)

int detect_physical_cores(void)

float dot_f16(const uint16_t * w_f16, const float * x, int K)

float dot_fp32_q5_0_block(const float * x, const block_q5_0 * block)

Compute dot product of FP32 input with Q5_0 weight block, with online Q8 quantization.

float dot_fp32_q8_0_block(const float * x, const block_q8_0 * block)

Compute dot product of FP32 input with Q8_0 weight block, with online Q8 quantization.

float dot_q4_0(const void * w_q4_0, const float * x, int K)

float dot_q4_1(const void * w_q4_1, const float * x, int K)

float dot_q4_k(const void * w_q4k, const float * x, int K)

Compute dot product of Q4_K row with FP32 vector.

float dot_q4_k_q8_k_ref(const block_q4_K * w, const block_q8_K * x, int k)

float dot_q5_0(const void * w_q5_0, const float * x, int K)

float dot_q5_0_q8_k_32_sse(const block_q5_0 * bw, const block_q8_K * ba, int q8_offset)

float dot_q5_1(const void * w_q5_1, const float * x, int K)

float dot_q6_k_q8_k_256_sse(const block_q6_K * bw, const block_q8_K * ba)

SSE Optimized dot product for Q6_K x Q8_K Q6_K layout: ql: 128 bytes (low 4 bits) qh: 64 bytes (high 2 bits) scales: 16 bytes (int8 scales) d: fp16 super-scale

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

Scalar dot product for Q6_K x Q8_K.

float dot_q6_k_ref(const block_q6_K * w, const float * x, int K)

float dot_q8_0(const void * w_q8_0, const float * x, int K)

void embedding_backward(const int32_t * token_ids, int token_count, const float * d_output, float * d_token_embeddings, float * d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Backward pass / gradient computation

void embedding_backward_bf16(const int32_t * token_ids, int token_count, const uint16_t * d_output, uint16_t * d_token_embeddings, uint16_t * d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Backward pass / gradient computation

void embedding_forward(const int32_t * token_ids, int token_count, int vocab_size, const float * token_embeddings, const float * pos_embeddings, float * output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Forward pass computation

void embedding_forward_bf16(const int32_t * token_ids, int token_count, int vocab_size, const uint16_t * token_embeddings, const uint16_t * pos_embeddings, uint16_t * output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Forward pass computation

void embedding_forward_q4_k(const int32_t * token_ids, int token_count, int vocab_size, const void * token_embeddings, const float * pos_embeddings, float * output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Forward pass computation

void embedding_forward_q6_k(const int32_t * token_ids, int token_count, int vocab_size, const void * token_embeddings, const float * pos_embeddings, float * output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Forward pass computation

void embedding_forward_q8_0(const int32_t * token_ids, int token_count, int vocab_size, const void * token_embeddings, const float * pos_embeddings, float * output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

Forward pass computation

void emit_body_fields(FILE * out, const CKIRV2Graph * graph, CKBufferRole role_filter, int activation_group)

size_t emit_body_values(FILE * out, const CKIRV2Graph * graph, const CKMemPlan * plan, CKBufferRole role_filter, int activation_group)

void emit_bump_bytes_assignment(FILE * out, const char * indent, const char * struct_prefix, const char * name, const CKDimToken * shape)

void emit_bump_bytes_assignment_weight_dtype(FILE * out, const char * indent, const char * struct_prefix, const char * name, const CKDimToken * shape, const char * dtype_expr)

void emit_dim_expr(FILE * out, CKDimKind dim)

void emit_footer_fields(FILE * out, const CKIRV2Graph * graph, CKBufferRole role_filter, int activation_group)

void emit_footer_values(FILE * out, const CKIRV2Graph * graph, const CKMemPlan * plan, CKBufferRole role_filter, int activation_group, size_t * offset)

void emit_global_aliases_to_layer(FILE * out)

void emit_global_allocations(FILE * out)

void emit_global_offset_fields(FILE * out)

void emit_header_fields(FILE * out, const CKIRV2Graph * graph, CKBufferRole role_filter, int activation_group)

void emit_header_values(FILE * out, const CKIRV2Graph * graph, const CKMemPlan * plan, CKBufferRole role_filter, int activation_group, size_t * offset)

int emit_kernel_manifest(const CKIRGraph * forward, const char * runtime_path)

void emit_layer_allocations(FILE * out)

void emit_layer_offsets_struct(FILE * out)

void emit_library_api(FILE * out, const CKIRGraph * forward)

void emit_model_struct(FILE * out)

void emit_offset_field(FILE * out, const char * name)

int emit_plan_sources(FILE * f, const CKPlanStep * plan, size_t plan_count, const CKIRGraph * cfg, const char ** seen, size_t * seen_count, size_t seen_cap)

int emit_runtime_preamble(FILE * out)

void emit_schedule_block(FILE * out, const CKIRV2Graph * graph, const char * func_name, const char * label, const char * runtime_sym)

void emit_sgd_update(FILE * out)

void emit_shape_expr(FILE * out, const CKDimToken * shape)

void emit_span_field(FILE * out, const char * label)

void emit_span_value(FILE * out, const char * label, size_t offset, size_t size, int comma)

void emit_training_conditional_assignment(FILE * out, const char * indent, const char * struct_prefix, const char * name, const CKDimToken * shape)

int emit_unique_source(FILE * f, const char * path, const char ** seen, size_t * seen_count, size_t seen_cap)

void emit_zero_grad(FILE * out)

int encode_chunk(CKTrueBPE * bpe, const char * chunk, int chunk_len, int32_t * ids, int max_ids, CKBPETokenList * list)

uint8_t encode_int4_nibble(int8_t value)

int encode_text_segment(CKTrueBPE * bpe, const char * text, int text_len, int32_t * ids, int max_ids)

void * engine_thread_func(void * arg)

bool eos_is_potential_prefix(const char * token)

Check if token might be start of EOS pattern.

void eos_pattern_init(ChatTemplateType tmpl)

bool eos_pattern_process(const char * token_text, char * out_buf, size_t * out_len, void(*)(char *, size_t *, const char *) output_fn, ChatTemplateType tmpl)

Process a token for EOS pattern detection.

void eos_pattern_reset(void)

int find_best_merge(const CKTrueBPE * bpe, const CKBPETokenList * list, size_t * best_pos, const CKBPEMerge ** best_merge)

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

int32_t find_longest_match(const CKTokenizer * tok, const char * text, size_t text_len, size_t pos, size_t * match_len)

int32_t find_longest_match_hash(const CKTokenizer * tok, const char * text, size_t text_len, size_t pos, size_t * match_len)

int32_t find_longest_match_trie(const CKTokenizer * tok, const char * text, size_t text_len, size_t pos, size_t * match_len)

bool find_model_in_cache(const char * model_name, char * lib_out, char * weights_out, size_t out_size)

int find_object_range(const char * json, const char * key, const char ** out_start, size_t * out_len)

void flatten_head_major(const float * attn_out, float * dst, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

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

uint16_t float_to_bf16(float f)

const char * format_bandwidth(float bw_gbs, char * buf, size_t buf_size)

const char * format_size(uint64_t size_mb, char * buf, size_t buf_size)

void free_entry(CKTokenizerHashEntry * entry, bool free_value)

int fused_kernels_compute_kv_tile(int l1_size, int head_dim, int bytes_per_elem)

Compute optimal KV tile size for flash attention.

void fused_kernels_report_stats(int hidden, int num_layers, int seq_len)

Report memory savings from mega-fusion.

int fused_kernels_validate_constraints(int l1_size, int head_dim, int kv_tile_size, int bytes_per_elem)

Validate cache constraints for fusion.

void fused_output_projection_residual(float * output, const float * o_all, const float * W_o, const float * b_o, const float * residual, int hidden, int num_heads, int head_dim)

Fused output projection with residual add.

void geglu_backward_fp32(const float * x, const float * d_out, float * d_x, int n_tokens, int dim)

GeGLU backward pass (fp32) Testtest_geglu.py::TestGeGLU::test_geglu_backward_fp32

void geglu_forward_bf16(const uint16_t * x, uint16_t * out, int tokens, int dim, float * scratch)

GeGLU forward pass (bf16) Testtest_geglu.py::TestGeGLU::test_geglu_forward_bf16

void geglu_forward_fp32(const float * x, float * out, int tokens, int dim)

GeGLU forward pass (fp32) Testtest_geglu.py::TestGeGLU::test_geglu_forward_fp32

void gemv_f16(float * y, const uint16_t * W, const float * x, int M, int K)

Auto-dispatch GEMV based on available SIMD.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient (scalar reference)

void gemv_f16_ref(float * y, const uint16_t * W, const float * x, int M, int K)

Matrix-vector multiply with FP16 weights (scalar reference)

void gemv_fused_q5_0_bias(float * y, const void * W, const float * x, const float * bias, int M, int K)

void gemv_fused_q5_0_bias_dispatch(float * y, const void * W, const float * x, const float * bias, int M, int K)

void gemv_fused_q5_0_bias_parallel_omp(float * y, const void * W, const float * x, const float * bias, int M, int K)

void gemv_fused_q8_0_bias(float * y, const void * W, const float * x, const float * bias, int M, int K)

void gemv_fused_q8_0_bias_dispatch(float * y, const void * W, const float * x, const float * bias, int M, int K)

void gemv_nt_q5_0_head_major_output(float * output, const float * attn_out, const void * wo, const float * bias, int tokens, int embed_dim, int num_heads, int head_dim)

Output projection reading head-major attention output (Q5_0 weights)

void gemv_q4_0(float * y, const void * W, const float * x, int M, int K)

Auto-dispatch GEMV.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient.

void gemv_q4_0_ref(float * y, const void * W, const float * x, int M, int K)

Matrix-vector multiply with Q4_0 weights (scalar reference)

void gemv_q4_1(float * y, const void * W, const float * x, int M, int K)

Auto-dispatch GEMV.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient.

void gemv_q4_1_ref(float * y, const void * W, const float * x, int M, int K)

Matrix-vector multiply with Q4_1 weights (scalar reference)

void gemv_q4_k(float * y, const void * W, const float * x, int M, int K)

Auto-dispatch GEMV based on available SIMD.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient (scalar reference)

void gemv_q4_k_q8_k(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_q8_k_amx(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_q8_k_avx(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_q8_k_avx2(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_q8_k_parallel(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

void gemv_q4_k_q8_k_parallel_simd(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

void gemv_q4_k_q8_k_ref(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_q8_k_sse(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_q8_k_vnni(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q4_k_ref(float * y, const void * W, const float * x, int M, int K)

Matrix-vector multiply with Q4_K weights (scalar reference)

void gemv_q5_0(float * y, const void * W, const float * x, int M, int K)

Auto-dispatch GEMV for Q5_0 weights based on CPU features.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient.

void gemv_q5_0_from_fp32(float * out, const void * W_q5_0, const float * x_fp32, const float * bias, int M, int K, block_q8_0 * x_q8_scratch)

void gemv_q5_0_parallel(float * y, const void * W, const float * x, int M, int K, int ith, int nth)

Parallel reference GEMV for Q5_0 × FP32.

void gemv_q5_0_parallel_simd(float * y, const void * W, const float * x, int M, int K, int ith, int nth)

Parallel SIMD GEMV for Q5_0 × FP32 with prefetching.

void gemv_q5_0_q8_0(float * y, const void * W, const void * x_q8, int M, int K)

Matrix-vector multiply with Q5_0 weights and Q8_0 input.

void gemv_q5_0_q8_0_parallel_omp(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q5_0_q8_0_parallel_simd(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

Parallel SIMD GEMV for Q5_0 x Q8_0 with prefetching.

void gemv_q5_0_ref(float * y, const void * W, const float * x, int M, int K)

Matrix-vector multiply with Q5_0 weights (scalar reference)

void gemv_q5_1(float * y, const void * W, const float * x, int M, int K)

Auto-dispatch GEMV.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient.

void gemv_q5_1_ref(float * y, const void * W, const float * x, int M, int K)

Matrix-vector multiply with Q5_1 weights (scalar reference)

void gemv_q5_k(float * y, const void * W, const float * x, int M, int K)

void gemv_q5_k_ref(float * y, const void * W, const float * x, int M, int K)

void gemv_q6_k(float * y, const void * W, const float * x, int M, int K)

void gemv_q6_k_q8_k(float * y, const void * W, const void * x_q8, int M, int K)

GEMV: y = W @ x where W is Q6_K and x is Q8_K.

void gemv_q6_k_q8_k_avx(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q6_k_q8_k_avx2(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q6_k_q8_k_avx512(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q6_k_q8_k_avx512_vbmi(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q6_k_q8_k_parallel(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

Parallel reference GEMV for Q6_K × Q8_K.

void gemv_q6_k_q8_k_parallel_simd(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

Parallel SIMD GEMV for Q6_K × Q8_K.

void gemv_q6_k_q8_k_ref(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q6_k_q8_k_sse(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q8_0(float * y, const void * W, const float * x, int M, int K)

Auto-dispatch GEMV for Q8_0 weights based on CPU features.

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

Auto-dispatch backward.

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

Backward pass: compute input gradient (scalar reference)

void gemv_q8_0_from_fp32(float * out, const void * W_q8_0, const float * x_fp32, const float * bias, int M, int K, block_q8_0 * x_q8_scratch)

void gemv_q8_0_parallel_simd(float * y, const void * W, const float * x, int M, int K, int ith, int nth)

Parallel SIMD GEMV for Q8_0 weights x FP32 input with prefetching.

void gemv_q8_0_q8_0(float * y, const void * W, const void * x_q8, int M, int K)

Matrix-vector multiply with Q8_0 weights and Q8_0 input.

void gemv_q8_0_q8_0_parallel(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

Parallel reference GEMV for Q8_0 x Q8_0.

void gemv_q8_0_q8_0_parallel_omp(float * y, const void * W, const void * x_q8, int M, int K)

void gemv_q8_0_q8_0_parallel_simd(float * y, const void * W, const void * x_q8, int M, int K, int ith, int nth)

Parallel SIMD GEMV for Q8_0 x Q8_0 with prefetching.

void gemv_q8_0_ref(float * y, const void * W, const float * x, int M, int K)

Matrix-vector multiply with Q8_0 weights (scalar reference)

const char * get_cache_dir(void)

const CPUInfo * get_cpu_info(void)

int get_current_cpu(void)

int get_numa_node_for_cpu(int cpu)

int get_optimal_decode_threads(void)

void get_q5_k_scale_min(int j, const uint8_t * scales, uint8_t * scale, uint8_t * min)

double get_time_ms(void)

void global_pool_init(void)

int gpt2_decode_byte(const unsigned char * s, int len)

int gpt2_pretokenize(const char * text, int text_len, PretokChunk * chunks, int max_chunks)

void gradient_accumulate_bf16(uint16_t * dst, const uint16_t * src, size_t numel)

Accumulate gradients: dst += src (bf16)

void gradient_accumulate_f32(float * dst, const float * src, size_t numel)

Accumulate gradients: dst += src (fp32)

float gradient_clip_norm_bf16(uint16_t * grad, size_t numel, float max_norm)

Clip gradient norm (bf16)

float gradient_clip_norm_f32(float * grad, size_t numel, float max_norm)

Clip gradient norm (fp32)

void gradient_scale_bf16(uint16_t * grad, size_t numel, float scale)

Scale gradients: grad *= scale (bf16)

void gradient_scale_f32(float * grad, size_t numel, float scale)

Scale gradients by a constant: grad *= scale (fp32)

void handle_sigint(int sig)

int has_cpu_flag(const char * flags, const char * flag)

uint32_t hash_pair(int32_t left, int32_t right)

uint32_t hash_string(const char * s, int len)

float hmax256_ps_fused(__m256 v)

float hsum256_ps_fused(__m256 v)

int32_t hsum_epi32_sse(__m128i v)

void im2patch(const float * image, float * patches, int C, int H, int W, int P)

im2patch: Transforms an image into a sequence of flattened patches.

void im2patch_bf16(const uint16_t * image, uint16_t * patches, int C, int H, int W, int P)

int init_tokens_from_text(CKTrueBPE * bpe, CKBPETokenList * list, const char * text, int text_len)

int is_activation_role(CKBufferRole role)

bool is_bpe_digit(const char * s, int len)

bool is_bpe_letter(const char * s, int len)

bool is_bpe_newline(const char * s, int len)

bool is_bpe_punct(const char * s, int len)

bool is_digit(unsigned char c)

bool is_eos_token(const CLIOptions * opt, int token)

int is_footer_global(const char * name)

bool is_gpt2_space(const char * s, int len)

bool is_letter(unsigned char c)

bool is_whitespace(unsigned char c)

bool is_word_prefix_char(const char * s, int len)

int json_match_char(JSONParser * p, char c)

int json_parse_int(JSONParser * p, int * out)

int json_parse_string(JSONParser * p, char * buf, int max_len)

void json_skip_value(JSONParser * p)

void json_skip_whitespace(JSONParser * p)

void kv_cache_repack_head_major_inplace(float * buf, int num_heads, int tokens, int cache_capacity, int aligned_head_dim)

void kv_cache_store(float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)

void kv_cache_write_head_major(const float *__restrict k_token, const float *__restrict v_token, float *__restrict k_cache, float *__restrict v_cache, int num_kv_heads, int token_index, int cache_capacity, int head_dim, int aligned_head_dim)

void layout_transformer_from_ir(TransformerModel * m, const CKIRGraph * ir)

Compute a simple forward-only layout for TransformerModel based on: CKModelConfig (dims, heads, vocab, context) The IR graph structure (number of layers, op types)

void list_available_models(void)

bool load_eos_from_vocab_json(const char * weights_path, CLIOptions * opt)

int load_manifest(const char * path, ManifestEntry ** entries, int * num_entries)

bool load_model_api(const char * lib_path, ModelAPI * api)

int load_weights(QWEN2_DECODEModel * model, const char * bump_path, const char * manifest_path)

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

void logits_copy_to_position(const float *__restrict src, float *__restrict dst, int position, int vocab_size)

Copy logits to position-indexed location in output buffer.

int main(int argc, char ** argv)

CKOpType map_forward_to_backward(CKOpType op)

Backward pass / gradient computation

int match_special_token(const CKTrueBPE * bpe, const char * text, int text_len, int pos)

int max_k_for_query(int t_q, int T_q, int T_k)

void mega_fuse_get_optimal_tiles(int * q_tile, int * kv_tile, int head_dim)

Get optimal tile sizes for current CPU.

void mega_fuse_output_proj_residual(const float * attn_token, const float * wo, const float * bo, const float * residual, float * output, int embed_dim, int aligned_embed_dim, int num_heads, int head_dim, int aligned_head_dim)

void mega_fuse_report_stats(int hidden, int num_layers, int seq_len)

Report memory savings from mega-fusion.

size_t merge_hash(uint64_t key, size_t num_buckets)

uint64_t merge_key(int32_t left_id, int32_t right_id)

CKMergeTable * merge_table_create(size_t num_buckets)

void merge_table_free(CKMergeTable * table)

int merge_table_insert(CKMergeTable * table, const CKBPEMerge * merge)

const CKBPEMerge * merge_table_lookup(const CKMergeTable * table, int32_t left_id, int32_t right_id)

int model_align_elems(int elems, int elem_bytes, int align_bytes)

void model_decode(MODELModel * model, const int * token, int token_index)

void model_decode_token(MODELModel * model, const int * token, int token_index)

void model_forward(MODELModel * model, const int * tokens, int num_tokens)

Forward pass computation

void model_forward_prefill_impl(MODELModel * model, const int * tokens, int num_tokens)

Forward pass computation

void model_layer_0_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_0_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_10_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_10_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_11_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_11_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_12_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_12_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_13_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_13_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_14_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_14_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_15_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_15_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_16_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_16_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_17_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_17_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_18_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_18_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_19_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_19_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_1_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_1_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_20_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_20_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_21_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_21_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_22_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_22_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_23_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_23_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_2_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_2_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_3_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_3_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_4_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_4_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_5_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_5_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_6_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_6_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_7_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_7_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_8_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_8_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_9_decode(MODELModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void model_layer_9_prefill(MODELModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

int model_model_allocate(MODELModel * model)

void model_model_free(MODELModel * model)

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

int model_verify_canaries(MODELModel * model)

void moe_accumulate_expert_f32(float * output, const float * expert_output, float routing_weight, int hidden_dim)

Accumulate expert output: output += routing_weight * expert_output.

const char * op_name(CKOpType op)

void out_proj_head_major_q5_0_q8_0(const uint8_t * attn_q8, const void * wo, const float * bias, float * output, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void out_proj_head_major_q8_0_q8_0(const uint8_t * attn_q8, const void * wo, const float * bias, float * output, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)

void output_append(char * buf, size_t * len, const char * text)

void output_flush(char * buf, size_t * len)

void output_token(char * buf, size_t * len, const char * token)

void pack_a_panel(const float * A, int lda, float * Ap, int mc, int kc, int mr)

void pack_b_panel(const float * B, int ldb, float * Bp, int kc, int nc, int nr)

bool parse_args(int argc, char ** argv, CLIOptions * opt)

bool parse_eos_ids(const char * arg, CLIOptions * opt)

int parse_float_field_any(const char * json, size_t len, const char *const * keys, float * out_value)

int parse_float_field_in_range(const char * json, size_t len, const char * key, float * out_value)

int parse_int_field(const char * json, const char * key, int * out_value)

int parse_int_field_any(const char * json, size_t len, const char *const * keys, int * out_value)

int parse_int_field_in_range(const char * json, size_t len, const char * key, int * out_value)

bool parse_manifest_entry(const char * json, const char * name, size_t * offset, size_t * size)

int parse_manifest_int(const char * json, const char * key)

CKOpType parse_op(const char * s)

unsigned long long parse_u64(const char * s)

void patch2im(const float * d_patches, float * d_image, int C, int H, int W, int P)

patch2im: Accumulates gradients from patches back into the image. (Backward pass)

void patch2im_bf16(const uint16_t * d_patches, uint16_t * d_image, int C, int H, int W, int P)

float pcie_bandwidth_gbs(int gen, int width)

size_t plan_size(const CKMemPlan * plan, int idx)

CKPoolBlock * pool_new_block(size_t capacity)

int preprocess_bpe_spaces(const char * text, int text_len, char * out, int out_max, CKSpacePrefixStyle style)

int preprocess_spm_llama_text(const char * text, int text_len, char * out, int out_max, bool add_space_prefix)

int preprocess_spm_text(const char * text, int text_len, char * out, int out_max, bool add_space_prefix)

int preprocess_text(const CKTrueBPE * bpe, const char * text, int text_len, char * out, int out_max)

void print_banner(void)

void print_cpu_info(void)

void print_header(const char * title)

void print_help(const char * prog)

void print_ok(const char * msg)

void print_progress(int token_id, float token_per_sec)

void print_section(const char * title)

void print_tree_item(int level, int is_last, const char * fmt, ...)

void print_usage(const char * prog)

void print_version(void)

void print_warning(const char * msg)

bool process_repl_command(const char * line, CLIOptions * opt, ModelAPI * api)

void qk_norm_forward(float * q, float * k, const float * q_gamma, const float * k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

Per-head RMSNorm on Q and K.

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

void qkv_projection_parallel(const void * ln1_q8, const void * WQ, const void * WK, const void * WV, float * q_out, float * k_out, float * v_out, int H, int H_kv, int head_dim, int embed_dim, int num_threads)

Parallel Q/K/V projection for single token decode.

int qkv_q8_0_dtype_supported(CKDataType dt)

int qkv_q8_k_dtype_supported(CKDataType dt)

void quantize_attn_out_head_major_q8_0(const float * attn_out, uint8_t * dst, int tokens, int num_heads, int aligned_head_dim)

void quantize_batch_q8_0(const float * x, void * vy, int num_rows, int k)

Batch quantize FP32 to Q8_0 format (row-major output)

void quantize_batch_q8_k(const float * x, void * vy, int num_rows, int k)

Batch quantize FP32 to Q8_K format (row-major output)

void quantize_row_q8_0(const float * x, void * vy, int k)

Quantize FP32 to Q8_0 format (scalar reference)

void quantize_row_q8_k(const float * x, void * vy, int k)

void quantize_row_q8_k_ref(const float * x, void * vy, int k)

void quantize_row_q8_k_sse(const float * x, void * vy, int k)

int qwen2_0_5b_decode_align_elems(int elems, int elem_bytes, int align_bytes)

void qwen2_0_5b_decode_decode(QWEN2_0_5B_DECODEModel * model, const int * token, int token_index)

void qwen2_0_5b_decode_decode_token(QWEN2_0_5B_DECODEModel * model, const int * token, int token_index)

void qwen2_0_5b_decode_forward(QWEN2_0_5B_DECODEModel * model, const int * tokens, int num_tokens)

Forward pass computation

void qwen2_0_5b_decode_forward_prefill_impl(QWEN2_0_5B_DECODEModel * model, const int * tokens, int num_tokens)

Forward pass computation

void qwen2_0_5b_decode_layer_0_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_0_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_10_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_10_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_11_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_11_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_12_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_12_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_13_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_13_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_14_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_14_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_15_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_15_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_16_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_16_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_17_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_17_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_18_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_18_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_19_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_19_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_1_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_1_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_20_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_20_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_21_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_21_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_22_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_22_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_23_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_23_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_2_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_2_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_3_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_3_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_4_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_4_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_5_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_5_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_6_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_6_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_7_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_7_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_8_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_8_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_9_decode(QWEN2_0_5B_DECODEModel * model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

void qwen2_0_5b_decode_layer_9_prefill(QWEN2_0_5B_DECODEModel * model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

int qwen2_0_5b_decode_model_allocate(QWEN2_0_5B_DECODEModel * model)

void qwen2_0_5b_decode_model_free(QWEN2_0_5B_DECODEModel * model)

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

int qwen2_0_5b_decode_verify_canaries(QWEN2_0_5B_DECODEModel * model)

int read_file_int(const char * path)

int read_file_string(const char * path, char * buf, size_t buf_size)

uint64_t read_file_uint64(const char * path)

int read_floats(FILE * f, float * dst, size_t count)

bool read_manifest_entry(const char * json_path, const char * entry_name, size_t * out_offset, size_t * out_size)

char * read_prompt_file(const char * path)

int record_allocation(void * ptr, size_t len, int was_mmap)

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

Backward pass / gradient computation

void relu_backward_bf16(const uint16_t * input, const uint16_t * d_output, uint16_t * d_input, size_t n)

Backward pass / gradient computation

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

Forward pass computation

void relu_forward_bf16(const uint16_t * input, uint16_t * output, size_t n)

Forward pass computation

void relu_forward_inplace(float * data, size_t n)

Forward pass computation

void relu_forward_inplace_bf16(uint16_t * data, size_t n)

Forward pass computation

void residual_add(float * residual, float * addend, int n)

void residual_add_parallel(const float * a, const float * b, float * out, int n, int ith, int nth)

Single-token decode with parallel SIMD kernels.

size_t resolve_dim(const CKModelConfig * cfg, const CKIRV2AlignInfo * align, CKDimKind kind, int tokens_override)

size_t resolve_shape_elems(const CKModelConfig * cfg, const CKIRV2AlignInfo * align, const CKDimToken * shape, int tokens_override)

bool resolve_symbol(void * handle, const char * name, void ** out_ptr, bool required)

void run_benchmark(void * model, int num_tokens)

int run_command(const char * cmd, char * output, size_t output_size)

void run_generation_test(void * model, int num_tokens)

int run_inference(const char * bump_path, const char * manifest_path, const char * tokenizer_path, const char * prompt, int max_tokens, float temperature, int topk)

int run_prompt(ModelAPI * api, CKTrueBPE * tokenizer, CLIOptions * opt, const char * input)

int sample_argmax(const float * logits, int vocab_size)

int sample_token(float * logits, int vocab_size, float temp, int top_k)

int sample_top_p(float * logits, int vocab_size, float temperature, float top_p)

int sample_topk(float * probs, int vocab_size, int topk)

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

Scaled copy: y = alpha * x.

size_t score_index(int h, int i, int j, int aligned_context_window)

void sgd_momentum_update_bf16(const uint16_t * grad, uint16_t * weight, float * velocity, size_t numel, float lr, float momentum, float weight_decay)

SGD with momentum (bf16 weights/gradients)

void sgd_momentum_update_f32(const float * grad, float * weight, float * velocity, size_t numel, float lr, float momentum, float weight_decay)

SGD with momentum optimizer update (fp32 version)

void silu(float * x, int n)

float silu_prefill(float x)

float silu_scalar(float x)

int simd_strcmp(const char * s1, const char * s2)

void simple_embedding(const int32_t * tokens, int num_tokens, const float * weight, float * output, int vocab_size, int embed_dim)

void spm_build_byte_lookup(CKTokenizer * tok, const char * strings, const int32_t * offsets, int vocab_size)

int spm_count_unknown_run(const CKTokenizer * tok, const char * text, int text_len, size_t pos)

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

int spm_find_candidates_at_pos(const CKTokenizer * tok, const char * text, int text_len, size_t pos, int32_t * candidates, int max_candidates)

int32_t spm_get_byte_token(const CKTokenizer * tok, unsigned char byte_val)

bool spm_is_byte_token(const CKTokenizer * tok, int32_t token_id)

int spm_llama_resegment_node(const CKTokenizer * tok, const SpmLlamaNode * nodes, int node_id, int32_t * ids, int max_ids, int out_idx)

bool spm_token_allowed_in_dp(const CKTokenizer * tok, int32_t token_id)

bool spm_token_is_byte_format(const char * token)

int starts_with(const char * s, const char * prefix)

int token_list_append(CKBPETokenList * list, const char * str, size_t len, int32_t id)

void token_list_clear(CKBPETokenList * list)

CKBPETokenList * token_list_create(size_t initial_capacity)

void token_list_free(CKBPETokenList * list)

int token_list_merge_at(CKBPETokenList * list, size_t pos, const char * merged_str, size_t merged_len, int32_t merged_id)

int32_t * tokenize(const char * text, int * num_tokens)

void topk_batched_f32(const float * scores, int num_tokens, int n_experts, int k, int * indices, float * weights)

Batched top-K selection for multiple tokens.

void topk_f32(const float * scores, int n, int k, int * indices, float * values)

Find top-K indices and values from a score vector.

int topology_discover(SystemTopology * topo)

int topology_discover_affinity(AffinityInfo * aff)

int topology_discover_cache(CacheTopology * cache)

int topology_discover_cpu(CPUInfo * cpu)

int topology_discover_memory(MemoryInfo * mem)

int topology_discover_network(NetworkTopology * net)

int topology_discover_numa(NUMATopology * numa)

int topology_discover_pcie(PCIeTopology * pcie)

int topology_estimate_channels_from_bandwidth(float measured_bw_gbs, int memory_speed_mhz, const char * memory_type)

float topology_estimate_memory_bandwidth(const MemoryInfo * mem)

float topology_estimate_network_training_time(const NetworkTopology * net, uint64_t model_size_mb)

int topology_generate_recommendations(const SystemTopology * topo, RecommendationList * recs)

float topology_measure_memory_bandwidth(void)

float topology_measure_memory_bandwidth_ex(int * numa_node_out, int * num_threads_out)

void topology_print_affinity(const AffinityInfo * aff)

void topology_print_cache(const CacheTopology * cache, int logical_cores)

void topology_print_cpu(const CPUInfo * cpu)

void topology_print_distributed_potential(const SystemTopology * topo)

void topology_print_memory(const MemoryInfo * mem)

void topology_print_network(const NetworkTopology * net)

void topology_print_numa(const NUMATopology * numa, int sockets)

void topology_print_pcie(const PCIeTopology * pcie)

void topology_print_recommendations(const RecommendationList * recs)

void topology_print_summary(const SystemTopology * topo)

void trim_string(char * str)

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

Unpack Q4_K sub-block scales and mins.

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

Unpack Q5_K sub-block scales and mins.

int utf8_char_len(unsigned char c)

int utf8_len(unsigned char c)

void v6_prefill(const float * embed_weight, const int32_t * tokens, int num_tokens, float * logits)

void vec_dot_q5_0_q8_0(int n, float * s, const void * vx, const void * vy)

Auto-dispatch quantized dot product Q5_0 x Q8_0.

void vec_dot_q5_0_q8_0_ref(int n, float * s, const void * vx, const void * vy)

Quantized dot product: Q5_0 weights x Q8_0 input (scalar reference)

void vec_dot_q6_k_q8_k(int n, float * s, const void * vx, const void * vy)

Q6_K x Q8_K dot product (single row)

void vec_dot_q8_0_q8_0(int n, float * s, const void * vx, const void * vy)

Auto-dispatch quantized dot product Q8_0 x Q8_0.

void vec_dot_q8_0_q8_0_ref(int n, float * s, const void * vx, const void * vy)

Quantized dot product: Q8_0 weights x Q8_0 input (scalar reference)

void vec_scale_parallel(float * y, float scale, int n, int ith, int nth)

void vec_zero_parallel(float * y, int n, int ith, int nth)

void weighted_sum_f32(float * y, const float ** vectors, const float * weights, int k, int n)

Weighted sum of k vectors: y = sum_i(weights[i] * vectors[i])

void * worker_main(void * arg)

void zero_gradients_bf16(uint16_t * grad, size_t numel)

Zero out gradient buffer (bf16)

void zero_gradients_f32(float * grad, size_t numel)

Zero out gradient buffer (fp32)