[maxtext] Add reusable block-diffusion rollout transitions

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):

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
@@ -790,6 +791,11 @@ olmo_apply_ngram_filter: true # mask instances with repetitive n-grams (OLMo-cor
 # Training loop
 steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps
 log_period: 100 # The frequency of Tensorboard flush, gcs metrics writing, and managed profiler metrics updating.
+training_objective: 'causal_lm' # Supported objectives: causal_lm, block_diffusion
+block_diffusion_mask_id: -1 # Tokenizer mask-token id; required for training_objective='block_diffusion'
+block_diffusion_min_noise: 0.001 # Minimum per-block corruption probability for block-diffusion training
+block_diffusion_logit_alignment: 'same_position' # Supported alignments: same_position, shifted
+block_diffusion_canvas_policy: 'all_masked' # Supported canvases: all_masked, seed_and_mask
 
 jax_distributed_initialization_timeout: 300 # This is the default timeout in https://github.com/jax-ml/jax/blob/main/jax/_src/distributed.py
 # Note there are two separate initializations - the jax coordination service (aka jax.distributed.initialize) and the backend (e.g. PjRT), the timeout above refers

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."
   )
@@ -1579,6 +1583,28 @@ class Distillation(BaseModel):
 class TrainingLoop(BaseModel):
   """Configuration for the main training loop, evaluation, and reproducibility."""
 
