kv_cache_kernels.c File Reference

KV-cache helper kernels (head-major layout) More...

#include "ckernel_engine.h"
#include "bf16_utils.h"
#include <stddef.h>
#include <string.h>

Go to the source code of this file.

Functions
static void	ck_local_fp32_to_bf16_row (const float *src, uint16_t *dst, int n)
static void	ck_local_fp32_to_fp16_row (const float *src, uint16_t *dst, int n)
void	kv_cache_repack_head_major_inplace (float *buf, int num_heads, int tokens, int cache_capacity, int aligned_head_dim)
void	kv_cache_store (float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)
void	kv_cache_store_batch_bf16 (uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)
void	kv_cache_store_batch_f16 (uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)
void	kv_cache_store_batch_f32 (float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)
void	kv_cache_store_bf16 (uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)
void	kv_cache_store_f16 (uint16_t *__restrict kv_cache_k, uint16_t *__restrict kv_cache_v, const float *__restrict k, const float *__restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)
void	kv_cache_store_shared_q (float *__restrict kv_cache_k, float *__restrict kv_cache_v, const float *__restrict q, int layer, int pos, int num_heads, int head_dim, int max_seq_len)
void	kv_cache_write_head_major (const float *__restrict k_token, const float *__restrict v_token, float *__restrict k_cache, float *__restrict v_cache, int num_kv_heads, int token_index, int cache_capacity, int head_dim, int aligned_head_dim)
void	logits_copy_to_position (const float *__restrict src, float *__restrict dst, int position, int vocab_size)
	Copy logits to position-indexed location in output buffer.

Detailed Description

KV-cache helper kernels (head-major layout)

CK-ENGINE KERNEL RULES:

1. NO malloc/free - memory via bump allocator, pointers passed in
2. NO OpenMP - parallelization at orchestrator/codegen layer
3. API must define: inputs, outputs, workspace, and memory layouts
4. Pure computation - deterministic, no side effects

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

Small, explicit helpers used by the runtime/orchestrator to maintain per-layer KV caches during autoregressive decoding.

Layout: k_cache[kv_head, token, aligned_head_dim] v_cache[kv_head, token, aligned_head_dim] with contiguous row-major storage and stride aligned_head_dim.

Definition in file kv_cache_kernels.c.

Function Documentation

ck_local_fp32_to_bf16_row()

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

inlinestatic

Definition at line 39 of file kv_cache_kernels.c.

{
    if (!src || !dst || n <= 0) {
        return;
    }
    for (int i = 0; i < n; ++i) {
        dst[i] = float_to_bf16(src[i]);
    }
}

References float_to_bf16().

Referenced by kv_cache_store_batch_bf16(), and kv_cache_store_bf16().

ck_local_fp32_to_fp16_row()

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

inlinestatic

Definition at line 29 of file kv_cache_kernels.c.

{
    if (!src || !dst || n <= 0) {
        return;
    }
    for (int i = 0; i < n; ++i) {
        dst[i] = CK_FP32_TO_FP16(src[i]);
    }
}

References CK_FP32_TO_FP16.

Referenced by kv_cache_store_batch_f16(), and kv_cache_store_f16().

kv_cache_repack_head_major_inplace()

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

Definition at line 49 of file kv_cache_kernels.c.

{
    if (!buf) {
        return;
    }
    if (num_heads <= 0 || tokens <= 0 || cache_capacity <= 0 || aligned_head_dim <= 0) {
        return;
    }
    if (tokens > cache_capacity) {
        tokens = cache_capacity;
    }
    if (tokens == cache_capacity) {
        return;
    }
 
    const size_t old_head_stride = (size_t)tokens * (size_t)aligned_head_dim;
    const size_t new_head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    const size_t bytes = (size_t)tokens * (size_t)aligned_head_dim * sizeof(float);
 
    // Move head blocks from high to low to avoid overwriting source data
    // for heads that have not yet been moved.
    for (int h = num_heads - 1; h >= 0; --h) {
        float *src = buf + (size_t)h * old_head_stride;
        float *dst = buf + (size_t)h * new_head_stride;
        memmove(dst, src, bytes);
    }
}

Referenced by qwen2_0_5b_decode_forward_prefill_impl(), qwen2_0_5b_decode_forward_prefill_impl(), and qwen2_0_5b_decode_forward_prefill_impl().

kv_cache_store()

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

Definition at line 122 of file kv_cache_kernels.c.

