DSV4: Route compressed and sliding train attention onto the splash kernels

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)

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(

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,14 +621,18 @@ 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",
         "vllm_rpa",
         "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)

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