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
def generic_fused_qk_rope_bshd_paged[dtype: DType, //, *, interleaved: Bool, target: StringSpan[ImmStaticOrigin]](q_proj: TileTensor[dtype, Engine=q_proj.Engine, address_space=q_proj.address_space, linear_idx_type=q_proj.linear_idx_type], kv_collection: PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], freqs_cis: TileTensor[dtype, Engine=freqs_cis.Engine, address_space=freqs_cis.address_space, linear_idx_type=freqs_cis.linear_idx_type], layer_idx: UInt32, valid_lengths: TileTensor[.uint32, Engine=valid_lengths.Engine, address_space=valid_lengths.address_space, linear_idx_type=valid_lengths.linear_idx_type], output: TileTensor[dtype, Engine=output.Engine, address_space=output.address_space, linear_idx_type=output.linear_idx_type], context: DeviceContext)
Performs a fused RoPE projection for Q and K with paged KV cache.
Applies RoPE to both Q (returned) and K (in the paged cache) to ensure proper dependency ordering after fused_qkv_padded_matmul.
Args:
- q_proj (
TileTensor[dtype, Engine=q_proj.Engine, address_space=q_proj.address_space, linear_idx_type=q_proj.linear_idx_type]): Query projection tensor of shape [batch, seq_len, n_heads, head_dim]. - kv_collection (
PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): The paged KV cache collection. - freqs_cis (
TileTensor[dtype, Engine=freqs_cis.Engine, address_space=freqs_cis.address_space, linear_idx_type=freqs_cis.linear_idx_type]): Frequency tensor for RoPE of shape [max_seq_len, head_dim]. - layer_idx (
UInt32): The layer index for accessing the correct cache. - valid_lengths (
TileTensor[.uint32, Engine=valid_lengths.Engine, address_space=valid_lengths.address_space, linear_idx_type=valid_lengths.linear_idx_type]): Tensor of shape [batch] containing the valid length for each sequence. RoPE is only applied to positions within these lengths. - output (
TileTensor[dtype, Engine=output.Engine, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Output tensor for Q with RoPE applied, same shape as q_proj. - context (
DeviceContext): Device context pointer for execution.