Fix transformer sharding, flash block sizing, and tests

Author: csgoogle

src/maxdiffusion/models/attention_flax.py

@@ -190,6 +190,49 @@ def _pad_data_for_flash(tensor, heads, flash_block_size, num_shards: int = 1):
   return tensor, kv_size, seq_len
 
 
+def _flash_sequence_length(tensor: Array) -> int:
+  if tensor.ndim == 3:
+    return tensor.shape[1]
+  if tensor.ndim == 4:
+    return tensor.shape[2]
+  raise ValueError(f"Flash attention expects rank-3 or rank-4 inputs, got rank {tensor.ndim}.")
+
+
+def _select_flash_block_sizes(
+    query: Array,
+    key: Array,
+    flash_block_sizes: BlockSizes,
+    dtype: jnp.dtype,
+    attention_kernel: str,
+) -> BlockSizes:
+  query_seq_len = _flash_sequence_length(query)
+  key_seq_len = _flash_sequence_length(key)
+
+  q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
+  if key_seq_len != query_seq_len:
+    kv_max_block_size = ((key_seq_len + 127) // 128) * 128
+  else:
+    kv_max_block_size = q_max_block_size
+
+  # keep configured block sizes for self-attention, but let
+  # cross-attention derive safe KV-aware sizes when q_len != kv_len.
+  if flash_block_sizes and key_seq_len == query_seq_len:
+    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, key_seq_len),
+      block_kv=min(kv_max_block_size, key_seq_len),
+      block_q_dkv=block_size_q,
+      block_kv_dkv=min(kv_max_block_size, key_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 convert_to_tokamax_splash_config(
     block_sizes: BlockSizes,
     q_layout: tokamax_splash_attention_kernel.QKVLayout = tokamax_splash_attention_kernel.QKVLayout.HEAD_DIM_MINOR,
@@ -244,28 +287,7 @@ def _tpu_flash_attention(
 ) -> jax.Array:
   """TPU Flash Attention"""
 
-  q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
-  # This is the case for cross-attn.
-  if key.shape[1] != query.shape[1]:
-    kv_max_block_size = ((key.shape[1] + 127) // 128) * 128
-  else:
-    kv_max_block_size = q_max_block_size
-  # ensure that for cross attention we override the block sizes.
-  if flash_block_sizes and key.shape[1] == query.shape[1]:
-    block_sizes = flash_block_sizes
-  else:
-    block_size_q = flash_block_sizes.block_q if flash_block_sizes else q_max_block_size
-    block_sizes = splash_attention_kernel.BlockSizes(
-        block_q=block_size_q,
-        block_kv_compute=min(kv_max_block_size, key.shape[2]),
-        block_kv=min(kv_max_block_size, key.shape[2]),
-        block_q_dkv=block_size_q,
-        block_kv_dkv=min(kv_max_block_size, key.shape[2]),
-        block_kv_dkv_compute=min(kv_max_block_size, query.shape[2]),
-        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.shape[2]),
-        use_fused_bwd_kernel=True if attention_kernel == "tokamax_flash" else False,
-    )
+  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)
@@ -979,7 +1001,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("embed",),
+            ("heads",),
         ),
     )
 
@@ -993,7 +1015,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("embed",),
+            ("heads",),
         ),
     )
 
@@ -1007,7 +1029,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("embed",),
+            ("heads",),
         ),
     )
 
@@ -1021,7 +1043,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("heads",),
+            ("embed",),
         ),
     )
 
@@ -1263,8 +1285,7 @@ def __call__(
 
     with jax.named_scope("proj_attn"):
       hidden_states = self.proj_attn(attn_output)
-      if self.drop_out.rate > 0:
-        hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs)
+      hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs)
     return hidden_states
 
 
@@ -1333,11 +1354,13 @@ def setup(self):
         precision=self.precision,
     )
 
+    proj_attn_kernel_axes = ("heads", "embed")
+
     self.proj_attn = nn.Dense(
         self.query_dim,
-        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), kernel_axes),
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), proj_attn_kernel_axes),
         use_bias=True,
-        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("heads",)),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("embed",)),
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
         name="i_proj",
@@ -1346,9 +1369,9 @@ def setup(self):
 
     self.encoder_proj_attn = nn.Dense(
         self.query_dim,
-        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), kernel_axes),
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), proj_attn_kernel_axes),
         use_bias=True,
-        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("heads",)),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("embed",)),
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
         name="e_proj",

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -193,11 +193,11 @@ def __init__(
         kernel_init=nnx.with_partitioning(
             nnx.initializers.xavier_uniform(),
             (
-                "mlp",
                 "embed",
+                "mlp",
             ),
         ),
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("mlp",)),
     )
 
   def __call__(self, x: jax.Array) -> jax.Array:
