← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
ck_parity_api.c File Reference

C-Kernel-Engine Parity Testing API Implementation. More...

#include "ck_parity_api.h"
#include "ckernel_quant.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>

Go to the source code of this file.

Functions

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)
 
void attn_gate_sigmoid_mul_forward (const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
 
int ck_get_block_q4_k_size (void)
 Get Q4_K block size in bytes.
 
int ck_get_block_q6_k_size (void)
 Get Q6_K block size in bytes.
 
int ck_get_block_q8_k_size (void)
 Get Q8_K block size in bytes.
 
int ck_get_qk_k (void)
 Get QK_K (elements per super-block)
 
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_attn_gate_sigmoid_mul (const float *x, const float *gate, float *out, int rows, int dim)
 Multiply attention output rows by sigmoid(gate) elementwise.
 
void ck_test_dequant_q4_0 (const void *src, float *dst, int n)
 Dequantize Q4_0 data to FP32.
 
void ck_test_dequant_q4_k (const void *src, float *dst, int n)
 Dequantize Q4_K data to FP32.
 
void ck_test_dequant_q6_k (const void *src, float *dst, int n)
 Dequantize Q6_K data to FP32.
 
void ck_test_gated_deltanet_autoregressive (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)
 Gated DeltaNet autoregressive update.
 
void ck_test_gemm_q4_k (const void *weight_q4k, const float *input_f32, float *output, int rows, int cols, int n_tokens)
 Q4_K GEMM - batched matrix multiply with quantized weights.
 
void ck_test_gemm_q5_0 (const void *weight_q5_0, const float *input_f32, float *output, int rows, int cols, int n_tokens)
 Test Q5_0 x Q8_0 GEMM (batch matrix multiply)
 
void ck_test_gemm_q6_k (const void *weight_q6k, const float *input_f32, float *output, int rows, int cols, int n_tokens)
 Test Q6_K x Q8_K GEMM (batch matrix multiply)
 
void ck_test_gemm_q8_0 (const void *weight_q8_0, const float *input_f32, float *output, int rows, int cols, int n_tokens)
 Test Q8_0 x Q8_0 GEMM (batch matrix multiply)
 
void ck_test_gemv_q4_k (const void *weight_q4k, const float *input_f32, float *output, int cols)
 Q4_K GEMV - dot product of quantized weights and FP32 input.
 
void ck_test_gemv_q5_0 (const void *weight_q5_0, const float *input_f32, float *output, int rows, int cols)
 Q5_0 GEMV - matrix-vector multiply with Q5_0 weights.
 
void ck_test_gemv_q5_0_q8_0 (const void *weight_q5_0, const float *input_f32, float *output, int rows, int cols)
 Q5_0 x Q8_0 quantized GEMV - matches llama.cpp's approach.
 
void ck_test_gemv_q6_k (const void *weight_q6k, const float *input_f32, float *output, int cols)
 Q6_K GEMV.
 
void ck_test_gemv_q8_0 (const void *weight_q8_0, const float *input_f32, float *output, int rows, int cols)
 Q8_0 GEMV - matrix-vector multiply with Q8_0 weights.
 
void ck_test_gemv_q8_0_q8_0 (const void *weight_q8_0, const float *input_f32, float *output, int rows, int cols)
 Q8_0 x Q8_0 quantized GEMV - matches llama.cpp's approach.
 
void ck_test_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 ck_test_quantize_q8_k (const float *src, void *dst, int n)
 Quantize FP32 to Q8_K (for activations)
 
void ck_test_recurrent_conv_state_update (const float *state_in, const float *q, const float *k, const float *v, float *conv_x, float *state_out, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
 Build the recurrent convolution input history window.
 
void ck_test_recurrent_dt_gate (const float *alpha, const float *dt_bias, const float *a, float *gate, int rows, int dim)
 Transform recurrent alpha rows into the DeltaNet gate.
 
void ck_test_recurrent_norm_gate (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
 Per-head RMSNorm followed by SiLU(z) gating for recurrent outputs.
 
void ck_test_recurrent_qk_l2_norm (float *q, float *k, int rows, int q_dim, int k_dim, int head_dim, float eps)
 Apply per-head L2 normalization to recurrent Q/K rows in-place.
 
void ck_test_recurrent_silu (const float *x, float *out, int rows, int dim)
 Apply SiLU elementwise to recurrent rows.
 
void ck_test_recurrent_split_conv_qkv (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
 Split the post-convolution recurrent packed QKV rows.
 
void ck_test_recurrent_split_qkv (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
 Split a packed recurrent QKV matrix into explicit Q, K, and V outputs.
 
void ck_test_rmsnorm (const float *input, const float *weight, float *output, int n_tokens, int dim, float eps)
 RMSNorm.
 
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 ck_test_softmax (const float *input, float *output, int n)
 Softmax (simple, non-causal)
 
void ck_test_split_q_gate (const float *packed_qg, float *q, float *gate, int rows, int q_dim, int gate_dim, int group_dim)
 Split a packed full-attention Q+gate matrix into Q rows and gate rows.
 
void ck_test_ssm_conv1d (const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
 qwen3next/Qwen3.5 SSM causal depthwise convolution.
 
void ck_test_swiglu (const float *gate_up, float *output, int n_tokens, int intermediate_dim)
 SwiGLU activation.
 
void ck_test_vec_dot_q4_k_q8_k (const void *weight_q4_k, const void *input_q8_k, float *output, int cols)
 
void ck_test_vec_dot_q5_0_q8_0 (const void *weight_q5_0, const void *input_q8_0, float *output, int cols)
 Direct Q5_0 x Q8_0 dot product test (takes pre-quantized Q8_0 input)
 
void ck_test_vec_dot_q6_k_q8_k (const void *weight_q6_k, const void *input_q8_k, float *output, int cols)
 
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 dequant_q4_0_row (const void *src, float *dst, size_t n_elements)
 Dequantize Q4_0 row (multiple blocks)
 
void dequant_q4_k_row (const void *src, float *dst, size_t n_elements)
 Dequantize Q4_K row (multiple blocks)
 
void dequant_q6_k_row (const void *src, float *dst, size_t n_elements)
 Dequantize Q6_K row (multiple blocks)
 
void gated_deltanet_autoregressive_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)
 
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_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_q6_k_q8_k (const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
 
void gemm_nt_q8_0_q8_0 (const void *A_q8, const void *B_q8, const float *bias, float *C, int M, int N, int K)
 gemm_nt_q8_0_q8_0 with optional bias (matches header signature)
 
void gemv_q4_k_q8_k (float *y, const void *W, const void *x_q8, int M, int K)
 
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_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_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_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_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 mega_fused_outproj_mlp_prefill (float *output, const float *attn_out, const float *residual, const float *ln2_gamma, const void *wo, const float *bo, int wo_dt, const void *w1, const float *b1, int w1_dt, const void *w2, const float *b2, int w2_dt, int tokens, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim, int intermediate_dim, int aligned_intermediate_dim, float eps, void *scratch)
 
size_t mega_fused_outproj_mlp_prefill_scratch_size (int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim, int aligned_intermediate_dim)
 Get scratch buffer size for mega_fused_outproj_mlp_prefill.
 
void 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 recurrent_conv_state_update_forward (const float *state_in, const float *q, const float *k, const float *v, float *conv_x, float *state_out, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
 
void recurrent_dt_gate_forward (const float *alpha, const float *dt_bias, const float *a, float *gate, int rows, int num_heads, int state_dim)
 
void recurrent_norm_gate_forward (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
 
void recurrent_qk_l2_norm_forward (float *q, float *k, int rows, int q_dim, int k_dim, int head_dim, float eps)
 
void recurrent_silu_forward (const float *x, float *out, int rows, int dim)
 
void recurrent_split_conv_qkv_forward (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
 
void recurrent_split_qkv_forward (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
 
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)
 
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)
 
void rope_precompute_cache (float *cos_cache, float *sin_cache, int max_seq_len, int head_dim, float base, int rotary_dim, const char *scaling_type, float scaling_factor)
 
void split_q_gate_forward (const float *packed_qg, float *q, float *gate, int rows, int q_dim, int gate_dim, int group_dim)
 
void ssm_conv1d_forward (const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
 
void swiglu_forward (const float *input, float *output, int tokens, int dim)
 
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_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.
 

Detailed Description

C-Kernel-Engine Parity Testing API Implementation.

Wraps CK kernels for parity testing against llama.cpp/ggml.

Definition in file ck_parity_api.c.

Function Documentation

◆ attention_forward_causal_head_major_gqa_flash_strided()

void attention_forward_causal_head_major_gqa_flash_strided ( const float *  q,
const float *  k,
const float *  v,
float *  output,
int  num_heads,
int  num_kv_heads,
int  num_tokens,
int  head_dim,
int  aligned_head_dim,
int  kv_stride_tokens 
)
extern

Flash attention forward with custom KV stride (for KV cache)

Test:

test_flash_attention.py::TestFlashAttention::test_flash_strided

test_kv_cache_attention.py::TestKVCacheAttention::test_flash_attention

Variant with configurable kv_stride_tokens for KV cache layouts where K/V may not be contiguous in memory.

After changes: make test

Definition at line 4496 of file attention_kernels.c.

4506{
4508 num_heads, num_kv_heads,
4509 num_tokens, head_dim,
4510 aligned_head_dim,
4511 kv_stride_tokens,
4512 /*causal=*/1,
4513 /*round_full_kv_fp16=*/0,
4514 /*output_token_major=*/0,
4515 1.0f / sqrtf((float)head_dim));
4516}
static void attention_forward_head_major_gqa_flash_impl(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int causal, int round_full_kv_fp16, int output_token_major, float scale)

Referenced by ck_test_attention_causal().

◆ attn_gate_sigmoid_mul_forward()

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

Definition at line 154 of file hybrid_attention_kernels.c.

159 {
160 if (!x || !gate || !out || rows <= 0 || num_heads <= 0 ||
161 state_dim <= 0 || num_heads > INT_MAX / state_dim) {
162 return;
163 }
164 const int dim = num_heads * state_dim;
165 const size_t elements = (size_t)rows * (size_t)dim;
166 ck_attn_gate_sigmoid_mul_args_t args = {x, gate, out, dim};
167 const size_t minimum_elements_per_thread = 32768u;
168
169 const int partial_x_overlap = out != x &&
170 hybrid_ranges_overlap(out, x, elements);
171 const int gate_overlap = hybrid_ranges_overlap(out, gate, elements);
172 if (partial_x_overlap || gate_overlap ||
173 elements < 2u * minimum_elements_per_thread) {
174 attn_gate_sigmoid_mul_row_range(0, rows, &args);
175 return;
176 }
177
178 ck_threadpool_t *pool = ck_threadpool_global();
179 int active = pool ? ck_threadpool_n_threads(pool) : 1;
180 size_t useful_threads =
181 (elements + minimum_elements_per_thread - 1u) /
182 minimum_elements_per_thread;
183 if ((size_t)active > useful_threads) {
184 active = (int)useful_threads;
185 }
186 if (active > rows) {
187 active = rows;
188 }
189
190 if (active > 1 && ck_threadpool_thread_id(pool) <= 0) {
192 pool, active, 0, rows, 4,
194 } else {
195 attn_gate_sigmoid_mul_row_range(0, rows, &args);
196 }
197}
void ck_threadpool_parallel_for_n(ck_threadpool_t *pool, int active_threads, int begin, int end, int grain_size, ck_range_fn_t fn, void *args)
ck_threadpool_t * ck_threadpool_global(void)
int ck_threadpool_thread_id(const ck_threadpool_t *pool)
int ck_threadpool_n_threads(const ck_threadpool_t *pool)
static int hybrid_ranges_overlap(const float *lhs, const float *rhs, size_t elements)
static void attn_gate_sigmoid_mul_row_range(int begin, int end, void *opaque)

Referenced by ck_test_attn_gate_sigmoid_mul().

◆ ck_get_block_q4_k_size()

int ck_get_block_q4_k_size ( void  )

Get Q4_K block size in bytes.

Definition at line 1026 of file ck_parity_api.c.

1027{
1028 return sizeof(block_q4_K);
1029}

◆ ck_get_block_q6_k_size()

int ck_get_block_q6_k_size ( void  )

Get Q6_K block size in bytes.

Definition at line 1031 of file ck_parity_api.c.

1032{
1033 return sizeof(block_q6_K);
1034}

◆ ck_get_block_q8_k_size()

int ck_get_block_q8_k_size ( void  )

Get Q8_K block size in bytes.

Definition at line 1036 of file ck_parity_api.c.

1037{
1038 return sizeof(block_q8_K);
1039}

◆ ck_get_qk_k()

int ck_get_qk_k ( void  )

Get QK_K (elements per super-block)

Definition at line 1041 of file ck_parity_api.c.

1042{
1043 return QK_K;
1044}
#define QK_K

References QK_K.

◆ ck_test_attention_causal()

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)

Layout (head-major, matches llama.cpp test): Q: [num_heads, tokens, head_dim] K: [num_kv_heads, seq_len, head_dim] V: [num_kv_heads, seq_len, head_dim] out: [num_heads, tokens, head_dim]

Supports GQA (grouped-query attention) where num_heads > num_kv_heads. Causal masking: token t can only attend to positions 0..t (inclusive).

Parameters
qQuery [num_heads, tokens, head_dim]
kKey [num_kv_heads, seq_len, head_dim]
vValue [num_kv_heads, seq_len, head_dim]
outOutput [num_heads, tokens, head_dim]
num_headsNumber of query heads
num_kv_headsNumber of key/value heads (for GQA)
tokensNumber of query tokens
seq_lenKey/value sequence length (for prefill: seq_len == tokens)
head_dimDimension per head

Definition at line 884 of file ck_parity_api.c.

893{
894 /* For prefill, seq_len == tokens, and kv_stride == tokens.
895 * The CK kernel expects strided KV layout with kv_stride_tokens parameter.
896 * For parity testing with contiguous tensors, kv_stride = seq_len.
897 */
899 q, k, v, out,
900 num_heads, num_kv_heads, tokens,
901 head_dim, head_dim, /* aligned_head_dim = head_dim for testing */
902 seq_len /* kv_stride_tokens = seq_len for contiguous KV */
903 );
904}
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)

References attention_forward_causal_head_major_gqa_flash_strided().

◆ ck_test_attn_gate_sigmoid_mul()

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

Multiply attention output rows by sigmoid(gate) elementwise.

Layout: x : [rows, dim] gate : [rows, dim] out : [rows, dim]

Definition at line 859 of file ck_parity_api.c.

864{
865 attn_gate_sigmoid_mul_forward(x, gate, out, rows, 1, dim);
866}
void attn_gate_sigmoid_mul_forward(const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)

References attn_gate_sigmoid_mul_forward().

◆ ck_test_dequant_q4_0()

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

Dequantize Q4_0 data to FP32.

Definition at line 202 of file ck_parity_api.c.

203{
204 dequant_q4_0_row(src, dst, (size_t)n);
205}
void dequant_q4_0_row(const void *src, float *dst, size_t n_elements)
Dequantize Q4_0 row (multiple blocks)

References dequant_q4_0_row().

◆ ck_test_dequant_q4_k()

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

Dequantize Q4_K data to FP32.

Parameters
srcInput Q4_K blocks
dstOutput FP32 values
nNumber of elements (must be multiple of 256)

Definition at line 192 of file ck_parity_api.c.

193{
194 dequant_q4_k_row(src, dst, (size_t)n);
195}
void dequant_q4_k_row(const void *src, float *dst, size_t n_elements)
Dequantize Q4_K row (multiple blocks)

References dequant_q4_k_row().

◆ ck_test_dequant_q6_k()

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

Dequantize Q6_K data to FP32.

Definition at line 197 of file ck_parity_api.c.

198{
199 dequant_q6_k_row(src, dst, (size_t)n);
200}
void dequant_q6_k_row(const void *src, float *dst, size_t n_elements)
Dequantize Q6_K row (multiple blocks)

References dequant_q6_k_row().

◆ ck_test_gated_deltanet_autoregressive()

void ck_test_gated_deltanet_autoregressive ( const float *  q,
const float *  k,
const float *  v,
const float *  g,
const float *  beta,
const float *  state_in,
float *  state_out,
float *  out,
int  num_heads,
int  state_dim,
float  norm_eps 
)

Gated DeltaNet autoregressive update.

Layout: q, k, v [num_heads, state_dim] g, beta [num_heads] state_* [num_heads, state_dim, state_dim] row-major per head out [num_heads, state_dim]

This mirrors the single-token recurrent update used by qwen3next in llama.cpp after projections/convolution but before output projection.

Definition at line 741 of file ck_parity_api.c.

752{
754 q,
755 k,
756 v,
757 g,
758 beta,
759 state_in,
760 state_out,
761 out,
762 num_heads,
763 state_dim,
764 norm_eps);
765}
void gated_deltanet_autoregressive_forward(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)

References gated_deltanet_autoregressive_forward().

◆ ck_test_gemm_q4_k()

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.

Computes: output[t,r] = sum_k(weight[r,k] * input[t,k])

Parameters
weight_q4kQ4_K quantized weights [rows, cols]
input_f32FP32 input [n_tokens, cols]
outputFP32 output [n_tokens, rows]
rowsNumber of output rows
colsNumber of columns (must be multiple of 256)
n_tokensBatch size

Definition at line 434 of file ck_parity_api.c.

438{
439 /* Allocate Q8_K buffer for quantized activations */
440 int n_blocks_per_row = cols / CK_QK_K;
441 block_q8_K *q8_data = (block_q8_K *)malloc(n_tokens * n_blocks_per_row * sizeof(block_q8_K));
442 if (!q8_data) {
443 memset(output, 0, n_tokens * rows * sizeof(float));
444 return;
445 }
446
447 /* Quantize all input tokens */
448 for (int t = 0; t < n_tokens; t++) {
449 quantize_row_q8_k(input_f32 + t * cols,
450 q8_data + t * n_blocks_per_row, cols);
451 }
452
453 /* Use gemm_nt_q4_k_q8_k: C[M,N] = A[M,K] * B[N,K]^T
454 * Our layout: output[n_tokens, rows] = input[n_tokens, cols] * weight[rows, cols]^T
455 * So: M = n_tokens, N = rows, K = cols
456 */
457 gemm_nt_q4_k_q8_k(q8_data, weight_q4k, NULL, output, n_tokens, rows, cols);
458
459 free(q8_data);
460}
void quantize_row_q8_k(const float *x, void *vy, 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)
#define CK_QK_K

References CK_QK_K, gemm_nt_q4_k_q8_k(), and quantize_row_q8_k().

◆ ck_test_gemm_q5_0()

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)

Q5_0 GEMM - batched matrix multiply with Q5_0 weights (32-element blocks)

Used for MLP W1 (gate/up projection) and attention Q/K with Q5_0 weights.

Definition at line 533 of file ck_parity_api.c.

537{
538 /* Allocate Q8_0 buffer for quantized activations */
539 int n_blocks_per_row = cols / CK_QK8_0;
540 block_q8_0 *q8_data = (block_q8_0 *)malloc(n_tokens * n_blocks_per_row * sizeof(block_q8_0));
541 if (!q8_data) {
542 memset(output, 0, n_tokens * rows * sizeof(float));
543 return;
544 }
545
546 /* Quantize all input tokens */
547 for (int t = 0; t < n_tokens; t++) {
548 quantize_row_q8_0(input_f32 + t * cols,
549 q8_data + t * n_blocks_per_row, cols);
550 }
551
552 /* Use gemm_nt_q5_0_q8_0: C[M,N] = A[M,K] * B[N,K]^T
553 * Our layout: output[n_tokens, rows] = input[n_tokens, cols] * weight[rows, cols]^T
554 * So: M = n_tokens, N = rows, K = cols
555 */
556 gemm_nt_q5_0_q8_0(q8_data, weight_q5_0, NULL, output, n_tokens, rows, cols);
557
558 free(q8_data);
559}
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 quantize_row_q8_0(const float *x, void *vy, int k)
Quantize FP32 to Q8_0 format (scalar reference)
#define CK_QK8_0

References CK_QK8_0, gemm_nt_q5_0_q8_0(), and quantize_row_q8_0().

◆ ck_test_gemm_q6_k()

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)

Q6_K GEMM - batched matrix multiply with Q6_K weights.

Used for MLP W2 (down projection) with Q6_K weights.

Definition at line 467 of file ck_parity_api.c.

471{
472 /* Allocate Q8_K buffer for quantized activations */
473 int n_blocks_per_row = cols / CK_QK_K;
474 block_q8_K *q8_data = (block_q8_K *)malloc(n_tokens * n_blocks_per_row * sizeof(block_q8_K));
475 if (!q8_data) {
476 memset(output, 0, n_tokens * rows * sizeof(float));
477 return;
478 }
479
480 /* Quantize all input tokens */
481 for (int t = 0; t < n_tokens; t++) {
482 quantize_row_q8_k(input_f32 + t * cols,
483 q8_data + t * n_blocks_per_row, cols);
484 }
485
486 /* Use gemm_nt_q6_k_q8_k: C[M,N] = A[M,K] * B[N,K]^T
487 * Our layout: output[n_tokens, rows] = input[n_tokens, cols] * weight[rows, cols]^T
488 * So: M = n_tokens, N = rows, K = cols
489 */
490 gemm_nt_q6_k_q8_k(q8_data, weight_q6k, NULL, output, n_tokens, rows, cols);
491
492 free(q8_data);
493}
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)

References CK_QK_K, gemm_nt_q6_k_q8_k(), and quantize_row_q8_k().

◆ ck_test_gemm_q8_0()

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)

