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

src/maxtext/common/common_types.py

@@ -123,6 +123,7 @@ class AttentionType(enum.Enum):
   MLA = "mla"
   COMPRESSED = "compressed"
   FULL = "full"
+  BLOCK_DIFFUSION = "block_diffusion"
 
 
 class ShardMode(enum.Enum):

src/maxtext/configs/base.yml

@@ -374,12 +374,21 @@ 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 attention (Fast_dLLM / Diffusion-GR2): bidirectional within a block of bd_size tokens, block-causal across blocks. Weight-preserving.
+enable_block_diffusion: false
+bd_size: 32 # Block size for block-diffusion attention (must be > 0 when enabled; max_target_length must be divisible by it)
+mask_id: 151669 # Token id of the [MASK] token used by block-diffusion (must be < vocab_size when enabled)
+# On-policy distillation (OPD): the block-diffusion student generates its own completion by denoising, the frozen teacher scores it, and forward-KL is applied at the committed positions. Requires enable_block_diffusion on the student.
+opd_on_policy: false
+opd_threshold: 0.9 # Confidence threshold for committing a denoised position during the on-policy rollout.
+opd_temperature: 0.0 # Softmax temperature for the on-policy commit step (0.0 = greedy/argmax).
+opd_max_denoise_steps: 0 # Per-block denoise iteration cap for the on-policy rollout; <= 0 means use bd_size.
 attn_logits_soft_cap: 0.0
 final_logits_soft_cap: 0.0
 z_loss_multiplier: 0.0

src/maxtext/configs/post_train/rl.yml

@@ -72,6 +72,17 @@ rl:
   # If null, the entire model is resharded at once.
   reshard_chunk_size: null
 
+  # ====== Block-diffusion RL rollout (Fast_dLLM / Diffusion-GR2, TraceRL-style) ======
+  # If true, the policy is a block-diffusion model: rollouts generate by block-diffusion
+  # denoising (MaxTextDiffusionRollout) and the GRPO/DAPO importance ratio uses the shared
+  # block-diffusion per-token logprob instead of the AR next-token logprob. Requires
+  # enable_block_diffusion=true (top level). When false, RL is autoregressive and unchanged.
+  rl_diffusion_rollout: false
+  # Confidence threshold for committing a denoised position during the diffusion rollout.
+  rl_diffusion_threshold: 0.9
+  # Per-block denoise iteration cap for the diffusion rollout; <= 0 means use bd_size.
+  rl_diffusion_max_denoise_steps: 0
+
 
 # ====== Models ======
 # for MaxText

src/maxtext/configs/types.py

@@ -559,7 +559,7 @@ class Attention(BaseModel):
       "autoselected",
       description="The attention algorithm to use (dot_product, flash, 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(
@@ -606,6 +606,17 @@ 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.")
+  enable_block_diffusion: bool = Field(
+      False,
+      description=(
+          "If True, use block-diffusion attention: bidirectional within a block of `bd_size` tokens "
+          "and block-causal across blocks (Fast_dLLM / Diffusion-GR2). Weight-preserving."
+      ),
+  )
+  bd_size: int = Field(32, description="Block size for block-diffusion attention. Must be > 0 when enabled.")
+  mask_id: int = Field(
+      151669, description="Token id of the [MASK] token used by block-diffusion. Must be < vocab_size when enabled."
+  )
   attn_logits_soft_cap: None | NonNegativeFloat = Field(
       None, description="Soft-cap value for attention logits. None means no cap."
   )
@@ -1478,6 +1489,26 @@ class Distillation(BaseModel):
       "The other parameters will be frozen if this attribute is non empty)",
   )
 
+  # --- On-policy distillation (OPD) for block-diffusion --
+  opd_on_policy: bool = Field(
+      False,
+      description=(
+          "If True, run on-policy block-diffusion distillation: the student generates its own completion "
+          "by block-diffusion denoising, the frozen teacher scores that sequence, and forward-KL is applied "
+          "only at the committed positions. Requires enable_block_diffusion on the student. When False, "
+          "distillation is unchanged (off-policy)."
+      ),
+  )
+  opd_threshold: float = Field(
+      0.9, description="Confidence threshold for committing a denoised position during the on-policy rollout."
+  )
+  opd_temperature: float = Field(
+      0.0, description="Softmax temperature for the on-policy commit step (0.0 = greedy/argmax)."
+  )
+  opd_max_denoise_steps: int = Field(
+      0, description="Per-block denoise iteration cap for the on-policy rollout; <= 0 means use bd_size."
+  )
+
 
 class TrainingLoop(BaseModel):
   """Configuration for the main training loop, evaluation, and reproducibility."""
