From ee90fa17e916fa4e01de9ebc09e68b1c5ffa299b Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Thu, 15 Oct 2026 05:38:39 +0000 Subject: [PATCH] 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. --- end_to_end/tpu/run_wan_stack_tests.sh | 3 + src/maxdiffusion/models/attention_flax.py | 50 +++- .../tests/dot_fallback_layout_test.py | 224 ++++++++++++++++++ .../tests/tile_size_grid_search_test.py | 17 ++ .../utils/tile_size_grid_search.py | 7 +- 5 files changed, 293 insertions(+), 8 deletions(-) create mode 100644 src/maxdiffusion/tests/dot_fallback_layout_test.py diff --git a/end_to_end/tpu/run_wan_stack_tests.sh b/end_to_end/tpu/run_wan_stack_tests.sh index 886c8f74a..ea53d0778 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/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[@]}" "$@" diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 8e19536da..514cfa032 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/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, diff --git a/src/maxdiffusion/tests/dot_fallback_layout_test.py b/src/maxdiffusion/tests/dot_fallback_layout_test.py new file mode 100644 index 000000000..a431f52de --- /dev/null +++ b/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() diff --git a/src/maxdiffusion/tests/tile_size_grid_search_test.py b/src/maxdiffusion/tests/tile_size_grid_search_test.py index 3419c142f..ff2f250cc 100644 --- a/src/maxdiffusion/tests/tile_size_grid_search_test.py +++ b/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() diff --git a/src/maxdiffusion/utils/tile_size_grid_search.py b/src/maxdiffusion/utils/tile_size_grid_search.py index 3b7afb90b..bc6c74573 100644 --- a/src/maxdiffusion/utils/tile_size_grid_search.py +++ b/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