tilelang.tileop.gemm.gemm_tcgen05

Classes

GemmTCGEN5

GEMM operator for Blackwell (SM100) TCGEN5MMA instructions.

Module Contents

class tilelang.tileop.gemm.gemm_tcgen05.GemmTCGEN5

Bases: tilelang.tileop.gemm.gemm_base.GemmBase

GEMM operator for Blackwell (SM100) TCGEN5MMA instructions.

Supports the SS (Shared-Shared) and TS (TensorMemory-Shared) variants, as well as block-scaled MXFP8 GEMM when SFA/SFB scale factors are present. Layout inference and lowering are dispatched based on the memory scopes of operands A and B.

infer_shared_layout(continuity)

Infer a standard shared-memory swizzle layout for TCGEN05 operands.

参数:

continuity (int)

返回类型:

Callable[[tvm.tir.Buffer], tilelang.layout.Layout]

infer_layout(target, thread_nums)

Infer swizzled layouts for operands and accumulator.

For SS: both A and B get swizzled shared-memory layouts. For TS: A and C get TMEM store layouts, B gets a swizzled shared-memory layout. For block-scaled: same as SS (A and B get swizzle, C gets TMEM store layout).

参数:

  • target (tvm.target.Target)

  • thread_nums (int)

lower(layout_map, target, thread_bounds, thread_var, mbar_phase_expr=None)

Lower the GEMM tile-op into a TIR prim_func containing TCGEN5MMA calls.

参数:

  • layout_map (dict)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_var (tvm.tir.Var)

  • mbar_phase_expr (tvm.tir.PrimExpr | None)