Mojo function

fused_qk_rope_ragged

fused_qk_rope_ragged[type: DType, collection_t: KVCollectionT, //, cache_t: KVCacheT, *, interleaved: Bool, target: StringSlice[StaticConstantOrigin]](q_proj: NDBuffer[type, 3, origin, shape, strides], input_row_offsets: NDBuffer[uint32, 1, origin, shape, strides], kv_collection: collection_t, freqs_cis: NDBuffer[type, 2, origin, shape, strides], layer_idx: SIMD[uint32, 1], output: NDBuffer[type, 3, origin, shape, strides], context: Optional[DeviceContext])

Applies RoPE (Rotary Position Embedding) to query and key tensors.

This function can applies RoPE only to the last rope_dim elements of each head, leaving the first unroped_dim elements unchanged. This is required for DeepSeek models where only part of each head undergoes rotary transformation.