fix(attention): dot-product fallback layout for 4-D inputs; tile_size_grid_search iters/bkv_compute guards

Two small fixes peeled out of the Ulysses x Ring PR so that they can be
reviewed on their own.

Dot-product fallback (`_apply_attention_dot`)
- Sequences below `flash_min_seq_length` bypass the flash/ulysses kernels and
  run the dot-product fallback. Callers that apply rotary embeddings hand the
  dispatcher `[B, H, S, D]` (the dispatcher reads S from axis 2 when
  `ndim == 4`), but the `split_head_dim` path reshaped those tensors as if
  they were `[B, S, H*D]`. Both layouts hold the same number of elements, so
  the reshape succeeded and returned a wrong answer silently. 4-D inputs are
  now transposed to `[B, S, H, D]` (`_to_bshd`); the flat 3-D contract is
  unchanged.
- `kv_heads` (optional, default None) is accepted so that K/V with fewer heads
  than Q are repeated on the head axis after the transpose, on both the
  `split_head_dim` and batched-heads paths; `heads % kv_heads != 0` raises
  ValueError instead of silently truncating the repeat factor.
  `dot_product_kernel` forwards `context.get("kv_heads")`.

tile_size_grid_search
- `time_callable` rejects `iters < 1` instead of dividing by zero.
- `_broadcast_winner` encodes `bkv_compute=None` as -1 in the broadcast
  payload and decodes it back, so a winner without an explicit compute block
  is found again on the non-source hosts.

Tests
- dot_fallback_layout_test.py (6, new, CPU): the reviewer's counterexample
  (uniform softmax must give per-head token means), random 4-D inputs vs a
  dense reference and GQA repeat on 4-D (both over split_head_dim in
  {True, False}), 3-D/4-D agreement, the heads % kv_heads guard, and the
  dispatcher seq-axis contract: `_apply_attention` on a [B, H=128, S=32, D]
  query with flash_min_seq_length=64 must route to a spied dot_product kernel
  (S read from axis 2) that then returns the dense reference.
- tile_size_grid_search_test.py (16, +2): `iters=0` raises; `_broadcast_winner`
  round-trips `bkv_compute=None`.
- Both files are added to end_to_end/tpu/run_wan_stack_tests.sh.
Verified (this head): TPU v6e-8 (jax 0.11.2), every file of
end_to_end/tpu/run_wan_stack_tests.sh at this head run directly with unittest:
custom_splash_fixed_m 31, ring_fixed_m 18, dot_fallback_layout 6,
tile_size_grid_search 16 -- 71 tests, all OK.

Author: Perseus14

end_to_end/tpu/run_wan_stack_tests.sh

@@ -28,5 +28,8 @@ 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"
+  # Dot-product fallback layout and tile-size grid-search fixes (feat/attn-dot-fallback-fixes)
+  "$T/dot_fallback_layout_test.py"
+  "$T/tile_size_grid_search_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/models/attention_flax.py

@@ -1725,17 +1725,52 @@ def _apply_attention_dot(
     float32_qk_product: bool,
     use_memory_efficient_attention: bool,
     attention_mask: Array = None,
+    kv_heads: int | None = None,
 ):
   """Apply Attention."""
+  effective_kv_heads = kv_heads if kv_heads is not None else heads
+  if heads <= 0 or effective_kv_heads <= 0:
+    raise ValueError(f"heads ({heads}) and kv_heads ({effective_kv_heads}) must be positive integers.")
+  if heads % effective_kv_heads:
+    raise ValueError(f"heads ({heads}) must be a multiple of kv_heads ({effective_kv_heads}).")
   if split_head_dim:
