Mojo function

kv_matmul_ragged_paged

kv_matmul_ragged_paged[dtype: DType, params: KVCacheStaticParams, page_size: Int, //, target: StringSlice[StaticConstantOrigin]](hidden_state: LayoutTensor[dtype, hidden_state.layout, hidden_state.origin, element_layout=hidden_state.element_layout, layout_int_type=hidden_state.layout_int_type, linear_idx_type=hidden_state.linear_idx_type, masked=hidden_state.masked, alignment=hidden_state.alignment], 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], weight: LayoutTensor[dtype, weight.layout, weight.origin, element_layout=weight.element_layout, layout_int_type=weight.layout_int_type, linear_idx_type=weight.linear_idx_type, masked=weight.masked, alignment=weight.alignment], kv_collection: PagedKVCacheCollection[dtype, params, page_size], layer_idx: UInt32, ctx: DeviceContextPtr)

Performs a matmul, writing the output into a mutable ContinuousBatchingKVCacheCollection object.

Args:

hidden_state (LayoutTensor): Tensor with shape (sum(seq_lens), num_heads * head_size).
input_row_offsets (LayoutTensor): Tensor with shape (batch_size + 1,) denoting the start of each sequence along the seq_len dimension.
weight (LayoutTensor): Tensor with shape (num_heads * head_size, num_kv_heads * head_size).
kv_collection (PagedKVCacheCollection): The historical KVCache for keys and values. The KVCache for this layer is retrieved via layer_idx.
layer_idx (UInt32): 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.