diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 9dd4aff3dc..8ed18bfd12 100644 --- a/src/maxtext/common/common_types.py +++ b/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): diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 55925ffd58..4d0bf560e8 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 03640ec758..8e42109894 100644 --- a/src/maxtext/configs/types.py +++ b/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 ( diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 8e6c96a72e..af48802cc9 100644 --- a/src/maxtext/layers/attention_op.py +++ b/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) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 5a65b35ba4..32f233cbeb 100644 --- a/src/maxtext/layers/attentions.py +++ b/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 diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 20522bbd3b..0c7aa03005 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -44,9 +44,12 @@ from maxtext.layers import attention_op from maxtext.layers.attention_op import ( AttentionOp, + BlockCausalMask, ChunkedCausalMask, + _generate_block_causal_attention_mask, _generate_chunk_attention_mask, _make_bidirectional_block_mask, + _resolve_attention_type, ) from maxtext.layers.attentions import Attention from maxtext.layers import embeddings @@ -330,6 +333,143 @@ def test_value_error_on_zero_chunk_size(self): _generate_chunk_attention_mask(mask_shape=(4, 4), chunk_size=0) +class BlockCausalMaskTest(unittest.TestCase): + """Tests the shared dense and Splash block-causal masks.""" + + def _make_op(self, sequence_length, *, attention_type=AttentionType.BLOCK_DIFFUSION): + """Builds a minimal dot-product attention operator.""" + config = types.SimpleNamespace( + block_diffusion_block_size=4, + context_parallel_load_balance=False, + context_sharding="context", + ) + mesh = types.SimpleNamespace(shape={}) + kwargs = {} + if attention_type is not None: + kwargs["attention_type"] = attention_type + return AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=sequence_length, + mesh=mesh, + attention_kernel="dot_product", + **kwargs, + ) + + def test_dense_and_splash_masks_match(self): + sequence_length = 10 + block_size = 4 + query_positions = np.arange(sequence_length)[:, None] + key_positions = np.arange(sequence_length)[None, :] + expected = query_positions // block_size >= key_positions // block_size + + splash_mask = BlockCausalMask((sequence_length, sequence_length), block_size) + dense_mask = _generate_block_causal_attention_mask((sequence_length, sequence_length), block_size) + + np.testing.assert_array_equal(splash_mask[:, :], expected) + np.testing.assert_array_equal(dense_mask, expected) + self.assertTrue(expected[1, 3]) + self.assertFalse(expected[3, 4]) + self.assertTrue(expected[4, 3]) + self.assertTrue(expected[8, 9]) + + def test_dot_product_mask_respects_packed_segments(self): + sequence_length = 12 + query = jnp.zeros((1, sequence_length, 1, 8)) + key = jnp.zeros((1, sequence_length, 1, 8)) + segment_ids = jnp.asarray([[1] * 8 + [2] * 4], dtype=jnp.int32) + + mask = self._make_op(sequence_length).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_TRAIN, + ) + + positions = np.arange(sequence_length) + expected = (positions[:, None] // 4 >= positions[None, :] // 4) & ( + np.asarray(segment_ids[0])[:, None] == np.asarray(segment_ids[0])[None, :] + ) + np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected) + + def test_autoregressive_mask_is_unchanged(self): + key_length = 8 + query = jnp.zeros((1, 1, 1, 8)) + key = jnp.zeros((1, key_length, 1, 8)) + segment_ids = jnp.asarray([[1, 1, 0, 1, 0, 0, 1, 0]], dtype=jnp.int32) + + default_mask = self._make_op(key_length, attention_type=None).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_AUTOREGRESSIVE, + ) + block_diffusion_mask = self._make_op(key_length).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(block_diffusion_mask, default_mask) + np.testing.assert_array_equal( + np.asarray(default_mask == 0.0)[0, 0, 0, 0], + np.asarray(segment_ids[0] == DECODING_ACTIVE_SEQUENCE_INDICATOR), + ) + + def test_invalid_block_size_raises(self): + with self.assertRaises(ValueError): + BlockCausalMask((4, 4), 0) + with self.assertRaises(ValueError): + _generate_block_causal_attention_mask((4, 4), -1) + + +class AttentionTypeResolutionTest(unittest.TestCase): + + def test_config_selects_block_diffusion_without_model_dispatch(self): + config = types.SimpleNamespace(attention_type=AttentionType.BLOCK_DIFFUSION.value) + + self.assertEqual(_resolve_attention_type(config, None), AttentionType.BLOCK_DIFFUSION) + self.assertEqual(_resolve_attention_type(config, AttentionType.GLOBAL), AttentionType.BLOCK_DIFFUSION) + self.assertEqual(_resolve_attention_type(config, AttentionType.FULL), AttentionType.FULL) + self.assertEqual(_resolve_attention_type(config, AttentionType.LOCAL_SLIDING), AttentionType.LOCAL_SLIDING) + self.assertEqual( + _resolve_attention_type(types.SimpleNamespace(attention_type=AttentionType.GLOBAL.value), AttentionType.FULL), + AttentionType.FULL, + ) + self.assertEqual(_resolve_attention_type(types.SimpleNamespace(), None), AttentionType.GLOBAL) + + def test_attention_op_honors_config_without_overriding_specialized_layers(self): + config = types.SimpleNamespace( + attention_type=AttentionType.BLOCK_DIFFUSION.value, + block_diffusion_block_size=4, + ) + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=8, + mesh=types.SimpleNamespace(shape={}), + attention_kernel="dot_product", + attention_type=AttentionType.GLOBAL, + ) + + self.assertEqual(op.attention_type, AttentionType.BLOCK_DIFFUSION) + + full_op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=8, + mesh=types.SimpleNamespace(shape={}), + attention_kernel="dot_product", + attention_type=AttentionType.FULL, + ) + + self.assertEqual(full_op.attention_type, AttentionType.FULL) + + class LoadBalancedMaskTest(unittest.TestCase): """Tests for load-balanced Splash masks.""" @@ -370,6 +510,20 @@ def test_load_balanced_chunk_window(self): np.testing.assert_array_equal((causal_mask & chunk_mask)[:, :], expected_mask) + def test_load_balanced_block_causal_mask(self): + sequence_length = 8 + block_size = 2 + mask = attention_op.LoadBalancedBlockCausalMask( + shape=(sequence_length, sequence_length), + block_size=block_size, + cp_size=2, + ) + query_positions = mask.q_sequence[:, None] + key_positions = np.arange(sequence_length)[None, :] + expected = query_positions // block_size >= key_positions // block_size + + np.testing.assert_array_equal(mask[:, :], expected) + def test_dot_product_local_mask_uses_segment_positions(self): config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context") mesh = types.SimpleNamespace(shape={"context": 4}) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index f0504e4dd8..f333748192 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -192,6 +192,54 @@ def test_load_balanced_chunk_context_parallel_config(self): self.assertEqual(config.attention_type, "chunk") self.assertTrue(config.context_parallel_load_balance) + def test_block_diffusion_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=7", + "vocab_size=256", + "max_target_length=2048", + "hardware=cpu", + ] + config = pyconfig.initialize(argv) + + self.assertEqual(config.attention_type, "block_diffusion") + self.assertEqual(config.block_diffusion_block_size, 7) + self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0) + self.assertFalse(hasattr(config, "enable_block_diffusion")) + + def test_block_diffusion_invalid_configs_raise(self): + base_overrides = { + "run_name": "test", + "steps": 1, + "attention": "dot_product", + "attention_type": "block_diffusion", + "block_diffusion_block_size": 32, + "vocab_size": 256, + "max_target_length": 2048, + "hardware": "cpu", + } + cases = [ + {"block_diffusion_block_size": 0}, + {"attention": "cudnn_flash_te"}, + {"attention": "flash", "hardware": "gpu"}, + ] + for overrides in cases: + with self.subTest(overrides=overrides): + config = base_overrides | overrides + argv = ["", _BASE_CONFIG_PATH, *(f"{key}={value}" for key, value in config.items())] + with self.assertRaises((ValueError, pydantic.ValidationError)): + pyconfig.initialize(argv) + + def test_default_attention_remains_global(self): + config = pyconfig.initialize(["", _BASE_CONFIG_PATH, "run_name=test", "steps=1"]) + + self.assertEqual(config.attention_type, "global") + @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"}) def test_env_override(self): """Tests that environment variables override YAML values.""" diff --git a/tests/unit/tokamax_splash_attention_mask_test.py b/tests/unit/tokamax_splash_attention_mask_test.py index 590e8046b8..b1fe154029 100644 --- a/tests/unit/tokamax_splash_attention_mask_test.py +++ b/tests/unit/tokamax_splash_attention_mask_test.py @@ -20,6 +20,7 @@ from maxtext.kernels.tokamax_splash_attention import splash_attention_mask from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info +from maxtext.layers.attention_op import BlockCausalMask class TokamaxSplashAttentionMaskTest(absltest.TestCase): @@ -56,6 +57,22 @@ def test_process_causal_mask_keeps_lazy_mask_function(self): self.assertIsNone(mask_info.partial_mask_blocks) np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + def test_process_block_causal_mask_keeps_lazy_mask_function(self): + mask = BlockCausalMask((8, 8), block_size=4) + + mask_info, mask_function = splash_attention_mask_info.process_mask( + mask, + block_shape=(2, 2), + q_seq_shards=1, + kv_seq_shards=1, + return_dynamic_grid=True, + ) + + self.assertIs(mask_function, mask.mask_function) + self.assertIsNotNone(mask_info.q_sequence) + self.assertIsNone(mask_info.partial_mask_blocks) + np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + if __name__ == "__main__": absltest.main()