Fixed-m softmax for 2D-ring x Ulysses attention (WAN 2.2)

[syhuang22]

Replaces the online-softmax running-max and rescale chain in the 2D-ring x
Ulysses attention kernel with a per-query Cauchy-Schwarz bound known before
streaming starts. Heads whose bound cannot be proven safe keep online softmax,
so the output is exact either way.

Result

720p 81f T2V, 40 steps, v7x-8, dp2 x (ring2 x uly2):

attention	block_q / block_kv	Denoise total	s/step
ulysses_ring_custom (baseline)	9472 / 1024	109.3s	2.733
ulysses_ring_custom_fixed_m	6400 / 2048	105.3s	2.633

-3.7%. In-kernel the ring hop goes 22.65 -> 21.15 ms

Each mode uses its own optimal tile. Plain is MXU-bound and peaks at
block_q=9472, degrading to ~3.14 s/step at 6400/2048; fixed-m is flat across
both (2.63-2.69 s/step). That flatness is the mechanism working -- with the
softmax VPU chain gone the kernel stops caring which tile it was given.

Reproduce

One TPU v7x-8 host. run_wan_fast_inference.sh selects the attention kernel
and the matching tile:

./end_to_end/tpu/run_wan_fast_inference.sh 22 40            # fixed-m
FIXEDM=0 ./end_to_end/tpu/run_wan_fast_inference.sh 22 40   # baseline

Read Denoise Total from the timing summary. The first run per variant
compiles (~200s) and already reports steady-state latency; later runs reuse the
AOT cache. Give each variant its own WAN_CACHE_ROOT if you change kernel
source between runs -- the AOT fingerprint covers config, not source.

Notes

• 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 hop partials combine by accumulation; a gate failure
  falls back to an LSE-space merge, which is invariant to the overshoot.
• Emitting two bodies in the ragged last KV block degrades the whole grid's
  instruction schedule (measured 3x), so the all-heads-fixed branch compiles a
  single branch-free body.

Tests

src/maxdiffusion/tests/ring_fixed_m_test.py (5 cases, TPU) checks against an
f32 dense-softmax reference: the online path, an all-eligible call, a sink head
that falls back everywhere, a ragged KV tail, and a caller-supplied per-shard
max||k|| being promoted to the global bound.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.