warpgroup submodule#

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

Bases: Enum

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

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

Bases: Enum

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

ACCUMULATE = 'accum_c'#

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

Bases: MmaOp

F16/BF16 warpgroup MMA Operation.

See the PTX documentation. This Operation covers the instructions using the .f16 or .bf16 qualifiers for the input operands.

Supported data type combinations:

A Data Type	B Data Type	Acc Type	Mma-K
F16	F16	F16, F32	16
BF16	BF16	F32	16

Supported architectures: sm_90a

Constraints:

Mma-M = 64
8 <= Mma-N <= 256, step 8
A and B support both K-major and MN-major (transpose) when A is in shared memory (descriptor). When A is in registers, only B can be transposed.

Execution Model:

WGMMA is asynchronous and collective at warpgroup scope (4 contiguous warps). In user code, cute.gemm(...) should be issued warpgroup-uniformly.
Before issuing cute.gemm(...), call cute.nvgpu.warpgroup.fence() to order prior register writes to accumulator/A fragments with subsequent WGMMA reads.
After issuing cute.gemm(...), call cute.nvgpu.warpgroup.commit_group(). Use cute.nvgpu.warpgroup.wait_group(N) before consuming or reusing accumulator values from pending WGMMA groups.

cute.nvgpu.warpgroup.fence()
cute.gemm(tiled_mma, acc, tCrA[tile_crd], tCrB[tile_crd], acc)
cute.nvgpu.warpgroup.commit_group()
cute.nvgpu.warpgroup.wait_group(1)
# ... pipeline continues ...
cute.nvgpu.warpgroup.wait_group(0)

descriptive_name: ClassVar[str] = 'warpgroup 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, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

class cutlass.cute.nvgpu.warpgroup.MmaF8Op( 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, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, )#

Bases: MmaOp

F8 warpgroup MMA Operation.

See the PTX documentation. This Operation covers the instructions using the .e4m3 or .e5m2 qualifiers for the input operands.

Supported data type combinations:

A Data Type	B Data Type	Acc Type	Mma-K
E4M3, E5M2	E4M3, E5M2	F16, F32	32

Supported architectures: sm_90a

Constraints:

Mma-M = 64
8 <= Mma-N <= 256, step 8
A and B data types are independent (mixed FP8 allowed)
Transpose (MN-major) is not supported for A or B. Both operands must be K-major.

Execution Model:

WGMMA is asynchronous and collective at warpgroup scope (4 contiguous warps). In user code, cute.gemm(...) should be issued warpgroup-uniformly.
Before issuing cute.gemm(...), call cute.nvgpu.warpgroup.fence() to order prior register writes to accumulator/A fragments with subsequent WGMMA reads.
After issuing cute.gemm(...), call cute.nvgpu.warpgroup.commit_group(). Use cute.nvgpu.warpgroup.wait_group(N) before consuming or reusing accumulator values from pending WGMMA groups.

cute.nvgpu.warpgroup.fence()
cute.gemm(tiled_mma, acc, tCrA[tile_crd], tCrB[tile_crd], acc)
cute.nvgpu.warpgroup.commit_group()
cute.nvgpu.warpgroup.wait_group(1)
# ... pipeline continues ...
cute.nvgpu.warpgroup.wait_group(0)

descriptive_name: ClassVar[str] = 'warpgroup F8 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, a_src: OperandSource, a_major_mode: OperandMajorMode | OperandMajorMode, b_major_mode: OperandMajorMode | OperandMajorMode, ) → None#

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

Bases: Enum

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

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#

K_INTER = 5#

K_SW32 = 6#

K_SW64 = 7#

K_SW128 = 8#

cutlass.cute.nvgpu.warpgroup.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.warpgroup.fence( *, loc: cutlass._mlir.ir.Location | None = None, ip: cutlass._mlir.ir.InsertionPoint | None = None, ) → None#: See the PTX documentation.

cutlass.cute.nvgpu.warpgroup.commit_group( *, loc: cutlass._mlir.ir.Location | None = None, ip: cutlass._mlir.ir.InsertionPoint | None = None, ) → None#: See the PTX documentation.

cutlass.cute.nvgpu.warpgroup.wait_group( group: Any, *, loc: cutlass._mlir.ir.Location | None = None, ip: cutlass._mlir.ir.InsertionPoint | None = None, ) → None#: See the PTX documentation.