Mojo struct
BlockwiseFP8TileStorage
struct BlockwiseFP8TileStorage[a_type: DType, b_type: DType, c_type: DType, a_scales_type: DType, a_dim0: Int, a_dim1: Int, b_dim0: Int, b_dim1: Int, c_dim0: Int, c_dim1: Int, a_scales_dim0: Int, a_scales_dim1: Int, num_pipeline_stages: Int, num_output_stages: Int]
Storage for blockwise FP8 matmul tiles (A, B, C, A-scales).
Single source of truth for blockwise FP8 tile arrays and storage. B-scales are read directly from global memory during epilogue.
All tiles use TileTensor natively. LayoutTensor conversion is kept for SMemEpilogueWriter compatibility.
IMPORTANT: Field order preserves SMEM layout compatibility: a, b, c, a_scales.
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. - a_scales_type (
DType): Data type for A scale tiles. - a_dim0 (
Int): First dimension for A tiles. - a_dim1 (
Int): Second dimension for A tiles. - b_dim0 (
Int): First dimension for B tiles. - b_dim1 (
Int): Second dimension for B tiles. - c_dim0 (
Int): First dimension for C tiles (OutputM). - c_dim1 (
Int): Second dimension for C tiles (OutputN). - a_scales_dim0 (
Int): First dimension for A scale tiles. - a_scales_dim1 (
Int): Second dimension for A scale tiles. - num_pipeline_stages (
Int): Number of input pipeline stages. - num_output_stages (
Int): Number of output pipeline stages.
Fields
- a_tiles_storage (
BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].ATileArray.Storage): - b_tiles_storage (
BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].BTileArray.Storage): - c_tiles_storage (
BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].CTileArray.Storage): - a_scales_tiles_storage (
BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].AScalesTileArray.Storage):
Implemented traits
AnyType,
ImplicitlyDestructible
comptime members
__del__is_trivial
comptime __del__is_trivial = True
AScalesTileArray
comptime AScalesTileArray = SMemTileArray2DRowMajor[a_scales_type, a_scales_dim0, a_scales_dim1, num_pipeline_stages]
ATileArray
comptime ATileArray = SMemTileArray2D[a_type, a_dim0, a_dim1, num_pipeline_stages]
BTileArray
comptime BTileArray = SMemTileArray2D[b_type, b_dim0, b_dim1, 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]
Methods
a_tiles
a_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].ATileArray
Get A tile array accessor.
Returns:
BlockwiseFP8TileStorage
b_tiles
b_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].BTileArray
Get B tile array accessor.
Returns:
BlockwiseFP8TileStorage
c_tiles
c_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].CTileArray
Get C tile array accessor.
Returns:
BlockwiseFP8TileStorage
a_scales_tiles
a_scales_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockwiseFP8TileStorage[a_type, b_type, c_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages, num_output_stages].AScalesTileArray
Get A-scales tile array accessor.
Returns:
BlockwiseFP8TileStorage
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!