@@ -2173,6 +2204,21 @@ class RL(BaseModel):
           "If None, no chunking is applied, which may lead to OOM errors if tensors are too large."
       ),
   )
+  rl_diffusion_rollout: bool = Field(
+      False,
+      description=(
+          "If True, the policy is a block-diffusion model: rollouts generate by block-diffusion "
+          "denoising (MaxTextDiffusionRollout) and the GRPO/DAPO importance ratio uses the shared "
+          "block-diffusion (TraceRL-style) per-token logprob instead of the AR next-token logprob. "
+          "Requires enable_block_diffusion=True. When False, RL is autoregressive and unchanged."
+      ),
+  )
+  rl_diffusion_threshold: float = Field(
+      0.9, description="Confidence threshold for committing a denoised position during the RL diffusion rollout."
+  )
+  rl_diffusion_max_denoise_steps: int = Field(
+      0, description="Per-block denoise iteration cap for the RL diffusion rollout; <= 0 means use bd_size."
+  )
 
 
 class RLDataset(BaseModel):
@@ -3510,3 +3556,28 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
             f"must be equal to attention_output_dim ({self.attention_output_dim})"
         )
     return self
+
+  @model_validator(mode="after")
+  def _validate_block_diffusion(self) -> "MaxTextConfig":
+    """Validates block-diffusion attention settings.
+
+    Only enforced when `enable_block_diffusion` is True, because the block-diffusion
+    defaults (e.g. `mask_id`) are intentionally out of range for the default vocab.
+    Defined after `set_derived_and_validate_values`; `vocab_size` and
+    `max_target_length` are never mutated by that validator, so their final values
+    are visible here regardless of validator ordering.
+    """
+    if self.enable_block_diffusion:
+      if not isinstance(self.bd_size, int) or self.bd_size <= 0:
+        raise ValueError("`bd_size` must be an integer > 0 when `enable_block_diffusion` is True.")
+      if self.mask_id >= self.vocab_size:
+        raise ValueError(
+            f"`mask_id` ({self.mask_id}) must be < `vocab_size` ({self.vocab_size}) "
+            "when `enable_block_diffusion` is True."
+        )
+      if self.max_target_length % self.bd_size != 0:
+        raise ValueError(
+            f"`max_target_length` ({self.max_target_length}) must be divisible by `bd_size` "
+            f"({self.bd_size}) when `enable_block_diffusion` is True."
+        )
+    return self

src/maxtext/input_pipeline/hf_data_processing.py

@@ -357,7 +357,19 @@ def preprocessing_pipeline(
       operations.append(input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id))
     operations.append(grain.Batch(batch_size=batch_size, drop_remainder=drop_remainder))
 
-  if shift and not use_dpo:
+  if config.enable_block_diffusion and not use_dpo:
+    # Masked-diffusion (CFT) corruption for block-diffusion SFT (Fast_dLLM / Diffusion-GR2).
+    # Replaces ShiftData: block-diffusion targets are aligned, not next-token shifted.
+    if packing:
+      raise ValueError("enable_block_diffusion is not supported with packing=True; set packing=False.")
+    operations.append(
+        input_pipeline_utils.BlockDiffusionMasking(
+            bd_size=config.bd_size,
+            mask_id=config.mask_id,
+            axis=1,
+        )
+    )
+  elif shift and not use_dpo:
     operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1))
 
   # Since HuggingFace IterableDataset does not support access through index

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -976,6 +976,95 @@ def map(self, element):
     return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis)
 
 
