tilelang.cuda.op.gemm.gemm_mma_sm120¶

Attributes¶

GEMM_INST_MMA_BLOCK_SCALED

Classes¶

GemmMMASm120BlockScaled

SM120 warp-level block-scaled MMA lowering.

Module Contents¶

tilelang.cuda.op.gemm.gemm_mma_sm120.GEMM_INST_MMA_BLOCK_SCALED = 'cuda.mma.blockscaled'¶
class tilelang.cuda.op.gemm.gemm_mma_sm120.GemmMMASm120BlockScaled¶

Bases: tilelang.tileop.gemm_blockscaled.gemm_blockscaled_base.GemmBlockScaledMixin, tilelang.cuda.op.gemm.gemm_mma.GemmMMA

SM120 warp-level block-scaled MMA lowering.

A may be shared or a full fragment; B remains shared. Fragment A and fragment scales use row-major packed uint32 scales. Compact K-major scales retain the shared/shared ping-pong path, including odd atom grids.

supported_sf_layouts = ('rowmajor', 'blockscaled_chunk_kmajor')¶
intrin_emitter_cls¶
infer_layout(target, thread_nums)¶
Parameters:

  • target (tvm.target.Target)

  • thread_nums (int)

lower(layout_map, target, thread_bounds, thread_index, mbar_phase_expr=None)¶
Parameters:

  • layout_map (dict)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_index (tvm.tirx.PrimExpr)

  • mbar_phase_expr (tvm.tirx.PrimExpr | None)