tcgen05 submodule#

class cutlass.cute.nvgpu.tcgen05.Repetition(value)#

Bases: Enum

An enumeration for the number of repetitions of a given TMEM copy within the instruction.

x1 = 1#

x2 = 2#

x4 = 4#

x8 = 8#

x16 = 16#

x32 = 32#

x64 = 64#

x128 = 128#

class cutlass.cute.nvgpu.tcgen05.TmemLoadRedOp(value)#

Bases: Enum

An enumeration for the possible reduce operations for TMEM load operations.

class cutlass.cute.nvgpu.tcgen05.Pack(value)#

Bases: Enum

An enumeration for the possible packing patterns for TMEM to RMEM copies.

NONE = 1#

PACK_16b_IN_32b = 2#

class cutlass.cute.nvgpu.tcgen05.Unpack(value)#

Bases: Enum

An enumeration for the possible unpacking patterns for RMEM to TMEM copies.

NONE = 1#

UNPACK_32b_IN_16b = 2#

class cutlass.cute.nvgpu.tcgen05.Ld16x64bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, )#

Bases: _LdBase

16x64b TMEM load Operation.

See the PTX documentation. This Operation corresponds to the .16x64b qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.Ld16x128bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, )#

Bases: _LdBase

16x128b TMEM load Operation.

See the PTX documentation. This Operation corresponds to the .16x128b qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.Ld16x256bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, )#

Bases: _LdBase

16x256b TMEM load Operation.

See the PTX documentation. This Operation corresponds to the .16x256b qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.Ld16x32bx2Op( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, )#

Bases: _LdBase

16x32bx2 TMEM load Operation.

See the PTX documentation. This Operation corresponds to the .16x32bx2 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.Ld32x32bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, )#

Bases: _LdBase

32x32b TMEM load Operation.

See the PTX documentation. This Operation corresponds to the .32x32 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition = <Repetition.x1>, pack: ~cutlass.cute.nvgpu.tcgen05.copy.Pack = <Pack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.St16x64bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, )#

Bases: _StBase

16x64b TMEM store Operation.

See the PTX documentation. This Operation corresponds to the .16x64 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.St16x128bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, )#

Bases: _StBase

16x128b TMEM store Operation.

See the PTX documentation. This Operation corresponds to the .16x128 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.St16x256bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, )#

Bases: _StBase

16x256b TMEM store Operation.

See the PTX documentation. This Operation corresponds to the .16x256 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.St16x32bx2Op( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, )#

Bases: _StBase

16x32x2b TMEM store Operation.

See the PTX documentation. This Operation corresponds to the .16x32x2 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.St32x32bOp( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, )#

Bases: _StBase

32x32b TMEM store Operation.

See the PTX documentation. This Operation corresponds to the .32x32 qualifier.

__init__( repeat: ~cutlass.cute.nvgpu.tcgen05.copy.Repetition, unpack: ~cutlass.cute.nvgpu.tcgen05.copy.Unpack = <Unpack.NONE>, ) → None#

class cutlass.cute.nvgpu.tcgen05.OperandSource(value)#

Bases: Enum

An enumeration for the source memory location of the A input operand of the MMA.

class cutlass.cute.nvgpu.tcgen05.CtaGroup(value)#

Bases: Enum

An enumeration for the cta_group qualifier of the MMA.

ONE = 1#

TWO = 2#

class cutlass.cute.nvgpu.tcgen05.Field(value)#

Bases: Enum

An enumeration for the fields of the MMA Atom that can be modified at runtime.

NEGATE_A = 'neg_a'#

NEGATE_B = 'neg_b'#

ACCUMULATE = 'accum_c'#

SFA = 'sf_a'#

SFB = 'sf_b'#

class cutlass.cute.nvgpu.tcgen05.MmaTF32Op( instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, )#

Bases: MmaOp

TF32 tcgen05 MMA Operation.

See the PTX documentation. This Operation corresponds to the .kind::tf32 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 TF32 MMA Operation'#

