Ring attention integration and other optimizations#359

Draft

eltsai wants to merge 27 commits intomainfrom

elisatsai_ring_attention

Commits on Dec 15, 2025

Reapply "Attention bug fixes, tokamax splash defaulting logic (#282)" (#287)
Elisa Tsai
committed
Reapply "Cross self attention switch (#251)" (#288)
Elisa Tsai
committed
Disable unsafe rng
Elisa Tsai
committed

Commits on Dec 17, 2025

Integrate tokamax ring attention as optional attention kernel for WAN 2.1
Elisa Tsai
committed

Commits on Dec 29, 2025

Merge branch 'main' into elisatsai_disable_unsafe_rng
eltsai
authored

Commits on Dec 30, 2025

Fixed formatting issue
Elisa Tsai
committed
Updated scheduler test values
Elisa Tsai
committed
Updated values based on v5p-8 tests
Elisa Tsai
committed

Commits on Jan 5, 2026

Fixing ring attention
Elisa Tsai
committed

Commits on Feb 10, 2026

moving kernel init outside the sharding map
Elisa Tsai
committed

Commits on Feb 15, 2026

Revert "moving kernel init outside the sharding map"
Elisa Tsai
committed

Commits on Feb 23, 2026

jitting and sharding vae, refactored for loops in jitted VAE, 132 sec on 16 TPUs
Elisa Tsai
committed

Commits on Feb 26, 2026

Renaming VAE sharding axis to vae_spatial
Elisa Tsai
committed
Renaming VAE sharding axis to vae_spatial
Elisa Tsai
committed

Commits on Mar 2, 2026

ring-attention
coolkp
committed

Commits on Mar 4, 2026

Commits on Mar 5, 2026

Commits on Mar 10, 2026

Added sharding on ROPE
Elisa Tsai
committed

Commits on Mar 11, 2026

cfg cache

James Huang
authored and
Elisa Tsai
committed
Merged CFG cache, 220 sec using tokamax_flash
Elisa Tsai
committed

Commits on Mar 12, 2026

Changed profiling logic
Elisa Tsai
committed

Commits on Mar 16, 2026

Format fix
Elisa Tsai
committed
Merge remote-tracking branch 'origin/main' into elisatsai_ring_attention
Elisa Tsai
committed

Commits on Mar 19, 2026

updated vae config logic to be the consistent, update xprof logic
Elisa Tsai
committed