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

matmul_gpu_qint4

def matmul_gpu_qint4[c_type: DType, a_type: DType, //, group_size: Int, target: StringSlice[ImmStaticOrigin], elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c_tt: TileTensor[c_type, Storage=c_tt.Storage, linear_idx_type=c_tt.linear_idx_type], a_tt: TileTensor[a_type, Storage=a_tt.Storage, linear_idx_type=a_tt.linear_idx_type], b_tt: TileTensor[DType.uint8, Storage=b_tt.Storage, linear_idx_type=b_tt.linear_idx_type], ctx: Optional[DeviceContext] = None)

Launches a GPU int4 quantized matrix multiplication for the given tile tensors.

Constraints:

Requires an NVIDIA GPU target. a_type and c_type must both be bfloat16.

Parameters:

c_type (DType): The dtype of the output matrix.
a_type (DType): The dtype of the A matrix elements.
group_size (Int): The number of K elements sharing a single scale.
target (StringSlice[ImmStaticOrigin]): The target platform string, which must identify a GPU.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): An optional elementwise epilogue applied per output element.

Args:

c_tt (TileTensor[c_type, Storage=c_tt.Storage, linear_idx_type=c_tt.linear_idx_type]): The output tile tensor in global memory.
a_tt (TileTensor[a_type, Storage=a_tt.Storage, linear_idx_type=a_tt.linear_idx_type]): The left-hand (activation) tile tensor in global memory.
b_tt (TileTensor[DType.uint8, Storage=b_tt.Storage, linear_idx_type=b_tt.linear_idx_type]): The packed quantized weight tile tensor in global memory.
ctx (Optional[DeviceContext]): The device context used to enqueue the kernel.

Raises:

An error if the input tensors are not rank-2 or the target is not a GPU.