+@dataclasses.dataclass
+class BlockDiffusionMasking(grain.RandomMapTransform):
+  """Masked-diffusion (conversion fine-tuning / CFT) corruption for block-diffusion SFT.
+
+  Turns an autoregressive SFT batch into a masked-diffusion training example
+  (Fast_dLLM / Diffusion-GR2). For each example a random subset of the RESPONSE
+  tokens is corrupted by replacing them with ``mask_id`` in ``inputs``; the clean
+  tokens are kept as ``targets`` ALIGNED (NOT next-token shifted); and
+  ``targets_segmentation`` is set to 1 only at the masked positions. MaxText's loss
+  already gates the per-position cross-entropy on ``targets_segmentation != 0``, so
+  writing the "predict-here" mask into ``targets_segmentation`` makes the loss
+  compute exactly (and only) at the masked positions with NO loss-side change.
+  ``inputs_segmentation`` (the attention mask) is left untouched so attention still
+  sees every real prompt+response token, including the masked ones being denoised.
+
+  This transform REPLACES ``ShiftData`` in the block-diffusion pipeline: block
+  diffusion predicts each masked position from the (bidirectional-within-block)
+  context at the SAME position, so targets must stay aligned rather than shifted.
+
+  Response span: the maskable span is exactly the completion tokens, identified as
+  ``targets_segmentation != 0`` on the incoming (post ``PadOrTrimToMaxLength``,
+  pre-shift) batch. With ``sft_train_on_completion_only=True`` this is the
+  assistant/completion span (prompt and padding have ``targets_segmentation == 0``
+  and are never masked); with completion-only disabled (or non-SFT) it is every
+  real token, matching a LLaDA-style whole-sequence masked-diffusion objective.
+
+  Masking scheme (v1, per-block; matches the Fast_dLLM reference): the sequence is
+  split into contiguous blocks of ``bd_size`` tokens, aligned to position 0 so the
+  blocks coincide with the block-diffusion attention blocks (``position // bd_size``).
+  For each block a ratio ``t ~ U(0, 1)`` is drawn and the per-token mask probability
+  is ``p = (1 - eps) * t + eps``. The ``eps`` floor (1e-3) keeps ``p`` strictly
+  positive so a block is never guaranteed fully unmasked. Every RESPONSE position in
+  the block is then masked independently with probability ``p``. Drawing ``t``
+  per-block (rather than one ratio for the whole sequence) gives a mix of lightly and
+  heavily corrupted blocks within a single example, which is what the block-diffusion
+  decoder sees at inference (block-by-block denoising).
+
+  Randomness uses grain's per-record RNG (``grain.RandomMapTransform``) so each batch
+  draws a fresh, reproducible mask; masks for the examples within a batch are drawn
+  independently. For a fixed ``rng`` the output is deterministic.
+
+  TODO(block-diffusion): add the optional complementary-mask pass from the Fast_dLLM
+  reference (also train on the un-masked response tokens using the complement of the
+  sampled mask) for a lower-variance objective. Skipped in v1.
+  """
+
+  def __init__(self, bd_size, mask_id, eps=1e-3, axis=1):
+    if bd_size <= 0:
+      raise ValueError(f"bd_size must be positive, got {bd_size}")
+    self.bd_size = bd_size
+    self.mask_id = mask_id
+    self.eps = eps
+    self.axis = axis
+
+  def random_map(self, element, rng: np.random.Generator):
+    inputs = element["inputs"]
+    targets = element["targets"]
+    targets_segmentation = element["targets_segmentation"]
+
+    seq_len = inputs.shape[self.axis]
+    if seq_len % self.bd_size != 0:
+      raise ValueError(f"sequence length ({seq_len}) must be divisible by bd_size ({self.bd_size})")
+    num_blocks = seq_len // self.bd_size
+
+    # Maskable span = completion tokens (targets_segmentation != 0). Prompt and
+    # padding carry targets_segmentation == 0 and are therefore never masked.
+    response = targets_segmentation != 0
+
+    # Per-block ratio t ~ U(0,1) -> per-token mask probability p = (1-eps)*t + eps,
+    # broadcast from one value per block to every position in that block.
+    block_shape = list(inputs.shape)
+    block_shape[self.axis] = num_blocks
+    t = rng.random(size=tuple(block_shape))
+    p_block = (1.0 - self.eps) * t + self.eps
+    p_pos = np.repeat(p_block, self.bd_size, axis=self.axis)
+
+    mask = (rng.random(size=inputs.shape) < p_pos) & response
+
+    element["inputs"] = np.where(mask, inputs.dtype.type(self.mask_id), inputs)
+    # targets stay aligned to inputs (clean tokens, no shift). Values outside the
+    # masked positions are ignored by the loss (targets_segmentation == 0 there).
+    element["targets"] = targets
+    # Loss is computed exactly at the masked positions.
+    element["targets_segmentation"] = mask.astype(targets_segmentation.dtype)
+    # inputs_segmentation / *_position are intentionally left unchanged: attention
+    # must still cover every real prompt+response token, including masked ones.
+    return element
+
+
 @dataclasses.dataclass
 class ComputeQwen3OmniPositions(grain.MapTransform):
   """Computes 3D position IDs for Qwen3-Omni multimodal sequences.