Skip to main content

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: NDBuffer[dtype, 2, origin, shape], input_row_offsets: NDBuffer[DType.uint32, 1, origin, shape, strides], weight: NDBuffer[weight_dtype, 2, origin, shape], input_scale: NDBuffer[scale_dtype, 2, origin, shape], weight_scale: NDBuffer[scale_dtype, 2, origin, shape], 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 (NDBuffer): Tensor with shape (sum(seq_lens), num_heads * head_size).
  • input_row_offsets (NDBuffer): Tensor with shape (batch_size + 1,) denoting the start of each sequence along the seq_len dimension.
  • weight (NDBuffer): Tensor with shape (num_heads * head_size, num_kv_heads * head_size).
  • input_scale (NDBuffer): Scale to be multiplied to the input Tensor.
  • weight_scale (NDBuffer): 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?