__init__( instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaF16BF16Op( ab_dtype: Type[cutlass.cute.typing.Numeric], acc_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, )#

Bases: MmaOp

F16/BF16 tcgen05 MMA Operation.

See the PTX documentation. This Operation corresponds to the .kind::f16 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 F16/BF16 MMA Operation'#

__init__( ab_dtype: Type[cutlass.cute.typing.Numeric], acc_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaI8Op( ab_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, )#

Bases: MmaOp

I8 tcgen05 MMA Operation.

See the PTX documentation. This Operation corresponds to the .kind::i8 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 I8 MMA Operation'#

__init__( ab_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaFP8Op(**kwargs)#

Bases: MmaOp

F8 tcgen05 MMA Operation.

Deprecated since version Use: MmaF8F6F4Op instead.

See the PTX documentation.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 F8 MMA Operation'#

__init__( ab_dtype: Type[cutlass.cute.typing.Numeric], acc_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaF8F6F4Op( a_dtype: Type[cutlass.cute.typing.Numeric], b_dtype: Type[cutlass.cute.typing.Numeric], acc_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, )#

Bases: MmaOp

F8F6F4 tcgen05 MMA Operation.

See the PTX documentation.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 F8F6F4 MMA Operation'#

__init__( a_dtype: Type[cutlass.cute.typing.Numeric], b_dtype: Type[cutlass.cute.typing.Numeric], acc_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaMXF8Op(**kwargs)#

Bases: BlockScaledMmaOp

MXF8 tcgen05 BlockScaled MMA Operation.

Deprecated since version Use: MmaMXF8F6F4Op instead.

See the PTX documentation. This Operation corresponds to the .kind::mxf8f6f4 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 MXF8 BlockScaled MMA Operation'#

__init__( ab_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaMXF8F6F4Op( a_dtype: Type[cutlass.cute.typing.Numeric], b_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, )#

Bases: BlockScaledMmaOp

MXF8F6F4 tcgen05 BlockScaled MMA Operation.

See the PTX documentation. This Operation corresponds to the .kind::mxf8f6f4 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 MXF8F6F4 BlockScaled MMA Operation'#

__init__( a_dtype: Type[cutlass.cute.typing.Numeric], b_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaMXF4Op( instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, )#

Bases: BlockScaledMmaOp

MXF4 tcgen05 BlockScaled MMA Operation.

See the PTX documentation. This Operation corresponds to the .kind::mxf4 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 MXF4 BlockScaled MMA Operation'#

__init__( instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, ) → None#

class cutlass.cute.nvgpu.tcgen05.MmaMXF4NVF4Op( sf_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, )#

Bases: BlockScaledMmaOp

MXF4NVF4 tcgen05 BlockScaled MMA Operation.

See the PTX documentation. This Operation corresponds to the .kind::mxf4nvf4 qualifier.

MMA operations should be issued by a single thread. The DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: MMA without elect_one
cute.gemm(mma_atom, d, a, b, c)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.gemm(mma_atom, d, a, b, c)

descriptive_name = 'tcgen05 MXF4NVF4 BlockScaled MMA Operation'#

__init__( sf_dtype: Type[cutlass.cute.typing.Numeric], instruction_shape: cutlass.cute.typing.Shape, cta_group: CtaGroup, a_src: OperandSource, ) → None#

class cutlass.cute.nvgpu.tcgen05.SmemLayoutAtomKind(value)#

Bases: Enum

Enum class for the kinds of SMEM layout atoms for SM100.

Given a swizzle kind, an SMEM layout atom is the compact layout of smallest size that can be used to construct an SMEM layout using blocked product for operand A or B such that the resulting layout is legal for both TMA and UMMA.

Note that there are other ways of creating legal layouts for operand A and B.

MN_INTER = 1#

MN_SW32 = 2#

MN_SW64 = 3#

MN_SW128 = 4#

MN_SW128_32B = 5#

K_INTER = 6#

K_SW32 = 7#

K_SW64 = 8#

K_SW128 = 9#

cutlass.cute.nvgpu.tcgen05.make_smem_layout_atom( kind: SmemLayoutAtomKind, element_type: Type[cutlass.cute.typing.Numeric], *, loc: cutlass._mlir.ir.Location | None = None, ip: cutlass._mlir.ir.InsertionPoint | None = None, ) → cutlass.cute.typing.ComposedLayout#

Makes a SMEM layout Atom.

This function creates a composed layout in unit of elements consistent with the requested layout Atom kind and element data type.

Parameters:

kind (SmemLayoutAtomKind) – The kind of layout Atom
element_type (Type[Numeric]) – The element data type to construct the layout for

Returns:

The SMEM layout atom

Return type:

ComposedLayout

cutlass.cute.nvgpu.tcgen05.tile_to_mma_shape( atom: cutlass.cute.typing.Layout | cutlass.cute.typing.ComposedLayout, mma_tile_shape: cutlass.cute.typing.Shape, order: cutlass.cute.typing.IntTuple | None = None, *, loc: cutlass._mlir.ir.Location | None = None, ip: cutlass._mlir.ir.InsertionPoint | None = None, ) → cutlass.cute.typing.Layout | cutlass.cute.typing.ComposedLayout#: Tiles a layout to an MMA shape.

cutlass.cute.nvgpu.tcgen05.commit( mbar_ptr: cutlass.cute.typing.Pointer, mask: ~typing.Any | None = None, cta_group: ~cutlass.cute.nvgpu.tcgen05.mma.CtaGroup = <CtaGroup.ONE>, *, loc: cutlass._mlir.ir.Location | None = None, ip: cutlass._mlir.ir.InsertionPoint | None = None, ) → None#

Perform an arrive operation on a mbarrier upon completion of previous MMA operations.

Single-Thread Execution Required - DSL Does NOT Handle Automatically: This operation must be wrapped in cute.arch.elect_one(). Without elect_one(), all 32 threads in the warp will execute the commit, causing 32x redundant tcgen05.commit PTX instructions.

# CORRECT: Wrap tcgen05.commit in elect_one
with cute.arch.elect_one():
    tcgen05.commit(barrier_ptr, None, cta_group)

# WRONG: Without elect_one, all threads execute (32x redundant)
tcgen05.commit(barrier_ptr, None, cta_group)

Parameters:

mbar_ptr (Pointer) – A pointer to the mbarrier in SMEM
mask (Int) – An optional multicast mask for the CTAs in the cluster to signal arrival to
cta_group (CtaGroup) – The CTA group size for the operation (ONE or TWO)