Q8_0 GEMM - batched matrix multiply with Q8_0 weights (32-element blocks)

Used for attention V projection with Q8_0 weights.

Definition at line 500 of file ck_parity_api.c.

504{
505 /* Allocate Q8_0 buffer for quantized activations */
506 int n_blocks_per_row = cols / CK_QK8_0;
507 block_q8_0 *q8_data = (block_q8_0 *)malloc(n_tokens * n_blocks_per_row * sizeof(block_q8_0));
508 if (!q8_data) {
509 memset(output, 0, n_tokens * rows * sizeof(float));
510 return;
511 }
512
513 /* Quantize all input tokens */
514 for (int t = 0; t < n_tokens; t++) {
515 quantize_row_q8_0(input_f32 + t * cols,
516 q8_data + t * n_blocks_per_row, cols);
517 }
518
519 /* Use gemm_nt_q8_0_q8_0: C[M,N] = A[M,K] * B[N,K]^T
520 * Our layout: output[n_tokens, rows] = input[n_tokens, cols] * weight[rows, cols]^T
521 * So: M = n_tokens, N = rows, K = cols
522 */
523 gemm_nt_q8_0_q8_0(q8_data, weight_q8_0, NULL, output, n_tokens, rows, cols);
524
525 free(q8_data);
526}
void gemm_nt_q8_0_q8_0(const void *A_q8, const void *B_q8, const float *bias, float *C, int M, int N, int K)
gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

References CK_QK8_0, gemm_nt_q8_0_q8_0(), and quantize_row_q8_0().

◆ ck_test_gemv_q4_k()

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.

Internally quantizes input to Q8_K, then computes dot product.

Parameters
weight_q4kQ4_K quantized weights [cols]
input_f32FP32 input vector [cols]
outputOutput scalar [1]
colsNumber of columns (must be multiple of 256)

Definition at line 220 of file ck_parity_api.c.

224{
225 /* Allocate Q8_K buffer for quantized activations */
226 int n_blocks = cols / CK_QK_K;
227 block_q8_K *q8_data = (block_q8_K *)malloc(n_blocks * sizeof(block_q8_K));
228 if (!q8_data) {
229 *output = 0.0f;
230 return;
231 }
232
233 /* Quantize input to Q8_K */
234 quantize_row_q8_k(input_f32, q8_data, cols);
235
236 /* Compute dot product using GEMV with M=1 */
237 gemv_q4_k_q8_k(output, weight_q4k, q8_data, 1, cols);
238
239 free(q8_data);
240}
void gemv_q4_k_q8_k(float *y, const void *W, const void *x_q8, int M, int K)

References CK_QK_K, gemv_q4_k_q8_k(), and quantize_row_q8_k().

◆ ck_test_gemv_q5_0()

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.

Parameters
weight_q5_0Q5_0 quantized weights [rows * cols]
input_f32FP32 input vector [cols]
outputFP32 output vector [rows]
rowsNumber of output rows
colsNumber of columns (must be multiple of 32)

Definition at line 263 of file ck_parity_api.c.

