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
BlockScaledTileStorage
struct BlockScaledTileStorage[a_type: DType, b_type: DType, c_type: DType, sfa_type: DType, sfb_type: DType, a_shape: IndexList[Int(2)], b_shape: IndexList[Int(2)], c_dim0: Int, c_dim1: Int, sfa_shape: IndexList[Int(2)], sfb_shape: IndexList[Int(2)], num_pipeline_stages: Int, num_output_stages: Int]
Storage for block-scaled matmul tiles (A, B, C, SFA, SFB).
Single source of truth for block-scaled tile arrays and storage.
All tiles use TileTensor natively. LayoutTensor conversion is kept for SMemEpilogueWriter compatibility.
IMPORTANT: Field order preserves SMEM layout compatibility: a, b, c, sfa, sfb.
Parameters
- a_type (
DType): Data type for A matrix tiles. - b_type (
DType): Data type for B matrix tiles. - c_type (
DType): Data type for C matrix tiles. - sfa_type (
DType): Data type for A scale factor tiles. - sfb_type (
DType): Data type for B scale factor tiles. - a_shape (
IndexList[Int(2)]): Tile shape for A tiles as (rows, cols). - b_shape (
IndexList[Int(2)]): Tile shape for B tiles as (rows, cols). - c_dim0 (
Int): First dimension for C tiles (OutputM). - c_dim1 (
Int): Second dimension for C tiles (OutputN). - sfa_shape (
IndexList[Int(2)]): Tile shape for SFA tiles. - sfb_shape (
IndexList[Int(2)]): Tile shape for SFB tiles. - num_pipeline_stages (
Int): Number of input pipeline stages. - num_output_stages (
Int): Number of output pipeline stages.
Fields
- a_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_stages].ATileArray.Storage): - b_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_stages].BTileArray.Storage): - c_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_stages].CTileArray.Storage): - sfa_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_stages].SFATileArray.Storage): - sfb_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_stages].SFBTileArray.Storage):
Implemented traits
comptime members
ATileArray
comptime ATileArray = SMemTileArray2D[a_type, a_shape[Int(0)], a_shape[Int(1)], num_pipeline_stages]
BTileArray
comptime BTileArray = SMemTileArray2D[b_type, b_shape[Int(0)], b_shape[Int(1)], num_pipeline_stages]
c_tile_layout
comptime c_tile_layout = Layout.row_major(c_dim0, c_dim1)
CTileArray
comptime CTileArray = SMemTileArray2DRowMajor[c_type, c_dim0, c_dim1, num_output_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())))
SFATileArray
comptime SFATileArray = SMemTileArrayWithLayout[sfa_type, BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_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())))
SFBTileArray
comptime SFBTileArray = SMemTileArrayWithLayout[sfb_type, BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_shape, b_shape, c_dim0, c_dim1, sfa_shape, sfb_shape, num_pipeline_stages, num_output_stages].sfb_layout, num_pipeline_stages]
Methods
a_tiles
def a_tiles(ref[AddressSpace._value] self) -> Self.ATileArray
Get A tile array accessor.
Returns:
Self.ATileArray
b_tiles
def b_tiles(ref[AddressSpace._value] self) -> Self.BTileArray
Get B tile array accessor.
Returns:
Self.BTileArray
c_tiles
def c_tiles(ref[AddressSpace._value] self) -> Self.CTileArray
Get C tile array accessor.
Returns:
Self.CTileArray
sfa_tiles
def sfa_tiles(ref[AddressSpace._value] self) -> Self.SFATileArray
Get SFA tile array accessor.
Returns:
Self.SFATileArray
sfb_tiles
def sfb_tiles(ref[AddressSpace._value] self) -> Self.SFBTileArray
Get SFB tile array accessor.
Returns:
Self.SFBTileArray