rewrite logic

Author: eltsai

src/maxdiffusion/models/attention_flax.py

@@ -21,14 +21,17 @@
 import jax
 from jax.ad_checkpoint import checkpoint_name
 import jax.numpy as jnp
+from jax.experimental import shard_map
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
+from maxdiffusion.kernels.splash_attention import splash_attention_mask as tokamax_splash_attention_mask
 from maxdiffusion.kernels.splash_attention import splash_attention_kernel as tokamax_splash_attention_kernel
+from maxdiffusion.kernels.splash_attention import ring_attention_kernel as tokamax_ring_attention_kernel
+from maxdiffusion.kernels.splash_attention import base as tokamax_splash_base
 from einops import rearrange
 from .. import common_types, max_logging
 
 from . import custom_splash_attention as custom_splash
-
 from . import quantizations
 from .modeling_flax_utils import get_activation
 
@@ -311,42 +314,191 @@ def _tpu_flash_attention(
     flash_block_sizes: BlockSizes,
     dtype: jnp.dtype = jnp.float32,
     attention_kernel: str = "flash",
-    mask_padding_tokens: bool = False,
+    mask_padding_tokens: bool = True,
     residual_checkpoint_name: str | None = None,
     attention_mask: jax.Array = None,
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
 ) -> jax.Array:
-  """Torchax-Style Tensor Parallel TPU Flash Attention"""
-
-  # 1. Reshape from (Batch, Seq_Len, Heads*Dim) -> (Batch, Heads, Seq_Len, Dim)
-  # We pass num_context_shards=1 because we hold the full sequence locally.
-  query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards=1)
-  key, _ = _reshape_data_for_flash(key, heads, num_context_shards=1)
-  value, _ = _reshape_data_for_flash(value, heads, num_context_shards=1)
-
-  # 2. Call the built-in Torchax shard_map wrapper!
-  # This automatically handles the SPMD boundaries, the min() padding logic for
-  # cross-attention, and the axis swapping.
-  x = custom_splash.tpu_custom_attention(
-      query=query,
-      key=key,
-      value=value,
+  """TPU Flash Attention"""
+
+  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 = _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)
+
+  @functools.partial(
+      shard_map.shard_map,
       mesh=mesh,
-      # CRITICAL: JAX already scaled `key` by 1/sqrt(d) in the router.
-      # Passing 1.0 ensures the kernel only applies the log2(e) multiplier to `query`.
-      scale=1.0,
-      block_q=2048,
-      block_kv=2048,
-      block_kv_compute=1024,
-      block_kv_compute_in=256,
-      heads_per_tile=1,
+      in_specs=(q_axis_names, kv_axis_names, kv_axis_names),
+      out_specs=q_axis_names,
+      check_rep=False,
   )
+  def wrap_flash_attention(query, key, value):
+    uses_fused_kernel = block_sizes.use_fused_bwd_kernel
+    block_q_sizes = (
+        block_sizes.block_q,
+        block_sizes.block_q_dkv,
+    )
+    block_kv_sizes = (
+        block_sizes.block_kv,
+        block_sizes.block_kv_dkv,
+    )
+    if uses_fused_kernel:
+      block_q_sizes += (block_sizes.block_q_dkv,)
+      block_kv_sizes += (block_sizes.block_kv_dkv,)
+    else:
+      block_q_sizes += (block_sizes.block_q_dq,)
+      block_kv_sizes += (block_sizes.block_kv_dq,)
 
-  # 3. Trim back to original sequence length
-  x = x[:, :, :orig_q_seq_len, :]
+    block_q = max(*block_q_sizes)
+    query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, block_q)
 
