[maxtext] Add generalized block-causal attention

src/maxtext/common/common_types.py

@@ -135,6 +135,7 @@ class AttentionType(enum.Enum):
   MLA = "mla"
   COMPRESSED = "compressed"
   FULL = "full"
+  BLOCK_DIFFUSION = "block_diffusion"
 
 
 class ShardMode(enum.Enum):

src/maxtext/configs/base.yml

@@ -380,12 +380,13 @@ param_scan_axis: 1
 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores
 # The attention_type parameter determines the variants of attention, e.g. global or local_sliding
 attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te
-attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla
+attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla, full, compressed, block_diffusion
 share_kv_projections: false # Note: Not compatible with attention_type='mla'
 attention_bias: false # If true, adds a learnable bias to the query, key, and value projections
 attention_sink: false
 sliding_window_size: 0
 chunk_attn_window_size: 0
+block_diffusion_block_size: 32 # Size of each bidirectional block when attention_type='block_diffusion'
 attn_logits_soft_cap: 0.0
 final_logits_soft_cap: 0.0
 z_loss_multiplier: 0.0

src/maxtext/configs/types.py

@@ -572,7 +572,7 @@ class Attention(BaseModel):
       "autoselected",
       description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, etc).",
   )
-  attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed"] = Field(
+  attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field(
       "global", description="The variant of attention to use."
   )
   share_kv_projections: bool = Field(
@@ -619,6 +619,10 @@ class Attention(BaseModel):
   )
   sliding_window_size: NonNegativeInt = Field(0, description="The size of the sliding window for local attention.")
   chunk_attn_window_size: NonNegativeInt = Field(0, description="The window size for chunked attention.")
+  block_diffusion_block_size: PositiveInt = Field(
+      32,
+      description="The size of each bidirectional block when attention_type is 'block_diffusion'.",
+  )
   attn_logits_soft_cap: None | NonNegativeFloat = Field(
       None, description="Soft-cap value for attention logits. None means no cap."
   )
@@ -3283,6 +3287,14 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         not isinstance(self.sliding_window_size, int) or self.sliding_window_size <= 0
     ):
       raise ValueError("`sliding_window_size` must be an integer > 0 for 'local_sliding' attention.")
+    if self.attention_type == AttentionType.BLOCK_DIFFUSION.value:
+      if self.attention not in ("autoselected", "dot_product", "flash"):
+        raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.")
+      if self.attention in ("autoselected", "flash") and self.hardware != "tpu":
+        raise ValueError(
+            "Block-diffusion attention with attention='autoselected' or attention='flash' requires hardware='tpu'; "
+            "use attention='dot_product' on other hardware."
+        )
     if self.quantize_kvcache and not self.kv_quant_axis:
       raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.")
     if (

src/maxtext/layers/attention_op.py

@@ -84,6 +84,16 @@
 dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None))
 
 
+def _resolve_attention_type(config: Config, attention_type: AttentionType | str | None) -> AttentionType:
+  configured_attention_type = AttentionType(getattr(config, "attention_type", AttentionType.GLOBAL.value))
+  if attention_type is None:
+    return configured_attention_type
+  resolved_attention_type = AttentionType(attention_type)
+  if configured_attention_type == AttentionType.BLOCK_DIFFUSION and resolved_attention_type == AttentionType.GLOBAL:
+    return configured_attention_type
+  return resolved_attention_type
+
+
 def validate_compute_axis_order(s: AxisIdxes) -> None:
   valid_compute_axis_order = ((0, 1, 2, 3), (0, 2, 1, 3))
   if s not in valid_compute_axis_order:  # currently supported compute_axis_order
@@ -204,6 +214,50 @@ def __hash__(self):
     )
 
 
