[maxtext] Add generalized block-causal attention

[ethannnnnn]

Motivation

Block-diffusion models require bidirectional attention within a block and causal visibility across blocks. This should be a declared attention capability, not a model-name branch, and it must leave ordinary causal models unchanged.

Scope

• Add block_diffusion as a default-off attention type with a namespaced block-size setting.
• Implement matching masks for dense attention, TPU Splash/Tokamax, packed sequences, and load-balanced context parallelism.
• Preserve explicitly specialized attention layers and reject unsupported kernel combinations.

Design

For logical query position q, key position k, and block size B, attention is allowed when both positions belong to the same packed segment and q // B >= k // B. This produces a dense rectangle within the current block while preventing visibility into future blocks. The shared attention boundary applies this rule only to default/global attention; local, full, vision, and other explicit layer policies retain their existing behavior.

Logical positions and segment IDs remain the source of truth after context-parallel reordering. Partial final blocks are valid.

Compatibility

attention_type=global remains the default, so causal masks and autoregressive decode are unchanged unless block diffusion is explicitly selected. This PR adds no training objective, mask token, rollout, or model-specific activation.

Extensibility

Later training and generation integrations can consume one attention capability and block-size contract. Supporting another block-diffusion model does not require adding its name to the attention stack.

Tests

• Attention suite: 42 passed; 37 device-dependent cases skipped.
• Configuration suite: 18 passed plus 21 subtests.
• Tokamax mask tests: 3 passed.
• Pyink, YAML lint, and git diff --check passed.

Known limitations

GPU flash/CuDNN combinations that cannot express the block-causal mask are rejected rather than silently falling back to different semantics.

Stack

First PR in this repository stack; no preceding PR dependency.

MaxText block-diffusion design document

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.