Mojo struct
ScalingFactorLoader
@register_passable(trivial)
struct ScalingFactorLoader[tma_origin: ImmutOrigin, sf_dtype: DType, sf_layout: Layout, sf_desc_layout: Layout, /, *, cta_group: Int, BM_or_MMA_N: Int]
TMA loader for 5D MXFP8 scaling factor tensors (SFA or SFB).
Fields
- tma_op (
ScalingFactorLoader[tma_origin, sf_dtype, sf_layout, sf_desc_layout, cta_group=cta_group, BM_or_MMA_N=BM_or_MMA_N].TmaOpPtr):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
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
TmaOp
comptime TmaOp = TMATensorTile[sf_dtype, sf_layout, sf_desc_layout]
TmaOpPtr
comptime TmaOpPtr = Pointer[ScalingFactorLoader[tma_origin, sf_dtype, sf_layout, sf_desc_layout, cta_group=cta_group, BM_or_MMA_N=BM_or_MMA_N].TmaOp, tma_origin]
Methods
__init__
__init__(tma_op: Pointer[ScalingFactorLoader[tma_origin, sf_dtype, sf_layout, sf_desc_layout, cta_group=cta_group, BM_or_MMA_N=BM_or_MMA_N].TmaOp, tma_origin]) -> Self
Initialize the scaling factor loader.
Args:
- tma_op (
Pointer): Pointer to TMA descriptor for scaling factors.
load
load(self, dest: LayoutTensor[sf_dtype, layout, origin, address_space=AddressSpace.SHARED, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], ref [3] barrier: SharedMemBarrier, k_iter: UInt, mn_block_idx: UInt, batch_idx: UInt)
Load scaling factors using TMA 5D copy.
Issues an async load from global memory to shared memory for scaling factor tiles.
Args:
- dest (
LayoutTensor): Destination SMEM tile for scaling factors. - barrier (
SharedMemBarrier): Memory barrier for TMA completion signaling. - k_iter (
UInt): K iteration index. - mn_block_idx (
UInt): M block index (for A) or N block index (for B). - batch_idx (
UInt): Batch dimension index.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!