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

BlockwiseFP8TilePayload

struct BlockwiseFP8TilePayload[a_type: DType, b_type: DType, a_scales_type: DType, a_shape: IndexList[Int(2)], b_shape: IndexList[Int(2)], a_scales_shape: IndexList[Int(2)], num_pipeline_stages: Int]

Tile payload for blockwise FP8 matmul (A, B, A-scales tiles).

Unlike BlockScaledTilePayload, this only stores A-scales in SMEM. B-scales are read directly from global memory during the epilogue phase.

Parameters

• ​a_type (DType): Element type of the A operand tiles.
• ​b_type (DType): Element type of the B operand tiles.
• ​a_scales_type (DType): Element type of the A-scales tiles.
• ​a_shape (IndexList[Int(2)]): A tile dimensions as (BM, BK).
• ​b_shape (IndexList[Int(2)]): B tile dimensions as (BN, BK).
• ​a_scales_shape (IndexList[Int(2)]): A-scales tile dimensions as (1, BM).
• ​num_pipeline_stages (Int): Number of pipeline buffer stages.

Fields

• ​a_tiles (BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_shape, b_shape, a_scales_shape, num_pipeline_stages].ATileArray):
• ​b_tiles (BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_shape, b_shape, a_scales_shape, num_pipeline_stages].BTileArray):
• ​a_scales_tiles (BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_shape, b_shape, a_scales_shape, num_pipeline_stages].AScalesTileArray):

Implemented traits

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

comptime members

AScalesTile

comptime AScalesTile = BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_shape, b_shape, a_scales_shape, num_pipeline_stages].AScalesTileArray.Tile

AScalesTileArray

comptime AScalesTileArray = SMemTileArray2DRowMajor[a_scales_type, a_scales_shape[Int(0)], a_scales_shape[Int(1)], num_pipeline_stages]

ATile

comptime ATile = BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_shape, b_shape, a_scales_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 = BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_shape, b_shape, a_scales_shape, num_pipeline_stages].BTileArray.Tile

BTileArray

comptime BTileArray = SMemTileArray2D[b_type, b_shape[Int(0)], b_shape[Int(1)], 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], a_scales_tiles: SMemTileArray2DRowMajor[a_scales_type, a_scales_shape[Int(0)], a_scales_shape[Int(1)], 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[a_scales_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]]

Get A, B, A-scales 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[a_scales_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.

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:

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.

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:

Self.BTile

get_a_scales_tile

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

Get A-scales tile 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:

Self.AScalesTile