Skip to content

Commit d16343e

Browse files
author
Toshi Pahadia
committed
perf(wan2.2): Optimize attention kernels, rematerialization, and Ulysses parallelism for TPU v7x
- Implement custom Pallas splash & ring attention forward and fused backward kernels - Optimize rematerialization policy with custom checkpointing rules - Implement Ulysses sequence parallelism with 2-phase all-to-all communications - Eliminate device-to-host sync stall via deterministic host-side expert routing - Support multi-host evaluation metric gathering with multihost_utils.process_allgather - Specify sequential dimension semantics in Pallas backward kernels for safe accumulation
1 parent 3b6ee39 commit d16343e

6 files changed

Lines changed: 1796 additions & 124 deletions

File tree

0 commit comments

Comments
 (0)