From a891c8fda6d3be7be9e8059dedecb779bedf72d6 Mon Sep 17 00:00:00 2001 From: Fareed Sheriff Date: Fri, 1 May 2026 10:48:19 -0700 Subject: [PATCH] Add Flash Attention option to `MultiHeadAttention` to scale to longer sequence lengths without materializing large attention dot product intermediates in HBM. PiperOrigin-RevId: 908804155 --- .../lib/architecture/attention.py | 149 +++++++++++++++--- .../lib/architecture/attention_test.py | 110 +++++++++++++ .../lib/architecture/dit_blocks.py | 6 +- .../lib/architecture/dit_blocks_test.py | 37 +++-- 4 files changed, 268 insertions(+), 34 deletions(-) diff --git a/hackable_diffusion/lib/architecture/attention.py b/hackable_diffusion/lib/architecture/attention.py index 69f1830..3fa9e80 100644 --- a/hackable_diffusion/lib/architecture/attention.py +++ b/hackable_diffusion/lib/architecture/attention.py @@ -15,6 +15,7 @@ """Attention layers and utils.""" import dataclasses +import math from typing import Literal import warnings @@ -22,6 +23,7 @@ from hackable_diffusion.lib import hd_typing from hackable_diffusion.lib.architecture import sequence_embedders import jax +import jax.experimental.pallas.ops.tpu.flash_attention as flash import jax.numpy as jnp import kauldron.ktyping as kt @@ -51,9 +53,10 @@ ################################################################################ +# TODO: b/539954103 - Rename to AttentionSpec. @dataclasses.dataclass(frozen=True) class AttentionHeadsSpec: - """Configuration for multi-head attention dimensionality. + """Configuration for multi-head attention. Specify at least one of `num_heads` or `head_dim`. When only one is set, the other is inferred from the embedding dimension at call time via `resolve()`. @@ -63,10 +66,14 @@ class AttentionHeadsSpec: Attributes: num_heads: Fixed number of attention heads. head_dim: Fixed dimension per head. + use_flash_attention: Whether to use Flash Attention. + block_sizes: Block sizes for FlashAttention. """ num_heads: int | None = None head_dim: int | None = None + use_flash_attention: bool = False + block_sizes: flash.BlockSizes | None = None def __post_init__(self): if not self.num_heads and not self.head_dim: @@ -126,7 +133,7 @@ def _dot_product_attention( q: Float["batch head sequence_query dim"], # pyrefly: ignore[not-a-type] k: Float["batch head sequence_key dim"], # pyrefly: ignore[not-a-type] v: Float["batch head sequence_key dim"], # pyrefly: ignore[not-a-type] - rescale: Float["..."], # pyrefly: ignore[bad-index, not-a-type] + rescale: float | Float["..."], # pyrefly: ignore[bad-index, not-a-type] *, mask: Bool["batch sequence_key"] | None = None, dropout_rate: float = 0.0, @@ -157,7 +164,9 @@ def _dot_product_attention( # for masked tokens. if mask is not None: bcast_mask = jnp.expand_dims(mask, axis=(1, 2)) - attention_logits = jnp.where(bcast_mask, attention_logits, MASK_LOGITS_VALUE) + attention_logits = jnp.where( + bcast_mask, attention_logits, MASK_LOGITS_VALUE + ) # Softmax and attention weights attention_weights = _stable_softmax(logits=attention_logits) @@ -176,6 +185,77 @@ def _dot_product_attention( return attention_output +@kt.typechecked +def _flash_dot_product_attention( + q: Float["batch head sequence_query dim"], # pyrefly: ignore[not-a-type] + k: Float["batch head sequence_key dim"], # pyrefly: ignore[not-a-type] + v: Float["batch head sequence_key dim"], # pyrefly: ignore[not-a-type] + rescale: float | Float["..."], # pyrefly: ignore[bad-index, not-a-type] + *, + mask: Bool["batch sequence_key"] | None = None, + block_sizes: flash.BlockSizes | None = None, +) -> Float["batch sequence_query head_dim_concat"]: + """Broadcasts mask to SegmentIds and performs flash attention. + + Prevents quadratic attention memory usage and full intermediate and mask + materialization in HBM as done by _dot_product_attention. + + Args: + q: Query tensor. + k: Key tensor. + v: Value tensor. + rescale: Rescale factor for the attention scores. + mask: Mask tensor. Mask is True for tokens we want to keep and False for + tokens we want to mask. If None, no masking is performed. + block_sizes: Block sizes for FlashAttention. If None, default block size of + 128 for Q and K is used. Block sizes must be divisible by 128 for TPU + Flash Attention; they can be increased from 128 to decrease inter-block + communication latency and increase MFU contingent on VMEM capacity and + sequence length. + + Returns: + The output tensor. + """ + b, seq_len_q, seq_len_kv = q.shape[0], q.shape[2], k.shape[2] + bq, bkm = ( + (block_sizes.block_q, block_sizes.block_k_major) + if block_sizes + else (128, 128) + ) + pad_q, pad_k = (bq - (seq_len_q % bq)) % bq, (bkm - (seq_len_kv % bkm)) % bkm + + q_padded = jnp.pad(q, ((0, 0), (0, 0), (0, pad_q), (0, 0))) + k_padded = jnp.pad(k, ((0, 0), (0, 0), (0, pad_k), (0, 0))) + v_padded = jnp.pad(v, ((0, 0), (0, 0), (0, pad_k), (0, 0))) + + segment_ids = None + if mask is not None or pad_k > 0: + m = mask if mask is not None else jnp.ones((b, seq_len_kv), dtype=bool) + mask_padded = jnp.pad(m, ((0, 0), (0, pad_k)), constant_values=False) + kv_seg = jnp.where( + mask_padded, + jnp.int32(0), + jnp.arange(1, mask_padded.shape[1] + 1, dtype=jnp.int32), + ) + q_seg = jnp.zeros((b, seq_len_q + pad_q), dtype=jnp.int32) + segment_ids = flash.SegmentIds(q=q_seg, kv=kv_seg) + + attn_output = flash.flash_attention( + q_padded, + k_padded, + v_padded, + segment_ids=segment_ids, + causal=False, + sm_scale=float(rescale), + block_sizes=block_sizes, + ) + return ( + attn_output[:, :, :seq_len_q, :] + .transpose(0, 2, 1, 3) + .reshape(b, seq_len_q, -1) + ) + + ################################################################################ # MARK: Multi-Head Attention ################################################################################ @@ -194,8 +274,8 @@ class MultiHeadAttention(nn.Module): It supports RoPE for positional embeddings and QK normalization. Attributes: - attention_heads_spec: An AttentionHeadsSpec instance that specifies num_heads - and/or head_dim for the attention layer. + attention_heads_spec: An AttentionHeadsSpec instance that specifies + num_heads and/or head_dim for the attention layer. normalize_qk: Whether to normalize query and key before attention. use_rope: Whether to use rotary positional embeddings on query and key. rope_positions_fn: The position function of rotary positional embeddings to @@ -205,6 +285,9 @@ class MultiHeadAttention(nn.Module): is initialized to zeros. dropout_rate: The dropout rate for the attention weights. dtype: The data type of the computation. + use_flash: Whether to use Flash Attention, currently only supported for TPU. + block_sizes: Block sizes for FlashAttention. If None, default block sizes + are used. Block sizes must be divisible by 128 for TPU Flash Attention. """ attention_heads_spec: AttentionHeadsSpec @@ -262,6 +345,14 @@ def __call__( f"In cross-attention, mask shape {mask.shape} does not match" f" expected shape {c.shape[:2]}." ) + + if ( + self.attention_heads_spec.use_flash_attention + and self.dropout_rate > 0.0 + and is_training + ): + raise ValueError("Flash attention is not supported with dropout.") + b, _, d = x.shape # batch size, sequence length, embedding dim head_d, num_heads = self.attention_heads_spec.resolve(d) @@ -304,7 +395,6 @@ def __call__( if self.qk_norm_method == "rms_norm": q = nn.RMSNorm(name="RMSNorm_Q")(q) k = nn.RMSNorm(name="RMSNorm_K")(k) - scale = 1.0 / jnp.sqrt(jnp.float32(head_d)) # QK L2 normalization: https://arxiv.org/abs/2010.04245 elif self.qk_norm_method == "l2": scale = self.param( @@ -317,14 +407,22 @@ def __call__( norm_q = jnp.linalg.norm(q, ord=2, axis=-1, keepdims=True) norm_k = jnp.linalg.norm(k, ord=2, axis=-1, keepdims=True) - q = q / (norm_q + SAFETY_EPSILON) + # we pre-scale Q here instead of within the attention function to avoid + # passing in a differentiable parameter to the attention; the scale + # parameter still receives the same gradient as it would if it were + # inside the attention + q = q * (scale / (norm_q + SAFETY_EPSILON)) k = k / (norm_k + SAFETY_EPSILON) else: raise ValueError( f"Unsupported QK normalization method: {self.qk_norm_method}." ) - else: - scale = 1.0 / jnp.sqrt(jnp.float32(head_d)) + + # Downstream attention dot-product scaling factor + is_l2 = self.normalize_qk and self.qk_norm_method == "l2" + # 'rescale=1.0' is passed to _dot_product_attention for functional + # equivalence. + rescale = 1.0 if is_l2 else (1.0 / math.sqrt(head_d)) # RoPE: https://arxiv.org/abs/2104.09864 if self.use_rope: @@ -336,15 +434,30 @@ def __call__( )(k) # shape is [batch, num_heads, sequence_length, head_dim] - attention_output = _dot_product_attention( - q=q, - k=k, - v=v, - rescale=scale, - mask=mask, - dropout_rate=self.dropout_rate, - is_training=is_training, - ) + if self.attention_heads_spec.use_flash_attention and not ( + self.dropout_rate > 0.0 and is_training + ): + # exception is raised when the block sizes are not divisible by 128 + # (number of TPU lanes) or when enabling TPU Flash Attention on a non-TPU + # backend. + attention_output = _flash_dot_product_attention( + q=q, + k=k, + v=v, + rescale=rescale, + mask=mask, + block_sizes=self.attention_heads_spec.block_sizes, + ) + else: + attention_output = _dot_product_attention( + q=q, + k=k, + v=v, + rescale=rescale, + mask=mask, + dropout_rate=self.dropout_rate, + is_training=is_training, + ) attention_output = nn.Dense( features=d, diff --git a/hackable_diffusion/lib/architecture/attention_test.py b/hackable_diffusion/lib/architecture/attention_test.py index 1523fa5..8dd3186 100644 --- a/hackable_diffusion/lib/architecture/attention_test.py +++ b/hackable_diffusion/lib/architecture/attention_test.py @@ -19,6 +19,8 @@ from hackable_diffusion.lib.architecture import attention from hackable_diffusion.lib.architecture import sequence_embedders import jax +from jax.experimental.pallas import tpu as pltpu +import jax.experimental.pallas.ops.tpu.flash_attention as flash import jax.numpy as jnp import kauldron.ktyping as kt import numpy as np @@ -641,6 +643,114 @@ def test_qk_norm_invalid_method_raises_error(self): ): module.init(self.rng, self.x, c=None) + @parameterized.named_parameters( + # Format: (name, normalize_qk, qk_norm_method, seq_len_q, seq_len_kv, dim) + # 1. Aligned Baseline + ("aligned_no_norm", False, "l2", 128, 128, 128), + ("aligned_l2_norm", True, "l2", 128, 128, 128), + # 2. Mixed Boundaries (covers 1 below/above and 10 above/below in + # cross-attention) + ("mixed_boundary_1_l2", True, "l2", 127, 129, 128), + ("mixed_boundary_10_rms", True, "rms_norm", 138, 118, 128), + # 3. Extreme / Small Sizes (also covers dim 64) + ("minimum_size", True, "l2", 1, 1, 64), + ("small_context_dim_64", True, "l2", 8, 32, 64), + # 4. Large Context + ("large_context", True, "l2", 256, 512, 128), + ) + def test_flash_attention_correctness( + self, normalize_qk, qk_norm_method, seq_len_q, seq_len_kv, dim + ): + """Verifies that the Flash Attention forward and backward passes match Naive.""" + + # Setup inputs + rng_init, rng_eval, rng_grad = jax.random.split(self.rng, 3) + x = jax.random.normal(rng_eval, (self.batch_size, seq_len_q, dim)) + c = jax.random.normal(rng_eval, (self.batch_size, seq_len_kv, dim)) + + # Compile two modules: one with naive, one with flash + module_naive = attention.MultiHeadAttention( + attention_heads_spec=attention.AttentionHeadsSpec( + num_heads=self.num_heads, use_flash_attention=False + ), + normalize_qk=normalize_qk, + qk_norm_method=qk_norm_method, + ) + module_flash = attention.MultiHeadAttention( + attention_heads_spec=attention.AttentionHeadsSpec( + num_heads=self.num_heads, use_flash_attention=True + ), + normalize_qk=normalize_qk, + qk_norm_method=qk_norm_method, + ) + + variables = module_naive.init(rng_init, x, c) + + # Define forward wrapper to use with jax.vjp + def forward_fn(module, q, k): + return module.apply(variables, q, k, is_training=False) + + # 1. Forward and VJP: Naive Attention + out_naive, vjp_fn_naive = jax.vjp( + lambda q, k: forward_fn(module_naive, q, k), x, c + ) + + # 2. Forward and VJP: Flash Attention under TPU interpret mode + with pltpu.force_tpu_interpret_mode(): + out_flash, vjp_fn_flash = jax.vjp( + lambda q, k: forward_fn(module_flash, q, k), x, c + ) + + # Assert Forward Correctness + np.testing.assert_allclose(out_flash, out_naive, atol=1e-5, rtol=1e-5) + + # 3. Propagate incoming gradients backward + g = jax.random.normal(rng_grad, out_naive.shape) + dx_naive, dc_naive = vjp_fn_naive(g) + + with pltpu.force_tpu_interpret_mode(): + dx_flash, dc_flash = vjp_fn_flash(g) + + # Assert VJP Correctness (Gradients of inputs) + np.testing.assert_allclose(dx_flash, dx_naive, atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(dc_flash, dc_naive, atol=1e-5, rtol=1e-5) + + def test_flash_attention_fails_without_tpu_or_interpreter(self): + """Verifies that Flash Attention fails when run on CPU without interpreter.""" + spec = attention.AttentionHeadsSpec( + num_heads=self.num_heads, + use_flash_attention=True, + ) + module = attention.MultiHeadAttention(attention_heads_spec=spec) + x = jnp.ones((self.batch_size, self.seq_len_q, self.dim)) + # We force interpret mode to None, which disables CPU interpreter emulation. + with pltpu.force_tpu_interpret_mode(None): # type: ignore[wrong-arg-types] + with self.assertRaises((ValueError, RuntimeError)): + variables = module.init(self.rng, x, None) + module.apply(variables, x, None) + + def test_flash_attention_invalid_block_k_raises_error(self): + """Verifies that block_k not divisible by 128 raises error in Flash Attention.""" + invalid_block_sizes = flash.BlockSizes( + block_q=128, + block_k_major=128, + block_k=64, # 64 is not a multiple of 128 + block_b=1, + ) + spec = attention.AttentionHeadsSpec( + num_heads=self.num_heads, + use_flash_attention=True, + block_sizes=invalid_block_sizes, + ) + + module = attention.MultiHeadAttention(attention_heads_spec=spec) + x = jnp.ones((self.batch_size, self.seq_len_q, self.dim)) + mask = jnp.ones((self.batch_size, self.seq_len_q), dtype=jnp.bool_) + + with self.assertRaises((NotImplementedError, ValueError, RuntimeError)): + variables = module.init(self.rng, x, None, mask=mask) + module.apply(variables, x, None, mask=mask) + if __name__ == "__main__": absltest.main() diff --git a/hackable_diffusion/lib/architecture/dit_blocks.py b/hackable_diffusion/lib/architecture/dit_blocks.py index c4737c0..9ea102c 100644 --- a/hackable_diffusion/lib/architecture/dit_blocks.py +++ b/hackable_diffusion/lib/architecture/dit_blocks.py @@ -234,7 +234,9 @@ def __call__( # Attention Branch. x_normed = self.conditional_norm_attention(x, c=cond_activated) - attention_out = self.attention(x_normed, c=None, mask=mask, is_training=is_training) + attention_out = self.attention( + x_normed, c=None, mask=mask, is_training=is_training + ) if use_gates: gate_msa = self.gate_msa(cond_activated) attention_out = gate_msa[..., None, :] * attention_out @@ -405,7 +407,7 @@ def __call__( return einops.rearrange( x, - "... (hn wn) (hp wp c) -> ... (hn hp) (wn wp) c", + '... (hn wn) (hp wp c) -> ... (hn hp) (wn wp) c', hn=hn, wn=wn, hp=hp, diff --git a/hackable_diffusion/lib/architecture/dit_blocks_test.py b/hackable_diffusion/lib/architecture/dit_blocks_test.py index 4343bd3..65e3c06 100644 --- a/hackable_diffusion/lib/architecture/dit_blocks_test.py +++ b/hackable_diffusion/lib/architecture/dit_blocks_test.py @@ -125,9 +125,7 @@ def test_variable_shapes_ada_rms_norm(self): module = dit_blocks.DiTBlock( hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), - norm_strategy=normalization.ConditionalRMSNormStrategy( - use_shift=False - ), + norm_strategy=normalization.ConditionalRMSNormStrategy(use_shift=False), use_gates=False, ffn_type='swiglu', ) @@ -242,9 +240,7 @@ def test_use_gates_false_without_zero_init_output_raises(self): module = dit_blocks.DiTBlock( hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), - norm_strategy=normalization.ConditionalRMSNormStrategy( - use_shift=False - ), + norm_strategy=normalization.ConditionalRMSNormStrategy(use_shift=False), use_gates=False, zero_init_output=False, ) @@ -270,9 +266,7 @@ def test_ffn_use_bias(self, ffn_type, ffn_use_bias): module = dit_blocks.DiTBlock( hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), - norm_strategy=normalization.ConditionalRMSNormStrategy( - use_shift=False - ), + norm_strategy=normalization.ConditionalRMSNormStrategy(use_shift=False), use_gates=False, ffn_type=ffn_type, ffn_use_bias=ffn_use_bias, @@ -310,7 +304,10 @@ def test_preset_output_shape(self, block_cls): cond_shape = (self.batch, self.c) x = jnp.ones(input_shape) cond = jnp.ones(cond_shape) - module = block_cls(hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4)) + module = block_cls( + hidden_size=self.d, + attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), + ) variables = module.init(self.key, x, cond, is_training=False) output = module.apply(variables, x, cond, is_training=False) self.assertEqual(output.shape, input_shape) @@ -326,7 +323,10 @@ def test_preset_identity_at_init(self, block_cls): cond_shape = (self.batch, self.c) x = jax.random.normal(self.key, input_shape) cond = jnp.zeros(cond_shape) - module = block_cls(hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4)) + module = block_cls( + hidden_size=self.d, + attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), + ) variables = module.init(self.key, x, cond, is_training=False) output = module.apply(variables, x, cond, is_training=False) self.assertTrue(jnp.allclose(output, x, atol=1e-5)) @@ -335,7 +335,10 @@ def test_flux_has_no_gates(self): """Verifies DiTBlockFlux has no gate parameters.""" x = jnp.ones((self.batch, self.n, self.d)) cond = jnp.ones((self.batch, self.c)) - module = dit_blocks.DiTBlockFlux(hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4)) + module = dit_blocks.DiTBlockFlux( + hidden_size=self.d, + attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), + ) variables = module.init(self.key, x, cond, is_training=False) leaves_with_paths = test_helpers.get_leaves_with_paths(variables) gate_paths = [p for p in leaves_with_paths if 'Gate' in p] @@ -345,7 +348,10 @@ def test_sd3_has_gates(self): """Verifies DiTBlockSD3 has gate parameters.""" x = jnp.ones((self.batch, self.n, self.d)) cond = jnp.ones((self.batch, self.c)) - module = dit_blocks.DiTBlockSD3(hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4)) + module = dit_blocks.DiTBlockSD3( + hidden_size=self.d, + attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), + ) variables = module.init(self.key, x, cond, is_training=False) leaves_with_paths = test_helpers.get_leaves_with_paths(variables) gate_paths = [p for p in leaves_with_paths if 'Gate' in p] @@ -355,7 +361,10 @@ def test_ada_ln_zero_has_gates(self): """Verifies DiTBlockAdaLNZero has gate parameters.""" x = jnp.ones((self.batch, self.n, self.d)) cond = jnp.ones((self.batch, self.c)) - module = dit_blocks.DiTBlockAdaLNZero(hidden_size=self.d, attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4)) + module = dit_blocks.DiTBlockAdaLNZero( + hidden_size=self.d, + attention_heads_spec=attention.AttentionHeadsSpec(num_heads=4), + ) variables = module.init(self.key, x, cond, is_training=False) leaves_with_paths = test_helpers.get_leaves_with_paths(variables) gate_paths = [p for p in leaves_with_paths if 'Gate' in p]