jax.experimental.pallas.mosaic_gpu.try_cluster_cancel#

jax.experimental.pallas.mosaic_gpu.try_cluster_cancel(result_ref, barrier)[source]#

Initiates an async request to claim a new work unit from the grid.

It allows an SM to dynamically acquire work by atomically canceling the launch of a pending cluster from the grid and claiming its CTA ID as the next unit of work.

Parameters:

  • result_ref (_Ref) – An SMEM ref where the 16-byte result will be stored.

  • barrier (_Ref) – A barrier used to coordinate the completion of the query.

Return type:

None

See also

jax.experimental.pallas.mosaic_gpu.query_cluster_cancel()