From 92da629b17d0891122d78984bfe5d68461834c84 Mon Sep 17 00:00:00 2001 From: jcarin-sys Date: Mon, 17 Aug 2026 13:53:59 +0000 Subject: [PATCH] Route DeepSeek-V4 compressed and sliding train attention onto the splash kernels Opt-in TPU path (compressed_use_dynamic_splash) that avoids materializing dense seq^2 logits. At seq 16384 the dense path needs 96.8 GiB per chip and does not compile on v6e. Co-authored-by: Sudarsanan Co-authored-by: Armin Co-authored-by: utlz --- src/maxtext/configs/base.yml | 3 + src/maxtext/configs/types.py | 10 + src/maxtext/layers/attention_op.py | 165 +++++++++- tests/unit/attention_test.py | 5 +- tests/unit/dynamic_splash_mask_test.py | 436 +++++++++++++++++++++++++ 5 files changed, 613 insertions(+), 6 deletions(-) create mode 100644 tests/unit/dynamic_splash_mask_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index a4cf6877e3..5ce7ff1f88 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -457,6 +457,9 @@ o_lora_rank: 0 # Output LoRA rank for Compressed Attention. o_groups: 0 # Output groups for Compressed Attention. compress_ratios: [] # Per-layer compression ratios (0, 4, 128, etc). compressed_rope_max_timescale: 160_000 # If positive, used for Compressed Sparse/Heavy Attention. +# Route COMPRESSED and LOCAL_SLIDING train attention to tokamax splash on TPU. +# Requires use_tokamax_splash: true. +compressed_use_dynamic_splash: false # QK-Clip (Muon Clip) Configuration use_qk_clip: false # Enable QK-Clip (supported in MLA with DotProduct or Tokamax Splash) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1f8659f7e0..99ea6debf2 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -708,6 +708,11 @@ class CompressedAttention(BaseModel): compressed_rope_max_timescale: int = Field( 160000, description="If positive, used for Compressed Sparse/Heavy Attention." ) + compressed_use_dynamic_splash: bool = Field( + False, + description="Route COMPRESSED and LOCAL_SLIDING train attention to tokamax splash on TPU. Requires " + "use_tokamax_splash=true.", + ) class AttentionIndexer(BaseModel): @@ -3678,6 +3683,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("MoBA is only supported with dot_product attention.") if self.decoder_block == DecoderBlockType.DEEPSEEK4 and self.attention != "dot_product": raise ValueError("DeepSeek4 decoder block currently only supports dot_product attention.") + if self.compressed_use_dynamic_splash and not self.use_tokamax_splash: + raise ValueError( + "`compressed_use_dynamic_splash` requires `use_tokamax_splash=true`; without it the boolean mask " + "would be silently ignored by the non-tokamax splash branches." + ) if self.mla_qk_head_chunk_size > 0: if self.mla_qk_head_chunk_size > self.num_query_heads or self.num_query_heads % self.mla_qk_head_chunk_size != 0: raise ValueError( diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 81beb97020..1958126d4b 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -134,6 +134,59 @@ def apply_mask_to_logits(logits: Array, mask: Array): return jnp.where((mask >= DEFAULT_MASK_VALUE * 0.5), logits, DEFAULT_MASK_VALUE) +def build_local_sliding_splash_mask( + batch: int, + decoder_segment_ids: Array | None, + segment_positions: Array | None, + q_seq_len: int, + kv_seq_len: int, + sliding_window_size: int | None, +) -> Array: + """Builds the causal sliding mask for the uncompressed COMPRESSED prefix. + + To match the dense path under packing, query positions may reset by segment while key + positions remain physical offsets. + """ + if segment_positions is not None: + row_ids = segment_positions[:, :, None] + else: + row_ids = jnp.arange(q_seq_len)[None, :, None] + col_ids = jnp.arange(kv_seq_len)[None, None, :] + distance = row_ids - col_ids + mask = distance >= 0 + if sliding_window_size is not None: + mask &= distance < sliding_window_size + mask = jnp.broadcast_to(mask, (batch, q_seq_len, kv_seq_len)) + if decoder_segment_ids is not None: + segment = decoder_segment_ids[:, :, None] == decoder_segment_ids[:, None, :] + mask = jnp.logical_and(mask, segment[..., :kv_seq_len]) + return mask + + +def build_compressed_splash_mask( + compressed_mask: Array, + decoder_segment_ids: Array | None, + segment_positions: Array | None, + q_seq_len: int, + kv_seq_len: int, + sliding_window_size: int | None, +) -> Array: + """Builds the boolean COMPRESSED mask for dynamic splash attention. + + The result has shape `[batch, q_seq_len, kv_seq_len]` and matches the dense path's + `apply_mask_to_logits` keep predicate. + """ + c_len = compressed_mask.shape[-1] + s_len = kv_seq_len - c_len + b = compressed_mask.shape[0] + + uncompressed = build_local_sliding_splash_mask( + b, decoder_segment_ids, segment_positions, q_seq_len, s_len, sliding_window_size + ) + compressed_keep = compressed_mask.reshape(b, q_seq_len, c_len) >= DEFAULT_MASK_VALUE * 0.5 + return jnp.concatenate([uncompressed, compressed_keep], axis=-1) + + def validate_gpu_flash_attention(sinks: Array | None, record_max_logits: bool) -> None: """Helper function to check for unsupported features with flash attention on GPU.""" if sinks is not None: @@ -568,7 +621,8 @@ def __init__( raise ValueError("causal_block_size must be positive for block-diffusion attention") if self.attention_kernel not in ("autoselected", "dot_product", "flash"): raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") - # Block sizes are only used by TPU splash attention kernels. Exclude non-splash kernels + # The opt-in COMPRESSED and LOCAL_SLIDING routes need splash block sizes despite using + # attention_kernel="dot_product". if self.attention_kernel not in ( "dot_product", "paged", @@ -576,6 +630,9 @@ def __init__( "vllm_batched_rpa", "cudnn_flash_te", "cudnn_flash_jax", + ) or ( + self.attention_type in (AttentionType.COMPRESSED, AttentionType.LOCAL_SLIDING) + and self.config.compressed_use_dynamic_splash ): if self.attention_type == AttentionType.LOCAL_SLIDING: self.block_q = self.config.local_sa_block_q @@ -1387,6 +1444,44 @@ def apply_attention( self.max_logits = nnx.Intermediate(local_max) return local_out, local_max, local_sum + # LOCAL_SLIDING has a static splash mask; COMPRESSED needs the dynamic mask from its + # compressor. Dense-only mask modifiers must continue through the dot-product path. + elif ( + self.config.compressed_use_dynamic_splash + and model_mode == MODEL_MODE_TRAIN + and target_hardware == "tpu" + and previous_chunk is None + and bidirectional_mask is None + and ( + self.attention_type == AttentionType.LOCAL_SLIDING + or (self.attention_type == AttentionType.COMPRESSED and compressed_mask is not None) + ) + ): + if self.attention_type == AttentionType.LOCAL_SLIDING: + out, max_logits = self.tpu_flash_attention( + query, + key, + value, + decoder_segment_ids, + self.attn_logits_soft_cap, + sinks, + model_mode=model_mode, + record_max_logits=record_max_logits, + ) + if max_logits is not None: + self.max_logits = nnx.Intermediate(max_logits) + return out, None, None + return self.dynamic_splash_attention( + query, + key, + value, + decoder_segment_ids, + segment_positions, + compressed_mask, + sinks, + record_max_logits, + ) + # 'vllm_rpa' uses the same dot-attention wrapper but routes to the vLLM # ragged paged attention kernel in `Attention.__call__`. elif ( @@ -1597,6 +1692,65 @@ def wrap_ragged_attention(query, key, value, lengths, block_size): return wrap_ragged_attention(query, key, value, lengths, block_size) + def dynamic_splash_attention( + self, + query: Array, + key: Array, + value: Array, + decoder_segment_ids: Array | None, + segment_positions: Array | None, + compressed_mask: Array, + sinks: Array | None, + record_max_logits: bool = False, + ) -> tuple[Array, None, None]: + """Runs COMPRESSED train attention with the tokamax dynamic splash kernel. + + Packing is folded into the boolean mask because no segment ids exist for the compressed + KV columns. KV and mask columns are padded to satisfy all splash block sizes. + """ + q_seq_len = query.shape[1] + kv_seq_len = key.shape[1] + + if self.mesh.shape.get(self.config.context_sharding, 1) > 1: + raise NotImplementedError( + "compressed_use_dynamic_splash does not support context parallelism for COMPRESSED " + "attention (the dynamic-mask splash path has no load-balanced reorder for the " + "compressed kv concat). LOCAL_SLIDING layers take the static splash path and do." + ) + for name, blk in (("block_q", self.block_q), ("block_q_dkv", self.block_q_dkv)): + eff = min(blk, q_seq_len) + if q_seq_len % eff: + raise ValueError( + f"compressed_use_dynamic_splash: query length {q_seq_len} must be a multiple of " + f"min({name}={blk}, q_len) = {eff}." + ) + + bool_mask = build_compressed_splash_mask( + compressed_mask, decoder_segment_ids, segment_positions, q_seq_len, kv_seq_len, self.sliding_window_size + ) + + lcm = math.lcm(self.block_kv, self.block_kv_compute, self.block_kv_dkv, self.block_kv_dkv_compute) + padded_kv_len = math.ceil(kv_seq_len / lcm) * lcm + if padded_kv_len != kv_seq_len: + pad = padded_kv_len - kv_seq_len + key = jnp.pad(key, ((0, 0), (0, pad), (0, 0), (0, 0))) + value = jnp.pad(value, ((0, 0), (0, pad), (0, 0), (0, 0))) + bool_mask = jnp.pad(bool_mask, ((0, 0), (0, 0), (0, pad))) + + out, max_logits = self.tpu_flash_attention( + query, + key, + value, + decoder_segment_ids=None, + attn_logits_soft_cap=self.attn_logits_soft_cap, + sinks=sinks, + indexer_mask=bool_mask, + record_max_logits=record_max_logits, + ) + if max_logits is not None: + self.max_logits = nnx.Intermediate(max_logits) + return out, None, None + def tpu_flash_attention( self, query: Array, @@ -2068,8 +2222,9 @@ def wrap_flash_attention( decoder_segment_ids_tuple = None if self.config.use_tokamax_splash: - if self.config.use_indexer and indexer_mask is not None: - # Construct the splash kernel call with dynamic mask + if indexer_mask is not None: + # Pallas interpret mode does not propagate input_output_aliases across grid steps, so + # dynamic-splash GQA gradients must be checked on TPU. def dynamic_mask_splash_kernel(q, k, v, segment, sinks, indexer_mask): splash_kernel = tokamax_splash_kernel.make_dynamic_splash_mha( mask=indexer_mask, @@ -2083,9 +2238,9 @@ def dynamic_mask_splash_kernel(q, k, v, segment, sinks, indexer_mask): else: return kernel(q, k, v, segment, sinks=sinks), None - # Iterate over batch dimension for (query, key, value, segment, sinks, mask) attn_fn = jax.vmap(dynamic_mask_splash_kernel, (0, 0, 0, 0, None, 0)) - indexer_mask = jnp.isclose(indexer_mask, 0.0) + if indexer_mask.dtype != jnp.bool_: + indexer_mask = jnp.isclose(indexer_mask, 0.0) if record_max_logits: attention_output, max_logits = attn_fn(query, key, value, decoder_segment_ids_tuple, sinks, indexer_mask) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index e04dc80cc8..ec745f6191 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -391,6 +391,7 @@ def _make_flash_op( sa_use_base2_exp=False, use_tokamax_splash=False, use_jax_splash=False, + compressed_use_dynamic_splash=False, ) device = types.SimpleNamespace(platform="cpu") mesh = types.SimpleNamespace( @@ -713,7 +714,9 @@ def test_load_balanced_block_causal_mask(self): np.testing.assert_array_equal(mask[:, :], expected) def test_dot_product_local_mask_uses_segment_positions(self): - config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context") + config = types.SimpleNamespace( + context_parallel_load_balance=True, context_sharding="context", compressed_use_dynamic_splash=False + ) mesh = types.SimpleNamespace(shape={"context": 4}) seq_len = 16 sliding_window_size = 4 diff --git a/tests/unit/dynamic_splash_mask_test.py b/tests/unit/dynamic_splash_mask_test.py new file mode 100644 index 0000000000..6db31e9f95 --- /dev/null +++ b/tests/unit/dynamic_splash_mask_test.py @@ -0,0 +1,436 @@ +# 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. + +"""Tests for the COMPRESSED / LOCAL_SLIDING dynamic-splash boolean masks.""" + +import math +import types +import unittest +from unittest import mock + +from absl.testing import parameterized +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.common.common_types import ( + AttentionType, + DEFAULT_MASK_VALUE, + MODEL_MODE_AUTOREGRESSIVE, + MODEL_MODE_PREFILL, + MODEL_MODE_TRAIN, +) +from maxtext.layers.attention_op import ( + AttentionOp, + build_compressed_splash_mask, + build_local_sliding_splash_mask, +) + + +def _stub_config(compressed_use_dynamic_splash=False, block=8): + """A minimal stand-in for the dynamic-splash configuration.""" + kernel_fields = { + f"sa_block_{n}": block for n in ("q", "kv", "kv_compute", "q_dkv", "kv_dkv", "kv_dkv_compute", "q_dq", "kv_dq") + } + kernel_fields.update({f"sa_{n}_layout": "HEAD_DIM_MINOR" for n in ("q", "k", "v")}) + kernel_fields.update(sa_use_fused_bwd_kernel=True, sa_fuse_reciprocal=False, sa_use_base2_exp=False) + kernel_fields["use_splash_scheduler"] = False + kernel_fields.update({f"local_{name}": value for name, value in kernel_fields.items()}) + return types.SimpleNamespace( + compressed_use_dynamic_splash=compressed_use_dynamic_splash, + context_parallel_load_balance=False, + context_parallel_strategy="all_gather", + context_sharding="context", + **kernel_fields, + ) + + +def _make_op(attention_type, sliding_window_size, *, config=None, platform="cpu", context_parallel_size=1): + """Builds an AttentionOp whose only job is mask generation / dispatch.""" + device = types.SimpleNamespace(platform=platform) + mesh = types.SimpleNamespace( + devices=np.asarray([device], dtype=object), + shape={"context": context_parallel_size}, + ) + return AttentionOp( + config=config if config is not None else _stub_config(), + num_query_heads=1, + num_kv_heads=1, + max_target_length=32, + mesh=mesh, + attention_kernel="dot_product", + attention_type=attention_type, + sliding_window_size=sliding_window_size, + ) + + +def _dense_keep_set(additive_mask): + """Converts a dense additive mask to its keep set.""" + return np.asarray(additive_mask)[:, 0, 0, :, :] >= DEFAULT_MASK_VALUE * 0.5 + + +def _packing(seq_len, batch, packed): + """Returns segment ids and per-segment positions, which reset under packing.""" + if not packed: + return None, jnp.asarray(np.arange(seq_len)[None, :]).repeat(batch, axis=0) + segment_len = seq_len // 2 + segment_ids = jnp.asarray(np.arange(seq_len)[None, :] // segment_len + 1).repeat(batch, axis=0) + reset_positions = jnp.asarray((np.arange(seq_len) % segment_len)[None, :]).repeat(batch, axis=0) + return segment_ids, reset_positions + + +def _additive_lattice(rng, shape): + """Builds a compressed-block mask, including the keep predicate's boundary.""" + choices = np.asarray([0.0, DEFAULT_MASK_VALUE * 0.5, DEFAULT_MASK_VALUE, -np.inf], dtype=np.float32) + return jnp.asarray(choices[rng.integers(0, len(choices), size=shape)]) + + +class DynamicSplashMaskParityTest(parameterized.TestCase): + """The boolean splash mask must keep exactly what the dense additive mask keeps.""" + + @pytest.mark.cpu_only + @parameterized.product( + seq_len=(16, 32), + sliding_window_size=(4, 16), + packed=(False, True), + ) + def test_local_sliding_mask_matches_dense(self, seq_len, sliding_window_size, packed): + batch = 2 + segment_ids, _ = _packing(seq_len, batch, packed) + op = _make_op(AttentionType.LOCAL_SLIDING, sliding_window_size) + dense = op.generate_attention_mask( + jnp.zeros((batch, seq_len, 1, 1)), + jnp.zeros((batch, seq_len, 1, 1)), + segment_ids, + MODEL_MODE_TRAIN, + ) + expected = np.broadcast_to(_dense_keep_set(dense), (batch, seq_len, seq_len)) + actual = build_local_sliding_splash_mask(batch, segment_ids, None, seq_len, seq_len, sliding_window_size) + np.testing.assert_array_equal(np.asarray(actual), expected) + + @pytest.mark.cpu_only + @parameterized.product( + seq_len=(16, 32), + sliding_window_size=(4, 16, None), + packed=(False, True), + compress_ratio=(4, 8), + positions=("none", "physical", "reset"), + ) + def test_compressed_mask_matches_dense(self, seq_len, sliding_window_size, packed, compress_ratio, positions): + batch = 2 + compressed_len = seq_len // compress_ratio + kv_seq_len = seq_len + compressed_len + segment_ids, reset_positions = _packing(seq_len, batch, packed) + segment_positions = { + "none": None, + "physical": jnp.asarray(np.arange(seq_len)[None, :]).repeat(batch, axis=0), + "reset": reset_positions, + }[positions] + compressed_mask = _additive_lattice(np.random.default_rng(0), (batch, 1, 1, seq_len, compressed_len)) + op = _make_op(AttentionType.COMPRESSED, sliding_window_size) + dense = op.generate_attention_mask( + jnp.zeros((batch, seq_len, 1, 1)), + jnp.zeros((batch, kv_seq_len, 1, 1)), + segment_ids, + MODEL_MODE_TRAIN, + compressed_mask=compressed_mask, + segment_positions=segment_positions, + ) + expected = np.broadcast_to(_dense_keep_set(dense), (batch, seq_len, kv_seq_len)) + actual = build_compressed_splash_mask( + compressed_mask, segment_ids, segment_positions, seq_len, kv_seq_len, sliding_window_size + ) + np.testing.assert_array_equal(np.asarray(actual), expected) + + @pytest.mark.cpu_only + def test_packed_reset_positions_retain_no_uncompressed_columns(self): + """Pins the dense path's reset-query/physical-key asymmetry under packing.""" + batch, seq_len, window, compressed_len = 1, 16, 4, 4 + kv_seq_len = seq_len + compressed_len + segment_ids, reset_positions = _packing(seq_len, batch, packed=True) + compressed_mask = jnp.zeros((batch, 1, 1, seq_len, compressed_len), dtype=jnp.float32) + op = _make_op(AttentionType.COMPRESSED, window) + dense = op.generate_attention_mask( + jnp.zeros((batch, seq_len, 1, 1)), + jnp.zeros((batch, kv_seq_len, 1, 1)), + segment_ids, + MODEL_MODE_TRAIN, + compressed_mask=compressed_mask, + segment_positions=reset_positions, + ) + keep = np.broadcast_to(_dense_keep_set(dense), (batch, seq_len, kv_seq_len)) + self.assertEqual(keep[:, seq_len // 2 :, :seq_len].sum(), 0) + self.assertEqual(keep[:, : seq_len // 2, :seq_len].sum(), 26) + actual = build_compressed_splash_mask(compressed_mask, segment_ids, reset_positions, seq_len, kv_seq_len, window) + np.testing.assert_array_equal(np.asarray(actual), keep) + + +class DynamicSplashAttentionTest(unittest.TestCase): + """Checks what dynamic_splash_attention hands to the kernel, without running the kernel.""" + + def _capture(self, op, batch, q_seq_len, kv_seq_len, compressed_mask, decoder_segment_ids=None, segment_positions=None): + """Runs dynamic_splash_attention with a stubbed kernel and returns the captured call kwargs.""" + captured = {} + + def fake_tpu_flash_attention(query, key, value, **kwargs): + captured.update(kwargs, key=key, value=value) + return jnp.zeros_like(query), None + + with mock.patch.object(op, "tpu_flash_attention", fake_tpu_flash_attention): + out, exp_max, exp_sum = op.dynamic_splash_attention( + jnp.zeros((batch, q_seq_len, 1, 4)), + jnp.zeros((batch, kv_seq_len, 1, 4)), + jnp.zeros((batch, kv_seq_len, 1, 4)), + decoder_segment_ids, + segment_positions, + compressed_mask, + None, + ) + self.assertIsNone(exp_max) + self.assertIsNone(exp_sum) + self.assertEqual(out.shape, (batch, q_seq_len, 1, 4)) + return captured + + @pytest.mark.cpu_only + def test_kernel_call_pads_kv_folds_packing_into_the_mask_and_guards_its_assumptions(self): + batch, seq_len, compressed_len, block = 2, 32, 8, 16 + kv_seq_len = seq_len + compressed_len + segment_ids, reset_positions = _packing(seq_len, batch, packed=True) + op = _make_op(AttentionType.COMPRESSED, 4, config=_stub_config(True, block=block)) + compressed_mask = _additive_lattice(np.random.default_rng(1), (batch, 1, 1, seq_len, compressed_len)) + captured = self._capture(op, batch, seq_len, kv_seq_len, compressed_mask, segment_ids, reset_positions) + + self.assertIsNone(captured["decoder_segment_ids"]) + padded_kv_len = math.ceil(kv_seq_len / block) * block + self.assertGreater(padded_kv_len, kv_seq_len) + mask = np.asarray(captured["indexer_mask"]) + self.assertEqual(mask.dtype, np.bool_) + self.assertEqual(mask.shape, (batch, seq_len, padded_kv_len)) + self.assertEqual(captured["key"].shape[1], padded_kv_len) + self.assertEqual(captured["value"].shape[1], padded_kv_len) + self.assertFalse(mask[:, :, kv_seq_len:].any()) + unpadded = build_compressed_splash_mask(compressed_mask, segment_ids, reset_positions, seq_len, kv_seq_len, 4) + np.testing.assert_array_equal(mask[:, :, :kv_seq_len], np.asarray(unpadded)) + + small = _additive_lattice(np.random.default_rng(2), (1, 1, 1, 24, compressed_len)) + with self.assertRaisesRegex(ValueError, "must be a multiple of"): + self._capture(op, 1, 24, 24 + compressed_len, small) + cp_op = _make_op(AttentionType.COMPRESSED, 4, config=_stub_config(True, block=block), context_parallel_size=2) + with self.assertRaisesRegex(NotImplementedError, "context parallelism"): + self._capture(cp_op, batch, seq_len, kv_seq_len, compressed_mask) + + +class DynamicSplashRoutingTest(parameterized.TestCase): + """Tests dynamic, static, and dense routing.""" + + def _dispatch( + self, + *, + flag, + model_mode, + platform, + attention_type, + compressed_mask, + previous_chunk=None, + bidirectional_mask=None, + ): + """Returns 'dynamic' (dynamic-mask splash), 'static' (static splash) or 'dot' (dense path).""" + op = _make_op(attention_type, 4, config=_stub_config(flag), platform=platform) + with ( + mock.patch.object(op, "dynamic_splash_attention", return_value=(None, None, None)) as dynamic, + mock.patch.object(op, "tpu_flash_attention", return_value=(None, None)) as static, + mock.patch.object(op, "apply_attention_dot", return_value=(None, None, None)) as dot, + ): + op.apply_attention( + jnp.zeros((1, 32, 1, 4)), + jnp.zeros((1, 32, 1, 4)), + jnp.zeros((1, 32, 1, 4)), + None, + None, + None, + model_mode, + previous_chunk=previous_chunk, + bidirectional_mask=bidirectional_mask, + compressed_mask=compressed_mask, + qk_product_einsum=jnp.einsum, + wv_product_einsum=jnp.einsum, + ) + self.assertEqual(dynamic.called + static.called + dot.called, 1) + if dynamic.called: + return "dynamic" + return "static" if static.called else "dot" + + @pytest.mark.cpu_only + @parameterized.named_parameters( + ("compressed", AttentionType.COMPRESSED, "dynamic"), + ("local_sliding", AttentionType.LOCAL_SLIDING, "static"), + ) + def test_gate(self, attention_type, enabled_target): + compressed_mask = jnp.zeros((1, 1, 1, 32, 8)) if attention_type == AttentionType.COMPRESSED else None + kwargs = {"attention_type": attention_type, "compressed_mask": compressed_mask} + on_tpu_train = {"flag": True, "model_mode": MODEL_MODE_TRAIN, "platform": "tpu", **kwargs} + self.assertEqual(self._dispatch(**on_tpu_train), enabled_target) + self.assertEqual(self._dispatch(flag=False, model_mode=MODEL_MODE_TRAIN, platform="tpu", **kwargs), "dot") + self.assertEqual(self._dispatch(flag=True, model_mode=MODEL_MODE_PREFILL, platform="tpu", **kwargs), "dot") + self.assertEqual(self._dispatch(flag=True, model_mode=MODEL_MODE_AUTOREGRESSIVE, platform="tpu", **kwargs), "dot") + self.assertEqual(self._dispatch(flag=True, model_mode=MODEL_MODE_TRAIN, platform="cpu", **kwargs), "dot") + self.assertEqual(self._dispatch(previous_chunk=object(), **on_tpu_train), "dot") + self.assertEqual(self._dispatch(bidirectional_mask=jnp.ones((1, 32), dtype=bool), **on_tpu_train), "dot") + + +def _v4_test_config(compressed_use_dynamic_splash, *extra_overrides): + """A deepseek4 config small enough to instantiate in a test, with the flag set either way.""" + from maxtext.configs.pyconfig import initialize # pylint: disable=import-outside-toplevel + from tests.utils.test_helpers import get_test_config_path # pylint: disable=import-outside-toplevel + + overrides = ( + "model_name=deepseek4-284b attention=dot_product qk_rope_head_dim=16 v_head_dim=16 qk_nope_head_dim=16 " + "use_tokamax_splash=True override_model_config=True " + f"compressed_use_dynamic_splash={compressed_use_dynamic_splash}" + ).split() + return initialize([None, get_test_config_path(), *overrides, *extra_overrides]) + + +class DynamicSplashEndToEndTest(parameterized.TestCase): + """Dense-vs-splash parity for a whole compressed attention layer (needs the Pallas TPU kernel).""" + + def _layer_output(self, compressed_use_dynamic_splash, compress_ratio, packed): + """Runs one CompressedAttention layer with the flag set either way and returns its output.""" + from flax import nnx # pylint: disable=import-outside-toplevel + from jax.sharding import Mesh # pylint: disable=import-outside-toplevel + from maxtext.layers.attention_compressed import CompressedAttention # pylint: disable=import-outside-toplevel + + seq_len = 128 + layer = CompressedAttention( + config=_v4_test_config(compressed_use_dynamic_splash), + num_query_heads=4, + num_kv_heads=1, + head_dim=512, + max_target_length=seq_len, + mesh=Mesh(np.array(jax.devices()[:1]), ("data",)), + attention_kernel="dot_product", + inputs_q_shape=(1, seq_len, 4096), + inputs_kv_shape=(1, seq_len, 4096), + compress_ratio=compress_ratio, + sliding_window_size=seq_len // 2, + q_lora_rank=1024, + rngs=nnx.Rngs(0), + ) + inputs = jax.random.normal(jax.random.PRNGKey(1), (1, seq_len, 4096), dtype=jnp.float32) + segment_ids = jnp.ones((1, seq_len), dtype=jnp.int32) + positions = jnp.arange(seq_len)[None, :] + if packed: + segment_ids = segment_ids.at[:, seq_len // 2 :].set(2) + positions = positions % (seq_len // 2) + return layer(inputs, inputs, segment_ids, positions, deterministic=True)[0] + + @pytest.mark.tpu_only + @parameterized.product( + compress_ratio=(0, 4, 128), + packed=(False, True), + ) + def test_dense_and_dynamic_splash_agree(self, compress_ratio, packed): + dense = self._layer_output(False, compress_ratio, packed) + splash = self._layer_output(True, compress_ratio, packed) + np.testing.assert_allclose(np.asarray(splash), np.asarray(dense), rtol=2e-2, atol=2e-2) + + +class DynamicSplashBackwardParityTest(unittest.TestCase): + """Tests COMPRESSED GQA gradients on TPU. + + Pallas interpret mode does not propagate the kernel's aliased gradient accumulators. + """ + + def _forward_and_grads(self, compressed_use_dynamic_splash, q, k, v, compressed_mask, decoder_segment_ids, cotangent): + """Returns (out, dq, dk, dv) for one COMPRESSED AttentionOp with the flag set either way.""" + from jax.sharding import Mesh # pylint: disable=import-outside-toplevel + + blocks = [f"sa_block_{n}=128" for n in ("q", "kv", "kv_compute", "q_dkv", "kv_dkv", "kv_dkv_compute")] + op = AttentionOp( + config=_v4_test_config(compressed_use_dynamic_splash, *blocks), + mesh=Mesh(np.array(jax.devices()[:1]), ("data",)), + attention_kernel="dot_product", + max_target_length=q.shape[1], + num_query_heads=q.shape[2], + num_kv_heads=k.shape[2], + attention_type=AttentionType.COMPRESSED, + sliding_window_size=64, + ) + + def loss_fn(q, k, v): + out, _, exp_sum = op.apply_attention( + q, + k, + v, + decoder_segment_ids, + None, + None, + MODEL_MODE_TRAIN, + compressed_mask=compressed_mask, + qk_product_einsum=jnp.einsum, + wv_product_einsum=jnp.einsum, + ) + if exp_sum is not None: + out = out / exp_sum + return jnp.sum(out * cotangent), out + + (_, out), grads = jax.value_and_grad(loss_fn, argnums=(0, 1, 2), has_aux=True)(q, k, v) + return (out, *grads) + + @pytest.mark.tpu_only + def test_gradients_match_the_dense_path(self): + batch, q_seq_len, compressed_len, heads, kv_heads, head_dim = 1, 128, 32, 4, 1, 128 + kv_seq_len = q_seq_len + compressed_len + keys = jax.random.split(jax.random.PRNGKey(0), 4) + q = jax.random.normal(keys[0], (batch, q_seq_len, heads, head_dim), dtype=jnp.float32) + k = jax.random.normal(keys[1], (batch, kv_seq_len, kv_heads, head_dim), dtype=jnp.float32) + v = jax.random.normal(keys[2], (batch, kv_seq_len, kv_heads, head_dim), dtype=jnp.float32) + cotangent = jax.random.normal(keys[3], (batch, q_seq_len, heads, head_dim), dtype=jnp.float32) + block_of = jnp.arange(compressed_len) * (q_seq_len // compressed_len) + keep = block_of[None, :] <= jnp.arange(q_seq_len)[:, None] + compressed_mask = jnp.where(keep, 0.0, DEFAULT_MASK_VALUE)[None, None, None, :, :] + segment_ids = jnp.ones((batch, q_seq_len), dtype=jnp.int32) + + args = (q, k, v, compressed_mask, segment_ids, cotangent) + dense = self._forward_and_grads(False, *args) + splash = self._forward_and_grads(True, *args) + for name, got, want in zip(("out", "dq", "dk", "dv"), splash, dense): + got = np.asarray(got, dtype=np.float64).ravel() + want = np.asarray(want, dtype=np.float64).ravel() + rel_l2 = np.linalg.norm(got - want) / np.linalg.norm(want) + cosine = np.dot(got, want) / (np.linalg.norm(got) * np.linalg.norm(want)) + self.assertLess(rel_l2, 5e-2, f"{name}: relative L2 distance {rel_l2:.3e}") + self.assertGreater(cosine, 0.999, f"{name}: cosine similarity {cosine:.6f}") + + +class DynamicSplashConfigValidatorTest(unittest.TestCase): + """compressed_use_dynamic_splash without use_tokamax_splash must be rejected at config time.""" + + @pytest.mark.cpu_only + def test_flag_requires_tokamax_splash(self): + from maxtext.configs.pyconfig import initialize # pylint: disable=import-outside-toplevel + from tests.utils.test_helpers import get_test_config_path # pylint: disable=import-outside-toplevel + + with self.assertRaisesRegex(Exception, "use_tokamax_splash"): + initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + compressed_use_dynamic_splash=True, + use_tokamax_splash=False, + ) + + +if __name__ == "__main__": + unittest.main()