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: ImmutOrigin]
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.
Fields
- b (
TileTensor[b_type, b_layout, origin]): - b_tile_stack_ptr (
UnsafePointer[Scalar[b_type], MutUntrackedOrigin]): - 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.
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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!