Mojo function

generic_fused_qk_rope_bshd_paged_ragged

generic_fused_qk_rope_bshd_paged_ragged[dtype: DType, freq_dtype: DType, //, *, interleaved: Bool, has_position_ids: Bool, target: StringSlice[StaticConstantOrigin], mrope_types: TypeList[mrope_types.values] = *?, mrope_section: Optional[Coord[mrope_types]] = None](q_proj: TileTensor[dtype, linear_idx_type=q_proj.linear_idx_type, element_size=q_proj.element_size], input_row_offsets: TileTensor[DType.uint32, linear_idx_type=input_row_offsets.linear_idx_type, element_size=input_row_offsets.element_size], kv_collection: PagedKVCacheCollection, freqs_cis: TileTensor[freq_dtype, linear_idx_type=freqs_cis.linear_idx_type, element_size=freqs_cis.element_size], position_ids: TileTensor[DType.uint32, linear_idx_type=position_ids.linear_idx_type, element_size=position_ids.element_size], layer_idx: UInt32, output: TileTensor[dtype, linear_idx_type=output.linear_idx_type, element_size=output.element_size], context: DeviceContextPtr = DeviceContextPtr())

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.