Mojo function

generic_flare_mla_prefill_ragged_paged_plan

generic_flare_mla_prefill_ragged_paged_plan[target: StringSlice[StaticConstantOrigin]](input_row_offsets: LayoutTensor[DType.uint32, input_row_offsets.layout, input_row_offsets.origin, element_layout=input_row_offsets.element_layout, layout_int_type=input_row_offsets.layout_int_type, linear_idx_type=input_row_offsets.linear_idx_type, masked=input_row_offsets.masked, alignment=input_row_offsets.alignment], kv_collection: PagedKVCacheCollection[kv_collection.dtype_, kv_collection.kv_params_, kv_collection.page_size, kv_collection.scale_dtype_, kv_collection.quantization_granularity_], layer_idx: UInt32, buffer_token_size: UInt32, buffer_row_offsets: LayoutTensor[DType.uint32, buffer_row_offsets.layout, buffer_row_offsets.origin, element_layout=buffer_row_offsets.element_layout, layout_int_type=buffer_row_offsets.layout_int_type, linear_idx_type=buffer_row_offsets.linear_idx_type, masked=buffer_row_offsets.masked, alignment=buffer_row_offsets.alignment], cache_offsets: LayoutTensor[DType.uint32, cache_offsets.layout, cache_offsets.origin, element_layout=cache_offsets.element_layout, layout_int_type=cache_offsets.layout_int_type, linear_idx_type=cache_offsets.linear_idx_type, masked=cache_offsets.masked, alignment=cache_offsets.alignment], buffer_lengths: LayoutTensor[DType.int32, buffer_lengths.layout, buffer_lengths.origin, element_layout=buffer_lengths.element_layout, layout_int_type=buffer_lengths.layout_int_type, linear_idx_type=buffer_lengths.linear_idx_type, masked=buffer_lengths.masked, alignment=buffer_lengths.alignment], context: DeviceContextPtr)