Commits on Sep 13, 2026

1. feat(attention,wan): fixed-m splash and 2D-ring attention with Wan2.2… …

   … TPU optimizations
   
   Implements exact fixed-m splash attention with Global Virtual K-Centering,
   2D Ulysses+Ring attention, and Wan2.2 fast serving optimizations for Cloud TPUs.
   
   Key Features & Safety:
   1. Fixed-m Splash Attention:
      - Dynamic C(N) headroom constants guaranteeing FP32 accumulator safety.
      - Fixed-m dtype safety checks (rejecting FP16/FP8 exponent overflow).
      - Global V-magnitude bounds (|V| <= 256) with safe online softmax fallback.
   2. 2D Ulysses + Ring Attention:
      - Global Virtual K-Centering via cross-ring pmean(k).
      - Cross-ring v_ok reduction via jax.lax.pmin preventing unsafe hop-level fixed-m.
      - Per-Q-block fixed-m accumulation with tuned tile size (block_q=9472).
      - Full GQA support across arbitrary Ulysses and Ring splits.
   3. Wan2.2 Fast Serving:
      - Persistent AOT compilation caching and zero-exec warmup.
      - Tuned TPU inference recipes in run_wan_fast_inference.sh.
   
   Hardware Verification (Cloud TPU v6e-8):
   - 44 tests passed in ring_fixed_m_test and custom_splash_fixed_m_test.
   - 25-prompt VBench evaluation: zero quality regression vs main.
   - Step latency: 3.445 s/step (tile 9472) vs 3.526 s/step (main), delivering a +2.32% denoise speedup.

   

   Perseus14 committed Sep 13, 2026
   Configuration menu

   • View commit details
   • Copy full SHA for 66c8eb6
   • Browse repository at this point

   Copy the full SHA

   66c8eb6 View commit details

   Browse the repository at this point in the history