267{
268 /* Match llama.cpp's test_gemv_q5_0:
269 * 1. Quantize input to Q8_0 format
270 * 2. Use quantized dot product (vec_dot_q5_0_q8_0)
271 *
272 * This ensures parity with llama.cpp which always uses the
273 * quantized path, NOT the FP32 dequantization path.
274 */
275 int n_blocks = cols / CK_QK8_0;
276 block_q8_0 *q8_data = (block_q8_0 *)malloc(n_blocks * sizeof(block_q8_0));
277 if (!q8_data) {
278 for (int r = 0; r < rows; r++) output[r] = 0.0f;
279 return;
280 }
281
282 /* Quantize input to Q8_0 */
283 quantize_row_q8_0(input_f32, q8_data, cols);
284
285 /* Call the quantized GEMV kernel (same as ck_test_gemv_q5_0_q8_0) */
286 gemv_q5_0_q8_0(output, weight_q5_0, q8_data, rows, cols);
287
288 free(q8_data);
289}
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.

References CK_QK8_0, gemv_q5_0_q8_0(), and quantize_row_q8_0().

◆ ck_test_gemv_q5_0_q8_0()

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.

This version quantizes the input to Q8_0 first, then uses integer dot products (like llama.cpp does). Use this for parity testing.

Parameters
weight_q5_0Q5_0 quantized weights [rows * cols]
input_f32FP32 input vector [cols] - will be quantized to Q8_0
outputFP32 output vector [rows]
rowsNumber of output rows
colsNumber of columns (must be multiple of 32)

Definition at line 319 of file ck_parity_api.c.

323{
324 /* This matches llama.cpp's approach:
325 * 1. Quantize input to Q8_0 format
326 * 2. Use quantized dot product (integer math)
327 * 3. Scale at the end
328 */
329 int n_blocks = cols / CK_QK8_0;
330 block_q8_0 *q8_data = (block_q8_0 *)malloc(n_blocks * sizeof(block_q8_0));
331 if (!q8_data) {
332 for (int r = 0; r < rows; r++) output[r] = 0.0f;
333 return;
334 }
335
336 /* Quantize input to Q8_0 */
337 quantize_row_q8_0(input_f32, q8_data, cols);
338
339 /* Call the quantized GEMV kernel */
340 gemv_q5_0_q8_0(output, weight_q5_0, q8_data, rows, cols);
341
342 free(q8_data);
343}

References CK_QK8_0, gemv_q5_0_q8_0(), and quantize_row_q8_0().

◆ ck_test_gemv_q6_k()

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

Q6_K GEMV.

Definition at line 242 of file ck_parity_api.c.

246{
247 /* Match runtime decode path:
248 * 1) quantize FP32 activation to Q8_K
249 * 2) run Q6_K x Q8_K GEMV kernel with M=1
250 */
251 int n_blocks = cols / CK_QK_K;
252 block_q8_K *q8_data = (block_q8_K *)malloc(n_blocks * sizeof(block_q8_K));
253 if (!q8_data) {
254 *output = 0.0f;
255 return;
256 }
257
258 quantize_row_q8_k(input_f32, q8_data, cols);
259 gemv_q6_k_q8_k(output, weight_q6k, q8_data, 1, cols);
260 free(q8_data);
261}
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.

References CK_QK_K, gemv_q6_k_q8_k(), and quantize_row_q8_k().

◆ ck_test_gemv_q8_0()

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.

Parameters
weight_q8_0Q8_0 quantized weights [rows * cols]
input_f32FP32 input vector [cols]
outputFP32 output vector [rows]
rowsNumber of output rows
colsNumber of columns (must be multiple of 32)

Definition at line 291 of file ck_parity_api.c.

295{
296 /* Match llama.cpp's test_gemv_q8_0:
297 * 1. Quantize input to Q8_0 format
298 * 2. Use quantized dot product (vec_dot_q8_0_q8_0)
299 *
300 * This ensures parity with llama.cpp which always uses the
301 * quantized path, NOT the FP32 dequantization path.
302 */
303 int n_blocks = cols / CK_QK8_0;
304 block_q8_0 *q8_data = (block_q8_0 *)malloc(n_blocks * sizeof(block_q8_0));
305 if (!q8_data) {
306 for (int r = 0; r < rows; r++) output[r] = 0.0f;
307 return;
308 }
309
310 /* Quantize input to Q8_0 */
311 quantize_row_q8_0(input_f32, q8_data, cols);
312
313 /* Call the quantized GEMV kernel (same as ck_test_gemv_q8_0_q8_0) */
314 gemv_q8_0_q8_0(output, weight_q8_0, q8_data, rows, cols);
315
316 free(q8_data);
317}
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.

References CK_QK8_0, gemv_q8_0_q8_0(), and quantize_row_q8_0().

◆ ck_test_gemv_q8_0_q8_0()

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.

This version quantizes the input to Q8_0 first, then uses integer dot products (like llama.cpp does). Use this for parity testing.

Parameters
weight_q8_0Q8_0 quantized weights [rows * cols]
input_f32FP32 input vector [cols] - will be quantized to Q8_0
outputFP32 output vector [rows]
rowsNumber of output rows
colsNumber of columns (must be multiple of 32)

Definition at line 345 of file ck_parity_api.c.

349{
350 /* This matches llama.cpp's approach:
351 * 1. Quantize input to Q8_0 format
352 * 2. Use quantized dot product (integer math)
353 * 3. Scale at the end
354 */
355 int n_blocks = cols / CK_QK8_0;
356 block_q8_0 *q8_data = (block_q8_0 *)malloc(n_blocks * sizeof(block_q8_0));
357 if (!q8_data) {
358 for (int r = 0; r < rows; r++) output[r] = 0.0f;
359 return;
360 }
361
362 /* Quantize input to Q8_0 */
363 quantize_row_q8_0(input_f32, q8_data, cols);
364
365 /* Call the quantized GEMV kernel */
366 gemv_q8_0_q8_0(output, weight_q8_0, q8_data, rows, cols);
367
368 free(q8_data);
369}

References CK_QK8_0, gemv_q8_0_q8_0(), and quantize_row_q8_0().

◆ ck_test_outproj_mlp_fused_q5_0()

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)

This is a simplified wrapper for parity testing that:

  • Uses Q5_0 for W_o and W1 weights
  • Uses Q4_K for W2 weights
  • Allocates scratch internally
Parameters
attn_outAttention output [num_heads, tokens, head_dim] (FP32, head-major)
residualResidual input [tokens, embed_dim] (FP32)
ln2_gammaRMSNorm gamma [embed_dim] (FP32)
woOutProj weights [embed_dim, embed_dim] (Q5_0)
w1MLP W1 weights [2*intermediate, embed_dim] (Q5_0)
w2MLP W2 weights [embed_dim, intermediate] (Q4_K or Q6_K)
outputOutput [tokens, embed_dim] (FP32)
tokensNumber of tokens
num_headsNumber of attention heads
head_dimDimension per head
embed_dimEmbedding dimension (= num_heads * head_dim)
intermediateMLP intermediate dimension
epsRMSNorm epsilon
w2_is_q6kIf true, W2 is Q6_K; if false, W2 is Q4_K

Definition at line 959 of file ck_parity_api.c.

974{
975 /* CK uses dtype enum: CK_DT_Q5_0 = 11, CK_DT_Q4_K = 7, CK_DT_Q6_K = 8 */
976 const int CK_DT_Q5_0_VAL = 11;
977 const int CK_DT_Q4_K_VAL = 7;
978 const int CK_DT_Q6_K_VAL = 8;
979
980 /* For parity testing, aligned = actual (no padding) */
981 int aligned_embed_dim = embed_dim;
982 int aligned_head_dim = head_dim;
983 int aligned_intermediate = intermediate;
984
985 /* Ensure intermediate is multiple of 256 (QK_K) for K-quants */
986 if ((intermediate % 256) != 0) {
987 aligned_intermediate = ((intermediate + 255) / 256) * 256;
988 }
989
990 /* Allocate scratch */
992 tokens, aligned_embed_dim, num_heads, aligned_head_dim, aligned_intermediate);
993
994 void *scratch = malloc(scratch_size);
995 if (!scratch) {
996 return;
997 }
998
999 /* Call the mega-fused kernel */
1001 output,
1002 attn_out,
1003 residual,
1004 ln2_gamma,
1005 wo, NULL, CK_DT_Q5_0_VAL, /* W_o with Q5_0 */
1006 w1, NULL, CK_DT_Q5_0_VAL, /* W1 with Q5_0 */
1007 w2, NULL, w2_is_q6k ? CK_DT_Q6_K_VAL : CK_DT_Q4_K_VAL, /* W2 with Q4_K or Q6_K */
1008 tokens,
1009 embed_dim,
1010 aligned_embed_dim,
1011 num_heads,
1012 aligned_head_dim,
1013 intermediate,
1014 aligned_intermediate,
1015 eps,
1016 scratch
1017 );
1018
1019 free(scratch);
1020}
void mega_fused_outproj_mlp_prefill(float *output, const float *attn_out, const float *residual, const float *ln2_gamma, const void *wo, const float *bo, int wo_dt, const void *w1, const float *b1, int w1_dt, const void *w2, const float *b2, int w2_dt, int tokens, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim, int intermediate_dim, int aligned_intermediate_dim, float eps, void *scratch)
size_t mega_fused_outproj_mlp_prefill_scratch_size(int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim, int aligned_intermediate_dim)
Get scratch buffer size for mega_fused_outproj_mlp_prefill.

References mega_fused_outproj_mlp_prefill(), and mega_fused_outproj_mlp_prefill_scratch_size().

◆ ck_test_quantize_q8_k()

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

Quantize FP32 to Q8_K (for activations)

Parameters
srcInput FP32 values
dstOutput Q8_K blocks
nNumber of elements (must be multiple of 256)

Definition at line 211 of file ck_parity_api.c.

212{
213 quantize_row_q8_k(src, dst, n);
214}

References quantize_row_q8_k().

◆ ck_test_recurrent_conv_state_update()

void ck_test_recurrent_conv_state_update ( const float *  state_in,
const float *  q,
const float *  k,
const float *  v,
float *  conv_x,
float *  state_out,
int  history_len,
int  num_seqs,
int  num_tokens,
int  q_dim,
int  k_dim,
int  v_dim 
)

Build the recurrent convolution input history window.

Layout: state_in [num_seqs, channels, history_len] q [num_seqs * num_tokens, q_dim] k [num_seqs * num_tokens, k_dim] v [num_seqs * num_tokens, v_dim] conv_x [num_seqs, channels, history_len + num_tokens] state_out [num_seqs, channels, history_len]

Definition at line 811 of file ck_parity_api.c.

823{
825 state_in, q, k, v, conv_x, state_out, history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim);
826}
void recurrent_conv_state_update_forward(const float *state_in, const float *q, const float *k, const float *v, float *conv_x, float *state_out, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)

References recurrent_conv_state_update_forward().

◆ ck_test_recurrent_dt_gate()

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

Transform recurrent alpha rows into the DeltaNet gate.

Layout: alpha : [rows, dim] dt_bias : [dim] a : [dim] gate : [rows, dim]

Definition at line 801 of file ck_parity_api.c.

807{
808 recurrent_dt_gate_forward(alpha, dt_bias, a, gate, rows, 1, dim);
809}
void recurrent_dt_gate_forward(const float *alpha, const float *dt_bias, const float *a, float *gate, int rows, int num_heads, int state_dim)

References recurrent_dt_gate_forward().

◆ ck_test_recurrent_norm_gate()

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

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

Definition at line 868 of file ck_parity_api.c.

876{
877 recurrent_norm_gate_forward(x, gate, weight, out, rows, num_heads, head_dim, eps);
878}
void recurrent_norm_gate_forward(const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)

References recurrent_norm_gate_forward().

◆ ck_test_recurrent_qk_l2_norm()

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

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

Layout: q : [rows, q_dim] k : [rows, k_dim] where q_dim and k_dim are multiples of head_dim.

Definition at line 848 of file ck_parity_api.c.

855{
856 recurrent_qk_l2_norm_forward(q, k, rows, q_dim, k_dim, head_dim, eps);
857}
void recurrent_qk_l2_norm_forward(float *q, float *k, int rows, int q_dim, int k_dim, int head_dim, float eps)

References recurrent_qk_l2_norm_forward().

◆ ck_test_recurrent_silu()

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

Apply SiLU elementwise to recurrent rows.

Definition at line 828 of file ck_parity_api.c.

832{
833 recurrent_silu_forward(x, out, rows, dim);
834}
void recurrent_silu_forward(const float *x, float *out, int rows, int dim)

References recurrent_silu_forward().

