cpasync submodule#

class cutlass.cute.nvgpu.cpasync.LoadCacheMode(value)#

Bases: Enum

An enumeration for the possible cache modes of a non-bulk cp.async instruction.

See the PTX documentation.

class cutlass.cute.nvgpu.cpasync.CopyG2SOp(
cache_mode: LoadCacheMode | LoadCacheMode = cutlass._mlir.dialects.cute_nvgpu.LoadCacheMode.always,
)#

Bases: CopyOp

Non-bulk asynchronous GMEM to SMEM Copy Operation.

See the PTX documentation.

__init__(
cache_mode: LoadCacheMode | LoadCacheMode = cutlass._mlir.dialects.cute_nvgpu.LoadCacheMode.always,
)#
class cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(
cta_group: ~cutlass.cute.nvgpu.tcgen05.mma.CtaGroup = <CtaGroup.ONE>,
)#

Bases: TmaCopyOp

Bulk tensor asynchronous GMEM to SMEM Copy Operation using the TMA unit.

TMA copy operations are issued by a single thread within a warp, but the DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: TMA copy without elect_one
cute.copy(
    tma_atom,
    gmem_tensor,  # TMA partition ensures single-thread automatically
    smem_tensor,
    tma_bar_ptr=barrier_ptr
)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.copy(tma_atom, gmem_tensor, smem_tensor, tma_bar_ptr=barrier_ptr)

While the TMA copy itself does not need elect_one(), barrier initialization and transaction byte setup must use elect_one():

# Barrier setup requires elect_one
with cute.arch.elect_one():
    cute.arch.mbarrier_init(barrier_ptr, arrival_count)
    cute.arch.mbarrier_expect_tx(barrier_ptr, num_tma_bytes)

# TMA copy does NOT need elect_one
cute.copy(tma_atom, gmem_tensor, smem_tensor, tma_bar_ptr=barrier_ptr)

PTX Programming Model: In PTX, TMA operations (cp.async.bulk.tensor) must be issued by a single thread. The DSL automatically handles this.

See the PTX documentation. This Operation uses TMA in the .tile mode.

See also

  • cute.arch.elect_one() - NOT needed for TMA copy, but needed for barrier setup

  • cute.arch.mbarrier_init() - Requires elect_one

  • cute.arch.mbarrier_expect_tx() - Requires elect_one

  • Tutorial example: examples/blackwell/tutorial_tma/tma_v0.py

cta_group: CtaGroup = 1#
__init__(
cta_group: ~cutlass.cute.nvgpu.tcgen05.mma.CtaGroup = <CtaGroup.ONE>,
) → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
cta_group: ~cutlass.cute.nvgpu.tcgen05.mma.CtaGroup = <CtaGroup.ONE>,
)#

Bases: TmaCopyOp

Bulk tensor asynchronous multicast GMEM to SMEM Copy Operation using the TMA unit.

TMA multicast operations are issued by a single thread within a warp, but the DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: TMA multicast without elect_one
cute.copy(
    tma_atom.with_(mcast_mask=cluster_mask),
    gmem_tensor,
    smem_tensor,
    tma_bar_ptr=barrier_ptr
)

# WRONG: Do NOT wrap in elect_one (can cause deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.copy(tma_atom.with_(mcast_mask=mask), gmem_tensor, smem_tensor)

PTX Programming Model: In PTX, TMA multicast operations (cp.async.bulk.tensor.multicast) must be issued by a single thread.

See the PTX documentation. This Operation uses TMA in the .tile mode.

See also

cta_group: CtaGroup = 1#
__init__(
cta_group: ~cutlass.cute.nvgpu.tcgen05.mma.CtaGroup = <CtaGroup.ONE>,
) → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp#

Bases: TmaCopyOp

Bulk tensor asynchronous SMEM to GMEM Copy Operation using the TMA unit.

TMA store operations are issued by a single thread within a warp, but the DSL automatically handles this by implicitly adding elect_one() around the copy operation.

# CORRECT: TMA store without elect_one
cute.copy(
    tma_atom,
    smem_tensor,  # Source: shared memory
    gmem_tensor,  # Destination: global memory
)

# WRONG: Do NOT wrap in elect_one (causes deadlock)
with cute.arch.elect_one():  # INCORRECT
    cute.copy(tma_atom, smem_tensor, gmem_tensor)

PTX Programming Model: In PTX, TMA store operations must be issued by a single thread. The DSL automatically handles this.

See the PTX documentation. This Operation uses TMA in the .tile mode.

