[maxtext] Add block-diffusion training primitives

[ethannnnnn]

Motivation

MaxText needs small, model-independent primitives for discrete block-diffusion
training before attention, input-pipeline, trainer, or post-training support is
introduced.

Scope

This PR adds only:

• block corruption under maxtext.diffusion.block_diffusion;
• explicit target alignment for same_position and shifted logits; and
• focused unit tests for both supported contracts.

It does not change configuration, attention, model execution, training loops,
serving, SFT, distillation, or RL.

Design

The corruption API consumes clean tokens, an explicit eligibility mask, a
logical block size, and one of two declared model contracts:

• same_position with an all_masked canvas; or
• shifted with a seed_and_mask canvas.

Each nonempty eligible block is guaranteed to retain at least one supervised
position. Padding remains untouched. Target alignment follows logical token
positions, including reordered physical layouts, instead of assuming that
array index equals sequence position.

Compatibility

There is no behavior change for existing MaxText training or inference paths.
The new package is not activated by any configuration in this PR.

Tests

python -m pytest -q tests/unit/block_diffusion_test.py

The tests cover partial blocks, per-block noise sampling, deterministic RNG,
all-invalid rows, seeded anchors, mask/loss provenance, reordered logical
positions, padding, and invalid contracts and shapes.

Current clean-head result: 13 passed, 5 subtests passed; Pyink, compileall,
and git diff --check also pass.

Checklist

• I have performed a self-review of my code.
• I have added comments or docstrings where the contracts are not self-evident.
• End-to-end testing is not applicable to this default-off, pure primitive layer; focused CPU tests are listed above.
• No standalone documentation page or toctree change is needed; the public design document is linked below.

Follow-ups

Block-causal attention and pre-training integration will be proposed only after
this primitive contract is reviewed and landed. CFT/SFT, OPD, exact replay,
and RL remain separate later changes.

Design document:
https://docs.google.com/document/d/1N7KcCoAIErB2CV9EJ2G1u_mdN-AQgqwNUYSvM0mtqMI/edit

[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.

[hengtaoguo]

Thanks for the feature contribution!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[shralex]

Since block diffusion deviates significantly from the standard autoregressive (AR) pipeline (introducing concepts like canvas_policy, seed_and_mask, and target_alignment), future contributors would greatly benefit from a README.md in the src/maxtext/diffusion/block_diffusion/ directory.

You could include the following in the README or module-level docstring:

Literature Links: Link to the Block Diffusion (https://arxiv.org/abs/2503.09573 ?) paper so readers understand the theoretical backing.

Terminology definitions: Briefly explain the core concepts implemented in this PR:
Canvas: The fixed-length sequence of token slots being refined via denoising.
Logit Alignment (same_position vs shifted): Explain how logits map to targets, as standard causal AR models shift targets by 1, whereas diffusion models might align them directly.
Corruption / Masking: Explain that corrupt_tokens handles the forward noising process by masking valid tokens independently within each bounded bidirectional block.

Adding this context will make the codebase much more approachable for anyone maintaining the diffusion or speculative decoding pipelines in MaxText!

[ethannnnnn]

Thanks for the documentation suggestion. I added src/maxtext/diffusion/block_diffusion/README.md with the block-diffusion paper, canvas and block terminology, corruption semantics, logit alignment, and the two currently supported training contracts.

[ethannnnnn]

@hengtaoguo @shralex Thanks for the approvals! All required checks are green and the PR is pull-ready. GitHub prevents me from merging into protected main. Could one of you please squash-merge it or enable auto-merge?