tilelang.cuda.op.gemm.gemm_tcgen05_blockscaled¶

Blackwell TCGEN5MMA block-scaled GEMM lowering.

Attributes¶

GEMM_INST_TCGEN05_BLOCK_SCALED

Classes¶

GemmTCGEN5BlockScaled

kind::mxf8f6f4.block_scale TCGEN5MMA with A/B in shared memory,

Module Contents¶

tilelang.cuda.op.gemm.gemm_tcgen05_blockscaled.GEMM_INST_TCGEN05_BLOCK_SCALED = 'cuda.tcgen05.blockscaled'¶
class tilelang.cuda.op.gemm.gemm_tcgen05_blockscaled.GemmTCGEN5BlockScaled¶

Bases: tilelang.tileop.gemm_blockscaled.gemm_blockscaled_base.GemmBlockScaledMixin, tilelang.cuda.op.gemm.gemm_tcgen05.GemmTCGEN5

kind::mxf8f6f4.block_scale TCGEN5MMA with A/B in shared memory, the accumulator in tensor memory and SFA/SFB already resident in tensor memory.

Shares the shared-operand swizzle inference and emitter plumbing with the dense GemmTCGEN5; differs in the fixed 1x1 warp partition (kept even under cta_group::2), the dense-only (no .ws) instruction shapes and the block-scaled MMA issue.

tcgen05_allow_ws = False¶
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.

Parameters:

  • target (tvm.target.Target)

  • thread_nums (int)

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

Lower to TIR containing block-scaled TCGEN5MMA calls.

Follows the same completion protocol as the dense TCGEN05 lowering: the synchronous T.gemm_blockscaled posts completion to mbar and waits on it right after issue, while the explicit is_tcgen05 op never waits and may omit mbar when the caller or WS schedule emits a later completion arrival.

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)