{
    (void)layer;
    kv_cache_write_head_major(k, v,
                              kv_cache_k, kv_cache_v,
                              num_kv_heads,
                              pos,
                              max_seq_len,
                              head_dim,
                              head_dim);
}

References kv_cache_write_head_major().

kv_cache_store_batch_bf16()

void kv_cache_store_batch_bf16	(	uint16_t *__restrict	kv_cache_k,
		uint16_t *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	start_pos,
		int	num_tokens,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 291 of file kv_cache_kernels.c.

{
    if (!kv_cache_k || !kv_cache_v || !k || !v ||
        start_pos < 0 || num_tokens <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || max_seq_len <= 0 ||
        start_pos > max_seq_len - num_tokens) {
        return;
    }
 
    const size_t compact_head_stride = (size_t)num_tokens * (size_t)head_dim;
    const size_t cache_head_stride = (size_t)max_seq_len * (size_t)head_dim;
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_head = k + (size_t)h * compact_head_stride;
        const float *v_head = v + (size_t)h * compact_head_stride;
        uint16_t *k_head_cache = kv_cache_k + (size_t)h * cache_head_stride;
        uint16_t *v_head_cache = kv_cache_v + (size_t)h * cache_head_stride;
        for (int t = 0; t < num_tokens; ++t) {
            const size_t src_offset = (size_t)t * (size_t)head_dim;
            const size_t dst_offset = (size_t)(start_pos + t) * (size_t)head_dim;
            ck_local_fp32_to_bf16_row(k_head + src_offset, k_head_cache + dst_offset, head_dim);
            ck_local_fp32_to_bf16_row(v_head + src_offset, v_head_cache + dst_offset, head_dim);
        }
    }
}

References ck_local_fp32_to_bf16_row().

kv_cache_store_batch_f16()

void kv_cache_store_batch_f16	(	uint16_t *__restrict	kv_cache_k,
		uint16_t *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	start_pos,
		int	num_tokens,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 255 of file kv_cache_kernels.c.

{
    if (!kv_cache_k || !kv_cache_v || !k || !v) {
        return;
    }
    if (start_pos < 0 || num_tokens <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || max_seq_len <= 0 ||
        start_pos > max_seq_len - num_tokens) {
        return;
    }
 
    const size_t compact_head_stride = (size_t)num_tokens * (size_t)head_dim;
    const size_t cache_head_stride = (size_t)max_seq_len * (size_t)head_dim;
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_head = k + (size_t)h * compact_head_stride;
        const float *v_head = v + (size_t)h * compact_head_stride;
        uint16_t *k_head_cache = kv_cache_k + (size_t)h * cache_head_stride;
        uint16_t *v_head_cache = kv_cache_v + (size_t)h * cache_head_stride;
        for (int t = 0; t < num_tokens; ++t) {
            const size_t src_offset = (size_t)t * (size_t)head_dim;
            const size_t dst_offset = (size_t)(start_pos + t) * (size_t)head_dim;
            ck_local_fp32_to_fp16_row(k_head + src_offset, k_head_cache + dst_offset, head_dim);
            ck_local_fp32_to_fp16_row(v_head + src_offset, v_head_cache + dst_offset, head_dim);
        }
    }
}

References ck_local_fp32_to_fp16_row().

kv_cache_store_batch_f32()

void kv_cache_store_batch_f32	(	float *__restrict	kv_cache_k,
		float *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	start_pos,
		int	num_tokens,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 223 of file kv_cache_kernels.c.

{
    if (!kv_cache_k || !kv_cache_v || !k || !v ||
        start_pos < 0 || num_tokens <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || max_seq_len <= 0 ||
        start_pos > max_seq_len - num_tokens) {
        return;
    }
 
    const size_t compact_head_stride = (size_t)num_tokens * (size_t)head_dim;
    const size_t cache_head_stride = (size_t)max_seq_len * (size_t)head_dim;
    const size_t token_bytes = (size_t)num_tokens * (size_t)head_dim * sizeof(float);
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_head = k + (size_t)h * compact_head_stride;
        const float *v_head = v + (size_t)h * compact_head_stride;
        float *k_head_cache = kv_cache_k + (size_t)h * cache_head_stride
                            + (size_t)start_pos * (size_t)head_dim;
        float *v_head_cache = kv_cache_v + (size_t)h * cache_head_stride
                            + (size_t)start_pos * (size_t)head_dim;
        memcpy(k_head_cache, k_head, token_bytes);
        memcpy(v_head_cache, v_head, token_bytes);
    }
}

