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_dual_qk_rms_norm_rope_ragged_paged

def fused_dual_qk_rms_norm_rope_ragged_paged[dtype: DType, freq_dtype: DType, main_params: KVCacheStaticParams, main_page_size: Int, main_cache_dtype: DType, index_params: KVCacheStaticParams, index_page_size: Int, index_cache_dtype: DType, //, target: StringSlice[ImmStaticOrigin], multiply_before_cast: Bool, interleaved: Bool, main_q_input_fn: def[width: Int, alignment: Int](token: Int, head: Int, col: Int) capturing thin -> SIMD[dtype, width], index_q_input_fn: def[width: Int, alignment: Int](token: Int, head: Int, col: Int) capturing thin -> SIMD[dtype, width]](main_kv_collection: PagedKVCacheCollection[main_cache_dtype, main_params, main_page_size, scale_dtype_=main_kv_collection.scale_dtype_, quantization_granularity_=main_kv_collection.quantization_granularity_], index_kv_collection: PagedKVCacheCollection[index_cache_dtype, index_params, index_page_size, scale_dtype_=index_kv_collection.scale_dtype_, quantization_granularity_=index_kv_collection.quantization_granularity_], q_main_gamma: TileTensor[dtype, Storage=q_main_gamma.Storage, address_space=q_main_gamma.address_space, linear_idx_type=q_main_gamma.linear_idx_type], k_main_gamma: TileTensor[dtype, Storage=k_main_gamma.Storage, address_space=k_main_gamma.address_space, linear_idx_type=k_main_gamma.linear_idx_type], q_index_gamma: TileTensor[dtype, Storage=q_index_gamma.Storage, address_space=q_index_gamma.address_space, linear_idx_type=q_index_gamma.linear_idx_type], k_index_gamma: TileTensor[dtype, Storage=k_index_gamma.Storage, address_space=k_index_gamma.address_space, linear_idx_type=k_index_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], main_epsilon: Float32, index_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_main_output: TileTensor[dtype, Storage=q_main_output.Storage, address_space=q_main_output.address_space, linear_idx_type=q_main_output.linear_idx_type], q_index_output: TileTensor[dtype, Storage=q_index_output.Storage, address_space=q_index_output.address_space, linear_idx_type=q_index_output.linear_idx_type], context: DeviceContext)

Fuses two fused_qk_rms_norm_rope_ragged_paged launches into one.

MiniMax-M3 sparse layers fire the fused per-head RMSNorm+RoPE op twice back to back: once for the main GQA Q / K cache and once for the lightning indexer's IndexQ / index-K cache. Both read (disjoint) slices of the same combined QKV+IndexQ matmul output, share one input_row_offsets, and share one freqs_cis table, so they can run in a single launch. The grid is a four-band concatenation [ q_main | k_main | q_index | k_index ]; each row selects its band's cache, gamma, DPS output, Q read lambda, and epsilon at runtime. main_q_input_fn / index_q_input_fn read each band's Q; the respective key caches are updated in place.

Because the two paged caches can be different types (the main GQA cache and the indexer's single-head K-only cache differ in KV-heads-per-device under tensor parallelism), the kernel is parameterized on two independent cache_t types. All compile-time RoPE geometry (dtype, rope_dim via freqs_cis.static_shape[1], interleaved, head_size) must be identical across both bands; a divergence (e.g. a future main full-128 rope while the indexer stays partial-64) trips a compile-time assert rather than silently mis-roping a band.