Fixed-m softmax for 2D-ring x Ulysses attention (WAN 2.2)
Removes the online-softmax running-max and rescale chain from the ring attention kernel, replacing it with a per-query Cauchy-Schwarz bound that is known before streaming starts. Heads whose bound cannot be proven safe keep online softmax, so the result is exact either way. 720p 81f T2V, 40 steps, v7x-8, dp2 x (ring2 x uly2), denoise total: attention=ulysses_ring_custom 9472/1024 109.3s 2.733 s/step attention=ulysses_ring_custom_fixed_m 6400/2048 105.3s 2.633 s/step The tile differs per mode because the optimum moves. Plain is MXU-bound and peaks at bq=9472, degrading to ~3.14 s/step at 6400/2048; fixed-m is flat across both (2.63-2.69), which is the mechanism working -- with the softmax VPU chain gone the kernel stops caring which tile it was given. In-kernel the ring hop goes 22.65 -> 21.15 ms, i.e. 85% -> 91% of the K=128 MXU roofline. Ring specifics, over the flat fixed-m design: - No k-smoothing is possible on a ring (no rank holds the full K, and RoPE does not preserve a zero mean), so the gate uses the two-sided halved bound 106.5 rather than the flat 213. - The bound is global over every K shard, so a head's pinned m is identical on every hop and the partials combine by accumulation. When any head fails the gate the call falls back to an LSE-space merge, which is invariant to the bound's overshoot; a naive (m, l) merge would flush the other partials. - The all-heads-fixed branch compiles a branch-free kernel. Per-head dispatch in the ragged last KV block degrades the whole grid's instruction schedule (measured 3x), so exactly one body is emitted there. Also in this change: - Warmup now runs one real step per weight set. A 2-step warmup puts both steps on the high-noise transformer, leaving the low-noise weights untouched until mid-generation, where their first call cost 26.4s against a 2.6s step. The first generation now reports steady-state latency instead of needing a throwaway run (denoise 118.5s -> 105.3s on a cold process). - _ulysses_attention no longer folds the CFG batch into heads when data/fsdp already shard it; the folded size-1 batch is not divisible by those axes. - Cross-attention demotes to ulysses_custom for the fixed-m ring variant too, matching the other two ring variants. - run_wan_fast_inference.sh takes FIXEDM=0/1 and picks the matching tile. Tests: src/maxdiffusion/tests/ring_fixed_m_test.py covers the online path, an all-eligible call, a sink head that falls back everywhere, a ragged KV tail, and a caller-supplied local max||k|| being promoted to the global bound.
J
James Huang committed
8dd630549ce1cb55e167a66102dcc3aef4cc1926
Parent: bda41e1