For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

generic_fused_qkv_matmul_kv_cache_paged_ragged

def generic_fused_qkv_matmul_kv_cache_paged_ragged[dtype: DType, weight_dtype: DType, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), group_size: Optional[Int] = None, has_zp: Optional[Bool] = None](hidden_state: LayoutTensor[dtype, 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, 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[weight_dtype, 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[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], layer_idx: UInt32, output: LayoutTensor[dtype, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], ctx: DeviceContext)

Performs a fused QKV matmul. Q outputs are written to the output argument while K and V outputs are written in-place into k_cache and v_cache.

Parameters:

dtype (DType): Element type of the hidden_state input and output tensors.
weight_dtype (DType): Element type of the weight tensor.
target (StringSlice[ImmStaticOrigin]): Compilation target string used to dispatch GPU versus CPU paths (defaults to "cpu").
group_size (Optional[Int]): Block size for GPTQ-style quantization of weight; when set, weight must be uint8 (defaults to None for no quantization).
has_zp (Optional[Bool]): Whether the weight quantization uses a zero point; currently must be falsy when group_size is set (defaults to None).

Args:

hidden_state (LayoutTensor[dtype, 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]): Tensor with shape (sum(seq_lens), num_heads * head_size).
input_row_offsets (LayoutTensor[DType.uint32, 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]): Tensor with shape (batch_size + 1). The value at each index is the start_idx of the corresponding batch in hidden_state.
weight (LayoutTensor[weight_dtype, 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]): Tensor with shape (num_heads * head_size, num_kv_heads * head_size).
kv_collection (PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): The object storing the KVCache for this layer.
layer_idx (UInt32): The current layer, used to retrieve the KVCache object from kv_collection.
output (LayoutTensor[dtype, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment]): The pre-allocated output buffer for Q projections. K and V projections are written in-place to k_cache and v_cache. Shape: (sum(seq_lens), num_heads * head_size).
ctx (DeviceContext): The call context pointer, passed by the graph compiler.