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[Int(2)], b_shape: IndexList[Int(2)], sfa_shape: IndexList[Int(2)], sfb_shape: IndexList[Int(2)], num_pipeline_stages: Int]
Tile payload for block-scaled matmul (A, B, SFA, SFB tiles).
Parameters
- a_type (
DType): Element type of the A operand tiles. - b_type (
DType): Element type of the B operand tiles. - sfa_type (
DType): Element type of the A scale factor (SFA) tiles. - sfb_type (
DType): Element type of the B scale factor (SFB) tiles. - a_shape (
IndexList[Int(2)]): A tile dimensions as(BM, BK). - b_shape (
IndexList[Int(2)]): B tile dimensions as(BN, BK). - sfa_shape (
IndexList[Int(2)]): Dimensions of the A scale factor (SFA) tile. - sfb_shape (
IndexList[Int(2)]): Dimensions of the B scale factor (SFB) tile. - num_pipeline_stages (
Int): Number of pipeline buffer stages.
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[Int(0)], a_shape[Int(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[Int(0)], b_shape[Int(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[Int(0)], a_shape[Int(1)], num_pipeline_stages], b_tiles: SMemTileArray2D[b_type, b_shape[Int(0)], b_shape[Int(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.
Parameters:
- k_group_size (
Int): Number of K-group tiles stored per pipeline stage.
Args:
- stage (
UInt32): Pipeline stage index into the circular buffer. - k_idx (
Int): K-group index within the current pipeline stage.
Returns:
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
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!