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

Mojo module

kv_cache_ragged

Implements KV-cache kernels for ragged (variable-length) sequences used in continuous batching.

Functions

generic_cross_attention_kv_cache: Dispatches cross-attention flash attention over a ragged batch against a paged KV cache.
generic_flare_mla_decode_kv_cache_ragged: Dispatches MLA decode attention over a ragged batch against a paged KV cache.
generic_flare_mla_decompress_k_cache_ragged_paged: Decompresses MLA latent K cache rows into a full K buffer via a matmul.
generic_flare_mla_prefill_kv_cache_ragged: Dispatches MLA prefill attention over a ragged batch against a paged KV cache.
generic_flare_mla_prefill_ragged_paged_plan: Computes the MLA prefill plan for a ragged paged KV cache.
generic_flash_attention_kv_cache_ragged: Dispatches flash attention over a ragged batch against a paged KV cache.
generic_flash_attention_kv_cache_ragged_sink: Dispatches flash attention over a ragged batch with attention sink weights.
generic_fused_qk_rope_bshd_paged_ragged: Performs a fused RoPE projection for Q and K projections.
generic_fused_qkv_index_matmul_kv_cache_paged_ragged: Performs a fused QKV + index-QK matmul (BF16, non-scaled). Q and IndexQ are written to the combined output buffer; K/V are scattered into the MAIN kv_collection and IndexK into the INDEX index_kv_collection.
generic_fused_qkv_index_matmul_kv_cache_paged_ragged_scale_float4: Performs a fused QKV + index-QK matmul. Q is written to q_output and IndexQ to iq_output; K/V are scattered into the MAIN kv_collection and IndexK into the INDEX index_kv_collection.
generic_fused_qkv_matmul_kv_cache_paged_ragged: 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_paged_ragged_bias: 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_paged_ragged_scale: 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_paged_ragged_scale_float4: 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_kv_cache_radd_dispatch: Adds an input tensor elementwise into the paged KV cache in-place.
k_matmul_ragged_paged: Performs a matmul, writing the output into a mutable PagedKVCacheCollection object.
k_matmul_ragged_paged_scale: Performs a matmul, writing the output into a mutable PagedKVCacheCollection object.
kv_cache_2m_iadd_dispatch: In-place add to paged KV cache with concatenated K/V layout. This kernel is only used for LoRA.
kv_cache_row_offsets_ragged_paged: Builds cumulative valid-cache row offsets for a ragged prefill batch.
kv_cache_row_offsets_ragged_paged_kernel: Computes cumulative valid-cache row offsets for one batch index in a ragged prefill batch.
kv_cache_store_padded: Stores padded input values into a paged KV cache via an elementwise kernel.
kv_cache_store_ragged: Stores ragged input values into a paged KV cache via an elementwise kernel.
kv_matmul_ragged_paged: Performs a matmul, writing the output into a mutable ContinuousBatchingKVCacheCollection object.
unfused_qkv_matmul_ragged_paged_gguf_quantized: Performs a quantized matmul, writing the output into a mutable PagedKVCacheCollection object.

Functions​

Functions