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 struct

BTileGenerator

struct BTileGenerator[config: KernelConfig, a_type: DType, b_type: DType, c_type: DType, b_layout: TensorLayout, transpose_b: Bool, b_packed: Bool, origin: ImmOrigin]

Struct to encapsulate a tile of B that supports prepacking.

If b_packed is true, calls to get_tile will return a buffer view from B. Otherwise, calls to get_tile will copy a tile from B into a stack allocated scratch buffer and return a view of that.

Parameters

config (KernelConfig): Kernel configuration supplying simd_size and kernel_cols used by the packing routines.
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.
b_layout (TensorLayout): Layout of the B TileTensor.
transpose_b (Bool): True if the B operand is transposed, stored as [N, K].
b_packed (Bool): True if B is already pre-packed into the mlas layout.
origin (ImmOrigin): Origin of the B TileTensor.

Fields

b (TileTensor[b_type, b_layout, origin]):
b_tile_stack_ptr (Pointer[Scalar[b_type], MutUntrackedOrigin, _safe=False]):
tile_n_k (IndexList[Int(2)]):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

`comptime` members

`PackedTileLayout`

comptime PackedTileLayout = Layout[*?, *?]

Methods

`get`

static def get(b: TileTensor[b_type, b_layout, origin], tile_n_k: IndexList[Int(2)]) -> Self

`get_tile`

def get_tile[inner_size: Int](self, global_offset: GemmShape, tile_dim_nk: IndexList[Int(2)], valid_data_dim_nk: IndexList[Int(2)]) -> TileTensor[b_type, Layout[*?, *?], ImmutAnyOrigin]

Get a packed matrix (B) tile.

valid_data_tile_nk is ignored for pre-packing, where the tile is padded to have shape of tile_dim_nk.

Parameters:

inner_size (Int): Size of the inner dimension along N in the packed tile; must be a multiple of simd_size.

Args:

global_offset (GemmShape): Offset in the global M, N, K dimensions.
tile_dim_nk (IndexList[Int(2)]): Tile shape based on cache size and matrix dimensions.
valid_data_dim_nk (IndexList[Int(2)]): The upper bounds for N and K dimensions.

Returns:

TileTensor[b_type, Layout[*?, *?], ImmutAnyOrigin]: A view of the packed tile.

Parameters​

Fields​

Implemented traits​

comptime members​

PackedTileLayout​

Methods​

get​

get_tile​