IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /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, Deinitable, ImplicitlyCopyable, 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:

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