Repository navigation
Commit d16343e
Toshi Pahadia
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 accumulation1 parent 3b6ee39 commit d16343e
6 files changed
Lines changed: 1796 additions & 124 deletions
File tree
- src/maxdiffusion
- kernels
- splash_attention
- models
- tests
- trainers
0 commit comments