jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem

jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem#

jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem(src, dst, barrier=None, *, collective_axes=None, leader_tracked=None, oob_mode=None, predicate=None)[source]#

Asynchronously copies a GMEM reference to a SMEM reference.

The underlying copy implementation is selected automatically based on the target architecture: TMA on Hopper and newer, and cp.async on older architectures.

When collective_axes is specified, the copy involves multiple CUDA blocks in the cluster. The value of leader_tracked determines the behavior:

  • None (multicast): all CUDA blocks sharing the same index along the collective axes receive the same data from src.

  • CopyPartition.PARTITIONED(axis): each CUDA block in the cluster receives a transfer_size // cluster_size tile of src. E.g. for src of shape (256, 256) with cluster size 2 along axis 0: block 0 gets src[0:128, :], block 1 gets src[128:256, :].

  • CopyPartition.REPLICATED: all CUDA blocks in the cluster load the same data, but only the first block tracks progress via barrier arrivals.

Note

For leader-tracked copies, only the first CUDA block in the cluster arrives on the barrier. If other blocks need to consume the copied data, an additional cluster barrier is necessary to ensure all blocks have finished the copy.

Parameters:

  • src (_Ref) – The source Ref. Must be in GMEM.

  • dst (_Ref) – The destination Ref. Must be in SMEM.

  • barrier (_Ref | None) – The barrier to use for tracking completion of the copy. Must be omitted on pre-Hopper GPUs, which use the cp.async implementation. When omitted, completion of the copy has to be awaited with jax.experimental.pallas.mosaic_gpu.wait_gmem_to_smem() instead of jax.experimental.pallas.mosaic_gpu.barrier_wait().

  • collective_axes (str | tuple[str, ...] | None) – The collective axes to use for the copy. Only a single collective axis is supported when leader_tracked is specified (but the axis can be composite). Only supported with the TMA implementation (Hopper and newer).

  • leader_tracked (CopyPartition | None) – If specified, only the leader block in the cluster will observe the completion of the copy. If CopyPartition.PARTITIONED(axis), performs a partitioned collective copy along the given axis. If CopyPartition.REPLICATED, all blocks load the same data. Only supported with the TMA implementation (Hopper and newer).

  • oob_mode (OOBFillMode | None) – The optional out-of-bounds fill mode. Can be OOBFillMode.UNDEFINED, OOBFillMode.PROMISE_IN_BOUNDS or OOBFillMode.ZEROS for the TMA implementation, and only OOBFillMode.PROMISE_IN_BOUNDS for the cp.async one.

  • predicate (jax.Array | None) – A boolean indicating whether the copy should be performed. If None, the copy is always performed.

Return type:

None

See also

jax.experimental.pallas.mosaic_gpu.barrier_arrive() jax.experimental.pallas.mosaic_gpu.barrier_wait() jax.experimental.pallas.mosaic_gpu.wait_gmem_to_smem()