+  training_objective: Literal["causal_lm", "block_diffusion"] = Field(
+      "causal_lm",
+      description="The token-prediction objective used to prepare targets and compute loss.",
+  )
+  block_diffusion_mask_id: int = Field(
+      -1,
+      description="The tokenizer mask-token id required by the block-diffusion training objective.",
+  )
+  block_diffusion_min_noise: float = Field(
+      1.0e-3,
+      gt=0.0,
+      le=1.0,
+      description="The minimum corruption probability sampled independently for each block.",
+  )
+  block_diffusion_logit_alignment: Literal["same_position", "shifted"] = Field(
+      "same_position",
+      description="How model logits align to clean target-token positions.",
+  )
+  block_diffusion_canvas_policy: Literal["all_masked", "seed_and_mask"] = Field(
+      "all_masked",
+      description="Whether every block is fully maskable or begins with a clean anchor token.",
+  )
   steps: int = Field(
       150_001,
       ge=-1,
@@ -3283,6 +3309,44 @@ 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.training_objective == "block_diffusion":
+      if self.attention_type != AttentionType.BLOCK_DIFFUSION.value:
+        raise ValueError("`training_objective='block_diffusion'` requires `attention_type='block_diffusion'`.")
+      if self.block_diffusion_mask_id < 0 or self.block_diffusion_mask_id >= self.vocab_size:
+        raise ValueError(
+            f"`block_diffusion_mask_id` ({self.block_diffusion_mask_id}) must satisfy "
+            f"0 <= block_diffusion_mask_id < vocab_size ({self.vocab_size})."
+        )
+      if self.packing:
+        raise ValueError("`training_objective='block_diffusion'` requires `packing=False`.")
+      if self.mtp_num_layers > 0:
+        raise ValueError("`training_objective='block_diffusion'` is not compatible with MTP.")
+      if self.num_vocab_tiling > 1:
+        raise ValueError("`training_objective='block_diffusion'` is not compatible with vocabulary tiling.")
+      if self.dataset_type != "hf":
+        raise ValueError("`training_objective='block_diffusion'` currently requires `dataset_type='hf'`.")
+      if self.use_dpo:
+        raise ValueError("`training_objective='block_diffusion'` is not compatible with DPO.")
+      if self.use_multimodal or self.use_audio:
+        raise ValueError("`training_objective='block_diffusion'` currently supports text-only training.")
+      valid_model_contracts = {
+          ("same_position", "all_masked"),
+          ("shifted", "seed_and_mask"),
+      }
+      model_contract = (self.block_diffusion_logit_alignment, self.block_diffusion_canvas_policy)
+      if model_contract not in valid_model_contracts:
+        raise ValueError(
+            "Block-diffusion training supports only `same_position/all_masked` or `shifted/seed_and_mask`; "
+            f"received `{model_contract[0]}/{model_contract[1]}`."
+        )
     if self.quantize_kvcache and not self.kv_quant_axis:
       raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.")
     if (

src/maxtext/diffusion/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared block-diffusion model semantics."""

src/maxtext/diffusion/denoise.py

@@ -0,0 +1,208 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Model-independent block-diffusion rollout state transitions."""
+
+from collections.abc import Callable
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def _validate_shapes(initial_tokens, positions, validity_mask, completion_mask):
+  """Checks that all token-level rollout arrays share a batch-major shape."""
+  expected_shape = tuple(initial_tokens.shape)
+  if len(expected_shape) != 2:
+    raise ValueError(f"initial_tokens must have shape [batch, length]; received {expected_shape}")
+  for name, value in (
+      ("positions", positions),
+      ("validity_mask", validity_mask),
+      ("completion_mask", completion_mask),
+  ):
+    if tuple(value.shape) != expected_shape:
+      raise ValueError(f"{name} must match initial_tokens shape; received {tuple(value.shape)} and {expected_shape}")
+
+
+def _concrete_numpy(value):
+  if isinstance(value, jax.core.Tracer):
+    return None
+  if isinstance(value, jax.Array) and not value.is_fully_addressable:
+    return None
+  return np.asarray(value)
+
+
+def _validate_logical_positions(positions, validity_mask, completion_mask, *, shifted_seed):
+  """Checks logical sequence invariants on eager, host-addressable inputs."""
+  concrete_positions = _concrete_numpy(positions)
+  concrete_validity = _concrete_numpy(validity_mask)
+  concrete_completion = _concrete_numpy(completion_mask)
+  if concrete_positions is None or concrete_validity is None or concrete_completion is None:
+    return
+  concrete_validity = np.asarray(concrete_validity, dtype=bool)
+  concrete_completion = np.asarray(concrete_completion, dtype=bool) & concrete_validity
+  sequence_length = concrete_positions.shape[1]
+  for row in range(concrete_positions.shape[0]):
+    valid_positions = np.asarray(concrete_positions[row, concrete_validity[row]])
+    expected_positions = np.arange(valid_positions.size, dtype=valid_positions.dtype)
+    if valid_positions.size and not np.array_equal(np.sort(valid_positions), expected_positions):
+      raise ValueError(
+          "valid logical positions must be unique and contiguous from zero within the physical sequence length"
+      )
+    if np.any(valid_positions < 0) or np.any(valid_positions >= sequence_length):
+      raise ValueError("valid logical positions must be nonnegative and smaller than the physical sequence length")
+    if shifted_seed:
+      position_zero = concrete_validity[row] & (concrete_positions[row] == 0)
+      if np.count_nonzero(position_zero) != 1 or np.any(concrete_completion[row] & position_zero):
+        raise ValueError("shifted seed-and-mask rollout requires exactly one prompt token at logical position zero")
+
+
+def validate_completion_suffix(positions, validity_mask, completion_mask, *, shifted_seed=False):
+  """Validates the initial OPD rollout scope when eager values are available.
+
+  The first integration intentionally supports one prompt followed by one
+  completion per sequence. This prevents clean future turns from appearing in
+  the bidirectional attention block of an earlier generated completion.
+  """
+  _validate_logical_positions(positions, validity_mask, completion_mask, shifted_seed=shifted_seed)
+  concrete_positions = _concrete_numpy(positions)
+  concrete_validity = _concrete_numpy(validity_mask)
+  concrete_completion = _concrete_numpy(completion_mask)
+  if concrete_positions is None or concrete_validity is None or concrete_completion is None:
+    return
+  concrete_validity = np.asarray(concrete_validity, dtype=bool)
+  concrete_completion = np.asarray(concrete_completion, dtype=bool) & concrete_validity
+  for row in range(concrete_positions.shape[0]):
+    valid_indices = np.flatnonzero(concrete_validity[row])
+    if valid_indices.size == 0:
+      continue
+    ordered_indices = valid_indices[np.argsort(concrete_positions[row, valid_indices])]
+    ordered_completion = concrete_completion[row, ordered_indices]
+    completion_indices = np.flatnonzero(ordered_completion)
+    if completion_indices.size and not np.all(ordered_completion[completion_indices[0] :]):
+      raise ValueError("diffusion OPD rollout requires completion_mask to be a contiguous suffix")
+
+
+def low_confidence_generate(
+    logits_fn: Callable[[jax.Array], jax.Array],
+    initial_tokens: jax.Array,
+    positions: jax.Array,
+    validity_mask: jax.Array,
+    completion_mask: jax.Array,
+    *,
+    block_size: int,
+    mask_id: int,
+    logit_alignment: str,
+    canvas_policy: str,
+    confidence_threshold: float = 0.9,
+    temperature: float = 1.0,
+    max_denoise_steps: int | None = None,
+) -> jax.Array:
+  """Generates a completion block by block with confidence-based commits.
+
+  ``logits_fn`` must return target-aligned logits for the current token canvas.
+  Each denoising step commits every token at or above the confidence threshold;
+  if a row has no such token, its highest-confidence unresolved token is forced
+  to commit. This guarantees progress for heterogeneous and partial blocks.
+  """
+  _validate_shapes(initial_tokens, positions, validity_mask, completion_mask)
+  valid_contracts = {("same_position", "all_masked"), ("shifted", "seed_and_mask")}
+  if (logit_alignment, canvas_policy) not in valid_contracts:
+    raise ValueError(
+        "rollout supports only same_position/all_masked or shifted/seed_and_mask; "
+        f"received {logit_alignment}/{canvas_policy}"
+    )
+  if block_size <= 0:
+    raise ValueError(f"block_size must be positive; received {block_size}")
+  if not 0.0 <= confidence_threshold <= 1.0:
+    raise ValueError(f"confidence_threshold must be in [0, 1]; received {confidence_threshold}")
+  if temperature <= 0.0:
+    raise ValueError(f"temperature must be positive; received {temperature}")
+  if max_denoise_steps is None:
+    max_denoise_steps = block_size
+  if max_denoise_steps < block_size:
+    raise ValueError(
+        f"max_denoise_steps must be at least block_size ({block_size}) to guarantee completion; "
+        f"received {max_denoise_steps}"
+    )
+
+  shifted_seed = logit_alignment == "shifted"
+  validate_completion_suffix(positions, validity_mask, completion_mask, shifted_seed=shifted_seed)
+  validity_mask = jnp.asarray(validity_mask, dtype=jnp.bool_)
+  completion_mask = jnp.asarray(completion_mask, dtype=jnp.bool_) & validity_mask
+  positions = jnp.asarray(positions, dtype=jnp.int32)
+  canvas = jnp.where(completion_mask, jnp.asarray(mask_id, initial_tokens.dtype), initial_tokens)
+  block_ids = positions // block_size
+  num_blocks = (initial_tokens.shape[1] + block_size - 1) // block_size
+
+  def propose(current_canvas):
+    logits = logits_fn(current_canvas)
+    expected_prefix = tuple(current_canvas.shape)
+    if len(logits.shape) != 3 or tuple(logits.shape[:2]) != expected_prefix:
+      raise ValueError(
+          "logits_fn must return [batch, length, vocab] target-aligned logits; "
+          f"received {tuple(logits.shape)} for canvas {expected_prefix}"
+      )
+    vocab_size = logits.shape[-1]
+    if vocab_size < 2:
+      raise ValueError("logits_fn must expose at least two vocabulary entries so the mask token can be excluded")
+    if not 0 <= mask_id < vocab_size:
+      raise ValueError(f"mask_id must satisfy 0 <= mask_id < vocab_size ({vocab_size}); received {mask_id}")
+    scaled_logits = jnp.asarray(logits, dtype=jnp.float32) / temperature
+    scaled_logits = scaled_logits.at[..., mask_id].set(-jnp.inf)
+    probabilities = jax.nn.softmax(scaled_logits, axis=-1)
+    return jnp.argmax(scaled_logits, axis=-1).astype(initial_tokens.dtype), jnp.max(probabilities, axis=-1)
+
+  def generate_block(block_id, current_canvas):
+    in_block = completion_mask & (block_ids == block_id)
+
+    def run_active_block(active_canvas):
+      anchors = in_block & (positions % block_size == 0) if shifted_seed else jnp.zeros_like(in_block)
+
+      def generate_anchors(anchor_canvas):
+        anchor_tokens, _ = propose(anchor_canvas)
+        return jnp.where(anchors, anchor_tokens, anchor_canvas)
+
+      if shifted_seed:
+        active_canvas = jax.lax.cond(jnp.any(anchors), generate_anchors, lambda value: value, active_canvas)
+      unresolved = in_block & ~anchors
+
+      def continue_denoising(state):
+        step, _, remaining = state
+        return (step < max_denoise_steps) & jnp.any(remaining)
+
+      def denoise_step(state):
+        step, step_canvas, remaining = state
+        proposed_tokens, confidence = propose(step_canvas)
+        commits = remaining & (confidence >= confidence_threshold)
+        row_needs_fallback = jnp.any(remaining, axis=1) & ~jnp.any(commits, axis=1)
+        fallback_confidence = jnp.max(jnp.where(remaining, confidence, -jnp.inf), axis=1)
+        tied_for_fallback = remaining & (confidence == fallback_confidence[:, None])
+        fallback_positions = jnp.where(tied_for_fallback, positions, positions.shape[1])
+        fallback_indices = jnp.argmin(fallback_positions, axis=1)
+        fallback = jax.nn.one_hot(fallback_indices, remaining.shape[1], dtype=jnp.bool_)
+        commits |= fallback & row_needs_fallback[:, None]
+        step_canvas = jnp.where(commits, proposed_tokens, step_canvas)
+        return step + 1, step_canvas, remaining & ~commits
+
+      _, active_canvas, _ = jax.lax.while_loop(
+          continue_denoising,
+          denoise_step,
+          (jnp.asarray(0, dtype=jnp.int32), active_canvas, unresolved),
+      )
+      return active_canvas
+
+    return jax.lax.cond(jnp.any(in_block), run_active_block, lambda value: value, current_canvas)
+
+  return jax.lax.fori_loop(0, num_blocks, generate_block, canvas)