jax
08a8bf71 - [Pallas/Mosaic GPU] Expose `plgpu.wgmma_accumulator_load`.

Commit
237 days ago
[Pallas/Mosaic GPU] Expose `plgpu.wgmma_accumulator_load`. This function allows loading from a WGMMA accumulator while controlling the amount of `wgmma` instructions that are allowed to remain in flight before dereferencing occurs. If the `wait_n` parameter is set to `None`, no synchronization is done. PiperOrigin-RevId: 858186473
Author
Parents
Loading