Mojo function

unfused_qkv_matmul_ragged_continuous_batching_gguf_quantized

unfused_qkv_matmul_ragged_continuous_batching_gguf_quantized[type: DType, num_heads: Int, head_dim: Int, //, quantization_encoding_q: StringSlice[StaticConstantOrigin], quantization_encoding_k: StringSlice[StaticConstantOrigin], quantization_encoding_v: StringSlice[StaticConstantOrigin]](hidden_state: NDBuffer[float32, 2, origin, shape], input_row_offsets: NDBuffer[uint32, 1, origin, shape, strides], q_weight: NDBuffer[uint8, 2, origin, shape], k_weight: NDBuffer[uint8, 2, origin, shape], v_weight: NDBuffer[uint8, 2, origin, shape], kv_collection: ContinuousBatchingKVCacheCollection[type, KVCacheStaticParams(UInt(num_heads), UInt(head_dim))], layer_idx: SIMD[uint32, 1], output: NDBuffer[float32, 2, origin, shape], ctx: DeviceContextPtr)

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

Unlike the un-quantized version (kv_matmul_ragged_continuous_batching), this implementation does not concat the q, k, and v weights together. Instead, it performs three matmuls. This allows the q, k, and v weights to have different quantization encodings.

This is only supported on CPU.

Args:

hidden_state (NDBuffer[float32, 2, origin, shape]): 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.
q_weight (NDBuffer[uint8, 2, origin, shape]): Tensor with shape (num_heads * head_size, num_kv_heads * head_size).
k_weight (NDBuffer[uint8, 2, origin, shape]): Tensor with shape (num_heads * head_size, num_kv_heads * head_size).
v_weight (NDBuffer[uint8, 2, origin, shape]): Tensor with shape (num_heads * head_size, num_kv_heads * head_size).
kv_collection (ContinuousBatchingKVCacheCollection[type, KVCacheStaticParams(UInt(num_heads), UInt(head_dim))]): The Collection object storing KVCache entries.
layer_idx (SIMD[uint32, 1]): The index of the layer being executed. Used to retrieve the KVCache for the given layer from kv_collection.
output (NDBuffer[float32, 2, origin, shape]): Tensor with shape (sum(seq_lens), num_kv_heads * head_size). This is the output buffer for the Q matmul.
ctx (DeviceContextPtr): The call context pointer, passed by the graph compiler.