jax.experimental.pallas.mosaic_gpu.try_cluster_cancel

jax.experimental.pallas.mosaic_gpu.try_cluster_cancel(result_ref, barrier, *, collective_axes=())

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 retrieving its indices.

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.

• collective_axes (Hashable | tuple[Hashable, ...]) – The thread / cluster axis names across which the cancellation is collective. When specified, only the axis index 0 along these axes will issue the cancellation request. If more than one thread along an axis calls this function, the axis must be specified here to avoid inadvertently cancelling multiple clusters.

Return type:

None

See also

jax.experimental.pallas.mosaic_gpu.query_cluster_cancel()