-  # 4. Flatten back to (Batch, Seq_Len, Heads * Head_Dim)
+    block_kv = max(*block_kv_sizes)
+    key, _, key_seq_len = _pad_data_for_flash(key, heads, block_kv)
+    value, _, _ = _pad_data_for_flash(value, heads, block_kv)
+
+    mask = splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]))
+    multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
+
+    q_padded_len = query.shape[2]
+    q_indices = jax.lax.broadcasted_iota(jnp.int32, (q_padded_len,), 0)
+    q_segment_ids = (q_indices < query_seq_len).astype(jnp.int32)
+
+    kv_padded_len = key.shape[2]
+    kv_indices = jax.lax.broadcasted_iota(jnp.int32, (kv_padded_len,), 0)
+    kv_segment_ids = (kv_indices < key_seq_len).astype(jnp.int32)
+
+    # If attention_mask is provided, apply it to kv_segment_ids
+    if attention_mask is not None:
+      mask_len = min(key_seq_len, attention_mask.shape[1])
+      kv_mask_for_batch = attention_mask[0, :mask_len]  # (mask_len,)
+      # If key_seq_len > mask_len, pad the mask with 1s (assume remaining tokens are valid)
+      if key_seq_len > mask_len:
+        extra_valid = jnp.ones((key_seq_len - mask_len,), dtype=jnp.int32)
+        kv_mask_for_batch = jnp.concatenate([kv_mask_for_batch, extra_valid], axis=0)  # (key_seq_len,)
+      # Pad to kv_padded_len
+      if kv_padded_len > key_seq_len:
+        padding = jnp.zeros((kv_padded_len - key_seq_len,), dtype=jnp.int32)
+        kv_mask_padded = jnp.concatenate([kv_mask_for_batch, padding], axis=0)  # (kv_padded_len,)
+      else:
+        kv_mask_padded = kv_mask_for_batch
+      # Both are (kv_padded_len,) - element-wise multiplication
+      kv_segment_ids = (kv_segment_ids * kv_mask_padded).astype(jnp.int32)
+
+    if attention_kernel == "tokamax_ring":
+      segment_ids = tokamax_splash_base.SegmentIds(q=q_segment_ids, kv=kv_segment_ids)
+    else:
+      segment_ids = splash_attention_kernel.SegmentIds(q=q_segment_ids, kv=kv_segment_ids)
+
+    # make_splash_mha is wrapped around shardmap and seq and head is already
+    # sharded based on in_specs, therefore setting head_shards=1 and q_seq_shards=1.
+    if attention_kernel == "tokamax_flash":
+      mask = tokamax_splash_attention_mask.FullMask(
+          _shape=(query.shape[2], key.shape[2]),
+      )
+      splash_kernel = tokamax_splash_attention_kernel.make_splash_mha(
+          mask=mask,
+          q_seq_shards=1,  # the sizes of the axis is sharding over seq_len
+          config=convert_to_tokamax_splash_config(
+              block_sizes,
+              residual_checkpoint_name=residual_checkpoint_name,
+              use_base2_exp=use_base2_exp,
+              use_experimental_scheduler=use_experimental_scheduler,
+          ),
+          save_residuals=False,
+      )
+    elif attention_kernel == "tokamax_ring":
+      mask = tokamax_splash_attention_mask.FullMask(
+          _shape=(query.shape[2], key.shape[2]),
+      )
+      splash_kernel = tokamax_ring_attention_kernel.make_ring_attention(
+          mask=mask,
+          is_mqa=False,
+          config=convert_to_tokamax_splash_config(
+              block_sizes,
+              residual_checkpoint_name=residual_checkpoint_name,
+              use_base2_exp=use_base2_exp,
+              use_experimental_scheduler=use_experimental_scheduler,
+          ),
+          save_residuals=False,
+          ring_axis="context",
+          rotate_segment_ids=False,  # We don't rotate segment ids in tokamax ring attention because our segment ids is for padding each kv shard has same segment ids
+      )
+    else:
+      splash_kernel = splash_attention_kernel.make_splash_mha(
+          mask=multi_head_mask,
+          head_shards=1,  # the sizes of the axis is sharding over heads
+          q_seq_shards=1,  # the sizes of the axis is sharding over seq_len
+          block_sizes=block_sizes,
+          save_residuals=True if "ring" in attention_kernel else False,
+          residual_checkpoint_name=residual_checkpoint_name,
+      )
+
+    vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))
+
+    if not mask_padding_tokens:
+      segment_ids = None
+    if attention_kernel in ["flash", "tokamax_flash", "tokamax_ring"]:
+      attention_output = vmapped_splash(query, key, value, segment_ids)
+    else:
+      if num_context_shards > 1:
+        out, (lse,) = vmapped_splash(query, key, value, segment_ids)
+        m = lse.astype(jnp.float32)
+        l = jnp.exp(lse - m)
+        o = out.astype(jnp.float32) * l[..., None]
+
+        perm = [(j, (j + 1) % num_context_shards) for j in range(num_context_shards)]
+
+        k1 = jax.lax.ppermute(key, axis_name="context", perm=perm)
+        v1 = jax.lax.ppermute(value, axis_name="context", perm=perm)
+
+        def ring_scan_body(carry, _):
+          m, l, o, k_current, v_current = carry
+          k_next = jax.lax.ppermute(k_current, axis_name="context", perm=perm)
+          v_next = jax.lax.ppermute(v_current, axis_name="context", perm=perm)
+
+          out_chunk, (lse_chunk,) = vmapped_splash(query, k_current, v_current, segment_ids)
+
+          m_chunk = lse_chunk.astype(jnp.float32)
+          m_old = m
+          m = jnp.maximum(m_old, m_chunk)
+
+          exp_m_diff = jnp.exp(m_old - m)
+          exp_m_chunk_diff = jnp.exp(m_chunk - m)
+
+          l = l * exp_m_diff + jnp.exp(lse_chunk - m)
+          o = o * exp_m_diff[..., None]
+          o += exp_m_chunk_diff[..., None] * out_chunk.astype(jnp.float32)
+
+          # Return the updated state for the next iteration
+          return (m, l, o, k_next, v_next), None
+
+        initial_carry = (m, l, o, k1, v1)
+        (m_final, l_final, o_final, _, _), _ = jax.lax.scan(
+            ring_scan_body, initial_carry, None, length=num_context_shards - 1
+        )
+
+        attention_output = o_final / l_final[..., None]
+      else:
+        raise ValueError("ring attention requires context > 1")
+    return attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype)
+
+  devices_in_batch_sharding = mesh.shape["data"] * (mesh.shape["fsdp"] if "fsdp" in mesh.shape else 1)
+  # This warning might show up when doing model eval for example, when calculating model flops
+  # and that is expected.
+  if not (query.shape[0] / devices_in_batch_sharding).is_integer():
+    max_logging.log(
+        "Warning, batch dimension should be shardable among the devices in data and fsdp"
+        f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}"
+    )
+  x = wrap_flash_attention(query, key, value)
+  # Trim back to original sequence length after context-axis padding.
+  x = x[:, :, :orig_q_seq_len, :]
   x = _reshape_heads_to_head_dim(x)
 
   return x
