jax.experimental.pallas.mosaic_gpu.async_store_tmem

目录

jax.experimental.pallas.mosaic_gpu.async_store_tmem#

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

将值存储到 TMEM。

此存储是异步的，并且在调用 commit_tmem 之前，不保证其可见（例如，通过读取或 MMA 操作）。

参数:

ref (_Ref) – 要存储到的 TMEM 引用。
value – 要存储的值。