-    b = key.shape[0]
-    query_states = jnp.reshape(query, (b, -1, heads, dim_head))
-    key_states = jnp.reshape(key, (b, -1, heads, dim_head))
-    value_states = jnp.reshape(value, (b, -1, heads, dim_head))
+
+    def _to_bshd(x: Array, n_heads: int) -> Array:
+      """Normalise to [B, S, H, D].
+
+      Callers that apply rotary embeddings hand us [B, H, S, D] (see
+      `_unflatten_heads`), while the flat path supplies [B, S, H*D]. Only the
+      latter can be reshaped into [B, S, H, D]; reinterpreting [B, H, S, D]
+      that way keeps the shape legal but interleaves heads with tokens, so it
+      corrupts the output silently. Transpose the 4-D case instead.
+      """
+      if x.ndim == 4:
+        return jnp.swapaxes(x, 1, 2)
+      return jnp.reshape(x, (x.shape[0], -1, n_heads, dim_head))
+
+    query_states = _to_bshd(query, heads)
+    key_states = _to_bshd(key, effective_kv_heads)
+    value_states = _to_bshd(value, effective_kv_heads)
+    if heads != effective_kv_heads:
+      num_repeats = heads // effective_kv_heads
+      key_states = jnp.repeat(key_states, num_repeats, axis=2)
+      value_states = jnp.repeat(value_states, num_repeats, axis=2)
   else:
     query_states = _reshape_heads_to_batch_dim(query, heads)
-    key_states = _reshape_heads_to_batch_dim(key, heads)
-    value_states = _reshape_heads_to_batch_dim(value, heads)
+    key_states = _reshape_heads_to_batch_dim(key, effective_kv_heads)
+    value_states = _reshape_heads_to_batch_dim(value, effective_kv_heads)
+    if heads != effective_kv_heads:
+      num_repeats = heads // effective_kv_heads
+      b = query.shape[0]
+      s_k = key_states.shape[1]
+      key_states = jnp.repeat(key_states.reshape(b, effective_kv_heads, s_k, -1), num_repeats, axis=1).reshape(
+          b * heads, s_k, -1
+      )
+      value_states = jnp.repeat(
+          value_states.reshape(b, effective_kv_heads, s_k, -1),
+          num_repeats,
+          axis=1,
+      ).reshape(b * heads, s_k, -1)
 
   if float32_qk_product:
     query_states = query_states.astype(jnp.float32)
@@ -1841,6 +1876,7 @@ def dot_product_kernel(q, k, v, context):
       context["float32_qk_product"],
       context["use_memory_efficient_attention"],
       context["attention_mask"],
+      kv_heads=context.get("kv_heads", None),
   )
 
 
@@ -2165,6 +2201,7 @@ def _apply_attention(
     preserve_asymmetric_block_sizes: bool = False,
     spatiotemporal_config: Optional[dict] = None,
     spatiotemporal_shape: Optional[Tuple[int, int, int]] = None,
+    kv_heads: Optional[int] = None,
 ):
   """Routes to different attention kernels using a module-level registry."""
 
