Mojo function
k_matmul_ragged_paged
k_matmul_ragged_paged[type: DType, num_heads: Int, head_dim: Int, page_size: Int, //, target: StringSlice[StaticConstantOrigin]](hidden_state: NDBuffer[type, 2, origin, shape, strides], input_row_offsets: NDBuffer[uint32, 1, origin, shape, strides], weight: NDBuffer[type, 2, origin, shape, strides], kv_collection: PagedKVCacheCollection[type, KVCacheStaticParams(UInt(num_heads), UInt(head_dim)), page_size], layer_idx: SIMD[uint32, 1], ctx: DeviceContextPtr)
Performs a matmul, writing the output into a mutable PagedKVCacheCollection object.
Args:
- hidden_state (
NDBuffer[type, 2, origin, shape, strides]
): Tensor with shape (sum(seq_lens), num_heads * head_size). - input_row_offsets (
NDBuffer[uint32, 1, origin, shape, strides]
): Tensor with shape (batch_size + 1,) denoting the start of each sequence along the seq_len dimension. - weight (
NDBuffer[type, 2, origin, shape, strides]
): Tensor with shape (num_heads * head_size, num_kv_heads * head_size). - kv_collection (
PagedKVCacheCollection[type, KVCacheStaticParams(UInt(num_heads), UInt(head_dim)), page_size]
): The historical KVCache for keys and values. The KVCache for this layer is retrieved via layer_idx. - layer_idx (
SIMD[uint32, 1]
): The index of the layer being executed. Used to retrieve the KVCache for the given layer from kv_collection. - ctx (
DeviceContextPtr
): The call context pointer, passed by the graph compiler.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!