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

k_matmul_ragged_paged_scale

def k_matmul_ragged_paged_scale[dtype: DType, weight_dtype: DType, scale_dtype: DType, target: StringSlice[ImmStaticOrigin], scales_granularity_mnk: IndexList[Int(3)]](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], input_scale: LayoutTensor[scale_dtype, element_layout=input_scale.element_layout, layout_int_type=input_scale.layout_int_type, linear_idx_type=input_scale.linear_idx_type, masked=input_scale.masked, alignment=input_scale.alignment], weight_scale: LayoutTensor[scale_dtype, element_layout=weight_scale.element_layout, layout_int_type=weight_scale.layout_int_type, linear_idx_type=weight_scale.linear_idx_type, masked=weight_scale.masked, alignment=weight_scale.alignment], kv_collection: PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], layer_idx: UInt32, ctx: DeviceContext)

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

Parameters:

dtype (DType): Element type of the hidden_state input tensor and of the key entries written to the KV cache.
weight_dtype (DType): Element type of the weight tensor; must match dtype.
scale_dtype (DType): Element type of the input_scale and weight_scale tensors.
target (StringSlice[ImmStaticOrigin]): Compilation target string; must be a GPU target.
scales_granularity_mnk (IndexList[Int(3)]): Block sizes along the M, N, and K matmul dimensions used to tile the scale application; -1 selects per-tensor scaling, 1 selects per-channel scaling, and any other value selects blockwise scaling.

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,) denoting the start of each sequence along the seq_len dimension.
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).
input_scale (LayoutTensor[scale_dtype, element_layout=input_scale.element_layout, layout_int_type=input_scale.layout_int_type, linear_idx_type=input_scale.linear_idx_type, masked=input_scale.masked, alignment=input_scale.alignment]): Scale to be multiplied to the input Tensor.
weight_scale (LayoutTensor[scale_dtype, element_layout=weight_scale.element_layout, layout_int_type=weight_scale.layout_int_type, linear_idx_type=weight_scale.linear_idx_type, masked=weight_scale.masked, alignment=weight_scale.alignment]): Scale to be multiplied to the weight Tensor.
kv_collection (PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): 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 (DeviceContext): The call context pointer, passed by the graph compiler.