Mojo function
k_matmul_ragged_paged_scale
k_matmul_ragged_paged_scale[dtype: DType, weight_dtype: DType, scale_dtype: DType, target: StringSlice[StaticConstantOrigin], scales_granularity_mnk: IndexList[3]](hidden_state: LayoutTensor[dtype, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], input_row_offsets: LayoutTensor[DType.uint32, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], weight: LayoutTensor[weight_dtype, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], input_scale: LayoutTensor[scale_dtype, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], weight_scale: LayoutTensor[scale_dtype, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], kv_collection: PagedKVCacheCollection[dtype_, kv_params_, page_size], layer_idx: UInt32, ctx: DeviceContextPtr)
Performs a matmul, writing the output into a mutable PagedKVCacheCollection 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). - input_scale (
LayoutTensor
): Scale to be multiplied to the input Tensor. - weight_scale (
LayoutTensor
): Scale to be multiplied to the weight Tensor. - 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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!