@@ -2213,6 +2250,7 @@ def _apply_attention(
 
   context = {
       "heads": heads,
+      "kv_heads": kv_heads,
       "mesh": mesh,
       "axis_names_q": axis_names_q,
       "axis_names_kv": axis_names_kv,

src/maxdiffusion/tests/dot_fallback_layout_test.py

@@ -0,0 +1,224 @@
+"""
+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.
+"""
+
+"""Layout regression tests for the short-sequence dot-product fallback.
+
+Sequences below `flash_min_seq_length` bypass the flash/ulysses kernels and
+run `_apply_attention_dot`. Callers that apply rotary embeddings hand the
+dispatcher `[B, H, S, D]` -- the dispatcher says so itself, reading the
+sequence length from axis 2 when `ndim == 4` -- but the `split_head_dim` path
+reshaped those tensors as if they were `[B, S, H*D]`.
+
+Both layouts hold the same number of elements, so the reshape succeeded and
+returned a wrong answer with no error. That silent corruption is what these
+tests exist to prevent. They are pure JAX and run on CPU.
+"""
+
+import math
+import unittest
+from unittest import mock
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxdiffusion.models import attention_flax
+from maxdiffusion.models.attention_flax import _apply_attention, _apply_attention_dot
+
+
+def _call_dot(query, key, value, heads, dim_head, kv_heads=None, split_head_dim=True):
+  """Runs the fallback and returns [B, H, S, D].
+
+  `_apply_attention_dot` emits the flat `[B, S, H*D]` form, so unflatten it
+  here rather than at each call site.
+  """
+  out = _apply_attention_dot(
+      query=query,
+      key=key,
+      value=value,
+      dtype=jnp.float32,
+      heads=heads,
+      dim_head=dim_head,
+      scale=1.0 / math.sqrt(dim_head),
+      split_head_dim=split_head_dim,
+      float32_qk_product=True,
+      use_memory_efficient_attention=False,
+      attention_mask=None,
+      kv_heads=kv_heads,
+  )
+  batch, seq, _ = out.shape
+  return jnp.swapaxes(out.reshape(batch, seq, heads, dim_head), 1, 2)
+
+
+def _reference_attention(query, key, value, scale):
+  """Dense f32 attention on explicit [B, H, S, D] inputs."""
+  q, k, v = (x.astype(jnp.float32) for x in (query, key, value))
+  logits = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale
+  return jnp.einsum("bhqk,bhkd->bhqd", jax.nn.softmax(logits, axis=-1), v)
+
+
+class DotFallbackLayoutTest(unittest.TestCase):
+  """`_apply_attention_dot` must transpose, not reshape, 4-D inputs."""
+
+  def test_zero_logits_return_per_head_token_means(self):
+    """The reviewer's counterexample, reproduced exactly.
+
+    One active head dimension, three tokens, two heads. Q = K = 0 makes every
+    logit zero, so softmax is uniform and each head's output is the mean of
+    its values over tokens:
+
+        head 0 values [1, 2, 3]  -> 2
+        head 1 values [10, 20, 30] -> 20
+
+    Reinterpreting [B, H, S, D] as [B, S, H, D] instead walks the buffer
+    [1, 2, 3, 10, 20, 30] as three token-pairs (1,2), (3,10), (20,30),
+    producing [8, 14].
+    """
+    heads, seq, dim_head = 2, 3, 1
+    shape = (1, heads, seq, dim_head)
+    query = jnp.zeros(shape, jnp.float32)
+    key = jnp.zeros(shape, jnp.float32)
+    value = jnp.array([[[[1.0], [2.0], [3.0]], [[10.0], [20.0], [30.0]]]], dtype=jnp.float32)
+    self.assertEqual(value.shape, shape)
+
+    out = np.asarray(_call_dot(query, key, value, heads, dim_head))
+
+    np.testing.assert_allclose(out[0, 0], 2.0, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(out[0, 1], 20.0, rtol=1e-5, atol=1e-5)
+    # The exact wrong answer the reshape produced.
+    self.assertFalse(np.allclose(out[0, 0], 8.0))
+    self.assertFalse(np.allclose(out[0, 1], 14.0))
+
+  def test_matches_reference_on_random_4d_inputs(self):
+    """Both the [B, S, H, D] and the batched-heads ([B*H, S, D]) paths."""
+    heads, seq, dim_head = 4, 8, 16
+    shape = (2, heads, seq, dim_head)
+    query = jax.random.normal(jax.random.PRNGKey(0), shape, jnp.float32)
+    key = jax.random.normal(jax.random.PRNGKey(1), shape, jnp.float32)
+    value = jax.random.normal(jax.random.PRNGKey(2), shape, jnp.float32)
+    expected = np.asarray(_reference_attention(query, key, value, 1.0 / math.sqrt(dim_head)))
+
+    for split_head_dim in (True, False):
+      with self.subTest(split_head_dim=split_head_dim):
+        out = np.asarray(_call_dot(query, key, value, heads, dim_head, split_head_dim=split_head_dim))
+        np.testing.assert_allclose(out, expected, rtol=1e-4, atol=1e-4)
+
+  def test_three_dim_inputs_agree_with_four_dim(self):
+    """The flat [B, S, H*D] contract must keep working, and agree."""
+    heads, seq, dim_head = 4, 8, 16
+    batch = 2
+    shape = (batch, heads, seq, dim_head)
+    query = jax.random.normal(jax.random.PRNGKey(3), shape, jnp.float32)
+    key = jax.random.normal(jax.random.PRNGKey(4), shape, jnp.float32)
+    value = jax.random.normal(jax.random.PRNGKey(5), shape, jnp.float32)
+
+    def flatten(x):
+      return jnp.swapaxes(x, 1, 2).reshape(batch, seq, heads * dim_head)
+
+    out_4d = np.asarray(_call_dot(query, key, value, heads, dim_head))
+    out_3d = np.asarray(_call_dot(flatten(query), flatten(key), flatten(value), heads, dim_head))
+    np.testing.assert_allclose(out_4d, out_3d, rtol=1e-5, atol=1e-5)
+
+  def test_gqa_repeat_still_applies_on_4d(self):
+    """Head repetition must happen after the transpose, on the head axis."""
+    heads, kv_heads, seq, dim_head = 4, 2, 8, 16
+    q = jax.random.normal(jax.random.PRNGKey(6), (1, heads, seq, dim_head), jnp.float32)
+    k = jax.random.normal(jax.random.PRNGKey(7), (1, kv_heads, seq, dim_head), jnp.float32)
+    v = jax.random.normal(jax.random.PRNGKey(8), (1, kv_heads, seq, dim_head), jnp.float32)
+
+    expected = np.asarray(
+        _reference_attention(
+            q,
+            jnp.repeat(k, heads // kv_heads, axis=1),
+            jnp.repeat(v, heads // kv_heads, axis=1),
+            1.0 / math.sqrt(dim_head),
+        )
+    )
+    for split_head_dim in (True, False):
+      with self.subTest(split_head_dim=split_head_dim):
+        out = np.asarray(_call_dot(q, k, v, heads, dim_head, kv_heads=kv_heads, split_head_dim=split_head_dim))
+        np.testing.assert_allclose(out, expected, rtol=1e-4, atol=1e-4)
+
+  def test_rejects_kv_heads_that_do_not_divide_heads(self):
+    """A non-integer repeat factor or non-positive head count must raise ValueError."""
+    heads, kv_heads, seq, dim_head = 4, 3, 8, 16
+    q = jnp.zeros((1, heads, seq, dim_head), jnp.float32)
+    kv = jnp.zeros((1, kv_heads, seq, dim_head), jnp.float32)
+    for split_head_dim in (True, False):
+      with self.subTest(split_head_dim=split_head_dim):
+        with self.assertRaisesRegex(ValueError, "multiple of kv_heads"):
+          _call_dot(q, kv, kv, heads, dim_head, kv_heads=kv_heads, split_head_dim=split_head_dim)
+        with self.assertRaisesRegex(ValueError, "positive integers"):
+          _call_dot(q, kv, kv, heads, dim_head, kv_heads=0, split_head_dim=split_head_dim)
+        with self.assertRaisesRegex(ValueError, "positive integers"):
+          _call_dot(q, kv, kv, heads, dim_head, kv_heads=-2, split_head_dim=split_head_dim)
+
+
+class DispatcherLayoutContractTest(unittest.TestCase):
+  """The threshold check and the dot path must agree on where S lives."""
+
+  def test_dispatcher_reads_seq_len_from_axis_2_and_dot_path_honours_it(self):
+    """`_apply_attention` reads S from axis 2 for 4-D inputs, i.e. [B, H, S, D].
+
+    The shape is chosen so the two readings disagree: H = 128 is above the
+    threshold while S = 32 is below it. Reading S from axis 2 routes to the
+    dot-product fallback, which must then interpret the same tensor as
+    [B, H, S, D] (the transpose, not the reshape). Reading S from axis 1 would
+    try to run the configured flash kernel instead and never reach the spy.
+    """
+    batch, heads, seq, dim_head = 1, 128, 32, 8
+    flash_min_seq_length = 64
+    shape = (batch, heads, seq, dim_head)
+    query = jax.random.normal(jax.random.PRNGKey(9), shape, jnp.float32)
+    key = jax.random.normal(jax.random.PRNGKey(10), shape, jnp.float32)
+    value = jax.random.normal(jax.random.PRNGKey(11), shape, jnp.float32)
+    scale = 1.0 / math.sqrt(dim_head)
+
+    real_kernel = attention_flax.KERNEL_REGISTRY["dot_product"]
+    spy = mock.Mock(side_effect=real_kernel)
+    with mock.patch.dict(attention_flax.KERNEL_REGISTRY, {"dot_product": spy}):
+      out = _apply_attention(
+          query,
+          key,
+          value,
+          heads=heads,
+          dim_head=dim_head,
+          split_head_dim=True,
+          float32_qk_product=True,
+          attention_kernel="ulysses_custom",
+          flash_min_seq_length=flash_min_seq_length,
+          use_memory_efficient_attention=False,
+          scale=scale,
+          dtype=jnp.float32,
+          mesh=None,
+          axis_names_q=None,
+          axis_names_kv=None,
+          flash_block_sizes=None,
+          dpa_layer=None,
+      )
+
+    spy.assert_called_once()
+    seen_query = spy.call_args.args[0]
+    self.assertEqual(seen_query.shape, shape)
+    self.assertLess(seen_query.shape[2], flash_min_seq_length)
+    self.assertGreaterEqual(seen_query.shape[1], flash_min_seq_length)
+    out = jnp.swapaxes(jnp.reshape(out, (batch, seq, heads, dim_head)), 1, 2)
+    expected = _reference_attention(query, key, value, scale)
+    np.testing.assert_allclose(np.asarray(out), np.asarray(expected), rtol=1e-4, atol=1e-4)
+
+
+if __name__ == "__main__":
+  unittest.main()

src/maxdiffusion/tests/tile_size_grid_search_test.py

@@ -129,6 +129,10 @@ def fake_fn():
     self.assertLess(mean, 30.0)  # ...and NOT in the steady-state mean
     self.assertEqual(len(times), 5)
 
+  def test_rejects_nonpositive_iters(self):
+    with self.assertRaises(ValueError):
+      time_callable(lambda: 1, iters=0)
+
 
 class OrchestratorTest(unittest.TestCase):
 
@@ -167,6 +171,19 @@ def test_oom_configs_pruned_not_raised(self):
     )
     self.assertTrue(any(r.status == "oom" for r in res.results))
 
+  def test_broadcast_winner_handles_none_bkv_compute(self):
+    from unittest import mock
+    from maxdiffusion.utils.tile_size_grid_search import _broadcast_winner
+
+    cand = BenchResult(bq=1024, bkv=512, bkv_compute=None, status="ok", mean_ms=10.0)
+    with (
+        mock.patch("jax.process_count", return_value=2),
+        mock.patch("jax.process_index", return_value=0),
+        mock.patch("jax.experimental.multihost_utils.broadcast_one_to_all", side_effect=lambda x, is_source: x),
+    ):
+      out = _broadcast_winner(cand, [cand])
+    self.assertIs(out, cand)
+
 
 if __name__ == "__main__":
   unittest.main()

src/maxdiffusion/utils/tile_size_grid_search.py

@@ -259,6 +259,8 @@ def time_callable(fn, *, iters: int = 10, warmup: int = 2, sync=lambda x: x):
   import statistics
   import time
 
+  if iters < 1:
+    raise ValueError(f"iters must be >= 1, got {iters}.")
   t0 = time.perf_counter()
   sync(fn())  # call #1: compilation happens HERE, untimed
   compile_ms = (time.perf_counter() - t0) * 1e3
@@ -380,12 +382,13 @@ def _broadcast_winner(best: Optional[BenchResult], results: list[BenchResult]) -
   is_source = jax.process_index() == 0
   payload = np.zeros((4,), dtype=np.int64)
   if is_source and best is not None:
-    payload[:] = (1, best.bq, best.bkv, best.bkv_compute)
+    bkv_cmp = -1 if best.bkv_compute is None else best.bkv_compute
+    payload[:] = (1, best.bq, best.bkv, bkv_cmp)
   payload = np.asarray(multihost_utils.broadcast_one_to_all(payload, is_source=is_source))
   if payload[0] == 0:
     return None
 
-  winner_key = tuple(int(value) for value in payload[1:])
+  winner_key = (int(payload[1]), int(payload[2]), None if int(payload[3]) < 0 else int(payload[3]))
   for result in results:
     if (result.bq, result.bkv, result.bkv_compute) == winner_key and result.status == "ok":
       return result