◆ ck_test_recurrent_split_conv_qkv()

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

Split the post-convolution recurrent packed QKV rows.

Definition at line 836 of file ck_parity_api.c.

844{
845 recurrent_split_conv_qkv_forward(packed_qkv, q, k, v, rows, q_dim, k_dim, v_dim);
846}
void recurrent_split_conv_qkv_forward(const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)

References recurrent_split_conv_qkv_forward().

◆ ck_test_recurrent_split_qkv()

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

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

Layout: packed_qkv : [rows, q_dim + k_dim + v_dim] q : [rows, q_dim] k : [rows, k_dim] v : [rows, v_dim]

Definition at line 789 of file ck_parity_api.c.

797{
798 recurrent_split_qkv_forward(packed_qkv, q, k, v, rows, q_dim, k_dim, v_dim);
799}
void recurrent_split_qkv_forward(const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)

References recurrent_split_qkv_forward().

◆ ck_test_rmsnorm()

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

RMSNorm.

Computes: output = (input / rms(input)) * weight where rms(x) = sqrt(mean(x^2) + eps)

Parameters
inputInput tensor [n_tokens, dim]
weightNormalization weights [dim]
outputOutput tensor [n_tokens, dim]
n_tokensNumber of tokens
dimHidden dimension
epsEpsilon for numerical stability

Definition at line 565 of file ck_parity_api.c.

569{
570 /* CK rmsnorm_forward has aligned_embed_dim parameter
571 * For testing, use dim as aligned_embed_dim (no padding) */
572 rmsnorm_forward(input, weight, output, NULL, n_tokens, dim, dim, eps);
573}
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)

References rmsnorm_forward().

◆ ck_test_rope()

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)

Applies rotary position embeddings to Q and K tensors.

NOTE: CK uses rotate-half format (split first/second halves) while some implementations use interleaved format. The test harness should account for this.

Parameters
qQuery tensor [n_tokens, n_heads * head_dim], modified in-place
kKey tensor [n_tokens, n_heads_kv * head_dim], modified in-place
n_tokensNumber of tokens
n_headsNumber of query heads
n_heads_kvNumber of key/value heads
head_dimDimension per head
pos_offsetStarting position for RoPE
thetaRoPE base frequency (typically 10000.0)

Definition at line 575 of file ck_parity_api.c.

578{
579 /* Precompute cos/sin cache */
580 int half_dim = head_dim / 2;
581 int max_seq = pos_offset + n_tokens;
582
583 float *cos_cache = (float *)malloc(max_seq * half_dim * sizeof(float));
584 float *sin_cache = (float *)malloc(max_seq * half_dim * sizeof(float));
585 if (!cos_cache || !sin_cache) {
586 free(cos_cache);
587 free(sin_cache);
588 return;
589 }
590
591 rope_precompute_cache(cos_cache, sin_cache, max_seq, head_dim, theta,
592 head_dim, "none", 1.0f);
593
594 /* CK RoPE expects layout [num_heads, num_tokens, head_dim]
595 * Reshape from [n_tokens, n_heads * head_dim] to [n_heads, n_tokens, head_dim]
596 */
597 float *q_reorder = (float *)malloc(n_heads * n_tokens * head_dim * sizeof(float));
598 float *k_reorder = (float *)malloc(n_heads_kv * n_tokens * head_dim * sizeof(float));
599
600 if (q_reorder && k_reorder) {
601 /* Reorder Q: [T, H*D] -> [H, T, D] */
602 for (int t = 0; t < n_tokens; t++) {
603 for (int h = 0; h < n_heads; h++) {
604 for (int d = 0; d < head_dim; d++) {
605 q_reorder[h * n_tokens * head_dim + t * head_dim + d] =
606 q[t * n_heads * head_dim + h * head_dim + d];
607 }
608 }
609 }
610
611 /* Reorder K: [T, H_kv*D] -> [H_kv, T, D] */
612 for (int t = 0; t < n_tokens; t++) {
613 for (int h = 0; h < n_heads_kv; h++) {
614 for (int d = 0; d < head_dim; d++) {
615 k_reorder[h * n_tokens * head_dim + t * head_dim + d] =
616 k[t * n_heads_kv * head_dim + h * head_dim + d];
617 }
618 }
619 }
620
621 /* Apply RoPE */
622 rope_forward_qk(q_reorder, k_reorder,
623 cos_cache, sin_cache,
624 n_heads, n_heads_kv, n_tokens,
625 head_dim, head_dim, pos_offset);
626
627 /* Reorder back: [H, T, D] -> [T, H*D] */
628 for (int t = 0; t < n_tokens; t++) {
629 for (int h = 0; h < n_heads; h++) {
630 for (int d = 0; d < head_dim; d++) {
631 q[t * n_heads * head_dim + h * head_dim + d] =
632 q_reorder[h * n_tokens * head_dim + t * head_dim + d];
633 }
634 }
635 }
636
637 for (int t = 0; t < n_tokens; t++) {
638 for (int h = 0; h < n_heads_kv; h++) {
639 for (int d = 0; d < head_dim; d++) {
640 k[t * n_heads_kv * head_dim + h * head_dim + d] =
641 k_reorder[h * n_tokens * head_dim + t * head_dim + d];
642 }
643 }
644 }
645 }
646
647 free(q_reorder);
648 free(k_reorder);
649 free(cos_cache);
650 free(sin_cache);
651}
void rope_precompute_cache(float *cos_cache, float *sin_cache, int max_seq_len, int head_dim, float base, int rotary_dim, const char *scaling_type, float scaling_factor)
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)

References rope_forward_qk(), and rope_precompute_cache().

◆ ck_test_rope_interleaved()

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)

Uses interleaved format: (x0, x1) -> (x0*cos - x1*sin, x0*sin + x1*cos)

Definition at line 653 of file ck_parity_api.c.

656{
657 /* Interleaved RoPE format (matches llama.cpp):
658 * (x0, x1) -> (x0*cos - x1*sin, x0*sin + x1*cos)
659 * Applied to consecutive pairs of elements
660 */
661
662 /* Precompute inverse frequencies */
663 float *inv_freq = (float *)malloc((head_dim / 2) * sizeof(float));
664 if (!inv_freq) return;
665
666 for (int i = 0; i < head_dim / 2; i++) {
667 inv_freq[i] = 1.0f / powf(theta, (float)(2 * i) / head_dim);
668 }
669
670 /* Apply RoPE to Q */
671 for (int t = 0; t < n_tokens; t++) {
672 int pos = pos_offset + t;
673 for (int h = 0; h < n_heads; h++) {
674 float *qh = q + t * n_heads * head_dim + h * head_dim;
675
676 for (int i = 0; i < head_dim / 2; i++) {
677 float freq = pos * inv_freq[i];
678 float cos_val = cosf(freq);
679 float sin_val = sinf(freq);
680
681 /* Interleaved format */
682 float x0 = qh[i * 2];
683 float x1 = qh[i * 2 + 1];
684 qh[i * 2] = x0 * cos_val - x1 * sin_val;
685 qh[i * 2 + 1] = x0 * sin_val + x1 * cos_val;
686 }
687 }
688 }
689
690 /* Apply RoPE to K */
691 for (int t = 0; t < n_tokens; t++) {
692 int pos = pos_offset + t;
693 for (int h = 0; h < n_heads_kv; h++) {
694 float *kh = k + t * n_heads_kv * head_dim + h * head_dim;
695
696 for (int i = 0; i < head_dim / 2; i++) {
697 float freq = pos * inv_freq[i];
698 float cos_val = cosf(freq);
699 float sin_val = sinf(freq);
700
701 float x0 = kh[i * 2];
702 float x1 = kh[i * 2 + 1];
703 kh[i * 2] = x0 * cos_val - x1 * sin_val;
704 kh[i * 2 + 1] = x0 * sin_val + x1 * cos_val;
705 }
706 }
707 }
708
709 free(inv_freq);
710}

◆ ck_test_softmax()

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

Softmax (simple, non-causal)

Computes: output[i] = exp(input[i]) / sum(exp(input))

Parameters
inputInput tensor [n]
outputOutput tensor [n]
nNumber of elements

Definition at line 719 of file ck_parity_api.c.

720{
721 /* Find max for numerical stability */
722 float max_val = input[0];
723 for (int i = 1; i < n; i++) {
724 if (input[i] > max_val) max_val = input[i];
725 }
726
727 /* Compute exp and sum */
728 float sum = 0.0f;
729 for (int i = 0; i < n; i++) {
730 output[i] = expf(input[i] - max_val);
731 sum += output[i];
732 }
733
734 /* Normalize */
735 float inv_sum = 1.0f / sum;
736 for (int i = 0; i < n; i++) {
737 output[i] *= inv_sum;
738 }
739}

◆ ck_test_split_q_gate()

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

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

Layout: packed_qg : [rows, q_dim + gate_dim] q : [rows, q_dim] gate : [rows, gate_dim]

Definition at line 778 of file ck_parity_api.c.

785{
786 split_q_gate_forward(packed_qg, q, gate, rows, q_dim, gate_dim, group_dim);
787}
void split_q_gate_forward(const float *packed_qg, float *q, float *gate, int rows, int q_dim, int gate_dim, int group_dim)

References split_q_gate_forward().

◆ ck_test_ssm_conv1d()

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

qwen3next/Qwen3.5 SSM causal depthwise convolution.

Layout: conv_x [num_seqs, num_channels, kernel_size - 1 + num_tokens] kernel [num_channels, kernel_size] out [num_seqs, num_tokens, num_channels]

This mirrors ggml's GGML_OP_SSM_CONV used immediately before the DeltaNet recurrent update in qwen3next/Qwen3.5.

Definition at line 767 of file ck_parity_api.c.

774{
775 ssm_conv1d_forward(conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs);
776}
void ssm_conv1d_forward(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

References ssm_conv1d_forward().

◆ ck_test_swiglu()

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

SwiGLU activation.

Computes: output = SiLU(gate) * up where SiLU(x) = x * sigmoid(x)

Parameters
gate_upInput tensor [n_tokens, 2 * intermediate_dim] Layout: [gate_0..gate_D-1, up_0..up_D-1] per token
outputOutput tensor [n_tokens, intermediate_dim]
n_tokensNumber of tokens
intermediate_dimIntermediate dimension

Definition at line 712 of file ck_parity_api.c.

715{
716 swiglu_forward(gate_up, output, n_tokens, intermediate_dim);
717}
void swiglu_forward(const float *input, float *output, int tokens, int dim)

References swiglu_forward().

◆ ck_test_vec_dot_q4_k_q8_k()

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

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

Definition at line 375 of file ck_parity_api.c.

379{
380 /* M=1 reaches the production Q4 dispatch without requantizing input. */
381 gemv_q4_k_q8_k(output, weight_q4_k, input_q8_k, 1, cols);
382}

References gemv_q4_k_q8_k().

◆ ck_test_vec_dot_q5_0_q8_0()

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)

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

This is a "direct" test that bypasses FP32-to-Q8_0 conversion. Useful for isolating kernel bugs from quantization bugs.

Parameters
weight_q5_0Q5_0 quantized weights [cols]
input_q8_0Q8_0 quantized input [cols] (pre-quantized!)
outputOutput scalar [1]
colsNumber of elements (must be multiple of 32)

Definition at line 406 of file ck_parity_api.c.

410{
411 vec_dot_q5_0_q8_0(cols, output, weight_q5_0, input_q8_0);
412}
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.

References vec_dot_q5_0_q8_0().

◆ ck_test_vec_dot_q6_k_q8_k()

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

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

Definition at line 384 of file ck_parity_api.c.

388{
389 /* Exercise the production M=1 provider. The scalar vec_dot helper is an
390 * internal architecture-neutral oracle and does not preserve the x86
391 * provider's declared lane reduction order. */
392 gemv_q6_k_q8_k(output, weight_q6_k, input_q8_k, 1, cols);
393}

References gemv_q6_k_q8_k().

◆ ck_test_vec_dot_q8_0_q8_0()

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)

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

Parameters
weight_q8_0Q8_0 quantized weights [cols]
input_q8_0Q8_0 quantized input [cols] (pre-quantized!)
outputOutput scalar [1]
colsNumber of elements (must be multiple of 32)

Definition at line 422 of file ck_parity_api.c.

426{
427 vec_dot_q8_0_q8_0(cols, output, weight_q8_0, input_q8_0);
428}
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.

References vec_dot_q8_0_q8_0().

◆ dequant_q4_0_row()

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

Dequantize Q4_0 row (multiple blocks)

