jax
6f06ab80 - [Mosaic:GPU] Simplify collective metadata initialization.

Commit
215 days ago
[Mosaic:GPU] Simplify collective metadata initialization. Since the output buffer can be reused by the buffer assigner in several different HLO ops, previously we had to copy the metadata to the device at every execution. Instead we can simply allocate small chunk of memory (~1-2KB) on a device at each mosaic custom call and use it to store collective metadata once during initialization. This also simplifies a custom call creation code in core.py. PiperOrigin-RevId: 875532483
Parents
Loading