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
BlockScaledSmem
struct BlockScaledSmem[a_type: DType, b_type: DType, c_type: DType, sfa_dtype: DType, sfb_dtype: DType, transpose_b: Bool, *, config: BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]]
SMEM struct for block-scaled matmul with CLC scheduler pipeline.
Thin wrapper over BlockScaledTileCore + SmemPipelineBundle.
Fieldsβ
- βcore (
BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core): - βpipelines (
BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Pipelines):
Implemented traitsβ
comptime membersβ
Coreβ
comptime Core = BlockScaledTileCore[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config]
Pipelinesβ
comptime Pipelines = SmemPipelineBundle[(config // config), config.num_accum_pipeline_stages, config.num_clc_pipeline_stages, BlockScaledTilePayload[a_type, b_type, sfa_dtype, sfb_dtype, IndexList(config.block_tile_shape[Int(0)], config.block_tile_shape[Int(2)], __list_literal__=NoneType(None)), IndexList(config.block_tile_shape[Int(1)], config.block_tile_shape[Int(2)], __list_literal__=NoneType(None)), IndexList(sfa_dim0[config](), sfa_dim1[config](), __list_literal__=NoneType(None)), IndexList(sfb_dim0[config](), sfb_dim1[config](), __list_literal__=NoneType(None)), config.num_pipeline_stages]]
Methodsβ
a_tilesβ
def a_tiles(ref[AddressSpace._value] self) -> Self.Core.ATileArray
Get A tile array accessor.
Returns:
Self.Core.ATileArray
b_tilesβ
def b_tiles(ref[AddressSpace._value] self) -> Self.Core.BTileArray
Get B tile array accessor.
Returns:
Self.Core.BTileArray
c_tilesβ
def c_tiles(ref[AddressSpace._value] self) -> Self.Core.CTileArray
Get C tile array accessor.
Returns:
Self.Core.CTileArray
sfa_tilesβ
def sfa_tiles(ref[AddressSpace._value] self) -> Self.Core.SFATileArray
Get SFA tile array accessor.
Returns:
Self.Core.SFATileArray
sfb_tilesβ
def sfb_tiles(ref[AddressSpace._value] self) -> Self.Core.SFBTileArray
Get SFB tile array accessor.
Returns:
Self.Core.SFBTileArray
ab_pipeline_sizeβ
static def ab_pipeline_size() -> Int
Total size of A+B tiles for all pipeline stages (in elements).
Returns:
sf_pipeline_sizeβ
static def sf_pipeline_size() -> Int
Total size of SFA+SFB tiles for all pipeline stages (in elements).
Returns:
c_output_sizeβ
total_tile_sizeβ
static def total_tile_size() -> Int
Total tile storage size (A+B+SFA+SFB+C) in elements.
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!