kv_cache_store_bf16()

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

Definition at line 195 of file kv_cache_kernels.c.

{
    (void)layer;
    if (!kv_cache_k || !kv_cache_v || !k || !v ||
        num_kv_heads <= 0 || pos < 0 || pos >= max_seq_len ||
        head_dim <= 0 || max_seq_len <= 0) {
        return;
    }
 
    const size_t head_stride = (size_t)max_seq_len * (size_t)head_dim;
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k + (size_t)h * (size_t)head_dim;
        const float *v_src = v + (size_t)h * (size_t)head_dim;
        uint16_t *k_dst = kv_cache_k + (size_t)h * head_stride + (size_t)pos * (size_t)head_dim;
        uint16_t *v_dst = kv_cache_v + (size_t)h * head_stride + (size_t)pos * (size_t)head_dim;
        ck_local_fp32_to_bf16_row(k_src, k_dst, head_dim);
        ck_local_fp32_to_bf16_row(v_src, v_dst, head_dim);
    }
}

References ck_local_fp32_to_bf16_row().

kv_cache_store_f16()

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

Definition at line 161 of file kv_cache_kernels.c.

{
    (void)layer;
    if (!kv_cache_k || !kv_cache_v || !k || !v) {
        return;
    }
    if (num_kv_heads <= 0 || pos < 0 || head_dim <= 0 || max_seq_len <= 0) {
        return;
    }
    if (pos >= max_seq_len) {
        return;
    }
 
    const size_t head_stride = (size_t)max_seq_len * (size_t)head_dim;
    const size_t token_stride = (size_t)head_dim;
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k + (size_t)h * token_stride;
        const float *v_src = v + (size_t)h * token_stride;
        uint16_t *k_dst = kv_cache_k + (size_t)h * head_stride + (size_t)pos * token_stride;
        uint16_t *v_dst = kv_cache_v + (size_t)h * head_stride + (size_t)pos * token_stride;
        ck_local_fp32_to_fp16_row(k_src, k_dst, head_dim);
        ck_local_fp32_to_fp16_row(v_src, v_dst, head_dim);
    }
}

References ck_local_fp32_to_fp16_row().

kv_cache_store_shared_q()

