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

generic_fused_qk_rope_bshd_paged_ragged

def generic_fused_qk_rope_bshd_paged_ragged[dtype: DType, freq_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: TileTensor[dtype, Storage=q_proj.Storage, linear_idx_type=q_proj.linear_idx_type], input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type], kv_collection: PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], freqs_cis: TileTensor[freq_dtype, Storage=freqs_cis.Storage, linear_idx_type=freqs_cis.linear_idx_type], position_ids: TileTensor[DType.uint32, Storage=position_ids.Storage, linear_idx_type=position_ids.linear_idx_type], layer_idx: UInt32, output: TileTensor[dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type], context: DeviceContext)

Performs a fused RoPE projection for Q and K projections.

We have a manually fused QKV projection with mo.opaque dtypes in our Llama model. Due to a limitation in custom op definitions, we can't declare both a tensor and opaque dtype as output from a custom kernel. This requires us to only note Q_proj as an output from the QKV projection. If we immediately follow the QKV proj kernel with a RoPE kernel applied to K, we'll get a race condition because the graph compiler doesn't know about the dependency between these kernels in the graph definition. Here we fuse the RoPE kernel applied to Q_proj with K_proj, so K_proj RoPE is only executed after QKV completes.

Parameters:

dtype (DType): Data type of the q_proj and output tensors (inferred).
freq_dtype (DType): Data type of the freqs_cis RoPE frequency table (inferred).
interleaved (Bool): Whether RoPE applies interleaved (GPT-NeoX style) rotation to adjacent element pairs.
has_position_ids (Bool): Whether per-token position_ids are provided; when False, position_ids is unused.
target (StringSlice[ImmStaticOrigin]): Target device string for kernel dispatch.
mrope_types (TypeList[mrope_types.values]): TypeList of coordinate element types constraining mrope_section for multimodal RoPE.
mrope_section (Optional[Coord[mrope_types]]): Optional section sizes splitting the head dimension into temporal, height, and width spans for multimodal RoPE (defaults to None).

Args:

q_proj (TileTensor[dtype, Storage=q_proj.Storage, linear_idx_type=q_proj.linear_idx_type]): Query projection tile tensor with shape (sum(seq_lens), num_heads, head_size).
input_row_offsets (TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type]): Tile tensor with shape (batch_size + 1,) denoting the start of each sequence along the ragged sequence dimension.
kv_collection (PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): The paged KV cache collection storing the K cache for this layer, retrieved via layer_idx.
freqs_cis (TileTensor[freq_dtype, Storage=freqs_cis.Storage, linear_idx_type=freqs_cis.linear_idx_type]): Precomputed RoPE frequency table applied to Q and K.
position_ids (TileTensor[DType.uint32, Storage=position_ids.Storage, linear_idx_type=position_ids.linear_idx_type]): Per-token position indices used to index into freqs_cis; ignored when has_position_ids is False.
layer_idx (UInt32): The index of the layer being executed, used to retrieve the K cache from kv_collection.
output (TileTensor[dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): The pre-allocated output tile tensor receiving the rotated Q projection.
context (DeviceContext): The call context pointer, passed by the graph compiler.