Parameters
srcQ4_0 data
dstFP32 output
n_elementsNumber of elements to dequantize

Definition at line 63 of file dequant_kernels.c.

64{
65 const block_q4_0 *blocks = (const block_q4_0 *)src;
66 const size_t n_blocks = n_elements / QK4_0;
67
68 for (size_t b = 0; b < n_blocks; b++) {
69 dequant_q4_0_block(&blocks[b], &dst[b * QK4_0]);
70 }
71}
#define QK4_0
void dequant_q4_0_block(const block_q4_0 *block, float *output)
Dequantize a single Q4_0 block to FP32.

Referenced by ck_test_dequant_q4_0(), and dequant_row().

◆ dequant_q4_k_row()

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

Dequantize Q4_K row (multiple blocks)

Definition at line 372 of file dequant_kernels.c.

373{
374 const block_q4_K *blocks = (const block_q4_K *)src;
375 const size_t n_blocks = n_elements / QK_K;
376
377 for (size_t b = 0; b < n_blocks; b++) {
378 dequant_q4_k_block(&blocks[b], &dst[b * QK_K]);
379 }
380}
void dequant_q4_k_block(const block_q4_K *block, float *output)
Dequantize a single Q4_K block to FP32.

Referenced by ck_test_dequant_q4_k(), and dequant_row().

◆ dequant_q6_k_row()

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

Dequantize Q6_K row (multiple blocks)

Definition at line 422 of file dequant_kernels.c.

423{
424 const block_q6_K *blocks = (const block_q6_K *)src;
425 const size_t n_blocks = n_elements / QK_K;
426
427 for (size_t b = 0; b < n_blocks; b++) {
428 dequant_q6_k_block(&blocks[b], &dst[b * QK_K]);
429 }
430}
void dequant_q6_k_block(const block_q6_K *block, float *output)
Dequantize a single Q6_K block to FP32.

Referenced by ck_test_dequant_q6_k(), and dequant_row().

◆ gated_deltanet_autoregressive_forward()

void gated_deltanet_autoregressive_forward ( const float *  q,
const float *  k,
const float *  v,
const float *  g,
const float *  beta,
const float *  state_in,
float *  state_out,
float *  out,
int  num_heads,
int  state_dim,
float  norm_eps 
)
extern

Definition at line 1956 of file deltanet_kernels.c.

1967{
1968 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
1969 return;
1970 }
1971 if (num_heads <= 0 || state_dim <= 0) {
1972 return;
1973 }
1974
1975 /*
1976 * q and k arrive pre-normalized by recurrent_qk_l2_norm, so the
1977 * ISA-specialized kernels can follow the same contract as the scalar ref.
1978 */
1981 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1982 return;
1983 }
1984#if defined(__AVX512F__)
1985 gated_deltanet_autoregressive_forward_avx512(
1986 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1987#elif defined(__AVX2__)
1988 gated_deltanet_autoregressive_forward_avx2(
1989 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1990#elif defined(__AVX__)
1991 gated_deltanet_autoregressive_forward_avx(
1992 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1993#else
1995 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1996#endif
1997}
int ck_strict_parity_enabled(void)
void gated_deltanet_autoregressive_forward_ref(const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)
static int ck_deltanet_force_ref(void)

Referenced by ck_test_gated_deltanet_autoregressive(), and gated_deltanet_prefill_forward().

◆ gemm_nt_q4_k_q8_k()

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

Definition at line 648 of file gemm_kernels_q4k_q8k.c.

653{
654 if (!A_q8 || !B || !C) {
655 return;
656 }
657 if (M <= 0 || N <= 0 || K <= 0) {
658 return;
659 }
660
661 gemm_q4_k_q8_k(C, B, A_q8, /*M_out=*/N, /*N_batch=*/M, K);
662
663 if (!bias) {
664 return;
665 }
666
667 for (int i = 0; i < M; ++i) {
668 float *row = C + (size_t)i * (size_t)N;
669 for (int j = 0; j < N; ++j) {
670 row[j] += bias[j];
671 }
672 }
673}
void gemm_q4_k_q8_k(float *Y, const void *W, const void *X_q8, int M, int N, int K)
#define C(color)
Definition show_config.c:39

Referenced by ck_test_gemm_q4_k().

◆ gemm_nt_q5_0_q8_0()

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

Batch GEMM with Q5_0 weights and Q8_0 activations for prefill.

Computes C = A @ B^T + bias where: A: [M x K] Q8_0 quantized activations (M tokens, K features) B: [N x K] Q5_0 quantized weights (N outputs, K features) C: [M x N] FP32 output

This is the INT8 batch kernel for prefill, using pre-quantized activations to avoid FP32->Q8_0 conversion overhead per operation.

Parameters
A_q8Input activations in Q8_0 format [M rows of K/32 blocks each]
B_q5Weights in Q5_0 format [N rows of K/32 blocks each]
biasOptional bias vector [N], NULL if not used
COutput matrix [M x N], row-major FP32
MBatch size (number of tokens)
NOutput dimension (number of output features)
KInput dimension (must be multiple of 32)

Definition at line 1728 of file gemm_kernels_q5_0.c.

1736{
1737 const block_q5_0 *weights = (const block_q5_0 *)B_q5;
1738 const block_q8_0 *inputs = (const block_q8_0 *)A_q8;
1739 const int blocks_per_row = K / QK5_0;
1740
1741 for (int m = 0; m < M; m++) {
1742 const block_q8_0 *input_row = &inputs[m * blocks_per_row];
1743
1744 for (int n = 0; n < N; n++) {
1745 const block_q5_0 *weight_row = &weights[n * blocks_per_row];
1746 float *out = &C[m * N + n];
1747
1748 /* Dispatches to vec_dot_q5_0_q8_0_avx (2x block unrolled) on AVX */
1749 vec_dot_q5_0_q8_0(K, out, weight_row, input_row);
1750
1751 if (bias) {
1752 *out += bias[n];
1753 }
1754 }
1755 }
1756}
#define QK5_0
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.

Referenced by ck_test_gemm_q5_0(), and gemm_nt_q5_0_q8_0_m2n4_tile().

◆ gemm_nt_q6_k_q8_k()

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

Definition at line 1548 of file gemm_kernels_q6k_q8k.c.

1553{
1554 if (!A_q8 || !B || !C) {
1555 return;
1556 }
1557 if (M <= 0 || N <= 0 || K <= 0) {
1558 return;
1559 }
1560#if defined(__AVX2__) && defined(__FMA__)
1561 if (M >= 8 && ck_strict_parity_enabled()) {
1562 gemm_nt_q6_k_q8_k_tiled_parity(A_q8, B, bias, C, M, N, K);
1563 return;
1564 }
1565#endif
1566
1567 /* Prefill GEMM uses the pinned x86 leaf reduction for short row groups
1568 * and the tiled reduction for strict batches of at least eight rows.
1569 * CK_DEBUG_Q6K_Q8K_REF=1 still forces the scalar reference leaf. */
1570 const block_q8_K *A = (const block_q8_K *)A_q8;
1571 const block_q6_K *W = (const block_q6_K *)B;
1572 const int blocks_per_vec = K / QK_K;
1573 const int blocks_per_row = K / QK_K;
1574
1575 for (int m = 0; m < M; ++m) {
1576 const block_q8_K *a_row = A + (size_t)m * (size_t)blocks_per_vec;
1577 float *c_row = C + (size_t)m * (size_t)N;
1578 for (int n = 0; n < N; ++n) {
1579 const block_q6_K *w_row = W + (size_t)n * (size_t)blocks_per_row;
1580 const float b = bias ? bias[n] : 0.0f;
1581 c_row[n] = ck_dot_q6_k_q8_k_fast_or_ref(w_row, a_row, K) + b;
1582 }
1583 }
1584}
void gemm_nt_q6_k_q8_k_tiled_parity(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.
static float ck_dot_q6_k_q8_k_fast_or_ref(const block_q6_K *w, const block_q8_K *x, int K)

Referenced by ck_test_gemm_q6_k().

◆ gemm_nt_q8_0_q8_0()

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

gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

C[m,n] = A[m,K] @ B[n,K]^T + bias[n]

Definition at line 552 of file gemm_batch_int8.c.

558{
559 /* First compute GEMM */
560#if defined(__AVX2__)
561 /*
562 * The production contract is bit-exact with llama.cpp's eight-lane
563 * FP32 accumulation tree. AVX-512/VNNI changes the integer dot, but it
564 * must not silently replace that FP32 reduction with the scalar-per-block
565 * candidate above. The AVX2-named entry point delegates each activation
566 * row to the certified x4 provider, which also uses VNNI instructions when
567 * they are available while preserving the declared reduction order.
568 */
569 gemm_nt_q8_0_q8_0_avx2(A, B, C, M, N, K);
570#elif defined(__AVX__)
571 gemm_nt_q8_0_q8_0_avx(A, B, C, M, N, K);
572#else
573 gemm_nt_q8_0_q8_0_ref(A, B, C, M, N, K);
574#endif
575
576 /* Add bias if provided */
577 if (bias != NULL) {
578 for (int m = 0; m < M; m++) {
579 for (int n = 0; n < N; n++) {
580 C[(size_t)m * N + n] += bias[n];
581 }
582 }
583 }
584}
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.

Referenced by ck_test_gemm_q8_0().

◆ gemv_q4_k_q8_k()

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

Definition at line 497 of file gemm_kernels_q4k_q8k.c.

501{
502 if (ck_q4k_q8k_force_ref()) {
503 gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
504 return;
505 }
506#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && !defined(CK_NO_AVX512_VNNI)
507 /* VNNI: Best for decode (single token) - INT8 dot product acceleration */
508 gemv_q4_k_q8_k_vnni(y, W, x_q8, M, K);
509#elif defined(__AVX2__)
510 gemv_q4_k_q8_k_avx2(y, W, x_q8, M, K);
511#elif defined(__AVX__)
512 /* AVX version uses maddubs_epi16 (more efficient than SSE) */
513 gemv_q4_k_q8_k_avx(y, W, x_q8, M, K);
514#elif defined(__SSE4_1__)
515 gemv_q4_k_q8_k_sse(y, W, x_q8, M, K);
516#elif defined(__ARM_FEATURE_DOTPROD)
517 /* SDOT NEON is integer-exact against the scalar reference (identical
518 * block arithmetic and FP32 epilogue); CK_DEBUG_Q4K_Q8_REF=1 still
519 * forces the reference above. Baseline armv8-a builds (no dot-product
520 * extension) keep the reference path, which the compiler already
521 * autovectorizes well. */
522 gemv_q4_k_q8_k_neon(y, W, x_q8, M, K);
523#else
524 gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
525#endif
526}
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_vnni(float *y, const void *W, const void *x_q8, int M, int K)
void gemv_q4_k_q8_k_ref(float *y, const void *W, const void *x_q8, int M, int K)
static int ck_q4k_q8k_force_ref(void)
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_sse(float *y, const void *W, const void *x_q8, int M, int K)

Referenced by ck_test_gemv_q4_k(), and ck_test_vec_dot_q4_k_q8_k().

◆ gemv_q5_0()

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

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

Dispatch priority (best available):

  1. AVX-512 (512-bit vectors) - Intel Skylake-X+
  2. AVX2+FMA (256-bit vectors) - Intel Haswell+
  3. AVX (256-bit vectors) - Intel Sandy Bridge+
  4. SSE4.1 (128-bit vectors) - Intel Nehalem+
  5. Reference (scalar) - Fallback

Uses ck_features.h for standardized feature detection.

Parameters
yOutput vector [M]
WWeight matrix in Q5_0 format [M x K]
xInput vector [K]
MNumber of output rows
KNumber of input columns (hidden dimension)

Definition at line 550 of file gemm_kernels_q5_0.c.

554{
555// Dispatch order: AVX512 > AVX2 > AVX > SSE > ref
556#if defined(__AVX512F__)
557 gemv_q5_0_avx512(y, W, x, M, K);
558#elif defined(__AVX2__)
559 gemv_q5_0_avx2(y, W, x, M, K);
560#elif defined(__AVX__)
561 gemv_q5_0_avx(y, W, x, M, K);
562#elif defined(__SSE4_1__)
563 gemv_q5_0_ref(y, W, x, M, K);
564#else
565 gemv_q5_0_ref(y, W, x, M, K);
566#endif
567}
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)

Referenced by dot_q5_0(), gemm_nt_q5_0(), and gemm_q5_0().

◆ gemv_q5_0_q8_0()

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

Matrix-vector multiply with Q5_0 weights and Q8_0 input.

Parameters
yOutput vector [M]
WWeight matrix in Q5_0 format [M x K]
x_q8Input vector in Q8_0 format [K]
MNumber of output rows
KNumber of columns (must be multiple of 32)

Definition at line 1640 of file gemm_kernels_q5_0.c.

1644{
1645 const block_q5_0 *w_blocks = (const block_q5_0 *)W;
1646 const block_q8_0 *x_blocks = (const block_q8_0 *)x_q8;
1647 const int blocks_per_row = K / QK5_0;
1648
1649 for (int row = 0; row < M; row++) {
1650 vec_dot_q5_0_q8_0(K, &y[row],
1651 &w_blocks[row * blocks_per_row],
1652 x_blocks);
1653 }
1654}

Referenced by ck_test_gemv_q5_0(), and ck_test_gemv_q5_0_q8_0().

◆ gemv_q6_k_q8_k()

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

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

Definition at line 1356 of file gemm_kernels_q6k_q8k.c.

1360{
1362 gemv_q6_k_q8_k_ref(y, W, x_q8, M, K);
1363 return;
1364 }
1365
1366#if defined(__AVX2__)
1367 /* llama.cpp's x86 Q6_K production graph keeps the AVX2 reduction order
1368 * even when AVX-512 is available. Wider ISA availability is not a license
1369 * to change this numerical contract. */
1370 gemv_q6_k_q8_k_avx2(y, W, x_q8, M, K);
1371 return;
1372#elif defined(__AVX__)
1373 gemv_q6_k_q8_k_avx(y, W, x_q8, M, K);
1374 return;
1375#elif defined(__SSE4_1__)
1376 gemv_q6_k_q8_k_sse(y, W, x_q8, M, K);
1377 return;
1378#elif defined(__ARM_NEON) || defined(__aarch64__)
1379 /* NEON keeps the same integer-exact per-block dot as the scalar
1380 * reference but accumulates the FP32 block contributions sequentially
1381 * instead of through eight partial lanes, so strict parity mode (handled
1382 * above) still uses the reference reduction order. */
1383 gemv_q6_k_q8_k_neon(y, W, x_q8, M, K);
1384 return;
1385#endif
1386 gemv_q6_k_q8_k_ref(y, W, x_q8, M, K);
1387}
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_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)
static int ck_q6k_q8k_force_ref(void)