@@ -429,6 +581,7 @@ def wrap_ulysses_attention(query, key, value):
     block_kv = max(*block_kv_sizes)
     key, _, key_seq_len = _pad_data_for_flash(key, heads, block_kv)
     value, _, _ = _pad_data_for_flash(value, heads, block_kv)
+
     mask = splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]))
     multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
 
@@ -577,10 +730,6 @@ def wrap_ulysses_attention(query, key, value):
   return x
 
 
-def torchax_attention():
-  pass
-
-
 def _apply_attention_dot(
     query: Array,
     key: Array,
@@ -692,7 +841,7 @@ def _apply_attention(
     axis_names_kv: AxisNames,
     flash_block_sizes: BlockSizes,
     dpa_layer: Callable,
-    mask_padding_tokens: bool = False,
+    mask_padding_tokens: bool = True,
     residual_checkpoint_name: str | None = None,
     attention_mask: Array = None,
     use_base2_exp: bool = False,
@@ -954,7 +1103,7 @@ def __init__(
       flash_block_sizes: BlockSizes = None,
       dtype: DType = jnp.float32,
       quant: Quant = None,
-      mask_padding_tokens: bool = False,
+      mask_padding_tokens: bool = True,
       residual_checkpoint_name: str | None = None,
       use_base2_exp: bool = False,
       use_experimental_scheduler: bool = False,
@@ -1122,7 +1271,7 @@ def __init__(
       qkv_bias: bool = False,
       quant: Quant = None,
       is_self_attention: bool = True,
-      mask_padding_tokens: bool = False,
+      mask_padding_tokens: bool = True,
       residual_checkpoint_name: str | None = None,
       enable_jax_named_scopes: bool = False,
       added_kv_proj_dim: Optional[int] = None,  # New for I2V