@@ -249,8 +249,8 @@ def __init__(
         kernel_init=nnx.with_partitioning(
             nnx.initializers.xavier_uniform(),
             (
+                 "mlp",
                 "embed",
-                "mlp",
             ),
         ),
     )
@@ -262,8 +262,7 @@ def conditional_named_scope(self, name: str):
   def __call__(self, hidden_states: jax.Array, deterministic: bool = True, rngs: nnx.Rngs = None) -> jax.Array:
     hidden_states = self.act_fn(hidden_states)  # Output is (4, 75600, 13824)
     hidden_states = checkpoint_name(hidden_states, "ffn_activation")
-    if self.drop_out.rate > 0:
-      hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs)
+    hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs)
     with jax.named_scope("proj_out"):
       return self.proj_out(hidden_states)  # output is (4, 75600, 5120)
 

src/maxdiffusion/tests/attention_test.py

@@ -20,8 +20,9 @@
 import jax
 from jax.sharding import Mesh
 import jax.numpy as jnp
-from ..models.attention_flax import FlaxAttention
+from ..common_types import BlockSizes
 from .. import max_utils
+from ..models.attention_flax import FlaxAttention, _select_flash_block_sizes
 from .. import pyconfig
 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -35,6 +36,8 @@ def setUp(self):
 
   def test_splash_attention(self):
     """Test numerics of splash attention are equivalent to dot_product"""
+    if jax.devices()[0].platform != "tpu":
+      self.skipTest("TPU splash attention test requires a TPU backend.")
 
     pyconfig.initialize(
         [
@@ -92,6 +95,62 @@ def test_splash_attention(self):
 
     assert diff_norm < 1.0
 
+  def test_select_flash_block_sizes_keeps_self_attention_config(self):
+    flash_block_sizes = BlockSizes(
+        block_q=2048,
+        block_kv=1024,
+        block_kv_compute=1024,
+        block_q_dkv=2048,
+        block_kv_dkv=1024,
+        block_kv_dkv_compute=1024,
+        block_q_dq=2048,
+        block_kv_dq=1024,
+    )
+
+    query = jnp.ones((1, 4096, 128), dtype=jnp.bfloat16)
+    key = jnp.ones((1, 4096, 128), dtype=jnp.bfloat16)
+
+    selected = _select_flash_block_sizes(query, key, flash_block_sizes, jnp.bfloat16, "flash")
+
+    self.assertEqual(selected, flash_block_sizes)
+
+  def test_select_flash_block_sizes_overrides_cross_attention_kv_blocks(self):
+    flash_block_sizes = BlockSizes(
+        block_q=2048,
+        block_kv=2048,
+        block_kv_compute=1024,
+        block_q_dkv=2048,
+        block_kv_dkv=2048,
+        block_kv_dkv_compute=1024,
+        block_q_dq=2048,
+        block_kv_dq=1024,
+    )
+
+    query = jnp.ones((1, 4096, 128), dtype=jnp.bfloat16)
+    key = jnp.ones((1, 512, 128), dtype=jnp.bfloat16)
+
+    selected = _select_flash_block_sizes(query, key, flash_block_sizes, jnp.bfloat16, "flash")
+
+    self.assertEqual(selected.block_q, flash_block_sizes.block_q)
+    self.assertEqual(selected.block_q_dkv, flash_block_sizes.block_q)
+    self.assertEqual(selected.block_q_dq, flash_block_sizes.block_q)
+    self.assertEqual(selected.block_kv, 512)
+    self.assertEqual(selected.block_kv_compute, 512)
+    self.assertEqual(selected.block_kv_dkv, 512)
+    self.assertEqual(selected.block_kv_dkv_compute, 512)
+    self.assertEqual(selected.block_kv_dq, 512)
+
+  def test_select_flash_block_sizes_uses_sequence_axis_for_rank_4_inputs(self):
+    query = jnp.ones((1, 4, 4096, 128), dtype=jnp.bfloat16)
+    key = jnp.ones((1, 4, 512, 128), dtype=jnp.bfloat16)
+
+    selected = _select_flash_block_sizes(query, key, None, jnp.bfloat16, "flash")
+
+    self.assertEqual(selected.block_q, 1024)
+    self.assertEqual(selected.block_kv, 512)
+    self.assertEqual(selected.block_kv_compute, 512)
+    self.assertEqual(selected.block_kv_dkv_compute, 512)
+
 
 if __name__ == "__main__":
   absltest.main()