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 function

rope_k_cache

def rope_k_cache[freq_dtype: DType, cache_t: KVCacheT, width: SIMDLength, //, *, interleaved: Bool, has_nope_prefix: Bool = False, rope_prefix_dim: Int = Int(0)](k_cache: cache_t, b_idx: Int, h_idx: Int, s_idx: Int, d_idx: Int, freq_val: SIMD[freq_dtype, width], head_size: Int)

Applies RoPE to a key cache entry and stores the result back in the cache.

Loads the key cache values at the given batch, head, sequence, and dimension indices, applies the rotary position embedding via rope_value, and writes the transformed result back to the cache. Supports both interleaved and split layouts, and optionally leaves a nope prefix region unrotated.

Parameters:

freq_dtype (DType): Element type of the freq_val frequency coefficients (inferred).
cache_t (KVCacheT): KV cache view type used to load and store key cache entries (inferred).
width (SIMDLength): Number of elements per SIMD vector (inferred).
interleaved (Bool): Whether the RoPE weights use interleaved real and imaginary layout.
has_nope_prefix (Bool): Whether a leading prefix of the head dimension is left unrotated (defaults to False).
rope_prefix_dim (Int): Number of trailing head dimensions that undergo RoPE when has_nope_prefix is set (defaults to 0).

Args:

k_cache (cache_t): The KV cache key view to read from and write to.
b_idx (Int): The batch index into the cache.
h_idx (Int): The head index into the cache.
s_idx (Int): The sequence position index into the cache.
d_idx (Int): The head dimension index into the cache.
freq_val (SIMD[freq_dtype, width]): The RoPE frequency coefficients for this position.
head_size (Int): The size of each attention head dimension.