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

fused_qk_rms_norm_rope_ragged_paged

def fused_qk_rms_norm_rope_ragged_paged[dtype: DType, freq_dtype: DType, params: KVCacheStaticParams, page_size: Int, cache_dtype: DType, //, target: StringSlice[ImmStaticOrigin], multiply_before_cast: Bool, interleaved: Bool, q_input_fn: def[width: Int, alignment: Int](token: Int, head: Int, col: Int) capturing thin -> SIMD[dtype, width]](kv_collection: PagedKVCacheCollection[cache_dtype, params, page_size, scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], q_gamma: TileTensor[dtype, Storage=q_gamma.Storage, address_space=q_gamma.address_space, linear_idx_type=q_gamma.linear_idx_type], k_gamma: TileTensor[dtype, Storage=k_gamma.Storage, address_space=k_gamma.address_space, linear_idx_type=k_gamma.linear_idx_type], freqs_cis: TileTensor[freq_dtype, Storage=freqs_cis.Storage, address_space=freqs_cis.address_space, linear_idx_type=freqs_cis.linear_idx_type], epsilon: Float32, weight_offset: Scalar[dtype], layer_idx: UInt32, input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, address_space=input_row_offsets.address_space, linear_idx_type=input_row_offsets.linear_idx_type], q_output: TileTensor[dtype, Storage=q_output.Storage, address_space=q_output.address_space, linear_idx_type=q_output.linear_idx_type], context: DeviceContext)

Fuses per-head RMSNorm and RoPE for Q and new K-cache entries.

Applies per-head RMSNorm to Q and the newly written key-cache entries, then RoPE to the normalized values, in a single GPU launch (the fusion of fused_qk_rms_norm_ragged_paged and fused_qk_rope_ragged).

Q is read through q_input_fn(token, head, col), so the caller reads Q from wherever it lives: a rank-3 projection, or a slice + reshape of a combined matmul output that the graph compiler folds into the read. K comes from the paged key cache.

The RoPE dimension is taken from freqs_cis.static_shape[1]. When it is smaller than the head dimension, RoPE is applied only to the prefix [0, rope_dim) of each head (non-interleaved layout) and the suffix is left un-roped, matching fused_qk_rope_ragged's has_nope_prefix path.