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, *, impl='tma', collective_axes=None, leader_tracked=None, oob_mode=None)[source]#

Asynchronously copies a GMEM reference to a SMEM reference.

When impl="tma" and 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) – The barrier to use for tracking completion of the copy.

  • impl (Literal['tma', 'cp_async']) – The underlying copy implementation to use: "cp_async" or "tma". Defaults to "tma".

  • 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 when impl="tma".

  • 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 when impl="tma".

  • oob_mode (OOBFillMode | None) – The optional out-of-bounds fill mode. Can be OOBFillMode.UNDEFINED, OOBFillMode.PROMISE_IN_BOUNDS or OOBFillMode.ZEROS. Only supported when impl="tma".

Return type:

None

See also

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