void kv_cache_store_shared_q	(	float *__restrict	kv_cache_k,
		float *__restrict	kv_cache_v,
		const float *__restrict	q,
		int	layer,
		int	pos,
		int	num_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 142 of file kv_cache_kernels.c.

{
    (void)layer;
    kv_cache_write_head_major(q, q,
                              kv_cache_k, kv_cache_v,
                              num_heads,
                              pos,
                              max_seq_len,
                              head_dim,
                              head_dim);
}

References kv_cache_write_head_major().

kv_cache_write_head_major()

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

Definition at line 81 of file kv_cache_kernels.c.

{
    if (!k_token || !v_token || !k_cache || !v_cache) {
        return;
    }
    if (num_kv_heads <= 0 || token_index < 0 || cache_capacity <= 0) {
        return;
    }
    if (token_index >= cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    const size_t token_stride = (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k_token + (size_t)h * token_stride;
        const float *v_src = v_token + (size_t)h * token_stride;
 
        float *k_dst = k_cache + (size_t)h * head_stride + (size_t)token_index * token_stride;
        float *v_dst = v_cache + (size_t)h * head_stride + (size_t)token_index * token_stride;
 
        for (int d = 0; d < head_dim; ++d) {
            k_dst[d] = k_src[d];
            v_dst[d] = v_src[d];
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) {
            k_dst[d] = 0.0f;
            v_dst[d] = 0.0f;
        }
    }
}

Referenced by ck_layer_forward_rmsnorm_swiglu_decode(), ck_layer_forward_rmsnorm_swiglu_decode_fused(), ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl(), ck_layer_forward_rmsnorm_swiglu_decode_q4_k(), ck_layer_forward_rmsnorm_swiglu_decode_quant(), kv_cache_store(), kv_cache_store_shared_q(), mega_fused_attention_decode_workspace(), qwen2_0_5b_decode_layer_0_decode(), qwen2_0_5b_decode_layer_0_decode(), qwen2_0_5b_decode_layer_0_decode(), qwen2_0_5b_decode_layer_10_decode(), qwen2_0_5b_decode_layer_10_decode(), qwen2_0_5b_decode_layer_10_decode(), qwen2_0_5b_decode_layer_11_decode(), qwen2_0_5b_decode_layer_11_decode(), qwen2_0_5b_decode_layer_11_decode(), qwen2_0_5b_decode_layer_12_decode(), qwen2_0_5b_decode_layer_12_decode(), qwen2_0_5b_decode_layer_12_decode(), qwen2_0_5b_decode_layer_13_decode(), qwen2_0_5b_decode_layer_13_decode(), qwen2_0_5b_decode_layer_13_decode(), qwen2_0_5b_decode_layer_14_decode(), qwen2_0_5b_decode_layer_14_decode(), qwen2_0_5b_decode_layer_14_decode(), qwen2_0_5b_decode_layer_15_decode(), qwen2_0_5b_decode_layer_15_decode(), qwen2_0_5b_decode_layer_15_decode(), qwen2_0_5b_decode_layer_16_decode(), qwen2_0_5b_decode_layer_16_decode(), qwen2_0_5b_decode_layer_16_decode(), qwen2_0_5b_decode_layer_17_decode(), qwen2_0_5b_decode_layer_17_decode(), qwen2_0_5b_decode_layer_17_decode(), qwen2_0_5b_decode_layer_18_decode(), qwen2_0_5b_decode_layer_18_decode(), qwen2_0_5b_decode_layer_18_decode(), qwen2_0_5b_decode_layer_19_decode(), qwen2_0_5b_decode_layer_19_decode(), qwen2_0_5b_decode_layer_19_decode(), qwen2_0_5b_decode_layer_1_decode(), qwen2_0_5b_decode_layer_1_decode(), qwen2_0_5b_decode_layer_1_decode(), qwen2_0_5b_decode_layer_20_decode(), qwen2_0_5b_decode_layer_20_decode(), qwen2_0_5b_decode_layer_20_decode(), qwen2_0_5b_decode_layer_21_decode(), qwen2_0_5b_decode_layer_21_decode(), qwen2_0_5b_decode_layer_21_decode(), qwen2_0_5b_decode_layer_22_decode(), qwen2_0_5b_decode_layer_22_decode(), qwen2_0_5b_decode_layer_22_decode(), qwen2_0_5b_decode_layer_23_decode(), qwen2_0_5b_decode_layer_23_decode(), qwen2_0_5b_decode_layer_23_decode(), qwen2_0_5b_decode_layer_2_decode(), qwen2_0_5b_decode_layer_2_decode(), qwen2_0_5b_decode_layer_2_decode(), qwen2_0_5b_decode_layer_3_decode(), qwen2_0_5b_decode_layer_3_decode(), qwen2_0_5b_decode_layer_3_decode(), qwen2_0_5b_decode_layer_4_decode(), qwen2_0_5b_decode_layer_4_decode(), qwen2_0_5b_decode_layer_4_decode(), qwen2_0_5b_decode_layer_5_decode(), qwen2_0_5b_decode_layer_5_decode(), qwen2_0_5b_decode_layer_5_decode(), qwen2_0_5b_decode_layer_6_decode(), qwen2_0_5b_decode_layer_6_decode(), qwen2_0_5b_decode_layer_6_decode(), qwen2_0_5b_decode_layer_7_decode(), qwen2_0_5b_decode_layer_7_decode(), qwen2_0_5b_decode_layer_7_decode(), qwen2_0_5b_decode_layer_8_decode(), qwen2_0_5b_decode_layer_8_decode(), qwen2_0_5b_decode_layer_8_decode(), qwen2_0_5b_decode_layer_9_decode(), qwen2_0_5b_decode_layer_9_decode(), and qwen2_0_5b_decode_layer_9_decode().

logits_copy_to_position()

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

Copy logits to position-indexed location in output buffer.

Used in decode mode to copy single-token logits from position 0 to the correct sequence position. This moves buffer management logic from codegen to the IR layer, making codegen "dumb" - just emit kernel calls, no runtime if-statements.

Parameters

src	Source logits buffer (single token) [vocab_size]
dst	Destination logits buffer [max_seq_len, vocab_size]
position	Token position index (0-based)
vocab_size	Number of logits per token

Definition at line 337 of file kv_cache_kernels.c.

{
    if (!src || !dst || position < 0 || vocab_size <= 0) {
        return;
    }
 
    // Copy logits to dst[position * vocab_size : (position+1) * vocab_size]
    // Use memmove for safety in case src and dst overlap (e.g., src == dst)
    float *dst_pos = dst + (size_t)position * (size_t)vocab_size;
    memmove(dst_pos, src, (size_t)vocab_size * sizeof(float));
}

References vocab_size.