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

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/models/ltx2/attention_ltx2.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")