IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

BlockScaledTilePayload

struct BlockScaledTilePayload[a_type: DType, b_type: DType, sfa_type: DType, sfb_type: DType, a_shape: IndexList[2], b_shape: IndexList[2], sfa_shape: IndexList[2], sfb_shape: IndexList[2], num_pipeline_stages: Int]

Tile payload for block-scaled matmul (A, B, SFA, SFB tiles).

Fields​

  • ​a_tiles (BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].ATileArray):
  • ​b_tiles (BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].BTileArray):
  • ​sfa_tiles (BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].SFATileArray):
  • ​sfb_tiles (BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].SFBTileArray):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TilePayload, TrivialRegisterPassable

comptime members​

ATile​

comptime ATile = BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].ATileArray.Tile

ATileArray​

comptime ATileArray = SMemTileArray2D[a_type, a_shape[0], a_shape[1], num_pipeline_stages]

BTile​

comptime BTile = BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].BTileArray.Tile

BTileArray​

comptime BTileArray = SMemTileArray2D[b_type, b_shape[0], b_shape[1], num_pipeline_stages]

sfa_layout​

comptime sfa_layout = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(Coord(ComptimeInt(), ComptimeInt()), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(Coord(ComptimeInt(), ComptimeInt()), ComptimeInt())))

SFATile​

comptime SFATile = BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].SFATileArray.Tile

SFATileArray​

comptime SFATileArray = SMemTileArrayWithLayout[sfa_type, BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].sfa_layout, num_pipeline_stages]

sfb_layout​

comptime sfb_layout = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(Coord(ComptimeInt(), ComptimeInt()), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(Coord(ComptimeInt(), ComptimeInt()), ComptimeInt())))

SFBTile​

comptime SFBTile = BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].SFBTileArray.Tile

SFBTileArray​

comptime SFBTileArray = SMemTileArrayWithLayout[sfb_type, BlockScaledTilePayload[a_type, b_type, sfa_type, sfb_type, a_shape, b_shape, sfa_shape, sfb_shape, num_pipeline_stages].sfb_layout, num_pipeline_stages]

Methods​

__init__​

def __init__(a_tiles: SMemTileArray2D[a_type, a_shape[0], a_shape[1], num_pipeline_stages], b_tiles: SMemTileArray2D[b_type, b_shape[0], b_shape[1], num_pipeline_stages], sfa_tiles: SMemTileArrayWithLayout[sfa_type, Self.sfa_layout, num_pipeline_stages], sfb_tiles: SMemTileArrayWithLayout[sfb_type, Self.sfb_layout, num_pipeline_stages]) -> Self

get_tile​

def get_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Tuple[TileTensor[a_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED], TileTensor[b_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED], TileTensor[sfa_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED], TileTensor[sfb_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]]

Get A, B, SFA, SFB tiles at the specified stage and k-group index.

Returns:

Tuple[TileTensor[a_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED], TileTensor[b_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED], TileTensor[sfa_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED], TileTensor[sfb_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]]

get_a_tile​

def get_a_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Self.ATile

Get A tile at the specified stage and k-group index.

Returns:

Self.ATile

get_b_tile​

def get_b_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Self.BTile

Get B tile at the specified stage and k-group index.

Returns:

Self.BTile

get_sfa_tile​

def get_sfa_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Self.SFATile

Get SFA tile at the specified stage and k-group index.

Returns:

Self.SFATile

get_sfb_tile​

def get_sfb_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Self.SFBTile

Get SFB tile at the specified stage and k-group index.

Returns:

Self.SFBTile

Models

Cookbook

Puzzles

Blog

Community

Contact