See also

  • cute.arch.elect_one() - NOT needed for TMA store

  • CopyBulkTensorTileG2SOp - TMA load operation

  • Tutorial example: examples/blackwell/tutorial_tma/tma_v0.py

__init__() → None#
class cutlass.cute.nvgpu.cpasync.CopyReduceBulkTensorTileS2GOp(
reduction_kind: cutlass._mlir.dialects.cute_nvgpu.ReductionKind = cutlass._mlir.dialects.cute_nvgpu.ReductionKind.ADD,
)#

Bases: TmaCopyOp

Bulk tensor asynchronous SMEM to GMEM Reduction Operation using the TMA unit.

See the PTX documentation. This Operation uses TMA in the .tile mode.

__init__(
reduction_kind: cutlass._mlir.dialects.cute_nvgpu.ReductionKind = cutlass._mlir.dialects.cute_nvgpu.ReductionKind.ADD,
) → None#
class cutlass.cute.nvgpu.cpasync.CopyDsmemStoreOp#

Bases: CopyOp

Asynchronous Store operation to DSMEM with explicit synchronization.

See the PTX documentation.

__init__() → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkG2SOp#

Bases: CopyOp

Bulk copy asynchrnous GMEM to SMEM Copy Operation.

See the PTX documentation.

__init__() → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkG2SMulticastOp#

Bases: CopyOp

Bulk multicast copy asynchrnous GMEM to SMEM Copy Operation.

See the PTX documentation.

__init__() → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkS2GOp#

Bases: CopyOp

Bulk copy asynchrnous SMEM to GMEM Copy Operation.

See the PTX documentation.

__init__() → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkS2GByteMaskOp#

Bases: CopyOp

Bulk copy asynchrnous SMEM to GMEM Copy Operation with mask. The i-th bit in the 16-bit wide byteMask operand specifies whether the i-th byte of each 16-byte wide chunk of source data is copied to the destination.

See the PTX documentation.

__init__() → None#
class cutlass.cute.nvgpu.cpasync.CopyBulkS2SOp#

Bases: CopyOp

Bulk copy asynchrnous SMEM CTA to Cluster Copy Operation.

See the PTX documentation.

__init__() → None#
class cutlass.cute.nvgpu.cpasync.TmaCopyOp#

Bases: CopyOp

Base class for all TMA copy operations.

class cutlass.cute.nvgpu.cpasync.TmaInfo(
copy_atom: CopyAtom,
tma_tensor: Any,
smem_layout: Any | None = None,
)#

Bases: object

Container for TMA Copy Atom and related data.

This class uses software composition to bundle a CopyAtom with the SMEM layout and TMA tensor.

Supports tuple unpacking for backward compatibility:

atom, tma_tensor = make_tiled_tma_atom(...)

Access smem_layout via the container:

tma_info = make_tiled_tma_atom(...)
layout = tma_info.smem_layout
Parameters:

  • atom (CopyAtom) – The TMA Copy Atom

  • tma_tensor – The TMA tensor for coordinate mapping

  • smem_layout – The SMEM layout used to construct the TMA descriptor

__init__(
copy_atom: CopyAtom,
tma_tensor: Any,
smem_layout: Any | None = None,
) → None#
property atom: CopyAtom#

The TMA Copy Atom.

property tma_tensor: Any#

The TMA tensor for coordinate mapping.

property smem_layout: Any#

The SMEM layout used to construct the TMA descriptor.

cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom(
op: CopyBulkTensorTileG2SOp | CopyBulkTensorTileG2SMulticastOp | CopyBulkTensorTileS2GOp | CopyBulkTensorIm2ColG2SOp | CopyBulkTensorIm2ColS2GOp | CopyReduceBulkTensorTileS2GOp,
gmem_tensor: cutlass.cute.typing.Tensor,
smem_layout_: cutlass.cute.typing.Layout | cutlass.cute.typing.ComposedLayout,
cta_tiler: cutlass.cute.typing.Tiler,
num_multicast: int = 1,
*,
internal_type: Type[cutlass.cute.typing.Numeric] | None = None,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → TmaInfo#

Makes a TMA Copy Atom in the .tile mode to copy tiles of a GMEM tensor to/from SMEM buffer with the given Layout.

Given

  • a GMEM tensor

  • a SMEM layout

  • a CTA-level Tiler

this function figures out the bulk tensor asynchronous copy instruction to use with the maximum “TMA vector length” to copy tiles of the GMEM tensor to/from an SMEM buffer with the provided layout while maintaining consistency with the provided Tiler.

This function returns two results:

  1. the Copy Atom

  2. a TMA tensor that maps logical coordinates of the GMEM tensor to coordinates consumed by the TMA unit. TMA tensors contain basis stride elements that enable their associated layout to compute coordinates. Like other CuTe tensors, TMA tensors can be partitioned.

Parameters:

  • op (TMAOp) – The TMA Copy Operation to construct an Atom

  • gmem_tensor (Tensor) – The GMEM tensor involved in the Copy

  • smem_layout (Union[Layout, ComposedLayout]) – The SMEM layout to construct the Copy Atom, either w/ or w/o the stage mode

  • cta_tiler (Tiler) – The CTA Tiler to use

  • num_multicast (int) – The multicast factor

  • internal_type (Type[Numeric]) – Optional internal data type to use when the tensor data type is not supported by the TMA unit

Returns:

A TmaInfo containing the Copy Atom, TMA tensor, and SMEM layout

Return type:

TmaInfo

cutlass.cute.nvgpu.cpasync.tma_partition(
atom: CopyAtom,
cta_coord: cutlass.cute.typing.Coord,
cta_layout: cutlass.cute.typing.Layout,
smem_tensor: cutlass.cute.typing.Tensor,
gmem_tensor: cutlass.cute.typing.Tensor,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → Tuple[cutlass.cute.typing.Tensor, cutlass.cute.typing.Tensor]#

Tiles the GMEM and SMEM tensors for the provided TMA Copy Atom.

cutlass.cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk: cutlass.cute.typing.Layout,
cta_coord_vmnk: cutlass.cute.typing.Coord,
mcast_mode: int,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → cutlass.cute.typing.Int16#

Computes a multicast mask for a TMA load Copy.

Parameters:

  • cta_layout_vmnk (Layout) – The VMNK layout of the cluster

  • cta_coord_vmnk (Coord) – The VMNK coordinate of the current CTA

  • mcast_mode (int) – The tensor mode in which to multicast

Returns:

The resulting mask

Return type:

Int16

cutlass.cute.nvgpu.cpasync.prefetch_descriptor(
tma_atom: CopyAtom,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → None#

Prefetches the TMA descriptor associated with the TMA Atom.

cutlass.cute.nvgpu.cpasync.copy_tensormap(
tma_atom: CopyAtom,
tensormap_ptr: cutlass.cute.typing.Pointer,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → None#

Copies the tensormap held by a TMA Copy Atom to the memory location pointed to by the provided pointer.

Parameters:

  • tma_atom (CopyAtom) – The TMA Copy Atom

  • tensormap_ptr (Pointer) – The pointer to the memory location to copy the tensormap to

cutlass.cute.nvgpu.cpasync.update_tma_descriptor(
tma_atom: CopyAtom,
gmem_tensor: cutlass.cute.typing.Tensor,
tma_desc_ptr: cutlass.cute.typing.Pointer,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → None#

Updates the TMA descriptor in the memory location pointed to by the provided pointer using information from a TMA Copy Atom and the provided GMEM tensor.

Specifically, the following fields of the TMA descriptor will be updated:

  1. the GMEM tensor base address

  2. the GMEM tensor shape

  3. the GMEM tensor stride

Other fields of the TMA descriptor are left unchanged.

Parameters:

  • tma_atom (CopyAtom) – The TMA Copy Atom

  • gmem_tensor (Tensor) – The GMEM tensor

  • tensormap_ptr (Pointer) – The pointer to the memory location of the descriptor to udpate

cutlass.cute.nvgpu.cpasync.fence_tma_desc_acquire(
tma_desc_ptr: cutlass.cute.typing.Pointer,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → None#

See the PTX documentation.

cutlass.cute.nvgpu.cpasync.cp_fence_tma_desc_release(
tma_desc_global_ptr: cutlass.cute.typing.Pointer,
tma_desc_shared_ptr: cutlass.cute.typing.Pointer,
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → None#

See the PTX documentation.

cutlass.cute.nvgpu.cpasync.fence_tma_desc_release(
*,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → None#

See the PTX documentation.

cutlass.cute.nvgpu.cpasync.group_bulk_copy_modes(
src: cutlass.cute.typing.Tensor,
dst: cutlass.cute.typing.Tensor,
loc: cutlass._mlir.ir.Location | None = None,
ip: cutlass._mlir.ir.InsertionPoint | None = None,
) → Tuple[cutlass.cute.typing.Tensor, cutlass.cute.typing.Tensor]#

Copy async bulk need group mode 0, acquiring whole tensor for bulk copy