+class BlockCausalMask(splash_attention_mask._ComputableMask):  # pylint: disable=protected-access,abstract-method
+  """Lazy mask with bidirectional attention within causal blocks."""
+
+  block_size: int
+
+  def __init__(
+      self,
+      shape: tuple[int, int],
+      block_size: int,
+      shard_count: int = 1,
+  ):
+    if block_size <= 0:
+      raise ValueError("block_size must be positive")
+    self.block_size = block_size
+
+    def block_causal_mask_function(q_ids, kv_ids):
+      return (q_ids // self.block_size) >= (kv_ids // self.block_size)
+
+    super().__init__(
+        shape=shape,
+        mask_function=block_causal_mask_function,
+        shard_count=shard_count,
+    )
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+    return (
+        self.shape == other.shape
+        and self.block_size == other.block_size
+        and np.array_equal(self.q_sequence, other.q_sequence)
+    )
+
+  def __hash__(self):
+    return hash(
+        (
+            type(self),
+            self.shape,
+            self.block_size,
+            self.q_sequence.tobytes() if self.q_sequence is not None else None,
+        )
+    )
+
+
 def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, q_offset: int = 0) -> jax.Array:
   """Generates an explicit boolean mask for chunked causal attention.
 
@@ -234,6 +288,16 @@ def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int,
   return chunk_mask
 
 
+def _generate_block_causal_attention_mask(mask_shape: tuple[int, int], block_size: int, q_offset: int = 0) -> jax.Array:
+  """Generates a block-causal mask, bidirectional within each block."""
+  if block_size <= 0:
+    raise ValueError("block_size must be positive")
+
+  row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + q_offset
+  col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
+  return (row_ids // block_size) >= (col_ids // block_size)
+
+
 def _make_block_mask_indices(bidirectional_mask):
   """Creates block mask identifying segments based on a bidirectional mask.
 
