Remove _resolve_tpu_attention_block_sizes, consolidate into _select_flash_block_sizes

Author: csgoogle

src/maxdiffusion/models/attention_flax.py

@@ -270,38 +270,6 @@ def convert_to_tokamax_splash_config(
   )
 
 
-def _resolve_tpu_attention_block_sizes(
-    query_seq_len: int,
-    kv_seq_len: int,
-    flash_block_sizes: BlockSizes,
-    dtype: jnp.dtype,
-    attention_kernel: str = "flash",
-) -> BlockSizes:
-  """Resolve TPU splash attention block sizes for self- and cross-attention."""
-  q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
-  is_cross_attention = kv_seq_len != query_seq_len
-  if is_cross_attention:
-    kv_max_block_size = ((kv_seq_len + 127) // 128) * 128
-  else:
-    kv_max_block_size = q_max_block_size
-
-  if flash_block_sizes and not is_cross_attention:
-    return flash_block_sizes
-
-  block_size_q = flash_block_sizes.block_q if flash_block_sizes else q_max_block_size
-  return splash_attention_kernel.BlockSizes(
-      block_q=block_size_q,
-      block_kv_compute=min(kv_max_block_size, kv_seq_len),
-      block_kv=min(kv_max_block_size, kv_seq_len),
-      block_q_dkv=block_size_q,
-      block_kv_dkv=min(kv_max_block_size, kv_seq_len),
-      block_kv_dkv_compute=min(kv_max_block_size, query_seq_len),
-      block_q_dq=None if attention_kernel == "tokamax_flash" else block_size_q,
-      block_kv_dq=None if attention_kernel == "tokamax_flash" else min(kv_max_block_size, query_seq_len),
-      use_fused_bwd_kernel=True if attention_kernel == "tokamax_flash" else False,
-  )
-
-
 def _tpu_flash_attention(
     query: jax.Array,
     key: jax.Array,
@@ -319,18 +287,11 @@ def _tpu_flash_attention(
 ) -> jax.Array:
   """TPU Flash Attention"""
 
-  block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, attention_kernel)
   num_context_shards = mesh.shape["context"]
   query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards)
   key, _ = _reshape_data_for_flash(key, heads, num_context_shards)
   value, _ = _reshape_data_for_flash(value, heads, num_context_shards)
-  block_sizes = _resolve_tpu_attention_block_sizes(
-      query_seq_len=query.shape[2],
-      kv_seq_len=key.shape[2],
-      flash_block_sizes=flash_block_sizes,
-      dtype=dtype,
-      attention_kernel=attention_kernel,
-  )
+  block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, attention_kernel)
 
   q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
   kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
@@ -530,12 +491,7 @@ def _ulysses_attention(
   q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
   kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
 
-  block_sizes = _resolve_tpu_attention_block_sizes(
-      query_seq_len=query.shape[2],
-      kv_seq_len=key.shape[2],
-      flash_block_sizes=flash_block_sizes,
-      dtype=dtype,
-  )
+  block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "flash")
 
   @functools.partial(
       jax.shard_map,

src/maxdiffusion/tests/attention_test.py

@@ -186,21 +186,30 @@ def test_default_flash_block_sizes_use_sequence_axis_for_3d_inputs(self):
     assert block_sizes.block_q_dq == 1024
     assert block_sizes.block_kv_dq == 128
 
-  def test_resolve_tpu_attention_block_sizes(self):
-    """Shared block-size selection should keep self-attn overrides and derive cross-attn defaults."""
+  def test_select_flash_block_sizes_returns_configured_for_self_attention(self):
+    """Block-size selection should return the configured sizes unchanged for self-attention."""
     custom_block_sizes = self._ulysses_block_sizes(block_size=16)
+    query = jnp.zeros((1, 128, 1), dtype=jnp.float32)
+    key = jnp.zeros((1, 128, 1), dtype=jnp.float32)
 
-    self_attention_block_sizes = attention_flax._resolve_tpu_attention_block_sizes(
-        query_seq_len=128,
-        kv_seq_len=128,
+    self_attention_block_sizes = _select_flash_block_sizes(
+        query=query,
+        key=key,
         flash_block_sizes=custom_block_sizes,
         dtype=jnp.float32,
+        attention_kernel="flash",
     )
     self.assertIs(self_attention_block_sizes, custom_block_sizes)
 
-    cross_attention_block_sizes = attention_flax._resolve_tpu_attention_block_sizes(
-        query_seq_len=257,
-        kv_seq_len=513,
+  def test_select_flash_block_sizes_derives_cross_attn_defaults_for_tokamax(self):
+    """Block-size selection should derive cross-attn defaults and set tokamax_flash flags."""
+    custom_block_sizes = self._ulysses_block_sizes(block_size=16)
+    query = jnp.zeros((1, 257, 1), dtype=jnp.float32)
+    key = jnp.zeros((1, 513, 1), dtype=jnp.float32)
+
+    cross_attention_block_sizes = _select_flash_block_sizes(
+        query=query,
+        key=key,
         flash_block_sizes=custom_block_sizes,
         dtype=jnp.float32,
         attention_kernel="tokamax_flash",