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

pack_b

def pack_b[transpose_b: Bool, simd_size: Int, inner_size: Int, a_type: DType, b_type: DType, c_type: DType](dst: TileTensor[b_type, Storage=dst.Storage, linear_idx_type=dst.linear_idx_type], src: TileTensor[b_type, Storage=src.Storage, linear_idx_type=src.linear_idx_type], tile_n: Int, tile_k: Int)

Utility function to pack the entire B matrix, such that each [tile_n // inner_size, tile_k, inner_size] tile of src is contiguous in dst.

Tiles (not tile contents) are stored in rowmajor order, so tile[i, j] is tile_n * tile_k bytes away from tile[i, j+1].

Parameters:

transpose_b (Bool): True if the B operand is transposed, stored as [N, K] instead of [K, N].
simd_size (Int): SIMD vector width for b_type.
inner_size (Int): Size of the inner dimension along N in the packed tile; must be a multiple of simd_size.
a_type (DType): Element type of the A operand of the matmul.
b_type (DType): Element type of the B operand being packed.
c_type (DType): Element type of the C output of the matmul.

Args:

dst (TileTensor[b_type, Storage=dst.Storage, linear_idx_type=dst.linear_idx_type]): Pre-allocated mutable buffer that receives the packed B matrix.
src (TileTensor[b_type, Storage=src.Storage, linear_idx_type=src.linear_idx_type]): Read-only buffer containing the original B matrix to pack.
tile_n (Int): Tile size along the N dimension of the matmul.
tile_k (Int): Tile size along the K dimension of the matmul.