Referenced by ck_test_gemv_q6_k(), ck_test_vec_dot_q6_k_q8_k(), and gemm_q6_k_q8_k().

◆ gemv_q8_0()

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

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

Dispatch priority (best available):

  1. AVX-512 (512-bit vectors) - Intel Skylake-X+
  2. AVX2+FMA (256-bit vectors) - Intel Haswell+
  3. AVX (256-bit vectors) - Intel Sandy Bridge+
  4. SSE4.1 (128-bit vectors) - Intel Nehalem+
  5. Reference (scalar) - Fallback

Uses ck_features.h for standardized feature detection.

Parameters
yOutput vector [M]
WWeight matrix in Q8_0 format [M x K]
xInput vector [K]
MNumber of output rows
KNumber of input columns (hidden dimension)

Definition at line 694 of file gemm_kernels_q8_0.c.

698{
699 if (ck_q8_0_debug_ref()) {
700 gemv_q8_0_ref(y, W, x, M, K);
701 return;
702 }
703
704// Dispatch order: AVX512 > AVX2 > AVX > SSE > ref
705#if defined(__AVX512F__)
706 gemv_q8_0_avx512(y, W, x, M, K);
707#elif defined(__AVX2__)
708 gemv_q8_0_avx2(y, W, x, M, K);
709#elif defined(__AVX__)
710 gemv_q8_0_avx(y, W, x, M, K);
711#elif defined(__SSE4_1__)
712 gemv_q8_0_sse(y, W, x, M, K);
713#else
714 gemv_q8_0_ref(y, W, x, M, K);
715#endif
716}
static int ck_q8_0_debug_ref(void)
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)

Referenced by dot_q8_0(), gemm_nt_q8_0_rowloop(), and gemm_q8_0().

◆ gemv_q8_0_q8_0()

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

Matrix-vector multiply with Q8_0 weights and Q8_0 input.

Parameters
yOutput vector [M]
WWeight matrix in Q8_0 format [M x K]
x_q8Input vector in Q8_0 format [K]
MNumber of output rows
KNumber of columns (must be multiple of 32)

Definition at line 1405 of file gemm_kernels_q8_0.c.

1409{
1410 const block_q8_0 *w_blocks = (const block_q8_0 *)W;
1411 const block_q8_0 *x_blocks = (const block_q8_0 *)x_q8;
1412 const int blocks_per_row = K / QK8_0;
1413
1414 for (int row = 0; row < M; row++) {
1415 vec_dot_q8_0_q8_0(K, &y[row],
1416 &w_blocks[row * blocks_per_row],
1417 x_blocks);
1418 }
1419}
#define QK8_0
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.

Referenced by ck_test_gemv_q8_0(), ck_test_gemv_q8_0_q8_0(), gemm_q8_0_q8_0_m2n4_strided(), and gemv_q8_0_q8_0_x4().

◆ mega_fused_outproj_mlp_prefill()

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

◆ mega_fused_outproj_mlp_prefill_scratch_size()

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

Get scratch buffer size for mega_fused_outproj_mlp_prefill.

Definition at line 159 of file mega_fused_outproj_mlp_prefill.c.

164{
165 if (tokens <= 0 || aligned_embed_dim <= 0 || num_heads <= 0 ||
166 aligned_head_dim <= 0 || aligned_intermediate_dim <= 0) {
167 return 0;
168 }
169
170 const size_t q8_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0,
171 (size_t)aligned_head_dim);
172 const size_t attn_q8_bytes = (size_t)num_heads * (size_t)tokens * q8_row_bytes;
173 const size_t h1_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
174 const size_t ln2_bytes = h1_bytes;
175 const size_t mlp_scratch = fused_mlp_swiglu_prefill_w1w2_quant_scratch_size(
176 aligned_embed_dim, aligned_intermediate_dim);
177
178 return align_up_size(attn_q8_bytes, 64) +
179 align_up_size(h1_bytes, 64) +
180 align_up_size(ln2_bytes, 64) +
181 align_up_size(mlp_scratch, 64);
182}
@ CK_DT_Q8_0
static size_t ck_dtype_row_bytes(CKDataType dt, size_t n_elements)
Calculate total bytes for n_elements of given dtype.
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.
static size_t align_up_size(size_t value, size_t align)

Referenced by ck_test_outproj_mlp_fused_q5_0().

◆ quantize_row_q8_0()

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

Quantize FP32 to Q8_0 format (scalar reference)

Parameters
xInput FP32 values
vyOutput Q8_0 blocks
kNumber of elements (must be multiple of 32)

Definition at line 125 of file gemm_kernels_q8_0.c.

126{
127 block_q8_0 *y = (block_q8_0 *)vy;
128 const int nb = k / QK8_0; /* QK8_0 = 32 */
129
130#if defined(__AVX__)
131 const __m256 sign_bit = _mm256_set1_ps(-0.0f);
132
133 for (int i = 0; i < nb; i++) {
134 __m256 v0 = _mm256_loadu_ps(x + 0);
135 __m256 v1 = _mm256_loadu_ps(x + 8);
136 __m256 v2 = _mm256_loadu_ps(x + 16);
137 __m256 v3 = _mm256_loadu_ps(x + 24);
138 x += QK8_0;
139
140 __m256 max_abs = _mm256_andnot_ps(sign_bit, v0);
141 max_abs = _mm256_max_ps(max_abs, _mm256_andnot_ps(sign_bit, v1));
142 max_abs = _mm256_max_ps(max_abs, _mm256_andnot_ps(sign_bit, v2));
143 max_abs = _mm256_max_ps(max_abs, _mm256_andnot_ps(sign_bit, v3));
144
145 __m128 max4 = _mm_max_ps(_mm256_extractf128_ps(max_abs, 1),
146 _mm256_castps256_ps128(max_abs));
147 max4 = _mm_max_ps(max4, _mm_movehl_ps(max4, max4));
148 max4 = _mm_max_ss(max4, _mm_movehdup_ps(max4));
149 const float max_scalar = _mm_cvtss_f32(max4);
150
151#if defined(__INTEL_LLVM_COMPILER)
152 const float d = ck_q8_0_div_rounded_f32(max_scalar, 127.0f);
153 const float id = max_scalar != 0.0f
154 ? ck_q8_0_div_rounded_f32(127.0f, max_scalar)
155 : 0.0f;
156#else
157 const float d = max_scalar / 127.0f;
158 const float id = max_scalar != 0.0f ? 127.0f / max_scalar : 0.0f;
159#endif
160 y[i].d = CK_FP32_TO_FP16(d);
161
162 const __m256 mul = _mm256_set1_ps(id);
163 v0 = _mm256_mul_ps(v0, mul);
164 v1 = _mm256_mul_ps(v1, mul);
165 v2 = _mm256_mul_ps(v2, mul);
166 v3 = _mm256_mul_ps(v3, mul);
167
168 /* Match llama.cpp x86 Q8 quantization: nearest-even rounding. */
169 v0 = _mm256_round_ps(v0, _MM_ROUND_NEAREST);
170 v1 = _mm256_round_ps(v1, _MM_ROUND_NEAREST);
171 v2 = _mm256_round_ps(v2, _MM_ROUND_NEAREST);
172 v3 = _mm256_round_ps(v3, _MM_ROUND_NEAREST);
173
174 __m256i i0 = _mm256_cvtps_epi32(v0);
175 __m256i i1 = _mm256_cvtps_epi32(v1);
176 __m256i i2 = _mm256_cvtps_epi32(v2);
177 __m256i i3 = _mm256_cvtps_epi32(v3);
178
179#if defined(__AVX2__)
180 i0 = _mm256_packs_epi32(i0, i1);
181 i2 = _mm256_packs_epi32(i2, i3);
182 i0 = _mm256_packs_epi16(i0, i2);
183
184 const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7);
185 i0 = _mm256_permutevar8x32_epi32(i0, perm);
186 _mm256_storeu_si256((__m256i *)y[i].qs, i0);
187#else
188 __m128i ni0 = _mm256_castsi256_si128(i0);
189 __m128i ni1 = _mm256_extractf128_si256(i0, 1);
190 __m128i ni2 = _mm256_castsi256_si128(i1);
191 __m128i ni3 = _mm256_extractf128_si256(i1, 1);
192 __m128i ni4 = _mm256_castsi256_si128(i2);
193 __m128i ni5 = _mm256_extractf128_si256(i2, 1);
194 __m128i ni6 = _mm256_castsi256_si128(i3);
195 __m128i ni7 = _mm256_extractf128_si256(i3, 1);
196
197 ni0 = _mm_packs_epi32(ni0, ni1);
198 ni2 = _mm_packs_epi32(ni2, ni3);
199 ni4 = _mm_packs_epi32(ni4, ni5);
200 ni6 = _mm_packs_epi32(ni6, ni7);
201
202 ni0 = _mm_packs_epi16(ni0, ni2);
203 ni4 = _mm_packs_epi16(ni4, ni6);
204
205 _mm_storeu_si128((__m128i *)(y[i].qs + 0), ni0);
206 _mm_storeu_si128((__m128i *)(y[i].qs + 16), ni4);
207#endif
208 }
209#else
210 for (int i = 0; i < nb; i++) {
211 const float *xb = x + i * QK8_0;
212
213 /* Find max absolute value in block */
214 float amax = 0.0f;
215 for (int j = 0; j < QK8_0; j++) {
216 float av = xb[j] >= 0 ? xb[j] : -xb[j];
217 if (av > amax) amax = av;
218 }
219
220 /* Compute scale: d = max / 127 */
221 float d = amax / 127.0f;
222 float id = d != 0.0f ? 127.0f / amax : 0.0f;
223
224 /* Store scale as FP16 */
225 y[i].d = CK_FP32_TO_FP16(d);
226
227 /* Quantize values */
228 for (int j = 0; j < QK8_0; j++) {
229 float v = xb[j] * id;
230 int q = ck_nearest_int_q8_0(v);
231 if (q > 127) q = 127;
232 if (q < -127) q = -127;
233 y[i].qs[j] = (int8_t)q;
234 }
235 }
236#endif
237}
#define CK_FP32_TO_FP16(x)
static int ck_nearest_int_q8_0(float fval)
int8_t qs[32]
int32_t id
Definition tokenizer.h:316

