Mojo struct
BlockwiseFP8TilePayload
@register_passable(trivial)
struct BlockwiseFP8TilePayload[a_type: DType, b_type: DType, a_scales_type: DType, a_dim0: Int, a_dim1: Int, b_dim0: Int, b_dim1: Int, a_scales_dim0: Int, a_scales_dim1: Int, 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.
Fields
- a_tiles (
BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].ATileArray): - b_tiles (
BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].BTileArray): - a_scales_tiles (
BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].AScalesTileArray):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TilePayload,
TrivialRegisterPassable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
AScalesTile
comptime AScalesTile = BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].AScalesTileArray.Tile
AScalesTileArray
comptime AScalesTileArray = SMemTileArray2D[a_scales_type, a_scales_dim0, a_scales_dim1, num_pipeline_stages]
ATile
comptime ATile = BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].ATileArray.Tile
ATileArray
comptime ATileArray = SMemTileArray2D[a_type, a_dim0, a_dim1, num_pipeline_stages]
BTile
comptime BTile = BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].BTileArray.Tile
BTileArray
comptime BTileArray = SMemTileArray2D[b_type, b_dim0, b_dim1, num_pipeline_stages]
Methods
__init__
__init__(a_tiles: SMemTileArray2D[a_type, a_dim0, a_dim1, num_pipeline_stages], b_tiles: SMemTileArray2D[b_type, b_dim0, b_dim1, num_pipeline_stages], a_scales_tiles: SMemTileArray2D[a_scales_type, a_scales_dim0, a_scales_dim1, num_pipeline_stages]) -> Self
get_tile
get_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Tuple[BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].ATile, BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].BTile, BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].AScalesTile]
Get A, B, A-scales tiles at the specified stage and k-group index.
Returns:
get_a_tile
get_a_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].ATile
Get A tile at the specified stage and k-group index.
Returns:
BlockwiseFP8TilePayload
get_b_tile
get_b_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].BTile
Get B tile at the specified stage and k-group index.
Returns:
BlockwiseFP8TilePayload
get_a_scales_tile
get_a_scales_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> BlockwiseFP8TilePayload[a_type, b_type, a_scales_type, a_dim0, a_dim1, b_dim0, b_dim1, a_scales_dim0, a_scales_dim1, num_pipeline_stages].AScalesTile
Get A-scales tile at the specified stage and k-group index.
Returns:
BlockwiseFP8TilePayload
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!