Mojo module

kv_cache

Functions

copy_kv_pages_d2h: Copy selected pages for a single layer from device to host KV cache.
generic_flash_attention_kv_cache_padded:
generic_flash_attention_kv_cache_padded_materialized_mask:
generic_fused_qk_rope_bshd_continuous_batch: Performs a fused RoPE projection for Q and K projections.
generic_fused_qk_rope_bshd_paged: Performs a fused RoPE projection for Q and K with paged KV cache.
generic_fused_qkv_matmul_kv_cache_bshd_continuous_batch: Performs a fused QKV matmul. Q outputs are written to the output argument while K and V outputs are written in-place into k_cache and v_cache.
generic_fused_qkv_matmul_kv_cache_bshd_paged: Performs a fused QKV matmul. Q outputs are written to the output argument while K and V outputs are written in-place into k_cache and v_cache.
generic_get_continuous_cache:
generic_get_paged_cache:
print_kv_cache_cont_batch_generic_cpu:
print_kv_cache_cont_batch_generic_gpu:
print_kv_cache_paged_generic_cpu:
print_kv_cache_paged_generic_gpu:
rms_norm_kv_cache_ragged_paged: Performs RMSNorm in place on new entries in the key cache.