Referenced by ck_test_gemm_q5_0(), ck_test_gemm_q8_0(), ck_test_gemv_q5_0(), ck_test_gemv_q5_0_q8_0(), ck_test_gemv_q8_0(), and ck_test_gemv_q8_0_q8_0().

◆ quantize_row_q8_k()

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

Definition at line 229 of file gemm_kernels_q4k_q8k.c.

229 {
230 const char *ref_env = getenv("CK_DEBUG_Q8K_REF");
231 if (ref_env && atoi(ref_env) != 0) {
232 quantize_row_q8_k_ref(x, vy, k);
233 return;
234 }
235#if defined(__AVX512F__) && defined(__AVX512BW__)
236 quantize_row_q8_k_avx512(x, vy, k);
237#elif defined(__AVX2__)
238 quantize_row_q8_k_avx2(x, vy, k);
239#elif defined(__AVX__)
240 quantize_row_q8_k_avx(x, vy, k);
241#elif defined(__SSE4_1__)
242 quantize_row_q8_k_sse(x, vy, k);
243#elif defined(__ARM_NEON) || defined(__aarch64__)
244 quantize_row_q8_k_neon(x, vy, k);
245#else
246 quantize_row_q8_k_ref(x, vy, k);
247#endif
248}
void quantize_row_q8_k_avx512(const float *x, void *vy, int k)
void quantize_row_q8_k_avx2(const float *x, void *vy, int k)
void quantize_row_q8_k_avx(const float *x, void *vy, int k)
void quantize_row_q8_k_sse(const float *x, void *vy, int k)
void quantize_row_q8_k_ref(const float *x, void *vy, int k)

Referenced by ck_test_gemm_q4_k(), ck_test_gemm_q6_k(), ck_test_gemv_q4_k(), ck_test_gemv_q6_k(), and ck_test_quantize_q8_k().

◆ recurrent_conv_state_update_forward()

void recurrent_conv_state_update_forward ( const float *  state_in,
const float *  q,
const float *  k,
const float *  v,
float *  conv_x,
float *  state_out,
int  history_len,
int  num_seqs,
int  num_tokens,
int  q_dim,
int  k_dim,
int  v_dim 
)
extern

Definition at line 100 of file recurrent_state_kernels.c.

111 {
112 if (!state_in || !q || !k || !v || !conv_x || !state_out ||
113 history_len < 0 || num_seqs <= 0 || num_tokens < 0 ||
114 q_dim < 0 || k_dim < 0 || v_dim < 0 ||
115 q_dim > INT_MAX - k_dim || q_dim + k_dim > INT_MAX - v_dim ||
116 history_len > INT_MAX - num_tokens) {
117 return;
118 }
119 const int channels = q_dim + k_dim + v_dim;
120 const int total_len = history_len + num_tokens;
121 if (channels <= 0 || total_len <= 0 || num_seqs > INT_MAX / channels ||
122 (num_tokens > 0 && num_seqs > INT_MAX / num_tokens)) {
123 return;
124 }
125
126 ck_recurrent_conv_state_args_t args = {
127 state_in, q, k, v, conv_x, state_out,
128 history_len, num_tokens, q_dim, k_dim, v_dim, channels, total_len,
129 };
130 const int jobs = num_seqs * channels;
131 const size_t elements = (size_t)jobs * (size_t)num_tokens;
132 const size_t minimum_elements_per_thread = 32768u;
133 if (num_tokens <= 1 || elements < 2u * minimum_elements_per_thread) {
135 return;
136 }
137
138 ck_threadpool_t *pool = ck_threadpool_global();
139 int active = pool ? ck_threadpool_n_threads(pool) : 1;
140 size_t useful_threads =
141 (elements + minimum_elements_per_thread - 1u) /
142 minimum_elements_per_thread;
143 if ((size_t)active > useful_threads) active = (int)useful_threads;
144 if (active > jobs) active = jobs;
145
146 if (active > 1 && ck_threadpool_thread_id(pool) <= 0) {
148 pool, active, 0, jobs, 32,
150 } else {
152 }
153}
static void ck_recurrent_conv_state_channel_range(int begin, int end, void *opaque)

Referenced by ck_test_recurrent_conv_state_update().

◆ recurrent_dt_gate_forward()

void recurrent_dt_gate_forward ( const float *  alpha,
const float *  dt_bias,
const float *  a,
float *  gate,
int  rows,
int  num_heads,
int  state_dim 
)
extern

Definition at line 66 of file recurrent_gate_kernels.c.

72 {
73 const int dim = num_heads * state_dim;
74 for (int row = 0; row < rows; ++row) {
75 const float *alpha_row = alpha + (size_t) row * (size_t) dim;
76 float *gate_row = gate + (size_t) row * (size_t) dim;
77 for (int col = 0; col < dim; ++col) {
78 const float x = alpha_row[col] + dt_bias[col];
79 gate_row[col] = recurrent_softplus(x) * a[col];
80 }
81 }
82}
static float recurrent_softplus(float x)

Referenced by ck_test_recurrent_dt_gate().

◆ recurrent_norm_gate_forward()

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

Definition at line 18 of file recurrent_norm_kernels.c.

25 {
26 const int inner_dim = num_heads * head_dim;
27 for (int row = 0; row < rows; ++row) {
28 const float *x_row = x + (size_t) row * (size_t) inner_dim;
29 const float *gate_row = gate + (size_t) row * (size_t) inner_dim;
30 float *out_row = out + (size_t) row * (size_t) inner_dim;
31
32 for (int head = 0; head < num_heads; ++head) {
33 const float *x_head = x_row + (size_t) head * (size_t) head_dim;
34 const float *gate_head = gate_row + (size_t) head * (size_t) head_dim;
35 float *out_head = out_row + (size_t) head * (size_t) head_dim;
36
37 float ms = 0.0f;
38 for (int col = 0; col < head_dim; ++col) {
39 ms += x_head[col] * x_head[col];
40 }
41 ms /= (float) head_dim;
42 const float inv_rms = 1.0f / sqrtf(ms + eps);
43
44 for (int col = 0; col < head_dim; ++col) {
45 const float g = gate_head[col];
46 const float silu = g * recurrent_sigmoid_local(g);
47 out_head[col] = x_head[col] * inv_rms * weight[col] * silu;
48 }
49 }
50 }
51}
static float recurrent_sigmoid_local(float x)
static void silu(float *x, int n)

Referenced by ck_test_recurrent_norm_gate().

◆ recurrent_qk_l2_norm_forward()

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

Definition at line 124 of file recurrent_qk_norm_kernels.c.

130 {
131 if (!q || !k || rows <= 0 || q_dim <= 0 || k_dim <= 0 ||
132 head_dim <= 0 || q_dim % head_dim != 0 || k_dim % head_dim != 0) {
133 return;
134 }
135 const int q_heads = q_dim / head_dim;
136 const int k_heads = k_dim / head_dim;
137 if (q_heads > INT_MAX - k_heads) {
138 return;
139 }
140 const int heads_per_row = q_heads + k_heads;
141 if (rows > INT_MAX / heads_per_row) {
142 return;
143 }
144
145 ck_recurrent_qk_l2_args_t args = {
146 q, k, q_heads, k_heads, head_dim, eps,
147 };
148 const int jobs = rows * heads_per_row;
149 const size_t elements = (size_t)jobs * (size_t)head_dim;
150 const size_t minimum_elements_per_thread = 32768u;
151 const uintptr_t q_begin = (uintptr_t)q;
152 const uintptr_t q_end = q_begin +
153 (size_t)rows * (size_t)q_dim * sizeof(float);
154 const uintptr_t k_begin = (uintptr_t)k;
155 const uintptr_t k_end = k_begin +
156 (size_t)rows * (size_t)k_dim * sizeof(float);
157 const int ranges_overlap = q_begin < k_end && k_begin < q_end;
158 if (ranges_overlap || elements < 2u * minimum_elements_per_thread) {
159 recurrent_l2_norm_rows_forward_one(q, rows, q_dim, head_dim, eps);
160 recurrent_l2_norm_rows_forward_one(k, rows, k_dim, head_dim, eps);
161 return;
162 }
163
164 ck_threadpool_t *pool = ck_threadpool_global();
165 int active = pool ? ck_threadpool_n_threads(pool) : 1;
166 size_t useful_threads =
167 (elements + minimum_elements_per_thread - 1u) /
168 minimum_elements_per_thread;
169 if ((size_t)active > useful_threads) {
170 active = (int)useful_threads;
171 }
172 if (active > jobs) {
173 active = jobs;
174 }
175
176 if (active > 1 && ck_threadpool_thread_id(pool) <= 0) {
178 pool, active, 0, jobs, 8,
180 } else {
181 recurrent_l2_norm_rows_forward_one(q, rows, q_dim, head_dim, eps);
182 recurrent_l2_norm_rows_forward_one(k, rows, k_dim, head_dim, eps);
183 }
184}
static void recurrent_qk_l2_norm_head_range(int begin, int end, void *opaque)
static void recurrent_l2_norm_rows_forward_one(float *x, int rows, int dim, int head_dim, float eps)

Referenced by ck_test_recurrent_qk_l2_norm().

◆ recurrent_silu_forward()

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

Definition at line 138 of file recurrent_gate_kernels.c.

141 {
142 for (int row = 0; row < rows; ++row) {
143 const float *x_row = x + (size_t) row * (size_t) dim;
144 float *out_row = out + (size_t) row * (size_t) dim;
145 for (int col = 0; col < dim; ++col) {
146 const float xv = x_row[col];
147 out_row[col] = xv * recurrent_sigmoid(xv);
148 }
149 }
150}
static float recurrent_sigmoid(float x)

Referenced by ck_test_recurrent_silu().

◆ recurrent_split_conv_qkv_forward()

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

Definition at line 97 of file recurrent_split_kernels.c.

104 {
105 recurrent_split_qkv_forward(packed_qkv, q, k, v, rows, q_dim, k_dim, v_dim);
106}
void recurrent_split_qkv_forward(const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)

Referenced by ck_test_recurrent_split_conv_qkv().

◆ recurrent_split_qkv_forward()

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

Definition at line 5 of file recurrent_split_kernels.c.

12 {
13 const int packed_dim = q_dim + k_dim + v_dim;
14 for (int row = 0; row < rows; ++row) {
15 const float *src = packed_qkv + (size_t) row * (size_t) packed_dim;
16 float *q_dst = q + (size_t) row * (size_t) q_dim;
17 float *k_dst = k + (size_t) row * (size_t) k_dim;
18 float *v_dst = v + (size_t) row * (size_t) v_dim;
19 memcpy(q_dst, src, (size_t) q_dim * sizeof(float));
20 memcpy(k_dst, src + q_dim, (size_t) k_dim * sizeof(float));
21 memcpy(v_dst, src + q_dim + k_dim, (size_t) v_dim * sizeof(float));
22 }
23}

Referenced by ck_test_recurrent_split_qkv(), and recurrent_split_conv_qkv_forward().

◆ rmsnorm_forward()

void rmsnorm_forward ( const float *  input,
const float *  gamma,
float *  output,
float *  rstd_cache,
int  tokens,
int  d_model,
int  aligned_embed_dim,
float  eps 
)
extern

Definition at line 735 of file rmsnorm_kernels.c.

