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 function

generic_fused_qk_rope_bshd_paged_ragged_kernel_api

def generic_fused_qk_rope_bshd_paged_ragged_kernel_api[dtype: DType, freq_dtype: DType, cache_dtype: DType, //, *, interleaved: Bool, has_position_ids: Bool, target: StringSlice[ImmStaticOrigin], mrope_types: TypeList[mrope_types.values] = *?, mrope_section: Optional[Coord[mrope_types]] = None](q_proj: ManagedTensorSlice[static_spec=q_proj.static_spec], input_row_offsets: ManagedTensorSlice[static_spec=input_row_offsets.static_spec], kv_collection: PagedKVCacheCollection[cache_dtype, scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], freqs_cis: ManagedTensorSlice[static_spec=freqs_cis.static_spec], position_ids: ManagedTensorSlice[static_spec=position_ids.static_spec], layer_idx: UInt32, output: ManagedTensorSlice[static_spec=output.static_spec], context: DeviceContext)

Applies fused rotary position embedding to Q and K, updating keys in the paged KV cache.

Parameters:

• ​dtype (DType): Element type of the query projection and output tensors (inferred).
• ​freq_dtype (DType): Element type of the freqs_cis RoPE frequency table (inferred).
• ​cache_dtype (DType): Element type of the paged KV cache entries (inferred).
• ​interleaved (Bool): Whether RoPE pairs adjacent real and imaginary components rather than splitting them into halves.
• ​has_position_ids (Bool): Whether per-token position_ids are supplied and used for the rotation.
• ​target (StringSlice[ImmStaticOrigin]): Target backend identifier for kernel dispatch.
• ​mrope_types (TypeList[mrope_types.values]): Coordinate element types for the M-RoPE section split (defaults to an empty TypeList).
• ​mrope_section (Optional[Coord[mrope_types]]): Optional M-RoPE section coordinate describing how the head dimension is partitioned across modalities (defaults to None).

Args:

• ​q_proj (ManagedTensorSlice[static_spec=q_proj.static_spec]): Query projection tensor of shape [batch, seq, head_dim].
• ​input_row_offsets (ManagedTensorSlice[static_spec=input_row_offsets.static_spec]): Cumulative row offsets per batch, shape [batch_size + 1].
• ​kv_collection (PagedKVCacheCollection[cache_dtype, scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): Paged KV cache collection whose keys are read and updated in place.
• ​freqs_cis (ManagedTensorSlice[static_spec=freqs_cis.static_spec]): RoPE frequency table applied to the query and key projections.
• ​position_ids (ManagedTensorSlice[static_spec=position_ids.static_spec]): Per-token position ids used to index freqs_cis.
• ​layer_idx (UInt32): Index of the layer whose keys are updated in the cache.
• ​output (ManagedTensorSlice[static_spec=output.static_spec]): Output tensor for the rotation-applied query projection, shape [batch, seq, head_dim].
• ​context (DeviceContext): Device context for kernel dispatch.