@@ -479,7 +543,13 @@ def __init__(
     self.dtype = dtype
     self.quant = quant
     self.kv_quant = kv_quant
-    self.attention_type = attention_type
+    self.attention_type = _resolve_attention_type(self.config, attention_type)
+    self.block_diffusion_block_size = getattr(self.config, "block_diffusion_block_size", None)
+    if self.attention_type == AttentionType.BLOCK_DIFFUSION:
+      if self.block_diffusion_block_size is None or self.block_diffusion_block_size <= 0:
+        raise ValueError("block_diffusion_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
     if self.attention_kernel not in (
         "dot_product",
@@ -782,6 +852,7 @@ def generate_attention_mask(
     if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in (
         AttentionType.FULL,
         AttentionType.COMPRESSED,
+        AttentionType.BLOCK_DIFFUSION,
     ):
       if use_segment_positions:
         causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :]
@@ -861,6 +932,19 @@ def generate_attention_mask(
         )
       output_mask = chunk_mask * output_mask
 
+    elif self.attention_type == AttentionType.BLOCK_DIFFUSION and model_mode != MODEL_MODE_AUTOREGRESSIVE:
+      if use_segment_positions:
+        block_mask = (
+            (position_row_ids // self.block_diffusion_block_size) >= (position_col_ids // self.block_diffusion_block_size)
+        )[:, None, None, :, :]
+      else:
+        block_mask = _generate_block_causal_attention_mask(
+            mask_shape=(q_seq_len, kv_seq_len),
+            block_size=self.block_diffusion_block_size,
+            q_offset=next_pos,
+        )[None, None, None, :, :]
+      output_mask = block_mask if output_mask is None else jnp.logical_and(output_mask, block_mask)
+
     if bidirectional_mask is not None:
       image_mask = _make_bidirectional_block_mask(bidirectional_mask)
       output_mask = output_mask | image_mask[:, None, None, ...]
@@ -1126,6 +1210,11 @@ def apply_attention(
         return out, None, None
 
       else:
+        if self.attention_type == AttentionType.BLOCK_DIFFUSION:
+          raise ValueError(
+              "Block-diffusion flash attention is supported only by TPU Splash; "
+              "use attention='dot_product' on other hardware."
+          )
         if model_mode == MODEL_MODE_AUTOREGRESSIVE:
           # fallback to dot_product as pallas gpu flash attention doesn't support decode stage
           return self.apply_attention_dot(
@@ -1430,13 +1519,24 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
       sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap)
       mask_shape = (query.shape[2], key.shape[2])  # (q_seq_len, kv_seq_len)
       mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask
+      use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel
       if self.attention_type == AttentionType.FULL:
         mask = mask_module.FullMask(mask_shape)
+      elif self.attention_type == AttentionType.BLOCK_DIFFUSION:
+        mask_type = LoadBalancedBlockCausalMask if use_load_balanced_cp else BlockCausalMask
+        mask_kwargs = {"cp_size": cp_size} if use_load_balanced_cp else {}
+        mask = mask_type(
+            shape=mask_shape,
+            block_size=self.block_diffusion_block_size,
+            **mask_kwargs,
+        )
       else:
         mask = mask_module.CausalMask(shape=mask_shape)
 
-      use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel
-      if use_load_balanced_cp and self.attention_type != AttentionType.FULL:
+      if use_load_balanced_cp and self.attention_type not in (
+          AttentionType.FULL,
+          AttentionType.BLOCK_DIFFUSION,
+      ):
         mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size)
 
       # Apply local masking if local sliding attention is enabled.
@@ -2447,3 +2547,21 @@ def __init__(
         shard_count=shard_count,
     )
     self.q_sequence = _load_balanced_q_sequence(shape, cp_size)
+
+
+class LoadBalancedBlockCausalMask(BlockCausalMask):  # pylint: disable=abstract-method
+  """Lazy block-causal mask with load-balanced query positions."""
+
+  def __init__(
+      self,
+      shape: tuple[int, int],
+      block_size: int,
+      cp_size: int,
+      shard_count: int = 1,
+  ):
+    super().__init__(
+        shape=shape,
+        block_size=block_size,
+        shard_count=shard_count,
+    )
+    self.q_sequence = _load_balanced_q_sequence(shape, cp_size)

src/maxtext/layers/attentions.py

@@ -51,7 +51,7 @@
     AttentionType,
 )
 from maxtext.layers import nnx_wrappers
-from maxtext.layers.attention_op import AttentionOp
+from maxtext.layers.attention_op import AttentionOp, _resolve_attention_type
 from maxtext.layers.embeddings import (
     LLaMARotaryEmbedding,
     LlamaVisionRotaryEmbedding,
@@ -120,7 +120,7 @@ def attention_as_linen(
     float32_logits: bool = False,  # cast logits in float32 for stability.
     quant: Optional[Quant] = None,
     kv_quant: Optional[KVQuant] = None,
-    attention_type: AttentionType = AttentionType.GLOBAL,  # Default to global attention
+    attention_type: AttentionType | None = None,
     attn_logits_soft_cap: float | None = None,
     sliding_window_size: int | None = None,
     use_ragged_attention: bool = False,
@@ -277,7 +277,7 @@ def __init__(
       float32_logits: bool = False,  # cast logits in float32 for stability.
       quant: Optional[Quant] = None,
       kv_quant: Optional[KVQuant] = None,
-      attention_type: AttentionType = AttentionType.GLOBAL,  # Default to global attention
+      attention_type: AttentionType | None = None,
       attn_logits_soft_cap: float | None = None,
       sliding_window_size: int | None = None,
       use_ragged_attention: bool = False,
@@ -388,7 +388,7 @@ def __init__(
     self.float32_logits = float32_logits
     self.quant = quant
     self.kv_quant = kv_quant
-    self.attention_type = attention_type
+    self.attention_type = _resolve_attention_type(self.config, attention_type)
     self.attn_logits_soft_cap = attn_logits_soft_cap
     self.sliding_window_size = sliding_window_size
     self.use_ragged_attention = use_ragged_attention