Mojo module
mla_decode_sm100_utils
comptime values
MBarType
comptime MBarType = UnsafePointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED]
QOTMATile
comptime QOTMATile[dtype: DType, BM: Int, BK: Int, swizzle_mode: TensorMapSwizzle] = TMATensorTile[dtype, tile_layout_k_major[dtype, BM, BK, swizzle_mode](), _tma_desc_tile_layout[dtype, 2, IndexList(BM, BK, Tuple()), swizzle_mode]()]
Parameters
- dtype (
DType): - BM (
Int): - BK (
Int): - swizzle_mode (
TensorMapSwizzle):
SharedMemPointer
comptime SharedMemPointer[type: AnyType] = UnsafePointer[type, MutAnyOrigin, address_space=AddressSpace.SHARED]
Parameters
- type (
AnyType):
SharedMemTensor
comptime SharedMemTensor[dtype: DType, layout: Layout] = LayoutTensor[dtype, layout, MutAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=DType.int32, linear_idx_type=DType.int32, alignment=128]
Parameters
Structs
-
DecodeCConsumer: -
DecodeCProducer: -
DecodeKVConsumer: -
DecodeKVProducer: -
DecodeOConsumer: -
DecodeOProducer: -
DecodeOutConsumer: -
DecodeOutProducer: -
DecodePConsumer: -
DecodePProducer: -
DecodeSConsumer: -
DecodeSM100MiscMBars: -
DecodeSM100PVSS: -
DecodeSM100QKTSS: -
DecodeSProducer: -
KVCvt2MmaConsumer: -
KVCvt2MmaProducer: -
KVLoad2CvtConsumer: -
KVLoad2CvtProducer: -
KVPipelineGeneric: KVPipeline hasnum_kv_stages * num_qk_stagesstages.num_kv_stagesrefers to how manyKandVtiles we pipeline for performing theS = Q@K'andO += P@VMMAs. Each of these MMAs is broken up intonum_qk_stagespipelined MMAs. We setstep=Falsefor all but the last MMA that completes the operation. An alternative implementation would separate the two, and potentially allow for more overall stages at the cost of slightly more bookkeeping. -
MLA_Decode_Pack: -
MLA_SM100_Decode_Common: -
MLA_SM100_Decode_Config: -
OffsetPosition: -
OutPipeline: OutPipeline hasnum_out_stagesstages.num_out_stagesrefers to how many output stages we pipeline for performing the output store.
Functions
-
build_mma_ss_ws: -
bulk_mma_ws: -
clamped_index_coordinate: -
cvt_fp8x8_from_2xu32_to_bf16x8_packed_u32x4: -
e8m0_to_bf16_broadcast: Convert an e8m0 scale byte to a bf16 value broadcast into both halves of a uint32. -
hmul2_bf16x8_by_scalar: Multiply 8 packed bf16 values (in 4 uint32 registers) by a bf16x2 scalar broadcast. -
ld_shared_v4_u32: -
num_matrix_view_rows_decode: -
st_shared_v4_b32_at_bf16_elem_off: -
tma_tile_qo: -
write_bf16x2_row_to_smem_chunked: Chunked write with optional scaling. Reduces register pressure.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!