feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks

Implements exact fixed-m splash attention in Pallas on TPU. Instead of an
online running max, a head / Q-block whose Cauchy-Schwarz bound
max||q|| * max||k - k_bar|| is within a length-dependent window uses a
precomputed m = ceil(bound) - C(N) (+ q . k_bar with virtual K-centering) and
skips the per-block rescales; everything else keeps online softmax.

- `get_fixed_m_constants(N)`: length-dependent recenter C(N) and safe bound
  that keep the f32 numerator/denominator from overflowing or flushing.
- `fixed_m_dtype_is_safe`: rejects fp16/fp8 (exponent range too small for the
  recentered exponent); bf16 and fp32 pass.
- V-magnitude gate: |V| <= DEFAULT_MAX_V_BOUND (256), otherwise online.
- Per-(head, Q-block) dispatch in the kernel; the Ulysses and R=1 ring callers
  compile a uniform (all-fixed) and a hybrid kernel and pick one with
  `lax.cond` on `all_fixed`. Both kernels are compiled; the extra compile
  time / program size has not been measured.
- Virtual K-centering: the kernel receives k_mean and adds q . k_bar to the
  fixed bound in VMEM (an MXU dot_general on an (Hkv, 8, d) tiled k_mean,
  precision=HIGHEST because the gate has no slack at its edge). K itself is
  never re-written in HBM. k_mean is padded to the lane-padded head_dim after
  `_pad_data_for_flash`, so head_dim < 128 works (regression test at d=64).

Explicit metadata contracts on the fixed-m ring path
----------------------------------------------------
1. Norm representation is declared, not inferred. The gate previously guessed
   whether `fixed_m_norms` were squared with `(qn.max() * mk.max()) < 1000.0`.
   Magnitude cannot answer that: legacy unsquared norms of (1000, 2) have a
   true bound of 2000, but read as squared give sqrt(2000) ~= 44.7, a ~45x
   under-estimate that admits fixed-m where it must fall back. Replaced by a
   REQUIRED `fixed_m_norms_squared` argument (no default; omitting it with
   fixed-m on raises). Tests use squared norms.

2. The V-safety predicate is required, not assumed. The V-magnitude and dtype
   verdict cannot be re-derived from one hop's Q/K, so the ring kernel raises
   unless `v_ok` is passed (`value=None` in the metadata helper gives
   v_ok = 0, fail-closed). `_ulysses_ring_custom_attention` (R > 1) computes it
   on the PRE-all-to-all shards -- dtype safety plus |V| <= 256 -- and reduces
   it with pmin over BOTH internal axes (ring, ulysses): before the a2a each
   device holds a sequence shard of every head, and the ring's
   accumulate-vs-LSE branch must be taken uniformly by every ppermute
   participant. Only Q and K go through the optimization_barrier that keeps
   XLA from duplicating their producer chains into the norm reductions; V is
   reduced outside it so the Q/K a2a does not wait on V's producer.

3. Norm shape is validated against per_q_block. Both gates compute
   `qn * mk[:, None]`, so a (num_heads,) array supplied with per_q_block=True
   broadcasts to (num_heads, num_heads) and pairs head j's query norm with head
   h's key norm; a sink head then inherits a small bound and the kernel
   evaluates exp2(large_logit - small_m) -> inf. The kernel now rejects the
   mismatch.

Known limitations
- On the uncentered R > 1 ring path the fixed bound can still flush the
  numerator for heads far below the bound; this is not measured here.
- The ring kernel runs its own pmin over global-bound eligibility (added by
  this commit) unless the caller promises `uniform_fixed_m=True`; on the
  Ulysses+Ring path that pmin is redundant with the caller's (not measured).

Tests
- custom_splash_fixed_m_test.py (31): kernel unit tests with relative-L2
  checks, k-norm (mk) shape and V-metadata fail-closed checks, integration
  through `_ulysses_attention` and the R=1 ring (uniform branch, hybrid
  branch via a sink head, d=64 -- the test re-derives the gate from the
  inputs and asserts it admits fixed-m, it does not observe the cond -- ragged
  KV tails with N % 8 == 0 and N % block_kv != 0 whose tail is an odd multiple
  of 8), and a U=2/R=2 check that a single V outlier on one pre-a2a device
  makes v_ok False on every device (ring kernel stubbed; needs >= 4 devices,
  explicit skip otherwise).
- ring_fixed_m_test.py (18): 8 backend-independent contract tests (including
  the required fixed_m_norms_squared and v_ok arguments), 6 ring numerical
  tests vs an f32 reference (2 devices), 1 test that declared-unsquared and
  squared norms give the same result, and 3 tests showing the raw-K bound is
  unsound where the centered bound is not.
- CI: to keep GitHub Actions short, 18 of these 49 tests (the kernel
  numerics grids and the multi-device integration tests) skip when
  GITHUB_ACTIONS=true. The 9 tests from main and the contract, dtype and
  V-metadata tests still run there. The new
  end_to_end/tpu/run_wan_stack_tests.sh runs all of them on a TPU VM; later
  PRs in this stack add their test files to it.

Verified (final tree): TPU v6e-8 (jax 0.11.2, end_to_end/tpu/run_wan_stack_tests.sh): 49 passed, 6 subtests passed in 170.4s.

Author: Perseus14

end_to_end/tpu/run_wan_stack_tests.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Runs the tests of the Wan fast-inference kernels and serving path on a TPU VM
+# (validated on v6e-8 and tpu7x-8), including the Pallas kernel grids and
+# multi-device ring/Ulysses tests that GitHub Actions skips (they check
+# GITHUB_ACTIONS=true) to keep CI short.
+#
+# Usage, from the repo root of an installed checkout:
+#   bash end_to_end/tpu/run_wan_stack_tests.sh [extra pytest args, e.g. -x --durations=20]
+set -euo pipefail
+cd "$(dirname "$0")/../.."
+unset GITHUB_ACTIONS
+T=src/maxdiffusion/tests
+TESTS=(
+  # Fixed-m custom splash kernel and ring (feat/fixed-m-kernel)
+  "$T/custom_splash_fixed_m_test.py"
+  "$T/ring_fixed_m_test.py"
+)
+PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -3,7 +3,8 @@ hardware: 'tpu'
 skip_jax_distributed_system: False
 # Supported attention kernels:
 # dot_product, flash, tokamax_flash, tokamax_ring, tokamax_ring_custom,
-# ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_ring,
+# ulysses, ulysses_custom, ulysses_custom_fixed_m,
+# ulysses_custom_fixed_m_per_q_block, ulysses_ring,
 # ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_bidir,
 # and cudnn_flash_te (GPU only).
 attention: 'flash'

src/maxdiffusion/configs/ltx2_video.yml

@@ -3,7 +3,8 @@ hardware: 'tpu'
 skip_jax_distributed_system: False
 # Supported attention kernels:
 # dot_product, flash, tokamax_flash, tokamax_ring, tokamax_ring_custom,
-# ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_ring,
+# ulysses, ulysses_custom, ulysses_custom_fixed_m,
+# ulysses_custom_fixed_m_per_q_block, ulysses_ring,
 # ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_bidir,
 # and cudnn_flash_te (GPU only).
 attention: 'flash'

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -468,6 +468,7 @@ def __init__(
         "ulysses_ring_custom_bidir",
         "ulysses_custom",
         "ulysses_custom_fixed_m",
+        "ulysses_custom_fixed_m_per_q_block",
     )
     cross_attention_uses_local_kv = not is_self_attention and (
         cross_attention_remapped_to_flash or attention_kernel in ("flash", "tokamax_flash", "cudnn_flash_te")