jax.experimental.pallas.mosaic_gpu.async_store_tmem

Contents

jax.experimental.pallas.mosaic_gpu.async_store_tmem#

jax.experimental.pallas.mosaic_gpu.async_store_tmem(ref, value)[source]#

Stores the value to TMEM.

The store is asynchronous and is not guaranteed to be visible (e.g. by reads or MMA operations) until commit_tmem has been called.

Parameters:

ref (_Ref) – The TMEM reference to store to.
value – The value to store.