Mojo function

mla_prefill_plan_kernel

mla_prefill_plan_kernel[buffer_row_offsets_shape: DimList, buffer_row_offsets_stride: DimList, cache_offsets_shape: DimList, cache_offsets_stride: DimList, buffer_lengths_shape: DimList, buffer_lengths_stride: DimList, input_row_offsets_shape: DimList, input_row_offsets_stride: DimList, cache_t: KVCacheT](buffer_row_offsets: NDBuffer[DType.uint32, 2, MutableAnyOrigin, buffer_row_offsets_shape, buffer_row_offsets_stride], cache_offsets: NDBuffer[DType.uint32, 2, MutableAnyOrigin, cache_offsets_shape, cache_offsets_stride], buffer_lengths: NDBuffer[DType.int32, 1, MutableAnyOrigin, buffer_lengths_shape, buffer_lengths_stride], input_row_offsets: NDBuffer[DType.uint32, 1, MutableAnyOrigin, input_row_offsets_shape, input_row_offsets_stride], k_cache: cache_t, buffer_token_size: UInt32)