[maxtext] Add generalized block-diffusion CFT and SFT

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/scoring.py

@@ -0,0 +1,43 @@
+# 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.
+
+"""Target-alignment utilities shared by diffusion training objectives."""
+
+import jax.numpy as jnp
+
+
+def align_logits_to_targets(logits, alignment, positions=None, validity_mask=None):
+  """Aligns model output positions with clean block-diffusion targets."""
+  if alignment == "same_position":
+    return logits
+  if alignment == "shifted":
+    if positions is None:
+      indices = jnp.maximum(jnp.arange(logits.shape[1], dtype=jnp.int32) - 1, 0)
+      return logits[:, indices, :]
+    positions = jnp.asarray(positions, dtype=jnp.int32)
+    if positions.shape != logits.shape[:2]:
+      raise ValueError(f"positions must match logits [batch, length]; got {positions.shape} and {logits.shape[:2]}")
+    if validity_mask is None:
+      validity_mask = jnp.ones_like(positions, dtype=jnp.bool_)
+    sequence_length = logits.shape[1]
+    array_indices = jnp.broadcast_to(jnp.arange(sequence_length, dtype=jnp.int32), positions.shape)
+    row_indices = jnp.broadcast_to(jnp.arange(logits.shape[0], dtype=jnp.int32)[:, None], positions.shape)
+    safe_positions = jnp.clip(positions, 0, sequence_length - 1)
+    updates = jnp.where(validity_mask, array_indices + 1, 0)
+    position_to_index = jnp.zeros_like(positions).at[row_indices, safe_positions].max(updates)
+    previous_positions = jnp.maximum(positions - 1, 0)
+    source_indices = jnp.take_along_axis(position_to_index, jnp.clip(previous_positions, 0, sequence_length - 1), axis=1)
+    source_indices = jnp.maximum(source_indices - 1, 0)
+    return jnp.take_along_axis(logits, source_indices[..., None], axis=1)
+  raise ValueError(f"Unsupported block-diffusion logit alignment: {alignment}")

src/maxtext/input_pipeline/hf_data_processing.py

@@ -42,6 +42,40 @@ def _get_pad_id(tokenizer):
   return pad_id
 
 
+def _get_training_objective_transform(
+    config,
+    *,
+    shift,
+    use_dpo,
+    use_sft,
+    completion_only,
+    packing,
+    pad_id,
+    bos_token_id,
+):
+  """Selects the aligned or next-token target transform for an HF pipeline."""
+  objective = getattr(config, "training_objective", "causal_lm")
+  if objective == "block_diffusion":
+    if packing:
+      raise ValueError("Block-diffusion training requires packing=False.")
+    if use_dpo:
+      raise ValueError("Block-diffusion training is not compatible with DPO.")
+    return input_pipeline_utils.BlockDiffusionCorruption(
+        block_size=config.block_diffusion_block_size,
+        mask_id=config.block_diffusion_mask_id,
+        min_noise=config.block_diffusion_min_noise,
+        axis=1,
+        completion_only=bool(use_sft and completion_only),
+        seed_first_token=config.block_diffusion_canvas_policy == "seed_and_mask",
+        include_seed_in_loss=config.block_diffusion_logit_alignment == "shifted",
+    )
+  if objective != "causal_lm":
+    raise ValueError(f"Unsupported training objective: {objective}")
+  if shift and not use_dpo:
+    return input_pipeline_utils.ShiftData(ignored_ids=[pad_id, bos_token_id], axis=1)
+  return None
+
+
 def vision_sft_preprocessing_pipeline(
     dataset,
     config,
@@ -321,12 +355,15 @@ def preprocessing_pipeline(
   )
   operations = []
   if use_sft:
+    is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion"
     operations.append(
         input_pipeline_utils.SFTPromptMasking(
             text_column_name=data_column_names[0],
             completion_only=sft_train_on_completion_only,
             max_target_length=max_target_length,
             unk_id=pad_id,
+            target_aligned=is_block_diffusion,
+            emit_completion_mask=is_block_diffusion,
         )
     )
     data_column_names = ("inputs", "targets")
@@ -357,8 +394,18 @@ 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:
-    operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1))
+  target_transform = _get_training_objective_transform(
+      config,
+      shift=shift,
+      use_dpo=use_dpo,
+      use_sft=use_sft,
+      completion_only=sft_train_on_completion_only,
+      packing=packing,
+      pad_id=pad_id,
+      bos_token_id=tokenizer.bos_token_id,
+  )
+  if target_transform is not None:
+    operations.append(target_transform)
 
   # Since HuggingFace IterableDataset does not support access through index
   # Indexes generated by dummy_index_sampler is not used.