743{
745 input,
746 gamma,
747 output,
748 rstd_cache,
749 tokens,
750 d_model,
751 aligned_embed_dim,
752 aligned_embed_dim,
753 eps
754 );
755}
void rmsnorm_forward_strided_f32(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

Referenced by ck_test_rmsnorm().

◆ rope_forward_qk()

void rope_forward_qk ( float *  q,
float *  k,
const float *  cos_cache,
const float *  sin_cache,
int  num_heads,
int  num_kv_heads,
int  num_tokens,
int  head_dim,
int  aligned_head_dim,
int  pos_offset 
)
extern

RoPE forward for both Q and K (common inference pattern)

Test:

test_rope.py::TestRoPEForward::test_rope_forward_qk

test_fused_attention_decode.py::TestFusedAttentionDecode::test_qk_rope

test_parity.py::test_rope_qk_parity

Combined RoPE forward for both Q and K in one call. q: [num_heads, num_tokens, head_dim] k: [num_kv_heads, num_tokens, head_dim]

After changes: make test && make llamacpp-parity-full

Definition at line 1094 of file rope_kernels.c.

1104{
1105 rope_forward_qk_with_rotary_dim(q, k, cos_cache, sin_cache, num_heads, num_kv_heads,
1106 num_tokens, head_dim, aligned_head_dim, pos_offset, head_dim);
1107}
void rope_forward_qk_with_rotary_dim(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

Referenced by ck_test_rope().

◆ rope_precompute_cache()

void rope_precompute_cache ( float *  cos_cache,
float *  sin_cache,
int  max_seq_len,
int  head_dim,
float  base,
int  rotary_dim,
const char *  scaling_type,
float  scaling_factor 
)
extern

Precompute RoPE cos/sin cache with rotary_dim and scaling support

Test:

test_rope.py::TestRoPECache::test_cache_computation

test_rope.py::TestRoPECache::test_cache_values

Precomputes cos(m * theta_i) and sin(m * theta_i) for positions 0..max_seq_len-1. Only computes for first rotary_dim channels; remaining head_dim - rotary_dim channels are NOT rotated (pass through unchanged).

Scaling types:

  • "none": Standard RoPE
  • "linear": Scale positions by 1/scaling_factor
  • "dynamic": NTK-aware dynamic scaling
  • "yarn": YaRN scaling (beta-based)
Parameters
cos_cacheOutput: [max_seq_len, rotary_dim/2] cos values
sin_cacheOutput: [max_seq_len, rotary_dim/2] sin values
max_seq_lenMaximum sequence length for cache
head_dimFull head dimension (for frequency computation)
baseRoPE base frequency (theta)
rotary_dimNumber of dimensions to rotate (0 = use head_dim)
scaling_typeScaling type string: "none", "linear", "dynamic", "yarn"
scaling_factorScaling factor (1.0 = no scaling)

After changes: make test

Definition at line 384 of file rope_kernels.c.

392{
393 // Use rotary_dim = head_dim if not specified
394 if (rotary_dim <= 0 || rotary_dim > head_dim) {
395 rotary_dim = head_dim;
396 }
397
398 // Use no scaling if not specified
399 int is_linear_scaling = 0;
400 if (scaling_type != NULL && strcmp(scaling_type, "linear") == 0 && scaling_factor > 0.0f && scaling_factor != 1.0f) {
401 is_linear_scaling = 1;
402 }
403
404 int rotary_half = rotary_dim / 2;
405
406 for (int pos = 0; pos < max_seq_len; ++pos) {
407 // Apply linear scaling to position if needed
408 float effective_pos = (float)pos;
409 if (is_linear_scaling) {
410 effective_pos = (float)pos / scaling_factor;
411 }
412
413 for (int i = 0; i < rotary_half; ++i) {
414 // Match the FP32 reference contract directly. Computing this via
415 // long-double log/exp makes the final float depend on the host
416 // libm and long-double ABI.
417 const float exponent = ((float)(2 * i)) / (float)rotary_dim;
418 const float freq_f = 1.0f / powf(base, exponent);
419 float angle_f = effective_pos * freq_f;
420 cos_cache[pos * rotary_half + i] = cosf(angle_f);
421 sin_cache[pos * rotary_half + i] = sinf(angle_f);
422 }
423 }
424}

Referenced by ck_test_rope().

◆ split_q_gate_forward()

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

Definition at line 41 of file hybrid_attention_kernels.c.

47 {
48 const int packed_dim = q_dim + gate_dim;
49 if (!packed_qg || !q || !gate || rows <= 0 || q_dim <= 0 || gate_dim <= 0) {
50 return;
51 }
52 if (group_dim <= 0) {
53 group_dim = q_dim;
54 }
55 const int q_groups = q_dim / group_dim;
56 const int gate_group_dim = (q_groups > 0 && gate_dim % q_groups == 0) ? (gate_dim / q_groups) : gate_dim;
57 for (int row = 0; row < rows; ++row) {
58 const float *src = packed_qg + (size_t) row * (size_t) packed_dim;
59 float *q_dst = q + (size_t) row * (size_t) q_dim;
60 float *gate_dst = gate + (size_t) row * (size_t) gate_dim;
61 if (q_groups > 0 && q_groups * group_dim == q_dim && q_groups * gate_group_dim == gate_dim) {
62 for (int group = 0; group < q_groups; ++group) {
63 const size_t src_group_off = (size_t) group * (size_t) (group_dim + gate_group_dim);
64 memcpy(
65 q_dst + (size_t) group * (size_t) group_dim,
66 src + src_group_off,
67 (size_t) group_dim * sizeof(float));
68 memcpy(
69 gate_dst + (size_t) group * (size_t) gate_group_dim,
70 src + src_group_off + (size_t) group_dim,
71 (size_t) gate_group_dim * sizeof(float));
72 }
73 } else {
74 memcpy(q_dst, src, (size_t) q_dim * sizeof(float));
75 memcpy(gate_dst, src + q_dim, (size_t) gate_dim * sizeof(float));
76 }
77 }
78}

Referenced by ck_test_split_q_gate().

◆ ssm_conv1d_forward()

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

Definition at line 128 of file ssm_kernels.c.

135{
136 ssm_conv1d_forward_ref(conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs);
137}
void ssm_conv1d_forward_ref(const float *conv_x, const float *kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)
Definition ssm_kernels.c:39

Referenced by ck_test_ssm_conv1d().

◆ swiglu_forward()

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

SwiGLU forward pass

Test:

test_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

SwiGLU: y = silu(gate) * up where silu(x) = x * sigmoid(x)

After changes: make test && make llamacpp-parity-full

Definition at line 283 of file swiglu_kernels.c.

287{
288 const char *fast_env = getenv("CK_SWIGLU_FAST");
289 const char *exact_env = getenv("CK_SWIGLU_EXACT");
291 !(fast_env && atoi(fast_env) != 0) ||
292 (exact_env && atoi(exact_env) != 0)) {
293 swiglu_forward_exact(input, output, tokens, dim);
294 return;
295 }
296
297 int T = tokens;
298 int D = dim;
299
300 for (int t = 0; t < T; ++t) {
301 const float *row = input + (size_t)t * (2 * D);
302 float *out_row = output + (size_t)t * D;
303 int d = 0;
304
305#if defined(__AVX512F__)
306 // AVX-512: Process 16 floats at a time
307 for (; d + 16 <= D; d += 16) {
308 __m512 a = _mm512_loadu_ps(&row[d]); // gate
309 __m512 b = _mm512_loadu_ps(&row[D + d]); // value
310
311 __m512 s = sigmoid512_fast(a); // sigmoid(a)
312 __m512 silu = _mm512_mul_ps(a, s); // silu(a) = a * sigmoid(a)
313 __m512 y = _mm512_mul_ps(silu, b); // y = silu(a) * b
314
315 _mm512_storeu_ps(&out_row[d], y);
316 }
317#elif defined(__AVX2__)
318 // AVX2: Process 8 floats at a time
319 for (; d + 8 <= D; d += 8) {
320 __m256 a = _mm256_loadu_ps(&row[d]); // gate
321 __m256 b = _mm256_loadu_ps(&row[D + d]); // value
322
323 __m256 s = sigmoid256_fast(a); // sigmoid(a)
324 __m256 silu = _mm256_mul_ps(a, s); // silu(a) = a * sigmoid(a)
325 __m256 y = _mm256_mul_ps(silu, b); // y = silu(a) * b
326
327 _mm256_storeu_ps(&out_row[d], y);
328 }
329#elif defined(__AVX__)
330 // AVX1: Vectorize arithmetic, use scalar sigmoid
331 float a_arr[8] __attribute__((aligned(32)));
332 float s_arr[8] __attribute__((aligned(32)));
333
334 for (; d + 8 <= D; d += 8) {
335 __m256 a = _mm256_loadu_ps(&row[d]); // gate
336 __m256 b = _mm256_loadu_ps(&row[D + d]); // value
337
338 // Compute sigmoid scalarly
339 _mm256_store_ps(a_arr, a);
340 for (int j = 0; j < 8; ++j) {
341 s_arr[j] = sigmoid_scalar(a_arr[j]);
342 }
343 __m256 s = _mm256_load_ps(s_arr);
344
345 __m256 silu = _mm256_mul_ps(a, s); // silu(a) = a * sigmoid(a)
346 __m256 y = _mm256_mul_ps(silu, b); // y = silu(a) * b
347
348 _mm256_storeu_ps(&out_row[d], y);
349 }
350#endif
351
352 // Scalar fallback for remaining elements
353 for (; d < D; ++d) {
354 float a = row[d]; // gate
355 float b = row[D + d]; // value
356
357 float s = sigmoid_scalar(a); // sigmoid(a)
358 float silu = a * s; // silu(a) = a * sigmoid(a)
359
360 out_row[d] = silu * b;
361 }
362 }
363}
float sigmoid_scalar(float x)
void swiglu_forward_exact(const float *input, float *output, int tokens, int dim)
__attribute__((visibility("default"))) CKTokenizer *ck_tokenizer_create(CKTokenizerType type)

Referenced by ck_test_swiglu().

◆ vec_dot_q5_0_q8_0()

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

Auto-dispatch quantized dot product Q5_0 x Q8_0.

Dispatch priority:

  1. AVX512 (best performance on modern Intel/AMD)
  2. AVX (256-bit float ops, works on Sandy/Ivy Bridge and newer)
  3. SSSE3 (128-bit fallback)
  4. Reference scalar (last resort)

Definition at line 1602 of file gemm_kernels_q5_0.c.

1603{
1604#if defined(__AVX2__)
1605 /* llama.cpp uses the packed AVX2 dot on AVX-512 hosts as well. It keeps
1606 * Q5/Q8 data in byte lanes and avoids the per-block 32-bit lane expansion
1607 * overhead of the baseline AVX-512 path. */
1608 vec_dot_q5_0_q8_0_avx2(n, s, vx, vy);
1609#elif defined(__AVX512F__)
1610 vec_dot_q5_0_q8_0_avx512(n, s, vx, vy);
1611#elif defined(__ARM_NEON) || defined(__aarch64__)
1612 vec_dot_q5_0_q8_0_neon(n, s, vx, vy);
1613#elif defined(__AVX__)
1614 /* AVX for 256-bit float ops (works on Ivy Bridge and newer) */
1615 vec_dot_q5_0_q8_0_avx(n, s, vx, vy);
1616#elif defined(__SSSE3__)
1617 /* SSSE3 - most efficient on older CPUs */
1618 vec_dot_q5_0_q8_0_sse(n, s, vx, vy);
1619#else
1620 vec_dot_q5_0_q8_0_ref(n, s, vx, vy);
1621#endif
1622}
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)

Referenced by ck_test_vec_dot_q5_0_q8_0().

◆ vec_dot_q6_k_q8_k()

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

Q6_K x Q8_K dot product (single row)

Definition at line 1339 of file gemm_kernels_q6k_q8k.c.

1340{
1341 if (!s || !vx || !vy || n <= 0) {
1342 return;
1343 }
1344
1345 const block_q6_K *x = (const block_q6_K *)vx;
1346 const block_q8_K *y = (const block_q8_K *)vy;
1347
1348 /* This is the architecture-neutral scalar oracle. Production x86 dispatch
1349 * uses the separately parity-tested SIMD reduction tree. */
1350 *s = dot_q6_k_q8_k_ref(x, y, n);
1351}
static 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.

◆ vec_dot_q8_0_q8_0()

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

Auto-dispatch quantized dot product Q8_0 x Q8_0.

Definition at line 1368 of file gemm_kernels_q8_0.c.

1369{
1370 if (ck_q8_0_q8_0_debug_ref()) {
1371 vec_dot_q8_0_q8_0_ref(n, s, vx, vy);
1372 return;
1373 }
1374#ifdef __AVX512F__
1375 vec_dot_q8_0_q8_0_avx512(n, s, vx, vy);
1376#elif defined(__AVX2__)
1377 vec_dot_q8_0_q8_0_avx2(n, s, vx, vy);
1378#elif defined(__ARM_NEON) || defined(__aarch64__)
1379 vec_dot_q8_0_q8_0_neon(n, s, vx, vy);
1380#elif defined(__AVX__)
1381 vec_dot_q8_0_q8_0_avx(n, s, vx, vy);
1382#elif defined(__SSE4_1__)
1383 vec_dot_q8_0_q8_0_sse(n, s, vx, vy);
1384#else
1385 vec_dot_q8_0_q8_0_ref(n, s, vx, vy);
1386#endif
1387}
static int ck_q8_0_q8_0_debug_ref(void)
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)

Referenced by ck_test_vec_dot_q8_0_q8_0().