Optional block-diffusion support for Qwen3 (attention + masked-diffusion SFT + on-policy distillation + RL rollout)

[ethannnnnn]

Summary

Adds optional block-diffusion (masked-diffusion / discrete-diffusion) support for the Qwen3 family across the full lifecycle: the attention pattern, a masked-diffusion SFT objective, on-policy distillation from an AR teacher, and an RL rollout. Block-diffusion LMs (Fast-dLLM, BD3-LM) attend bidirectionally within a fixed block and causally across blocks, and generate by iterative denoising. Everything is gated (default off); autoregressive paths are unchanged. Conversion is weight-preserving — only the attention pattern, objective, and decode change.

This branch spans four logical pieces; happy to split into a stacked series if that's easier to review.

1. Block-diffusion attention

• New AttentionType.BLOCK_DIFFUSION (common/common_types.py): bidirectional within a block of bd_size, causal across blocks ((i // bd) >= (j // bd)).
• Implemented on both mask paths — a branch in the dense generate_attention_mask and a BlockDiffusionMask(splash_attention_mask._ComputableMask) set as the base mask on the splash path (layers/attention_op.py).
• Config enable_block_diffusion / bd_size / mask_id (configs/types.py, configs/base.yml) with a validator; bd_size threaded into AttentionOp; wired on the Qwen3 attention module (models/qwen3.py), covering dense + MoE decoder layers.

2. Masked-diffusion SFT data transform

• BlockDiffusionMasking (input_pipeline/input_pipeline_utils.py): per-block noising — draw a masking ratio per block, replace response tokens with mask_id in the inputs, keep targets aligned (not shifted), and set the loss mask (targets_segmentation) only at masked positions. Wired into the HF preprocessing pipeline, gated, replacing the AR shift.
• The loss function is unchanged — it already restricts cross-entropy to targets_segmentation != 0.

3. On-policy distillation (OPD)

• trainers/post_train/distillation/diffusion_generate.py: a JAX block-diffusion generation helper (inner while_loop denoise + outer fori_loop across blocks), a shared per-token diffusion log-prob, and the commit rule.
• Wired into train_distill.py: the student generates on-policy under stop_gradient (outside value_and_grad, like the frozen teacher); the existing forward-KL distillation loss is reused, applied only at committed positions. Gated by opd_on_policy.

4. RL rollout (GRPO/DAPO)

• integration/vllm/maxtext_diffusion_rollout.py: MaxTextDiffusionRollout (a plain BaseRollout, since block-diffusion decode is not autoregressive) whose generate() denoises and whose get_per_token_logps() uses the shared diffusion log-prob.
• Wired as the rollout engine in trainers/post_train/rl/train_rl.py, gated by rl_diffusion_rollout. The GRPO loss's new-policy log-prob is supplied via a pluggable hook so the same diffusion log-prob function backs both the rollout (old policy) and the loss (new policy), keeping the importance ratio consistent.

Back-compat / gating

All four pieces gated (enable_block_diffusion, opd_on_policy, rl_diffusion_rollout, default off). No AR path changes.

Testing (CPU)

• Attention: mask logic (dense + splash _ComputableMask) vs a reference, config validator, and a tiny end-to-end Qwen3 forward confirming weight-preservation and that the flag is wired (38 tests incl. regressions).
• SFT transform: only response masked, mask_id set, targets aligned, loss-mask correct, determinism, AR path unchanged (8 tests).
• OPD: stub-model generation loop, committed-position forward-KL, and a jax.grad check confirming generation is stop-gradient (16 tests).
• RL: rollout generation and a consistency test asserting the rollout and the loss use the same log-prob function with identical values (13 tests).

Limitations / follow-ups

• Splash-kernel execution and full training/decode runs need TPU (CPU covers the dense path + all pure-JAX logic).
• Generation is currently cacheless (a full-sequence forward per denoise step); a KV-cached block variant is a follow-up.

Draft pending TPU validation.

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