For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

kv_cache

Functions

• ​copy_kv_pages_d2h: Copy selected pages for a single layer from device to host KV cache.
• ​fused_dual_qk_rms_norm_rope_ragged_paged: Fuses two fused_qk_rms_norm_rope_ragged_paged launches into one.
• ​fused_qk_rms_norm_ragged_paged: Applies per-head RMSNorm to Q and new K-cache entries in one GPU launch.
• ​fused_qk_rms_norm_rope_ragged_paged: Fuses per-head RMSNorm and RoPE for Q and new K-cache entries.
• ​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:
• ​generic_get_paged_cache_with_scales: Create a PagedKVCacheCollection with scales for MLA attention.
• ​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.
• ​rms_norm_value_cache_ragged_paged: Performs RMSNorm in place on new entries in the value cache.