Merge pull request #477 from AI-Hypercomputer:feat/fixed-m-kernel

PiperOrigin-RevId: 1007276103

Author: Google-ML-Automation

end_to_end/tpu/run_wan_stack_tests.sh

@@ -0,0 +1,33 @@
+#!/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). Convenience entry point for running the
+# Pallas kernel grids and multi-device ring/Ulysses tests natively on a TPU
+# generation other than the CI runner's; the same tests also run in CI (the
+# multi-device ones self-skip below their device count).
+#
+# 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/configuration_utils.py

@@ -496,9 +496,7 @@ def extract_init_dict(cls, config_dict, **kwargs):
     diffusers_library = importlib.import_module(__name__.split(".")[0])
 
     if cls.has_compatibles:
-      compatible_classes = [
-          c for c in cls._get_compatibles() if not isinstance(c, DummyObject)  # type: ignore
-      ]
+      compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]  # type: ignore
     else:
       compatible_classes = []
 

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -479,6 +479,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")